cont : avoid copy_state

Address review comments
Rebase and address review comments
2026-08-25 15:07:34 +02:00 · 2026-08-06 10:19:52 +03:00 · 2026-08-06 00:27:46 +05:30 · 2026-08-06 00:27:46 +05:30 · 2026-08-06 00:27:46 +05:30 · 2026-08-06 00:27:46 +05:30
188 changed files with 3319 additions and 33594 deletions
@@ -57,6 +57,7 @@ COPY --from=web /app/tools/ui/dist tools/ui/dist
 RUN HIPCXX="$(hipconfig -l)/clang" HIP_PATH="$(hipconfig -R)" \
    cmake -S . -B build \
        -DGGML_HIP=ON \
+        -DGGML_HIP_ROCWMMA_FATTN=ON \
        -DAMDGPU_TARGETS="$ROCM_DOCKER_ARCH" \
        -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON \
        -DCMAKE_BUILD_TYPE=Release -DLLAMA_BUILD_TESTS=OFF \
@@ -99,6 +99,7 @@ jobs:
        run: |
          cmake -B build -S . \
            -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang" \
+            -DGGML_HIP_ROCWMMA_FATTN=ON \
            -DGPU_TARGETS="gfx1030" \
            -DGGML_HIP=ON
          cmake --build build --config Release -j $(nproc)
@@ -150,6 +150,7 @@ jobs:
            -DLLAMA_BUILD_BORINGSSL=ON `
            -DROCM_DIR="${env:HIP_PATH}" `
            -DGGML_HIP=ON `
+            -DGGML_HIP_ROCWMMA_FATTN=ON `
            -DGPU_TARGETS="gfx1100"  `
            -DGGML_RPC=ON
          cmake --build build -j ${env:NUMBER_OF_PROCESSORS}
@@ -71,26 +71,6 @@ jobs:
          nvidia-smi
          GG_BUILD_CUDA=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp

-  gpu-rocm:
-    runs-on: [self-hosted, Linux, AMD]
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: Test
-        id: ggml-ci
-        # HIP_LAUNCH_BLOCKING=1: workaround for an async-execution correctness
-        # issue on integrated RDNA3.5 (gfx1151) where batched inference returns
-        # incorrect output (perplexity ~88 vs ~9.4). Serializing kernel launches
-        # restores correctness. Remove once the underlying ROCm/HIP issue is fixed.
-        env:
-          HIP_LAUNCH_BLOCKING: "1"
-        run: |
-          rocminfo
-          GG_BUILD_ROCM=1 GG_BUILD_AMDGPU_TARGETS=gfx1151 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
-
  gpu-vulkan-nvidia-cm:
    runs-on: [self-hosted, Linux, NVIDIA]

@@ -1229,6 +1229,7 @@ jobs:
            -DGPU_TARGETS="${{ matrix.gpu_targets }}" \
            -DGGML_HIP=ON \
            -DHIP_PLATFORM=amd \
+            -DGGML_HIP_ROCWMMA_FATTN=ON \
            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
            ${{ env.CMAKE_ARGS }}
          cmake --build build --config Release -j $(nproc)
@@ -1352,6 +1353,7 @@ jobs:
            -DGGML_NATIVE=OFF `
            -DGGML_CPU=OFF `
            -DGPU_TARGETS="${{ matrix.gpu_targets }}" `
+            -DGGML_HIP_ROCWMMA_FATTN=ON `
            -DGGML_HIP=ON `
            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} `
            -DLLAMA_BUILD_BORINGSSL=ON
@@ -10,9 +10,6 @@
 # # with CUDA support
 # GG_BUILD_CUDA=1 bash ./ci/run.sh ./tmp/results ./tmp/mnt
 #
-# # with ROCm support
-# GG_BUILD_ROCM=1 GG_BUILD_AMDGPU_TARGETS=gfx1151 bash ./ci/run.sh ./tmp/results ./tmp/mnt
-#
 # # with SYCL support
 # GG_BUILD_SYCL=1 bash ./ci/run.sh ./tmp/results ./tmp/mnt
 #
@@ -92,7 +89,7 @@ if [ ! -z ${GG_BUILD_CUDA} ]; then
 fi

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

 function gg_check_build_requirements {
    if ! command -v git &> /dev/null; then
-        gg_printf 'git not found, please install\n'
-        exit 1
+        gg_printf 'git not found, please install'
    fi

    if ! command -v git-lfs &> /dev/null; then
-        gg_printf 'git-lfs not found, please install\n'
-        exit 1
-    fi
-
-    if ! git config --get filter.lfs.clean &> /dev/null; then
-        gg_printf 'git-lfs not initialized, please run `git lfs install`\n'
-        exit 1
+        gg_printf 'git-lfs not found, please install'
    fi

    if ! command -v wget &> /dev/null; then
-        gg_printf 'wget not found, please install\n'
-        exit 1
+        gg_printf 'wget not found, please install'
    fi

    if ! command -v python3 &> /dev/null; then
-        gg_printf 'python3 not found, please install\n'
-        exit 1
+        gg_printf 'python3 not found, please install'
    fi

    if ! command -v pip3 &> /dev/null; then
-        gg_printf 'pip3 not found, please install\n'
-        exit 1
+        gg_printf 'pip3 not found, please install'
    fi

    if ! python3 -m ensurepip --help &> /dev/null; then
-        gg_printf 'ensurepip not found, please install python3-venv package\n'
-        exit 1
+        gg_printf 'ensurepip not found, please install python3-venv package'
    fi

    if ! command -v cmake &> /dev/null; then
-        gg_printf 'cmake not found, please install\n'
-        exit 1
+        gg_printf 'cmake not found, please install'
    fi

    if ! command -v ccache &> /dev/null; then
-        gg_printf 'ccache not found, please consider installing for faster builds\n'
+        gg_printf 'ccache not found, please consider installing for faster builds'
    fi

    if ! command -v ctest &> /dev/null; then
-        gg_printf 'ctest not found, please install\n'
-        exit 1
+        gg_printf 'ctest not found, please install'
    fi
 }

@@ -3308,16 +3308,6 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            params.server_tools = parse_csv_row(value);
        }
    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TOOLS"));
-    add_opt(common_arg(
-        {"--tools-runtime"}, "OPTION",
-        "experimental: run tools in a separate runtime environment (default: none, use host environment)\n"
-        "available options:\n"
-        "  'docker:<image>': spin up a new Docker container and reuse it for all invocations, clean up on server exit\n"
-        "  'docker-container:<id>': use an existing Docker container by ID, won't stop on server exit\n",
-        [](common_params & params, const std::string & value) {
-            params.server_tools_runtime = value;
-        }
-    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TOOLS_RUNTIME"));
    add_opt(common_arg(
        {"--mcp-servers-config"}, "PATH",
        "experimental: path to JSON file with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n"
@@ -1639,6 +1639,7 @@ struct llama_context_params common_context_params_to_llama(const common_params &
    cparams.n_seq_max         = params.n_parallel;
    cparams.n_rs_seq          = params.speculative.need_n_rs_seq();
    cparams.n_outputs_max     = std::max(params.n_outputs_max, 0);
+    cparams.n_outputs_max_per_seq = std::max(params.n_outputs_max_per_seq, 0);
    cparams.n_batch           = params.n_batch;
    cparams.n_ubatch          = params.n_ubatch;
    cparams.n_threads         = params.cpuparams.n_threads;
@@ -447,6 +447,7 @@ struct common_params {
    int32_t n_parallel            =     1; // number of parallel sequences to decode
    int32_t n_sequences           =     1; // number of sequences to decode
    int32_t n_outputs_max         =     0; // max outputs in a batch (0 = n_batch)
+    int32_t n_outputs_max_per_seq =     1; // max outputs per sequence
    int32_t grp_attn_n            =     1; // group-attention factor
    int32_t grp_attn_w            =   512; // group-attention width
    int32_t n_print               =    -1; // print token count every n tokens (-1 = disabled)
@@ -655,7 +656,6 @@ struct common_params {

    // enable built-in tools
    std::vector<std::string> server_tools;
-    std::string server_tools_runtime;

    // MCP server configs (Cursor-compatible JSON)
    std::string mcp_servers_config;   // path to JSON file with MCP server definitions
@@ -116,6 +116,8 @@ static llama_sampler_i llama_sampler_llg_i = {
    /* .backend_accept    = */ NULL,
    /* .backend_apply     = */ NULL,
    /* .backend_set_input = */ NULL,
+    /* .backend_reset     = */ NULL,
+    /* .copy_state        = */ NULL,
 };

 static size_t llama_sampler_llg_tokenize_fn(const void * user_data, const uint8_t * bytes, size_t bytes_len,
@@ -217,6 +217,7 @@ static struct llama_sampler_i common_reasoning_budget_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ nullptr,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl) {
@@ -518,6 +518,26 @@ struct common_sampler * common_sampler_clone(common_sampler * gsmpl) {
    };
 }

+void common_sampler_copy(const common_sampler * src, common_sampler * dst) {
+    if (!src || !dst || src == dst) {
+        return;
+    }
+
+    GGML_ASSERT((src->grmr == nullptr) == (dst->grmr == nullptr));
+    GGML_ASSERT((src->rbudget == nullptr) == (dst->rbudget == nullptr));
+
+    llama_sampler_copy(src->grmr,    dst->grmr);
+    llama_sampler_copy(src->rbudget, dst->rbudget);
+    llama_sampler_copy(src->chain,   dst->chain);
+
+    dst->params     = src->params;
+    dst->prev       = src->prev;
+    dst->cur        = src->cur;
+    dst->cur_p      = src->cur_p;
+    dst->cur_p.data = src->cur_p.data ? dst->cur.data() : nullptr; // re-point to dst's buffer
+    dst->t_total_us = src->t_total_us;
+}
+
 void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl) {
    // TODO: measure grammar performance

@@ -47,6 +47,7 @@ void common_sampler_free(struct common_sampler * gsmpl);
 void                    common_sampler_accept(struct common_sampler * gsmpl, llama_token token, bool is_generated);
 void                    common_sampler_reset (struct common_sampler * gsmpl);
 struct common_sampler * common_sampler_clone (struct common_sampler * gsmpl);
+void                    common_sampler_copy  (const struct common_sampler * src, struct common_sampler * dst);

 // arguments can be nullptr to skip printing
 void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl);
@@ -2292,6 +2292,7 @@ common_params common_base_params_to_speculative(const common_params & params) {
    result.cache_type_k  = params_spec.cache_type_k;
    result.cache_type_v  = params_spec.cache_type_v;
    result.n_outputs_max = params.n_parallel;
+    result.n_outputs_max_per_seq = 1;

    return result;
 }
@@ -2377,6 +2378,17 @@ common_speculative_init_result_ptr common_speculative_init_from_params(common_pa
    return std::make_unique<common_speculative_init_result>(params, model_tgt, ctx_tgt);
 }

+common_speculative_output_limits common_speculative_get_output_limits(
+        int32_t n_batch, int32_t n_parallel, int32_t n_draft) {
+    const int64_t per_seq = 1 + (int64_t) std::max(0, n_draft);
+    const int64_t total   = (int64_t) n_parallel * per_seq;
+
+    return {
+        /* .total   = */ (int32_t) std::min<int64_t>(n_batch, total),
+        /* .per_seq = */ (int32_t) std::min<int64_t>(n_batch, per_seq),
+    };
+}
+
 // initialization of the speculative decoding system
 //
 common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq) {
@@ -25,6 +25,15 @@ int32_t common_speculative_n_max(const common_params_speculative * spec);

 common_params common_base_params_to_speculative(const common_params & params);

+struct common_speculative_output_limits {
+    int32_t total;
+    int32_t per_seq;
+};
+
+// return the output limits needed for speculative decoding
+common_speculative_output_limits common_speculative_get_output_limits(
+        int32_t n_batch, int32_t n_parallel, int32_t n_draft);
+
 common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq);

 void common_speculative_free(common_speculative * spec);
@@ -70,7 +70,6 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "Exaone4ForCausalLM": "exaone",
    "ExaoneForCausalLM": "exaone",
    "ExaoneMoEForCausalLM": "exaone",
-    "ExaoneMoeForCausalLM": "exaone",
    "FalconForCausalLM": "falcon",
    "FalconH1ForCausalLM": "falcon_h1",
    "FalconMambaForCausalLM": "mamba",
@@ -533,13 +533,6 @@ class DeepseekV4Model(TextModel):
        for key, value in raw_hparams.items():
            self.hparams.setdefault(key, value)

-        # workaround for special rope_parameters (main/compress) in transformers 5.x
-        if self.rope_parameters.get("full_attention", self.rope_parameters).get("rope_type") is None:
-            if (rope_scaling := raw_hparams.get("rope_scaling")) is not None:
-                if "rope_type" not in rope_scaling and (rope_type := rope_scaling.get("type")) is not None:
-                    rope_scaling["rope_type"] = rope_type
-                self.rope_parameters.update(**rope_scaling)
-
        self.block_count = self.hparams["num_hidden_layers"]
        if self.mtp_only:
            self.block_count += self.hparams.get("num_nextn_predict_layers", 0)
@@ -123,9 +123,7 @@ class Exaone4Model(TextModel):
                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))


-# note: transformers >= 5.1 renamed the class to "ExaoneMoeForCausalLM" (lowercase 'e'),
-#       so accept both spellings - LG AI have updated the configs of already-released models
-@ModelBase.register("ExaoneMoEForCausalLM", "ExaoneMoeForCausalLM")
+@ModelBase.register("ExaoneMoEForCausalLM")
 class ExaoneMoEModel(Exaone4Model):
    model_arch = gguf.MODEL_ARCH.EXAONE_MOE

@@ -449,8 +449,6 @@ Or
 use 1 SYCL GPUs: [0] with Max compute units:512
 ```

-User can use the device management in [docs/multi-gpu.md](https://github.com/ggml-org/llama.cpp/blob/master/docs/multi-gpu.md), like parameter `--device SYCL0,SYCL1` to assign one or more devices.
-
 ## Windows

 ### Install GPU driver
@@ -765,7 +763,6 @@ Or
 use 1 SYCL GPUs: [0] with Max compute units:512
 ```

-User can use the device management in [docs/multi-gpu.md](https://github.com/ggml-org/llama.cpp/blob/master/docs/multi-gpu.md), like parameter `--device SYCL0,SYCL1` to assign one or more devices.

 ## Environment Variable

@@ -898,45 +895,6 @@ Pass these via `CXXFLAGS` or add a one-off `#define` to enable a flag on the spo
    set UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS=1
  ```

- When I set `SYCL_CACHE_PERSISTENT=1` in running time, I meet crash.
-
-  `SYCL_CACHE_PERSISTENT=1` is not recommended by llama.cpp SYCL backend.
-  When cache is enabled, SYCL runtime will try to cache and reuse JIT-compiled binaries.
-
-  We find some AI will tell user this cmd to speed up SYCL backend. It only speeds up the startup to skip the JIT process, instead of running speed.
-
-  It will bring negative impact when the SYCL binary file is changed frequently in your running environment. The new & old codes mix will lead to crash.
-
-  Compare to the benefit, it has brought more failed cases.
-  If you are not familiar with the SYCL compiler principle of JIT and AOT, please don't use it.
-
-  To restore, you need to remove the local cache: `~/.cache/libsycl_cache/` and execute `unset SYCL_CACHE_PERSISTENT` in running time.
-
- How to use iGPU and dGPU in same time?
-
-  1. Detect the devices in your running time.
-  ```
-  source /opt/intel/oneapi/setvars.sh
-  ./build/bin/llama-server --list-devices
-
-  or
-  ./build/bin/llama-cli --list-devices
-  ./build/bin/llama-bench --list-devices
-  ./build/bin/llama-completion --list-devices
-
-  Available devices:
-    SYCL0: Intel(R) Arc(TM) A770 Graphics (15473 MiB, 15473 MiB free)
-    SYCL1: Intel(R) UHD Graphics 770 (59675 MiB, 44986 MiB free)
-  ```
-
-  The dGPU will be in the head of this list and iGPU will be the end.
-  If not all GPUs are listed, please check the env var: ONEAPI_DEVICE_SELECTOR and unset it.
-
-  2. Set the iGPU and dGPU
-
-  Set the iGPU and dGPU by `./build/bin/llama-server --device SYCL0,SYCL1,SYCLxxx`.
-
-
 ### **GitHub contribution**:
 Please add the `[SYCL]` prefix/tag in issues/PRs titles to help the SYCL contributors to check/address them without delay.

@@ -15,7 +15,7 @@ Legend:
 | Operation | BLAS | CANN | CPU | CUDA | ET | MTL | OpenCL | SYCL | Vulkan | WebGPU | ZenDNN | zDNN |
 |-----------|------|------|------|------|------|------|------|------|------|------|------|------|
 |                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              ACC | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                              ACC | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ |
 |                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             ADD1 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                           ADD_ID | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -41,9 +41,9 @@ Legend:
 |                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              DIV | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                     DSV4_HC_COMB | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                     DSV4_HC_POST | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                      DSV4_HC_PRE | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     DSV4_HC_COMB | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                     DSV4_HC_POST | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                      DSV4_HC_PRE | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                              DUP | ❌ | ✅ | ✅ | 🟡 | ❌ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              ELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              EXP | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -59,7 +59,7 @@ Legend:
 |                             GELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         GELU_ERF | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                       GELU_QUICK | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                         GET_ROWS | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ❌ | ❌ |
+|                         GET_ROWS | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                    GET_ROWS_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ |
 |                       GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -68,7 +68,7 @@ Legend:
 |                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                LIGHTNING_INDEXER | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                LIGHTNING_INDEXER | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -202,6 +202,12 @@ Example Video:

 If a draft model is combined with a draftless decoding the draftless decoding has higher precedence.

+### Backend Sampling
+
+Use `--backend-sampling` to run supported target-model samplers on the model backend. Draft-model sampling uses the backend by default and can be controlled with `--spec-draft-backend-sampling` and `--no-spec-draft-backend-sampling`.
+
+Unsupported samplers and device layouts fall back to CPU sampling. Tensor split mode does not support backend sampling. A fixed seed produces repeatable random draws, but stochastic CPU and backend sampling can still select different tokens because floating-point operations can differ between implementations and devices. Use greedy sampling when exact output matching is required.
+
 ### General Speculative Parameters

 ```
@@ -3,9 +3,11 @@
 #include "common.h"
 #include "ngram-cache.h"
 #include "sampling.h"
+#include "speculative.h"
 #include "log.h"
 #include "llama.h"

+#include <algorithm>
 #include <clocale>
 #include <cstdint>
 #include <cstdio>
@@ -27,6 +29,10 @@ int main(int argc, char ** argv){
    // max. number of additional tokens to draft if match is found
    const int n_draft = params.speculative.draft.n_max;

+    const auto output_limits = common_speculative_get_output_limits(params.n_batch, params.n_parallel, n_draft);
+    params.n_outputs_max = output_limits.total;
+    params.n_outputs_max_per_seq = output_limits.per_seq;
+
    // init llama.cpp
    llama_backend_init();
    llama_numa_init(params.numa);
@@ -47,7 +47,6 @@ CMD_ARGS+=("../../convert_hf_to_gguf.py" "--verbose")
 CMD_ARGS+=("${MODEL_PATH}")
 CMD_ARGS+=("--outfile" "${CONVERTED_MODEL}")
 CMD_ARGS+=("--outtype" "${TYPE}")
-CMD_ARGS+=("--model-name" "${MODEL_NAME}")
 [[ -n "$METADATA_OVERRIDE" ]] && CMD_ARGS+=("--metadata" "${METADATA_OVERRIDE}")
 [[ -n "$MMPROJ" ]] && CMD_ARGS+=("${MMPROJ}")

@@ -31,7 +31,6 @@ python ../../convert_hf_to_gguf.py --verbose \
    ${EMBEDDING_MODEL_PATH} \
    --outfile ${CONVERTED_MODEL} \
    --outtype ${TYPE} \
-    --model-name ${MODEL_NAME} \
    ${SENTENCE_TRANSFORMERS}

 echo ""
@@ -5,6 +5,7 @@
 #include "log.h"
 #include "llama.h"

+#include <algorithm>
 #include <clocale>
 #include <cstdio>
 #include <cstring>
@@ -29,6 +30,11 @@ int main(int argc, char ** argv) {
        return 1;
    }

+    const auto output_limits = common_speculative_get_output_limits(
+            params.n_batch, params.n_parallel, common_speculative_n_max(&params.speculative));
+    params.n_outputs_max = output_limits.total;
+    params.n_outputs_max_per_seq = output_limits.per_seq;
+
    // init llama.cpp
    llama_backend_init();
    llama_numa_init(params.numa);
@@ -55,6 +61,9 @@ int main(int argc, char ** argv) {

        auto params_dft = params;

+        params_dft.n_outputs_max = params.n_parallel;
+        params_dft.n_outputs_max_per_seq = 1;
+
        params_dft.devices      = params_spec.devices;
        params_dft.model        = params_spec.mparams;
        params_dft.n_gpu_layers = params_spec.n_gpu_layers;
@@ -1,6 +1,7 @@
 #include "arg.h"
 #include "common.h"
 #include "sampling.h"
+#include "speculative.h"
 #include "log.h"
 #include "llama.h"

@@ -57,6 +58,11 @@ int main(int argc, char ** argv) {
    // max number of parallel drafting sequences (i.e. tree branches)
    const int n_seq_dft = params.n_parallel;

+    const auto output_limits = common_speculative_get_output_limits(
+            params.n_batch, params.n_parallel, params.speculative.draft.n_max);
+    params.n_outputs_max = output_limits.total;
+    params.n_outputs_max_per_seq = output_limits.per_seq;
+
    // probability threshold for splitting a draft branch (only for n_seq_dft > 1)
    const float p_draft_split = params.speculative.draft.p_split;

@@ -83,6 +89,8 @@ int main(int argc, char ** argv) {
    params.devices = params.speculative.draft.devices;
    params.model = params.speculative.draft.mparams;
    params.n_gpu_layers = params.speculative.draft.n_gpu_layers;
+    params.n_outputs_max = params.n_parallel;
+    params.n_outputs_max_per_seq = 1;
    if (params.speculative.draft.cpuparams.n_threads > 0) {
        params.cpuparams.n_threads = params.speculative.draft.cpuparams.n_threads;
    }
@@ -12,7 +12,6 @@ This script processes files with specified options.

 Options:
  -h, --help    Display this help message and exit.
-  -d, --device  <value>    Set SYCL devices (default: SYCL0).
  -c, --context <value>    Set context length. Bigger need more memory.
  -p, --promote <value>    Prompt to start generation with.
  -m, --model   <value>    Full model file path.
@@ -42,16 +41,10 @@ MODEL_FILE=../models/Qwen3.5-4B-Q4_0.gguf
 NGL=99
 CONTEXT=4096
 GGML_SYCL_DEVICE=-1
-SYCL_DEVICES="SYCL0"
 SPLIT_MODE=layer
 LOG_VERBOSE=3
 while [[ $# -gt 0 ]]; do
    case "$1" in
-        -d|--device)
-            SYCL_DEVICES="$2"
-            shift
-            shift
-            ;;
        -c|--context)
            CONTEXT=$2
            # Shift twice to consume both the option flag and its value
@@ -102,6 +95,8 @@ while [[ $# -gt 0 ]]; do
    esac
 done

+
+
 source /opt/intel/oneapi/setvars.sh

 #export GGML_SYCL_DEBUG=1
@@ -112,19 +107,17 @@ source /opt/intel/oneapi/setvars.sh
 export UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS=1
 echo "UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS=${UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS}"

-echo "ONEAPI_DEVICE_SELECTOR=${ONEAPI_DEVICE_SELECTOR}"
-
-
 if [ $GGML_SYCL_DEVICE -ne -1 ]; then
    echo "Use $GGML_SYCL_DEVICE as main GPU"
    #use signle GPU only
    GPUS_SETTING="-mg $GGML_SYCL_DEVICE -sm ${SPLIT_MODE}"
+    echo "ONEAPI_DEVICE_SELECTOR=${ONEAPI_DEVICE_SELECTOR}"
 else
-    echo "Use Intel GPUs: ${SYCL_DEVICES}"
+    echo "Use all Intel GPUs, including iGPU & dGPU"
    GPUS_SETTING="-sm ${SPLIT_MODE}"
-fi
+ fi

-echo "run cmd: ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --device ${SYCL_DEVICES} --mmap --host 0.0.0.0 --port 8000"
-ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --device ${SYCL_DEVICES} --mmap --host 0.0.0.0 --port 8000
+echo "run cmd: ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE}  --mmap --host 0.0.0.0 --port 8000"
+ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --mmap --host 0.0.0.0 --port 8000


@@ -12,7 +12,6 @@ This script processes files with specified options.

 Options:
  -h, --help    Display this help message and exit.
-  -d, --device  <value>    Set SYCL devices (default: SYCL0).
  -c, --context <value>    Set context length. Bigger need more memory.
  -p, --promote <value>    Prompt to start generation with.
  -m, --model   <value>    Full model file path.
@@ -43,16 +42,10 @@ MODEL_FILE=../models/llama-2-7b.Q4_0.gguf
 NGL=99
 CONTEXT=4096
 GGML_SYCL_DEVICE=-1
-SYCL_DEVICES="SYCL0"
 SPLIT_MODE=layer
 LOG_VERBOSE=3
 while [[ $# -gt 0 ]]; do
    case "$1" in
-        -d|--device)
-            SYCL_DEVICES="$2"
-            shift
-            shift
-            ;;
        -c|--context)
            CONTEXT=$2
            # Shift twice to consume both the option flag and its value
@@ -122,17 +115,16 @@ source /opt/intel/oneapi/setvars.sh
 export UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS=1
 echo "UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS=${UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS}"

-echo "ONEAPI_DEVICE_SELECTOR=${ONEAPI_DEVICE_SELECTOR}"
-
 if [ $GGML_SYCL_DEVICE -ne -1 ]; then
    echo "Use $GGML_SYCL_DEVICE as main GPU"
    #use signle GPU only
    GPUS_SETTING="-mg $GGML_SYCL_DEVICE -sm ${SPLIT_MODE}"
+    echo "ONEAPI_DEVICE_SELECTOR=${ONEAPI_DEVICE_SELECTOR}"
 else
-    echo "Use Intel GPUs: ${SYCL_DEVICES}"
+    echo "Use all Intel GPUs, including iGPU & dGPU"
    GPUS_SETTING="-sm ${SPLIT_MODE}"
 fi

-echo "run cmd: ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -no-cnv -p "${INPUT_PROMPT}" -n 200 -e -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --device ${SYCL_DEVICES} --mmap "
-ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -no-cnv -p "${INPUT_PROMPT}" -n 200 -e -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --device ${SYCL_DEVICES} --mmap
+echo "run cmd: ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -no-cnv -p "${INPUT_PROMPT}" -n 200 -e -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE}  --mmap "
+ZES_ENABLE_SYSMAN=1 ${BIN_FILE} -m ${MODEL_FILE} -no-cnv -p "${INPUT_PROMPT}" -n 200 -e -ngl ${NGL} -s ${SEED} -c ${CONTEXT} ${GPUS_SETTING} -lv ${LOG_VERBOSE} --mmap

@@ -13,7 +13,6 @@ set "MODEL_FILE=..\models\Qwen3.5-4B-Q4_0.gguf"
 set "NGL=99"
 set "CONTEXT=4096"
 set "GGML_SYCL_DEVICE=-1"
-set "SYCL_DEVICES=SYCL0"
 set "SPLIT_MODE=layer"
 set "LOG_VERBOSE=3"

@@ -37,21 +36,6 @@ if /I "%~1"=="--context" (
  goto parse_args
 )

-if /I "%~1"=="-d" (
-  if "%~2"=="" goto missing_value
-  set "SYCL_DEVICES=%~2"
-  shift
-  shift
-  goto parse_args
-)
-if /I "%~1"=="--device" (
-  if "%~2"=="" goto missing_value
-  set "SYCL_DEVICES=%~2"
-  shift
-  shift
-  goto parse_args
-)
-
 if /I "%~1"=="-m" (
  if "%~2"=="" goto missing_value
  set "MODEL_FILE=%~2"
@@ -146,7 +130,6 @@ echo This script processes files with specified options.
 echo.
 echo Options:
 echo   -h, --help    Display this help message and exit.
-echo   -d, --device ^<value^>    Set SYCL devices (default: SYCL0).
 echo   -c, --context ^<value^>    Set context length. Bigger need more memory.
 echo   -m, --model   ^<value^>    Full model file path.
 echo   -mg,--main-gpu ^<value^>   Set main GPU ID (0 - n) for single GPU mode.
@@ -177,20 +160,19 @@ REM Support malloc device memory more than 4GB.
 set "UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS=1"
 echo UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS=%UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS%

-echo ONEAPI_DEVICE_SELECTOR=%ONEAPI_DEVICE_SELECTOR%
-
 if not "%GGML_SYCL_DEVICE%"=="-1" (
  echo Use %GGML_SYCL_DEVICE% as main GPU
  REM Use single GPU only.
  set "GPUS_SETTING=-mg %GGML_SYCL_DEVICE% -sm %SPLIT_MODE%"
-  ) else (
-  echo Use Intel GPUs: %SYCL_DEVICES%
+  echo ONEAPI_DEVICE_SELECTOR=%ONEAPI_DEVICE_SELECTOR%
+) else (
+  echo Use all Intel GPUs, including iGPU ^& dGPU
  set "GPUS_SETTING=-sm %SPLIT_MODE%"
 )

-echo run cmd: ZES_ENABLE_SYSMAN=1 %BIN_FILE% -m "%MODEL_FILE%" -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --device %SYCL_DEVICES% --mmap --host 0.0.0.0 --port 8000
+echo run cmd: ZES_ENABLE_SYSMAN=1 %BIN_FILE% -m "%MODEL_FILE%" -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --mmap --host 0.0.0.0 --port 8000
 set "ZES_ENABLE_SYSMAN=1"
-%BIN_FILE% -m "%MODEL_FILE%" -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --device "%SYCL_DEVICES%" --mmap --host 0.0.0.0 --port 8000
+%BIN_FILE% -m "%MODEL_FILE%" -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --mmap --host 0.0.0.0 --port 8000

 endlocal

@@ -19,7 +19,6 @@ set "MODEL_FILE=..\models\llama-2-7b.Q4_0.gguf"
 set "NGL=99"
 set "CONTEXT=4096"
 set "GGML_SYCL_DEVICE=-1"
-set "SYCL_DEVICES=SYCL0"
 set "SPLIT_MODE=layer"
 set "LOG_VERBOSE=3"

@@ -43,21 +42,6 @@ if /I "%~1"=="--context" (
  goto parse_args
 )

-if /I "%~1"=="-d" (
-  if "%~2"=="" goto missing_value
-  set "SYCL_DEVICES=%~2"
-  shift
-  shift
-  goto parse_args
-)
-if /I "%~1"=="--device" (
-  if "%~2"=="" goto missing_value
-  set "SYCL_DEVICES=%~2"
-  shift
-  shift
-  goto parse_args
-)
-
 if /I "%~1"=="-p" (
  if "%~2"=="" goto missing_value
  set "INPUT_PROMPT=%~2"
@@ -167,7 +151,6 @@ echo This script processes files with specified options.
 echo.
 echo Options:
 echo   -h, --help    Display this help message and exit.
-echo   -d, --device ^<value^>    Set SYCL devices (default: SYCL0).
 echo   -c, --context ^<value^>    Set context length. Bigger need more memory.
 echo   -p, --promote ^<value^>    Prompt to start generation with.
 echo   -m, --model   ^<value^>    Full model file path.
@@ -199,21 +182,19 @@ REM Support malloc device memory more than 4GB.
 set "UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS=1"
 echo UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS=%UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS%

-echo ONEAPI_DEVICE_SELECTOR=%ONEAPI_DEVICE_SELECTOR%
-
 if not "%GGML_SYCL_DEVICE%"=="-1" (
  echo Use %GGML_SYCL_DEVICE% as main GPU
  REM Use single GPU only.
  set "GPUS_SETTING=-mg %GGML_SYCL_DEVICE% -sm %SPLIT_MODE%"
-  )
-else (
-  echo Use Intel GPUs: %SYCL_DEVICES%
+  echo ONEAPI_DEVICE_SELECTOR=%ONEAPI_DEVICE_SELECTOR%
+) else (
+  echo Use all Intel GPUs, including iGPU ^& dGPU
  set "GPUS_SETTING=-sm %SPLIT_MODE%"
 )

-echo run cmd: ZES_ENABLE_SYSMAN=1 %BIN_FILE% -m %MODEL_FILE% -no-cnv -p "%INPUT_PROMPT%" -n 200 -e -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --device %SYCL_DEVICES% --mmap
+echo run cmd: ZES_ENABLE_SYSMAN=1 %BIN_FILE% -m %MODEL_FILE% -no-cnv -p "%INPUT_PROMPT%" -n 200 -e -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --mmap
 set "ZES_ENABLE_SYSMAN=1"
-%BIN_FILE% -m "%MODEL_FILE%" -no-cnv -p "%INPUT_PROMPT%" -n 200 -e -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --device "%SYCL_DEVICES%" --mmap
+%BIN_FILE% -m "%MODEL_FILE%" -no-cnv -p "%INPUT_PROMPT%" -n 200 -e -ngl %NGL% -s %SEED% -c %CONTEXT% %GPUS_SETTING% -lv %LOG_VERBOSE% --mmap

 endlocal

@@ -4,8 +4,8 @@ project("ggml" C CXX ASM)

 ### GGML Version
 set(GGML_VERSION_MAJOR 0)
-set(GGML_VERSION_MINOR 19)
-set(GGML_VERSION_PATCH 0)
+set(GGML_VERSION_MINOR 18)
+set(GGML_VERSION_PATCH 1)
 set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")

 list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/")
@@ -2788,12 +2788,6 @@ extern "C" {
            struct ggml_cgraph * cgraph,
            struct ggml_tensor * tensor);

-    // add the tensor and its parents to the graph without marking them for compute
-    // the flag is set later, when the tensor is reached from a node that computes
-    GGML_API void ggml_build_forward_order(
-            struct ggml_cgraph * cgraph,
-            struct ggml_tensor * tensor);
-
    GGML_API void ggml_build_backward_expand(
        struct ggml_context *  ctx,        // context for gradient computation
        struct ggml_cgraph  *  cgraph,
@@ -8,22 +8,6 @@
 #include <sys/sysctl.h>
 #endif

-#if !defined(HWCAP_FPHP)
-#define HWCAP_FPHP (1 << 9)
-#endif
-
-#if !defined(HWCAP_ASIMDHP)
-#define HWCAP_ASIMDHP (1 << 10)
-#endif
-
-#if !defined(HWCAP_ASIMDDP)
-#define HWCAP_ASIMDDP (1 << 20)
-#endif
-
-#if !defined(HWCAP_SVE)
-#define HWCAP_SVE (1 << 22)
-#endif
-
 #if !defined(HWCAP2_SVE2)
 #define HWCAP2_SVE2 (1 << 1)
 #endif
@@ -39,7 +23,7 @@
 struct aarch64_features {
    // has_neon not needed, aarch64 has NEON guaranteed
    bool has_dotprod     = false;
-    bool has_fp16        = false;
+    bool has_fp16_va     = false;
    bool has_sve         = false;
    bool has_sve2        = false;
    bool has_i8mm        = false;
@@ -52,7 +36,7 @@ struct aarch64_features {
        uint32_t hwcap2 = getauxval(AT_HWCAP2);

        has_dotprod = !!(hwcap & HWCAP_ASIMDDP);
-        has_fp16    = !!(hwcap & HWCAP_FPHP) && !!(hwcap & HWCAP_ASIMDHP);
+        has_fp16_va = !!(hwcap & HWCAP_FPHP);
        has_sve     = !!(hwcap & HWCAP_SVE);
        has_sve2    = !!(hwcap2 & HWCAP2_SVE2);
        has_i8mm    = !!(hwcap2 & HWCAP2_I8MM);
@@ -91,7 +75,7 @@ static int ggml_backend_cpu_aarch64_score() {
    score += 1<<1;
 #endif
 #ifdef GGML_USE_FP16_VECTOR_ARITHMETIC
-    if (!af.has_fp16) { return 0; }
+    if (!af.has_fp16_va) { return 0; }
    score += 1<<2;
 #endif
 #ifdef GGML_USE_SVE
@@ -195,7 +195,6 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
                    case GGML_TYPE_Q4_K:
                    case GGML_TYPE_Q6_K:
                    case GGML_TYPE_Q8_0:
-                    case GGML_TYPE_Q5_0:
                    case GGML_TYPE_Q5_1:
                    case GGML_TYPE_Q5_K:
                        //case GGML_TYPE_MXFP4:
@@ -215,7 +214,6 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
                    case GGML_TYPE_Q4_K:
                    case GGML_TYPE_Q6_K:
                    case GGML_TYPE_Q8_0:
-                    case GGML_TYPE_Q5_0:
                    case GGML_TYPE_Q5_1:
                    case GGML_TYPE_Q5_K:
                        //case GGML_TYPE_MXFP4:
@@ -253,9 +253,9 @@ static void ggml_cpy_f32_q8_0_cuda(
    const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {

    GGML_ASSERT(ne % QK8_0 == 0);
-    const int64_t num_blocks = (ne/QK8_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
+    const int64_t num_blocks = ne / QK8_0;
    GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_f32_q<cpy_blck_f32_q8_0, QK8_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
+    cpy_f32_q<cpy_blck_f32_q8_0, QK8_0><<<num_blocks, 1, 0, stream>>>
        (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }

@@ -264,9 +264,9 @@ static void ggml_cpy_q8_0_f32_cuda(
    const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t nb00, const int64_t nb01, const int64_t nb02,
    const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {

-    const int64_t num_blocks = (ne/QK8_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
+    const int64_t num_blocks = ne;
    GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_q_f32<cpy_blck_q8_0_f32, QK8_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
+    cpy_q_f32<cpy_blck_q8_0_f32, QK8_0><<<num_blocks, 1, 0, stream>>>
        (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }

@@ -276,9 +276,9 @@ static void ggml_cpy_f32_q4_0_cuda(
    const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {

    GGML_ASSERT(ne % QK4_0 == 0);
-    const int64_t num_blocks = (ne/QK4_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
+    const int64_t num_blocks = ne / QK4_0;
    GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_f32_q<cpy_blck_f32_q4_0, QK4_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
+    cpy_f32_q<cpy_blck_f32_q4_0, QK4_0><<<num_blocks, 1, 0, stream>>>
        (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }

@@ -289,9 +289,9 @@ static void ggml_cpy_q4_0_f32_cuda(
    const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12,
    const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
    cudaStream_t stream) {
-    const int64_t num_blocks = (ne/QK4_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
+    const int64_t num_blocks = ne;
    GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_q_f32<cpy_blck_q_f32<dequantize_q4_0, QK4_0>, QK4_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>(
+    cpy_q_f32<cpy_blck_q_f32<dequantize_q4_0, QK4_0>, QK4_0><<<num_blocks, 1, 0, stream>>>(
        cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
         ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
@@ -302,9 +302,9 @@ static void ggml_cpy_f32_q4_1_cuda(
    const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {

    GGML_ASSERT(ne % QK4_1 == 0);
-    const int64_t num_blocks = (ne/QK4_1 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
+    const int64_t num_blocks = ne / QK4_1;
    GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_f32_q<cpy_blck_f32_q4_1, QK4_1><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
+    cpy_f32_q<cpy_blck_f32_q4_1, QK4_1><<<num_blocks, 1, 0, stream>>>
        (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }

@@ -315,9 +315,9 @@ static void ggml_cpy_q4_1_f32_cuda(
    const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12,
    const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
    cudaStream_t stream) {
-    const int64_t num_blocks = (ne/QK4_1 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
+    const int64_t num_blocks = ne;
    GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_q_f32<cpy_blck_q_f32<dequantize_q4_1, QK4_1>, QK4_1><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>(
+    cpy_q_f32<cpy_blck_q_f32<dequantize_q4_1, QK4_1>, QK4_1><<<num_blocks, 1, 0, stream>>>(
        cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
         ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
@@ -328,9 +328,9 @@ static void ggml_cpy_f32_q5_0_cuda(
    const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {

    GGML_ASSERT(ne % QK5_0 == 0);
-    const int64_t num_blocks = (ne/QK5_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
+    const int64_t num_blocks = ne / QK5_0;
    GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_f32_q<cpy_blck_f32_q5_0, QK5_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
+    cpy_f32_q<cpy_blck_f32_q5_0, QK5_0><<<num_blocks, 1, 0, stream>>>
        (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }

@@ -341,9 +341,9 @@ static void ggml_cpy_q5_0_f32_cuda(
    const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12,
    const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
    cudaStream_t stream) {
-    const int64_t num_blocks = (ne/QK5_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
+    const int64_t num_blocks = ne;
    GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_q_f32<cpy_blck_q_f32<dequantize_q5_0, QK5_0>, QK5_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>(
+    cpy_q_f32<cpy_blck_q_f32<dequantize_q5_0, QK5_0>, QK5_0><<<num_blocks, 1, 0, stream>>>(
        cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
        ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
@@ -354,9 +354,9 @@ static void ggml_cpy_f32_q5_1_cuda(
    const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {

    GGML_ASSERT(ne % QK5_1 == 0);
-    const int64_t num_blocks = (ne/QK5_1 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
+    const int64_t num_blocks = ne / QK5_1;
    GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_f32_q<cpy_blck_f32_q5_1, QK5_1><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
+    cpy_f32_q<cpy_blck_f32_q5_1, QK5_1><<<num_blocks, 1, 0, stream>>>
        (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }

@@ -367,9 +367,9 @@ static void ggml_cpy_q5_1_f32_cuda(
    const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12,
    const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
    cudaStream_t stream) {
-    const int64_t num_blocks = (ne/QK5_1 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
+    const int64_t num_blocks = ne;
    GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_q_f32<cpy_blck_q_f32<dequantize_q5_1, QK5_1>, QK5_1><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>(
+    cpy_q_f32<cpy_blck_q_f32<dequantize_q5_1, QK5_1>, QK5_1><<<num_blocks, 1, 0, stream>>>(
        cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
        ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
@@ -380,9 +380,9 @@ static void ggml_cpy_f32_iq4_nl_cuda(
    const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {

    GGML_ASSERT(ne % QK4_NL == 0);
-    const int64_t num_blocks = (ne/QK4_NL + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
+    const int64_t num_blocks = ne / QK4_NL;
    GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_f32_q<cpy_blck_f32_iq4_nl, QK4_NL><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
+    cpy_f32_q<cpy_blck_f32_iq4_nl, QK4_NL><<<num_blocks, 1, 0, stream>>>
        (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }

@@ -2651,52 +2651,6 @@ static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope,
    return true;
 }

-static bool ggml_cuda_should_fuse_rms_norm_mul_rope(const ggml_tensor * rms_norm,
-                                                    const ggml_tensor * mul,
-                                                    const ggml_tensor * rope) {
-    if (rms_norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || rope->op != GGML_OP_ROPE) {
-        return false;
-    }
-
-    if (rms_norm->src[0]->type != GGML_TYPE_F32 || rms_norm->type != GGML_TYPE_F32 ||
-        mul->src[0]->type != GGML_TYPE_F32 || mul->src[1]->type != GGML_TYPE_F32 ||
-        mul->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32) {
-        return false;
-    }
-
-    if (rope->src[0] != mul) {
-        return false;
-    }
-
-    //if rms norm is the B operand, then we don't handle broadcast
-    if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) {
-        return false;
-    }
-
-    if (!ggml_are_same_shape(rms_norm, mul)) {
-        return false;
-    }
-
-    //rms_norm kernel assumes contiguous rows
-    if (!ggml_is_contiguous_rows(rms_norm->src[0]) ||
-        !ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) {
-        return false;
-    }
-
-    // the fused kernel handles the norm/neox rope modes only
-    const int mode = ((const int32_t *) rope->op_params)[2];
-    if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX) {
-        return false;
-    }
-
-    const int n_dims = ((const int32_t *) rope->op_params)[1];
-    if (n_dims % 2 != 0 || rope->src[0]->ne[0] % 2 != 0) {
-        return false;
-    }
-
-    return true;
-}
-
 // match gated_delta_net + the strided cpy that scatters its state snapshots into the cache
 // (slot i -> rollback group i, slot 0 newest), so the kernel can write them and skip the cpy.
 static int ggml_cuda_try_gdn_cache_fusion(
@@ -3026,36 +2980,6 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
        }
    }

-    std::initializer_list<enum ggml_op> rms_norm_mul_rope_ops          = { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE };
-    std::initializer_list<enum ggml_op> rms_norm_mul_rope_set_rows_ops = { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS };
-
-    if (is_equal(rms_norm_mul_rope_set_rows_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 4 })) {
-        const ggml_tensor * rms_norm = cgraph->nodes[node_idx];
-        const ggml_tensor * mul      = cgraph->nodes[node_idx + 1];
-        const ggml_tensor * rope     = cgraph->nodes[node_idx + 2];
-        const ggml_tensor * view     = cgraph->nodes[node_idx + 3];
-        const ggml_tensor * set_rows = cgraph->nodes[node_idx + 4];
-
-        if (ggml_check_edges(cgraph, node_idx, {{1, 0, 0}, {2, 0, 1}, {3, 0, 2}, {4, 0, 3}}) &&
-            ggml_cuda_should_fuse_rms_norm_mul_rope(rms_norm, mul, rope) &&
-            ggml_cuda_should_fuse_rope_set_rows(rope, view, set_rows)) {
-            int out_nodes[] = { node_idx + 4 };
-            return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1);
-        }
-    }
-
-    if (is_equal(rms_norm_mul_rope_ops, ops) && ggml_can_fuse(cgraph, node_idx, ops)) {
-        const ggml_tensor * rms_norm = cgraph->nodes[node_idx];
-        const ggml_tensor * mul      = cgraph->nodes[node_idx + 1];
-        const ggml_tensor * rope     = cgraph->nodes[node_idx + 2];
-
-        if (ggml_cuda_should_fuse_rms_norm_mul_rope(rms_norm, mul, rope)) {
-            int out_nodes[] = { node_idx + 2 };
-            return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1);
-        }
-        return false;
-    }
-
    std::initializer_list<enum ggml_op> rope_set_rows_ops = { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS };

    if (is_equal(rope_set_rows_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) {
@@ -3064,8 +2988,7 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
        const ggml_tensor * set_rows = cgraph->nodes[node_idx + 2];

        if (ggml_cuda_should_fuse_rope_set_rows(rope, view, set_rows)) {
-            int out_nodes[] = { node_idx + 2 };
-            return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1);
+            return true;
        }
    }

@@ -3917,16 +3840,6 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
        return fused_node_count - 1;
    }

-    if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) {
-        ggml_cuda_op_rms_norm_mul_rope_fused(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 4]);
-        return 4;
-    }
-
-    if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE }, {})) {
-        ggml_cuda_op_rms_norm_mul_rope_fused(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2], nullptr);
-        return 2;
-    }
-
    if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) {
        ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]);
        return 2;
@@ -4120,11 +4033,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
                    continue;
                }
 #ifndef NDEBUG
-                // On integrated GPUs (APUs, e.g. RDNA3.5) the scheduler may place a
-                // node's output on the host-visible buffer, which the compute path
-                // handles. Allow that here, mirroring the src-tensor check below.
-                assert(node->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) ||
-                       (integrated && ggml_backend_buft_is_cuda_host(node->buffer->buft)));
+                assert(node->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device));
                for (int j = 0; j < GGML_MAX_SRC; j++) {
                    if (node->src[j] != nullptr) {
                        assert(node->src[j]->buffer);
@@ -5296,7 +5205,6 @@ static bool ggml_backend_cuda_device_offload_op(ggml_backend_dev_t dev, const gg

 static ggml_backend_event_t ggml_backend_cuda_device_event_new(ggml_backend_dev_t dev) {
 #ifdef GGML_CUDA_NO_PEER_COPY
-    GGML_UNUSED(dev);
    return nullptr;
 #else
    ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *)dev->context;
@@ -8,6 +8,7 @@ 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(
@@ -48,7 +49,6 @@ static __device__ __forceinline__ float nvfp4_native_scale_error(
    return err;
 }
 #endif // CUDART_VERSION >= 12080
-#endif // defined(BLACKWELL_MMA_AVAILABLE)

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

 void ggml_cuda_op_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * set_rows);
-
-void ggml_cuda_op_rms_norm_mul_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * rms_norm, ggml_tensor * mul, ggml_tensor * rope, ggml_tensor * set_rows);
@@ -3816,7 +3816,7 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) {
    }

    nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
-    nth = std::min(nth, (args.ne00_t + 31)/32*32);
+    nth = std::min(nth, args.ne00_t);

    const size_t smem = pipeline.smem;

@@ -11328,8 +11328,8 @@ kernel void kernel_lightning_indexer(
    const int i_kv_0   = tgpig.x*NK;            // first key of this threadgroup
    const int i_kv     = i_kv_0 + sgitg*NKPSG;  // first key of this simdgroup

-    threadgroup half sk[NK * DK16 * 16];
-    threadgroup half4x4 * sk4x4 = (threadgroup half4x4 *) sk;
+    threadgroup half4x4   sk4x4[NK*DK16];
+    threadgroup half    * sk = (threadgroup half *) sk4x4;

    for (short i = tiitg; i < NK*DK16; i += NTG) {
        const short ik  = i/DK16;
@@ -1022,20 +1022,9 @@ static T block_reduce(T val, T * shared_vals, int block_size_template) {
 }

 static __dpct_inline__ float ggml_sycl_ue4m3_to_fp32(uint8_t x) {
-    // UE4M3 is unsigned: 4 exp bits (bias 7), 3 mantissa bits, no sign, no NaN.
-    // exp == 0xF is a valid exponent (256-448 range), not NaN.
-    if (x == 0 || x == 0x7F) {
-        return 0.0f;
-    }
-    const int exp = (x >> 3) & 0xF;
-    const int man = x & 0x7;
-    float raw;
-    if (exp == 0) {
-        raw = man * (1.0f / 8.0f) * sycl::pow(2.0f, -6.0f);
-    } else {
-        raw = (1.0f + man / 8.0f) * sycl::pow(2.0f, (float) exp - 7.0f);
-    }
-    return raw * 0.5f;
+    const uint32_t bits = x * (x != 0x7F && x != 0xFF);
+    const __nv_fp8_e4m3 xf = *reinterpret_cast<const __nv_fp8_e4m3 *>(&bits);
+    return static_cast<float>(xf) / 2;
 }

 #endif // GGML_SYCL_COMMON_HPP
@@ -1,280 +0,0 @@
-#include "ggml-impl.h"
-#include "dsv4-hc.hpp"
-
-#include <cmath>
-
-static constexpr int DSV4_HC = 4;
-
-static void dsv4_hc_pre_f32_sycl(
-        const float * x, const float * weights, float * dst,
-        int64_t n_embd, int64_t hc, int64_t n_tokens,
-        int64_t sx0, int64_t sx1, int64_t sx2,
-        int64_t sw0, int64_t sw1,
-        int64_t sd0, int64_t sd1,
-        queue_ptr stream) {
-    const int64_t nr = n_embd * n_tokens;
-    const int64_t block_size = 256;
-    const int64_t num_blocks = (nr + block_size - 1) / block_size;
-
-    stream->parallel_for(
-        sycl::nd_range<1>(sycl::range<1>(num_blocks * block_size), sycl::range<1>(block_size)),
-        [=](sycl::nd_item<1> item) {
-            const int64_t ir = item.get_global_id(0);
-            if (ir >= nr) {
-                return;
-            }
-
-            const int64_t i0 = ir % n_embd;
-            const int64_t it = ir / n_embd;
-
-            float sum = x[i0*sx0 + it*sx2] * weights[it*sw1];
-            for (int64_t ih = 1; ih < hc; ++ih) {
-                const float xv = x[i0*sx0 + ih*sx1 + it*sx2];
-                const float wv = weights[ih*sw0 + it*sw1];
-                sum += xv * wv;
-            }
-
-            dst[i0*sd0 + it*sd1] = sum;
-        });
-}
-
-static void dsv4_hc_comb_norm_cols(float * comb, float eps) {
-    for (int idst = 0; idst < DSV4_HC; ++idst) {
-        float sum = eps;
-        for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
-            sum += comb[idst + DSV4_HC*isrc];
-        }
-
-        const float inv_sum = 1.0f / sum;
-        for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
-            comb[idst + DSV4_HC*isrc] *= inv_sum;
-        }
-    }
-}
-
-static void dsv4_hc_comb_norm_rows(float * comb, float eps) {
-    for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
-        float sum = eps;
-        for (int idst = 0; idst < DSV4_HC; ++idst) {
-            sum += comb[idst + DSV4_HC*isrc];
-        }
-
-        const float inv_sum = 1.0f / sum;
-        for (int idst = 0; idst < DSV4_HC; ++idst) {
-            comb[idst + DSV4_HC*isrc] *= inv_sum;
-        }
-    }
-}
-
-static void dsv4_hc_comb_f32_sycl(
-        const float * mixes,
-        const float * scale,
-        const float * base,
-        float * dst,
-        int64_t n_tokens,
-        int64_t sm0,
-        int64_t sm1,
-        int64_t ss0,
-        int64_t sb0,
-        int64_t sd0,
-        int64_t sd1,
-        int64_t sd2,
-        float eps,
-        int32_t n_iter,
-        queue_ptr stream) {
-    constexpr int comb_offset = 2*DSV4_HC;
-
-    const int64_t block_size = 256;
-    const int64_t num_blocks = (n_tokens + block_size - 1) / block_size;
-
-    stream->parallel_for(
-        sycl::nd_range<1>(sycl::range<1>(num_blocks * block_size), sycl::range<1>(block_size)),
-        [=](sycl::nd_item<1> item_ct1) {
-            const int64_t it = item_ct1.get_global_id(0);
-
-            if (it >= n_tokens) {
-                return;
-            }
-
-            const float scale_comb = scale[2*ss0];
-            float comb[DSV4_HC*DSV4_HC];
-
-            for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
-                float max = -INFINITY;
-                for (int idst = 0; idst < DSV4_HC; ++idst) {
-                    const int idx = idst + DSV4_HC*isrc;
-                    const float v = mixes[(comb_offset + idx)*sm0 + it*sm1] * scale_comb + base[(comb_offset + idx)*sb0];
-                    comb[idx] = v;
-                    max = fmaxf(max, v);
-                }
-
-                float sum = 0.0f;
-                for (int idst = 0; idst < DSV4_HC; ++idst) {
-                    const int idx = idst + DSV4_HC*isrc;
-                    const float v = expf(comb[idx] - max);
-                    comb[idx] = v;
-                    sum += v;
-                }
-
-                const float inv_sum = 1.0f / sum;
-                for (int idst = 0; idst < DSV4_HC; ++idst) {
-                    const int idx = idst + DSV4_HC*isrc;
-                    comb[idx] = comb[idx] * inv_sum + eps;
-                }
-            }
-
-            dsv4_hc_comb_norm_cols(comb, eps);
-            for (int32_t i = 1; i < n_iter; ++i) {
-                dsv4_hc_comb_norm_rows(comb, eps);
-                dsv4_hc_comb_norm_cols(comb, eps);
-            }
-
-            for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
-                for (int idst = 0; idst < DSV4_HC; ++idst) {
-                    const int idx = idst + DSV4_HC*isrc;
-                    dst[idst*sd0 + isrc*sd1 + it*sd2] = comb[idx];
-                }
-            }
-        });
-}
-
-static void dsv4_hc_post_f32_sycl(
-        const float * x, const float * residual, const float * post, const float * comb, float * dst,
-        int64_t n_embd, int64_t hc, int64_t n_tokens,
-        int64_t sx0, int64_t sx1,
-        int64_t sr0, int64_t sr1, int64_t sr2,
-        int64_t sp0, int64_t sp1,
-        int64_t sc0, int64_t sc1, int64_t sc2,
-        int64_t sd0, int64_t sd1, int64_t sd2,
-        queue_ptr stream) {
-    const int64_t nr = n_embd * hc * n_tokens;
-    const int64_t block_size = 256;
-    const int64_t num_blocks = (nr + block_size - 1) / block_size;
-
-    stream->parallel_for(
-        sycl::nd_range<1>(sycl::range<1>(num_blocks * block_size), sycl::range<1>(block_size)),
-        [=](sycl::nd_item<1> item) {
-            const int64_t ir = item.get_global_id(0);
-            if (ir >= nr) {
-                return;
-            }
-
-            const int64_t i0   = ir % n_embd;
-            const int64_t idst = (ir / n_embd) % hc;
-            const int64_t it   = ir / (n_embd * hc);
-
-            float sum = x[i0*sx0 + it*sx1] * post[idst*sp0 + it*sp1];
-            for (int64_t isrc = 0; isrc < hc; ++isrc) {
-                sum += residual[i0*sr0 + isrc*sr1 + it*sr2] * comb[idst*sc0 + isrc*sc1 + it*sc2];
-            }
-
-            dst[i0*sd0 + idst*sd1 + it*sd2] = sum;
-        });
-}
-
-void ggml_sycl_op_dsv4_hc_pre(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
-    const ggml_tensor * x       = dst->src[0];
-    const ggml_tensor * weights = dst->src[1];
-
-    GGML_ASSERT(x->type == GGML_TYPE_F32);
-    GGML_ASSERT(weights->type == GGML_TYPE_F32);
-    GGML_ASSERT(dst->type == GGML_TYPE_F32);
-
-    GGML_TENSOR_LOCALS(size_t, nbx, x,       nb);
-    GGML_TENSOR_LOCALS(size_t, nbw, weights, nb);
-    GGML_TENSOR_LOCALS(size_t, nbd, dst,     nb);
-
-    const int64_t n_embd   = x->ne[0];
-    const int64_t hc       = x->ne[1];
-    const int64_t n_tokens = x->ne[2];
-
-    queue_ptr stream = ctx.stream();
-
-    dsv4_hc_pre_f32_sycl(
-            (const float *) x->data, (const float *) weights->data, (float *) dst->data,
-            n_embd, hc, n_tokens,
-            nbx0 / sizeof(float), nbx1 / sizeof(float), nbx2 / sizeof(float),
-            nbw0 / sizeof(float), nbw1 / sizeof(float),
-            nbd0 / sizeof(float), nbd1 / sizeof(float),
-            stream);
-}
-
-void ggml_sycl_op_dsv4_hc_comb(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/3);
-
-    const ggml_tensor * mixes = dst->src[0];
-    const ggml_tensor * scale = dst->src[1];
-    const ggml_tensor * base  = dst->src[2];
-
-    GGML_ASSERT(mixes->type == GGML_TYPE_F32);
-    GGML_ASSERT(scale->type == GGML_TYPE_F32);
-    GGML_ASSERT(base->type == GGML_TYPE_F32);
-    GGML_ASSERT(dst->type == GGML_TYPE_F32);
-
-    constexpr int64_t hc_mix_dim = (2 + DSV4_HC)*DSV4_HC;
-
-    GGML_ASSERT(mixes->ne[0] == hc_mix_dim);
-    GGML_ASSERT(dst->ne[0] == DSV4_HC);
-    GGML_ASSERT(dst->ne[1] == DSV4_HC);
-    GGML_ASSERT(dst->ne[2] == mixes->ne[1]);
-    GGML_ASSERT(scale->ne[0] >= 3);
-    GGML_ASSERT(base->ne[0] == hc_mix_dim);
-
-    GGML_TENSOR_LOCALS(size_t, nbm, mixes, nb);
-    GGML_TENSOR_LOCALS(size_t, nbs, scale, nb);
-    GGML_TENSOR_LOCALS(size_t, nbb, base,  nb);
-    GGML_TENSOR_LOCALS(size_t, nbd, dst,   nb);
-
-    const int64_t n_tokens = mixes->ne[1];
-    const float eps = ggml_get_op_params_f32(dst, 0);
-    const int32_t n_iter = ggml_get_op_params_i32(dst, 1);
-
-    queue_ptr stream = ctx.stream();
-
-    dsv4_hc_comb_f32_sycl(
-            (const float *) mixes->data, (const float *) scale->data, (const float *) base->data, (float *) dst->data,
-            n_tokens,
-            nbm0 / sizeof(float), nbm1 / sizeof(float),
-            nbs0 / sizeof(float),
-            nbb0 / sizeof(float),
-            nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float),
-            eps, n_iter, stream);
-}
-
-void ggml_sycl_op_dsv4_hc_post(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/4);
-    const ggml_tensor * x        = dst->src[0];
-    const ggml_tensor * residual = dst->src[1];
-    const ggml_tensor * post     = dst->src[2];
-    const ggml_tensor * comb     = dst->src[3];
-
-    GGML_ASSERT(x->type == GGML_TYPE_F32);
-    GGML_ASSERT(residual->type == GGML_TYPE_F32);
-    GGML_ASSERT(post->type == GGML_TYPE_F32);
-    GGML_ASSERT(comb->type == GGML_TYPE_F32);
-    GGML_ASSERT(dst->type == GGML_TYPE_F32);
-
-    GGML_TENSOR_LOCALS(size_t, nbx, x,        nb);
-    GGML_TENSOR_LOCALS(size_t, nbr, residual, nb);
-    GGML_TENSOR_LOCALS(size_t, nbp, post,     nb);
-    GGML_TENSOR_LOCALS(size_t, nbc, comb,     nb);
-    GGML_TENSOR_LOCALS(size_t, nbd, dst,      nb);
-
-    const int64_t n_embd   = x->ne[0];
-    const int64_t n_tokens = x->ne[1];
-    const int64_t hc       = residual->ne[1];
-
-    queue_ptr stream = ctx.stream();
-
-    dsv4_hc_post_f32_sycl(
-            (const float *) x->data, (const float *) residual->data,
-            (const float *) post->data, (const float *) comb->data, (float *) dst->data,
-            n_embd, hc, n_tokens,
-            nbx0 / sizeof(float), nbx1 / sizeof(float),
-            nbr0 / sizeof(float), nbr1 / sizeof(float), nbr2 / sizeof(float),
-            nbp0 / sizeof(float), nbp1 / sizeof(float),
-            nbc0 / sizeof(float), nbc1 / sizeof(float), nbc2 / sizeof(float),
-            nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float),
-            stream);
-}
@@ -1,10 +0,0 @@
-#ifndef GGML_SYCL_DSV4_HC_HPP
-#define GGML_SYCL_DSV4_HC_HPP
-
-#include "common.hpp"
-
-void ggml_sycl_op_dsv4_hc_pre(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
-void ggml_sycl_op_dsv4_hc_comb(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
-void ggml_sycl_op_dsv4_hc_post(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
-
-#endif // GGML_SYCL_DSV4_HC_HPP
@@ -420,31 +420,53 @@ static void clamp(const T * x, T * dst, const float min, const float max, const
    }
 }

-template<typename T, typename F>
-static void unary_gated_op_flat_kernel(const T * x, const T * g, T * dst, const uint64_t k, const sycl::nd_item<1> & item_ct1, F func) {
-    SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
-        dst[i] = func(x[i]) * g[i];
-    }
-}
-
-template<typename T, typename F>
-static void unary_gated_op_generic_kernel(
-        const T * x,
-        const T * g,
-        T * dst,
-        const uint64_t k,
-        const sycl::uint3 n_fd,
-        const uint64_t o0,
-        const uint64_t o1,
-        const sycl::nd_item<1> & item_ct1,
-        F func) {
-
-    // rows of n columns at strides o0 and o1: two halves of one fused tensor, or two tensors
+template<typename T>
+static void gated_op_fused_geglu(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
    SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
        const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
        const int64_t j0 = rc.x() * o0 + rc.y();
        const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
-        dst[i] = func(x[j0]) * g[j1];
+        dst[i] = op_gelu(x[j0]) * g[j1];
+    }
+}
+
+template<typename T>
+static void gated_op_fused_reglu(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+    SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
+        const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+        const int64_t j0 = rc.x() * o0 + rc.y();
+        const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
+        dst[i] = op_relu(x[j0]) * g[j1];
+    }
+}
+
+template<typename T>
+static void gated_op_fused_swiglu(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+    SYCL_GLOBAL_ID_LOOP(k, item_ct1)  {
+        const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+        const int64_t j0 = rc.x() * o0 + rc.y();
+        const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
+        dst[i] = op_silu(x[j0]) * g[j1];
+    }
+}
+
+template<typename T>
+static void gated_op_fused_geglu_erf(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+    SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
+        const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+        const int64_t j0 = rc.x() * o0 + rc.y();
+        const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
+        dst[i] = op_gelu_erf(x[j0]) * g[j1];
+    }
+}
+
+template<typename T>
+static void gated_op_fused_geglu_quick(const T * x, const T * g, T * dst, const uint64_t k, const sycl::uint3 n_fd, const uint64_t o0, const uint64_t o1, const sycl::nd_item<1> &item_ct1) {
+    SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
+        const sycl::uint2 rc = fast_div_modulo((uint32_t) i, n_fd);
+        const int64_t j0 = rc.x() * o0 + rc.y();
+        const int64_t j1 = o0 == o1 ? j0 : rc.x() * o1 + rc.y();
+        dst[i] = op_gelu_quick(x[j0]) * g[j1];
    }
 }

@@ -648,35 +670,6 @@ static inline void ggml_sycl_op_unary(
        });
 }

-template<typename F>
-static inline void ggml_sycl_op_unary_gated(
-        ggml_backend_sycl_context & ctx, ggml_tensor * dst, F func) {
-
-    dispatch_ggml_sycl_op_fused_glu(ctx, dst,
-        [func](const auto * x_ptr, const auto * g_ptr, auto * dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
-
-            const uint32_t num_blocks = (uint32_t) ceil_div(k, SYCL_GLU_BLOCK_SIZE);
-            const sycl::nd_range<1> launch_range(num_blocks * sycl::range<1>(SYCL_GLU_BLOCK_SIZE),
-                                                 sycl::range<1>(SYCL_GLU_BLOCK_SIZE));
-
-            // o0 == n and o1 == n make the index math the identity, so index flat
-            // note: not ggml_is_contiguous - a fused [gate|up] src0 is contiguous with o0 == 2n
-            if (o0 == n && o1 == n) {
-                main_stream->parallel_for(launch_range,
-                    [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                        unary_gated_op_flat_kernel(x_ptr, g_ptr, dst_ptr, k, item_ct1, func);
-                    });
-            } else {
-                // launch-invariant divisor, and only this path needs it
-                const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
-                main_stream->parallel_for(launch_range,
-                    [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                        unary_gated_op_generic_kernel(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1, func);
-                    });
-            }
-        });
-}
-

 static inline void ggml_sycl_op_arange(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
    GGML_ASSERT(dst->type == GGML_TYPE_F32);
@@ -974,21 +967,42 @@ static inline void ggml_sycl_op_acc(ggml_backend_sycl_context & ctx, ggml_tensor
 }

 static inline void ggml_sycl_op_geglu(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    ggml_sycl_detail::ggml_sycl_op_unary_gated(ctx, dst, [](auto x) {
-        return op_gelu(x);
-    });
+    ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
+        [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
+            const uint32_t num_blocks = ceil_div(k, SYCL_GELU_BLOCK_SIZE);
+            const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
+            main_stream->parallel_for(
+                    sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_GELU_BLOCK_SIZE)),
+                    sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                gated_op_fused_geglu(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
+            });
+        });
 }

 static inline void ggml_sycl_op_reglu(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    ggml_sycl_detail::ggml_sycl_op_unary_gated(ctx, dst, [](auto x) {
-        return op_relu(x);
-    });
+    ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
+        [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
+            const uint32_t num_blocks = ceil_div((uint32_t)k, SYCL_RELU_BLOCK_SIZE); // Using RELU block size for reglu
+            const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
+            main_stream->parallel_for(
+                    sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_RELU_BLOCK_SIZE)),
+                    sycl::range<1>(SYCL_RELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                gated_op_fused_reglu(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
+            });
+        });
 }

 static inline void ggml_sycl_op_swiglu(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    ggml_sycl_detail::ggml_sycl_op_unary_gated(ctx, dst, [](auto x) {
-        return op_silu(x);
-    });
+    ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
+        [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
+            const uint32_t num_blocks = ceil_div((uint32_t)k, SYCL_SILU_BLOCK_SIZE); // Using SILU block size for swiglu
+            const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
+            main_stream->parallel_for(
+                    sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_SILU_BLOCK_SIZE)),
+                    sycl::range<1>(SYCL_SILU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                gated_op_fused_swiglu(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
+            });
+        });
 }

 __dpct_inline__ float ggml_sycl_op_swiglu_oai_single(float x, float g, float alpha = 1.702f, float limit = 7.0f) {
@@ -1083,15 +1097,29 @@ void ggml_sycl_op_swiglu_oai(ggml_backend_sycl_context & ctx, ggml_tensor * dst)
 }

 static inline void ggml_sycl_op_geglu_erf(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    ggml_sycl_detail::ggml_sycl_op_unary_gated(ctx, dst, [](auto x) {
-        return op_gelu_erf(x);
-    });
+    ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
+        [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
+            const uint32_t num_blocks = ceil_div(k, SYCL_GELU_BLOCK_SIZE);
+            const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
+            main_stream->parallel_for(
+                    sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_GELU_BLOCK_SIZE)),
+                    sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                gated_op_fused_geglu_erf(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
+            });
+        });
 }

 static inline void ggml_sycl_op_geglu_quick(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    ggml_sycl_detail::ggml_sycl_op_unary_gated(ctx, dst, [](auto x) {
-        return op_gelu_quick(x);
-    });
+    ggml_sycl_detail::dispatch_ggml_sycl_op_fused_glu(ctx, dst,
+        [](const auto* x_ptr, const auto* g_ptr, auto* dst_ptr, uint64_t k, uint64_t n, uint64_t o0, uint64_t o1, queue_ptr main_stream) {
+            const uint32_t num_blocks = ceil_div(k, SYCL_GELU_BLOCK_SIZE);
+            const sycl::uint3 n_fd = init_fastdiv_values((uint32_t) n);
+            main_stream->parallel_for(
+                    sycl::nd_range<1>((num_blocks * sycl::range<1>(SYCL_GELU_BLOCK_SIZE)),
+                    sycl::range<1>(SYCL_GELU_BLOCK_SIZE)), [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                gated_op_fused_geglu_quick(x_ptr, g_ptr, dst_ptr, k, n_fd, o0, o1, item_ct1);
+            });
+        });
 }


@@ -73,7 +73,6 @@ static void flash_attn_ext_vec(const char* __restrict__ Q,
                        const int32_t nb31,
                        const int32_t nb32,
                        const int64_t nb33) {
-
 #ifdef SYCL_FLASH_ATTN
    // Skip unused kernel variants for faster compilation:

@@ -470,6 +469,7 @@ static void flash_attn_ext_vec(const char* __restrict__ Q,
        }
    }

+
    item_ct1.barrier(sycl::access::fence_space::local_space);

 #pragma unroll
@@ -591,24 +591,22 @@ void ggml_sycl_flash_attn_ext_vec_case_impl(ggml_backend_sycl_context & ctx, ggm

    const auto arch = ggml_sycl_info().devices[ctx.device].hw_info.arch;
    const int nthreads = ggml_sycl_fattn_vec_get_nthreads_device(arch);
-    if constexpr (D <= 256) {
-        if (nthreads == 256) {
-            constexpr int nthreads_hw = 256;
-            constexpr int nwarps = nthreads_hw / warp_size;
-            launch_fattn<D, cols_per_block, 1,
-                         flash_attn_ext_vec<D, cols_per_block, type_K, type_V,
-                                            use_logit_softcap, warp_size, nthreads_hw>, warp_size>(
-                ctx, dst, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
-            return;
-        }
+    // 256 threads would overflow the 64 KB work-group local memory at D == 512, so keep 128 there.
+    if (D <= 256 && nthreads == 256) {
+        constexpr int nthreads_hw = 256;
+        constexpr int nwarps = nthreads_hw / warp_size;
+        launch_fattn<D, cols_per_block, 1,
+                     flash_attn_ext_vec<D, cols_per_block, type_K, type_V,
+                                        use_logit_softcap, warp_size, nthreads_hw>, warp_size>(
+            ctx, dst, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
+    } else {
+        constexpr int nthreads_hw = 128;
+        constexpr int nwarps = nthreads_hw / warp_size;
+        launch_fattn<D, cols_per_block, 1,
+                     flash_attn_ext_vec<D, cols_per_block, type_K, type_V,
+                                        use_logit_softcap, warp_size, nthreads_hw>, warp_size>(
+            ctx, dst, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
    }
-
-    constexpr int nthreads_hw = 128;
-    constexpr int nwarps = nthreads_hw / warp_size;
-    launch_fattn<D, cols_per_block, 1,
-                 flash_attn_ext_vec<D, cols_per_block, type_K, type_V,
-                                    use_logit_softcap, warp_size, nthreads_hw>, warp_size>(
-        ctx, dst, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
 }

 template <int D, int type_K, int type_V>
@@ -62,8 +62,6 @@
 #include "ggml-sycl/repeat_back.hpp"
 #include "ggml-sycl/set_rows.hpp"
 #include "ggml-sycl/set.hpp"
-#include "ggml-sycl/dsv4-hc.hpp"
-#include "ggml-sycl/lightning-indexer.hpp"
 #include "ggml-sycl/conv2d.hpp"
 #include "ggml-sycl/conv2d-dw.hpp"
 #include "ggml-sycl/conv2d-transpose.hpp"
@@ -4944,18 +4942,6 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg
        case GGML_OP_SET_ROWS:
            ggml_sycl_op_set_rows(ctx, dst);
            break;
-        case GGML_OP_DSV4_HC_PRE:
-            ggml_sycl_op_dsv4_hc_pre(ctx, dst);
-            break;
-        case GGML_OP_DSV4_HC_COMB:
-            ggml_sycl_op_dsv4_hc_comb(ctx, dst);
-            break;
-        case GGML_OP_DSV4_HC_POST:
-            ggml_sycl_op_dsv4_hc_post(ctx, dst);
-            break;
-        case GGML_OP_LIGHTNING_INDEXER:
-            ggml_sycl_op_lightning_indexer(ctx, dst);
-            break;
        case GGML_OP_DUP:
            ggml_sycl_dup(ctx, dst);
            break;
@@ -5809,33 +5795,17 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons

        case GGML_OP_SET_ROWS:
            {
-                auto res = (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16 ||
-                            op->src[0]->type == GGML_TYPE_BF16) &&
-                           (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32);
+
+                auto res = ((op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 ||
+                         op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q5_0 ||
+                         op->type == GGML_TYPE_Q1_0 ||
+                         op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_IQ4_NL ||
+                         op->type == GGML_TYPE_MXFP4 || op->type == GGML_TYPE_NVFP4) &&
+                        op->src[0]->type == GGML_TYPE_F32 &&
+                        (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32));
                return res;
            }
            break;
-        case GGML_OP_DSV4_HC_PRE:
-            return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
-                op->type == GGML_TYPE_F32;
-        case GGML_OP_DSV4_HC_COMB:
-            return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
-                op->src[2]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
-        case GGML_OP_DSV4_HC_POST:
-            return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
-                op->src[2]->type == GGML_TYPE_F32 && op->src[3]->type == GGML_TYPE_F32 &&
-                op->type == GGML_TYPE_F32;
-        case GGML_OP_LIGHTNING_INDEXER:
-            return op->src[0]->type == GGML_TYPE_F32 &&
-                (op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_F32 ||
-                 op->src[1]->type == GGML_TYPE_BF16 || op->src[1]->type == GGML_TYPE_Q8_0 ||
-                 op->src[1]->type == GGML_TYPE_Q5_1 || op->src[1]->type == GGML_TYPE_Q5_0 ||
-                 op->src[1]->type == GGML_TYPE_Q4_1 || op->src[1]->type == GGML_TYPE_Q4_0 ||
-                 op->src[1]->type == GGML_TYPE_IQ4_NL) &&
-                op->src[2]->type == GGML_TYPE_F32 &&
-                op->src[3]->type == GGML_TYPE_F16 &&
-                op->type == GGML_TYPE_F32 &&
-                op->src[0]->ne[0] == WARP_SIZE * 8;
        case GGML_OP_CPY:
            {
                ggml_type src0_type = op->src[0]->type;
@@ -1,197 +0,0 @@
-#include "lightning-indexer.hpp"
-#include "dequantize.hpp"
-
-static void lightning_indexer_f32_sycl(
-        const char * q, const char * k, const char * w, const char * m, float * dst,
-        int64_t n_embd, int64_t n_head, int64_t n_batch, int64_t n_stream, int64_t n_kv,
-        int64_t nem3,
-        int64_t nbq1, int64_t nbq2, int64_t nbq3,
-        int64_t nbk2, int64_t nbk3,
-        int64_t nbw1, int64_t nbw3,
-        int64_t nbm1, int64_t nbm3,
-        int64_t nb1, int64_t nb3,
-        ggml_type k_type,
-        queue_ptr stream) {
-
-    constexpr int64_t LANES = WARP_SIZE;
-    constexpr int64_t ELEMS_PER_LANE = 8;
-    constexpr int64_t ROWS_PER_BLOCK = 4;
-    constexpr int64_t BLOCK_SIZE = ROWS_PER_BLOCK * LANES;
-
-    const int64_t n_rows = n_batch * n_stream * n_kv;
-    const int64_t n_blocks = (n_rows + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK;
-
-    stream->parallel_for(
-        sycl::nd_range<1>(
-            sycl::range<1>(n_blocks * BLOCK_SIZE),
-            sycl::range<1>(BLOCK_SIZE)),
-        [=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            const int64_t ir   = item.get_global_id(0);
-            const int64_t lane = ir % LANES;
-            const int64_t row  = ir / LANES;
-            if (row >= n_rows) {
-                return;
-            }
-
-            const int64_t i_bs     = row / n_kv;
-            const int64_t i_kv     = row % n_kv;
-            const int64_t i_batch  = i_bs / n_stream;
-            const int64_t i_stream = i_bs % n_stream;
-
-            // load K row slice into registers (row is contiguous, nbk0 == type size)
-            const char * k_base = k + i_kv*nbk2 + i_stream*nbk3;
-            float k_local[ELEMS_PER_LANE];
-            if (k_type == GGML_TYPE_F16) {
-                const sycl::half * k_row = (const sycl::half *) k_base;
-#pragma unroll
-                for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) {
-                    k_local[j] = static_cast<float>(k_row[lane*ELEMS_PER_LANE + j]);
-                }
-            } else if (k_type == GGML_TYPE_F32) {
-                const float * k_row = (const float *) k_base;
-#pragma unroll
-                for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) {
-                    k_local[j] = k_row[lane*ELEMS_PER_LANE + j];
-                }
-            } else {
-                const int64_t lane_base = lane * ELEMS_PER_LANE;
-                switch (k_type) {
-                    case GGML_TYPE_BF16: {
-                        const sycl::ext::oneapi::bfloat16 * k_row = (const sycl::ext::oneapi::bfloat16 *) k_base;
-#pragma unroll
-                        for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) {
-                            k_local[j] = static_cast<float>(k_row[lane_base + j]);
-                        }
-                    } break;
-                    case GGML_TYPE_Q4_0:
-                    case GGML_TYPE_Q4_1:
-                    case GGML_TYPE_Q5_0:
-                    case GGML_TYPE_Q5_1: {
-#pragma unroll
-                        for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) {
-                            const int64_t idx = lane_base + j;
-                            const int64_t ib  = idx / QK4_0;
-                            const int iqs     = idx % (QK4_0/2);
-                            dfloat2 kv;
-                            if (k_type == GGML_TYPE_Q4_0) {
-                                dequantize_q4_0(k_base, ib, iqs, kv);
-                            } else if (k_type == GGML_TYPE_Q4_1) {
-                                dequantize_q4_1(k_base, ib, iqs, kv);
-                            } else if (k_type == GGML_TYPE_Q5_0) {
-                                dequantize_q5_0(k_base, ib, iqs, kv);
-                            } else {
-                                dequantize_q5_1(k_base, ib, iqs, kv);
-                            }
-                            k_local[j] = (idx % QK4_0) < (QK4_0/2) ? static_cast<float>(kv.x()) : static_cast<float>(kv.y());
-                        }
-                    } break;
-                    case GGML_TYPE_Q8_0: {
-#pragma unroll
-                        for (int64_t pair = 0; pair < ELEMS_PER_LANE / 2; ++pair) {
-                            const int64_t elem0 = lane_base + 2 * pair;
-                            dfloat2 kv;
-                            dequantize_q8_0(k_base, elem0 / QK8_0, elem0 % QK8_0, kv);
-                            k_local[2 * pair + 0] = static_cast<float>(kv.x());
-                            k_local[2 * pair + 1] = static_cast<float>(kv.y());
-                        }
-                    } break;
-                    case GGML_TYPE_IQ4_NL: {
-#pragma unroll
-                        for (int64_t pair = 0; pair < ELEMS_PER_LANE / 2; ++pair) {
-                            const int64_t elem0 = lane_base + 2 * pair;
-                            dfloat2 kv;
-                            dequantize_iq4_nl(k_base, elem0 / QK4_NL, elem0 % QK4_NL, kv);
-                            k_local[2 * pair + 0] = static_cast<float>(kv.x());
-                            k_local[2 * pair + 1] = static_cast<float>(kv.y());
-                        }
-                    } break;
-                    default:
-#pragma unroll
-                        for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) {
-                            k_local[j] = 0.0f;
-                        }
-                        break;
-                }
-            }
-
-            const char  * q_base = q + i_batch*nbq2 + i_stream*nbq3;
-            const float * w_base = (const float *) (w + i_batch*nbw1 + i_stream*nbw3);
-
-            float score = 0.0f;
-            for (int64_t h = 0; h < n_head; ++h) {
-                const float * q_row = (const float *) (q_base + h*nbq1);
-                float dot = 0.0f;
-#pragma unroll
-                for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) {
-                    const int64_t i = lane*ELEMS_PER_LANE + j;
-                    if (i < n_embd) {
-                        dot += q_row[i] * k_local[j];
-                    }
-                }
-                dot = sycl::reduce_over_group(item.get_sub_group(), dot, sycl::plus<float>());
-                if (lane == 0) {
-                    score += sycl::max(dot, 0.0f) * w_base[h];
-                }
-            }
-
-            if (lane == 0) {
-                const sycl::half * m_base = (const sycl::half *) (m + i_batch*nbm1 + (i_stream % nem3)*nbm3);
-                // flat-index store: storing through a strided base pointer
-                // hangs/misroutes writes on this stack when n_batch*n_stream > 1
-                const int64_t dst_idx = i_kv + i_batch*(nb1/sizeof(float)) + i_stream*(nb3/sizeof(float));
-                dst[dst_idx] = score + static_cast<float>(m_base[i_kv]);
-            }
-        });
-}
-
-void ggml_sycl_op_lightning_indexer(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
-    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/4);
-    const ggml_tensor * q = dst->src[0];
-    const ggml_tensor * k = dst->src[1];
-    const ggml_tensor * w = dst->src[2]; // weights
-    const ggml_tensor * m = dst->src[3]; // mask
-
-    GGML_ASSERT(dst->type == GGML_TYPE_F32);
-    GGML_ASSERT(  q->type == GGML_TYPE_F32);
-    GGML_ASSERT(  w->type == GGML_TYPE_F32);
-    GGML_ASSERT(  m->type == GGML_TYPE_F16);
-    GGML_ASSERT(k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_BF16 ||
-                k->type == GGML_TYPE_Q8_0 || k->type == GGML_TYPE_Q5_1 || k->type == GGML_TYPE_Q5_0 ||
-                k->type == GGML_TYPE_Q4_1 || k->type == GGML_TYPE_Q4_0 || k->type == GGML_TYPE_IQ4_NL);
-
-    GGML_TENSOR_LOCALS(int64_t, neq, q, ne);
-    GGML_TENSOR_LOCALS(size_t,  nbq, q, nb);
-    GGML_TENSOR_LOCALS(int64_t, nek, k, ne);
-    GGML_TENSOR_LOCALS(size_t,  nbk, k, nb);
-    GGML_TENSOR_LOCALS(size_t,  nbw, w, nb);
-    GGML_TENSOR_LOCALS(int64_t, nem, m, ne);
-    GGML_TENSOR_LOCALS(size_t,  nbm, m, nb);
-    GGML_TENSOR_LOCALS(int64_t, ne, dst, ne);
-    GGML_TENSOR_LOCALS(size_t,  nb, dst, nb);
-
-    // input rows must be contiguous
-    GGML_ASSERT(nbq0 == ggml_type_size(q->type));
-    GGML_ASSERT(nbk0 == ggml_type_size(k->type));
-    GGML_ASSERT(nbm0 == ggml_type_size(m->type));
-    GGML_ASSERT(nb0  == ggml_type_size(dst->type));
-
-    const int64_t n_embd   = neq0;
-    const int64_t n_head   = neq1;
-    const int64_t n_batch  = neq2;
-    const int64_t n_stream = neq3;
-    const int64_t n_kv     = nek2;
-
-    GGML_ASSERT(n_embd == WARP_SIZE * 8);
-
-    lightning_indexer_f32_sycl(
-            (const char *) q->data, (const char *) k->data,
-            (const char *) w->data, (const char *) m->data, (float *) dst->data,
-            n_embd, n_head, n_batch, n_stream, n_kv, nem3,
-            nbq1, nbq2, nbq3,
-            nbk2, nbk3,
-            nbw1, nbw3,
-            nbm1, nbm3,
-            nb1, nb3,
-            k->type,
-            ctx.stream());
-}
@@ -1,8 +0,0 @@
-#ifndef GGML_SYCL_LIGHTNING_INDEXER_HPP
-#define GGML_SYCL_LIGHTNING_INDEXER_HPP
-
-#include "common.hpp"
-
-void ggml_sycl_op_lightning_indexer(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
-
-#endif // GGML_SYCL_LIGHTNING_INDEXER_HPP
@@ -20,6 +20,8 @@
 #define MATRIX_ROW_PADDING 512 // last row of quant. matrices is a multiple of this to avoid out-of-bounds memory accesses

 #define SYCL_COL2IM_1D_BLOCK_SIZE 256
+#define SYCL_GELU_BLOCK_SIZE 256
+#define SYCL_SILU_BLOCK_SIZE 256
 #define SYCL_TANH_BLOCK_SIZE 256
 #define SYCL_RELU_BLOCK_SIZE 256
 #define SYCL_HARDSIGMOID_BLOCK_SIZE 256
@@ -1,10 +1,6 @@
 #include "set_rows.hpp"
 #include "cpy.hpp"

-#include "ggml-quants.h"
-
-#include <vector>
-
 namespace utils {
 template<typename T>
 static constexpr bool is_arithmetic_v() {
@@ -24,17 +20,7 @@ convert (const char* src, char* dst) {
   *reinterpret_cast<TOut*>(dst) = dst_val;
 }

-#ifdef GGML_SYCL_HAS_BF16
-// sycl::vec::convert does not provide a half -> bfloat16 path, so route through float.
-template<>
-inline void convert<sycl::half, sycl::ext::oneapi::bfloat16>(const char* src, char* dst) {
-    const float tmp = sycl::vec<sycl::half, 1>(*reinterpret_cast<const sycl::half*>(src))
-                          .template convert<float, sycl::rounding_mode::automatic>()[0];
-    *reinterpret_cast<sycl::ext::oneapi::bfloat16*>(dst) = sycl::ext::oneapi::bfloat16(tmp);
-}
-#endif
-
-template <typename TIn, typename TIdx, typename blockType, int qk, cpy_kernel_t cpyblck>
+template <typename TIdx, typename blockType, int qk, cpy_kernel_t cpyblck>
 static void set_rows_sycl_q(const char * __restrict__ src0_d,
                            const TIdx * __restrict__ src1_d,
                            blockType * __restrict__ dst_d,
@@ -82,22 +68,13 @@ static void set_rows_sycl_q(const char * __restrict__ src0_d,
        const int64_t i11         = i02 % ne11;
        const int64_t i10         = i01;
        const size_t  src_offset  = calculate_offset<3>({ nb01, nb02, nb03 }, { i01, i02, i03 });
-        const char *  src_block   = src0_d + src_offset + i00 * sizeof(TIn);
+        const char *  src_block   = src0_d + src_offset + i00 * sizeof(float);
        const size_t  src1_offset = calculate_offset<3>({ nb10, nb11, nb12 }, { i10, i11, i12 });
        const int64_t dst_row     = src1_d[src1_offset / sizeof(TIdx)];
        const size_t  dst_offset =
            calculate_offset<3>({ nb1, nb2, nb3 }, { dst_row, i02, i03 }) + (i00 / qk) * sizeof(blockType);
        char * dst_block = reinterpret_cast<char *>(reinterpret_cast<char *>(dst_d) + dst_offset);
-        if constexpr (std::is_same_v<TIn, float>) {
-            cpyblck(src_block, dst_block);
-        } else {
-            float src_block_f32[qk];
-            const TIn * src_block_t = reinterpret_cast<const TIn *>(src_block);
-            for (int j = 0; j < qk; ++j) {
-                src_block_f32[j] = (float) src_block_t[j];
-            }
-            cpyblck(reinterpret_cast<const char *>(src_block_f32), dst_block);
-        }
+        cpyblck(src_block, dst_block);
    });
    GGML_UNUSED(ne10);
    GGML_UNUSED(ne13);
@@ -105,139 +82,6 @@ static void set_rows_sycl_q(const char * __restrict__ src0_d,
    GGML_UNUSED(nb13);
 }

-template<typename blockType>
-using quantize_row_qk_t = void (*)(const float *, blockType *, int64_t);
-
-using quantize_rows_f_t = size_t (*)(const float *, void *, int64_t, int64_t, const float *);
-
-template <typename TIn, typename TIdx, typename blockType, int qk, quantize_row_qk_t<blockType> quantize_row>
-static void set_rows_sycl_qk_host(
-        const ggml_tensor * src0,
-        const ggml_tensor * src1,
-        ggml_tensor * dst,
-        const int64_t ne00,
-        const int64_t ne01,
-        const int64_t ne02,
-        const int64_t ne03,
-        const int64_t ne11,
-        const int64_t ne12,
-        const size_t nb01,
-        const size_t nb02,
-        const size_t nb03,
-        const size_t nb10,
-        const size_t nb11,
-        const size_t nb12,
-        const size_t nb1,
-        const size_t nb2,
-        const size_t nb3,
-        queue_ptr stream) {
-    GGML_ASSERT(ne00 % qk == 0);
-
-    const size_t src0_bytes = ggml_nbytes(src0);
-    const size_t src1_bytes = ggml_nbytes(src1);
-
-    std::vector<char> src0_host(src0_bytes);
-    std::vector<char> src1_host(src1_bytes);
-
-    stream->memcpy(src0_host.data(), src0->data, src0_bytes);
-    stream->memcpy(src1_host.data(), src1->data, src1_bytes);
-    stream->wait();
-
-    std::vector<float> src_row_f32(ne00);
-    const int64_t nblocks = ne00 / qk;
-    std::vector<blockType> dst_row_q(nblocks);
-
-    for (int64_t i03 = 0; i03 < ne03; ++i03) {
-        for (int64_t i02 = 0; i02 < ne02; ++i02) {
-            for (int64_t i01 = 0; i01 < ne01; ++i01) {
-                const int64_t i12 = i03 % ne12;
-                const int64_t i11 = i02 % ne11;
-                const int64_t i10 = i01;
-
-                const size_t src1_offset = calculate_offset<3>({ nb10, nb11, nb12 }, { i10, i11, i12 });
-                const int64_t dst_row = *(const TIdx *) (src1_host.data() + src1_offset);
-
-                const size_t src0_row_offset = calculate_offset<3>({ nb01, nb02, nb03 }, { i01, i02, i03 });
-                const TIn * src_row = reinterpret_cast<const TIn *>(src0_host.data() + src0_row_offset);
-
-                for (int64_t i00 = 0; i00 < ne00; ++i00) {
-                    src_row_f32[i00] = (float) src_row[i00];
-                }
-
-                quantize_row(src_row_f32.data(), dst_row_q.data(), ne00);
-
-                const size_t dst_offset = calculate_offset<3>({ nb1, nb2, nb3 }, { dst_row, i02, i03 });
-                stream->memcpy((char *) dst->data + dst_offset, dst_row_q.data(), nblocks * sizeof(blockType));
-                stream->wait();
-            }
-        }
-    }
-}
-
-template <typename TIn, typename TIdx, typename blockType, int qk, quantize_rows_f_t quantize_rows>
-static void set_rows_sycl_iq_host(
-        const ggml_tensor * src0,
-        const ggml_tensor * src1,
-        ggml_tensor * dst,
-        const int64_t ne00,
-        const int64_t ne01,
-        const int64_t ne02,
-        const int64_t ne03,
-        const int64_t ne11,
-        const int64_t ne12,
-        const size_t nb01,
-        const size_t nb02,
-        const size_t nb03,
-        const size_t nb10,
-        const size_t nb11,
-        const size_t nb12,
-        const size_t nb1,
-        const size_t nb2,
-        const size_t nb3,
-        queue_ptr stream) {
-    GGML_ASSERT(ne00 % qk == 0);
-
-    const size_t src0_bytes = ggml_nbytes(src0);
-    const size_t src1_bytes = ggml_nbytes(src1);
-
-    std::vector<char> src0_host(src0_bytes);
-    std::vector<char> src1_host(src1_bytes);
-
-    stream->memcpy(src0_host.data(), src0->data, src0_bytes);
-    stream->memcpy(src1_host.data(), src1->data, src1_bytes);
-    stream->wait();
-
-    std::vector<float> src_row_f32(ne00);
-    const int64_t nblocks = ne00 / qk;
-    std::vector<blockType> dst_row_q(nblocks);
-
-    for (int64_t i03 = 0; i03 < ne03; ++i03) {
-        for (int64_t i02 = 0; i02 < ne02; ++i02) {
-            for (int64_t i01 = 0; i01 < ne01; ++i01) {
-                const int64_t i12 = i03 % ne12;
-                const int64_t i11 = i02 % ne11;
-                const int64_t i10 = i01;
-
-                const size_t src1_offset = calculate_offset<3>({ nb10, nb11, nb12 }, { i10, i11, i12 });
-                const int64_t dst_row = *(const TIdx *) (src1_host.data() + src1_offset);
-
-                const size_t src0_row_offset = calculate_offset<3>({ nb01, nb02, nb03 }, { i01, i02, i03 });
-                const TIn * src_row = reinterpret_cast<const TIn *>(src0_host.data() + src0_row_offset);
-
-                for (int64_t i00 = 0; i00 < ne00; ++i00) {
-                    src_row_f32[i00] = (float) src_row[i00];
-                }
-
-                quantize_rows(src_row_f32.data(), dst_row_q.data(), 1, ne00, nullptr);
-
-                const size_t dst_offset = calculate_offset<3>({ nb1, nb2, nb3 }, { dst_row, i02, i03 });
-                stream->memcpy((char *) dst->data + dst_offset, dst_row_q.data(), nblocks * sizeof(blockType));
-                stream->wait();
-            }
-        }
-    }
-}
-
 template<typename TIn, typename TIdx, typename TOut>
 static void k_set_rows(
        const char * __restrict__ src0, const TIdx * __restrict__ src1, char * __restrict__ dst,
@@ -356,194 +200,31 @@ static void set_rows_sycl(ggml_backend_sycl_context & ctx, const ggml_tensor * s
            break;
 #endif
        case GGML_TYPE_Q8_0:
-            set_rows_sycl_q<TIn, TIdx, block_q8_0, QK8_0, cpy_blck_f32_q8_0>(
-                src0_d, src1_d, (block_q8_0 *) dst->data, ne00, ne01, ne02, ne03,
-                ne10, ne11, ne12, ne13, nb00, nb01,
-                nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
+            set_rows_sycl_q<TIdx, block_q8_0, QK8_0, cpy_blck_f32_q8_0>(src0_d, src1_d, (block_q8_0 *)dst->data, ne00, ne01, ne02, ne03, ne10, ne11, ne12, ne13, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
            break;
        case GGML_TYPE_Q1_0:
-            set_rows_sycl_q<TIn, TIdx, block_q1_0, QK1_0, cpy_blck_f32_q1_0>(
-                src0_d, src1_d, (block_q1_0 *) dst->data, ne00, ne01, ne02, ne03,
-                ne10, ne11, ne12, ne13, nb00, nb01,
-                nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
-            break;
-        case GGML_TYPE_Q2_0:
-            set_rows_sycl_q<TIn, TIdx, block_q2_0, QK2_0, cpy_blck_f32_q2_0>(
-                src0_d, src1_d, (block_q2_0 *) dst->data, ne00, ne01, ne02, ne03,
-                ne10, ne11, ne12, ne13, nb00, nb01,
-                nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
+            set_rows_sycl_q<TIdx, block_q1_0, QK1_0, cpy_blck_f32_q1_0>(src0_d, src1_d, (block_q1_0 *)dst->data, ne00, ne01, ne02, ne03, ne10, ne11, ne12, ne13, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
            break;
        case GGML_TYPE_Q5_1:
-            set_rows_sycl_q<TIn, TIdx, block_q5_1, QK5_1, cpy_blck_f32_q5_1>(
-                src0_d, src1_d, (block_q5_1 *) dst->data, ne00, ne01, ne02, ne03,
-                ne10, ne11, ne12, ne13, nb00, nb01,
-                nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
+            set_rows_sycl_q<TIdx, block_q5_1, QK5_1, cpy_blck_f32_q5_1>(src0_d, src1_d, (block_q5_1 *)dst->data, ne00, ne01, ne02, ne03, ne10, ne11, ne12, ne13, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
            break;
        case GGML_TYPE_Q5_0:
-            set_rows_sycl_q<TIn, TIdx, block_q5_0, QK5_0, cpy_blck_f32_q5_0>(
-                src0_d, src1_d, (block_q5_0 *) dst->data, ne00, ne01, ne02, ne03,
-                ne10, ne11, ne12, ne13, nb00, nb01,
-                nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
+            set_rows_sycl_q<TIdx, block_q5_0, QK5_0, cpy_blck_f32_q5_0>(src0_d, src1_d, (block_q5_0 *)dst->data, ne00, ne01, ne02, ne03, ne10, ne11, ne12, ne13, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
            break;
        case GGML_TYPE_Q4_1:
-            set_rows_sycl_q<TIn, TIdx, block_q4_1, QK4_1, cpy_blck_f32_q4_1>(
-                src0_d, src1_d, (block_q4_1 *) dst->data, ne00, ne01, ne02, ne03,
-                ne10, ne11, ne12, ne13, nb00, nb01,
-                nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
+            set_rows_sycl_q<TIdx, block_q4_1, QK4_1, cpy_blck_f32_q4_1>(src0_d, src1_d, (block_q4_1 *)dst->data, ne00, ne01, ne02, ne03, ne10, ne11, ne12, ne13, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
            break;
        case GGML_TYPE_Q4_0:
-            set_rows_sycl_q<TIn, TIdx, block_q4_0, QK4_0, cpy_blck_f32_q4_0>(
-                src0_d, src1_d, (block_q4_0 *) dst->data, ne00, ne01, ne02, ne03,
-                ne10, ne11, ne12, ne13, nb00, nb01,
-                nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
+            set_rows_sycl_q<TIdx, block_q4_0, QK4_0, cpy_blck_f32_q4_0>(src0_d, src1_d, (block_q4_0 *)dst->data, ne00, ne01, ne02, ne03, ne10, ne11, ne12, ne13, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
            break;
        case GGML_TYPE_IQ4_NL:
-            set_rows_sycl_q<TIn, TIdx, block_iq4_nl, QK4_NL, cpy_blck_f32_iq4_nl>(
-                src0_d, src1_d, (block_iq4_nl *) dst->data, ne00, ne01, ne02, ne03,
-                ne10, ne11, ne12, ne13, nb00, nb01,
-                nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
+            set_rows_sycl_q<TIdx, block_iq4_nl, QK4_NL, cpy_blck_f32_iq4_nl>(src0_d, src1_d, (block_iq4_nl *)dst->data, ne00, ne01, ne02, ne03, ne10, ne11, ne12, ne13, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
            break;
        case GGML_TYPE_MXFP4:
-            set_rows_sycl_q<TIn, TIdx, block_mxfp4, QK_MXFP4, cpy_blck_f32_mxfp4>(
-                src0_d, src1_d, (block_mxfp4 *) dst->data, ne00, ne01, ne02, ne03,
-                ne10, ne11, ne12, ne13, nb00, nb01,
-                nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
+            set_rows_sycl_q<TIdx, block_mxfp4, QK_MXFP4, cpy_blck_f32_mxfp4>(src0_d, src1_d, (block_mxfp4 *)dst->data, ne00, ne01, ne02, ne03, ne10, ne11, ne12, ne13, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
            break;
        case GGML_TYPE_NVFP4:
-            set_rows_sycl_q<TIn, TIdx, block_nvfp4, QK_NVFP4, cpy_blck_f32_nvfp4>(
-                src0_d, src1_d, (block_nvfp4 *) dst->data, ne00, ne01, ne02, ne03,
-                ne10, ne11, ne12, ne13, nb00, nb01,
-                nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
-            break;
-        case GGML_TYPE_Q2_K:
-            set_rows_sycl_qk_host<TIn, TIdx, block_q2_K, QK_K, quantize_row_q2_K_ref>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
-            break;
-        case GGML_TYPE_Q3_K:
-            set_rows_sycl_qk_host<TIn, TIdx, block_q3_K, QK_K, quantize_row_q3_K_ref>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
-            break;
-        case GGML_TYPE_Q4_K:
-            set_rows_sycl_qk_host<TIn, TIdx, block_q4_K, QK_K, quantize_row_q4_K_ref>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
-            break;
-        case GGML_TYPE_Q5_K:
-            set_rows_sycl_qk_host<TIn, TIdx, block_q5_K, QK_K, quantize_row_q5_K_ref>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
-            break;
-        case GGML_TYPE_Q6_K:
-            set_rows_sycl_qk_host<TIn, TIdx, block_q6_K, QK_K, quantize_row_q6_K_ref>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
-            break;
-        case GGML_TYPE_IQ2_XXS:
-            set_rows_sycl_iq_host<TIn, TIdx, block_iq2_xxs, QK_K, quantize_iq2_xxs>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
-            break;
-        case GGML_TYPE_IQ2_XS:
-            set_rows_sycl_iq_host<TIn, TIdx, block_iq2_xs, QK_K, quantize_iq2_xs>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
-            break;
-        case GGML_TYPE_IQ2_S:
-            set_rows_sycl_iq_host<TIn, TIdx, block_iq2_s, QK_K, quantize_iq2_s>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
-            break;
-        case GGML_TYPE_IQ3_XXS:
-            set_rows_sycl_qk_host<TIn, TIdx, block_iq3_xxs, QK_K, quantize_row_iq3_xxs_ref>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
-            break;
-        case GGML_TYPE_IQ3_S:
-            set_rows_sycl_qk_host<TIn, TIdx, block_iq3_s, QK_K, quantize_row_iq3_s_ref>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
-            break;
-        case GGML_TYPE_IQ1_S:
-            set_rows_sycl_iq_host<TIn, TIdx, block_iq1_s, QK_K, quantize_iq1_s>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
-            break;
-        case GGML_TYPE_IQ1_M:
-            set_rows_sycl_iq_host<TIn, TIdx, block_iq1_m, QK_K, quantize_iq1_m>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
-            break;
-        case GGML_TYPE_IQ4_XS:
-            set_rows_sycl_qk_host<TIn, TIdx, block_iq4_xs, QK_K, quantize_row_iq4_xs_ref>(
-                src0, src1, dst,
-                ne00, ne01, ne02, ne03,
-                ne11, ne12,
-                nb01, nb02, nb03,
-                nb10, nb11, nb12,
-                nb1, nb2, nb3,
-                stream);
+            set_rows_sycl_q<TIdx, block_nvfp4, QK_NVFP4, cpy_blck_f32_nvfp4>(src0_d, src1_d, (block_nvfp4 *)dst->data, ne00, ne01, ne02, ne03, ne10, ne11, ne12, ne13, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb1, nb2, nb3, stream);
            break;
        default:
            GGML_ABORT("Unsupported tensor type!");
@@ -556,21 +237,12 @@ void ggml_sycl_op_set_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
    const ggml_tensor * src0 = dst->src[0];
    const ggml_tensor * src1 = dst->src[1];

-    GGML_ASSERT(dst->src[0]->type == GGML_TYPE_F32 || dst->src[0]->type == GGML_TYPE_F16);
+    GGML_ASSERT(dst->src[0]->type == GGML_TYPE_F32);
    GGML_ASSERT(dst->src[1]->type == GGML_TYPE_I64 || dst->src[1]->type == GGML_TYPE_I32);

-    // dispatch on the index type (src1) and the source value type (src0)
-    if (src0->type == GGML_TYPE_F16) {
-        if (src1->type == GGML_TYPE_I64) {
-            set_rows_sycl<sycl::half, int64_t>(ctx, src0, src1, dst);
-        } else {
-            set_rows_sycl<sycl::half, int32_t>(ctx, src0, src1, dst);
-        }
+    if (src1->type == GGML_TYPE_I64) {
+        set_rows_sycl<float, int64_t>(ctx, src0, src1, dst);
    } else {
-        if (src1->type == GGML_TYPE_I64) {
-            set_rows_sycl<float, int64_t>(ctx, src0, src1, dst);
-        } else {
-            set_rows_sycl<float, int32_t>(ctx, src0, src1, dst);
-        }
+        set_rows_sycl<float, int32_t>(ctx, src0, src1, dst);
    }
 }
@@ -36,13 +36,9 @@ static void kernel_ssm_conv(
                    return;
                }

-                // src has the tokens of one channel contiguous, dst has the channels of one
-                // token contiguous, so either the loads or the store must be strided. Indexing
-                // token-fastest coalesces the d_conv loads, which measured faster except for
-                // short, cache-resident rows.
-                const int token   = static_cast<int>(idx % n_t);
-                const int channel = static_cast<int>((idx / n_t) % d_inner);
-                const int seq     = static_cast<int>(idx / (static_cast<size_t>(n_t) * static_cast<size_t>(d_inner)));
+                const int channel = static_cast<int>(idx % d_inner);
+                const int token   = static_cast<int>((idx / d_inner) % n_t);
+                const int seq     = static_cast<int>(idx / (static_cast<size_t>(d_inner) * static_cast<size_t>(n_t)));

                const float *s = src_data
                    + static_cast<size_t>(seq) * static_cast<size_t>(src_stride_seq)
@@ -186,22 +186,13 @@ static bool is_pow2(uint32_t x) { return x > 1 && (x & (x-1)) == 0; }

 #define VK_DEVICE_DESCRIPTOR_POOL_SIZE 256

-#define VK_CHECK(err, msg, dev)                                     \
+#define VK_CHECK(err, msg)                                          \
    do {                                                            \
-        vk::Result err_;                                            \
-        try {                                                       \
-            err_ = (err);                                           \
-        } catch (vk::DeviceLostError &) {                           \
-            ggml_vk_print_device_lost_info(dev);                    \
-            GGML_LOG_ERROR("ggml_vulkan: %s at %s:%d\n",            \
-                #err, __FILE__, __LINE__);                          \
-            throw;                                                  \
-        }                                                           \
+        vk::Result err_ = (err);                                    \
        if (err_ != vk::Result::eSuccess) {                         \
-            GGML_LOG_ERROR("ggml_vulkan: %s error %s at %s:%d\n",   \
+            fprintf(stderr, "ggml_vulkan: %s error %s at %s:%d\n",  \
                #err, to_string(err_).c_str(), __FILE__, __LINE__); \
-            throw vk::SystemError(vk::make_error_code(err_),        \
-                "ggml_vulkan: " msg);                               \
+            exit(1);                                                \
        }                                                           \
    } while (0)

@@ -311,13 +302,9 @@ struct vk_command_pool {
    }
 };

-static void ggml_vk_print_device_fault_info(const vk_device& device);
-static void ggml_vk_print_device_lost_info(const vk_device& device);
-
 // Prevent simultaneous submissions to the same queue.
 struct vk_queue_handle {
    vk::Queue queue;
-    vk_device_ref device;
    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() {}
@@ -328,14 +315,7 @@ 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);
-        try {
-            queue.submit(submits, fence);
-        } catch (vk::DeviceLostError &) {
-            if (auto dev = device.lock()) {
-                ggml_vk_print_device_lost_info(dev);
-            }
-            throw;
-        }
+        queue.submit(submits, fence);
    }
    void lock()   override { mutex.lock(); }
    void unlock() override { mutex.unlock(); }
@@ -344,14 +324,7 @@ struct vk_queue_handle_synchronized : vk_queue_handle {
 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
-        try {
-            queue.submit(submits, fence);
-        } catch (vk::DeviceLostError &) {
-            if (auto dev = device.lock()) {
-                ggml_vk_print_device_lost_info(dev);
-            }
-            throw;
-        }
+        queue.submit(submits, fence);
    }
    // lock()/unlock() inherited no-ops
 };
@@ -862,15 +835,6 @@ struct vk_device_struct {

    bool pipeline_executable_properties_support {};

-    bool device_fault {};
-    PFN_vkGetDeviceFaultInfoEXT pfn_vkGetDeviceFaultInfoEXT {};
-
-    bool serialize_submissions {};
-
-    const ggml_cgraph * diag_cgraph {};
-    int diag_prev_start = -1;
-    int diag_prev_end = -1;
-
    size_t idx;

    bool mul_mat_l[GGML_TYPE_COUNT];
@@ -1154,57 +1118,6 @@ void vk_command_pool::destroy(vk::Device& device) {
    cmd_buffers.clear();
 }

-static void ggml_vk_print_device_fault_info(const vk_device& device) {
-    if (!device->device_fault || !device->pfn_vkGetDeviceFaultInfoEXT) {
-        return;
-    }
-
-    VkDeviceFaultCountsEXT fault_counts {};
-    fault_counts.sType = VK_STRUCTURE_TYPE_DEVICE_FAULT_COUNTS_EXT;
-    VkResult res = device->pfn_vkGetDeviceFaultInfoEXT(device->device, &fault_counts, nullptr);
-    if (res != VK_SUCCESS) {
-        GGML_LOG_ERROR("ggml_vulkan: vkGetDeviceFaultInfoEXT (counts) failed: %d\n", res);
-        return;
-    }
-
-    std::vector<VkDeviceFaultAddressInfoEXT> address_infos(fault_counts.addressInfoCount);
-    std::vector<VkDeviceFaultVendorInfoEXT> vendor_infos(fault_counts.vendorInfoCount);
-
-    VkDeviceFaultInfoEXT fault_info {};
-    fault_info.sType = VK_STRUCTURE_TYPE_DEVICE_FAULT_INFO_EXT;
-    fault_info.pAddressInfos = address_infos.data();
-    fault_info.pVendorInfos = vendor_infos.data();
-
-    res = device->pfn_vkGetDeviceFaultInfoEXT(device->device, &fault_counts, &fault_info);
-    if (res != VK_SUCCESS) {
-        GGML_LOG_ERROR("ggml_vulkan: vkGetDeviceFaultInfoEXT (info) failed: %d\n", res);
-        return;
-    }
-
-    if (fault_counts.addressInfoCount == 0 && fault_counts.vendorInfoCount == 0 && fault_info.description[0] == '\0') {
-        return;
-    }
-
-    if (fault_info.description[0] != '\0') {
-        GGML_LOG_ERROR("ggml_vulkan: device fault on %s: %s\n", device->name.c_str(), fault_info.description);
-    }
-
-    for (uint32_t i = 0; i < fault_counts.addressInfoCount; i++) {
-        const auto& info = address_infos[i];
-        GGML_LOG_CONT("  address fault %u: type=%d address=0x%llx precision=0x%llx\n",
-                i, (int)info.addressType,
-                (unsigned long long)info.reportedAddress,
-                (unsigned long long)info.addressPrecision);
-    }
-    for (uint32_t i = 0; i < fault_counts.vendorInfoCount; i++) {
-        const auto& info = vendor_infos[i];
-        GGML_LOG_CONT("  vendor fault %u: %s (code=0x%llx data=0x%llx)\n",
-                i, info.description,
-                (unsigned long long)info.vendorFaultCode,
-                (unsigned long long)info.vendorFaultData);
-    }
-}
-
 struct vk_buffer_struct {
    vk::Buffer buffer = VK_NULL_HANDLE;
    vk::DeviceMemory device_memory = VK_NULL_HANDLE;
@@ -2146,36 +2059,6 @@ static uint64_t ggml_vk_get_node_flops(const ggml_tensor * node) {
    return 0;
 }

-static void ggml_vk_print_node_list(const ggml_cgraph * cgraph, int start, int end) {
-    uint64_t total_flops = 0;
-    int n_ops = 0;
-    for (int j = start; j <= end && j < cgraph->n_nodes; j++) {
-        uint64_t flops = ggml_vk_get_node_flops(cgraph->nodes[j]);
-        total_flops += flops;
-        n_ops++;
-        if (flops > 0) {
-            GGML_LOG_CONT("  node %d: %s (%s) [%.2f GFLOP]\n",
-                    j, cgraph->nodes[j]->name, ggml_op_name(cgraph->nodes[j]->op),
-                    flops / 1e9);
-        } else {
-            GGML_LOG_CONT("  node %d: %s (%s)\n",
-                    j, cgraph->nodes[j]->name, ggml_op_name(cgraph->nodes[j]->op));
-        }
-    }
-    GGML_LOG_CONT("  total: %d ops, %.2f GFLOP\n", n_ops, total_flops / 1e9);
-}
-
-static void ggml_vk_print_device_lost_info(const vk_device& device) {
-    ggml_vk_print_device_fault_info(device);
-    if (device->serialize_submissions && device->diag_cgraph != nullptr && device->diag_prev_start >= 0) {
-        GGML_LOG_ERROR("ggml_vulkan: device lost on %s, likely caused by previous submission (nodes %d to %d):\n",
-                device->name.c_str(), device->diag_prev_start, device->diag_prev_end);
-        ggml_vk_print_node_list(device->diag_cgraph, device->diag_prev_start, device->diag_prev_end);
-    } else {
-        GGML_LOG_ERROR("ggml_vulkan: device lost on %s\n", device->name.c_str());
-    }
-}
-
 class vk_perf_logger {
  public:
    void print_timings(bool force = false) {
@@ -2588,27 +2471,17 @@ static void ggml_vk_wait_for_fence(ggml_backend_vk_context * ctx) {
    // Use waitForFences while most of the graph executes. Hopefully the CPU can sleep
    // during this wait.
    if (ctx->almost_ready_fence_pending) {
-        VK_CHECK(ctx->device->device.waitForFences({ ctx->almost_ready_fence }, true, UINT64_MAX), "almost_ready_fence", ctx->device);
+        VK_CHECK(ctx->device->device.waitForFences({ ctx->almost_ready_fence }, true, UINT64_MAX), "almost_ready_fence");
        ctx->device->device.resetFences({ ctx->almost_ready_fence });
        ctx->almost_ready_fence_pending = false;
    }

    // Spin (w/pause) waiting for the graph to finish executing.
    vk::Result result;
-    for (;;) {
-        try {
-            result = ctx->device->device.getFenceStatus(ctx->fence);
-        } catch (vk::DeviceLostError &) {
-            ggml_vk_print_device_lost_info(ctx->device);
-            GGML_LOG_ERROR("ggml_vulkan: getFenceStatus at %s:%d\n", __FILE__, __LINE__);
-            throw;
-        }
-        if (result == vk::Result::eSuccess) {
-            break;
-        }
+    while ((result = ctx->device->device.getFenceStatus(ctx->fence)) != vk::Result::eSuccess) {
        if (result != vk::Result::eNotReady) {
-            GGML_LOG_ERROR("ggml_vulkan: error %s at %s:%d\n", to_string(result).c_str(), __FILE__, __LINE__);
-            throw vk::SystemError(vk::make_error_code(result), "ggml_vulkan: getFenceStatus");
+            fprintf(stderr, "ggml_vulkan: error %s at %s:%d\n", to_string(result).c_str(), __FILE__, __LINE__);
+            exit(1);
        }
        for (uint32_t i = 0; i < 100; ++i) {
            YIELD();
@@ -3299,7 +3172,6 @@ static std::unique_ptr<vk_queue> ggml_vk_create_queue(vk_device& device, uint32_
    }

    h->queue = device->device.getQueue2(queue_info2);
-    h->device = device;
    q->handle = h;

    q->cmd_pool.init(device, q.get());
@@ -6245,8 +6117,6 @@ static vk_device ggml_vk_get_device(size_t idx) {
 #endif
            } else if (strcmp(VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME, properties.extensionName) == 0) {
                internally_sync_support = true;
-            } else if (strcmp("VK_EXT_device_fault", properties.extensionName) == 0) {
-                device->device_fault = true;
            }
        }

@@ -6601,18 +6471,8 @@ static vk_device ggml_vk_get_device(size_t idx) {
        }
 #endif

-        VkPhysicalDeviceFaultFeaturesEXT fault_features {};
-        fault_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FAULT_FEATURES_EXT;
-        if (device->device_fault) {
-            last_struct->pNext = (VkBaseOutStructure *)&fault_features;
-            last_struct = (VkBaseOutStructure *)&fault_features;
-            device_extensions.push_back("VK_EXT_device_fault");
-        }
-
        vkGetPhysicalDeviceFeatures2(device->physical_device, &device_features2);

-        device->device_fault = device->device_fault && fault_features.deviceFault;
-
        device->has_internally_synchronized_queues = internally_synchronized_queues_features.internallySynchronizedQueues;

        // Build queue create infos only after querying whether internally synchronized queues are enabled.
@@ -6911,11 +6771,6 @@ static vk_device ggml_vk_get_device(size_t idx) {
        device_create_info.setPNext(&device_features2);
        device->device = device->physical_device.createDevice(device_create_info);

-        if (device->device_fault) {
-            device->pfn_vkGetDeviceFaultInfoEXT = (PFN_vkGetDeviceFaultInfoEXT)
-                vkGetDeviceProcAddr(device->device, "vkGetDeviceFaultInfoEXT");
-        }
-
        // Queues
        device->compute_queue = ggml_vk_create_queue(device, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false);

@@ -7038,8 +6893,6 @@ static vk_device ggml_vk_get_device(size_t idx) {

        device->idx = idx;

-        device->serialize_submissions = getenv("GGML_VK_SERIALIZE_SUBMISSIONS") != nullptr;
-
        device->disable_fusion = getenv("GGML_VK_DISABLE_FUSION") != nullptr;

        device->add_rms_fusion = !device->disable_fusion &&
@@ -8466,7 +8319,7 @@ static void ggml_vk_buffer_write_2d(vk_buffer& dst, size_t offset, const void *
        }

        ggml_vk_submit(subctx, dst->device->fence);
-        VK_CHECK(dst->device->device.waitForFences({ dst->device->fence }, true, UINT64_MAX), "vk_buffer_write_2d waitForFences", dst->device);
+        VK_CHECK(dst->device->device.waitForFences({ dst->device->fence }, true, UINT64_MAX), "vk_buffer_write_2d waitForFences");
        dst->device->device.resetFences({ dst->device->fence });
        ggml_vk_queue_command_pools_cleanup(dst->device);
    }
@@ -8578,7 +8431,7 @@ static void ggml_vk_buffer_read_2d(vk_buffer& src, size_t offset, void * dst, si
        ggml_vk_ctx_end(subctx);
        ggml_vk_submit(subctx, src->device->fence);
        VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX),
-                 "vk_buffer_read_2d uma waitForFences", src->device);
+                 "vk_buffer_read_2d uma waitForFences");
        src->device->device.resetFences({ src->device->fence });
        ggml_vk_queue_command_pools_cleanup(src->device);

@@ -8599,7 +8452,7 @@ static void ggml_vk_buffer_read_2d(vk_buffer& src, size_t offset, void * dst, si
        ggml_vk_ctx_end(subctx);

        ggml_vk_submit(subctx, src->device->fence);
-        VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX), "vk_buffer_read_2d waitForFences", src->device);
+        VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX), "vk_buffer_read_2d waitForFences");
        src->device->device.resetFences({ src->device->fence });
        ggml_vk_queue_command_pools_cleanup(src->device);

@@ -8634,7 +8487,7 @@ static void ggml_vk_buffer_copy(vk_buffer& dst, size_t dst_offset, vk_buffer& sr
        ggml_vk_buffer_copy_async(subctx, dst, dst_offset, src, src_offset, size);
        ggml_vk_ctx_end(subctx);
        ggml_vk_submit(subctx, src->device->fence);
-        VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX), "vk_buffer_copy waitForFences", src->device);
+        VK_CHECK(src->device->device.waitForFences({ src->device->fence }, true, UINT64_MAX), "vk_buffer_copy waitForFences");
        src->device->device.resetFences({ src->device->fence });
        ggml_vk_queue_command_pools_cleanup(src->device);
    } else {
@@ -8678,7 +8531,7 @@ static void ggml_vk_buffer_memset(vk_buffer& dst, size_t offset, uint32_t c, siz
    ggml_vk_ctx_end(subctx);

    ggml_vk_submit(subctx, dst->device->fence);
-    VK_CHECK(dst->device->device.waitForFences({ dst->device->fence }, true, UINT64_MAX), "vk_memset waitForFences", dst->device);
+    VK_CHECK(dst->device->device.waitForFences({ dst->device->fence }, true, UINT64_MAX), "vk_memset waitForFences");
    dst->device->device.resetFences({ dst->device->fence });
    ggml_vk_queue_command_pools_cleanup(dst->device);
 }
@@ -14413,7 +14266,7 @@ static void ggml_vk_test_matmul(ggml_backend_vk_context * ctx, size_t m, size_t

    auto begin = std::chrono::high_resolution_clock::now();
    ggml_vk_submit(subctx, ctx->fence);
-    VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_matmul waitForFences", ctx->device);
+    VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_matmul waitForFences");
    ctx->device->device.resetFences({ ctx->fence });
    ggml_vk_queue_command_pools_cleanup(ctx->device);

@@ -14615,7 +14468,7 @@ static void ggml_vk_test_dequant(ggml_backend_vk_context * ctx, size_t ne, ggml_
    auto begin = std::chrono::high_resolution_clock::now();

    ggml_vk_submit(subctx, ctx->fence);
-    VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_dequant waitForFences", ctx->device);
+    VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_dequant waitForFences");
    ctx->device->device.resetFences({ ctx->fence });
    ggml_vk_queue_command_pools_cleanup(ctx->device);

@@ -14901,7 +14754,7 @@ static void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m,
    auto begin = std::chrono::high_resolution_clock::now();

    ggml_vk_submit(subctx, ctx->fence);
-    VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_dequant waitForFences", ctx->device);
+    VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "ggml_vk_test_dequant waitForFences");
    ctx->device->device.resetFences({ ctx->fence });
    ggml_vk_queue_command_pools_cleanup(ctx->device);

@@ -15700,9 +15553,7 @@ static void ggml_vk_compute_forward(ggml_backend_vk_context * ctx, ggml_cgraph *
            memset(mset.dst, mset.val, mset.n);
        }

-        if (ctx->device->serialize_submissions) {
-            ggml_vk_submit(subctx, ctx->fence);
-        } else if (almost_ready && !ctx->almost_ready_fence_pending) {
+        if (almost_ready && !ctx->almost_ready_fence_pending) {
            ggml_vk_submit(subctx, ctx->almost_ready_fence);
            ctx->almost_ready_fence_pending = true;
        } else {
@@ -16313,20 +16164,12 @@ static void ggml_vk_synchronize(ggml_backend_vk_context * ctx) {
            memcpy(cpy.dst, cpy.src, cpy.n);
        }

-        if (ctx->device->serialize_submissions) {
-            ggml_vk_submit(compute_ctx, ctx->fence);
-            VK_CHECK(ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX), "synchronize waitForFences", ctx->device);
-            ctx->device->device.resetFences({ ctx->fence });
-        } else {
-            ggml_vk_submit(compute_ctx, {});
-        }
+        ggml_vk_submit(compute_ctx, {});
        ctx->submit_pending = true;
    }

    if (ctx->submit_pending) {
-        if (ctx->device->serialize_submissions) {
-            ctx->submit_pending = false;
-        } else if (ctx->device->async_use_transfer_queue && ctx->transfer_semaphore_last_submitted < ctx->transfer_semaphore.value) {
+        if (ctx->device->async_use_transfer_queue && ctx->transfer_semaphore_last_submitted < ctx->transfer_semaphore.value) {
            vk::TimelineSemaphoreSubmitInfo tl_info{
                1, &ctx->transfer_semaphore.value,
                0, nullptr,
@@ -16343,9 +16186,7 @@ static void ggml_vk_synchronize(ggml_backend_vk_context * ctx) {
        } else {
            ctx->device->compute_queue->handle->submit({}, ctx->fence);
        }
-        if (!ctx->device->serialize_submissions) {
-            ggml_vk_wait_for_fence(ctx);
-        }
+        ggml_vk_wait_for_fence(ctx);
        ctx->submit_pending = false;
        if (cmd_buf) {
            cmd_buf->in_use = false;
@@ -16917,10 +16758,6 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
    VK_LOG_DEBUG("ggml_backend_vk_graph_compute(" << cgraph->n_nodes << " nodes)");
    ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context;

-    ctx->device->diag_cgraph = nullptr;
-    ctx->device->diag_prev_start = -1;
-    ctx->device->diag_prev_end = -1;
-
    if (vk_instance.debug_utils_support) {
        vk::DebugUtilsLabelEXT dul = {};
        dul.pLabelName = "ggml_backend_vk_graph_compute";
@@ -17012,36 +16849,6 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
    }
    uint64_t flops_per_submit = std::min(flops_cap, ctx->last_total_flops / 40u);

-    auto const submit_after = [&](int start, int end) {
-        if (ctx->device->serialize_submissions) {
-            try {
-                auto res = ctx->device->device.waitForFences({ ctx->fence }, true, UINT64_MAX);
-                if (res != vk::Result::eSuccess) {
-                    GGML_LOG_ERROR("ggml_vulkan: waitForFences error during serialized submission\n");
-                    throw vk::SystemError(vk::make_error_code(res), "ggml_vulkan: waitForFences during serialized submission");
-                }
-            } catch (vk::DeviceLostError &) {
-                ggml_vk_print_device_fault_info(ctx->device);
-                GGML_LOG_ERROR("ggml_vulkan: device lost on %s waiting for submission (nodes %d to %d):\n",
-                        ctx->device->name.c_str(), start, end);
-                ggml_vk_print_node_list(cgraph, start, end);
-                throw;
-            }
-            ctx->device->device.resetFences({ ctx->fence });
-            ctx->submit_pending = false;
-            ctx->device->diag_cgraph = cgraph;
-            ctx->device->diag_prev_start = start;
-            ctx->device->diag_prev_end = end;
-        }
-        first_node_in_batch = true;
-        submitted_nodes = 0;
-        batch_flops = 0;
-        if (submit_count < 3) {
-            flops_per_submit *= 2;
-        }
-        submit_count++;
-    };
-
    for (int i = 0; i < cgraph->n_nodes; i++) {
        if (first_node_in_batch) {
            submit_node_idx = i;
@@ -17049,20 +16856,8 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg

        {
            auto node_flops = ggml_vk_get_node_flops(cgraph->nodes[i]);
-            total_flops += node_flops;
-
-            // Flush the current batch before recording a node that would push it over the flop threshold
-            if (flops_per_submit != 0 && submitted_nodes > 0 && batch_flops + node_flops >= flops_per_submit) {
-                vk_context flush_ctx = ggml_vk_get_compute_ctx(ctx);
-                ggml_vk_ctx_end(flush_ctx);
-                flush_ctx->exit_tensor_idx = -1;
-                ctx->compute_ctx.reset();
-                ggml_vk_compute_forward(ctx, cgraph, cgraph->nodes[submit_node_idx], submit_node_idx, false);
-                submit_after(submit_node_idx, i - 1);
-                submit_node_idx = i;
-            }
-
            batch_flops += node_flops;
+            total_flops += node_flops;
        }

        // op_srcs_fused_elementwise indicates whether an op's srcs all contribute to
@@ -17316,7 +17111,13 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
        }

        if (submit && enqueued) {
-            submit_after(submit_node_idx, i + (int)ctx->num_additional_fused_ops);
+            first_node_in_batch = true;
+            submitted_nodes = 0;
+            batch_flops = 0;
+            if (submit_count < 3) {
+                flops_per_submit *= 2;
+            }
+            submit_count++;
        }
        i += ctx->num_additional_fused_ops;
        ctx->num_additional_fused_ops = 0;
@@ -17332,13 +17133,13 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
        ggml_vk_ctx_end(compute_ctx);

        ggml_vk_submit(compute_ctx, ctx->device->fence);
-        VK_CHECK(ctx->device->device.waitForFences({ ctx->device->fence }, true, UINT64_MAX), "GGML_VULKAN_PERF waitForFences", ctx->device);
+        VK_CHECK(ctx->device->device.waitForFences({ ctx->device->fence }, true, UINT64_MAX), "GGML_VULKAN_PERF waitForFences");
        ctx->device->device.resetFences({ ctx->device->fence });
        ctx->compute_ctx.reset();

        // Get the results and pass them to the logger
        std::vector<uint64_t> timestamps(cgraph->n_nodes + 1);
-        VK_CHECK(ctx->device->device.getQueryPoolResults(ctx->query_pool, 0, ctx->query_idx, (cgraph->n_nodes + 1)*sizeof(uint64_t), timestamps.data(), sizeof(uint64_t), vk::QueryResultFlagBits::e64 | vk::QueryResultFlagBits::eWait), "get timestamp results", ctx->device);
+        VK_CHECK(ctx->device->device.getQueryPoolResults(ctx->query_pool, 0, ctx->query_idx, (cgraph->n_nodes + 1)*sizeof(uint64_t), timestamps.data(), sizeof(uint64_t), vk::QueryResultFlagBits::e64 | vk::QueryResultFlagBits::eWait), "get timestamp results");
        if (!vk_perf_logger_concurrent) {
            // Log each op separately
            for (int i = 1; i < ctx->query_idx; i++) {
@@ -18565,7 +18366,7 @@ static void ggml_backend_vk_device_event_synchronize(ggml_backend_dev_t dev, ggm
        vk::Semaphore sem = vkev->tl_semaphore.s;
        uint64_t val = vkev->tl_semaphore.value;
        vk::SemaphoreWaitInfo swi{vk::SemaphoreWaitFlags{}, sem, val};
-        VK_CHECK(device->device.waitSemaphores(swi, UINT64_MAX), "event_synchronize", device);
+        VK_CHECK(device->device.waitSemaphores(swi, UINT64_MAX), "event_synchronize");

        // Reset and move submitted events
        for (auto& event : vkev->events_submitted) {
@@ -7200,10 +7200,6 @@ void ggml_build_forward_expand(struct ggml_cgraph * cgraph, struct ggml_tensor *
    ggml_build_forward_impl(cgraph, tensor, true, true);
 }

-void ggml_build_forward_order(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor) {
-    ggml_build_forward_impl(cgraph, tensor, true, false);
-}
-
 void ggml_build_backward_expand(
        struct ggml_context *  ctx,
        struct ggml_cgraph  *  cgraph,
@@ -348,14 +348,15 @@ extern "C" {
    // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations
    //       https://github.com/ggml-org/llama.cpp/pull/7544
    struct llama_context_params {
-        uint32_t n_ctx;             // text context, 0 = from model
-        uint32_t n_batch;           // logical maximum batch size that can be submitted to llama_decode
-        uint32_t n_ubatch;          // physical maximum batch size
-        uint32_t n_seq_max;         // max number of sequences (i.e. distinct states for recurrent models)
-        uint32_t n_rs_seq;          // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL]
-        uint32_t n_outputs_max;     // max outputs in a ubatch (0 = n_batch)
-        int32_t  n_threads;         // number of threads to use for generation
-        int32_t  n_threads_batch;   // number of threads to use for batch processing
+        uint32_t n_ctx;                 // text context, 0 = from model
+        uint32_t n_batch;               // logical maximum batch size that can be submitted to llama_decode
+        uint32_t n_ubatch;              // physical maximum batch size
+        uint32_t n_seq_max;             // max number of sequences (i.e. distinct states for recurrent models)
+        uint32_t n_rs_seq;              // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL]
+        uint32_t n_outputs_max;         // max outputs in a ubatch (0 = n_batch)
+        uint32_t n_outputs_max_per_seq; // max outputs per sequence (0 = n_outputs_max)
+        int32_t  n_threads;             // number of threads to use for generation
+        int32_t  n_threads_batch;       // number of threads to use for batch processing

        enum llama_context_type      ctx_type;          // set the context type (e.g. MTP)
        enum llama_rope_scaling_type rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type`
@@ -1054,6 +1055,9 @@ extern "C" {
    //

    // Get the backend sampled token for the ith token.
+    // With multiple outputs, sampler state advances when the token is accepted,
+    // not when it is read through this function.
+    // When accepting multiple outputs, accept a contiguous prefix in output order.
    // Returns LLAMA_TOKEN_NULL if no token was sampled.
    LLAMA_API llama_token llama_get_sampled_token_ith(struct llama_context * ctx, int32_t i);

@@ -1270,9 +1274,12 @@ extern "C" {
        // [EXPERIMENTAL]
        // backend sampling interface:

-        // return true if the backend supports all ops needed by the sampler
+        // return true if the backend supports all ops needed by the sampler and can handle up to n_outputs_max_per_seq outputs per sequence
        // note: call once per sampler
-        bool (*backend_init)(struct llama_sampler * smpl, ggml_backend_buffer_type_t buft);
+        bool (*backend_init)(
+                struct llama_sampler       * smpl,
+                ggml_backend_buffer_type_t   buft,
+                uint32_t                     n_outputs_max_per_seq);

        // call after .backend_apply()
        void (*backend_accept)(
@@ -1290,6 +1297,9 @@ extern "C" {

        // called before graph execution to set inputs for the current ubatch
        void (*backend_set_input)(struct llama_sampler * smpl);
+
+        // called before rebuilding a sampling graph to clear any internal sampler state
+        void (*backend_reset)(struct llama_sampler * smpl);
    };

    struct llama_sampler {
@@ -1310,6 +1320,7 @@ extern "C" {
    LLAMA_API void                   llama_sampler_apply (      struct llama_sampler * smpl, llama_token_data_array * cur_p);
    LLAMA_API void                   llama_sampler_reset (      struct llama_sampler * smpl);
    LLAMA_API struct llama_sampler * llama_sampler_clone (const struct llama_sampler * smpl);
+    LLAMA_API void                   llama_sampler_copy  (const struct llama_sampler * src, struct llama_sampler * dst);
    // important: do not free if the sampler has been added to a llama_sampler_chain (via llama_sampler_chain_add)
    LLAMA_API void                   llama_sampler_free  (      struct llama_sampler * smpl);

@@ -1499,6 +1510,7 @@ extern "C" {
    LLAMA_API uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl);

    /// @details Sample and accept a token from the idx-th output of the last evaluation
+    // For multiple outputs from one sampler, call this function in output order without gaps.
    //
    // Shorthand for:
    //    const auto * logits = llama_get_logits_ith(ctx, idx);
@@ -1 +1 @@
-30bf8685ed4eb0a47f2b06229543327749904150
+90951f99af1fbebef3fbdd58ff5b8715b0bb9c43
@@ -10,6 +10,7 @@
 #include "llama-mmap.h"
 #include "llama-model.h"
 #include "llama-ext.h"
+#include "llama-sampler.h"
 #include "llama.h"

 #include <cinttypes>
@@ -159,25 +160,6 @@ llama_context::llama_context(
        }
    }

-    // Initialize backend samplers here so they are part of the sampling graph
-    // before the reserve passes run later in this function. This avoids a later
-    // re-reserve when graph nodes change.
-    if (params.samplers != nullptr && params.n_samplers > 0) {
-        for (size_t i = 0; i < params.n_samplers; ++i) {
-            const auto & config = params.samplers[i];
-
-            if (llama_sampler_chain_get(config.sampler, -1) == nullptr) {
-                throw std::runtime_error("the backend samplers must be of type llama_sampler_chain");
-            }
-
-            if (set_sampler(config.seq_id, config.sampler)) {
-                const int n_samplers = llama_sampler_chain_n(config.sampler);
-
-                LLAMA_LOG_INFO("%s: setting backend sampler for seq_id %d (n = %d)\n", __func__, config.seq_id, n_samplers);
-            }
-        }
-    }
-
    auto rope_scaling_type = params.rope_scaling_type;
    if (rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED) {
        rope_scaling_type = hparams.rope_scaling_type_train;
@@ -265,6 +247,27 @@ llama_context::llama_context(
    cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch);

    cparams.n_outputs_max = params.n_outputs_max == 0 || llama_model_has_encoder(&model) ? cparams.n_batch : params.n_outputs_max;
+    cparams.n_outputs_max_per_seq = params.n_outputs_max_per_seq == 0 ?
+            cparams.n_outputs_max : std::min(params.n_outputs_max_per_seq, cparams.n_outputs_max);
+
+    // Initialize backend samplers here so they are part of the sampling graph
+    // before the reserve passes run later in this function. This avoids a later
+    // re-reserve when graph nodes change.
+    if (params.samplers != nullptr && params.n_samplers > 0) {
+        for (size_t i = 0; i < params.n_samplers; ++i) {
+            const auto & config = params.samplers[i];
+
+            if (llama_sampler_chain_get(config.sampler, -1) == nullptr) {
+                throw std::runtime_error("the backend samplers must be of type llama_sampler_chain");
+            }
+
+            if (set_sampler(config.seq_id, config.sampler)) {
+                const int n_samplers = llama_sampler_chain_n(config.sampler);
+
+                LLAMA_LOG_INFO("%s: setting backend sampler for seq_id %d (n = %d)\n", __func__, config.seq_id, n_samplers);
+            }
+        }
+    }

    cparams.op_offload = params.op_offload;
    cparams.kv_unified = params.kv_unified;
@@ -300,18 +303,19 @@ llama_context::llama_context(
        }
    }

-    LLAMA_LOG_INFO("%s: n_seq_max     = %u\n",   __func__, cparams.n_seq_max);
-    LLAMA_LOG_INFO("%s: n_ctx         = %u\n",   __func__, cparams.n_ctx);
-    LLAMA_LOG_INFO("%s: n_ctx_seq     = %u\n",   __func__, cparams.n_ctx_seq);
-    LLAMA_LOG_INFO("%s: n_batch       = %u\n",   __func__, cparams.n_batch);
-    LLAMA_LOG_INFO("%s: n_ubatch      = %u\n",   __func__, cparams.n_ubatch);
-    LLAMA_LOG_INFO("%s: causal_attn   = %d\n",   __func__, cparams.causal_attn);
-    LLAMA_LOG_INFO("%s: flash_attn    = %s\n",   __func__, llama_flash_attn_type_name(params.flash_attn_type));
-    LLAMA_LOG_INFO("%s: kv_unified    = %s\n",   __func__, cparams.kv_unified ? "true" : "false");
-    LLAMA_LOG_INFO("%s: freq_base     = %.1f\n", __func__, cparams.rope_freq_base);
-    LLAMA_LOG_INFO("%s: freq_scale    = %g\n",   __func__, cparams.rope_freq_scale);
-    LLAMA_LOG_INFO("%s: n_rs_seq      = %u\n",   __func__, cparams.n_rs_seq);
-    LLAMA_LOG_INFO("%s: n_outputs_max = %u\n",   __func__, cparams.n_outputs_max);
+    LLAMA_LOG_INFO("%s: n_seq_max             = %u\n",   __func__, cparams.n_seq_max);
+    LLAMA_LOG_INFO("%s: n_ctx                 = %u\n",   __func__, cparams.n_ctx);
+    LLAMA_LOG_INFO("%s: n_ctx_seq             = %u\n",   __func__, cparams.n_ctx_seq);
+    LLAMA_LOG_INFO("%s: n_batch               = %u\n",   __func__, cparams.n_batch);
+    LLAMA_LOG_INFO("%s: n_ubatch              = %u\n",   __func__, cparams.n_ubatch);
+    LLAMA_LOG_INFO("%s: causal_attn           = %d\n",   __func__, cparams.causal_attn);
+    LLAMA_LOG_INFO("%s: flash_attn            = %s\n",   __func__, llama_flash_attn_type_name(params.flash_attn_type));
+    LLAMA_LOG_INFO("%s: kv_unified            = %s\n",   __func__, cparams.kv_unified ? "true" : "false");
+    LLAMA_LOG_INFO("%s: freq_base             = %.1f\n", __func__, cparams.rope_freq_base);
+    LLAMA_LOG_INFO("%s: freq_scale            = %g\n",   __func__, cparams.rope_freq_scale);
+    LLAMA_LOG_INFO("%s: n_rs_seq              = %u\n",   __func__, cparams.n_rs_seq);
+    LLAMA_LOG_INFO("%s: n_outputs_max         = %u\n",   __func__, cparams.n_outputs_max);
+    LLAMA_LOG_INFO("%s: n_outputs_max_per_seq = %u\n",   __func__, cparams.n_outputs_max_per_seq);

    if (cparams.n_ctx_seq < hparams.n_ctx_train) {
        LLAMA_LOG_INFO("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n",
@@ -1231,7 +1235,7 @@ bool llama_context::set_sampler(llama_seq_id seq_id, llama_sampler * sampler) {
    if (sampler && can_offload) {
        auto * buft = ggml_backend_dev_buffer_type(model.dev_output());

-        sampler->iface->backend_init(sampler, buft);
+        sampler->iface->backend_init(sampler, buft, cparams.n_outputs_max_per_seq);

        sampling.samplers[seq_id] = sampler;

@@ -1576,108 +1580,38 @@ int llama_context::encode(const llama_batch & batch_inp) {
    return 0;
 }

-static std::map<llama_seq_id, uint32_t> build_seq_to_output_row(const llama_ubatch & ubatch, uint32_t row_offset) {
-    std::map<llama_seq_id, uint32_t> seq_to_row;
-    // how many output tokens we have seen so far for this ubatch.
-    uint32_t local = 0;
-    for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
-        // skip tokens that are not output.
-        if (!ubatch.output[i]) {
-            continue;
-        }
-
-        const llama_seq_id seq_id = ubatch.seq_id[i][0];
-        // row_offset is the number of output tokens before this ubatch.
-        seq_to_row[seq_id] = row_offset + local;
-        ++local;
-    }
-    return seq_to_row;
-}
-
-static void copy_tensor_async_ints(
-    const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
-    const buffer_view<llama_token> & sampled,
-    const std::map<llama_seq_id, uint32_t> & seq_to_row,
-    ggml_backend_sched_t sched) {
-    if (!sampled.has_data()) {
-        return;
-    }
-
-    for (const auto & [seq_id, tensor] : tensor_map) {
-        auto it = seq_to_row.find(seq_id);
-        if (it == seq_to_row.end()) {
-            continue;
-        }
-
-        const uint32_t row = it->second;
-        GGML_ASSERT(row < sampled.size);
-
-        GGML_ASSERT(ggml_is_contiguous(tensor) && "sampled tokens tensor must be contiguous for async copy");
-
-        ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
-        ggml_backend_tensor_get_async(backend, tensor, sampled.data + row, 0, sizeof(sampled.data[row]));
-    }
-}
-
-static void copy_tensor_async_floats(
-    const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
-    const buffer_view<float> & dst,
+template<typename T>
+static void copy_tensor_async_rows(
+    const std::vector<ggml_tensor *> & tensors,
+    const buffer_view<T> & dst,
    size_t stride,
-    std::vector<uint32_t> & counts,
-    const std::map<llama_seq_id, uint32_t> & seq_to_row,
-    ggml_backend_sched_t sched) {
+    uint32_t row_offset,
+    ggml_backend_sched_t sched,
+    std::vector<uint32_t> * counts = nullptr) {
    if (!dst.has_data()) {
        return;
    }

-    for (const auto & [seq_id, tensor] : tensor_map) {
-        auto it = seq_to_row.find(seq_id);
-        if (it == seq_to_row.end()) {
+    for (size_t i = 0; i < tensors.size(); ++i) {
+        auto * tensor = tensors[i];
+        if (tensor == nullptr) {
            continue;
        }

-        const uint32_t row = it->second;
-        GGML_ASSERT(row < counts.size());
-
-        GGML_ASSERT(ggml_is_contiguous(tensor) && "logits/probs tensor must be contiguous for async copy");
+        const uint32_t row = row_offset + i;
+        const size_t n_elements = ggml_nelements(tensor);
+        GGML_ASSERT(ggml_is_contiguous(tensor) && "sampling tensor must be contiguous for async copy");
+        GGML_ASSERT(n_elements <= stride);
+        GGML_ASSERT((size_t) row * stride + n_elements <= dst.size);

        ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
-        float * row_ptr = dst.data + (size_t) row * stride;
+        T * row_ptr = dst.data + (size_t) row * stride;
        ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor));

-        // Update the actual number of logits/probabilities that were written for this row.
-        counts[row] = ggml_nelements(tensor);
-    }
-}
-
-static void copy_tensor_async_candidates(
-    const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
-    const buffer_view<llama_token> & dst,
-    size_t stride,
-    std::vector<uint32_t> & counts,
-    const std::map<llama_seq_id, uint32_t> & seq_to_row,
-    ggml_backend_sched_t sched) {
-    if (!dst.has_data()) {
-        return;
-    }
-
-    for (const auto & [seq_id, tensor] : tensor_map) {
-        auto it = seq_to_row.find(seq_id);
-        if (it == seq_to_row.end()) {
-            continue;
+        if (counts) {
+            GGML_ASSERT(row < counts->size());
+            (*counts)[row] = n_elements;
        }
-
-        const uint32_t row = it->second;
-        GGML_ASSERT(row < counts.size());
-
-        GGML_ASSERT(ggml_is_contiguous(tensor) && "candidates tensor must be contiguous for async copy");
-
-        ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
-        llama_token * row_ptr = dst.data + (size_t) row * stride;
-        ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor));
-
-        // Update the actual number of candidates that were written.
-        counts[row] = ggml_nelements(tensor);
    }
 }

@@ -1726,12 +1660,12 @@ int llama_context::decode(const llama_batch & batch_inp) {

    const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max;

-    // TODO: avoid this workaround in the future
-    if (has_samplers && batch_inp.logits) {
+    // embedding contexts output every token even when batch.logits is not set
+    if (has_samplers && (output_all || batch_inp.logits)) {
        std::vector<int32_t> seq_output_count(n_seq_max, 0);

        for (int32_t i = 0; i < batch_inp.n_tokens; ++i) {
-            if (batch_inp.logits[i] == 0) {
+            if (!output_all && batch_inp.logits[i] == 0) {
                continue;
            }

@@ -1740,10 +1674,17 @@ int llama_context::decode(const llama_batch & batch_inp) {
            for (int32_t s = 0; s < ns; ++s) {
                const llama_seq_id seq_id = batch_inp.seq_id ? batch_inp.seq_id[i][s] : 0;

+                if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) {
+                    continue;
+                }
+
                seq_output_count[seq_id]++;
-                if (seq_output_count[seq_id] > 1) {
-                    LLAMA_LOG_ERROR("%s: backend sampling requires at most one output token per sequence (seq_id %d had %d)\n",
-                            __func__, seq_id, seq_output_count[seq_id]);
+                auto sampler = sampling.samplers.find(seq_id);
+                if (sampler != sampling.samplers.end() &&
+                        seq_output_count[seq_id] > (int32_t) cparams.n_outputs_max_per_seq) {
+                    LLAMA_LOG_ERROR("%s: backend sampling supports at most %u outputs per sequence "
+                            "(seq_id %d had %d)\n", __func__, cparams.n_outputs_max_per_seq,
+                            seq_id, seq_output_count[seq_id]);
                    return -1;
                }
            }
@@ -1843,6 +1784,11 @@ int llama_context::decode(const llama_batch & batch_inp) {
        return -2;
    };

+    // start a new sampling transaction for this logical batch
+    for (const auto & entry : sampling.samplers) {
+        llama_sampler_backend_begin(entry.second);
+    }
+
    int64_t n_outputs_prev = 0;
    int64_t n_tokens_prev  = 0;

@@ -2009,17 +1955,14 @@ int llama_context::decode(const llama_batch & batch_inp) {
            }
        }

-        // Copy backend sampling output if this ubatch produced any sampling tensors.
-        if (has_samplers && (!res->t_sampled.empty() || !res->t_sampled_probs.empty() || !res->t_sampled_logits.empty())) {
-            const auto seq_to_output_row = build_seq_to_output_row(ubatch, n_outputs_prev);
+        if (has_samplers) {
            const auto stride = n_vocab;

            // async copy the sampling data from the backend to the host
-            copy_tensor_async_ints(res->t_sampled, sampling.sampled, seq_to_output_row, sched.get());
-
-            copy_tensor_async_floats    (res->t_sampled_logits, sampling.logits,     stride, sampling.logits_count,     seq_to_output_row, sched.get());
-            copy_tensor_async_floats    (res->t_sampled_probs,  sampling.probs,      stride, sampling.probs_count,      seq_to_output_row, sched.get());
-            copy_tensor_async_candidates(res->t_candidates,     sampling.candidates, stride, sampling.candidates_count, seq_to_output_row, sched.get());
+            copy_tensor_async_rows(res->t_sampled,        sampling.sampled,    1,      n_outputs_prev, sched.get());
+            copy_tensor_async_rows(res->t_sampled_logits, sampling.logits,     stride, n_outputs_prev, sched.get(), &sampling.logits_count);
+            copy_tensor_async_rows(res->t_sampled_probs,  sampling.probs,      stride, n_outputs_prev, sched.get(), &sampling.probs_count);
+            copy_tensor_async_rows(res->t_candidates,     sampling.candidates, stride, n_outputs_prev, sched.get(), &sampling.candidates_count);
        }

        n_outputs_prev += n_outputs;
@@ -2349,6 +2292,7 @@ void llama_context::output_reorder() {
 //

 uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
+    uint32_t res;
    if (model.arch == LLM_ARCH_QWEN3NEXT ||
        model.arch == LLM_ARCH_KIMI_LINEAR ||
        model.arch == LLM_ARCH_QWEN35 ||
@@ -2357,11 +2301,31 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
        (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) ||
        model.arch == LLM_ARCH_NANBEIGE ||
        model.arch == LLM_ARCH_MINIMAX_M3) {
-        return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
+        res = std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
+    } else {
+        res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
+        for (const auto & lora : model.loras) {
+            res += lora->get_n_nodes();
+        }
    }
-    uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
-    for (const auto & lora : model.loras) {
-        res += lora->get_n_nodes();
+
+    uint32_t n_sampling_nodes = 0;
+    uint32_t n_sampling_nodes_max = 0;
+    for (const auto & [seq_id, sampler] : sampling.samplers) {
+        const uint32_t n_nodes = llama_sampler_backend_n_nodes(sampler);
+        n_sampling_nodes += n_nodes;
+        if (cparams.n_outputs_max_per_seq > 1) {
+            n_sampling_nodes_max = std::max(n_sampling_nodes_max, n_nodes);
+        }
+    }
+
+    const uint32_t n_sampling_outputs_max = std::min<uint64_t>(
+            std::min(n_tokens, cparams.n_outputs_max),
+            (uint64_t) cparams.n_seq_max * cparams.n_outputs_max_per_seq);
+
+    res += n_sampling_nodes;
+    if (n_sampling_outputs_max > 1) {
+        res += (n_sampling_outputs_max - 1) * n_sampling_nodes_max;
    }
    return res;
 }
@@ -2394,13 +2358,57 @@ ggml_cgraph * llama_context::graph_reserve(
    llama_batch_allocr balloc(model.hparams.n_pos_per_embd());
    llama_ubatch ubatch = balloc.ubatch_reserve(n_tokens/n_seqs, n_seqs);

-    // set one output token per sequence in order to activate all backend samplers
+    // select sampler outputs first to reserve the largest valid sampling graph
    std::vector<llama_seq_id> seq_ids(n_seqs);
-    for (uint32_t i = 0; i < n_seqs; ++i) {
-        seq_ids[i] = i;
-        ubatch.n_seq_id[i] = 1;
-        ubatch.seq_id[i] = &seq_ids[i];
-        ubatch.output[i] = true;
+    for (uint32_t s = 0; s < n_seqs; ++s) {
+        seq_ids[s] = s;
+        for (uint32_t t = 0; t < ubatch.n_seq_tokens; ++t) {
+            const uint32_t i = s * ubatch.n_seq_tokens + t;
+            ubatch.n_seq_id[i] = 1;
+            ubatch.seq_id[i] = &seq_ids[s];
+        }
+    }
+
+    uint32_t n_outputs_set = 0;
+
+    std::vector<uint32_t> sampler_seqs;
+    std::vector<bool> has_sampler(n_seqs, false);
+    for (const auto & entry : sampling.samplers) {
+        const llama_seq_id seq_id = entry.first;
+        if (seq_id < 0 || (uint32_t) seq_id >= n_seqs) {
+            continue;
+        }
+
+        sampler_seqs.push_back(seq_id);
+        has_sampler[seq_id] = true;
+    }
+
+    const uint32_t n_sampling_outputs_per_seq = std::min(
+            ubatch.n_seq_tokens, cparams.n_outputs_max_per_seq);
+
+    // select sampling rows in round-robin order across sampler sequences
+    if (!sampler_seqs.empty()) {
+        const uint32_t n_sampler_seqs = sampler_seqs.size();
+        n_outputs_set = std::min<uint64_t>(
+                n_outputs, (uint64_t) n_sampler_seqs * n_sampling_outputs_per_seq);
+
+        for (uint32_t i = 0; i < n_outputs_set; ++i) {
+            const uint32_t s = sampler_seqs[i % n_sampler_seqs];
+            const uint32_t t = i / n_sampler_seqs;
+            ubatch.output[s * ubatch.n_seq_tokens + t] = true;
+        }
+    }
+
+    // use sequences without samplers for any remaining outputs
+    for (uint32_t t = 0; t < ubatch.n_seq_tokens && n_outputs_set < n_outputs; ++t) {
+        for (uint32_t s = 0; s < n_seqs && n_outputs_set < n_outputs; ++s) {
+            if (has_sampler[s]) {
+                continue;
+            }
+
+            ubatch.output[s * ubatch.n_seq_tokens + t] = true;
+            ++n_outputs_set;
+        }
    }

    auto * res = gf_res_reserve.get();
@@ -3488,6 +3496,7 @@ llama_context_params llama_context_default_params() {
        /*.n_seq_max                   =*/ 1,
        /*.n_rs_seq                    =*/ 0,
        /*.n_outputs_max               =*/ 0,
+        /*.n_outputs_max_per_seq       =*/ 1,
        /*.n_threads                   =*/ GGML_DEFAULT_N_THREADS, // TODO: better default
        /*.n_threads_batch             =*/ GGML_DEFAULT_N_THREADS,
        /*.ctx_type                    =*/ LLAMA_CONTEXT_TYPE_DEFAULT,
@@ -15,6 +15,7 @@ struct llama_cparams {
    uint32_t n_seq_max;
    uint32_t n_rs_seq;        // number of recurrent-state snapshots per seq for rollback
    uint32_t n_outputs_max;   // max outputs supported by the context
+    uint32_t n_outputs_max_per_seq;
    int32_t  n_threads;       // number of threads to use for generation
    int32_t  n_threads_batch; // number of threads to use for batch processing

@@ -4,6 +4,7 @@
 #include "llama-model.h"
 #include "llama-batch.h"
 #include "llama-cparams.h"
+#include "llama-sampler.h"

 #include "llama-kv-cache.h"
 #include "llama-kv-cache-iswa.h"
@@ -1353,24 +1354,24 @@ void llm_graph_result::set_outputs(const llm_graph_params & params) {
            }
        }
    }
-    for (auto & [seq_id, t] : t_sampled) {
-        if (t != nullptr) {
-            ggml_set_output(t);
+    for (auto * tensor : t_sampled) {
+        if (tensor != nullptr) {
+            ggml_set_output(tensor);
        }
    }
-    for (auto & [seq_id, t] : t_sampled_probs) {
-        if (t != nullptr) {
-            ggml_set_output(t);
+    for (auto * tensor : t_sampled_probs) {
+        if (tensor != nullptr) {
+            ggml_set_output(tensor);
        }
    }
-    for (auto & [seq_id, t] : t_sampled_logits) {
-        if (t != nullptr) {
-            ggml_set_output(t);
+    for (auto * tensor : t_sampled_logits) {
+        if (tensor != nullptr) {
+            ggml_set_output(tensor);
        }
    }
-    for (auto & [seq_id, t] : t_candidates) {
-        if (t != nullptr) {
-            ggml_set_output(t);
+    for (auto * tensor : t_candidates) {
+        if (tensor != nullptr) {
+            ggml_set_output(tensor);
        }
    }
 }
@@ -3649,77 +3650,102 @@ void llm_graph_context::build_sampling() const {
    auto inp_sampling = std::make_unique<llm_graph_input_sampling>(samplers);
    res->add_input(std::move(inp_sampling));

-    std::map<llama_seq_id, int32_t> seq_to_logit_row;
-    int32_t logit_row_idx = 0;
-
-    for (uint32_t i = 0; i < ubatch.n_tokens; i++) {
+    std::map<llama_seq_id, std::vector<uint32_t>> sampling_rows;
+    uint32_t n_rows = 0;
+    for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
        if (ubatch.output[i]) {
-            llama_seq_id seq_id = ubatch.seq_id[i][0];
-            seq_to_logit_row[seq_id] = logit_row_idx;
-            logit_row_idx++;
+            sampling_rows[ubatch.seq_id[i][0]].push_back(n_rows++);
        }
    }

+    res->t_sampled.resize(n_rows, nullptr);
+    res->t_sampled_probs.resize(n_rows, nullptr);
+    res->t_sampled_logits.resize(n_rows, nullptr);
+    res->t_candidates.resize(n_rows, nullptr);
+
    // res->t_logits will contain logits for all tokens that want the logits calculated (logits=1 or output=1)
    GGML_ASSERT(res->t_logits != nullptr && "missing t_logits tensor");

-    // add a dummy row of logits
-    // this trick makes the graph static, regardless of which samplers are activated
-    // this is important in order to minimize graph reallocations
+    // add a dummy row to keep the single-output graph static regardless of active samplers
+    // multi-output graphs can still vary with the number of output rows
    ggml_tensor * logits_t = ggml_pad(ctx0, res->t_logits, 0, 1, 0, 0);

-    for (const auto & [seq_id, sampler] : samplers) {
-        const auto it = seq_to_logit_row.find(seq_id);
-
-        // inactive samplers always work on the first row
-        const auto row_idx = it != seq_to_logit_row.end() ? it->second : 0;
-        const int i_out    = it != seq_to_logit_row.end() ? 1          : 0;
-
-        ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], row_idx * logits_t->nb[1]);
-        ggml_format_name(logits_seq, "logits_seq_%d", seq_id);
-
-        struct llama_sampler_data data = {
-            /*.logits      =*/ logits_seq,
-            /*.probs       =*/ nullptr,
-            /*.sampled     =*/ nullptr,
-            /*.candidates  =*/ nullptr,
-        };
-
-        assert(sampler->iface->backend_apply);
-        sampler->iface->backend_apply(sampler, ctx0, gf, &data);
-
-        if (data.sampled != nullptr) {
-            res->t_sampled[seq_id] = data.sampled;
-            outs[1] = data.sampled;
-            ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
-        }
-
-        if (data.probs != nullptr) {
-            res->t_sampled_probs[seq_id] = data.probs;
-            outs[1] = data.probs;
-            ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
-        }
-
-        if (data.logits != nullptr) {
-            res->t_sampled_logits[seq_id] = data.logits;
-            outs[1] = data.logits;
-            ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
-        }
-
-        if (data.candidates != nullptr) {
-            res->t_candidates[seq_id] = data.candidates;
-            outs[1] = data.candidates;
-            ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
+    for (const auto & entry : samplers) {
+        if (entry.second->iface->backend_reset) {
+            entry.second->iface->backend_reset(entry.second);
        }
    }

-    // TODO: Call llama_sampler_accept_ggml after all samplers have been applied.
+    static const std::vector<uint32_t> dummy_row = { 0 };
+
+    for (const auto & [seq_id, sampler] : samplers) {
+        const auto it = sampling_rows.find(seq_id);
+
+        // inactive samplers always work on the first row
+        const bool active = it != sampling_rows.end();
+        const auto & rows = active ? it->second : dummy_row;
+        const int i_out   = active ? 1          : 0;
+
+        for (uint32_t i = 0; i < rows.size(); ++i) {
+            ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], rows[i] * logits_t->nb[1]);
+            ggml_format_name(logits_seq, "logits_seq_%d_%u", seq_id, i);
+
+            struct llama_sampler_data data = {
+                /*.logits       =*/ logits_seq,
+                /*.probs        =*/ nullptr,
+                /*.sampled      =*/ nullptr,
+                /*.candidates   =*/ nullptr,
+            };
+
+            assert(sampler->iface->backend_apply);
+            sampler->iface->backend_apply(sampler, ctx0, gf, &data);
+
+            if (data.sampled != nullptr) {
+                if (active) {
+                    res->t_sampled[rows[i]] = data.sampled;
+                }
+                outs[1] = data.sampled;
+                ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
+            }
+
+            if (data.probs != nullptr) {
+                if (active) {
+                    res->t_sampled_probs[rows[i]] = data.probs;
+                }
+                outs[1] = data.probs;
+                ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
+            }
+
+            if (data.logits != nullptr) {
+                if (active) {
+                    res->t_sampled_logits[rows[i]] = data.logits;
+                }
+                outs[1] = data.logits;
+                ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
+            }
+
+            if (data.candidates != nullptr) {
+                if (active) {
+                    res->t_candidates[rows[i]] = data.candidates;
+                }
+                outs[1] = data.candidates;
+                ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
+            }
+        }
+    }
+
+    // TODO: Call backend_accept after all samplers have been applied.
    /*
    for (const auto & [seq_id, sampler] : samplers) {
-        if (auto it = res->t_sampled.find(seq_id); it != res->t_sampled.end()) {
-            ggml_tensor * selected_token = it->second;
-            if (selected_token != nullptr) {
-                llama_sampler_accept_ggml(sampler, ctx0, gf, selected_token);
+        const auto it = sampling_rows.find(seq_id);
+        if (it == sampling_rows.end()) {
+            continue;
+        }
+
+        for (uint32_t row : it->second) {
+            ggml_tensor * selected_token = res->t_sampled[row];
+            if (selected_token != nullptr && sampler->iface->backend_accept) {
+                sampler->iface->backend_accept(sampler, ctx0, gf, selected_token);
            }
        }
    }
@@ -904,10 +904,10 @@ public:

    std::vector<ggml_tensor *> t_layer_inp;

-    std::map<llama_seq_id, ggml_tensor *> t_sampled_logits;
-    std::map<llama_seq_id, ggml_tensor *> t_candidates;
-    std::map<llama_seq_id, ggml_tensor *> t_sampled;
-    std::map<llama_seq_id, ggml_tensor *> t_sampled_probs;
+    std::vector<ggml_tensor *> t_sampled;
+    std::vector<ggml_tensor *> t_sampled_probs;
+    std::vector<ggml_tensor *> t_sampled_logits;
+    std::vector<ggml_tensor *> t_candidates;

    std::vector<llm_graph_input_ptr> inputs;
    std::vector<llm_graph_fused_node> fused_nodes;
@@ -1249,13 +1249,7 @@ struct ggml_tensor * llama_model_loader::create_tensor(
        for (size_t dim = 0; dim < GGML_MAX_DIMS; dim++) {
            t_meta.ne[dim] = dim < ne.size() ? ne.begin()[dim] : 1;
            GGML_ASSERT(t_meta.ne[dim] >= 1);
-            if (dim == 0) {
-                t_meta.nb[dim] = ggml_type_size(type);
-            } else if (dim == 1) {
-                t_meta.nb[dim] = ggml_row_size(type, t_meta.ne[dim-1]);
-            } else {
-                t_meta.nb[dim] = t_meta.nb[dim-1]*t_meta.ne[dim-1];
-            }
+            t_meta.nb[dim] = dim == 0 ? ggml_type_size(type) : t_meta.ne[dim-1]*t_meta.nb[dim-1];
            GGML_ASSERT(t_meta.nb[dim] >= 1);
        }
        ggml_set_name(&t_meta, tn.str().c_str());
@@ -1278,18 +1272,10 @@ struct ggml_tensor * llama_model_loader::create_tensor(
    if (flags & TENSOR_ALLOW_RESHAPE) {
        for (size_t dim = 0; dim < GGML_MAX_DIMS; dim++) {
            t_meta.ne[dim] = dim < ne.size() ? ne.begin()[dim] : 1;
-            if (dim == 0) {
-                t_meta.nb[dim] = ggml_type_size(t_meta.type);
-            } else if (dim == 1) {
-                t_meta.nb[dim] = ggml_row_size(t_meta.type, t_meta.ne[dim-1]);
-            } else {
-                t_meta.nb[dim] = t_meta.ne[dim-1]*t_meta.nb[dim-1];
-            }
+            t_meta.nb[dim] = dim == 0 ? ggml_type_size(t_meta.type) : t_meta.ne[dim-1]*t_meta.nb[dim-1];
        }
    }

-    GGML_ASSERT(ggml_nbytes(&t_meta) == ggml_nbytes(cur));
-
    ggml_backend_buffer_type_t buft = buft_for_tensor(&t_meta);
    if (buft == nullptr) {
        return nullptr;
@@ -467,9 +467,11 @@ static void llama_sampler_empty_free(struct llama_sampler * smpl) {

 static bool llama_sampler_empty_backend_init(
        struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
    GGML_UNUSED(smpl);
    GGML_UNUSED(buft);
+    GGML_UNUSED(n_outputs_max_per_seq);

    return true;
 }
@@ -511,6 +513,7 @@ static struct llama_sampler_i llama_sampler_empty_i = {
    /* .backend_accept    = */ llama_sampler_empty_backend_accept,
    /* .backend_apply     = */ llama_sampler_empty_backend_apply,
    /* .backend_set_input = */ llama_sampler_empty_backend_set_input,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_empty(const char * name) {
@@ -559,6 +562,65 @@ private:
    bool support;
 };

+struct llama_sampler_backend_probe {
+    ggml_context_ptr ctx;
+    ggml_cgraph * gf;
+};
+
+static llama_sampler_backend_probe llama_sampler_backend_probe_graph(
+        llama_sampler * sampler,
+        int64_t         n_candidates,
+        uint32_t        max_nodes,
+        bool            with_candidates) {
+    ggml_init_params params = {
+        /*.mem_size   =*/ max_nodes * ggml_tensor_overhead() + ggml_graph_overhead_custom(max_nodes, false),
+        /*.mem_buffer =*/ nullptr,
+        /*.no_alloc   =*/ true,
+    };
+
+    ggml_context_ptr ctx_ptr { ggml_init(params) };
+    if (!ctx_ptr) {
+        throw std::runtime_error(format("failed to create ggml context"));
+    }
+
+    auto * ctx = ctx_ptr.get();
+    auto * gf = ggml_new_graph_custom(ctx, max_nodes, false);
+
+    llama_sampler_data data = {
+        /*.logits       =*/ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_candidates),
+        /*.probs        =*/ nullptr,
+        /*.sampled      =*/ nullptr,
+        /*.candidates   =*/ with_candidates ? ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_candidates) : nullptr,
+    };
+
+    if (sampler->iface->backend_reset) {
+        sampler->iface->backend_reset(sampler);
+    }
+    sampler->iface->backend_apply(sampler, ctx, gf, &data);
+
+    for (auto * output : { data.logits, data.probs, data.sampled, data.candidates }) {
+        if (output) {
+            ggml_build_forward_expand(gf, output);
+        }
+    }
+
+    if (sampler->iface->backend_reset) {
+        sampler->iface->backend_reset(sampler);
+    }
+
+    return { std::move(ctx_ptr), gf };
+}
+
+static uint32_t llama_sampler_backend_probe_n_nodes(const llama_sampler_backend_probe & probe) {
+    uint32_t n_tensors = 0;
+    for (auto * tensor = ggml_get_first_tensor(probe.ctx.get()); tensor;
+            tensor = ggml_get_next_tensor(probe.ctx.get(), tensor)) {
+        ++n_tensors;
+    }
+
+    return std::max<uint32_t>(ggml_graph_n_nodes(probe.gf), n_tensors);
+}
+
 // check if all ggml ops used by the sampler are supported by the backend
 static bool llama_sampler_backend_support(
        llama_sampler              * smpl,
@@ -569,50 +631,10 @@ static bool llama_sampler_backend_support(
        return true;
    }

-    ggml_init_params params = {
-        /*.mem_size   =*/ 128*ggml_tensor_overhead() + ggml_graph_overhead(),
-        /*.mem_buffer =*/ NULL,
-        /*.no_alloc   =*/ true,
-    };
+    auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, true);

-    ggml_context_ptr ctx_ptr { ggml_init(params) };
-    if (!ctx_ptr) {
-        throw std::runtime_error(format("failed to create ggml context"));
-    }
-
-    ggml_context * ctx = ctx_ptr.get();
-
-    const int64_t n = 1024*1024;
-
-    llama_sampler_data data = {
-        /*.logits     = */ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n),
-        /*.probs      = */ nullptr,
-        /*.sampled    = */ nullptr,
-        /*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n),
-    };
-
-    ggml_cgraph * gf = ggml_new_graph(ctx);
-
-    smpl->iface->backend_apply(smpl, ctx, gf, &data);
-
-    if (data.logits) {
-        ggml_build_forward_expand(gf, data.logits);
-    }
-
-    if (data.probs) {
-        ggml_build_forward_expand(gf, data.probs);
-    }
-
-    if (data.sampled) {
-        ggml_build_forward_expand(gf, data.sampled);
-    }
-
-    if (data.candidates) {
-        ggml_build_forward_expand(gf, data.candidates);
-    }
-
-    for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
-        struct ggml_tensor * op = ggml_graph_node(gf, i);
+    for (int i = 0; i < ggml_graph_n_nodes(probe.gf); i++) {
+        struct ggml_tensor * op = ggml_graph_node(probe.gf, i);

        if (!ggml_backend_dev_supports_op(device, op)) {
            LLAMA_LOG_WARN("%s: device '%s' does not have support for op %s needed for sampler '%s'\n",
@@ -678,8 +700,10 @@ static struct llama_sampler * llama_sampler_chain_clone(const struct llama_sampl

    auto * result = llama_sampler_chain_init(chain_src->params);

-    for (const auto & smpl : chain_src->samplers) {
-        llama_sampler_chain_add(result, llama_sampler_clone(smpl.ptr));
+    auto * chain_dst = (llama_sampler_chain *) result->ctx;
+    *chain_dst = *chain_src;
+    for (size_t i = 0; i < chain_src->samplers.size(); ++i) {
+        chain_dst->samplers[i].ptr = llama_sampler_clone(chain_src->samplers[i].ptr);
    }

    return result;
@@ -697,7 +721,8 @@ static void llama_sampler_chain_free(struct llama_sampler * smpl) {

 static bool llama_sampler_chain_backend_init(
        struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
    auto * chain = (llama_sampler_chain *) smpl->ctx;

    GGML_ASSERT(chain->is_init == false && "llama_sampler_chain_backend_init() called twice");
@@ -705,26 +730,32 @@ static bool llama_sampler_chain_backend_init(
    chain->is_init = true;

    bool res = true;
+    bool backend_prefix = true;

    for (auto & smpl : chain->samplers) {
-        bool res_cur = true;
+        bool cur_prefix = backend_prefix;

        // to be able to run a sampler on the backend, it has to:
        // - have the .backend_init() API implemented
        // - return true during .backend_init()
-        if (smpl.ptr->iface->backend_init) {
-            if (!smpl.ptr->iface->backend_init(smpl.ptr, buft)) {
-                res_cur = false;
+        // - support the requested per-sequence output limit
+        if (cur_prefix && smpl.ptr->iface->backend_init) {
+            if (!smpl.ptr->iface->backend_init(smpl.ptr, buft, n_outputs_max_per_seq)) {
+                cur_prefix = false;
            }
        } else {
-            res_cur = false;
+            cur_prefix = false;
        }

-        smpl.is_backend = res_cur;
+        smpl.is_backend = cur_prefix;
+        backend_prefix = cur_prefix;

-        res = res && res_cur;
+        res = res && cur_prefix;
    }

+    auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, false);
+    chain->n_nodes = llama_sampler_backend_probe_n_nodes(probe);
+
    return res;
 }

@@ -780,6 +811,19 @@ static void llama_sampler_chain_backend_set_input(struct llama_sampler * smpl) {
    }
 }

+static void llama_sampler_chain_backend_reset(struct llama_sampler * smpl) {
+    auto * chain = (llama_sampler_chain *) smpl->ctx;
+
+    for (auto & entry : chain->samplers) {
+        if (!entry.is_backend) {
+            break;
+        }
+        if (entry.ptr->iface->backend_reset) {
+            entry.ptr->iface->backend_reset(entry.ptr);
+        }
+    }
+}
+
 static struct llama_sampler_i llama_sampler_chain_i = {
    /* .name              = */ llama_sampler_chain_name,
    /* .accept            = */ llama_sampler_chain_accept,
@@ -791,22 +835,34 @@ static struct llama_sampler_i llama_sampler_chain_i = {
    /* .backend_accept    = */ llama_sampler_chain_backend_accept,
    /* .backend_apply     = */ llama_sampler_chain_backend_apply,
    /* .backend_set_input = */ llama_sampler_chain_backend_set_input,
+    /* .backend_reset     = */ llama_sampler_chain_backend_reset,
 };

 struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_params params) {
    return llama_sampler_init(
        /* .iface = */ &llama_sampler_chain_i,
        /* .ctx   = */ new llama_sampler_chain {
-            /* .params      = */ params,
-            /* .is_init     = */ false,
-            /* .samplers    = */ {},
-            /* .cur         = */ {},
-            /* .t_sample_us = */ 0,
-            /* .n_sample    = */ 0,
+            /* .params               = */ params,
+            /* .is_init              = */ false,
+            /* .n_nodes              = */ 0,
+            /* .samplers             = */ {},
+            /* .cur                  = */ {},
+            /* .t_sample_us          = */ 0,
+            /* .n_sample             = */ 0,
        }
    );
 }

+uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler) {
+    GGML_ASSERT(sampler != nullptr);
+    GGML_ASSERT(sampler->iface == &llama_sampler_chain_i);
+
+    const auto * chain = (const llama_sampler_chain *) sampler->ctx;
+    GGML_ASSERT(chain->is_init);
+
+    return chain->n_nodes;
+}
+
 llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_context * ctx, int32_t idx) {
    const llama_token   sampled_token  = llama_get_sampled_token_ith     (ctx, idx);
    const float *       sampled_probs  = llama_get_sampled_probs_ith     (ctx, idx);
@@ -816,6 +872,7 @@ llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_conte
    // If a backend sampler has already sampled a token, return it.
    if (sampled_token != LLAMA_TOKEN_NULL) {
        LLAMA_LOG_DEBUG("%s: Backend sampler selected token for idx %d. Skipping CPU samplers\n", __func__, idx);
+        llama_sampler_accept(smpl, sampled_token);
        return sampled_token;
    }

@@ -975,8 +1032,10 @@ static void llama_sampler_greedy_apply(struct llama_sampler * /*smpl*/, llama_to

 static bool llama_sampler_greedy_backend_init(
        struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
    auto * sctx = (llama_sampler_greedy *) smpl->ctx;
+    GGML_UNUSED(n_outputs_max_per_seq);

    const bool res = llama_sampler_backend_support(smpl, buft);

@@ -1012,6 +1071,7 @@ static struct llama_sampler_i llama_sampler_greedy_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ llama_sampler_greedy_backend_apply,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_greedy() {
@@ -1031,7 +1091,14 @@ struct llama_sampler_dist : public llama_sampler_backend {

    std::mt19937 rng;

-    ggml_tensor * inp_uniform;
+    // multi-output backend draws are committed as an accepted prefix
+    bool backend_transactional;
+    std::mt19937 rng_backend;
+    size_t n_backend_draws_generated;
+    size_t n_backend_draws_committed;
+
+    // inputs for the current sampling graph
+    std::vector<ggml_tensor *> inp_uniforms;
 };

 static const char * llama_sampler_dist_name(const struct llama_sampler * smpl) {
@@ -1050,7 +1117,11 @@ static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_da

    cur_p->selected = 0;

+    std::uniform_real_distribution<double> dist(0.0f, 1.0f);
+
    if (cur_p->size == 1) {
+        // keep the RNG state aligned with backend sampling, which draws once per output
+        dist(ctx->rng);
        cur_p->data[0].p = 1.0f;
        return;
    }
@@ -1075,7 +1146,6 @@ static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_da
    // sample from the obtained probabilities and normalize the probs in a single pass
    // this is ~3x faster on Mac with full gpt-oss vocab than the version below
    //
-    std::uniform_real_distribution<double> dist(0.0f, 1.0f);
    const double rnd = dist(ctx->rng);

          double sum_run = 0.0f;
@@ -1115,6 +1185,9 @@ static void llama_sampler_dist_reset(struct llama_sampler * smpl) {
    auto * ctx = (llama_sampler_dist *) smpl->ctx;
    ctx->seed_cur = get_rng_seed(ctx->seed);
    ctx->rng.seed(ctx->seed_cur);
+    ctx->rng_backend = ctx->rng;
+    ctx->n_backend_draws_generated = 0;
+    ctx->n_backend_draws_committed = 0;
 }

 static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sampler * smpl) {
@@ -1125,7 +1198,12 @@ static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sample
    {
        auto * result_ctx = (llama_sampler_dist *) result->ctx;

-        result_ctx->rng = ctx->rng;
+        result_ctx->seed_cur                  = ctx->seed_cur;
+        result_ctx->rng                       = ctx->rng;
+        result_ctx->backend_transactional     = ctx->backend_transactional;
+        result_ctx->rng_backend               = ctx->rng_backend;
+        result_ctx->n_backend_draws_generated = ctx->n_backend_draws_generated;
+        result_ctx->n_backend_draws_committed = ctx->n_backend_draws_committed;
    }

    return result;
@@ -1137,12 +1215,17 @@ static void llama_sampler_dist_free(struct llama_sampler * smpl) {

 static bool llama_sampler_dist_backend_init(
        struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
    auto * sctx = (llama_sampler_dist *) smpl->ctx;

    const bool res = llama_sampler_backend_support(smpl, buft);

    sctx->init(res);
+    sctx->backend_transactional = n_outputs_max_per_seq > 1;
+    sctx->rng_backend = sctx->rng;
+    sctx->n_backend_draws_generated = 0;
+    sctx->n_backend_draws_committed = 0;

    return res;
 }
@@ -1156,9 +1239,10 @@ static void llama_sampler_dist_backend_apply(

    auto * sctx = (llama_sampler_dist *) smpl->ctx;

-    sctx->inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-    ggml_set_name (sctx->inp_uniform, "uniform");
-    ggml_set_input(sctx->inp_uniform);
+    ggml_tensor * inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+    ggml_format_name(inp_uniform, "uniform_%zu", sctx->inp_uniforms.size());
+    ggml_set_input(inp_uniform);
+    sctx->inp_uniforms.push_back(inp_uniform);

    // flatten
    struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
@@ -1174,7 +1258,7 @@ static void llama_sampler_dist_backend_apply(
    // Recall that each entry in cumsum is the cumulative probability up to that
    // index so values stay negative while the cumulative total is below the
    // random value, and become zero/positive once the threshold is crossed.
-    struct ggml_tensor * diff = ggml_sub(ctx, cumsum, sctx->inp_uniform);
+    struct ggml_tensor * diff = ggml_sub(ctx, cumsum, inp_uniform);
    ggml_set_name(diff, "dist_cumsum");

    // The ggml_step function produces a tensor where entries are 1 if the
@@ -1189,6 +1273,9 @@ static void llama_sampler_dist_backend_apply(
    struct ggml_tensor * idxf = ggml_sum(ctx, mask);
    ggml_set_name(idxf, "dist_index_f32");

+    // Clamp to prevent out-of-bounds access when computing the index.
+    idxf = ggml_clamp(ctx, idxf, 1.0f, mask->ne[0]);
+
    // Use ggml_scale_bias to scale the index value by -1 and then add the size
    // of the mask to that value so we get the correct index ((-1 * idxf) + n).
    struct ggml_tensor * idx = ggml_cast(ctx, ggml_scale_bias(ctx, idxf, -1.0f, mask->ne[0]), GGML_TYPE_I32);
@@ -1210,22 +1297,52 @@ static void llama_sampler_dist_backend_apply(
 static void llama_sampler_dist_backend_set_input(struct llama_sampler * smpl) {
    auto * sctx = (llama_sampler_dist *) smpl->ctx;

-    GGML_ASSERT(sctx->inp_uniform != nullptr);
+    GGML_ASSERT(!sctx->inp_uniforms.empty());

    // We sample in double precision and cast to float to match rnd numbers of
-    // llama_dampler_dist which uses double precision (sampling from
+    // llama_sampler_dist which uses double precision (sampling from
    // std::uniform_real_distribution<double> and
    // std::uniform_real_distribution<float> with same rng will produce
    // different sequences).
    std::uniform_real_distribution<double> dist(0.0f, 1.0f);
-    const float rnd = dist(sctx->rng);

-    ggml_backend_tensor_set(sctx->inp_uniform, &rnd, 0, sizeof(float));
+    auto & rng = sctx->backend_transactional ? sctx->rng_backend : sctx->rng;
+
+    for (auto * inp_uniform : sctx->inp_uniforms) {
+        GGML_ASSERT(inp_uniform != nullptr);
+
+        const float rnd = dist(rng);
+        ggml_backend_tensor_set(inp_uniform, &rnd, 0, sizeof(float));
+
+        if (sctx->backend_transactional) {
+            ++sctx->n_backend_draws_generated;
+        }
+    }
+}
+
+static void llama_sampler_dist_backend_reset(struct llama_sampler * smpl) {
+    auto * sctx = (llama_sampler_dist *) smpl->ctx;
+    sctx->inp_uniforms.clear();
+}
+
+static void llama_sampler_dist_accept(struct llama_sampler * smpl, llama_token token) {
+    GGML_UNUSED(token);
+
+    auto * sctx = (llama_sampler_dist *) smpl->ctx;
+
+    if (!sctx->backend_transactional ||
+            sctx->n_backend_draws_committed >= sctx->n_backend_draws_generated) {
+        return;
+    }
+
+    std::uniform_real_distribution<double> dist(0.0f, 1.0f);
+    dist(sctx->rng);
+    ++sctx->n_backend_draws_committed;
 }

 static struct llama_sampler_i llama_sampler_dist_i = {
    /* .name              = */ llama_sampler_dist_name,
-    /* .accept            = */ nullptr,
+    /* .accept            = */ llama_sampler_dist_accept,
    /* .apply             = */ llama_sampler_dist_apply,
    /* .reset             = */ llama_sampler_dist_reset,
    /* .clone             = */ llama_sampler_dist_clone,
@@ -1234,6 +1351,7 @@ static struct llama_sampler_i llama_sampler_dist_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ llama_sampler_dist_backend_apply,
    /* .backend_set_input = */ llama_sampler_dist_backend_set_input,
+    /* .backend_reset     = */ llama_sampler_dist_backend_reset,
 };

 struct llama_sampler * llama_sampler_init_dist(uint32_t seed) {
@@ -1242,14 +1360,39 @@ struct llama_sampler * llama_sampler_init_dist(uint32_t seed) {
        /* .iface = */ &llama_sampler_dist_i,
        /* .ctx   = */ new llama_sampler_dist {
            ("dist"),
-            /* .seed        = */ seed,
-            /* .seed_cur    = */ seed_cur,
-            /* .rng         = */ std::mt19937(seed_cur),
-            /* .inp_uniform = */ nullptr,
+            /* .seed                      = */ seed,
+            /* .seed_cur                  = */ seed_cur,
+            /* .rng                       = */ std::mt19937(seed_cur),
+            /* .backend_transactional     = */ false,
+            /* .rng_backend               = */ std::mt19937(seed_cur),
+            /* .n_backend_draws_generated = */ 0,
+            /* .n_backend_draws_committed = */ 0,
+            /* .inp_uniforms              = */ {},
        }
    );
 }

+void llama_sampler_backend_begin(llama_sampler * sampler) {
+    GGML_ASSERT(sampler != nullptr);
+
+    if (sampler->iface == &llama_sampler_chain_i) {
+        auto * chain = (llama_sampler_chain *) sampler->ctx;
+        for (auto & entry : chain->samplers) {
+            if (!entry.is_backend) {
+                break;
+            }
+            llama_sampler_backend_begin(entry.ptr);
+        }
+    } else if (sampler->iface == &llama_sampler_dist_i) {
+        auto * ctx = (llama_sampler_dist *) sampler->ctx;
+        if (ctx->backend_transactional) {
+            ctx->rng_backend = ctx->rng;
+            ctx->n_backend_draws_generated = 0;
+            ctx->n_backend_draws_committed = 0;
+        }
+    }
+}
+
 // top-k

 struct llama_sampler_top_k : public llama_sampler_backend {
@@ -1277,8 +1420,10 @@ static void llama_sampler_top_k_free(struct llama_sampler * smpl) {

 static bool llama_sampler_top_k_backend_init(
        struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
    auto * sctx = (llama_sampler_top_k *) smpl->ctx;
+    GGML_UNUSED(n_outputs_max_per_seq);

    const bool res = llama_sampler_backend_support(smpl, buft);

@@ -1325,6 +1470,7 @@ static struct llama_sampler_i llama_sampler_top_k_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ llama_sampler_top_k_backend_apply,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_top_k(int32_t k) {
@@ -1423,8 +1569,10 @@ static void llama_sampler_top_p_free(struct llama_sampler * smpl) {

 static bool llama_sampler_top_p_backend_init(
        struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
    auto * sctx = (llama_sampler_top_p *) smpl->ctx;
+    GGML_UNUSED(n_outputs_max_per_seq);

    const bool res = llama_sampler_backend_support(smpl, buft);

@@ -1521,6 +1669,7 @@ static struct llama_sampler_i llama_sampler_top_p_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ llama_sampler_top_p_backend_apply,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_top_p(float p, size_t min_keep) {
@@ -1618,8 +1767,10 @@ static void llama_sampler_min_p_free(struct llama_sampler * smpl) {

 static bool llama_sampler_min_p_backend_init(
        struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
    auto * sctx = (llama_sampler_min_p *) smpl->ctx;
+    GGML_UNUSED(n_outputs_max_per_seq);

    const bool res = llama_sampler_backend_support(smpl, buft);

@@ -1680,6 +1831,7 @@ static struct llama_sampler_i llama_sampler_min_p_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ llama_sampler_min_p_backend_apply,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_min_p(float p, size_t min_keep) {
@@ -1790,6 +1942,7 @@ static struct llama_sampler_i llama_sampler_typical_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ nullptr,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_typical(float p, size_t min_keep) {
@@ -1866,8 +2019,10 @@ static void llama_sampler_backend_temp_sampling(

 static bool llama_sampler_temp_backend_init(
        struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
    auto * sctx = (llama_sampler_temp *) smpl->ctx;
+    GGML_UNUSED(n_outputs_max_per_seq);

    const bool res = llama_sampler_backend_support(smpl, buft);

@@ -1896,6 +2051,7 @@ static struct llama_sampler_i llama_sampler_temp_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ llama_sampler_temp_backend_apply,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_temp(float temp) {
@@ -2009,8 +2165,10 @@ static void llama_sampler_temp_ext_free(struct llama_sampler * smpl) {

 static bool llama_sampler_temp_ext_backend_init(
        struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
    auto * sctx = (llama_sampler_temp_ext *) smpl->ctx;
+    GGML_UNUSED(n_outputs_max_per_seq);

    const bool res = llama_sampler_backend_support(smpl, buft);

@@ -2095,6 +2253,7 @@ static struct llama_sampler_i llama_sampler_temp_ext_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ llama_sampler_temp_ext_backend_apply,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_temp_ext(float temp, float delta, float exponent) {
@@ -2202,6 +2361,7 @@ static struct llama_sampler_i llama_sampler_xtc_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ nullptr,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_xtc(float p, float t, size_t min_keep, uint32_t seed) {
@@ -2290,7 +2450,7 @@ static struct llama_sampler * llama_sampler_mirostat_clone(const struct llama_sa

    // copy the state
    {
-        auto * result_ctx = (llama_sampler_mirostat *) smpl->ctx;
+        auto * result_ctx = (llama_sampler_mirostat *) result->ctx;

        result_ctx->mu  = ctx->mu;
        result_ctx->rng = ctx->rng;
@@ -2321,6 +2481,7 @@ static struct llama_sampler_i llama_sampler_mirostat_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ nullptr,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_mirostat(int32_t n_vocab, uint32_t seed, float tau, float eta, int32_t m) {
@@ -2425,6 +2586,7 @@ static struct llama_sampler_i llama_sampler_mirostat_v2_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ nullptr,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_mirostat_v2(uint32_t seed, float tau, float eta) {
@@ -2546,6 +2708,7 @@ static struct llama_sampler_i llama_sampler_grammar_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ nullptr,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 static struct llama_sampler * llama_sampler_init_grammar_impl(
@@ -2790,9 +2953,15 @@ static void llama_sampler_penalties_free(struct llama_sampler * smpl) {

 static bool llama_sampler_penalties_backend_init(
        struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
    auto * sctx = (llama_sampler_penalties *) smpl->ctx;

+    if (n_outputs_max_per_seq > 1) {
+        sctx->init(false);
+        return false;
+    }
+
    const bool res = llama_sampler_backend_support(smpl, buft);

    sctx->init(res);
@@ -2952,6 +3121,12 @@ static void llama_sampler_penalties_backend_set_input(struct llama_sampler * smp
    ggml_backend_tensor_set(sctx->inp_counts,    sctx->host_counts.data(),    0, sctx->n_max * sizeof(int32_t));
 }

+static void llama_sampler_penalties_backend_reset(struct llama_sampler * smpl) {
+    auto * sctx = (llama_sampler_penalties *) smpl->ctx;
+    sctx->inp_token_ids = nullptr;
+    sctx->inp_counts    = nullptr;
+}
+
 static struct llama_sampler_i llama_sampler_penalties_i = {
    /* .name              = */ llama_sampler_penalties_name,
    /* .accept            = */ llama_sampler_penalties_accept,
@@ -2963,6 +3138,7 @@ static struct llama_sampler_i llama_sampler_penalties_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ llama_sampler_penalties_backend_apply,
    /* .backend_set_input = */ llama_sampler_penalties_backend_set_input,
+    /* .backend_reset     = */ llama_sampler_penalties_backend_reset,
 };

 struct llama_sampler * llama_sampler_init_penalties(
@@ -3058,6 +3234,7 @@ static struct llama_sampler_i llama_sampler_top_n_sigma_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ nullptr,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_top_n_sigma(float n) {
@@ -3395,6 +3572,7 @@ static struct llama_sampler_i llama_sampler_dry_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ nullptr,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const char** seq_breakers, size_t num_breakers) {
@@ -3614,6 +3792,7 @@ static struct llama_sampler_i llama_sampler_adaptive_p_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ nullptr,
    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_adaptive_p(
@@ -3715,13 +3894,17 @@ static void llama_sampler_logit_bias_backend_apply(

    const size_t n = sctx->logit_bias.size();

-    sctx->inp_logit_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n);
-    ggml_set_name(sctx->inp_logit_bias, "logit_bias");
-    ggml_set_input(sctx->inp_logit_bias);
+    if (sctx->inp_logit_bias == nullptr) {
+        GGML_ASSERT(sctx->inp_logit_idxs == nullptr);

-    sctx->inp_logit_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n);
-    ggml_set_name(sctx->inp_logit_idxs, "logit_idxs");
-    ggml_set_input(sctx->inp_logit_idxs);
+        sctx->inp_logit_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n);
+        ggml_set_name(sctx->inp_logit_bias, "logit_bias");
+        ggml_set_input(sctx->inp_logit_bias);
+
+        sctx->inp_logit_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n);
+        ggml_set_name(sctx->inp_logit_idxs, "logit_idxs");
+        ggml_set_input(sctx->inp_logit_idxs);
+    }

    ggml_tensor * cur = ggml_fill(ctx, data->logits, 0.0f);

@@ -3756,10 +3939,18 @@ static void llama_sampler_logit_bias_backend_set_input(struct llama_sampler * sm
    ggml_backend_tensor_set(sctx->inp_logit_idxs, data_logit_idxs.data(), 0, ggml_nbytes(sctx->inp_logit_idxs));
 }

+static void llama_sampler_logit_bias_backend_reset(struct llama_sampler * smpl) {
+    auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
+    sctx->inp_logit_bias = nullptr;
+    sctx->inp_logit_idxs = nullptr;
+}
+
 static bool llama_sampler_logit_bias_backend_init(
        struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
    GGML_UNUSED(buft);
+    GGML_UNUSED(n_outputs_max_per_seq);

    auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;

@@ -3783,6 +3974,7 @@ static struct llama_sampler_i llama_sampler_logit_bias_i = {
    /* .backend_accept    = */ nullptr,
    /* .backend_apply     = */ llama_sampler_logit_bias_backend_apply,
    /* .backend_set_input = */ llama_sampler_logit_bias_backend_set_input,
+    /* .backend_reset     = */ llama_sampler_logit_bias_backend_reset,
 };

 struct llama_sampler * llama_sampler_init_logit_bias(
@@ -4022,10 +4214,11 @@ static struct llama_sampler_i llama_sampler_infill_i = {
    /* .reset             = */ nullptr,
    /* .clone             = */ llama_sampler_infill_clone,
    /* .free              = */ llama_sampler_infill_free,
-    /* .backend_apply     = */ nullptr,
-    /* .backend_accept    = */ nullptr,
-    /* .backend_set_input = */ nullptr,
    /* .backend_init      = */ nullptr,
+    /* .backend_accept    = */ nullptr,
+    /* .backend_apply     = */ nullptr,
+    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
 };

 struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab) {
@@ -4039,6 +4232,27 @@ struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * voca
    );
 }

+void llama_sampler_copy(const struct llama_sampler * src, struct llama_sampler * dst) {
+    if (!src || !dst || src == dst) {
+        return;
+    }
+
+    GGML_ASSERT(src->iface == dst->iface && "llama_sampler_copy: cannot copy between different sampler types");
+
+    // build a temporary sampler carrying src's current state
+    llama_sampler * tmp = llama_sampler_clone(src);
+
+    // free dst's old state (frees dst->ctx, including children for a chain)
+    if (dst->iface->free) {
+        dst->iface->free(dst);
+    }
+
+    // transplant tmp's state into dst, then destroy the (now empty) temp shell
+    dst->ctx = tmp->ctx;
+    tmp->ctx = nullptr;
+    delete tmp;
+}
+
 // utils

 uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl) {
@@ -15,6 +15,8 @@ struct llama_sampler_chain {
    // has .backend_init() been called?
    bool is_init = false;

+    uint32_t n_nodes = 0;
+
    struct info {
        bool is_backend;

@@ -33,6 +35,9 @@ struct llama_sampler_chain {
    mutable int32_t n_sample;
 };

+uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler);
+void llama_sampler_backend_begin(llama_sampler * sampler);
+
 struct llama_sampler * llama_sampler_init_dry_testing(
        float   dry_multiplier,
        float   dry_base,
@@ -2,7 +2,9 @@
 #include "common.h"
 #include "download.h"
 #include "llama.h"
+#include "speculative.h"

+#include <limits>
 #include <string>
 #include <vector>
 #include <sstream>
@@ -14,6 +16,34 @@
 static void test(void) {
    common_params params;

+    auto assert_output_limits = [](int32_t n_batch, int32_t n_parallel, int32_t n_draft,
+                                   int32_t total, int32_t per_seq) {
+        const auto limits = common_speculative_get_output_limits(n_batch, n_parallel, n_draft);
+        assert(limits.total == total);
+        assert(limits.per_seq == per_seq);
+    };
+
+    assert_output_limits(16, 2,  3, 8, 4);
+    assert_output_limits(16, 2, -1, 2, 1);
+    assert_output_limits( 6, 2,  3, 6, 4);
+    assert_output_limits( 2, 1,  3, 2, 2);
+    assert_output_limits(
+            std::numeric_limits<int32_t>::max(),
+            std::numeric_limits<int32_t>::max(),
+            std::numeric_limits<int32_t>::max(),
+            std::numeric_limits<int32_t>::max(),
+            std::numeric_limits<int32_t>::max());
+
+    {
+        common_params base;
+        base.n_parallel = 4;
+        base.n_outputs_max_per_seq = 8;
+
+        const auto draft = common_base_params_to_speculative(base);
+        assert(draft.n_outputs_max == 4);
+        assert(draft.n_outputs_max_per_seq == 1);
+    }
+
    printf("test-arg-parser: make sure there is no duplicated arguments in any examples\n\n");
    for (int ex = 0; ex < LLAMA_EXAMPLE_COUNT; ex++) {
        try {
@@ -2584,7 +2584,6 @@ struct test_rms_norm_mul_rope : public test_case {
    const float eps;
    const bool multi_add; // test a sequence of adds feeding into rms_norm
    const bool set_rows;
-    const bool broadcast; // multiply by a 1D [ne0] weight, as model norm weights are
    int mode;

    std::string op_desc(ggml_tensor * t) override {
@@ -2595,12 +2594,12 @@ struct test_rms_norm_mul_rope : public test_case {
    bool run_whole_graph() override { return true; }

    std::string vars() override {
-        return VARS_TO_STR6(ne, eps, multi_add, set_rows, broadcast, mode);
+        return VARS_TO_STR5(ne, eps, multi_add, set_rows, mode);
    }

    test_rms_norm_mul_rope(std::array<int64_t, 4> ne, float eps = 1e-6f, bool multi_add = false,
-                           bool set_rows = false, bool broadcast = false, int mode = GGML_ROPE_TYPE_NORMAL)
-        : ne(ne), eps(eps), multi_add(multi_add), set_rows(set_rows), broadcast(broadcast), mode(mode) {}
+                           bool set_rows = false, int mode = GGML_ROPE_TYPE_NORMAL)
+        : ne(ne), eps(eps), multi_add(multi_add), set_rows(set_rows), mode(mode) {}

    ggml_tensor * build_graph(ggml_context * ctx) override {
        ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], 1);
@@ -2611,9 +2610,7 @@ struct test_rms_norm_mul_rope : public test_case {
            a = ggml_add(ctx, ggml_add(ctx, a, b), c);
        }

-        ggml_tensor * w = broadcast ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[0]) : b;
-
-        a = ggml_mul(ctx, ggml_rms_norm(ctx, a, eps), w);
+        a = ggml_mul(ctx, ggml_rms_norm(ctx, a, eps), b);

        ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne[2]);

@@ -8579,9 +8576,6 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
            test_cases.emplace_back(new test_cpy(type_src, type_dst, {256, 2, 3, 4}, {-1,-1,-1,-1}, {1, 0, 2, 3})); // cpy not-contiguous
        }
    }
-    // quant block count not a multiple of the kernel block size
-    test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_Q4_0, {96, 1, 1, 1}));
-    test_cases.emplace_back(new test_cpy(GGML_TYPE_Q4_0, GGML_TYPE_F32, {96, 1, 1, 1}));
    test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_I32, {256, 2, 3, 4}));
    test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_I32, {256, 2, 3, 4}, {-1,-1,-1,-1}, {1, 0, 2, 3}));
    test_cases.emplace_back(new test_cpy(GGML_TYPE_I32, GGML_TYPE_F32, {256, 2, 3, 4}));
@@ -8728,13 +8722,6 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
            test_cases.emplace_back(new test_l2_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, true));
            test_cases.emplace_back(new test_l2_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, false, true));
        }
-        // row lengths that are not a multiple of 32, for the scalar (33) and float4 (132, 260) paths
-        for (uint32_t n : { 33, 132, 260 }) {
-            for (bool v : { false, true }) {
-                test_cases.emplace_back(new test_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, v, eps));
-                test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, v, eps));
-            }
-        }
    }

    // in-place tests
@@ -8759,18 +8746,16 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {

    for (auto multi_add : {false, true}) {
        for (auto set_rows : {false, true}) {
-            for (auto broadcast : {false, true}) {
-                for (auto rope : {GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX}) {
-                    test_cases.emplace_back(new test_rms_norm_mul_rope({768, 1, 1, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
-                    test_cases.emplace_back(new test_rms_norm_mul_rope({768, 3, 1, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
-                    test_cases.emplace_back(new test_rms_norm_mul_rope({768, 3, 5, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
-                    test_cases.emplace_back(new test_rms_norm_mul_rope({128, 32, 2, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
-                    test_cases.emplace_back(new test_rms_norm_mul_rope({128, 4, 2, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
-                    test_cases.emplace_back(new test_rms_norm_mul_rope({128, 32, 50, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
-                    test_cases.emplace_back(new test_rms_norm_mul_rope({128, 4, 50, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
-                    test_cases.emplace_back(new test_rms_norm_mul_rope({8192, 2, 2, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
-                    test_cases.emplace_back(new test_rms_norm_mul_rope({8192, 2, 2, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
-                }
+            for (auto rope : {GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX}) {
+                test_cases.emplace_back(new test_rms_norm_mul_rope({768, 1, 1, 1}, 1e-6f, multi_add, set_rows, rope));
+                test_cases.emplace_back(new test_rms_norm_mul_rope({768, 3, 1, 1}, 1e-6f, multi_add, set_rows, rope));
+                test_cases.emplace_back(new test_rms_norm_mul_rope({768, 3, 5, 1}, 1e-6f, multi_add, set_rows, rope));
+                test_cases.emplace_back(new test_rms_norm_mul_rope({128, 32, 2, 1}, 1e-6f, multi_add, set_rows, rope));
+                test_cases.emplace_back(new test_rms_norm_mul_rope({128, 4, 2, 1}, 1e-6f, multi_add, set_rows, rope));
+                test_cases.emplace_back(new test_rms_norm_mul_rope({128, 32, 50, 1}, 1e-6f, multi_add, set_rows, rope));
+                test_cases.emplace_back(new test_rms_norm_mul_rope({128, 4, 50, 1}, 1e-6f, multi_add, set_rows, rope));
+                test_cases.emplace_back(new test_rms_norm_mul_rope({8192, 2, 2, 1}, 1e-6f, multi_add, set_rows, rope));
+                test_cases.emplace_back(new test_rms_norm_mul_rope({8192, 2, 2, 1}, 1e-6f, multi_add, set_rows, rope));
            }
        }
    }
@@ -9762,15 +9747,6 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
 static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
    std::vector<std::unique_ptr<test_case>> test_cases;

-    // SWIGLU at a 27B-class FFN width, fused [gate|up] vs split operands
-    // note: same bytes either way, so a backend that indexes them differently shows it here
-    for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {
-        for (int64_t n_tokens : {512, 2048}) {
-            test_cases.emplace_back(new test_glu(GGML_GLU_OP_SWIGLU, type, { 2*17408, n_tokens, 1, 1 }, 0, false));
-            test_cases.emplace_back(new test_glu_split(GGML_GLU_OP_SWIGLU, type, { 17408, n_tokens, 1, 1 }, 0));
-        }
-    }
-
    // Conv2d: K=CRS=NPQ=4096 matmul performance
    uint32_t                        iwh_idx  = 0;
    uint32_t                        kwh_idx  = 1;
@@ -14,6 +14,7 @@
 #include <fstream>
 #include <functional>
 #include <map>
+#include <random>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
@@ -80,7 +81,13 @@ struct test_context {
    std::unordered_map<llama_seq_id, int32_t> seq_positions;
    std::unordered_map<llama_seq_id, int32_t> last_batch_info;

-    test_context(const test_params & params, std::vector<llama_sampler_seq_config> & configs, int32_t n_seq_max = -1) {
+    test_context(
+            const test_params & params,
+            std::vector<llama_sampler_seq_config> & configs,
+            int32_t n_seq_max = -1,
+            uint32_t n_outputs_max = 0,
+            uint32_t n_ubatch = 0,
+            uint32_t n_outputs_max_per_seq = 1) {
        auto * model = params.model.get();

        GGML_ASSERT(model);
@@ -89,6 +96,11 @@ struct test_context {
        llama_context_params cparams = llama_context_default_params();
        cparams.n_ctx = 512;
        cparams.n_batch = 512;
+        if (n_ubatch > 0) {
+            cparams.n_ubatch = n_ubatch;
+        }
+        cparams.n_outputs_max = n_outputs_max;
+        cparams.n_outputs_max_per_seq = n_outputs_max_per_seq;
        cparams.samplers = configs.data();
        cparams.n_samplers = configs.size();
        cparams.kv_unified = true;
@@ -262,6 +274,65 @@ struct test_context {
    }
 };

+struct test_single_output_backend_sampler {
+    bool backend_initialized = false;
+    uint32_t backend_outputs_max_per_seq = 0;
+    int backend_apply_count = 0;
+    int apply_count = 0;
+};
+
+static const char * test_single_output_backend_sampler_name(const llama_sampler * /*smpl*/) {
+    return "single-output-backend";
+}
+
+static void test_single_output_backend_sampler_apply(
+        llama_sampler * smpl, llama_token_data_array * /*cur_p*/) {
+    auto * ctx = (test_single_output_backend_sampler *) smpl->ctx;
+    ctx->apply_count++;
+}
+
+static void test_single_output_backend_sampler_free(llama_sampler * smpl) {
+    delete (test_single_output_backend_sampler *) smpl->ctx;
+}
+
+static bool test_single_output_backend_sampler_backend_init(
+        llama_sampler * smpl, ggml_backend_buffer_type_t /*buft*/, uint32_t n_outputs_max_per_seq) {
+    auto * ctx = (test_single_output_backend_sampler *) smpl->ctx;
+    ctx->backend_outputs_max_per_seq = n_outputs_max_per_seq;
+    if (n_outputs_max_per_seq > 1) {
+        return false;
+    }
+    ctx->backend_initialized = true;
+    return true;
+}
+
+static void test_single_output_backend_sampler_backend_apply(
+        llama_sampler * smpl, ggml_context * /*ctx*/, ggml_cgraph * /*gf*/, llama_sampler_data * /*data*/) {
+    auto * ctx = (test_single_output_backend_sampler *) smpl->ctx;
+    ctx->backend_apply_count++;
+}
+
+static llama_sampler_i test_single_output_backend_sampler_i = {
+    /* .name              = */ test_single_output_backend_sampler_name,
+    /* .accept            = */ nullptr,
+    /* .apply             = */ test_single_output_backend_sampler_apply,
+    /* .reset             = */ nullptr,
+    /* .clone             = */ nullptr,
+    /* .free              = */ test_single_output_backend_sampler_free,
+    /* .backend_init      = */ test_single_output_backend_sampler_backend_init,
+    /* .backend_accept    = */ nullptr,
+    /* .backend_apply     = */ test_single_output_backend_sampler_backend_apply,
+    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+};
+
+static llama_sampler * test_single_output_backend_sampler_init(
+        test_single_output_backend_sampler ** sampler_ctx) {
+    auto * ctx = new test_single_output_backend_sampler;
+    *sampler_ctx = ctx;
+    return llama_sampler_init(&test_single_output_backend_sampler_i, ctx);
+}
+
 static void test_backend_greedy_sampling(const test_params & params) {
    const int seq_id = 0;

@@ -661,7 +732,7 @@ static void test_backend_multi_sequence_sampling(const test_params & params) {
 }

 static void test_backend_dist_sampling(const test_params & params) {
-    const int seq_id = 189;
+    const int seq_id = 0;
    const int32_t seed = 88;

    struct llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
@@ -1527,43 +1598,397 @@ static void test_backend_cpu_mixed_batch(const test_params & params) {
    printf("backend-cpu mixed batch test PASSED\n");
 }

-static void test_backend_max_outputs(const test_params & params) {
-    const int seq_id = 0;
-    const int32_t seed = 88;
+static void test_backend_multi_output_limit(const test_params & params) {
+    const llama_seq_id seq_id = 0;

-    llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
-    llama_sampler_ptr backend_sampler_chain(llama_sampler_chain_init(backend_chain_params));
-    llama_sampler_chain_add(backend_sampler_chain.get(), llama_sampler_init_dist(seed));
-    std::vector<llama_sampler_seq_config> backend_sampler_configs = {{ seq_id, backend_sampler_chain.get() }};
+    llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
+    llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(88));
+    std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
+    test_context test_ctx(params, configs, 1, 3, 0, 2);

-    test_context test_ctx(params, backend_sampler_configs);
-
-    llama_batch batch = llama_batch_init(512, 0, 1);
-    std::string prompt = "Hello";
-
-    std::vector<llama_token> tokens;
-    tokens.push_back(llama_vocab_bos(test_ctx.vocab));
-
-    std::vector<llama_token> prompt_tokens(32);
-    int n_tokens = llama_tokenize(test_ctx.vocab, prompt.c_str(), prompt.length(),
-                                   prompt_tokens.data(), prompt_tokens.size(),
-                                   false, false);
-    for (int i = 0; i < n_tokens; i++) {
-        tokens.push_back(prompt_tokens[i]);
+    llama_batch batch = llama_batch_init(3, 0, 1);
+    for (int i = 0; i < 3; ++i) {
+        common_batch_add(batch, llama_vocab_bos(test_ctx.vocab), i, { seq_id }, true);
    }

-    for (size_t i = 0; i < tokens.size(); i++) {
-        // set all tokens as output to trigger error
-        common_batch_add(batch, tokens[i], i, { seq_id }, true);
-    }
-
-    printf(">>> test_max_outputs expected error start:\n");
+    printf(">>> test_backend_multi_output_limit expected error start:\n");
    const int ret = llama_decode(test_ctx.ctx.get(), batch);
-    GGML_ASSERT(ret != 0 && "llama_decode should not succeed multiple outputs per sequence");
-    printf("<<< test_max_outputs expected error end.\n");
+    GGML_ASSERT(ret != 0 && "llama_decode should reject outputs above the per-sequence limit");
+    printf("<<< test_backend_multi_output_limit expected error end.\n");
+
    llama_batch_free(batch);

-    printf("backend max outputs test PASSED\n");
+    printf("backend multi-output limit test PASSED\n");
+}
+
+static void test_backend_multi_sequence_multi_output_dist(const test_params & params) {
+    const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
+    const int32_t n_vocab = llama_vocab_n_tokens(vocab);
+    const uint32_t seeds[] = { 88, 1337 };
+    // reduce the chance that swapped random inputs select the same token
+    const float temp = 10.0f;
+
+    llama_sampler_ptr chain_0(llama_sampler_chain_init(llama_sampler_chain_default_params()));
+    llama_sampler_ptr chain_1(llama_sampler_chain_init(llama_sampler_chain_default_params()));
+    llama_sampler_chain_add(chain_0.get(), llama_sampler_init_temp(temp));
+    llama_sampler_chain_add(chain_0.get(), llama_sampler_init_dist(seeds[0]));
+    llama_sampler_chain_add(chain_1.get(), llama_sampler_init_temp(temp));
+    llama_sampler_chain_add(chain_1.get(), llama_sampler_init_dist(seeds[1]));
+    std::vector<llama_sampler_seq_config> configs = {
+        { 0, chain_0.get() },
+        { 1, chain_1.get() },
+    };
+    test_context test_ctx(params, configs, 2, 4, 0, 2);
+
+    std::vector<llama_sampler_seq_config> reference_configs;
+    test_context reference_ctx(params, reference_configs, 2, 4);
+
+    const llama_token seq_tokens[2][2] = {
+        { llama_vocab_bos(vocab), llama_vocab_eos(vocab) },
+        { llama_vocab_eos(vocab), llama_vocab_bos(vocab) },
+    };
+
+    llama_batch batch = llama_batch_init(4, 0, 1);
+    for (int pos = 0; pos < 2; ++pos) {
+        common_batch_add(batch, seq_tokens[0][pos], pos, { 0 }, true);
+        common_batch_add(batch, seq_tokens[1][pos], pos, { 1 }, true);
+    }
+
+    GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
+    GGML_ASSERT(llama_decode(reference_ctx.ctx.get(), batch) == 0);
+
+    std::mt19937 reference_rngs[] = {
+        std::mt19937(seeds[0]),
+        std::mt19937(seeds[1]),
+    };
+    std::uniform_real_distribution<double> reference_dist(0.0, 1.0);
+
+    for (int i = 0; i < batch.n_tokens; ++i) {
+        const llama_seq_id seq_id = batch.seq_id[i][0];
+        GGML_ASSERT(seq_id == 0 || seq_id == 1);
+
+        llama_sampler * chain = seq_id == 0 ? chain_0.get() : chain_1.get();
+        const llama_token backend_token = llama_sampler_sample(chain, test_ctx.ctx.get(), i);
+        const float * sampled_logits = llama_get_sampled_logits_ith(test_ctx.ctx.get(), i);
+        const float * sampled_probs = llama_get_sampled_probs_ith(test_ctx.ctx.get(), i);
+        const uint32_t n_logits = llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i);
+        const uint32_t n_probs = llama_get_sampled_probs_count_ith(test_ctx.ctx.get(), i);
+        const float * reference_logits = llama_get_logits_ith(reference_ctx.ctx.get(), i);
+
+        GGML_ASSERT(backend_token >= 0 && backend_token < n_vocab);
+        GGML_ASSERT(sampled_logits != nullptr);
+        GGML_ASSERT(sampled_probs != nullptr);
+        GGML_ASSERT(reference_logits != nullptr);
+        GGML_ASSERT(n_logits == (uint32_t) n_vocab);
+        GGML_ASSERT(n_probs == (uint32_t) n_vocab);
+
+        float prob_sum = 0.0f;
+        float cumsum_before = 0.0f;
+        for (llama_token token = 0; token < n_vocab; ++token) {
+            const float expected_logit = reference_logits[token] / temp;
+            const float tolerance = 1e-4f * std::max(1.0f, std::fabs(expected_logit));
+            GGML_ASSERT(std::fabs(sampled_logits[token] - expected_logit) <= tolerance);
+            GGML_ASSERT(std::isfinite(sampled_probs[token]));
+            GGML_ASSERT(sampled_probs[token] >= 0.0f);
+
+            prob_sum += sampled_probs[token];
+            if (token < backend_token) {
+                cumsum_before += sampled_probs[token];
+            }
+        }
+
+        GGML_ASSERT(std::fabs(prob_sum - 1.0f) <= 1e-3f);
+
+        const float rnd = reference_dist(reference_rngs[seq_id]);
+        const float cumsum_sampled = cumsum_before + sampled_probs[backend_token];
+        GGML_ASSERT(rnd >= cumsum_before - 1e-4f);
+        GGML_ASSERT(rnd <= cumsum_sampled + 1e-4f);
+    }
+
+    llama_batch_free(batch);
+
+    printf("backend multi-sequence multi-output dist test PASSED\n");
+}
+
+static void test_backend_multi_output_dist_transaction(const test_params & params) {
+    const llama_seq_id seq_id = 0;
+    const uint32_t seed = 95;
+    const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
+
+    llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
+    llama_sampler_chain_add(chain.get(), llama_sampler_init_temp(10.0f));
+    llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(seed));
+    std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
+    test_context test_ctx(params, configs, 1, 3, 2, 3);
+
+    auto verify_random = [&](int32_t row, float rnd, bool accept = true) {
+        const llama_token token = accept ?
+            llama_sampler_sample(chain.get(), test_ctx.ctx.get(), row) :
+            llama_get_sampled_token_ith(test_ctx.ctx.get(), row);
+        const float * probs = llama_get_sampled_probs_ith(test_ctx.ctx.get(), row);
+
+        GGML_ASSERT(token >= 0 && token < llama_vocab_n_tokens(vocab));
+        GGML_ASSERT(probs != nullptr);
+
+        float cumsum_before = 0.0f;
+        for (llama_token i = 0; i < token; ++i) {
+            cumsum_before += probs[i];
+        }
+
+        const float cumsum_sampled = cumsum_before + probs[token];
+        GGML_ASSERT(rnd >= cumsum_before - 1e-4f);
+        GGML_ASSERT(rnd <= cumsum_sampled + 1e-4f);
+    };
+
+    std::mt19937 rng(seed);
+    std::uniform_real_distribution<double> dist(0.0, 1.0);
+    float randoms[3];
+    for (float & rnd : randoms) {
+        rnd = dist(rng);
+    }
+
+    int32_t pos = 0;
+    auto decode = [&]() {
+        llama_batch batch = llama_batch_init(3, 0, 1);
+        for (int32_t i = 0; i < 3; ++i) {
+            common_batch_add(batch, llama_vocab_bos(vocab), pos++, { seq_id }, true);
+        }
+        GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
+        return batch;
+    };
+
+    llama_batch batch = decode();
+    verify_random(0, randoms[0], false);
+    llama_batch_free(batch);
+
+    batch = decode();
+    verify_random(0, randoms[0]);
+    verify_random(1, randoms[1]);
+    llama_batch_free(batch);
+
+    batch = decode();
+    llama_sampler_ptr saved(llama_sampler_clone(chain.get()));
+    verify_random(0, randoms[2]);
+    llama_batch_free(batch);
+
+    llama_sampler_copy(saved.get(), chain.get());
+    batch = decode();
+    verify_random(0, randoms[2]);
+    llama_batch_free(batch);
+
+    printf("backend multi-output dist transaction test PASSED\n");
+}
+
+static void test_backend_multi_output_sampling_chain(const test_params & params) {
+    const llama_seq_id seq_id = 0;
+    const uint32_t seed = 88;
+    const float p = 0.9f;
+    const float temp = 0.8f;
+    const float cdf_epsilon = 1e-4f;
+    const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
+    const int32_t n_vocab = llama_vocab_n_tokens(vocab);
+    const uint32_t k = std::min<uint32_t>(512, n_vocab);
+    const llama_logit_bias bias = { llama_vocab_bos(vocab), -0.1f };
+
+    auto make_filter_chain = [&]() {
+        llama_sampler_ptr result(llama_sampler_chain_init(llama_sampler_chain_default_params()));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_logit_bias(n_vocab, 1, &bias));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_top_k(k));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_top_p(p, 1));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_min_p(0.01f, 1));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_temp(temp));
+        return result;
+    };
+
+    llama_sampler_ptr chain = make_filter_chain();
+    llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(seed));
+    std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
+    test_context test_ctx(params, configs, 1, 2, 2, 2);
+
+    std::vector<llama_sampler_seq_config> reference_configs;
+    test_context reference_ctx(params, reference_configs, 1, 2, 2);
+
+    llama_sampler_ptr reference_bias(llama_sampler_init_logit_bias(n_vocab, 1, &bias));
+    llama_sampler_ptr reference_top_k(llama_sampler_init_top_k(k));
+    llama_sampler_ptr reference_top_p(llama_sampler_init_top_p(p, 1));
+    llama_sampler_ptr reference_min_p(llama_sampler_init_min_p(0.01f, 1));
+    llama_sampler_ptr reference_temp(llama_sampler_init_temp(temp));
+    std::vector<llama_token_data> reference_data(n_vocab);
+
+    auto make_batch = [&](int32_t pos) {
+        llama_batch batch = llama_batch_init(2, 0, 1);
+        for (int i = 0; i < 2; ++i) {
+            common_batch_add(batch, llama_vocab_bos(vocab), pos + i, { seq_id }, true);
+        }
+        return batch;
+    };
+
+    llama_batch batch = make_batch(0);
+    GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
+    GGML_ASSERT(llama_decode(reference_ctx.ctx.get(), batch) == 0);
+
+    for (int i = 0; i < batch.n_tokens; ++i) {
+        const llama_token backend_token = llama_sampler_sample(chain.get(), test_ctx.ctx.get(), i);
+        const float * sampled_logits = llama_get_sampled_logits_ith(test_ctx.ctx.get(), i);
+        const float * sampled_probs = llama_get_sampled_probs_ith(test_ctx.ctx.get(), i);
+        const llama_token * sampled_candidates = llama_get_sampled_candidates_ith(test_ctx.ctx.get(), i);
+        const uint32_t n_logits = llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i);
+        const uint32_t n_probs = llama_get_sampled_probs_count_ith(test_ctx.ctx.get(), i);
+        const uint32_t n_candidates = llama_get_sampled_candidates_count_ith(test_ctx.ctx.get(), i);
+        const float * reference_logits = llama_get_logits_ith(reference_ctx.ctx.get(), i);
+
+        GGML_ASSERT(backend_token >= 0 && backend_token < n_vocab);
+        GGML_ASSERT(sampled_logits != nullptr);
+        GGML_ASSERT(sampled_probs != nullptr);
+        GGML_ASSERT(sampled_candidates != nullptr);
+        GGML_ASSERT(reference_logits != nullptr);
+        GGML_ASSERT(n_logits == k);
+        GGML_ASSERT(n_probs == n_logits);
+        GGML_ASSERT(n_candidates == n_logits);
+
+        for (llama_token token = 0; token < n_vocab; ++token) {
+            reference_data[token] = { token, reference_logits[token], 0.0f };
+        }
+
+        llama_token_data_array reference = {
+            /* .data     = */ reference_data.data(),
+            /* .size     = */ reference_data.size(),
+            /* .selected = */ LLAMA_TOKEN_NULL,
+            /* .sorted   = */ false,
+        };
+
+        llama_sampler_apply(reference_bias.get(), &reference);
+        llama_sampler_apply(reference_top_k.get(), &reference);
+        llama_sampler_apply(reference_top_p.get(), &reference);
+        GGML_ASSERT(reference.size > 0);
+
+        float cdf = 0.0f;
+        for (size_t j = 0; j < reference.size; ++j) {
+            cdf += reference.data[j].p;
+        }
+        const float cdf_before = cdf - reference.data[reference.size - 1].p;
+        const float boundary_distance = std::min(std::fabs(cdf_before - p), std::fabs(cdf - p));
+
+        llama_sampler_apply(reference_min_p.get(), &reference);
+        llama_sampler_apply(reference_temp.get(), &reference);
+
+        std::unordered_map<llama_token, float> reference_by_id;
+        for (size_t j = 0; j < reference.size; ++j) {
+            reference_by_id.emplace(reference.data[j].id, reference.data[j].logit);
+        }
+        size_t n_backend_only = 0;
+        int32_t sampled_index = -1;
+        float prob_sum = 0.0f;
+
+        for (uint32_t j = 0; j < n_logits; ++j) {
+            GGML_ASSERT(sampled_candidates[j] >= 0 && sampled_candidates[j] < n_vocab);
+            GGML_ASSERT(std::isfinite(sampled_probs[j]));
+            GGML_ASSERT(sampled_probs[j] >= 0.0f);
+            prob_sum += sampled_probs[j];
+
+            if (sampled_candidates[j] == backend_token) {
+                sampled_index = j;
+            }
+            if (!std::isfinite(sampled_logits[j])) {
+                GGML_ASSERT(std::isinf(sampled_logits[j]) && sampled_logits[j] < 0.0f);
+                GGML_ASSERT(sampled_probs[j] == 0.0f);
+                continue;
+            }
+
+            const auto match = reference_by_id.find(sampled_candidates[j]);
+            if (match == reference_by_id.end()) {
+                ++n_backend_only;
+                continue;
+            }
+
+            const float tolerance = 1e-4f * std::max(1.0f, std::fabs(match->second));
+            GGML_ASSERT(std::fabs(sampled_logits[j] - match->second) <= tolerance);
+            reference_by_id.erase(match);
+        }
+
+        const size_t n_reference_only = reference_by_id.size();
+
+        if (n_backend_only != 0 || n_reference_only != 0) {
+            GGML_ASSERT(n_backend_only <= 1);
+            GGML_ASSERT(n_reference_only <= 1);
+            GGML_ASSERT(boundary_distance <= cdf_epsilon);
+        }
+
+        GGML_ASSERT(sampled_index >= 0);
+        GGML_ASSERT(std::isfinite(sampled_logits[sampled_index]));
+        GGML_ASSERT(sampled_probs[sampled_index] > 0.0f);
+        GGML_ASSERT(std::fabs(prob_sum - 1.0f) <= 1e-3f);
+    }
+
+    llama_batch_free(batch);
+
+    batch = make_batch(2);
+    GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
+    llama_batch_free(batch);
+
+    printf("backend multi-output sampling chain test PASSED\n");
+}
+
+static void test_backend_multi_output_cpu_suffix(const test_params & params) {
+    const llama_seq_id seq_id = 0;
+    const int32_t k = 8;
+    const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
+
+    auto make_chain = [&](test_single_output_backend_sampler ** sampler_ctx) {
+        llama_sampler_ptr result(llama_sampler_chain_init(llama_sampler_chain_default_params()));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_top_k(k));
+        llama_sampler_chain_add(result.get(), test_single_output_backend_sampler_init(sampler_ctx));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_dist(88));
+        return result;
+    };
+
+    {
+        test_single_output_backend_sampler * sampler_ctx = nullptr;
+        llama_sampler_ptr chain = make_chain(&sampler_ctx);
+        std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
+        test_context test_ctx(params, configs, 1, 1, 0, 4);
+
+        llama_batch batch = llama_batch_init(1, 0, 1);
+        common_batch_add(batch, llama_vocab_bos(vocab), 0, { seq_id }, true);
+        GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
+
+        GGML_ASSERT(sampler_ctx->backend_initialized);
+        GGML_ASSERT(sampler_ctx->backend_outputs_max_per_seq == 1);
+        GGML_ASSERT(sampler_ctx->backend_apply_count > 0);
+        GGML_ASSERT(sampler_ctx->apply_count == 0);
+        GGML_ASSERT(llama_get_sampled_token_ith(test_ctx.ctx.get(), 0) != LLAMA_TOKEN_NULL);
+
+        llama_batch_free(batch);
+    }
+
+    {
+        test_single_output_backend_sampler * sampler_ctx = nullptr;
+        llama_sampler_ptr chain = make_chain(&sampler_ctx);
+        std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
+        test_context test_ctx(params, configs, 1, 2, 0, 0);
+
+        llama_batch batch = llama_batch_init(2, 0, 1);
+        for (int i = 0; i < 2; ++i) {
+            common_batch_add(batch, llama_vocab_bos(vocab), i, { seq_id }, true);
+        }
+        GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
+
+        GGML_ASSERT(!sampler_ctx->backend_initialized);
+        GGML_ASSERT(sampler_ctx->backend_outputs_max_per_seq == 2);
+        GGML_ASSERT(sampler_ctx->backend_apply_count == 0);
+        for (int i = 0; i < batch.n_tokens; ++i) {
+            GGML_ASSERT(llama_get_sampled_token_ith(test_ctx.ctx.get(), i) == LLAMA_TOKEN_NULL);
+            GGML_ASSERT(llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i) == (uint32_t) k);
+            GGML_ASSERT(llama_get_sampled_candidates_count_ith(test_ctx.ctx.get(), i) == (uint32_t) k);
+            const llama_token token = llama_sampler_sample(chain.get(), test_ctx.ctx.get(), i);
+            GGML_ASSERT(token >= 0 && token < llama_vocab_n_tokens(vocab));
+        }
+        GGML_ASSERT(sampler_ctx->apply_count == batch.n_tokens);
+
+        llama_batch_free(batch);
+    }
+
+    printf("backend multi-output CPU suffix test PASSED\n");
 }

 struct backend_test_case {
@@ -1583,7 +2008,11 @@ static const backend_test_case BACKEND_TESTS[] = {
    { "dist",            test_backend_dist_sampling,           true  },
    { "dist_and_cpu",    test_backend_dist_sampling_and_cpu,   true  },
    { "set_sampler",     test_backend_set_sampler,             true  },
-    { "max_outputs",     test_backend_max_outputs,             true  },
+    { "multi_output_limit",    test_backend_multi_output_limit,      true },
+    { "multi_sequence_multi_output_dist", test_backend_multi_sequence_multi_output_dist, true },
+    { "multi_output_dist_transaction", test_backend_multi_output_dist_transaction, true },
+    { "multi_output_sampling_chain", test_backend_multi_output_sampling_chain, true },
+    { "multi_output_cpu",      test_backend_multi_output_cpu_suffix, true },
    { "mixed",           test_backend_mixed_sampling,          true  },
    { "min_p",           test_backend_min_p_sampling,          true  },
    { "cpu_mixed",       test_backend_cpu_mixed_batch,         true  },
@@ -1668,18 +2097,9 @@ static std::vector<const backend_test_case *> collect_tests_to_run(const std::st
        }
    } else {
        for (const auto & test : BACKEND_TESTS) {
-            if (!test.enabled_by_default) {
-                continue;
+            if (test.enabled_by_default) {
+                selected.push_back(&test);
            }
-#ifdef GGML_USE_HIP
-            // TODO: remove this when https://github.com/ggml-org/llama.cpp/pull/26592 is merged
-            if (test.name == "penalties" || test.name == "set_sampler" ||
-                test.name == "mixed"     || test.name == "top_p") {
-                fprintf(stderr, "Skipping test '%s' on HIP backend (no backend TOP_K support)\n", test.name.c_str());
-                continue;
-            }
-#endif // GGML_USE_HIP
-            selected.push_back(&test);
        }
    }

@@ -437,14 +437,6 @@ static bool arch_supported(const llm_arch arch) {
    }
 #endif // GGML_USE_WEBGPU

-    // FIXME: jamba produces incorrect output (~0.55 NMSE vs CPU) on the HIP
-    // backend on RDNA3.5 (gfx1151); the SSM kernels need investigation.
-#ifdef GGML_USE_HIP
-    if (arch == LLM_ARCH_JAMBA) {
-        return false;
-    }
-#endif // GGML_USE_HIP
-
    return true;
 }

@@ -195,7 +195,7 @@ static const std::vector<std::string> dspark_dflash = {

 struct plan_case {
    const char * name;
-    const std::vector<std::string> files;
+    const std::vector<std::string> & files;
    const char * hf_repo;
    const char * hf_file;
    bool sidecars;        // request mmproj + mtp + dflash + eagle3 + dspark
@@ -1,6 +1,4 @@
 #include <stdio.h>
-#include <stdlib.h>
-#include <string.h>
 #include <assert.h>

 #include "mtmd.h"
@@ -64,72 +62,6 @@ int main(void) {
        }
    }

-    // test chunk save/load round-trip
-    for (size_t i = 0; i < n_chunks; i++) {
-        const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
-        assert(chunk != NULL);
-        enum mtmd_input_chunk_type type = mtmd_input_chunk_get_type(chunk);
-
-        // query the required buffer size (out_buf == NULL)
-        size_t expected_len = 0;
-        int32_t rc = mtmd_input_chunk_save(chunk, NULL, 0, &expected_len);
-        printf("    Chunk %zu: save query rc = %d, expected_len = %zu\n", i, rc, expected_len);
-        assert(rc == 0);
-        assert(expected_len > 0);
-
-        // saving into a too-small buffer must fail, not crash
-        char tiny_buf[1];
-        rc = mtmd_input_chunk_save(chunk, tiny_buf, sizeof(tiny_buf), NULL);
-        printf("    Chunk %zu: save into too-small buffer rc = %d (expect non-zero)\n", i, rc);
-        assert(rc != 0);
-
-        // save into a properly-sized buffer
-        char * buf = (char *) malloc(expected_len);
-        assert(buf != NULL);
-        rc = mtmd_input_chunk_save(chunk, buf, expected_len, NULL);
-        assert(rc == 0);
-
-        // loading from a truncated buffer must fail gracefully, not crash
-        if (expected_len > 1) {
-            mtmd_input_chunk * bad = mtmd_input_chunk_load(buf, expected_len - 1);
-            printf("    Chunk %zu: load from truncated buffer = %p (expect NULL)\n", i, (void *) bad);
-            assert(bad == NULL);
-        }
-
-        // load it back
-        mtmd_input_chunk * loaded = mtmd_input_chunk_load(buf, expected_len);
-        assert(loaded != NULL);
-
-        // metadata must match the original chunk
-        assert(mtmd_input_chunk_get_type(loaded) == type);
-        assert(mtmd_input_chunk_get_n_tokens(loaded) == mtmd_input_chunk_get_n_tokens(chunk));
-        assert(mtmd_input_chunk_get_n_pos(loaded) == mtmd_input_chunk_get_n_pos(chunk));
-
-        if (type == MTMD_INPUT_CHUNK_TYPE_TEXT) {
-            size_t n_tok_orig, n_tok_loaded;
-            const llama_token * tok_orig   = mtmd_input_chunk_get_tokens_text(chunk, &n_tok_orig);
-            const llama_token * tok_loaded = mtmd_input_chunk_get_tokens_text(loaded, &n_tok_loaded);
-            printf("    Chunk %zu: loaded %zu text tokens (orig %zu), first token %d (orig %d)\n",
-                i, n_tok_loaded, n_tok_orig,
-                n_tok_loaded > 0 ? tok_loaded[0] : -1,
-                n_tok_orig   > 0 ? tok_orig[0]   : -1);
-            assert(n_tok_orig == n_tok_loaded);
-            for (size_t j = 0; j < n_tok_orig; j++) {
-                assert(tok_orig[j] == tok_loaded[j]);
-            }
-        } else if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE || type == MTMD_INPUT_CHUNK_TYPE_AUDIO) {
-            const char * id_orig   = mtmd_input_chunk_get_id(chunk);
-            const char * id_loaded = mtmd_input_chunk_get_id(loaded);
-            printf("    Chunk %zu: loaded id '%s' (orig '%s')\n", i, id_loaded, id_orig);
-            assert(id_orig != NULL && id_loaded != NULL);
-            assert(strcmp(id_orig, id_loaded) == 0);
-        }
-
-        mtmd_input_chunk_free(loaded);
-        free(buf);
-    }
-    printf("Chunk save/load round-trip OK\n");
-
    // Free the chunks
    mtmd_input_chunks_free(chunks);

@@ -61,6 +61,35 @@ private:
    std::vector<llama_token_data> cur;
 };

+static llama_token sample_dist(llama_sampler * sampler, const std::vector<float> & logits) {
+    std::vector<llama_token_data> cur;
+    for (llama_token token_id = 0; token_id < (llama_token) logits.size(); ++token_id) {
+        cur.push_back({ token_id, logits[token_id], 0.0f });
+    }
+
+    llama_token_data_array cur_p = { cur.data(), cur.size(), -1, false };
+    llama_sampler_apply(sampler, &cur_p);
+    GGML_ASSERT(cur_p.selected >= 0);
+    GGML_ASSERT((size_t) cur_p.selected < cur_p.size);
+    return cur_p.data[cur_p.selected].id;
+}
+
+static void test_dist_singleton_rng() {
+    llama_sampler * singleton = llama_sampler_init_dist(4242);
+    llama_sampler * control   = llama_sampler_init_dist(4242);
+
+    sample_dist(singleton, { 0.0f });
+    sample_dist(control,   { 0.0f, 0.0f });
+
+    const std::vector<float> logits(256, 0.0f);
+    for (int i = 0; i < 4; ++i) {
+        GGML_ASSERT(sample_dist(singleton, logits) == sample_dist(control, logits));
+    }
+
+    llama_sampler_free(singleton);
+    llama_sampler_free(control);
+}
+
 static void test_temp(const std::vector<float> & probs, const std::vector<float> & probs_expected, float temp) {
    sampler_tester tester(probs, probs_expected);

@@ -308,6 +337,8 @@ static void test_perf() {
 int main(void) {
    ggml_time_init();

+    test_dist_singleton_rng();
+
    test_temp({0.1f, 0.2f, 0.3f, 0.4f}, {0.1f, 0.2f, 0.3f, 0.4f}, 1.0f);
    test_temp({0.1f, 0.2f, 0.3f, 0.4f}, {0.0f, 0.0f, 0.0f, 1.0f}, 0.0f);

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

-struct mtmd_serialization; // forward declaration
-
 // For images, buf.size() == nx*ny*3
 //     Memory layout: RGBRGBRGB...
 // For seq, buf.size() == nx*ny*3*nt
@@ -673,9 +671,6 @@ struct clip_image_f32 {
        return buf.empty();
    }

-    void serialize(struct mtmd_serialization & ser) const;
-    void deserialize(struct mtmd_serialization & ser);
-
  private:
    std::vector<float> buf;
    int nx_ = 0;
@@ -757,9 +752,6 @@ struct clip_image_f32_batch {
        }
        return new_batch;
    }
-
-    void serialize(struct mtmd_serialization & ser) const;
-    void deserialize(struct mtmd_serialization & ser);
 };

 //
@@ -170,17 +170,6 @@ struct clip_hparams {
        warmup_image_size = static_cast<int>(std::sqrt(image_max_pixels));
    }

-    // used by longest_edge preprocessor (no model-specific value for min/max tokens)
-    void set_limit_image_tokens() {
-        const int patch_area = patch_size * patch_size * n_merge * n_merge;
-        if (custom_image_min_tokens > 0) {
-            image_min_pixels = custom_image_min_tokens * patch_area;
-        }
-        if (custom_image_max_tokens > 0) {
-            image_max_pixels = custom_image_max_tokens * patch_area;
-        }
-    }
-
    void set_warmup_n_tokens(int n_tokens) {
        int n_tok_per_side = static_cast<int>(std::sqrt(n_tokens));
        GGML_ASSERT(n_tok_per_side * n_tok_per_side == n_tokens && "n_tokens must be n*n");
@@ -708,10 +708,9 @@ ggml_tensor * clip_graph::build_attn(
        ggml_tensor * sinks) const {
    // these nodes are added to the graph together so that they are not reordered
    // by doing so, the number of splits in the graph is reduced
-    // the order is fixed without the compute flag, so an unselected branch stays out of the compute set
-    ggml_build_forward_order(gf, q_cur);
-    ggml_build_forward_order(gf, k_cur);
-    ggml_build_forward_order(gf, v_cur);
+    ggml_build_forward_expand(gf, q_cur);
+    ggml_build_forward_expand(gf, k_cur);
+    ggml_build_forward_expand(gf, v_cur);

    ggml_tensor * q = ggml_permute(ctx0, q_cur, 0, 2, 1, 3);
    //cb(q, "q", il);
@@ -1434,7 +1433,6 @@ struct clip_model_loader {
                        // use default llava-uhd preprocessing params
                        get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false);
                        get_u32(KEY_PREPROC_IMAGE_SIZE, hparams.image_longest_edge, false);
-                        hparams.set_limit_image_tokens();
                    } break;
                case PROJECTOR_TYPE_LFM2:
                    {
@@ -1472,7 +1470,6 @@ struct clip_model_loader {
                        get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false);
                        hparams.image_longest_edge = hparams.image_size;
                        get_u32(KEY_PREPROC_IMAGE_SIZE, hparams.image_longest_edge, false);
-                        hparams.set_limit_image_tokens();
                        hparams.set_warmup_n_tokens(256); // avoid OOM on warmup
                    } break;
                case PROJECTOR_TYPE_DOTS_OCR:
@@ -1597,7 +1594,6 @@ struct clip_model_loader {
                        if (hparams.image_longest_edge == 0) {
                            hparams.image_longest_edge = 3024;
                        }
-                        // note: the step3vl preprocessor slices based on a fixed window grid, so it does not support custom min/max image tokens
                        hparams.warmup_image_size = hparams.image_size;
                    } break;
                case PROJECTOR_TYPE_YOUTUVL:
@@ -112,6 +112,7 @@ public:
        c2w_state.clear();
        audio_pcm.clear();
        overlay.clear();
+        overlay_idx = 0;
        h_state_buf.clear();
        out_buf.clear();
        prompt_embd_buf.clear();
@@ -204,9 +205,11 @@ public:
        top_p = inp->top_p > 0 ? inp->top_p : 1.0f;
        out_type = inp->out_type;

-        // the prompt above holds the whole text stream up to tts_eos, so every generated
-        // frame adds tts_pad on top of the codes embedding
-        overlay = row(tts_pad);
+        // the text stream keeps flowing during generation: after frame k, the input adds
+        // trailing text row k on top of the codes embedding, then tts_eos, then tts_pad
+        for (int i = 3; i < n_ids - 5; i++) overlay.push_back(row(ids[(size_t) i]));
+        overlay.push_back(row(tts_eos));
+        overlay.push_back(row(tts_pad));

        return 0;
    }
@@ -262,7 +265,9 @@ public:
        }

        std::vector<float> fb(out.embd, out.embd + n_embd);
-        for (int i = 0; i < n_embd; i++) fb[(size_t) i] += overlay[(size_t) i];
+        const auto & ov = overlay[std::min(overlay_idx, overlay.size() - 1)];
+        for (int i = 0; i < n_embd; i++) fb[(size_t) i] += ov[(size_t) i];
+        overlay_idx++;

        const int n_pos_per_embd = mrope ? 4 : 1;
        decode_embd_batch batch_embd(fb.data(), 1, n_pos_per_embd, n_embd);
@@ -432,7 +437,8 @@ private:
    std::vector<int32_t> codes_buf;
    std::vector<uint8_t> c2w_state;
    std::vector<float>   audio_pcm;
-    std::vector<float> overlay;
+    std::vector<std::vector<float>> overlay;
+    size_t overlay_idx = 0;
    std::vector<float> h_state_buf;
    mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
    std::vector<char> out_buf;
@@ -139,46 +139,50 @@ struct img_tool {
        }
    }

-    struct calc_size_opt {
-        int align_size = 1;
-        int min_pixels = 0;   // 0 = disabled
-        int max_pixels = 0;   // 0 = disabled
-        // applied before min/max_pixels, so min_pixels can push an edge back above longest_edge
-        int longest_edge = 0; // 0 = disabled
-    };
-
-    // calculate the size of the **resized** image, while preserving the aspect ratio and
-    // aligning to the nearest multiple of align_size ("smart_resize" in transformers code)
-    static clip_image_size calc_size_preserved_ratio(const clip_image_size & inp_size, const calc_size_opt & opts) {
-        GGML_ASSERT(opts.align_size > 0);
-        const int width  = inp_size.width;
-        const int height = inp_size.height;
-        if (width <= 0 || height <= 0) {
+    // calculate the size of the **resized** image, while preserving the aspect ratio
+    // the calculated size will be aligned to the nearest multiple of align_size
+    // if H or W size is larger than longest_edge, it will be resized to longest_edge
+    static clip_image_size calc_size_preserved_ratio(const clip_image_size & inp_size, const int align_size, const int longest_edge) {
+        GGML_ASSERT(align_size > 0);
+        if (inp_size.width <= 0 || inp_size.height <= 0 || longest_edge <= 0) {
            return {0, 0};
        }

-        auto round_by_factor = [f = opts.align_size](float x) { return static_cast<int>(std::round(x / static_cast<float>(f))) * f; };
-        auto ceil_by_factor  = [f = opts.align_size](float x) { return static_cast<int>(std::ceil(x / static_cast<float>(f))) * f; };
-        auto floor_by_factor = [f = opts.align_size](float x) { return static_cast<int>(std::floor(x / static_cast<float>(f))) * f; };
+        float scale = std::min(static_cast<float>(longest_edge) / inp_size.width,
+                               static_cast<float>(longest_edge) / inp_size.height);

-        int w_bar, h_bar;
-        if (opts.longest_edge > 0) {
-            const float scale = std::min(static_cast<float>(opts.longest_edge) / width,
-                                          static_cast<float>(opts.longest_edge) / height);
-            w_bar = ceil_by_factor(width  * scale);
-            h_bar = ceil_by_factor(height * scale);
-        } else {
-            // always align up first
-            w_bar = std::max(opts.align_size, round_by_factor(width));
-            h_bar = std::max(opts.align_size, round_by_factor(height));
-        }
+        float target_width_f  = static_cast<float>(inp_size.width)  * scale;
+        float target_height_f = static_cast<float>(inp_size.height) * scale;

-        if (opts.max_pixels > 0 && h_bar * w_bar > opts.max_pixels) {
-            const auto beta = std::sqrt(static_cast<float>(height) * width / opts.max_pixels);
-            h_bar = std::max(opts.align_size, floor_by_factor(height / beta));
-            w_bar = std::max(opts.align_size, floor_by_factor(width  / beta));
-        } else if (opts.min_pixels > 0 && h_bar * w_bar < opts.min_pixels) {
-            const auto beta = std::sqrt(static_cast<float>(opts.min_pixels) / (static_cast<float>(height) * width));
+        auto ceil_by_factor = [f = align_size](float x) { return static_cast<int>(std::ceil(x / static_cast<float>(f))) * f; };
+        int aligned_width  = ceil_by_factor(target_width_f);
+        int aligned_height = ceil_by_factor(target_height_f);
+
+        return {aligned_width, aligned_height};
+    }
+
+    // calculate the size of the **resized** image, while preserving the aspect ratio
+    // the calculated size will have min_pixels <= W*H <= max_pixels
+    // this is referred as "smart_resize" in transformers code
+    static clip_image_size calc_size_preserved_ratio(const clip_image_size & inp_size, const int align_size, const int min_pixels, const int max_pixels) {
+        GGML_ASSERT(align_size > 0);
+        const int width  = inp_size.width;
+        const int height = inp_size.height;
+
+        auto round_by_factor = [f = align_size](float x) { return static_cast<int>(std::round(x / static_cast<float>(f))) * f; };
+        auto ceil_by_factor  = [f = align_size](float x) { return static_cast<int>(std::ceil(x / static_cast<float>(f))) * f; };
+        auto floor_by_factor = [f = align_size](float x) { return static_cast<int>(std::floor(x / static_cast<float>(f))) * f; };
+
+        // always align up first
+        int h_bar = std::max(align_size, round_by_factor(height));
+        int w_bar = std::max(align_size, round_by_factor(width));
+
+        if (h_bar * w_bar > max_pixels) {
+            const auto beta = std::sqrt(static_cast<float>(height * width) / max_pixels);
+            h_bar = std::max(align_size, floor_by_factor(height / beta));
+            w_bar = std::max(align_size, floor_by_factor(width  / beta));
+        } else if (h_bar * w_bar < min_pixels) {
+            const auto beta = std::sqrt(static_cast<float>(min_pixels) / (height * width));
            h_bar = ceil_by_factor(height * beta);
            w_bar = ceil_by_factor(width * beta);
        }
@@ -933,12 +937,9 @@ mtmd_image_preproc_out mtmd_image_preprocessor_dyn_size::preprocess(const clip_i
    const int cur_merge = hparams.n_merge;
    const clip_image_size target_size = img_tool::calc_size_preserved_ratio(
        original_size,
-        {
-            /* align_size   */ hparams.patch_size * cur_merge,
-            /* min_pixels   */ hparams.image_min_pixels,
-            /* max_pixels   */ hparams.image_max_pixels,
-            /* longest_edge */ 0,
-        });
+        hparams.patch_size * cur_merge,
+        hparams.image_min_pixels,
+        hparams.image_max_pixels);
    img_tool::resize(img, resized_image, target_size,
                        hparams.image_resize_algo,
                        hparams.image_resize_pad,
@@ -960,12 +961,8 @@ mtmd_image_preproc_out mtmd_image_preprocessor_longest_edge::preprocess(const cl
    const int cur_merge = hparams.n_merge == 0 ? 1 : hparams.n_merge;
    const clip_image_size target_size = img_tool::calc_size_preserved_ratio(
        original_size,
-        {
-            /* align_size   */ hparams.patch_size * cur_merge,
-            /* min_pixels   */ std::max(0, hparams.image_min_pixels),
-            /* max_pixels   */ std::max(0, hparams.image_max_pixels),
-            /* longest_edge */ hparams.image_longest_edge,
-        });
+        hparams.patch_size * cur_merge,
+        hparams.image_longest_edge);
    img_tool::resize(img, resized_image, target_size,
                        hparams.image_resize_algo,
                        hparams.image_resize_pad,
@@ -1003,8 +1000,8 @@ mtmd_image_preprocessor_llava_uhd::slice_instructions mtmd_image_preprocessor_lf
    mtmd_image_preprocessor_llava_uhd::slice_instructions inst;
    const int align_size = hparams.patch_size * hparams.n_merge;
    inst.overview_size = img_tool::calc_size_preserved_ratio(
-                            original_size,
-                            { align_size, hparams.image_min_pixels, hparams.image_max_pixels, 0 });
+                            original_size, align_size,
+                            hparams.image_min_pixels, hparams.image_max_pixels);
    // tile if either dimension exceeds tile_size with tolerance
    const bool needs_tiling = original_size.width > tile_size * max_pixels_tolerance || original_size.height > tile_size * max_pixels_tolerance;

@@ -1112,8 +1109,7 @@ mtmd_image_preproc_out mtmd_image_preprocessor_idefics3::preprocess(const clip_i
    // CITE: https://github.com/huggingface/transformers/blob/main/src/transformers/models/idefics3/image_processing_idefics3.py#L737
    const clip_image_size original_size = img.get_size();
    const clip_image_size refined_size = img_tool::calc_size_preserved_ratio(
-        original_size,
-        { hparams.image_size, std::max(0, hparams.image_min_pixels), std::max(0, hparams.image_max_pixels), hparams.image_longest_edge });
+        original_size, hparams.image_size, hparams.image_longest_edge);
    // LOG_INF("%s: original size: %d x %d, refined size: %d x %d\n",
    //         __func__, original_size.width, original_size.height,
    //         refined_size.width, refined_size.height);
@@ -22,123 +22,8 @@
 #include <cstdlib>
 #include <cstring>
 #include <climits>
-#include <type_traits>
 #include <vector>

-// remember to bump this if the serialization format changes
-#define MTMD_SERIALIZATION_VERSION 1
-
-struct mtmd_serialization {
-    // note: using 64-bit here for future-proofing
-    uint64_t version = MTMD_SERIALIZATION_VERSION;
-    std::vector<char> data;
-    size_t read_pos = 0; // cursor used when reading
-
-    // for writing
-    mtmd_serialization(uint64_t version) : version(version) {
-        write(version);
-    }
-
-    // for reading
-    mtmd_serialization(uint64_t version, const char * buf, size_t len) {
-        // copy buf to data
-        data.assign(buf, buf + len);
-        uint64_t ver_in = read<uint64_t>();
-        if (ver_in != version) {
-            throw std::runtime_error("version mismatch");
-        }
-        this->version = ver_in;
-    }
-
-    template <typename T>
-    void write(T value) {
-        static_assert(std::is_trivially_copyable<T>::value && !std::is_same<T, bool>::value,
-            "T must be trivially copyable and not bool");
-        const char * p = reinterpret_cast<const char *>(&value);
-        data.insert(data.end(), p, p + sizeof(T));
-    }
-
-    template <typename T>
-    T read() {
-        static_assert(std::is_trivially_copyable<T>::value && !std::is_same<T, bool>::value,
-            "T must be trivially copyable and not bool");
-        if (read_pos + sizeof(T) > data.size()) {
-            throw std::runtime_error("read OOB");
-        }
-        T value;
-        std::memcpy(&value, data.data() + read_pos, sizeof(T));
-        read_pos += sizeof(T);
-        return value;
-    }
-
-};
-
-template <>
-void mtmd_serialization::write<bool>(bool value) {
-    write<uint8_t>(value ? 1 : 0);
-}
-template <>
-bool mtmd_serialization::read<bool>() {
-    return read<uint8_t>() != 0;
-}
-
-template <>
-void mtmd_serialization::write<std::string>(std::string value) {
-    write<uint64_t>(value.size());
-    data.insert(data.end(), value.begin(), value.end());
-}
-template <>
-std::string mtmd_serialization::read<std::string>() {
-    uint64_t len = read<uint64_t>();
-    if (read_pos + len > data.size()) {
-        throw std::runtime_error("read_string OOB");
-    }
-    std::string str(data.data() + read_pos, len);
-    read_pos += len;
-    return str;
-}
-
-// only mtmd.cpp needs these, so they're implemented here rather than in clip-impl.h
-void clip_image_f32::serialize(mtmd_serialization & ser) const {
-    // remember to bump MTMD_SERIALIZATION_VERSION if this is changed
-    // note: buf is intentionally NOT serialized; the loaded clip_image_f32 will always be a placeholder
-    ser.write(add_viewsep);
-    ser.write(add_newline);
-    ser.write((int32_t)nx_);
-    ser.write((int32_t)ny_);
-}
-void clip_image_f32::deserialize(mtmd_serialization & ser) {
-    add_viewsep = ser.read<bool>();
-    add_newline = ser.read<bool>();
-    nx_ = ser.read<int32_t>();
-    ny_ = ser.read<int32_t>();
-    buf.clear(); // always a placeholder after loading
-}
-
-void clip_image_f32_batch::serialize(mtmd_serialization & ser) const {
-    // remember to bump MTMD_SERIALIZATION_VERSION if this is changed
-    ser.write(is_audio);
-    ser.write<uint64_t>(entries.size());
-    for (const auto & entry : entries) {
-        entry.serialize(ser);
-    }
-}
-void clip_image_f32_batch::deserialize(mtmd_serialization & ser) {
-    is_audio = ser.read<bool>();
-    uint64_t n = ser.read<uint64_t>();
-    constexpr size_t min_entry_bytes = sizeof(uint8_t) * 2 + sizeof(int32_t) * 2;
-    if (n > (ser.data.size() - ser.read_pos) / min_entry_bytes) {
-        throw std::runtime_error("entries count exceeds buffer size");
-    }
-    entries.clear();
-    entries.reserve(n);
-    for (uint64_t i = 0; i < n; i++) {
-        clip_image_f32 entry;
-        entry.deserialize(ser);
-        entries.push_back(std::move(entry));
-    }
-}
-
 // for still image data, layout is RGBRGBRGB...
 // length of data must be nx * ny * 3 bytes
 //
@@ -198,7 +83,6 @@ enum mtmd_pos_type {
    MTMD_POS_TYPE_NORMAL,    // number of positions equals to number of tokens
    MTMD_POS_TYPE_MROPE,     // qwen-vl mrope style, each image takes max(t,h,w) position indexes
    MTMD_POS_TYPE_HUNYUANVL, // HunyuanVL mrope + BOI/EOI/newline layout with XD-RoPE dim-3
-    MTMD_POS_TYPE_COUNT,     // for validation
 };

 struct mtmd_image_tokens {
@@ -252,30 +136,6 @@ struct mtmd_image_tokens {
            id
        };
    }
-
-    void serialize(mtmd_serialization & ser) const {
-        // remember to bump MTMD_SERIALIZATION_VERSION if this is changed
-        ser.write(nx);
-        ser.write(ny);
-        ser.write((uint32_t)pos);
-        ser.write(image_idx);
-        ser.write(n_temporal_merge);
-        ser.write(id);
-        batch_f32.serialize(ser);
-    }
-    void deserialize(mtmd_serialization & ser) {
-        nx = ser.read<uint32_t>();
-        ny = ser.read<uint32_t>();
-        uint32_t pos_raw = ser.read<uint32_t>();
-        if (pos_raw >= MTMD_POS_TYPE_COUNT) {
-            throw std::runtime_error("invalid pos type");
-        }
-        pos = (mtmd_pos_type)pos_raw;
-        image_idx = ser.read<uint32_t>();
-        n_temporal_merge = ser.read<uint32_t>();
-        id = ser.read<std::string>();
-        batch_f32.deserialize(ser);
-    }
 };
 using mtmd_image_tokens_ptr = std::unique_ptr<mtmd_image_tokens>;

@@ -301,18 +161,6 @@ struct mtmd_audio_tokens {
            id
        };
    }
-
-    void serialize(mtmd_serialization & ser) const {
-        // remember to bump MTMD_SERIALIZATION_VERSION if this is changed
-        ser.write(n_tokens);
-        ser.write(id);
-        batch_f32.serialize(ser);
-    }
-    void deserialize(mtmd_serialization & ser) {
-        n_tokens = ser.read<uint32_t>();
-        id = ser.read<std::string>();
-        batch_f32.deserialize(ser);
-    }
 };
 using mtmd_audio_tokens_ptr = std::unique_ptr<mtmd_audio_tokens>;

@@ -344,66 +192,6 @@ struct mtmd_input_chunk {
        }
        return false;
    }
-
-    void serialize(mtmd_serialization & ser) const {
-        // remember to bump MTMD_SERIALIZATION_VERSION if this is changed
-        ser.write((uint32_t)type);
-
-        ser.write<uint64_t>(tokens_text.size());
-        for (llama_token tok : tokens_text) {
-            ser.write((int32_t)tok);
-        }
-
-        ser.write(tokens_image != nullptr);
-        if (tokens_image) {
-            tokens_image->serialize(ser);
-        }
-
-        ser.write(tokens_audio != nullptr);
-        if (tokens_audio) {
-            tokens_audio->serialize(ser);
-        }
-    }
-    void deserialize(mtmd_serialization & ser) {
-        uint32_t type_raw = ser.read<uint32_t>();
-        if (type_raw >= MTMD_INPUT_CHUNK_TYPE_COUNT) {
-            throw std::runtime_error("invalid chunk type");
-        }
-        type = (mtmd_input_chunk_type)type_raw;
-
-        uint64_t n_tokens_text = ser.read<uint64_t>();
-        // reject before resize() so a tiny corrupted/malicious buffer can't force a huge allocation
-        if (n_tokens_text > (ser.data.size() - ser.read_pos) / sizeof(int32_t)) {
-            throw std::runtime_error("tokens_text length exceeds buffer size");
-        }
-        tokens_text.resize(n_tokens_text);
-        for (uint64_t i = 0; i < n_tokens_text; i++) {
-            tokens_text[i] = (llama_token)ser.read<int32_t>();
-        }
-
-        if (ser.read<bool>()) {
-            tokens_image = std::make_unique<mtmd_image_tokens>();
-            tokens_image->deserialize(ser);
-        } else {
-            tokens_image.reset();
-        }
-
-        if (ser.read<bool>()) {
-            tokens_audio = std::make_unique<mtmd_audio_tokens>();
-            tokens_audio->deserialize(ser);
-        } else {
-            tokens_audio.reset();
-        }
-
-        // catch buffers where the declared type doesn't match which payload is actually present,
-        // so a mismatched chunk can't slip through and null-deref/abort later in an accessor
-        if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE && !tokens_image) {
-            throw std::runtime_error("type is IMAGE but tokens_image is missing");
-        }
-        if (type == MTMD_INPUT_CHUNK_TYPE_AUDIO && !tokens_audio) {
-            throw std::runtime_error("type is AUDIO but tokens_audio is missing");
-        }
-    }
 };

 struct mtmd_input_chunks {
@@ -2255,42 +2043,6 @@ void mtmd_input_chunk_free(mtmd_input_chunk * chunk) {
    }
 }

-int32_t mtmd_input_chunk_save(const mtmd_input_chunk * chunk, char * out_buf, size_t out_len, size_t * expected_out_len) {
-    try {
-        mtmd_serialization ser(MTMD_SERIALIZATION_VERSION);
-        chunk->serialize(ser);
-
-        if (expected_out_len) {
-            *expected_out_len = ser.data.size();
-        }
-        if (!out_buf) {
-            // caller is only querying the required size
-            return 0;
-        }
-        if (out_len < ser.data.size()) {
-            LOG_ERR("%s: out_buf is too small, need %zu bytes, got %zu\n", __func__, ser.data.size(), out_len);
-            return -1;
-        }
-        std::memcpy(out_buf, ser.data.data(), ser.data.size());
-        return 0;
-    } catch (const std::exception & e) {
-        LOG_ERR("%s: %s\n", __func__, e.what());
-        return -1;
-    }
-}
-
-mtmd_input_chunk * mtmd_input_chunk_load(const char * buf, size_t len) {
-    try {
-        mtmd_serialization ser(MTMD_SERIALIZATION_VERSION, buf, len);
-        mtmd::input_chunk_ptr chunk(new mtmd_input_chunk());
-        chunk->deserialize(ser);
-        return chunk.release();
-    } catch (const std::exception & e) {
-        LOG_ERR("%s: %s\n", __func__, e.what());
-        return nullptr;
-    }
-}
-
 // mtmd_image_tokens

 size_t mtmd_image_tokens_get_n_tokens(const mtmd_image_tokens * image_tokens) {
@@ -55,7 +55,6 @@ enum mtmd_input_chunk_type {
    MTMD_INPUT_CHUNK_TYPE_TEXT,
    MTMD_INPUT_CHUNK_TYPE_IMAGE,
    MTMD_INPUT_CHUNK_TYPE_AUDIO,
-    MTMD_INPUT_CHUNK_TYPE_COUNT, // for validation
 };

 // opaque types
@@ -233,15 +232,6 @@ MTMD_API llama_pos                  mtmd_input_chunk_get_n_pos       (const mtmd
 MTMD_API mtmd_input_chunk * mtmd_input_chunk_copy(const mtmd_input_chunk * chunk);
 MTMD_API void               mtmd_input_chunk_free(mtmd_input_chunk * chunk);

-// save/load an input chunk to/from a buffer (useful for KV save/load)
-// important: only chunk's metadata will be saved, the actual image/audio data will not be saved
-// the loaded chunk will always be a placeholder, cannot be used for mtmd_encode() or mtmd_batch_encode()
-// out_buf can be nullptr (to query expected_out_len)
-// returns 0 on success, non-zero on failure
-MTMD_API int32_t            mtmd_input_chunk_save(const mtmd_input_chunk * chunk, char * out_buf, size_t out_len, size_t * expected_out_len);
-// returns nullptr on failure
-MTMD_API mtmd_input_chunk * mtmd_input_chunk_load(const char * buf, size_t len);
-

 // mtmd_image_tokens
 //
@@ -201,7 +201,6 @@ Invoke a tool call, request body is a JSON object with:

 Headers:
 - `x-tool-cwd`: optional; if set, use as the CWD for tool; this is not part of tool's params because it's meant to be set by the runtime, not the LLM itself
- `x-tool-runtime`: optional; if set, run the tool inside this isolate instead of on the host. Only `docker-container:<id>` is supported for now, using an already-running container

 Returns JSON object. There are two response formats (MCP tools use the same two formats: their result content is concatenated into `plain_text_response`, and RPC or tool errors are surfaced as the `error` string):

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

 constexpr int HTTP_POLLING_SECONDS = 1;

-static uint32_t server_n_outputs_max(const common_params & params) {
-    const uint32_t n_batch  = params.n_batch;
-
+static common_speculative_output_limits server_output_limits(const common_params & params) {
    if (params.embedding ||
            (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) {
-        return n_batch;
+        return { params.n_batch, 1 };
    }

-    const uint32_t n_outputs_per_seq = 1 + common_speculative_n_max(&params.speculative);
+    auto result = common_speculative_get_output_limits(
+            params.n_batch, params.n_parallel, common_speculative_n_max(&params.speculative));

-    const uint64_t n_outputs = (uint64_t) params.n_parallel * n_outputs_per_seq;
-
-    return std::max<uint32_t>(1, std::min<uint64_t>(n_batch, n_outputs));
+    result.total   = std::max<int32_t>(1, result.total);
+    result.per_seq = std::max<int32_t>(1, result.per_seq);
+    return result;
 }

 // state diagram: https://github.com/ggml-org/llama.cpp/pull/9283
@@ -1063,7 +1062,9 @@ private:
        const bool is_resume = sleeping;

        params_base = params;
-        params_base.n_outputs_max = server_n_outputs_max(params_base);
+        const auto output_limits = server_output_limits(params_base);
+        params_base.n_outputs_max = output_limits.total;
+        params_base.n_outputs_max_per_seq = output_limits.per_seq;

        const bool has_mmproj = !params.mmproj.path.empty();
        const bool has_draft = params.speculative.has_dft();
@@ -1832,18 +1833,13 @@ private:

            const bool need_pre_sample_logits = task.params.sampling.n_probs > 0 && !task.params.post_sampling_probs;

-            bool backend_sampling = true;
-
-            backend_sampling &= task.params.sampling.backend_sampling;
-
-            // TODO: speculative decoding requires multiple samples per batch - not supported yet
-            backend_sampling &= !(slot.can_speculate());
+            bool use_backend_sampling = task.params.sampling.backend_sampling;

            // TODO: getting pre sampling logits is not yet supported with backend sampling
-            backend_sampling &= !need_pre_sample_logits;
+            use_backend_sampling &= !need_pre_sample_logits;

            // TODO: tmp until backend sampling is fully implemented
-            if (backend_sampling) {
+            if (use_backend_sampling) {
                llama_set_sampler(ctx_tgt, slot.id, common_sampler_get(slot.smpl.get()));
            } else {
                llama_set_sampler(ctx_tgt, slot.id, nullptr);
@@ -3865,7 +3861,8 @@ private:

        // speculative decoding - main model sample and accept
        iterate(slots, [&](server_slot & slot) {
-            if (slot.state != SLOT_STATE_GENERATING || !slot.can_speculate() || slot.spec_draft.empty()) {
+            if (slot.state != SLOT_STATE_GENERATING || !slot.can_speculate() ||
+                    slot.spec_draft.empty() || slot.spec_i_batch.empty()) {
                return;
            }

@@ -3876,7 +3873,6 @@ private:

            // verify and try to accept the draft
            {
-                // save the sampler sampler state in case we need to restore it
                common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get()));

                GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1);
@@ -3915,7 +3911,7 @@ private:
                        slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1);

                        slot.prompt.tokens.keep_first(ckpt.n_tokens);
-                        slot.smpl = std::move(smpl_save);
+                        common_sampler_copy(smpl_save.get(), slot.smpl.get());

                        return;
                    }
@@ -70,188 +70,6 @@ struct server_subproc {
    }
 };

-struct server_lru_sched {
-    server_lru_sched(server_models & models) : models(models) {}
-
-    bool has_capacity(std::unique_lock<std::mutex> & lk) {
-        check_lock(lk);
-        return models.base_params.models_max <= 0
-            || count_running() < (size_t) models.base_params.models_max;
-    }
-
-    // returns "" if no model can be given up
-    std::string pick_victim(std::unique_lock<std::mutex> & lk, const std::string & exclude) {
-        check_lock(lk);
-        std::string victim;
-        int64_t victim_last_used = 0;
-        for (const auto & m : models.mapping) {
-            if (m.first == exclude) {
-                continue;
-            }
-            // a busy model is mid-request, one still coming up has no request to finish
-            if (m.second.req_count != 0 || !m.second.meta.is_ready_or_sleep()) {
-                continue;
-            }
-            if (victim.empty() || m.second.meta.last_used < victim_last_used) {
-                victim           = m.first;
-                victim_last_used = m.second.meta.last_used;
-            }
-        }
-        return victim;
-    }
-
-    // requests wanting the same model share one entry, so they all need only one slot
-    // and all get unblocked by the single load that entry performs
-    void join(std::unique_lock<std::mutex> & lk, const std::string & model_id) {
-        check_lock(lk);
-        if (entry_t * e = find(model_id)) {
-            e->n_waiters++;
-            SRV_INF("request for name=%s joined the queue, %d waiting\n", model_id.c_str(), e->n_waiters);
-            return;
-        }
-        queue.push_back({ model_id, 1, false, false });
-        SRV_INF("models_max reached, request for name=%s queued at position %zu\n",
-                model_id.c_str(), queue.size());
-    }
-
-    void leave(std::unique_lock<std::mutex> & lk, const std::string & model_id) {
-        check_lock(lk);
-        for (auto it = queue.begin(); it != queue.end(); ++it) {
-            if (it->model_id == model_id) {
-                if (--it->n_waiters <= 0) {
-                    queue.erase(it); // last one waiting for this model went away
-                }
-                return;
-            }
-        }
-    }
-
-    bool queue_empty(std::unique_lock<std::mutex> & lk) {
-        check_lock(lk);
-        return queue.empty();
-    }
-
-    // true if it is this model's turn to load, and nobody is loading it yet
-    bool try_claim(std::unique_lock<std::mutex> & lk, const std::string & model_id) {
-        check_lock(lk);
-        if (queue.empty() || queue.front().model_id != model_id || queue.front().loading) {
-            return false;
-        }
-        if (!has_capacity(lk)) {
-            return false;
-        }
-        queue.front().loading = true;
-        return true;
-    }
-
-    // ok means the model is up: drop the entry, the other waiters just watch its status now
-    void claim_done(std::unique_lock<std::mutex> & lk, const std::string & model_id, bool ok) {
-        check_lock(lk);
-        for (auto it = queue.begin(); it != queue.end(); ++it) {
-            if (it->model_id == model_id) {
-                if (ok) {
-                    queue.erase(it);
-                } else {
-                    it->loading = false;
-                }
-                return;
-            }
-        }
-    }
-
-    // a model is on its way out for this entry, so other requests do not also give up one
-    void mark_slot_pending(std::unique_lock<std::mutex> & lk, const std::string & model_id) {
-        check_lock(lk);
-        if (entry_t * e = find(model_id)) {
-            e->slot_pending = true;
-        }
-    }
-
-    // model_id went idle: give up its slot if a queued request needs one
-    // thread-safe, caller must NOT hold models.mutex
-    void on_model_idle(const std::string & model_id) {
-        if (models.base_params.models_max <= 0) {
-            return; // no limit, nothing is ever queued
-        }
-        {
-            std::unique_lock<std::mutex> lk(models.mutex);
-            if (queue.empty()) {
-                return;
-            }
-            size_t promised     = 0;
-            bool   has_unserved = false;
-            for (const auto & e : queue) {
-                if (e.needs_slot()) {
-                    has_unserved = true;
-                } else {
-                    promised++;
-                }
-            }
-            if (!has_unserved) {
-                return;
-            }
-            if ((int) count_running() - (int) promised < models.base_params.models_max) {
-                return; // a slot is already on its way
-            }
-            // never give up a model that a queued request wants
-            for (const auto & e : queue) {
-                if (e.model_id == model_id) {
-                    return;
-                }
-            }
-            auto it = models.mapping.find(model_id);
-            if (it == models.mapping.end() || it->second.req_count != 0 || !it->second.meta.is_ready_or_sleep()) {
-                return;
-            }
-            for (auto & e : queue) {
-                if (!e.slot_pending) {
-                    e.slot_pending = true;
-                    break;
-                }
-            }
-        }
-        SRV_INF("model name=%s went idle, giving up its slot to a queued request\n", model_id.c_str());
-        models.unload(model_id);
-    }
-
-  private:
-    struct entry_t {
-        std::string model_id;
-        int  n_waiters;    // requests waiting for this model
-        bool slot_pending; // a model is already being evicted for this entry
-        bool loading;      // one of the waiters is doing the load right now
-
-        // a slot is already coming, or already taken by the load in flight
-        bool needs_slot() const { return !slot_pending && !loading; }
-    };
-
-    entry_t * find(const std::string & model_id) {
-        for (auto & e : queue) {
-            if (e.model_id == model_id) {
-                return &e;
-            }
-        }
-        return nullptr;
-    }
-
-    void check_lock(std::unique_lock<std::mutex> & lk) {
-        GGML_ASSERT(lk.owns_lock() && lk.mutex() == &models.mutex);
-    }
-
-    size_t count_running() {
-        size_t count = 0;
-        for (const auto & m : models.mapping) {
-            if (m.second.meta.is_running()) {
-                count++;
-            }
-        }
-        return count;
-    }
-
-    server_models & models;
-    std::deque<entry_t> queue;
-};
-
 // short loopback budget for the resumable stream router to child JSON calls (probe, lookup,
 // delete). distinct from params.timeout_read/write which only applies to the generation proxy
 static constexpr int STREAM_LOOKUP_TIMEOUT_MS = 250;
@@ -411,8 +229,7 @@ server_models::server_models(
            : ctx_preset(LLAMA_EXAMPLE_SERVER),
              base_params(params),
              base_env(get_environment()),
-              base_preset(ctx_preset.load_from_args(argc, argv)),
-              sched(std::make_unique<server_lru_sched>(*this)) {
+              base_preset(ctx_preset.load_from_args(argc, argv)) {
    // clean up base preset
    unset_reserved_args(base_preset, true);
    // set binary path
@@ -424,11 +241,8 @@ server_models::server_models(
        LOG_WRN("using original argv[0] as fallback: %s\n", argv[0]);
    }
    load_models();
-    debug_fake_timing = !common_get_env("LLAMA_SERVER_DEBUG_FAKE_TIMING").empty();
 }

-server_models::~server_models() = default;
-
 void server_models::add_model(server_model_meta && meta) {
    if (mapping.find(meta.name) != mapping.end()) {
        throw std::runtime_error(string_format("model '%s' appears multiple times", meta.name.c_str()));
@@ -899,15 +713,22 @@ void server_models::unload_lru() {
        return; // no limit
    }
    // remove one of the servers if we passed the models_max (least recently used - LRU)
-    std::string lru_model_name;
+    std::string lru_model_name = "";
+    int64_t lru_last_used = ggml_time_ms();
+    size_t count_active = 0;
    {
        std::unique_lock<std::mutex> lk(mutex);
-        if (sched->has_capacity(lk)) {
-            return;
+        for (const auto & m : mapping) {
+            if (m.second.meta.is_running()) {
+                count_active++;
+                if (m.second.meta.last_used < lru_last_used) {
+                    lru_model_name = m.first;
+                    lru_last_used = m.second.meta.last_used;
+                }
+            }
        }
-        lru_model_name = sched->pick_victim(lk, "");
    }
-    if (!lru_model_name.empty()) {
+    if (!lru_model_name.empty() && count_active >= (size_t)base_params.models_max) {
        SRV_INF("models_max limit reached, removing LRU name=%s\n", lru_model_name.c_str());
        unload(lru_model_name);
        // wait for unload to complete
@@ -925,11 +746,6 @@ void server_models::load(const std::string & name) {
 }

 void server_models::load(const std::string & name, const load_options & opts) {
-    if (debug_fake_timing) {
-        // do not hold the mutex here, other requests must keep making progress
-        std::this_thread::sleep_for(std::chrono::seconds(2));
-    }
-
    if (!opts.custom_meta.has_value()) {
        if (!has_model(name)) {
            throw std::runtime_error("model name=" + name + " is not found");
@@ -1322,7 +1138,7 @@ void server_models::wait(std::unique_lock<std::mutex> & lk, const std::string &
    });
 }

-bool server_models::ensure_model_ready(const std::string & name, const std::function<bool()> & should_stop) {
+bool server_models::ensure_model_ready(const std::string & name) {
    auto meta = get_meta(name);
    if (!meta.has_value()) {
        throw std::runtime_error("model name=" + name + " is not found");
@@ -1333,112 +1149,25 @@ bool server_models::ensure_model_ready(const std::string & name, const std::func
    if (meta->status == SERVER_MODEL_STATUS_SLEEPING) {
        return false; // child is sleeping but still running; new request will wake it up
    }
-
-    bool queued   = false;
-    bool did_load = false;
-    std::string victim;
-    {
-        std::unique_lock<std::mutex> lk(mutex);
-        auto it = mapping.find(name);
-        if (it != mapping.end() && it->second.meta.status == SERVER_MODEL_STATUS_UNLOADED) {
-            bool has_capacity = sched->has_capacity(lk);
-            if (has_capacity && sched->queue_empty(lk)) {
-                lk.unlock();
-                SRV_INF("model name=%s is not loaded, loading...\n", name.c_str());
-                load(name);
-                did_load = true;
-            } else {
-                // also queue when a slot looks free but others wait already, else they starve
-                sched->join(lk, name);
-                queued = true;
-                if (!has_capacity) {
-                    // an idle model may sit here right now, do not wait for a request to end
-                    victim = sched->pick_victim(lk, name);
-                    if (!victim.empty()) {
-                        sched->mark_slot_pending(lk, name);
-                    }
-                }
-            }
-        }
-    }
-    if (!victim.empty()) {
-        SRV_INF("evicting idle LRU name=%s to make room for name=%s\n", victim.c_str(), name.c_str());
-        unload(victim);
+    if (meta->status == SERVER_MODEL_STATUS_UNLOADED) {
+        SRV_INF("model name=%s is not loaded, loading...\n", name.c_str());
+        load(name);
    }

-    // while queued, this is also where the load happens: the head of the queue does it
+    // wait for loading to complete
    SRV_INF("waiting until model name=%s is fully loaded...\n", name.c_str());
-    std::unique_lock<std::mutex> lk(mutex);
-    auto leave_queue = [this, &queued, &lk, &name]() {
-        if (queued) {
-            sched->leave(lk, name);
-            queued = false;
+    wait(name, [&meta](const server_model_meta & new_meta) {
+        if (new_meta.status != SERVER_MODEL_STATUS_LOADING) {
+            meta = new_meta; // update meta for final check after wait
+            return true;
        }
-    };
+        return false;
+    });

-    try {
-        bool saw_loading = false;
-        while (true) {
-            auto it = mapping.find(name);
-            if (it == mapping.end()) {
-                break; // removed by another code path, nothing to wait for
-            }
-            const server_model_status status = it->second.meta.status;
-
-            if (status == SERVER_MODEL_STATUS_LOADED || status == SERVER_MODEL_STATUS_SLEEPING) {
-                break;
-            }
-            if (status == SERVER_MODEL_STATUS_DOWNLOADING || status == SERVER_MODEL_STATUS_DOWNLOADED) {
-                break; // do not wait on a download child
-            }
-            if (status == SERVER_MODEL_STATUS_LOADING) {
-                saw_loading = true;
-            } else if (status == SERVER_MODEL_STATUS_UNLOADED) {
-                if (did_load || saw_loading) {
-                    // a spawn happened and the instance came back down
-                    if (it->second.meta.is_failed()) {
-                        throw std::runtime_error("model name=" + name + " failed to load");
-                    }
-                    break; // unloaded by another code path, caller reports "not running"
-                }
-                if (!queued) {
-                    break; // not queued, and the load someone else started fell over
-                }
-            }
-
-            if (should_stop && should_stop()) {
-                // if a model was evicted for us, the free slot goes to the next waiter
-                throw std::runtime_error("request cancelled while waiting for model name=" + name);
-            }
-
-            // our turn: our model is at the head, and a slot really did free up
-            if (status == SERVER_MODEL_STATUS_UNLOADED && sched->try_claim(lk, name)) {
-                lk.unlock();
-                bool ok = true;
-                try {
-                    SRV_INF("slot available, loading queued model name=%s\n", name.c_str());
-                    load(name);
-                    did_load = true;
-                } catch (const std::exception & e) {
-                    // lost a race for the slot, stay in line and retry
-                    SRV_WRN("queued load of name=%s did not go through: %s\n", name.c_str(), e.what());
-                    ok = false;
-                }
-                lk.lock();
-                sched->claim_done(lk, name, ok);
-                if (ok) {
-                    queued = false; // entry is gone, the other waiters watch the status now
-                }
-                continue;
-            }
-
-            cv.wait_for(lk, std::chrono::milliseconds(200));
-        }
-    } catch (...) {
-        leave_queue();
-        throw;
+    // check final status
+    if (!meta.has_value() || meta->is_failed()) {
+        throw std::runtime_error("model name=" + name + " failed to load");
    }
-    leave_queue();

    return true;
 }
@@ -1451,16 +1180,9 @@ server_http_res_ptr server_models::proxy_request(const server_http_req & req, co
    if (!meta->is_running()) {
        throw std::invalid_argument("model name=" + name + " is not running");
    }
-    {
+    if (update_last_used) {
        std::unique_lock<std::mutex> lk(mutex);
-        if (update_last_used) {
-            mapping[name].meta.last_used = ggml_time_ms();
-        }
-        mapping[name].req_count++;
-    }
-    if (debug_fake_timing) {
-        // sleep after req_count++, so the model counts as busy while we wait here
-        std::this_thread::sleep_for(std::chrono::seconds(2));
+        mapping[name].meta.last_used = ggml_time_ms();
    }
    SRV_INF("proxying request to model %s on port %d\n", name.c_str(), meta->port);
    std::string proxy_path = req.path;
@@ -1476,29 +1198,13 @@ server_http_res_ptr server_models::proxy_request(const server_http_req & req, co
            req.headers,
            req.body,
            req.files,
-            // a detached request belongs to a replay session
-            detached
-                ? std::function<bool()>([]() { return false; })
-                : req.should_stop,
+            // a detached request belongs to a replay session that outlives the client socket:
+            // it reaches the child even when the downstream died during the load wait, the
+            // session buffer is the recipient and DELETE remains the stop
+            detached ? std::function<bool()>([]() { return false; }) : req.should_stop,
            base_params.timeout_read,
            base_params.timeout_write
            );
-
-    proxy->cleanup = [this, name]() {
-        bool went_idle = false;
-        {
-            std::unique_lock<std::mutex> lk(mutex);
-            auto it = mapping.find(name);
-            if (it != mapping.end() && it->second.req_count > 0) {
-                it->second.req_count--;
-                went_idle = it->second.req_count == 0;
-            }
-        }
-        if (went_idle) {
-            sched->on_model_idle(name);
-        }
-    };
-
    return proxy;
 }

@@ -1862,7 +1568,7 @@ void server_models_routes::init_routes() {
            return error_res;
        }
        if (autoload) {
-            models.ensure_model_ready(name, req.should_stop);
+            models.ensure_model_ready(name);
        }
        return models.proxy_request(req, method, name, false);
    };
@@ -1882,9 +1588,7 @@ void server_models_routes::init_routes() {
        // this request instead of leaving an orphan generation
        std::string conv_id = server_stream_conv_id_from_headers(req.headers);
        uint64_t ticket = models.conv_models.remember(conv_id, name);
-        // a dead socket must not cancel a session request, only a stop does (checked right below)
-        auto should_stop = ticket == 0 ? req.should_stop : nullptr;
-        bool waited = autoload && models.ensure_model_ready(name, should_stop);
+        bool waited = autoload && models.ensure_model_ready(name);
        if (ticket != 0 && !models.conv_models.alive(conv_id, ticket)) {
            SRV_INF("request for conv_id=%s cancelled while model name=%s was loading\n",
                    conv_id.c_str(), name.c_str());
@@ -2360,7 +2064,7 @@ server_http_proxy::server_http_proxy(
    cli->set_write_timeout(timeout_read, 0); // reversed for cli (client) vs srv (server)
    cli->set_read_timeout(timeout_write, 0);
    this->status = 500; // to be overwritten upon response
-    this->cleanup_pipes = [pipe]() {
+    this->cleanup = [pipe]() {
        pipe->close_read();
        pipe->close_write();
    };
@@ -2375,8 +2079,9 @@ server_http_proxy::server_http_proxy(
        return has_next; // false if EOF or pipe broken
    };

-    // build the header message forwarded to the reader thread, stripping internal proxy headers
-    auto make_header_msg = [](const httplib::Response & response) {
+    // wire up the HTTP client
+    // note: do NOT capture `this` pointer, as it may be destroyed before the thread ends
+    httplib::ResponseHandler response_handler = [pipe, cli](const httplib::Response & response) {
        msg_t msg;
        msg.status = response.status;
        for (const auto & [key, value] : response.headers) {
@@ -2390,17 +2095,7 @@ server_http_proxy::server_http_proxy(
            }
            msg.headers[key] = value;
        }
-        return msg;
-    };
-
-    // true once response_handler has already forwarded the headers
-    auto headers_sent = std::make_shared<std::atomic<bool>>(false);
-
-    // wire up the HTTP client
-    // note: do NOT capture `this` pointer, as it may be destroyed before the thread ends
-    httplib::ResponseHandler response_handler = [pipe, headers_sent, make_header_msg](const httplib::Response & response) {
-        headers_sent->store(true);
-        return pipe->write(make_header_msg(response)); // send headers first
+        return pipe->write(std::move(msg)); // send headers first
    };
    httplib::ContentReceiverWithProgress content_receiver = [pipe](const char * data, size_t data_length, size_t, size_t) {
        // send data chunks
@@ -2474,16 +2169,13 @@ server_http_proxy::server_http_proxy(

    // start the proxy thread
    SRV_DBG("start proxy thread %s %s\n", req.method.c_str(), req.path.c_str());
-    this->thread = std::thread([cli, pipe, req, headers_sent, make_header_msg]() {
+    this->thread = std::thread([cli, pipe, req]() {
        auto result = cli->send(std::move(req));
        if (result.error() != httplib::Error::Success) {
            auto err_str = httplib::to_string(result.error());
            SRV_ERR("http client error: %s\n", err_str.c_str());
            pipe->write({{}, 500, "", ""}); // header
            pipe->write({{}, 0, "proxy error: " + err_str, ""}); // body
-        } else if (!headers_sent->load()) {
-            // httplib skips response_handler for bodyless statuses like 204, send headers here instead
-            pipe->write(make_header_msg(*result));
        }
        pipe->close_write(); // signal EOF to reader
        SRV_DBG("%s", "client request thread ended\n");
@@ -84,6 +84,7 @@ struct server_model_meta {
    int exit_code = 0; // exit code of the model instance process (only valid if status == FAILED)
    int stop_timeout = 0; // seconds to wait before force-killing the model instance during shutdown
    mtmd_caps multimodal; // multimodal capabilities
+    // bool need_download = false; // whether the model needs to be downloaded before loading // TODO @ngxson: implement this

    bool is_ready() const {
        return status == SERVER_MODEL_STATUS_LOADED;
@@ -93,10 +94,6 @@ struct server_model_meta {
        return status == SERVER_MODEL_STATUS_LOADED || status == SERVER_MODEL_STATUS_LOADING || status == SERVER_MODEL_STATUS_SLEEPING;
    }

-    bool is_ready_or_sleep() const {
-        return status == SERVER_MODEL_STATUS_LOADED || status == SERVER_MODEL_STATUS_SLEEPING;
-    }
-
    bool is_failed() const {
        return status == SERVER_MODEL_STATUS_UNLOADED && exit_code != 0;
    }
@@ -106,19 +103,16 @@ struct server_model_meta {
 };

 struct server_models_routes;
-struct server_subproc;   // defined in server-models.cpp
-struct server_lru_sched; // defined in server-models.cpp
+struct server_subproc; // defined in server-models.cpp

 struct server_models {
    friend struct server_models_routes;
-    friend struct server_lru_sched;

 private:
    struct instance_t {
        std::shared_ptr<server_subproc> subproc; // shared between main thread and monitoring thread
        std::thread th;
        server_model_meta meta;
-        int req_count = 0; // number of active proxy requests
    };

    std::mutex mutex;
@@ -197,12 +191,6 @@ private:
    std::vector<std::string> base_env;
    common_preset base_preset; // base preset from llama-server CLI args

-    // queue of requests waiting for a models_max slot
-    std::unique_ptr<server_lru_sched> sched;
-
-    // if true, add some delay to simulate works (useful for testing)
-    bool debug_fake_timing = false;
-
    void update_meta(const std::string & name, const server_model_meta & meta);

    // unload least recently used models if the limit is reached
@@ -219,7 +207,6 @@ public:
    conv_model_tracker conv_models;

    server_models(const common_params & params, int argc, char ** argv);
-    ~server_models();

    server_response sse; // for real-time updates via SSE endpoint

@@ -276,9 +263,7 @@ public:
    // ensure the model is in ready state (thread-safe)
    // return false if model is ready
    // otherwise, load the model and blocking wait until it's ready, then return true (meta may need to be refreshed)
-    // if models_max is reached, the request waits in a queue until a slot frees up
-    // throws if the load fails, or if should_stop fires while waiting
-    bool ensure_model_ready(const std::string & name, const std::function<bool()> & should_stop = nullptr);
+    bool ensure_model_ready(const std::string & name);

    // proxy an HTTP request to the model instance
    server_http_res_ptr proxy_request(const server_http_req & req, const std::string & method, const std::string & name, bool update_last_used, bool detached = false);
@@ -358,6 +343,7 @@ struct server_models_routes {
 */
 struct server_http_proxy : server_http_res {
    std::function<void()> cleanup = nullptr;
+public:
    server_http_proxy(const std::string & method,
                      const std::string & scheme,
                      const std::string & host,
@@ -371,15 +357,11 @@ struct server_http_proxy : server_http_res {
                      int32_t timeout_write
                      );
    ~server_http_proxy() {
-        if (cleanup_pipes) {
-            cleanup_pipes();
-        }
        if (cleanup) {
            cleanup();
        }
    }
 private:
-    std::function<void()> cleanup_pipes = nullptr;
    std::thread thread;
    struct msg_t {
        std::map<std::string, std::string> headers;
@@ -519,7 +519,7 @@ task_params eval_llama_cmpl_schema(
                const json & data) {
    task_params params;

-    // Sampling parameter defaults are loaded from the global server context (but individual requests can still override them)
+    // Sampling parameter defaults are loaded from the global server context (but individual requests can still them)
    params.sampling      = params_base.sampling;
    params.speculative   = params_base.speculative;
    params.n_keep        = params_base.n_keep;
@@ -14,7 +14,6 @@ struct server_tool {
    std::string display_name;
    bool permission_write = false;
    bool support_stream = false; // if true, output can be streamed
-    bool uses_cwd = false;       // if true, the tool resolves paths and runs against the working directory

    virtual ~server_tool() = default;
    virtual json get_definition() const = 0;
@@ -31,8 +30,6 @@ struct server_tool {
    json to_json() const;
 };

-struct server_tools_docker_runtime; // impl detail, defined in server-tools.cpp
-
 struct server_tools {
    std::vector<std::unique_ptr<server_tool>> tools;

@@ -40,16 +37,9 @@ struct server_tools {
    server_response queue_res;
    std::atomic<int> res_id{0};

-    // set when --tools-runtime is configured; owns the docker container used to run tools, if any
-    std::unique_ptr<server_tools_docker_runtime> docker_runtime;
-
    void setup(const std::vector<std::string> & enabled_tools,
-               server_mcp & mcp_mgr,
-               const std::string & tools_runtime);
+               server_mcp & mcp_mgr);

    server_http_context::handler_t handle_get;
    server_http_context::handler_t handle_post;
-
-    server_tools();
-    ~server_tools();
 };
@@ -338,7 +338,7 @@ int llama_server(common_params & params, int argc, char ** argv) {

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


-@pytest.fixture(scope="session", autouse=True)
-def load_server_presets():
+@pytest.fixture(scope="module", autouse=True)
+def do_something():
    # this will be run once per test session, before any tests
    ServerPreset.load_all()
@@ -14,10 +14,10 @@ fi
 if [ $# -lt 1 ]
 then
    if [[ "${SLOW_TESTS:-0}" == 1 ]]; then
-        pytest --durations=30 -v -x
+        pytest -v -x
    else
-        pytest --durations=30 -v -x -m "not slow"
+        pytest -v -x -m "not slow"
    fi
 else
-    pytest --durations=30 "$@"
+    pytest "$@"
 fi
@@ -1,7 +1,5 @@
 import pytest
 from utils import *
-import threading
-from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

 server = ServerPreset.tinyllama2()

@@ -41,31 +39,3 @@ def test_mcp_proxy_custom_port():
    res = server.make_request("GET", f"/cors-proxy?url=http://{server.server_host}:{server.server_port}/models")
    assert res.status_code == 200
    assert "data" in res.body
-
-
-def test_mcp_proxy_no_content():
-    # note: see issue #26598
-    class NoContentHandler(BaseHTTPRequestHandler):
-        def do_POST(self):
-            self.send_response(204)
-            self.end_headers()
-
-        def log_message(self, format, *args):
-            pass
-
-    target = ThreadingHTTPServer(("127.0.0.1", 0), NoContentHandler)
-    target_thread = threading.Thread(target=target.serve_forever, daemon=True)
-    target_thread.start()
-
-    try:
-        global server
-        server.ui_mcp_proxy = True
-        server.start()
-
-        res = server.make_request("POST", f"/cors-proxy?url=http://127.0.0.1:{target.server_port}/", data={})
-        assert res.status_code == 204
-        assert res.body in (None, b"", "")
-    finally:
-        target.shutdown()
-        target.server_close()
-
@@ -85,7 +85,7 @@ def _wait_for_model_status(model_id: str, desired: set[str], timeout: int = 60)
        last_status = _get_model_status(model_id)
        if last_status in desired:
            return last_status
-        time.sleep(0.01)
+        time.sleep(1)
    raise AssertionError(
        f"Timed out waiting for {model_id} to reach {desired}, last status: {last_status}"
    )
@@ -145,156 +145,6 @@ def test_router_models_max_evicts_lru():
    assert _get_model_status(first) == "unloaded"


-# server_lru_sched tests (relying on LLAMA_SERVER_DEBUG_FAKE_TIMING)
-
-MODEL_A = "ggml-org/tinygemma3-GGUF:Q8_0"
-MODEL_B = "ggml-org/test-model-stories260K:F32"
-MODEL_C = "ggml-org/test-model-stories260K-infill:F32"
-
-
-def _tokenize(model_id: str, timeout: float | None = DEFAULT_REQUEST_TIMEOUT) -> ServerResponse:
-    return server.make_request(
-        "POST", "/tokenize", data={"model": model_id, "content": "hello world"}, timeout=timeout
-    )
-
-
-class _Bg:
-    """runs one request in a thread, keeps its result, error and finish time"""
-
-    def __init__(self, fn):
-        self.result = None
-        self.error: Exception | None = None
-        self.done_at: float = 0.0
-        self._thread = threading.Thread(target=self._run, args=(fn,), daemon=True)
-
-    def _run(self, fn):
-        try:
-            self.result = fn()
-        except Exception as e:
-            self.error = e
-        self.done_at = time.time()
-
-    def start(self):
-        self._thread.start()
-        return self
-
-    def join(self, timeout: int = 180):
-        self._thread.join(timeout)
-        assert not self._thread.is_alive(), "background request did not finish in time"
-        return self
-
-    def assert_ok(self, what: str):
-        assert self.error is None, f"{what} raised {self.error!r}"
-        assert self.result is not None and self.result.status_code == 200, \
-            f"{what} failed: {self.result.status_code if self.result else None} {self.result.body if self.result else None}"
-
-
-def test_router_queue_does_not_evict_busy_model():
-    """a request that finds no free slot waits, and the model serving a request survives it"""
-    global server
-    server.models_max = 1
-    server.start()
-
-    _load_model_and_wait(MODEL_A, timeout=120)
-
-    busy = _Bg(lambda: _tokenize(MODEL_A)).start()
-    time.sleep(0.5)  # let the request reach the child and take the only slot
-
-    # no slot free and MODEL_A is busy, so this queues instead of evicting mid-request
-    queued = _Bg(lambda: _tokenize(MODEL_B)).start()
-
-    busy.join()
-    queued.join()
-
-    # had MODEL_A been evicted while serving, its own request would have died
-    busy.assert_ok("request against the busy model")
-    queued.assert_ok("queued request")
-
-    _wait_for_model_status(MODEL_B, {"loaded"}, timeout=120)
-    assert _get_model_status(MODEL_A) == "unloaded"
-
-
-def test_router_queue_coalesces_requests_for_same_model():
-    """many requests for one missing model share a slot, so only one model is given up"""
-    global server
-    server.models_max = 2
-    server.start()
-
-    _load_model_and_wait(MODEL_A, timeout=120)
-    _load_model_and_wait(MODEL_B, timeout=120)
-
-    # keep MODEL_A busy so MODEL_B is the only model that can be given up
-    busy = _Bg(lambda: _tokenize(MODEL_A)).start()
-    time.sleep(0.5)
-
-    waiters = [_Bg(lambda: _tokenize(MODEL_C)).start() for _ in range(3)]
-
-    busy.join()
-    for w in waiters:
-        w.join()
-
-    busy.assert_ok("request against the busy model")
-    for i, w in enumerate(waiters):
-        w.assert_ok(f"queued request {i}")
-
-    _wait_for_model_status(MODEL_C, {"loaded"}, timeout=120)
-    # one entry for 3 requests means one eviction: MODEL_B goes, MODEL_A is left alone.
-    # without coalescing the leftover entries still ask for a slot,
-    # and MODEL_A is taken too as soon as it goes idle
-    assert _get_model_status(MODEL_A) == "loaded"
-    assert _get_model_status(MODEL_B) == "unloaded"
-
-
-def test_router_queue_client_disconnect_keeps_model():
-    """a client that leaves while queued must not cost a running model its slot"""
-    global server
-    server.models_max = 1
-    server.start()
-
-    _load_model_and_wait(MODEL_A, timeout=120)
-
-    busy = _Bg(lambda: _tokenize(MODEL_A)).start()
-    time.sleep(0.5)
-
-    # queues behind MODEL_A, then gives up long before MODEL_A goes idle
-    with pytest.raises(requests.exceptions.RequestException):
-        _tokenize(MODEL_B, timeout=1)
-
-    busy.join()
-    busy.assert_ok("request against the busy model")
-
-    # nobody is waiting anymore, so MODEL_A keeps its slot
-    time.sleep(3)
-    assert _get_model_status(MODEL_A) == "loaded"
-    assert _get_model_status(MODEL_B) == "unloaded"
-
-
-def test_router_queue_is_fifo():
-    """the queue is served in arrival order"""
-    global server
-    server.models_max = 1
-    server.start()
-
-    _load_model_and_wait(MODEL_A, timeout=120)
-
-    busy = _Bg(lambda: _tokenize(MODEL_A)).start()
-    time.sleep(0.5)
-
-    first = _Bg(lambda: _tokenize(MODEL_B)).start()
-    time.sleep(1)  # keep the arrival order unambiguous
-    second = _Bg(lambda: _tokenize(MODEL_C)).start()
-
-    busy.join()
-    first.join()
-    second.join()
-
-    busy.assert_ok("request against the busy model")
-    first.assert_ok("first queued request")
-    second.assert_ok("second queued request")
-
-    assert first.done_at < second.done_at, "queue was not served in arrival order"
-
-
 def test_router_no_models_autoload():
    global server
    server.no_models_autoload = True
@@ -460,7 +310,7 @@ def _wait_for_sse_event(collected: list, event_type: str, model: str, timeout: i
    while time.time() < deadline:
        if any(e.get("event") == event_type and e.get("model") == model for e in collected):
            return True
-        time.sleep(0.01)
+        time.sleep(0.5)
    return False


@@ -25,33 +25,34 @@ def fixture_create_server():

 def test_with_and_without_draft():
    global server
+    request = {
+        "prompt": "I believe the meaning of life is",
+        "temperature": 0.8,
+        "top_k": 40,
+        "seed": 4242,
+        "n_predict": 16,
+        "return_tokens": True,
+    }
+
    server.model_draft = None  # disable draft model
    server.spec_type = None
+    server.backend_sampling = True
    server.start()
-    res = server.make_request("POST", "/completion", data={
-        "prompt": "I believe the meaning of life is",
-        "temperature": 0.0,
-        "top_k": 1,
-        "n_predict": 16,
-    })
+    res = server.make_request("POST", "/completion", data=request)
    assert res.status_code == 200
-    content_no_draft = res.body["content"]
+    tokens_no_draft = res.body["tokens"]
    server.stop()

    # create new server with draft model
    create_server()
+    server.backend_sampling = True
    server.start()
-    res = server.make_request("POST", "/completion", data={
-        "prompt": "I believe the meaning of life is",
-        "temperature": 0.0,
-        "top_k": 1,
-        "n_predict": 16,
-    })
+    res = server.make_request("POST", "/completion", data=request)
    assert res.status_code == 200
    assert res.body["timings"]["draft_n"] > 0
-    content_draft = res.body["content"]
+    tokens_draft = res.body["tokens"]

-    assert content_no_draft == content_draft
+    assert tokens_no_draft == tokens_draft


 def test_different_draft_min_draft_max():
@@ -1,6 +1,4 @@
 import os
-import shutil
-import subprocess

 import pytest
 from utils import *
@@ -148,95 +146,6 @@ def test_tools_builtin_cwd_header():
            os.remove(marker_path)


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


-def test_tools_builtin_file_glob_search_junk_dirs(tmp_path):
-    global server
-    server.start()
-
-    (tmp_path / "build" / "nested").mkdir(parents=True)
-    (tmp_path / "build" / "artifact.txt").write_text("built")
-    (tmp_path / "src").mkdir()
-    (tmp_path / "src" / "main.cpp").write_text("int main() {}")
-
-    # a junk directory stays selectable as a working directory
-    res = call_tool("file_glob_search", {"path": str(tmp_path), "type": "dir", "max_depth": 1})
-    assert "build" in [e["path"] for e in res["entries"]]
-
-    # but it is never walked, so nothing inside it shows up
-    res = call_tool("file_glob_search", {"path": str(tmp_path), "type": "all"})
-    paths = [e["path"] for e in res["entries"]]
-    assert "src/main.cpp" in paths
-    assert "build/artifact.txt" not in paths
-    assert "build/nested" not in paths
-
-
 def test_tools_builtin_file_glob_search_rejects_invalid_type(tmp_path):
    global server
    server.start()
@@ -115,7 +115,6 @@ class ServerProcess:
    backend_sampling: bool = False
    gcp_compat: bool = False
    server_tools: str | None = None
-    server_tools_runtime: str | None = None
    mcp_servers_config: str | None = None
    mcp_servers_json: str | None = None
    cors_origins: str | None = None
@@ -133,10 +132,7 @@ class ServerProcess:
        self.external_server = "DEBUG_EXTERNAL" in os.environ

    def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None:
-        env = {
-            **os.environ,
-            "LLAMA_SERVER_DEBUG_FAKE_TIMING": "1",
-        }
+        env = {**os.environ}
        if "LLAMA_CACHE" not in os.environ:
            env["LLAMA_CACHE"] = "tmp"
        if self.external_server:
@@ -271,8 +267,6 @@ class ServerProcess:
            server_args.append("--ui-mcp-proxy")
        if self.server_tools:
            server_args.extend(["--tools", self.server_tools])
-        if self.server_tools_runtime:
-            server_args.extend(["--tools-runtime", self.server_tools_runtime])
        if self.mcp_servers_config:
            server_args.extend(["--mcp-servers-config", self.mcp_servers_config])
        if self.mcp_servers_json:
@@ -309,7 +303,6 @@ class ServerProcess:

        # wait for server to start
        start_time = time.time()
-        last_print_time = start_time
        while time.time() - start_time < timeout_seconds:
            try:
                response = self.make_request("GET", "/health", headers={
@@ -324,10 +317,8 @@ class ServerProcess:
            if self.process.poll() is not None:
                raise RuntimeError(f"Server process died with return code {self.process.returncode}")

-            if time.time() - last_print_time >= 1.0:
-                print(f"Waiting for server to start...")
-                last_print_time = time.time()
-            time.sleep(0.01)
+            print(f"Waiting for server to start...")
+            time.sleep(0.5)
        raise TimeoutError(f"Server did not start within {timeout_seconds} seconds")

    def stop(self) -> None:
@@ -179,20 +179,17 @@ int main(int argc, char ** argv) {
    const char * data        = nullptr;
    size_t       data_len    = 0;
    int64_t      n_samples   = 0;
-    const int64_t t_wav_start_us = ggml_time_us();
    if (gen.get_output(&sample_rate, &data, &data_len, &n_samples) != 0) {
        LOG_ERR("get_output failed\n");
        return 1;
    }
-    const double t_wav_s = (ggml_time_us() - t_wav_start_us) / 1e6;

    LOG_INF("generated %d frames, %zu bytes of WAV audio (%d Hz)\n", n_frames, data_len, sample_rate);

    const double t_prompt_s = (t_gen_start_us - t_prompt_start_us) / 1e6;
-    const double t_total_s  = t_prompt_s + t_gen_s + t_wav_s;
+    const double t_total_s  = t_prompt_s + t_gen_s;
    const double audio_s    = sample_rate > 0 ? (double) n_samples / sample_rate : 0.0;
-    LOG_INF("timings: prompt eval %.2fs + generation %.2fs + vocoder %.2fs = total %.2fs\n",
-            t_prompt_s, t_gen_s, t_wav_s, t_total_s);
+    LOG_INF("timings: prompt eval %.2fs + generation %.2fs = total %.2fs\n", t_prompt_s, t_gen_s, t_total_s);
    LOG_INF("         output audio = %.2fs (audio time = %.2fx process time)\n", audio_s, t_total_s > 0 ? audio_s / t_total_s : 0.0);
    FILE * f = fopen(params.out_file.c_str(), "wb");
    if (!f) {
@@ -1,3 +1,2 @@
 engine-strict=true
 ignore-scripts=true
-min-release-age=7
@@ -143,8 +143,10 @@ declare global {
 		idxThemeStyle?: number;
 		idxCodeBlock?: number;

-		// File System Access API - not in the DOM lib and unavailable in some browsers
-		showDirectoryPicker?: (options?: {
+		// File System Access API - missing from older DOM lib versions.
+		// Used by ChatFormWorkingDirectory's native folder picker. Feature availability
+		// is gated at runtime via `typeof window.showDirectoryPicker === 'function'`.
+		showDirectoryPicker: (options?: {
 			id?: string;
 			mode?: 'read' | 'readwrite';
 			startIn?: FileSystemHandle | string;
@@ -2,7 +2,6 @@
 	import {
 		ChatAttachmentsList,
 		ChatFormActions,
-		ChatFormContenteditable,
 		ChatFormFileInputInvisible,
 		ChatFormMcpResourcesList,
 		ChatFormPickers,
@@ -15,7 +14,9 @@
 		INPUT_CLASSES,
 		SETTING_CONFIG_DEFAULT,
 		INITIAL_FILE_SIZE,
-		PROMPT_CONTENT_SEPARATOR
+		PROMPT_CONTENT_SEPARATOR,
+		PROMPT_TRIGGER_PREFIX,
+		RESOURCE_TRIGGER_PREFIX
 	} from '$lib/constants';
 	import {
 		ContentPartType,
@@ -38,25 +39,8 @@
 		activeConversation,
 		pendingCwd
 	} from '$lib/stores/conversations.svelte';
-	import type {
-		FileMentionEntry,
-		GetPromptResult,
-		MCPPromptInfo,
-		MCPResourceInfo,
-		PromptMessage
-	} from '$lib/types';
-	import {
-		buildMentionInsertion,
-		containsCodeSpan,
-		containsFileMentionLink,
-		findCommandToken,
-		findMentionToken,
-		isIMEComposing,
-		isOffsetInCodeBlock,
-		parseClipboardContent,
-		uuid
-	} from '$lib/utils';
-	import { useChatFormPickers } from '$lib/hooks/use-chat-form-pickers.svelte';
+	import type { GetPromptResult, MCPPromptInfo, MCPResourceInfo, PromptMessage } from '$lib/types';
+	import { isIMEComposing, parseClipboardContent, uuid } from '$lib/utils';
 	import {
 		AudioRecorder,
 		convertToWav,
@@ -113,67 +97,29 @@
 	}: Props = $props();

 	// Component References
-	// Shared handle of the two input renderers (textarea + contenteditable).
-	type ChatInputHandle = {
-		focus(): void;
-		resetHeight(): void;
-		getElement(): HTMLElement | undefined;
-		getCaretOffset(): number;
-		setCaretOffset(offset: number): void;
-	};
-
 	let audioRecorder: AudioRecorder | undefined;
 	let chatFormActionsRef: ChatFormActions | undefined = $state(undefined);
 	let fileInputRef: ChatFormFileInputInvisible | undefined = $state(undefined);
 	let pickersRef: { handleKeydown: (event: KeyboardEvent) => boolean } | undefined =
 		$state(undefined);
-	let inputRef: ChatInputHandle | undefined = $state(undefined);
-
-	// Render-mode gate: the plain textarea by default, the contenteditable
-	// while the buffer carries a `file://` mention link or a complete code
-	// span (badges and code chips need a DOM the textarea cannot provide).
-	// Demotes back once neither remains.
-	let useContenteditable = $state(false);
+	let textareaRef: ChatFormTextarea | undefined = $state(undefined);

 	// Audio Recording State
 	let isRecording = $state(false);
 	let recordingSupported = $state(false);

-	// Invisible anchor at the form's top edge so the mention/WD popovers
-	// float above the box.
-	let mentionAnchor: HTMLDivElement | null = $state(null);
+	// Picker State
+	let isPromptPickerOpen = $state(false);
+	let promptSearchQuery = $state('');
+	let isInlineResourcePickerOpen = $state(false);
+	let resourceSearchQuery = $state('');

 	let cwd = $derived(activeConversation()?.cwd ?? pendingCwd());

-	const pickers = useChatFormPickers({
-		getValue: () => value,
-		setValue: (v) => {
-			value = v;
-			onValueChange?.(v);
-		},
-		getCaretOffset: () => inputRef?.getCaretOffset(),
-		setCaretOffset: (offset) => inputRef?.setCaretOffset(offset),
-		focusInput: refocusInput,
-		getShowModelSelector: () => showModelSelector,
-		hasPrompts: () => mcpStore.hasPromptsCapability(conversationsStore.getAllMcpServerOverrides()),
-		hasCwdTools: () => toolsStore.hasEnabledCwdTools,
-		getCwd: () => cwd,
-		getServerHome: () => toolsStore.serverHome ?? null,
-		openModelSelector: () => chatFormActionsRef?.openModelSelector(),
-		getPickersRef: () => pickersRef
-	});
-
-	async function handleWorkingDirectoryChange(newDir: string | null) {
-		// Committing a directory consumes the `/cwd` token; the chip's
-		// clear-X path has no token to consume.
-		const token = findCommandToken(value);
-		if (token && token.name === 'cwd') {
-			value = '';
-			onValueChange?.('');
-		}
-		await conversationsStore.setCwd(newDir);
+	async function handleWorkingDirectoryChange(value: string | null) {
+		await conversationsStore.setCwd(value);
 		if (conversationsStore.activeConversation) {
-			await chatStore.recordCwdChange(newDir?.trim() || null);
+			await chatStore.recordCwdChange(value?.trim() || null);
 		}
 	}

@@ -220,45 +166,23 @@
 	);
 	let canSubmit = $derived(value.trim().length > 0 || hasAttachments);

-	// Caret offset restored after a renderer swap. Callers that mutate
-	// `value` themselves (e.g. the mention picker) pin the target offset
-	// BEFORE the assignment; otherwise the swap effect snapshots the
-	// current caret.
-	let pendingCaretOffset = 0;
-	let caretOffsetPinned = false;
-
-	function queueCaretRestore() {
-		queueMicrotask(() => {
-			inputRef?.focus();
-			inputRef?.setCaretOffset(pendingCaretOffset);
-			caretOffsetPinned = false;
-		});
-	}
-
-	$effect(() => {
-		const wantContenteditable =
-			containsFileMentionLink(value ?? '') || containsCodeSpan(value ?? '');
-		if (useContenteditable === wantContenteditable) return;
-
-		if (!caretOffsetPinned) {
-			pendingCaretOffset = inputRef?.getCaretOffset() ?? (value ?? '').length;
-		}
-
-		useContenteditable = wantContenteditable;
-		queueCaretRestore();
-	});
-
 	onMount(() => {
 		recordingSupported = isAudioRecordingSupported();
 		audioRecorder = new AudioRecorder();
 	});

+	// Defer so the closing popover's focus scope tears down first - bits-ui
+	// yanks a synchronous focus() back into the still-mounted popover.
+	function refocusInput() {
+		queueMicrotask(() => textareaRef?.focus());
+	}
+
 	export function focus() {
-		inputRef?.focus();
+		textareaRef?.focus();
 	}

 	export function resetTextareaHeight() {
-		inputRef?.resetHeight();
+		textareaRef?.resetHeight();
 	}

 	export function openModelSelector() {
@@ -292,10 +216,46 @@
 		}
 	}

+	function handleInput() {
+		const perChatOverrides = conversationsStore.getAllMcpServerOverrides();
+		const hasServers = mcpStore.hasEnabledServers(perChatOverrides);
+
+		if (value.startsWith(PROMPT_TRIGGER_PREFIX) && hasServers) {
+			isPromptPickerOpen = true;
+			promptSearchQuery = value.slice(1);
+			isInlineResourcePickerOpen = false;
+			resourceSearchQuery = '';
+		} else if (
+			value.startsWith(RESOURCE_TRIGGER_PREFIX) &&
+			hasServers &&
+			mcpStore.hasResourcesCapability(perChatOverrides)
+		) {
+			isInlineResourcePickerOpen = true;
+			resourceSearchQuery = value.slice(1);
+			isPromptPickerOpen = false;
+			promptSearchQuery = '';
+		} else {
+			isPromptPickerOpen = false;
+			promptSearchQuery = '';
+			isInlineResourcePickerOpen = false;
+			resourceSearchQuery = '';
+		}
+	}
+
 	function handleKeydown(event: KeyboardEvent) {
-		// Pickers consume navigation/escape keys first; when consumed, skip
-		// the enter-to-submit logic below.
-		if (pickers.handleKeydown(event)) {
+		if (pickersRef?.handleKeydown(event)) {
+			return;
+		}
+
+		if (event.key === KeyboardKey.ESCAPE && isPromptPickerOpen) {
+			isPromptPickerOpen = false;
+			promptSearchQuery = '';
+			return;
+		}
+
+		if (event.key === KeyboardKey.ESCAPE && isInlineResourcePickerOpen) {
+			isInlineResourcePickerOpen = false;
+			resourceSearchQuery = '';
 			return;
 		}

@@ -303,15 +263,6 @@
 			const isModifier = event.ctrlKey || event.metaKey;
 			const sendOnEnter = currentConfig.sendOnEnter !== false;

-			// Caret inside a fenced code block (closed, or still open
-			// while being typed): Enter adds a line, never submits. The
-			// contenteditable consumes this case locally; this gate
-			// covers the plain textarea, where skipping submit lets the
-			// native newline through.
-			if (!isModifier && isOffsetInCodeBlock(value ?? '', inputRef?.getCaretOffset() ?? 0)) {
-				return;
-			}
-
 			if (sendOnEnter || isModifier) {
 				event.preventDefault();

@@ -381,7 +332,7 @@
 				}

 				setTimeout(() => {
-					inputRef?.focus();
+					textareaRef?.focus();
 				}, 10);

 				return;
@@ -408,7 +359,13 @@
 		promptInfo: MCPPromptInfo,
 		args?: Record<string, string>
 	) {
-		pickers.closePromptPicker();
+		// Only clear the value if the prompt was triggered by typing '/'
+		if (value.startsWith(PROMPT_TRIGGER_PREFIX)) {
+			value = '';
+			onValueChange?.('');
+		}
+		isPromptPickerOpen = false;
+		promptSearchQuery = '';

 		const promptName = promptInfo.title || promptInfo.name;
 		const placeholder: ChatUploadedFile = {
@@ -427,7 +384,7 @@

 		uploadedFiles = [...uploadedFiles, placeholder];
 		onUploadedFilesChange?.(uploadedFiles);
-		inputRef?.focus();
+		textareaRef?.focus();
 	}

 	function handlePromptLoadComplete(placeholderId: string, result: GetPromptResult) {
@@ -469,36 +426,39 @@
 		onUploadedFilesChange?.(uploadedFiles);
 	}

-	// Deferred so the closing popover's focus scope tears down first -
-	// bits-ui yanks a synchronous focus() back into the still-mounted popover.
-	function refocusInput() {
-		queueMicrotask(() => inputRef?.focus());
+	function handlePromptPickerClose() {
+		isPromptPickerOpen = false;
+		promptSearchQuery = '';
+		textareaRef?.focus();
 	}

-	// Splice the mention link in place of the `@<query>` token. Uses the
-	// live cursor, not a stale snapshot - the token may have been edited.
-	function handleMentionSelect(entry: FileMentionEntry) {
-		const cursor = inputRef?.getCaretOffset() ?? value.length;
-		const token = findMentionToken(value, cursor);
-		if (!token) return;
+	function handleInlineResourcePickerClose() {
+		isInlineResourcePickerOpen = false;
+		resourceSearchQuery = '';
+		textareaRef?.focus();
+	}

-		const built = buildMentionInsertion(entry, value, token);
-		if (!built) return;
-
-		// Pin the post-insertion caret BEFORE the swap effect runs;
-		// otherwise the effect clobbers it with the textarea's selection
-		// at promotion time (browser-dependent: usually reset to 0).
-		pendingCaretOffset = built.caretOffset;
-		caretOffsetPinned = true;
-
-		value = built.newValue;
-		onValueChange?.(built.newValue);
-
-		// Already in contenteditable mode: no renderer flip, so the swap
-		// effect's caret restore never runs.
-		if (useContenteditable) {
-			queueCaretRestore();
+	function handleInlineResourceSelect() {
+		if (value.startsWith(RESOURCE_TRIGGER_PREFIX)) {
+			value = '';
+			onValueChange?.('');
 		}
+
+		isInlineResourcePickerOpen = false;
+		resourceSearchQuery = '';
+		textareaRef?.focus();
+	}
+
+	function handleBrowseResources() {
+		isInlineResourcePickerOpen = false;
+		resourceSearchQuery = '';
+
+		if (value.startsWith(RESOURCE_TRIGGER_PREFIX)) {
+			value = '';
+			onValueChange?.('');
+		}
+
+		isResourceDialogOpen = true;
 	}

 	async function handleMicClick() {
@@ -543,32 +503,19 @@
 >
 	<ChatFormPickers
 		bind:this={pickersRef}
-		isCommandPickerOpen={pickers.isCommandPickerOpen}
-		commandQuery={pickers.commandQuery}
-		commands={pickers.availableCommands}
-		onCommandPickerClose={pickers.handleCommandPickerClose}
-		onCommandSelect={pickers.handleCommandSelect}
-		isPromptPickerOpen={pickers.isPromptPickerOpen}
-		promptSearchQuery={pickers.promptSearchQuery}
-		isMentionPickerOpen={pickers.isMentionPickerOpen}
-		mentionQuery={pickers.mentionQuery}
-		{mentionAnchor}
-		scopePath={pickers.mentionScopePath}
-		onPromptPickerClose={pickers.handlePromptPickerClose}
-		onMentionPickerClose={pickers.handleMentionPickerClose}
-		onMentionOpened={() => inputRef?.focus()}
-		onMentionSelect={handleMentionSelect}
+		{isPromptPickerOpen}
+		{promptSearchQuery}
+		{isInlineResourcePickerOpen}
+		{resourceSearchQuery}
+		onPromptPickerClose={handlePromptPickerClose}
+		onInlineResourcePickerClose={handleInlineResourcePickerClose}
+		onInlineResourceSelect={handleInlineResourceSelect}
 		onPromptLoadStart={handlePromptLoadStart}
 		onPromptLoadComplete={handlePromptLoadComplete}
 		onPromptLoadError={handlePromptLoadError}
+		onInlineResourceBrowse={handleBrowseResources}
 	/>

-	<div
-		bind:this={mentionAnchor}
-		class="pointer-events-none absolute top-0 right-0 left-0 h-px"
-		aria-hidden="true"
-	></div>
-
 	<div
 		class="{INPUT_CLASSES} overflow-hidden rounded-4xl md:rounded-3xl backdrop-blur-md {disabled
 			? 'cursor-not-allowed opacity-60'
@@ -587,36 +534,20 @@

 		<div
 			class="flex-column relative min-h-12 items-center rounded-4xl md:rounded-3xl py-2 pb-2.25 shadow-sm transition-all focus-within:shadow-md md:py-3!"
+			onpaste={handlePaste}
 		>
-			{#if useContenteditable}
-				<ChatFormContenteditable
-					class="px-5 py-1.5 md:pt-0 mb-0.5"
-					bind:this={inputRef}
-					bind:value
-					onKeydown={handleKeydown}
-					onInput={() => {
-						pickers.handleInput();
-						onValueChange?.(value);
-					}}
-					onPaste={handlePaste}
-					{disabled}
-					{placeholder}
-				/>
-			{:else}
-				<ChatFormTextarea
-					class="px-5 py-1.5 md:pt-0"
-					bind:this={inputRef}
-					bind:value
-					onKeydown={handleKeydown}
-					onInput={() => {
-						pickers.handleInput();
-						onValueChange?.(value);
-					}}
-					onPaste={handlePaste}
-					{disabled}
-					{placeholder}
-				/>
-			{/if}
+			<ChatFormTextarea
+				class="px-5 py-1.5 md:pt-0"
+				bind:this={textareaRef}
+				bind:value
+				onKeydown={handleKeydown}
+				onInput={() => {
+					handleInput();
+					onValueChange?.(value);
+				}}
+				{disabled}
+				{placeholder}
+			/>

 			{#if mcpHasResourceAttachments()}
 				<ChatFormMcpResourcesList
@@ -643,7 +574,7 @@
 				onMicClick={handleMicClick}
 				{onStop}
 				onSystemPromptClick={() => onSystemPromptClick?.({ message: value, files: uploadedFiles })}
-				onMcpPromptClick={showMcpPromptButton ? () => pickers.openPromptPicker() : undefined}
+				onMcpPromptClick={showMcpPromptButton ? () => (isPromptPickerOpen = true) : undefined}
 				onMcpResourcesClick={() => (isResourceDialogOpen = true)}
 			/>
 		</div>
@@ -651,15 +582,11 @@

 	<ContextGaugePopup />

-	{#if toolsStore.hasEnabledCwdTools}
+	{#if toolsStore.builtinTools.length > 0}
 		<ChatFormWorkingDirectory
 			directory={cwd}
-			isOpen={pickers.isWorkingDirectoryPickerOpen}
-			bind:query={pickers.workingDirectoryQuery}
-			customAnchor={mentionAnchor}
 			onChange={handleWorkingDirectoryChange}
-			onClose={pickers.handleWorkingDirectoryClose}
-			onOpen={pickers.handleWorkingDirectoryOpen}
+			onClose={refocusInput}
 			{disabled}
 		/>
 	{/if}
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Georgi Gerganov	6751477adf	cont : avoid copy_state	2026-08-06 10:19:52 +03:00
Gaurav Garg	9fcc9117cf	Address review comments	2026-08-06 00:27:46 +05:30
Gaurav Garg	8580cc0b81	Rebase and address review comments	2026-08-06 00:27:46 +05:30
Gaurav Garg	fea0c3d410	Rebase changes	2026-08-06 00:27:46 +05:30
Gaurav Garg	483169a47e	More test fixes	2026-08-06 00:27:46 +05:30
Gaurav Garg	d9fe17bdad	Fix tests on Vulkan	2026-08-06 00:27:46 +05:30
Gaurav Garg	9c22762603	Simpify some of the changes	2026-08-06 00:27:46 +05:30
Gaurav Garg	3e753aa27d	Fix CPU and backend sampling mismatches	2026-08-06 00:27:46 +05:30
Gaurav Garg	a24097dc57	Match dist between CPU and GPU	2026-08-06 00:27:46 +05:30
Gaurav Garg	813c39b225	Don't reuse memory for output views.	2026-08-06 00:27:46 +05:30
Gaurav Garg	f1f008b2ee	More fixes	2026-08-06 00:27:46 +05:30
Gaurav Garg	eab4e8444f	Add a numeric context parameter declaring the maximum outputs one sequence	2026-08-06 00:27:46 +05:30
Gaurav Garg	0eddfa0b4c	Clamp the mask sum before converting it into the sampled index	2026-08-06 00:27:46 +05:30
Gaurav Garg	167d5e141b	Enable backend sampling with token speculation	2026-08-06 00:27:46 +05:30