mtmd: fix gemma 4 ffn act (#30882)

Co-authored-by: TANGBUDU <TANGBUDU@users.noreply.github.com>

.github/workflows/build-and-test-snapdragon.yml

@@ -35,7 +35,7 @@ jobs:
   android-ndk-snapdragon:
     runs-on: ubuntu-24.04  # previously ubuntu-latest
     container:
-      image: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.7'
+      image: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.8'
     defaults:
       run:
         shell: bash
@@ -69,7 +69,7 @@ jobs:
   linux-iot-snapdragon:
     runs-on: ubuntu-24.04  # previously ubuntu-latest
     container:
-      image: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.7'
+      image: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.8'
     defaults:
       run:
         shell: bash

.github/workflows/build-webgpu.yml

@@ -7,15 +7,7 @@ on:
       - master
     paths: [
       '.github/workflows/build-webgpu.yml',
-      '**/CMakeLists.txt',
-      '**/.cmake',
-      '**/*.h',
-      '**/*.hpp',
-      '**/*.c',
-      '**/*.cpp',
-      '**/*.wgsl',
-      '**/*.tmpl',
-      'ggml/src/ggml-webgpu/wgsl-shaders/embed_wgsl.py'
+      'ggml/src/ggml-webgpu/**'
     ]
 
   pull_request:
@@ -63,135 +55,3 @@ jobs:
             -type f \( -name '*.cpp' -o -name '*.hpp' -o -name '*.h' \) \
             -print0 |
             xargs -0 clang-format-22 --dry-run --Werror
-
-  macos:
-    runs-on: macos-latest
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.24
-        with:
-          restore: false
-          save: false
-
-      - name: Dawn Dependency
-        id: dawn-depends
-        run: |
-          DAWN_VERSION="v20260908.214631"
-          DAWN_OWNER="google"
-          DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-macos-latest-Release"
-          echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
-          curl -L -o artifact.tar.gz \
-            "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
-          mkdir dawn
-          tar -xvf artifact.tar.gz -C dawn --strip-components=1
-
-      - name: ccache-buckets-restore
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
-        with:
-          key: webgpu-macos-latest
-          folder: llama.cpp
-          hf_bucket: ggml-org/cache
-
-      - name: Build
-        id: cmake_build
-        run: |
-          export CMAKE_PREFIX_PATH=dawn
-          cmake -B build -G "Ninja" -DCMAKE_BUILD_TYPE=Release -DGGML_WEBGPU=ON -DGGML_METAL=OFF -DGGML_BLAS=OFF
-          time cmake --build build --config Release -j $(sysctl -n hw.logicalcpu)
-
-      - name: ccache-buckets-save
-        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
-        with:
-          key: webgpu-macos-latest
-          folder: llama.cpp
-          evict-old-files: 1d
-          hf_bucket: ggml-org/cache
-          save: true
-
-      - name: Test
-        id: cmake_test
-        run: |
-          cd build
-          ctest -L main --verbose --timeout 900
-
-  ubuntu:
-    runs-on: ubuntu-24.04
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.24
-        with:
-          restore: false
-          save: false
-
-      - name: Dependencies
-        id: depends
-        run: |
-          sudo add-apt-repository -y ppa:kisak/kisak-mesa
-          sudo apt-get update -y
-          sudo apt-get install -y build-essential mesa-vulkan-drivers \
-            libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libssl-dev
-
-      - name: Dawn Dependency
-        id: dawn-depends
-        run: |
-          sudo apt-get install -y libxrandr-dev libxinerama-dev libxcursor-dev mesa-common-dev libx11-xcb-dev libxi-dev
-          DAWN_VERSION="v20260908.214631"
-          DAWN_OWNER="google"
-          DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-ubuntu-latest-Release"
-          echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
-          curl -L -o artifact.tar.gz \
-            "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
-          mkdir dawn
-          tar -xvf artifact.tar.gz -C dawn --strip-components=1
-
-      - name: ccache-buckets-restore
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
-        with:
-          key: webgpu-ubuntu-24.04
-          folder: llama.cpp
-          hf_bucket: ggml-org/cache
-
-      - name: Build
-        id: cmake_build
-        run: |
-          export Dawn_DIR=dawn/lib64/cmake/Dawn
-          cmake -B build \
-            -DGGML_WEBGPU=ON
-          time cmake --build build --config Release -j $(nproc)
-
-      - name: ccache-buckets-save
-        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
-        with:
-          key: webgpu-ubuntu-24.04
-          folder: llama.cpp
-          evict-old-files: 1d
-          hf_bucket: ggml-org/cache
-          save: true
-
-      - name: Test
-        id: cmake_test
-        run: |
-          cd build
-          ctest -L main --verbose --timeout 900

.github/workflows/release.yml

@@ -547,7 +547,7 @@ jobs:
     if: ${{ needs.check-release.outputs.should_release == 'true' }}
 
     runs-on: ubuntu-latest
-    container: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.7'
+    container: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.8'
 
     defaults:
       run:
@@ -600,7 +600,7 @@ jobs:
     if: ${{ needs.check-release.outputs.should_release == 'true' }}
 
     runs-on: ubuntu-latest
-    container: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.7'
+    container: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.8'
 
     defaults:
       run:

common/speculative.cpp

@@ -16,6 +16,7 @@
 #include <cassert>
 #include <cmath>
 #include <cstring>
+#include <fstream>
 #include <iomanip>
 #include <map>
 #include <cinttypes>
@@ -2152,9 +2153,6 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
 
     uint16_t n_draft;
 
-    bool save_dynamic;
-    bool save_static;
-
     struct seq_info {
         size_t cache_size = 0; // number of tokens in n-gram cache
 
@@ -2168,54 +2166,72 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
     common_speculative_impl_ngram_cache(
             const common_params_speculative & params,
             uint32_t n_seq,
-            uint16_t n_draft,
-            const std::string & path_static,
-            const std::string & path_dynamic,
-            bool save_dynamic,
-            bool save_static)
+            uint16_t n_draft)
         : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq, n_draft)
         , params(params.ngram_cache)
         , n_draft(n_draft)
-        , save_dynamic(save_dynamic)
-        , save_static(save_static)
     {
+        const auto & ngram = params.ngram_cache;
+
         SPC_TRC("%s", "adding speculative implementation 'ngram-cache'\n");
         SPC_TRC("- n_draft=%d, cache_static=%s, cache_dynamic=%s\n",
                 n_draft,
-                path_static.empty() ? "none" : path_static.c_str(),
-                path_dynamic.empty() ? "none" : path_dynamic.c_str());
+                ngram.lookup_cache_static.empty() ? "none" : ngram.lookup_cache_static.c_str(),
+                ngram.lookup_cache_dynamic.empty() ? "none" : ngram.lookup_cache_dynamic.c_str());
 
         sinfos.resize(n_seq);
 
-        if (!path_static.empty()) {
+        if (!ngram.lookup_cache_static.empty()) {
             try {
-                auto ngram_cache_static = common_ngram_cache_load(path_static);
+                auto ngram_cache_static = common_ngram_cache_load(ngram.lookup_cache_static);
 
                 for (auto & sinfo : sinfos) {
                     sinfo.ngram_cache_static = ngram_cache_static;
                 }
-            } catch (...) {
-                SPC_ERR("failed to open static lookup cache: %s", path_static.c_str());
+            } catch (const std::ifstream::failure &) {
+                SPC_ERR("failed to open static lookup cache: %s\n", ngram.lookup_cache_static.c_str());
                 GGML_ABORT("Couldn't read static lookup cache");
             }
         }
 
-        if (!path_dynamic.empty()) {
+        if (!ngram.lookup_cache_dynamic.empty()) {
             try {
-                auto ngram_cache_dynamic = common_ngram_cache_load(path_dynamic);
+                auto ngram_cache_dynamic = common_ngram_cache_load(ngram.lookup_cache_dynamic);
 
                 for (auto & sinfo : sinfos) {
                     sinfo.ngram_cache_dynamic = ngram_cache_dynamic;
                 }
-            } catch (...) {
-                SPC_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str());
-                GGML_ABORT("Couldn't read dynamic lookup cache");
+            } catch (const std::ifstream::failure &) {
+                std::error_code ec;
+                const bool exists = std::filesystem::exists(ngram.lookup_cache_dynamic, ec);
+                if (!exists && !ec) {
+                    // missing file is created on destruction
+                    SPC_WRN("dynamic lookup cache file does not exist: %s, starting with empty cache\n",
+                            ngram.lookup_cache_dynamic.c_str());
+                } else {
+                    SPC_ERR("failed to open dynamic lookup cache: %s\n", ngram.lookup_cache_dynamic.c_str());
+                    GGML_ABORT("Couldn't read dynamic lookup cache");
+                }
             }
         }
     }
 
-    void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override {
-        // noop
+    ~common_speculative_impl_ngram_cache() override {
+        if (params.lookup_cache_dynamic.empty() || sinfos.empty()) {
+            return;
+        }
+
+        auto & dynamic_cache = sinfos.front().ngram_cache_dynamic;
+        for (auto & sinfo : sinfos) {
+            common_ngram_cache_merge(dynamic_cache, sinfo.ngram_cache_context);
+        }
+
+        SPC_INF("saving dynamic lookup cache at %s\n", params.lookup_cache_dynamic.c_str());
+        common_ngram_cache_save(dynamic_cache, params.lookup_cache_dynamic);
+    }
+
+    void begin(llama_seq_id seq_id, const llama_tokens & /*prompt*/) override {
+        sinfos[seq_id].cache_size = 0;
     }
 
     void draft_one(
@@ -2226,22 +2242,6 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
 
         const auto & prompt = *dparams.prompt;
 
-        if (sinfo.cache_size < prompt.size() + 1) {
-            llama_tokens tokens_new;
-            tokens_new.reserve(prompt.size() + 1 - sinfo.cache_size);
-            for (size_t j = sinfo.cache_size; j < prompt.size(); ++j) {
-                tokens_new.push_back(prompt[j]);
-            }
-            tokens_new.push_back(dparams.id_last); // add the last token
-
-            // Update context ngram cache with new dparams.prompt:
-            common_ngram_cache_update(
-                    sinfo.ngram_cache_context,
-                    LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX,
-                    tokens_new, tokens_new.size(), false);
-            sinfo.cache_size = prompt.size() + 1;
-        }
-
         llama_tokens inp;
         inp.reserve(prompt.size() + 1);
         for (size_t j = 0; j < prompt.size(); ++j) {
@@ -2249,6 +2249,16 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
         }
         inp.push_back(dparams.id_last);
 
+        if (sinfo.cache_size < inp.size()) {
+            // Update context ngram cache with the full history so n-grams
+            // ending at the new tokens have their preceding context
+            common_ngram_cache_update(
+                    sinfo.ngram_cache_context,
+                    LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX,
+                    inp, inp.size() - sinfo.cache_size, false);
+            sinfo.cache_size = inp.size();
+        }
+
         result.push_back(dparams.id_last);
 
         common_ngram_cache_draft(
@@ -2309,20 +2319,12 @@ static common_ngram_map get_common_ngram_map(
     return common_ngram_map(size_key, size_value, key_only, min_hits);
 }
 
-static common_speculative_impl_ngram_cache create_state_ngram_cache(
+static std::unique_ptr<common_speculative_impl_ngram_cache> create_state_ngram_cache(
         const common_speculative_config & config,
-        uint32_t n_seq,
-        const std::string & path_static,
-        const std::string & path_dynamic) {
+        uint32_t n_seq) {
     uint16_t n_draft = 8; // TODO get from config?
 
-    // TODO bool param in common/common.h to set save_static/save_dynamic?
-    bool save_static = false;
-    bool save_dynamic = false;
-
-    common_speculative_impl_ngram_cache state(config.params, n_seq, n_draft, path_static, path_dynamic, save_static, save_dynamic);
-
-    return state;
+    return std::make_unique<common_speculative_impl_ngram_cache>(config.params, n_seq, n_draft);
 }
 
 std::string common_speculative_type_name_str(const std::vector<common_speculative_type> & types) {
@@ -2822,11 +2824,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
                 break;
             }
             case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: {
-                auto state = create_state_ngram_cache(
-                        config, n_seq,
-                        params.ngram_cache.lookup_cache_static,
-                        params.ngram_cache.lookup_cache_dynamic);
-                impls.push_back(std::make_unique<common_speculative_impl_ngram_cache>(state));
+                impls.push_back(create_state_ngram_cache(config, n_seq));
                 break;
             }
             default:

conversion/gemma.py

@@ -999,6 +999,8 @@ class Gemma4VisionAudioModel(MmprojModel):
         assert self.hparams_vision is not None
         self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.GEMMA4V)
         self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision.get("layer_norm_eps", 1e-6))
+        assert self.hparams_vision.get("hidden_activation", "gelu_pytorch_tanh") == "gelu_pytorch_tanh"
+        self.gguf_writer.add_vision_use_gelu(True)
 
         # audio params
         if self.has_audio_encoder:

conversion/lfm2.py

@@ -374,6 +374,16 @@ class LFM2VLModel(MmprojModel):
             if resize_algo is None:
                 raise ValueError(f"unsupported resample: {resample}")
             self.gguf_writer.add_vision_image_resize_algo(resize_algo)
+        # tiling params from processor_config.json (do_image_splitting=False -> single tile, i.e. no tiling)
+        preproc = self.preprocessor_config
+        if preproc.get("do_image_splitting", True):
+            self.gguf_writer.add_vision_preproc_min_tiles(preproc.get("min_tiles", 2))
+            self.gguf_writer.add_vision_preproc_max_tiles(preproc.get("max_tiles", 10))
+        else:
+            self.gguf_writer.add_vision_preproc_min_tiles(1)
+            self.gguf_writer.add_vision_preproc_max_tiles(1)
+        self.gguf_writer.add_vision_preproc_tile_size(preproc.get("tile_size", 512))
+        self.gguf_writer.add_vision_preproc_max_pixels_tolerance(preproc.get("max_pixels_tolerance", 2.0))
 
     @classmethod
     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:

docs/backend/snapdragon/README.md

@@ -6,8 +6,8 @@ The cross-compilation toolchain images are provided by the
 [Qualcomm Snapdragon Toolchain registry](https://github.com/snapdragon-toolchain).
 These Docker images include the Android NDK, OpenCL SDK, Hexagon SDK, CMake, and the necessary cross-compilers:
 
-* **Android toolchain**: `ghcr.io/snapdragon-toolchain/arm64-android:v0.7`
-* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7`
+* **Android toolchain**: `ghcr.io/snapdragon-toolchain/arm64-android:v0.8`
+* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.8`
 
 The unified build utility (`scripts/snapdragon/build.py`) automatically pulls
 and orchestrates these containers to perform target compilation.
@@ -40,7 +40,7 @@ Alternatively, you can build llama.cpp manually by entering the cross-compilatio
 
 ```bash
 # Start the cross-compilation container manually:
-~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-android:v0.7
+~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-android:v0.8
 
 # Inside the container, build the project using presets:
 [d]/workspace> cp docs/backend/snapdragon/CMakeUserPresets.json .
@@ -53,7 +53,7 @@ Preset CMake variables:
   GGML_HEXAGON="ON"
   GGML_OPENCL="ON"
   GGML_OPENMP="OFF"
-  HEXAGON_SDK_ROOT="/opt/hexagon/6.6.0.0"
+  HEXAGON_SDK_ROOT="/opt/hexagon/6.7.0.0"
 ...
 -- Including OpenCL backend
 -- Including Hexagon backend

docs/backend/snapdragon/linux.md

@@ -3,7 +3,7 @@
 The cross-compilation is performed using the Snapdragon Linux Docker toolchain image (see
 [github.com/snapdragon-toolchain](https://github.com/snapdragon-toolchain)):
 
-* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7`
+* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.8`
 
 The unified build utility (`scripts/snapdragon/build.py`) automatically pulls
 and orchestrates this container to perform target compilation. You only need to
@@ -29,7 +29,7 @@ Alternatively, you can build llama.cpp manually by entering the cross-compilatio
 
 ```bash
 # Start the cross-compilation container manually:
-~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-linux:v0.7
+~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-linux:v0.8
 
 # Inside the container, build the project using presets:
 [d]/workspace> cp docs/backend/snapdragon/CMakeUserPresets.json .

docs/backend/snapdragon/windows.md

@@ -8,7 +8,7 @@ Native Windows 11 arm64 builds have the following tool dependencies:
   - UCRT and Driver Kit
 - LLVM core libraries and Clang compiler (winget)
 - CMake, Git, Python (winget)
-- Hexagon SDK Community Edition 6.6 or later (see below)
+- Hexagon SDK Community Edition 6.7 or later (see below)
 - OpenCL SDK 2.3 or later (see below)
 
 Note: The rest of the **Windows** build process assumes that you're running natively in Powershell.
@@ -54,15 +54,15 @@ c:\Qualcomm\OpenCL_SDK\2.3.2
 
 Either use the trimmed down version (optimized for CI) from
 
-    https://github.com/snapdragon-toolchain/hexagon-sdk/releases/download/v6.6.0.0/hexagon-sdk-v6.6.0.0-arm64-wos.tar.xz
+    https://github.com/snapdragon-toolchain/hexagon-sdk/releases/download/v6.7.0.0/hexagon-sdk-v6.7.0.0-arm64-wos.tar.xz
 
 Or download the complete official version from
 
-    https://softwarecenter.qualcomm.com/catalog/item/Hexagon_SDK?version=6.6.0.0
+    https://softwarecenter.qualcomm.com/catalog/item/Hexagon_SDK?version=6.7.0.0
 
 Unzip/untar the archive into
 ```
-c:\Qualcomm\Hexagon_SDK\6.6.0.0
+c:\Qualcomm\Hexagon_SDK\6.7.0.0
 ```
 
 ## Install the latest Adreno GPU driver
@@ -150,8 +150,8 @@ The overall Hexagon backend build procedure for Windows on Snapdragon is the sam
 However, additional settings are required for generating and signing HTP Ops libraries.
 ```
 > $env:OPENCL_SDK_ROOT="C:\Qualcomm\OpenCL_SDK\2.3.2"
-> $env:HEXAGON_SDK_ROOT="C:\Qualcomm\Hexagon_SDK\6.6.0.0"
-> $env:HEXAGON_TOOLS_ROOT="C:\Qualcomm\Hexagon_SDK\6.6.0.0\tools\HEXAGON_Tools\19.0.07"
+> $env:HEXAGON_SDK_ROOT="C:\Qualcomm\Hexagon_SDK\6.7.0.0"
+> $env:HEXAGON_TOOLS_ROOT="C:\Qualcomm\Hexagon_SDK\6.7.0.0\tools\HEXAGON_Tools\21.0.04"
 > $env:HEXAGON_HTP_CERT="c:\Users\MyUsers\Certs\ggml-htp-v1.pfx"
 > $env:WINDOWS_SDK_BIN="C:\Program Files (x86)\Windows Kits\10\bin\10.0.26100.0"
 

docs/backend/zDNN.md

@@ -2,6 +2,7 @@
 
 > [!WARNING]
 > **Note:** zDNN is **not** the same as ZenDNN.
+>
 > - **zDNN** (this page): IBM's Deep Neural Network acceleration library for IBM Z & LinuxONE Mainframes
 > - **ZenDNN**: AMD's deep learning library for AMD EPYC CPUs ([see ZenDNN documentation](ZenDNN.md))
 
@@ -46,6 +47,14 @@ The IBM zDNN backend has the following CMake options that control the behaviour
 | --------------------- | ------------- | -------------------------------------------------------------------------------------- |
 | `GGML_ZDNN_HIGH_PREC` | `0`           | Set to `1` to requantize Q8_0 weights in smaller groups, for lower loss at lower speed |
 
+## Performance Tunables
+
+| Flag Option             | Default Value | Description                                            |
+| ----------------------- | ------------- | ------------------------------------------------------ |
+| `--flash-attn` or `-fa` | `auto`        | Set to `on` to run attention on the NNPA (Recommended) |
+
+Note: Set `-fa on` explicitly. Flash Attention fuses each layer's attention (the KQ matrix multiplication, softmax and KQV matrix multiplication) into one `FLASH_ATTN_EXT` operation, which zDNN runs on the NNPA instead of the CPU. The speedup grows with the context length. With `auto`, llama.cpp currently turns Flash Attention off, because the operation runs on a different device than the layer. zDNN runs it when the KV cache is F16 (the default), the model uses no attention sinks, ALiBi or logit softcapping, and there is one KV cache stream (one sequence, or `--kv-unified`). Otherwise it falls back to the CPU.
+
 ## 1. Install zDNN Library
 
 Note: A minimum of libzdnn 1.1.2 is required for the IBM zDNN backend to link correctly.

ggml/src/ggml-backend.cpp

@@ -2226,6 +2226,7 @@ bool ggml_op_alloc_size_may_expand(enum ggml_op op) {
         case GGML_OP_CUMSUM:
         case GGML_OP_ARGSORT:
         case GGML_OP_TOP_K:
+        case GGML_OP_ARGMAX:
             return true;
         default:
             return false;

ggml/src/ggml-cpu/simd-mappings.h

@@ -42,6 +42,7 @@ extern "C" {
     #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) neon_compute_fp32_to_fp16(x)
 
     #define GGML_CPU_FP16_TO_FP32(x) GGML_CPU_COMPUTE_FP16_TO_FP32(x)
+    #define GGML_CPU_FP32_TO_FP16(x) GGML_CPU_COMPUTE_FP32_TO_FP16(x)
 
     static inline float neon_compute_fp16_to_fp32(ggml_fp16_t h) {
         __fp16 tmp;
@@ -150,9 +151,15 @@ extern float ggml_table_f32_ue4m3[1 << 8];
 #define GGML_CPU_FP32_TO_FP16(x) GGML_CPU_COMPUTE_FP32_TO_FP16(x)
 #endif
 
-// On ARM NEON, it's quicker to directly convert x -> x instead of calling into ggml_lookup_fp16_to_fp32,
-// so we define GGML_CPU_FP16_TO_FP32 and GGML_CPU_FP32_TO_FP16 elsewhere for NEON.
-// This is also true for POWER9.
+// a branch that defines a GGML_CPU_COMPUTE_* macro must also select it, or the fallback takes over silently
+#if defined(GGML_CPU_COMPUTE_FP16_TO_FP32) && !defined(GGML_CPU_FP16_TO_FP32) && !defined(__F16C__) // see: https://github.com/ggml-org/llama.cpp/issues/30342
+#error "a branch defined GGML_CPU_COMPUTE_FP16_TO_FP32 without selecting it as GGML_CPU_FP16_TO_FP32"
+#endif
+#if defined(GGML_CPU_COMPUTE_FP32_TO_FP16) && !defined(GGML_CPU_FP32_TO_FP16)
+#error "a branch defined GGML_CPU_COMPUTE_FP32_TO_FP16 without selecting it as GGML_CPU_FP32_TO_FP16"
+#endif
+
+// fall back to the lookup table / the portable software conversion for whichever direction no branch selected
 #if !defined(GGML_CPU_FP16_TO_FP32)
 inline static float ggml_lookup_fp16_to_fp32(ggml_fp16_t f) {
     uint16_t s;

ggml/src/ggml-cuda/common.cuh

@@ -1551,11 +1551,11 @@ struct ggml_backend_cuda_context {
     // pool
     std::unique_ptr<ggml_cuda_pool> pools[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS];
 
-    static std::unique_ptr<ggml_cuda_pool> new_pool_for_device(int device, int stream_no);
+    static std::unique_ptr<ggml_cuda_pool> new_pool_for_device(ggml_backend_cuda_context & ctx, int device, int stream_no);
 
     ggml_cuda_pool & pool(int device) {
         if (pools[device][curr_stream_no] == nullptr) {
-            pools[device][curr_stream_no] = new_pool_for_device(device, curr_stream_no);
+            pools[device][curr_stream_no] = new_pool_for_device(*this, device, curr_stream_no);
         }
         return *pools[device][curr_stream_no];
     }

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

@@ -457,6 +457,7 @@ const ggml_cuda_device_info & ggml_cuda_info() {
 struct ggml_cuda_pool_leg : public ggml_cuda_pool {
     static const int MAX_BUFFERS = 256;
 
+    ggml_backend_cuda_context & ctx;
     int device;
     struct ggml_cuda_buffer {
         void * ptr = nullptr;
@@ -466,7 +467,8 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
     ggml_cuda_buffer buffer_pool[MAX_BUFFERS] = {};
     size_t pool_size = 0;
 
-    explicit ggml_cuda_pool_leg(int device) :
+    ggml_cuda_pool_leg(ggml_backend_cuda_context & ctx, int device) :
+        ctx(ctx),
         device(device) {
     }
 
@@ -488,6 +490,20 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
         }
     }
 
+    // captured CUDA graphs can still use freed buffers, so they must be captured again
+    void invalidate_cuda_graphs() {
+#ifdef USE_CUDA_GRAPH
+        int n_captured = 0;
+        for (const auto & [key, graph] : ctx.cuda_graphs) {
+            n_captured += graph->instance != nullptr;
+        }
+        if (n_captured > 0) {
+            GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: memory freed, %d captured CUDA graphs must be captured again\n", device, n_captured);
+        }
+        ctx.cuda_graphs.clear();
+#endif // USE_CUDA_GRAPH
+    }
+
     void * alloc(size_t size, size_t * actual_size) override {
 #ifdef DEBUG_CUDA_MALLOC
         int nnz = 0;
@@ -538,6 +554,7 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
                            device, look_ahead_size/1024.0/1024.0, cached_bytes/1024.0/1024.0);
             CUDA_CHECK(cudaDeviceSynchronize());
             clear_pool();
+            invalidate_cuda_graphs();
             err = ggml_cuda_device_malloc(&ptr, look_ahead_size, device);
             if (err == cudaSuccess) {
                 GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: retry succeeded\n", device);
@@ -566,6 +583,7 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
         ggml_cuda_set_device(device);
         CUDA_CHECK(cudaFree(ptr));
         pool_size -= size;
+        invalidate_cuda_graphs();
     }
 };
 
@@ -720,14 +738,15 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
 };
 #endif // defined(GGML_USE_VMM)
 
-std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(int                  device,
-                                                                               [[maybe_unused]] int stream_no) {
+std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(ggml_backend_cuda_context & ctx,
+                                                                               int                         device,
+                                                                               [[maybe_unused]] int        stream_no) {
 #if defined(GGML_USE_VMM)
     if (ggml_cuda_info().devices[device].vmm) {
         return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_vmm(device));
     }
 #endif // defined(GGML_USE_VMM)
-    return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_leg(device));
+    return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_leg(ctx, device));
 }
 
 // destroying a cuBLAS handle while a graph is being captured in a different thread can result in a CUDA error
@@ -2014,6 +2033,8 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c
     const ggml_tensor * src0 = dst->src[0];
     const ggml_tensor * src1 = dst->src[1];
 
+    const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
+
     if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
         return true;
     }
@@ -2032,7 +2053,7 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c
         return false;
     }
 
-    if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
+    if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
         return false;
     }
 
@@ -2050,6 +2071,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
     GGML_TENSOR_BINARY_OP_LOCALS
 
     const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+    const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
 
     // [TAG_MUL_MAT_ID_CUDA_GRAPHS]
     if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
@@ -2074,7 +2096,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
             return;
         }
 
-        if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
+        if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
             ggml_cuda_mul_mat_f(ctx, src0, src1, ids, dst);
             return;
         }

ggml/src/ggml-cuda/ssm-conv.cu

@@ -151,9 +151,10 @@ static void ssm_conv_f32_cuda(const float * src0, const float * src1, const floa
         case 3:  launch_kernel(std::integral_constant<int, 3 >{}); break;
         case 4:  launch_kernel(std::integral_constant<int, 4 >{}); break;
         case 5:  launch_kernel(std::integral_constant<int, 5 >{}); break;
+        case 7:  launch_kernel(std::integral_constant<int, 7 >{}); break;
         case 9:  launch_kernel(std::integral_constant<int, 9 >{}); break;
         case 15: launch_kernel(std::integral_constant<int, 15>{}); break;
-        default: GGML_ABORT("Only support kernel sizes 3, 4, 5, 9, 15 right now.");
+        default: GGML_ABORT("Only support kernel sizes 3, 4, 5, 7, 9, 15 right now.");
     }
 }
 

ggml/src/ggml-hexagon/CMakeLists.txt

@@ -21,6 +21,7 @@ message(STATUS "hexagon: using ${HEXAGON_SDK_ROOT} and ${HEXAGON_TOOLS_ROOT} for
 
 include(${HEXAGON_SDK_ROOT}/build/cmake/hexagon_fun.cmake)
 include(ExternalProject)
+find_package(Python3 REQUIRED COMPONENTS Interpreter)
 
 option(GGML_HEXAGON_HTP_DEBUG  "ggml-hexagon: enable HTP debug output" OFF)
 set(GGML_HEXAGON_HTP_BUILD_TYPE "Release" CACHE STRING "ggml-hexagon: HTP skel build type (Release, RelWithDebInfo, Debug)")
@@ -74,7 +75,8 @@ function(build_htp_skel V)
             -DHEXAGON_TOOLS_ROOT=${HEXAGON_TOOLS_ROOT}
             -DHEXAGON_HTP_DEBUG=${GGML_HEXAGON_HTP_DEBUG}
             -DDSP_VERSION=${V}
-            -DPREBUILT_LIB_DIR="toolv19_${V}")
+            -DPython3_EXECUTABLE=${Python3_EXECUTABLE}
+            -DPREBUILT_LIB_DIR="toolv21_${V}")
     list(APPEND HTP_SKELS ${CMAKE_CURRENT_BINARY_DIR}/libggml-htp-${V}.so)
     list(APPEND HTP_PROJECTS htp-${V})
     set(HTP_SKELS ${HTP_SKELS} PARENT_SCOPE)
@@ -85,6 +87,7 @@ build_htp_skel(v73)
 build_htp_skel(v75)
 build_htp_skel(v79)
 build_htp_skel(v81)
+build_htp_skel(v85)
 
 # Install Hexagon skels required at runtime
 install(FILES ${HTP_SKELS} TYPE LIB)

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

@@ -8399,9 +8399,9 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
             if (opt_arch < 73) {
                 GGML_LOG_WARN("ggml-hex: Hexagon arch v%d is under supported range, capping at v73\n", opt_arch);
                 opt_arch = 73;
-            } else if (opt_arch > 81) {
-                GGML_LOG_WARN("ggml-hex: Hexagon arch v%d is over supported range, capping at v81\n", opt_arch);
-                opt_arch = 81;
+            } else if (opt_arch > 85) {
+                GGML_LOG_WARN("ggml-hex: Hexagon arch v%d is over supported range, capping at v85\n", opt_arch);
+                opt_arch = 85;
             }
         }
     } else {

ggml/src/ggml-hexagon/htp/CMakeLists.txt

@@ -53,9 +53,17 @@ target_compile_definitions(${HTP_LIB} PRIVATE
     $<IF:$<BOOL:${HEXAGON_HTP_DEBUG}>,HTP_DEBUG=1,NDEBUG=1>
     $<IF:$<BOOL:${HEXAGON_HTP_DEBUG}>,FARF_HIGH=1,>)
 
-
 build_idl(htp_iface.idl ${HTP_LIB})
 
 set_target_properties(${HTP_LIB} PROPERTIES EXPORT_COMPILE_COMMANDS ON)
 
+if (NOT Python3_EXECUTABLE)
+    find_package(Python3 REQUIRED COMPONENTS Interpreter)
+endif()
+
+add_custom_command(TARGET ${HTP_LIB} POST_BUILD
+    COMMAND ${CMAKE_COMMAND} -E echo "Adding SWIV section $<TARGET_FILE_NAME:${HTP_LIB}>"
+    COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/htp-swiv.py -i $<TARGET_FILE:${HTP_LIB}> -o $<TARGET_FILE:${HTP_LIB}>
+    VERBATIM)
+
 install(TARGETS ${HTP_LIB})

ggml/src/ggml-hexagon/htp/htp-swiv.py

@@ -0,0 +1,512 @@
+#!/usr/bin/env python3
+
+import argparse
+import logging
+import os
+import shutil
+import struct
+import sys
+import tempfile
+import zlib
+
+ELF_MAGIC = b'\x7fELF'
+ELF_CLASS_INDEX = 4
+ELFCLASS32 = 1
+ELFCLASS64 = 2
+ELF_VERSION_INDEX = 6
+EV_CURRENT = 1
+MAX_PHDR_COUNT = 100
+
+LOAD_TYPE = 0x1
+
+SHT_NOBITS = 8
+
+SWIV_MAGIC = b'SWIV'
+SWIV_SEG_SIZE_MIN = 16
+# SHT_LOUSER plus ASCII "SWIV".
+SWIV_SECTION_TYPE = 0xD3574956
+
+logger = logging.getLogger("htp-swiv")
+
+
+class Elf_Ehdr_common:
+    s = struct.Struct('16sHHI')
+
+    def __init__(self, data):
+        unpacked_data = (Elf_Ehdr_common.s).unpack(data)
+        self.e_ident = unpacked_data[0]
+        self.e_type = unpacked_data[1]
+        self.e_machine = unpacked_data[2]
+        self.e_version = unpacked_data[3]
+
+
+class Elf32_Ehdr:
+    s = struct.Struct('16sHHIIIIIHHHHHH')
+
+    def __init__(self, data):
+        unpacked_data = (Elf32_Ehdr.s).unpack(data)
+        self.e_ident = unpacked_data[0]
+        self.e_type = unpacked_data[1]
+        self.e_machine = unpacked_data[2]
+        self.e_version = unpacked_data[3]
+        self.e_entry = unpacked_data[4]
+        self.e_phoff = unpacked_data[5]
+        self.e_shoff = unpacked_data[6]
+        self.e_flags = unpacked_data[7]
+        self.e_ehsize = unpacked_data[8]
+        self.e_phentsize = unpacked_data[9]
+        self.e_phnum = unpacked_data[10]
+        self.e_shentsize = unpacked_data[11]
+        self.e_shnum = unpacked_data[12]
+        self.e_shstrndx = unpacked_data[13]
+
+    def get_packed_data(self):
+        values = [self.e_ident,
+                  self.e_type,
+                  self.e_machine,
+                  self.e_version,
+                  self.e_entry,
+                  self.e_phoff,
+                  self.e_shoff,
+                  self.e_flags,
+                  self.e_ehsize,
+                  self.e_phentsize,
+                  self.e_phnum,
+                  self.e_shentsize,
+                  self.e_shnum,
+                  self.e_shstrndx
+                  ]
+
+        return (Elf32_Ehdr.s).pack(*values)
+
+
+class Elf32_Phdr:
+    s = struct.Struct('I' * 8)
+
+    def __init__(self, data):
+        unpacked_data = (Elf32_Phdr.s).unpack(data)
+        self.p_type = unpacked_data[0]
+        self.p_offset = unpacked_data[1]
+        self.p_vaddr = unpacked_data[2]
+        self.p_paddr = unpacked_data[3]
+        self.p_filesz = unpacked_data[4]
+        self.p_memsz = unpacked_data[5]
+        self.p_flags = unpacked_data[6]
+        self.p_align = unpacked_data[7]
+
+    def get_packed_data(self):
+        values = [self.p_type,
+                  self.p_offset,
+                  self.p_vaddr,
+                  self.p_paddr,
+                  self.p_filesz,
+                  self.p_memsz,
+                  self.p_flags,
+                  self.p_align
+                  ]
+
+        return (Elf32_Phdr.s).pack(*values)
+
+
+class Elf32_Shdr:
+    s = struct.Struct('I' * 10)
+
+    def __init__(self, data):
+        unpacked_data = (Elf32_Shdr.s).unpack(data)
+        self.sh_name = unpacked_data[0]
+        self.sh_type = unpacked_data[1]
+        self.sh_flags = unpacked_data[2]
+        self.sh_addr = unpacked_data[3]
+        self.sh_offset = unpacked_data[4]
+        self.sh_size = unpacked_data[5]
+        self.sh_link = unpacked_data[6]
+        self.sh_info = unpacked_data[7]
+        self.sh_addralign = unpacked_data[8]
+        self.sh_entsize = unpacked_data[9]
+        self.sh_name_str = ""
+
+    def get_packed_data(self):
+        values = [self.sh_name,
+                  self.sh_type,
+                  self.sh_flags,
+                  self.sh_addr,
+                  self.sh_offset,
+                  self.sh_size,
+                  self.sh_link,
+                  self.sh_info,
+                  self.sh_addralign,
+                  self.sh_entsize]
+
+        return (Elf32_Shdr.s).pack(*values)
+
+
+class Elf64_Ehdr:
+    s = struct.Struct('16sHHIQQQIHHHHHH')
+
+    def __init__(self, data):
+        unpacked_data = (Elf64_Ehdr.s).unpack(data)
+        self.e_ident = unpacked_data[0]
+        self.e_type = unpacked_data[1]
+        self.e_machine = unpacked_data[2]
+        self.e_version = unpacked_data[3]
+        self.e_entry = unpacked_data[4]
+        self.e_phoff = unpacked_data[5]
+        self.e_shoff = unpacked_data[6]
+        self.e_flags = unpacked_data[7]
+        self.e_ehsize = unpacked_data[8]
+        self.e_phentsize = unpacked_data[9]
+        self.e_phnum = unpacked_data[10]
+        self.e_shentsize = unpacked_data[11]
+        self.e_shnum = unpacked_data[12]
+        self.e_shstrndx = unpacked_data[13]
+
+    def get_packed_data(self):
+        values = [self.e_ident,
+                  self.e_type,
+                  self.e_machine,
+                  self.e_version,
+                  self.e_entry,
+                  self.e_phoff,
+                  self.e_shoff,
+                  self.e_flags,
+                  self.e_ehsize,
+                  self.e_phentsize,
+                  self.e_phnum,
+                  self.e_shentsize,
+                  self.e_shnum,
+                  self.e_shstrndx
+                  ]
+
+        return (Elf64_Ehdr.s).pack(*values)
+
+
+class Elf64_Phdr:
+    s = struct.Struct('IIQQQQQQ')
+
+    def __init__(self, data):
+        unpacked_data = (Elf64_Phdr.s).unpack(data)
+        self.p_type = unpacked_data[0]
+        self.p_flags = unpacked_data[1]
+        self.p_offset = unpacked_data[2]
+        self.p_vaddr = unpacked_data[3]
+        self.p_paddr = unpacked_data[4]
+        self.p_filesz = unpacked_data[5]
+        self.p_memsz = unpacked_data[6]
+        self.p_align = unpacked_data[7]
+
+    def get_packed_data(self):
+        values = [self.p_type,
+                  self.p_flags,
+                  self.p_offset,
+                  self.p_vaddr,
+                  self.p_paddr,
+                  self.p_filesz,
+                  self.p_memsz,
+                  self.p_align
+                  ]
+
+        return (Elf64_Phdr.s).pack(*values)
+
+
+class Elf64_Shdr:
+    s = struct.Struct('IIQQQQIIQQ')
+
+    def __init__(self, data):
+        unpacked_data = (Elf64_Shdr.s).unpack(data)
+        self.sh_name = unpacked_data[0]
+        self.sh_type = unpacked_data[1]
+        self.sh_flags = unpacked_data[2]
+        self.sh_addr = unpacked_data[3]
+        self.sh_offset = unpacked_data[4]
+        self.sh_size = unpacked_data[5]
+        self.sh_link = unpacked_data[6]
+        self.sh_info = unpacked_data[7]
+        self.sh_addralign = unpacked_data[8]
+        self.sh_entsize = unpacked_data[9]
+        self.sh_name_str = ""
+
+    def get_packed_data(self):
+        values = [self.sh_name,
+                  self.sh_type,
+                  self.sh_flags,
+                  self.sh_addr,
+                  self.sh_offset,
+                  self.sh_size,
+                  self.sh_link,
+                  self.sh_info,
+                  self.sh_addralign,
+                  self.sh_entsize]
+
+        return (Elf64_Shdr.s).pack(*values)
+
+
+class SWIV_Seg:
+    s = struct.Struct('4sIQ')
+
+    def __init__(self):
+        self.magic = SWIV_MAGIC
+        self.crc = 0
+        self.padding = 0
+
+    def get_packed_data(self):
+        values = [self.magic,
+                  self.crc,
+                  self.padding
+                  ]
+
+        return (SWIV_Seg.s).pack(*values)
+
+
+def is_elf64(elf_header):
+    return elf_header.e_ident[ELF_CLASS_INDEX] == ELFCLASS64
+
+
+def verify_elf_header(elf_header):
+    ident = elf_header.e_ident
+    return (ident.startswith(ELF_MAGIC)
+            and ident[ELF_CLASS_INDEX] in (ELFCLASS32, ELFCLASS64)
+            and ident[ELF_VERSION_INDEX] == EV_CURRENT)
+
+
+def file_copy_offset(in_fp, in_off, out_fp, out_off, num_bytes):
+    in_fp.seek(in_off)
+    read_in = in_fp.read(num_bytes)
+    out_fp.seek(out_off)
+    out_fp.write(read_in)
+
+    return num_bytes
+
+
+def preprocess_elf_file(elf_file_name):
+    with open(elf_file_name, 'rb') as elf_fp:
+        elf_header = Elf_Ehdr_common(elf_fp.read(Elf_Ehdr_common.s.size))
+
+        if verify_elf_header(elf_header) is False:
+            raise RuntimeError("Input ELF verification failed : " + elf_file_name)
+
+        elf_fp.seek(0)
+
+        if is_elf64(elf_header):
+            elf_header = Elf64_Ehdr(elf_fp.read(Elf64_Ehdr.s.size))
+        else:
+            elf_header = Elf32_Ehdr(elf_fp.read(Elf32_Ehdr.s.size))
+
+        phdr_table = []
+
+        phdr_size = elf_header.e_phentsize
+
+        file_offset = elf_header.e_phoff
+        elf_fp.seek(file_offset)
+
+        for _ in range(elf_header.e_phnum):
+            if is_elf64(elf_header):
+                phdr_table.append(Elf64_Phdr(elf_fp.read(phdr_size)))
+            else:
+                phdr_table.append(Elf32_Phdr(elf_fp.read(phdr_size)))
+
+    return [elf_header, phdr_table]
+
+
+def process_section_header(elf_file_name):
+    with open(elf_file_name, 'rb') as elf_fp:
+        elf_header = Elf_Ehdr_common(elf_fp.read(Elf_Ehdr_common.s.size))
+
+        if verify_elf_header(elf_header) is False:
+            raise RuntimeError("Input ELF verification failed : " + elf_file_name)
+
+        elf_fp.seek(0)
+
+        if is_elf64(elf_header):
+            elf_header = Elf64_Ehdr(elf_fp.read(Elf64_Ehdr.s.size))
+        else:
+            elf_header = Elf32_Ehdr(elf_fp.read(Elf32_Ehdr.s.size))
+
+        shdr_table = []
+
+        shdr_size = elf_header.e_shentsize
+        elf_fp.seek(elf_header.e_shoff)
+
+        for _ in range(elf_header.e_shnum):
+            if is_elf64(elf_header):
+                shdr_table.append(Elf64_Shdr(elf_fp.read(shdr_size)))
+            else:
+                shdr_table.append(Elf32_Shdr(elf_fp.read(shdr_size)))
+
+    return shdr_table
+
+
+def crc32_update_file(elf_fp, offset, size, crc):
+    elf_fp.seek(offset)
+    remaining = size
+    while remaining > 0:
+        chunk = elf_fp.read(min(1024 * 1024, remaining))
+        if not chunk:
+            raise RuntimeError("Unexpected EOF while calculating CRC")
+        crc = zlib.crc32(chunk, crc)
+        remaining -= len(chunk)
+    return crc
+
+
+def crc32_update_zeros(size, crc):
+    zero_chunk = b'\x00' * min(1024 * 1024, size)
+    remaining = size
+    while remaining > 0:
+        chunk = zero_chunk[:min(len(zero_chunk), remaining)]
+        crc = zlib.crc32(chunk, crc)
+        remaining -= len(chunk)
+    return crc
+
+
+def crc32_calculate(elf_fp, phdr_table):
+    crc = 0
+
+    # SWIV validates loadable segment contents, including BSS as zeros.
+    for curr_phdr in phdr_table:
+        if curr_phdr.p_type != LOAD_TYPE:
+            continue
+
+        crc = crc32_update_file(elf_fp, curr_phdr.p_offset, curr_phdr.p_filesz, crc)
+
+        zero_size = curr_phdr.p_memsz - curr_phdr.p_filesz
+        if zero_size > 0:
+            crc = crc32_update_zeros(zero_size, crc)
+
+    crc &= 0xffffffff
+    logger.info("calculated CRC: %#x", crc)
+
+    return crc
+
+
+def verify_swiv_seg(elf_in_fp, offset):
+
+    elf_in_fp.seek(offset)
+    swiv_seg = SWIV_Seg()
+    unpacked_data = (SWIV_Seg.s).unpack(elf_in_fp.read(SWIV_SEG_SIZE_MIN))
+    swiv_seg.magic = unpacked_data[0]
+
+    return swiv_seg.magic == SWIV_MAGIC
+
+
+def same_output_file(input_elf, output_elf):
+    try:
+        return os.path.exists(output_elf) and os.path.samefile(input_elf, output_elf)
+    except FileNotFoundError:
+        return False
+
+
+def make_output_target(input_elf, output_elf):
+    if not same_output_file(input_elf, output_elf):
+        return output_elf, None
+
+    output_dir = os.path.dirname(os.path.abspath(output_elf))
+    output_base = os.path.basename(output_elf)
+    fd, temp_output = tempfile.mkstemp(prefix=output_base + ".", suffix=".tmp", dir=output_dir)
+    os.close(fd)
+    shutil.copymode(input_elf, temp_output)
+    return temp_output, temp_output
+
+
+def swiv_generate(input_elf, output_elf):
+    output_target, temp_output = make_output_target(input_elf, output_elf)
+
+    try:
+        with open(input_elf, "rb") as elf_in_fp, open(output_target, "wb+") as elf_out_fp:
+            [elf_header, phdr_table] = preprocess_elf_file(input_elf)
+
+            num_phdrs = elf_header.e_phnum
+            num_shdrs = elf_header.e_shnum
+            phdr_size = elf_header.e_phentsize
+
+            if num_phdrs == 0 or num_phdrs > MAX_PHDR_COUNT:
+                raise RuntimeError("Input ELF has invalid number of program headers.")
+
+            if elf_header.e_shoff == 0 or num_shdrs == 0:
+                raise RuntimeError("Input ELF does not contain a section header table.")
+
+            shdr_table_input = process_section_header(input_elf)
+            shdr_table_output = process_section_header(input_elf)
+
+            last_phdr = phdr_table[num_phdrs - 1]
+            if last_phdr.p_memsz == SWIV_SEG_SIZE_MIN and verify_swiv_seg(elf_in_fp, last_phdr.p_offset):
+                raise RuntimeError("A SWIV segment is detected in the input ELF.")
+
+            last_shdr = shdr_table_input[elf_header.e_shnum - 1]
+            if last_shdr.sh_size == SWIV_SEG_SIZE_MIN and verify_swiv_seg(elf_in_fp, last_shdr.sh_offset):
+                raise RuntimeError("A SWIV section is detected in the input ELF.")
+
+            if is_elf64(elf_header):
+                swiv_shdr = Elf64_Shdr(b'\0' * Elf64_Shdr.s.size)
+            else:
+                swiv_shdr = Elf32_Shdr(b'\0' * Elf32_Shdr.s.size)
+
+            swiv_shdr.sh_name = 0
+            swiv_shdr.sh_type = SWIV_SECTION_TYPE
+            swiv_shdr.sh_flags = 0
+            swiv_shdr.sh_addr = 0
+            swiv_shdr.sh_size = SWIV_SEG_SIZE_MIN
+            swiv_shdr.sh_offset = elf_header.e_shoff + (elf_header.e_shnum + 1) * elf_header.e_shentsize
+            shdr_table_output.append(swiv_shdr)
+            elf_header.e_shnum += 1
+
+            if elf_header.e_phoff != 0:
+                for phdr in phdr_table:
+                    file_copy_offset(elf_in_fp, phdr.p_offset, elf_out_fp, phdr.p_offset, phdr.p_filesz)
+
+            elf_out_fp.seek(0)
+            elf_out_fp.write(elf_header.get_packed_data())
+
+            phdr_start = elf_header.e_phoff
+            for phdr in phdr_table:
+                elf_out_fp.seek(phdr_start)
+                elf_out_fp.write(phdr.get_packed_data())
+                phdr_start += phdr_size
+
+            for shdr in shdr_table_input[:num_shdrs]:
+                if shdr.sh_type == SHT_NOBITS or shdr.sh_size == 0:
+                    continue
+                file_copy_offset(elf_in_fp, shdr.sh_offset, elf_out_fp, shdr.sh_offset, shdr.sh_size)
+
+            shdr_start = elf_header.e_shoff
+            for shdr in shdr_table_output:
+                elf_out_fp.seek(shdr_start)
+                elf_out_fp.write(shdr.get_packed_data())
+                shdr_start += elf_header.e_shentsize
+
+            elf_out_fp.flush()
+
+            swiv_seg = SWIV_Seg()
+            swiv_seg.crc = crc32_calculate(elf_out_fp, phdr_table)
+            elf_out_fp.seek(swiv_shdr.sh_offset)
+            elf_out_fp.write(swiv_seg.get_packed_data())
+
+        if temp_output is not None:
+            os.replace(temp_output, output_elf)
+            temp_output = None
+
+        return 0
+    finally:
+        if temp_output is not None and os.path.exists(temp_output):
+            os.unlink(temp_output)
+
+
+def main(argv):
+    parser = argparse.ArgumentParser(description="Add a Software Image Validation (SWIV) section to an HTP ELF file.")
+    parser.add_argument("-i", "--input", dest="input_file", metavar="INPUT", required=True)
+    parser.add_argument("-o", "--output", dest="output_file", metavar="OUTPUT", required=True)
+    parser.add_argument("-v", "--verbose", action="store_true", help="show calculated CRC")
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(
+        level=logging.INFO if args.verbose else logging.WARNING,
+        format="%(message)s",
+    )
+
+    try:
+        return swiv_generate(args.input_file, args.output_file)
+    except RuntimeError as exc:
+        logger.error("ERROR: %s", exc)
+        return 2
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))

ggml/src/ggml-hexagon/libggml-htp.inf

@@ -18,6 +18,7 @@ libggml-htp-v73.so = 1
 libggml-htp-v75.so = 1
 libggml-htp-v79.so = 1
 libggml-htp-v81.so = 1
+libggml-htp-v85.so = 1
 
 [ControlFlags]
 ExcludeFromSelect = *
@@ -30,6 +31,7 @@ libggml-htp-v73.so,,,0x10 ;COPYFLG_NO_OVERWRITE
 libggml-htp-v75.so,,,0x10 ;COPYFLG_NO_OVERWRITE
 libggml-htp-v79.so,,,0x10 ;COPYFLG_NO_OVERWRITE
 libggml-htp-v81.so,,,0x10 ;COPYFLG_NO_OVERWRITE
+libggml-htp-v85.so,,,0x10 ;COPYFLG_NO_OVERWRITE
 
 [Strings]
 GGML   = 'GGML'

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -1567,6 +1567,28 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax(ggml_meta
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax_part(ggml_metal_library_t lib) {
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, "kernel_argmax_part_f32");
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, "kernel_argmax_part_f32", "kernel_argmax_part_f32", nullptr);
+    }
+
+    res.smem = 32*(sizeof(float) + sizeof(int32_t));
+
+    return res;
+}
+
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax_final(ggml_metal_library_t lib) {
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, "kernel_argmax_final_f32");
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, "kernel_argmax_final_f32", "kernel_argmax_final_f32", nullptr);
+    }
+
+    res.smem = 32*(sizeof(float) + sizeof(int32_t));
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort(ggml_metal_library_t lib, const ggml_tensor * op) {
     assert(op->op == GGML_OP_ARGSORT);
 

ggml/src/ggml-metal/ggml-metal-device.h

@@ -145,6 +145,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax            (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax_part       (ggml_metal_library_t lib);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax_final      (ggml_metal_library_t lib);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort           (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge     (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht              (ggml_metal_library_t lib, int n, enum ggml_type tsrc);

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -16,6 +16,9 @@
 
 #define N_MM_NPART_AMAX 256
 
+#define N_ARGMAX_NPART_MAX 256
+#define N_ARGMAX_CHUNK 4096
+
 // kernel parameters for mat-vec threadgroups
 //
 // N_R0: number of src0 rows to process per simdgroup
@@ -1379,8 +1382,10 @@ typedef struct {
 } ggml_metal_kargs_pool_1d;
 
 typedef struct {
-     int64_t ne00;
+    int64_t ne00;
     uint64_t nb01;
+    int64_t npart;
+    int64_t ne01;
 } ggml_metal_kargs_argmax;
 
 typedef struct {

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -2319,13 +2319,18 @@ int ggml_metal_op_pool_1d(ggml_metal_op_t ctx, int idx) {
 
     auto pipeline = ggml_metal_library_get_pipeline_pool_1d(lib, op, op_pool);
 
-    const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), (int) np);
-    const int ntg = (np + nth - 1) / nth;
+    const int nth = std::min(256, std::max(1, k0));
+    const int ntg = (int) np;
+
+    const size_t smem = op_pool == GGML_OP_POOL_AVG ?
+        2*32*sizeof(float) :
+        32*sizeof(float);
 
     ggml_metal_encoder_set_pipeline(enc, pipeline);
     ggml_metal_encoder_set_bytes   (enc, &args_pool_1d, sizeof(args_pool_1d),  0);
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         2);
+    ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
 
     ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1);
 
@@ -5421,6 +5426,25 @@ int ggml_metal_op_timestep_embedding(ggml_metal_op_t ctx, int idx) {
     return 1;
 }
 
+size_t ggml_metal_op_argmax_extra(const ggml_tensor * op) {
+    assert(op->op == GGML_OP_ARGMAX);
+
+    const int64_t ne00 = op->src[0]->ne[0];
+    const int64_t ne01 = op->src[0]->ne[1];
+
+    int npart = 1;
+    if (ne00 > N_ARGMAX_CHUNK) {
+        npart = (int) std::min<int64_t>(N_ARGMAX_NPART_MAX, (ne00 + N_ARGMAX_CHUNK - 1)/N_ARGMAX_CHUNK);
+    }
+
+    if (npart == 1) {
+        return 0;
+    }
+
+    // partial max + partial argmax for every row
+    return (sizeof(float) + sizeof(int32_t)) * npart * ne01;
+}
+
 int ggml_metal_op_argmax(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 
@@ -5432,22 +5456,61 @@ int ggml_metal_op_argmax(ggml_metal_op_t ctx, int idx) {
     GGML_TENSOR_LOCALS( int32_t, ne,  op,         ne);
     GGML_TENSOR_LOCALS(uint64_t, nb,  op,         nb);
 
+    const int64_t nrows = ggml_nrows(op->src[0]);
+
+    int npart = 1;
+    if (ne00 > N_ARGMAX_CHUNK) {
+        npart = (int) std::min<int64_t>(N_ARGMAX_NPART_MAX, (ne00 + N_ARGMAX_CHUNK - 1)/N_ARGMAX_CHUNK);
+    }
+
     ggml_metal_kargs_argmax args = {
-        /*.ne00 = */ ne00,
-        /*.nb01 = */ nb01,
+        /*.ne00  = */ ne00,
+        /*.nb01  = */ nb01,
+        /*.npart = */ npart,
+        /*.ne01  = */ ne01,
     };
 
-    auto pipeline = ggml_metal_library_get_pipeline_argmax(lib, op);
+    const size_t smem = 32*(sizeof(float) + sizeof(int32_t));
 
-    const int64_t nrows = ggml_nrows(op->src[0]);
+    if (npart > 1) {
+        auto pipeline_part = ggml_metal_library_get_pipeline_argmax_part(lib);
+        auto pipeline_final = ggml_metal_library_get_pipeline_argmax_final(lib);
+
+        ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op);
+        ggml_metal_buffer_id bid_tmp = bid_dst;
+        bid_tmp.offs += ggml_nbytes(op);
+
+        ggml_metal_encoder_set_pipeline(enc, pipeline_part);
+        ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
+        ggml_metal_encoder_set_buffer  (enc, bid_tmp, 2);
+        ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
+        ggml_metal_encoder_dispatch_threadgroups(enc, npart, nrows, 1, 256, 1, 1);
+
+        ggml_metal_op_concurrency_reset(ctx);
+
+        int nth = 32; // SIMD width
+        while (nth < npart && nth < 256) {
+            nth *= 2;
+        }
+
+        ggml_metal_encoder_set_pipeline(enc, pipeline_final);
+        ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+        ggml_metal_encoder_set_buffer  (enc, bid_tmp, 1);
+        ggml_metal_encoder_set_buffer  (enc, bid_dst, 2);
+        ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
+        ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, nth, 1, 1);
+
+        return 1;
+    }
+
+    auto pipeline = ggml_metal_library_get_pipeline_argmax(lib, op);
 
     int nth = 32; // SIMD width
     while (nth < ne00 && nth*ne01*ne02*ne03 < 256) {
         nth *= 2;
     }
 
-    const size_t smem = pipeline.smem;
-
     ggml_metal_encoder_set_pipeline(enc, pipeline);
     ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);

ggml/src/ggml-metal/ggml-metal-ops.h

@@ -47,6 +47,8 @@ size_t ggml_metal_op_flash_attn_ext_extra_tmp(const struct ggml_tensor * op);
 size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const struct ggml_tensor * op);
 size_t ggml_metal_op_flash_attn_ext_extra_idx(const struct ggml_tensor * op);
 
+size_t ggml_metal_op_argmax_extra(const struct ggml_tensor * op);
+
 int ggml_metal_op_concat            (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_repeat            (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_acc               (ggml_metal_op_t ctx, int idx);

ggml/src/ggml-metal/ggml-metal.cpp

@@ -245,6 +245,10 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_
             {
                 res = 2*sizeof(int32_t)*ggml_nelements(tensor->src[0]);
             } break;
+        case GGML_OP_ARGMAX:
+            {
+                res += ggml_metal_op_argmax_extra(tensor);
+            } break;
         default:
             break;
     }

ggml/src/ggml-metal/kernels/misc.metal

@@ -60,6 +60,140 @@ kernel void kernel_argmax_f32(
     dst_i32[tgpig] = arg_val;
 }
 
+kernel void kernel_argmax_part_f32(
+        constant ggml_metal_kargs_argmax & args,
+        device   const char * src0,
+        device         char * tmp,
+        threadgroup    char * shmem [[threadgroup(0)]],
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint  sgitg[[simdgroup_index_in_threadgroup]],
+        uint  tiisg[[thread_index_in_simdgroup]],
+        uint3  ntg[[threads_per_threadgroup]]) {
+    const int row  = tgpig.y;
+    const int part = tgpig.x;
+    const int ne00 = (int) args.ne00;
+    const int npart = (int) args.npart;
+    const int chunk = (ne00 + npart - 1) / npart;
+    const int start = part * chunk;
+    const int end   = min(start + chunk, ne00);
+
+    device const float * x_row = (device const float *) ((device const char *) src0 + row * args.nb01);
+
+    float   lmax = -INFINITY;
+    int32_t larg = -1;
+
+    for (int i00 = start + tpitg.x; i00 < end; i00 += ntg.x) {
+        if (x_row[i00] > lmax) {
+            lmax = x_row[i00];
+            larg = i00;
+        }
+    }
+
+    // find the argmax value in the block
+    float max_val = simd_max(lmax);
+    int32_t arg_val = simd_max(select(-1, larg, lmax == max_val));
+
+    threadgroup   float * shared_maxval = (threadgroup   float *) shmem;
+    threadgroup int32_t * shared_argmax = (threadgroup int32_t *) shmem + N_SIMDWIDTH;
+
+    if (ntg.x > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            shared_maxval[tiisg] = -INFINITY;
+            shared_argmax[tiisg] = -1;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            shared_maxval[sgitg] = max_val;
+            shared_argmax[sgitg] = arg_val;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        max_val = shared_maxval[tiisg];
+        arg_val = shared_argmax[tiisg];
+
+        float max_val_reduced   = simd_max(max_val);
+        int32_t arg_val_reduced = simd_max(select(-1, arg_val, max_val == max_val_reduced));
+
+        max_val = max_val_reduced;
+        arg_val = arg_val_reduced;
+    }
+
+    if (tpitg.x == 0) {
+        device float * tmp_max = (device float *) tmp;
+        device int32_t * tmp_arg = (device int32_t *) tmp + npart * (int) args.ne01;
+
+        tmp_max[row*npart + part] = max_val;
+        tmp_arg[row*npart + part] = arg_val;
+    }
+}
+
+kernel void kernel_argmax_final_f32(
+        constant ggml_metal_kargs_argmax & args,
+        device   const char * tmp,
+        device         char * dst,
+        threadgroup    char * shmem [[threadgroup(0)]],
+        uint3 tgpig[[threadgroup_position_in_grid]],
+        uint3 tpitg[[thread_position_in_threadgroup]],
+        uint  sgitg[[simdgroup_index_in_threadgroup]],
+        uint  tiisg[[thread_index_in_simdgroup]],
+        uint3  ntg[[threads_per_threadgroup]]) {
+    const int row = tgpig.x;
+    const int npart = (int) args.npart;
+
+    device const float * tmp_max = (device const float *) tmp;
+    device const int32_t * tmp_arg = (device const int32_t *) tmp + npart * (int) args.ne01;
+
+    float   lmax = -INFINITY;
+    int32_t larg = -1;
+
+    for (int i = tpitg.x; i < npart; i += ntg.x) {
+        if (tmp_max[row*npart + i] > lmax) {
+            lmax = tmp_max[row*npart + i];
+            larg = tmp_arg[row*npart + i];
+        }
+    }
+
+    // find the argmax value in the block
+    float max_val = simd_max(lmax);
+    int32_t arg_val = simd_max(select(-1, larg, lmax == max_val));
+
+    threadgroup   float * shared_maxval = (threadgroup   float *) shmem;
+    threadgroup int32_t * shared_argmax = (threadgroup int32_t *) shmem + N_SIMDWIDTH;
+
+    if (ntg.x > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            shared_maxval[tiisg] = -INFINITY;
+            shared_argmax[tiisg] = -1;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            shared_maxval[sgitg] = max_val;
+            shared_argmax[sgitg] = arg_val;
+        }
+
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        max_val = shared_maxval[tiisg];
+        arg_val = shared_argmax[tiisg];
+
+        float max_val_reduced   = simd_max(max_val);
+        int32_t arg_val_reduced = simd_max(select(-1, arg_val, max_val == max_val_reduced));
+
+        max_val = max_val_reduced;
+        arg_val = arg_val_reduced;
+    }
+
+    if (tpitg.x == 0) {
+        ((device int32_t *) dst)[row] = arg_val;
+    }
+}
+
 kernel void kernel_diag_f32(
         constant ggml_metal_kargs_diag & args,
         device   const char * src0,

ggml/src/ggml-metal/kernels/mul_mm.metal

@@ -593,6 +593,7 @@ template [[host_name("kernel_mul_mm_id_map0_ne20_4" )]] kernel kernel_mul_mm_id_
 template [[host_name("kernel_mul_mm_id_map0_ne20_5" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<5>;
 template [[host_name("kernel_mul_mm_id_map0_ne20_6" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<6>;
 template [[host_name("kernel_mul_mm_id_map0_ne20_8" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<8>;
+template [[host_name("kernel_mul_mm_id_map0_ne20_9" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<9>;
 template [[host_name("kernel_mul_mm_id_map0_ne20_10")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<10>;
 template [[host_name("kernel_mul_mm_id_map0_ne20_16")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<16>;
 template [[host_name("kernel_mul_mm_id_map0_ne20_22")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<22>;

ggml/src/ggml-metal/kernels/pool.metal

@@ -84,65 +84,117 @@ kernel void kernel_pool_1d_max_f32(
         constant        ggml_metal_kargs_pool_1d & args,
         device  const   float * src,
         device          float * dst,
-        uint            gid [[thread_position_in_grid]]
-) {
-
-    if (gid >= args.np) {
+        threadgroup    char * shmem [[threadgroup(0)]],
+        uint tgpig[[threadgroup_position_in_grid]],
+        uint tpitg[[thread_position_in_threadgroup]],
+        uint sgitg[[simdgroup_index_in_threadgroup]],
+        uint tiisg[[thread_index_in_simdgroup]],
+        uint ntg[[threads_per_threadgroup]]) {
+    if (tgpig >= args.np) {
         return;
     }
 
-    const int ow  = (int)gid % args.OW;
-    const int row = (int)gid / args.OW;
+    const int out = (int) tgpig;
+    const int ow  = out % (int) args.OW;
+    const int row = out / (int) args.OW;
 
-    const int base = ow * args.s0 - args.p0;
+    const int base     = ow * args.s0 - args.p0;
+    const int src_off  = row * args.IW;
+    const int dst_off  = row * args.OW;
 
-    float acc = -INFINITY;
+    float lmax = -INFINITY;
 
-    const int src_off = row * args.IW;
-    const int dst_off = row * args.OW;
-
-    for (int ki = 0; ki < args.k0; ++ki) {
-        int j = base + ki;
-        if (j < 0 || j >= args.IW){
-            continue;
+    for (int ki = tpitg; ki < args.k0; ki += ntg) {
+        const int j = base + ki;
+        if (j >= 0 && j < args.IW) {
+            lmax = max(lmax, src[src_off + j]);
         }
-        float v = src[src_off + j];
-        acc = max(acc, v);
     }
 
-    dst[dst_off + ow] = acc;
+    float max_val = simd_max(lmax);
+
+    threadgroup float * shared_max = (threadgroup float *) shmem;
+
+    if (ntg > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            shared_max[tiisg] = -INFINITY;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            shared_max[sgitg] = max_val;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        max_val = shared_max[tiisg];
+        max_val = simd_max(max_val);
+    }
+
+    if (tpitg == 0) {
+        dst[dst_off + ow] = max_val;
+    }
 }
 
 kernel void kernel_pool_1d_avg_f32(
         constant        ggml_metal_kargs_pool_1d & args,
         device  const   float * src,
         device          float * dst,
-        uint            gid [[thread_position_in_grid]]
-) {
-
-    if (gid >= args.np) {
+        threadgroup    char * shmem [[threadgroup(0)]],
+        uint tgpig[[threadgroup_position_in_grid]],
+        uint tpitg[[thread_position_in_threadgroup]],
+        uint sgitg[[simdgroup_index_in_threadgroup]],
+        uint tiisg[[thread_index_in_simdgroup]],
+        uint ntg[[threads_per_threadgroup]]) {
+    if (tgpig >= args.np) {
         return;
     }
 
-    const int ow  = (int)gid % args.OW;
-    const int row = (int)gid / args.OW;
+    const int out = (int) tgpig;
+    const int ow  = out % (int) args.OW;
+    const int row = out / (int) args.OW;
 
-    const int base = ow * args.s0 - args.p0;
+    const int base     = ow * args.s0 - args.p0;
+    const int src_off  = row * args.IW;
+    const int dst_off  = row * args.OW;
 
-    float acc = 0.0f;
-    int   cnt = 0;
+    float lsum = 0.0f;
+    int   lcnt = 0;
 
-    const int src_off = row * args.IW;
-    const int dst_off = row * args.OW;
-
-    for (int ki = 0; ki < args.k0; ++ki) {
+    for (int ki = tpitg; ki < args.k0; ki += ntg) {
         const int j = base + ki;
-        if (j < 0 || j >= args.IW) {
-            continue;
+        if (j >= 0 && j < args.IW) {
+            lsum += src[src_off + j];
+            lcnt += 1;
         }
-        acc += src[src_off + j];
-        cnt += 1;
     }
 
-    dst[dst_off + ow] = (cnt > 0) ? (acc / (float)cnt) : 0.0f;
+    float sum = simd_sum(lsum);
+    float cnt = simd_sum((float) lcnt);
+
+    threadgroup float * shared_sum = (threadgroup float *) shmem;
+    threadgroup float * shared_cnt = (threadgroup float *) shmem + N_SIMDWIDTH;
+
+    if (ntg > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            shared_sum[tiisg] = 0.0f;
+            shared_cnt[tiisg] = 0.0f;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            shared_sum[sgitg] = sum;
+            shared_cnt[sgitg] = cnt;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        sum = shared_sum[tiisg];
+        cnt = shared_cnt[tiisg];
+
+        sum = simd_sum(sum);
+        cnt = simd_sum(cnt);
+    }
+
+    if (tpitg == 0) {
+        dst[dst_off + ow] = (cnt > 0.0f) ? (sum / cnt) : 0.0f;
+    }
 }

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -831,6 +831,7 @@ struct ggml_backend_opencl_context {
     cl_program program_mul_mv_id_mxfp4_f32_flat;
     cl_program program_mul_mm_f32_f32_l4_lm;
     cl_program program_mul_mm_f16_f32_l4_lm;
+    cl_program program_mul_mm_f16_f32_l4_lm_n8;
     cl_program program_mul_mm_q8_0_f32_l4_lm;
 
     cl_kernel kernel_add, kernel_add_row, kernel_add_f16, kernel_add_row_f16;
@@ -1052,6 +1053,7 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_mul_mm_f32_f32_l4_lm;
     cl_kernel kernel_gemv_f32_f32_mc;  // multi-column (small-N) f32 GEMV for spec/MTP verify
     cl_kernel kernel_mul_mm_f16_f32_l4_lm;
+    cl_kernel kernel_mul_mm_f16_f32_l4_lm_n8;
     cl_kernel kernel_mul_mm_q1_0_f32_l4_lm;
     cl_kernel kernel_mul_mm_q4_0_f32_l4_lm;
     cl_kernel kernel_mul_mm_q4_1_f32_l4_lm;
@@ -1247,6 +1249,8 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_gemv_noshuffle_q8_0_f32_splitk;  // split-K across WGs (small-M decode)
     cl_kernel kernel_gemm_noshuffle_q1_0_f32;
     cl_kernel kernel_gemv_noshuffle_q1_0_f32;
+    cl_kernel kernel_gemv_noshuffle_q1_0_f32_nsg1 = nullptr;
+    int min_m_gemv_nsg1_q1_0 = 0;
     cl_kernel kernel_gemv_noshuffle_q4_k_f32;
     cl_kernel kernel_gemv_noshuffle_q4_k_f32_o4;  // 4-output-per-WI, long-vocab lm_head
     cl_kernel kernel_gemv_noshuffle_q4_k_f32_tiled;  // tiled-wide layout (opt-in)
@@ -1261,6 +1265,11 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_gemv_noshuffle_q4_k_f32_32b_trans;
     cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a = nullptr;  // dp4a (int8) dense prefill GEMM
     cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg = nullptr;  // dp4a dense prefill GEMM, weights via texture (X1 opt-in)
+    cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow = nullptr;
+    cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow_wimg = nullptr;
+    int tile_n_dp4a_q4_k = 32;
+    int tile_n_dp4a_narrow_q4_k = 32;
+    int max_n_dp4a_narrow_q4_k = 16;
     cl_kernel kernel_gemm_noshuffle_q5_k_q8_1_dp4a = nullptr;  // dp4a (int8) dense q5_K prefill GEMM
     cl_kernel kernel_gemm_noshuffle_q6_k_q8_1_dp4a = nullptr;  // dp4a (int8) dense q6_K prefill GEMM
     cl_kernel kernel_quant_a_q8_1;                    // plain activation q8_1 pre-pass
@@ -2781,6 +2790,11 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
             build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
 
         CL_CHECK((backend_ctx->kernel_mul_mm_f16_f32_l4_lm = clCreateKernel(backend_ctx->program_mul_mm_f16_f32_l4_lm, "kernel_mul_mm_f16_f32_l4_lm", &err), err));
+
+        const std::string narrow_opts = compile_opts + " -DBN=8 -DTN=1 -DKERNEL_NAME_LM=kernel_mul_mm_f16_f32_l4_lm_n8";
+        backend_ctx->program_mul_mm_f16_f32_l4_lm_n8 =
+            build_program_from_source(backend_ctx, kernel_src.c_str(), narrow_opts);
+        CL_CHECK((backend_ctx->kernel_mul_mm_f16_f32_l4_lm_n8 = clCreateKernel(backend_ctx->program_mul_mm_f16_f32_l4_lm_n8, "kernel_mul_mm_f16_f32_l4_lm_n8", &err), err));
         GGML_LOG_CONT(".");
     }
 
@@ -3983,6 +3997,16 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
 
         CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q1_0_f32 = clCreateKernel(prog, "kernel_gemv_noshuffle_q1_0_f32", &err), err));
         CL_CHECK(clReleaseProgram(prog));
+
+        const char * nsg1_min_m_env = getenv("GGML_OPENCL_Q1_0_GEMV_NSG1_MIN_M");
+        backend_ctx->min_m_gemv_nsg1_q1_0 = nsg1_min_m_env ? atoi(nsg1_min_m_env) :
+            (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ? 8192 : 0);
+        if (backend_ctx->min_m_gemv_nsg1_q1_0 > 0) {
+            const std::string nsg1_opts = CL_gemv_compile_opts + " -DN_SIMDGROUP=1";
+            cl_program prog_nsg1 = build_program_from_source(backend_ctx, kernel_src_CL_gemv_general.c_str(), nsg1_opts);
+            CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q1_0_f32_nsg1 = clCreateKernel(prog_nsg1, "kernel_gemv_noshuffle_q1_0_f32", &err), err));
+            CL_CHECK(clReleaseProgram(prog_nsg1));
+        }
         GGML_LOG_CONT(".");
     }
 
@@ -4453,16 +4477,32 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
 #else
         const std::string kernel_src = read_file("gemm_noshuffle_q4_k_q8_1_dp4a.cl");
 #endif
-        // Per-device dp4a dense tile. The X2-tuned TILESIZE_N=32 over-occupies LDS on
-        // X1 (1152 B/WG -> few resident WGs); TILESIZE_N=8 (288 B) lifts occupancy on
-        // X1, byte-identical. X2E keeps 32. Env override wins.
-        int q4k_dp4a_ts = (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) ? 8 : 32;
-        if (const char * e = getenv("GGML_OPENCL_Q4K_DP4A_TS")) { q4k_dp4a_ts = atoi(e); }
-        std::string dp4a_opts = compile_opts + " -DTILESIZE_N=" + std::to_string(q4k_dp4a_ts);
+        // X1 uses a smaller N tile to reduce local memory and register use.
+        int tile_n = (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) ? 8 : 32;
+        if (const char * e = getenv("GGML_OPENCL_Q4K_DP4A_TS")) { tile_n = atoi(e); }
+        GGML_ASSERT(tile_n >= 4 && tile_n <= 64 && tile_n % 4 == 0);
+        backend_ctx->tile_n_dp4a_q4_k = tile_n;
+        std::string dp4a_opts = compile_opts + " -DTILESIZE_N=" + std::to_string(tile_n);
         cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), dp4a_opts);
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a", &err), err));
         CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg = clCreateKernel(prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg", &err), err));
         CL_CHECK(clReleaseProgram(prog));
+
+        // Compile a smaller N tile to avoid unused columns in small batches.
+        int tile_n_narrow = MIN(tile_n, 16);
+        if (const char * e = getenv("GGML_OPENCL_Q4_K_DP4A_TS_NARROW")) { tile_n_narrow = atoi(e); }
+        GGML_ASSERT(tile_n_narrow >= 4 && tile_n_narrow <= 64 && tile_n_narrow % 4 == 0);
+        if (const char * e = getenv("GGML_OPENCL_Q4_K_DP4A_NARROW_MAX")) {
+            backend_ctx->max_n_dp4a_narrow_q4_k = atoi(e);
+        }
+        backend_ctx->tile_n_dp4a_narrow_q4_k = tile_n_narrow;
+        if (tile_n_narrow != tile_n) {
+            std::string narrow_opts = compile_opts + " -DTILESIZE_N=" + std::to_string(tile_n_narrow);
+            cl_program narrow_prog = build_program_from_source(backend_ctx, kernel_src.c_str(), narrow_opts);
+            CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow = clCreateKernel(narrow_prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a", &err), err));
+            CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow_wimg = clCreateKernel(narrow_prog, "kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg", &err), err));
+            CL_CHECK(clReleaseProgram(narrow_prog));
+        }
         GGML_LOG_CONT(".");
     }
 
@@ -9531,22 +9571,6 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
                     op->src[0]->ne[1] >= 32768) {   // vocab-scale weight; no FFN/attn weight is this tall
                     return false;
                 }
-                // The generic mul_mv (GEMV) kernels are wrong for large-batch prefill on
-                // Adreno. A quant mul_mat only avoids the GEMV when it reaches the Adreno
-                // trans-weight GEMM, which needs both a GEMM kernel for the type and
-                // use_adreno_kernels(). Decline the large-N shapes that would otherwise
-                // fall through to the GEMV.
-                {
-                    const ggml_type t = op->src[0]->type;
-                    const bool type_has_gemm = (t == GGML_TYPE_Q4_0 || t == GGML_TYPE_Q4_1 ||
-                                                t == GGML_TYPE_IQ4_NL || t == GGML_TYPE_Q8_0 ||
-                                                t == GGML_TYPE_Q4_K  || t == GGML_TYPE_Q5_K  ||
-                                                t == GGML_TYPE_Q6_K);
-                    const bool uses_gemm = type_has_gemm && use_adreno_kernels(backend_ctx, op->src[0]);
-                    if (!uses_gemm && op->src[1]->ne[1] >= 512) {
-                        return false;
-                    }
-                }
                 return op->src[1]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]);
             } else if (op->src[0]->type == GGML_TYPE_Q8_0) {
                 // ggml_cl_mul_mat_q8_0_f32_adreno now honors src1/dst view_offs (the
@@ -19834,7 +19858,9 @@ static void ggml_cl_mul_mat_q1_0_f32_adreno(ggml_backend_t backend, const ggml_t
         img_desc.buffer = b_sub_buf;
         CL_CHECK((b_img = clCreateImage(context, CL_MEM_READ_ONLY, &img_fmt, &img_desc, NULL, &err), err));
 
-        kernel = backend_ctx->kernel_gemv_noshuffle_q1_0_f32;
+        // Tall M has enough row parallelism to skip the K split and its reduction.
+        const bool use_nsg1 = backend_ctx->min_m_gemv_nsg1_q1_0 > 0 && M >= backend_ctx->min_m_gemv_nsg1_q1_0;
+        kernel = use_nsg1 ? backend_ctx->kernel_gemv_noshuffle_q1_0_f32_nsg1 : backend_ctx->kernel_gemv_noshuffle_q1_0_f32;
 
         int r2 = 1;
         int r3 = 1;
@@ -19856,8 +19882,9 @@ static void ggml_cl_mul_mat_q1_0_f32_adreno(ggml_backend_t backend, const ggml_t
         CL_CHECK(clSetKernelArg(kernel, 14, sizeof(int),      &r3));
 
         size_t wavesize = backend_ctx->adreno_wave_size;
-        size_t local_work_size[]  = { wavesize, 4, 1 };
-        size_t global_work_size[] = { CEIL_DIV(M, wavesize)*wavesize, 4, 1 };
+        const size_t nsg = use_nsg1 ? 1 : 4;
+        size_t local_work_size[]  = { wavesize, nsg, 1 };
+        size_t global_work_size[] = { CEIL_DIV(M, wavesize)*wavesize, nsg, 1 };
 
         backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
 
@@ -20321,7 +20348,14 @@ static void ggml_cl_mul_mat_q4_0_f32_adreno(ggml_backend_t backend, const ggml_t
         // per-lane K-walk) for small M. Layout stride is fixed (4 uints/block), so only
         // the K-split count changes; the mc3 kernel reads it via get_local_size(1). The
         // ne1==1 base kernel hardcodes N_SIMDGROUP=4, so it always stays at 4.
-        const int mc3_nsg = (use_q40_mc3 && ne01 < 4096) ? 8 : 4;
+        int mc3_nsg = (use_q40_mc3 && ne01 < 4096) ? 8 : 4;
+        if (use_q40_mc3) {
+            // Register pressure can lower the kernel limit below the device limit.
+            const size_t max_workgroup_size = backend_ctx->get_kernel_workgroup_size(kernel);
+            while (mc3_nsg > 1 && (size_t)(64 * mc3_nsg) > max_workgroup_size) {
+                mc3_nsg /= 2;
+            }
+        }
         size_t local_work_size[3] = {64, (size_t)mc3_nsg, 1};
         size_t global_work_size[3] = {(size_t)CEIL_DIV(ne01/2, 64)*64, (size_t)mc3_nsg, 1};
 
@@ -22667,8 +22701,13 @@ static void ggml_cl_mul_mat_q4_k_f32_adreno(ggml_backend_t backend, const ggml_t
                 }
             }
 
-            cl_kernel dk = use_wimg ? backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg
-                                    : backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a;
+            const bool use_narrow = backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow != nullptr
+                                 && N <= backend_ctx->max_n_dp4a_narrow_q4_k;
+            cl_kernel dk = use_narrow
+                ? (use_wimg ? backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow_wimg
+                            : backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_narrow)
+                : (use_wimg ? backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg
+                            : backend_ctx->kernel_gemm_noshuffle_q4_k_q8_1_dp4a);
             int ai = 0;
             if (use_wimg) {
                 CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_mem), &q4k_q_img));
@@ -22689,12 +22728,9 @@ static void ggml_cl_mul_mat_q4_k_f32_adreno(ggml_backend_t backend, const ggml_t
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_uchar), &mask_d6));
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_uchar), &mask_d4));
             CL_CHECK(clSetKernelArg(dk, ai++, sizeof(cl_uchar), &mask_hi2));
-            // Must match the compile-time TILESIZE_N chosen at program build (per-device,
-            // X1E=8 else 32; env override). Same inputs -> same value.
-            int q4k_dp4a_ts = (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) ? 8 : 32;
-            if (const char * e = getenv("GGML_OPENCL_Q4K_DP4A_TS")) q4k_dp4a_ts = atoi(e);
+            const int tile_n = use_narrow ? backend_ctx->tile_n_dp4a_narrow_q4_k : backend_ctx->tile_n_dp4a_q4_k;
             size_t d_local[3]  = { 64, 1, 1 };
-            size_t d_global[3] = { 64, (size_t)(M / 64), (size_t)CEIL_DIV(N, q4k_dp4a_ts) };
+            size_t d_global[3] = { 64, (size_t)(M / 64), (size_t)CEIL_DIV(N, tile_n) };
             backend_ctx->enqueue_ndrange_kernel(dk, 3, d_global, d_local, dst);
 
             if (q4k_q_img != nullptr) {
@@ -24815,7 +24851,9 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
                     return;
                 }
 #endif
-                kernel = backend_ctx->kernel_mul_mm_f16_f32_l4_lm;
+                static const char * mm_narrow_n_env = getenv("GGML_OPENCL_MM_NARROW_N");
+                const bool use_narrow_n = (mm_narrow_n_env == nullptr || mm_narrow_n_env[0] != '0') && ne11 >= 2 && ne11 <= 8;
+                kernel = use_narrow_n ? backend_ctx->kernel_mul_mm_f16_f32_l4_lm_n8 : backend_ctx->kernel_mul_mm_f16_f32_l4_lm;
                 nth0 = 128; // calculated as (BM*BN)/(TM*TN)
 
                 int batch_stride_a = ne00*ne01;
@@ -24874,8 +24912,9 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
                 CL_CHECK(clSetKernelArg(kernel, 17, sizeof(int),      &r2));
                 CL_CHECK(clSetKernelArg(kernel, 18, sizeof(int),      &r3));
 
-                // 64 is block tile size BM and BN - change here when BM and BN in the kernel are changed.
-                size_t global_work_size[] = {(size_t)(CEIL_DIV(ne01, 64)*nth0), (size_t)(CEIL_DIV(ne11, 64)), (size_t)ne12*ne13};
+                // BN=8/TN=1 keeps 128 threads for small batches, including attention KQ/KQV.
+                const int bn = use_narrow_n ? 8 : 64;
+                size_t global_work_size[] = {(size_t)(CEIL_DIV(ne01, 64)*nth0), (size_t)(CEIL_DIV(ne11, bn)), (size_t)ne12*ne13};
                 size_t local_work_size[] = {(size_t)nth0, 1, 1};
 
                 backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q1_0_f32.cl

@@ -8,7 +8,9 @@
 #endif
 
 #define QK1_0 128
+#ifndef N_SIMDGROUP
 #define N_SIMDGROUP 4
+#endif
 
 #define dequantizeBlockAccum_q1(total, bits, scale, regB, lb)                                       \
     total += (2.0f*(float)((bits >>  0) & 1u) - 1.0f) * scale * sub_group_broadcast(regB.s0, lb+0); \
@@ -106,15 +108,19 @@ __kernel void kernel_gemv_noshuffle_q1_0_f32(
         dequantizeBlockAccum_q1(totalSum, regA.s3, scale, regB, 12);
     }
 
-    // reduction in local memory, assumes #wave = N_SIMDGROUP = 4
-    local float reduceLM[SIMDGROUP_WIDTH * 3];
-    if (groupId == 1) reduceLM[SIMDGROUP_WIDTH * 0 + slid] = totalSum;
-    if (groupId == 2) reduceLM[SIMDGROUP_WIDTH * 1 + slid] = totalSum;
-    if (groupId == 3) reduceLM[SIMDGROUP_WIDTH * 2 + slid] = totalSum;
+#if N_SIMDGROUP > 1
+    // Sum the K slices from the other subgroups.
+    local float reduceLM[SIMDGROUP_WIDTH * (N_SIMDGROUP - 1)];
+    if (groupId > 0) {
+        reduceLM[SIMDGROUP_WIDTH * (groupId - 1) + slid] = totalSum;
+    }
     barrier(CLK_LOCAL_MEM_FENCE);
-    if (groupId == 0) totalSum += reduceLM[SIMDGROUP_WIDTH * 0 + slid];
-    if (groupId == 0) totalSum += reduceLM[SIMDGROUP_WIDTH * 1 + slid];
-    if (groupId == 0) totalSum += reduceLM[SIMDGROUP_WIDTH * 2 + slid];
+    if (groupId == 0) {
+        for (uint w = 0; w < N_SIMDGROUP - 1; ++w) {
+            totalSum += reduceLM[SIMDGROUP_WIDTH * w + slid];
+        }
+    }
+#endif
 
     if (groupId == 0) {
         dst = (global float*)((global char*)dst + offsetd);
@@ -122,6 +128,8 @@ __kernel void kernel_gemv_noshuffle_q1_0_f32(
         // so when ne01 is not a multiple of the wave size the tail work-items run past
         // row ne01 and would overrun dst into the adjacent tensor. No-op / byte-identical
         // when ne01 is wave-aligned (no padding).
-        if (gid < M) dst[gid] = totalSum;
+        if (gid < M) {
+            dst[gid] = totalSum;
+        }
     }
 }

ggml/src/ggml-opencl/kernels/mul_mm_f16_f32_l4_lm.cl

@@ -4,12 +4,20 @@
 #define LOAD_VEC_B 4
 
 #define BM 64
+#ifndef BN
 #define BN 64
+#endif
 #define BK 16
 #define TM 4
+#ifndef TN
 #define TN 8
+#endif
 
-kernel void kernel_mul_mm_f16_f32_l4_lm(
+#ifndef KERNEL_NAME_LM
+#define KERNEL_NAME_LM kernel_mul_mm_f16_f32_l4_lm
+#endif
+
+kernel void KERNEL_NAME_LM(
     global half4 * src0,
     ulong offset0,
     global float4 * src1,
@@ -93,7 +101,8 @@ kernel void kernel_mul_mm_f16_f32_l4_lm(
             }
         }
 
-        for (int l = 0; l < BN; l += loadstride_b) {
+        // BN=8 uses only the threads that load a column inside the tile.
+        for (int l = 0; loadc_b + l < BN; l += loadstride_b) {
             if (ic*BN + loadc_b + l < ne11) {
                 const int idx = pos_b + (loadc_b + l) * stride_b / LOAD_VEC_B + loadr_b;
                 buf_b[(loadr_b * LOAD_VEC_B + 0) * BN + loadc_b + l] = src1[idx].s0;

ggml/src/ggml-opencl/kernels/mul_mv_mxfp4_f32_flat.cl

@@ -132,7 +132,7 @@ kernel void kernel_mul_mv_mxfp4_f32_flat(
         global float4 * y4 = (global float4 *)yb;
 
         #pragma unroll
-        for (short row = 0; row < N_R0_MXFP4; row++) {
+        for (short row = 0; row < N_R0_MXFP4 && first_row + row < ne0; row++) {
             uchar xb_e = x_e[row * nb + ib];
 #ifdef SRC0Q_IMG
             ushort4 xb_q = as_ushort4(read_imageui(src0_q, (offset_q + row * nb + ib) * 2 + it).xy);

ggml/src/ggml-sycl/fattn-sparse.cpp

@@ -202,7 +202,8 @@ bool ggml_sycl_flash_attn_ext_sparse(ggml_backend_sycl_context & ctx, ggml_tenso
     uint8_t *    d_V    = v_alloc.get();
     sycl::half * d_mask = m_alloc.get();
 
-    SYCL_CHECK(CHECK_TRY_ERROR(stream->memset(d_cnt, 0, sizeof(int32_t))));
+    // reset with a kernel: on Level Zero a queue memset is not always ordered before the next kernel
+    stream->single_task([=]() { *d_cnt = 0; });
 
     sparse_fa_compact_mask(stream, (const sycl::half *) mask->data,
                            d_idx, d_cnt, n_kv, n_kv_g);

ggml/src/ggml-zdnn/common.hpp

@@ -6,6 +6,8 @@
 
 #include "zdnn.h"
 
+#include "scratch.hpp"
+
 #include <algorithm>
 #include <cstdlib>
 #include <future>
@@ -14,9 +16,8 @@
 #include <utility>
 #include <vector>
 
-#define GGML_ZDNN_NAME          "zDNN"
-#define GGML_ZDNN_VERSION       ZDNN_VERNUM
-#define GGML_ZDNN_SCRATCH_SLOTS 3  // int8 inputs, DFloat16 partial results, partial results for grouped matmul
+#define GGML_ZDNN_NAME    "zDNN"
+#define GGML_ZDNN_VERSION ZDNN_VERNUM
 
 #define ZDNN_CHECK(stmt)                \
     do {                                \
@@ -68,44 +69,100 @@ struct ggml_backend_zdnn_context {
     // zero bias for stacked quantized matmuls, keyed by (n_stacks, ne0)
     std::map<std::pair<int64_t, int64_t>, std::unique_ptr<ggml_backend_zdnn_buffer>> zero_bias;
 
-    std::vector<int8_t> scratch_q;
-    std::vector<float>  scratch_f32;
-
-    void * scratch_ztensor[GGML_ZDNN_SCRATCH_SLOTS]      = {};  // 4K aligned
-    size_t scratch_ztensor_size[GGML_ZDNN_SCRATCH_SLOTS] = {};
+    ggml_zdnn_pool pool;  // scratch memory of the ops
 };
 
-template <typename T>
-static inline T * ggml_zdnn_resize_scratch_max(std::vector<T> & scratch, int64_t n) {
-    if ((int64_t)scratch.size() < n) {
-        scratch.resize(n);
-    }
-
-    return scratch.data();
-}
-
 //
 // Scratch buffers
 //
 
-static inline void ggml_zdnn_init_scratch_ztensor(ggml_backend_zdnn_context * ctx,
-                                                  zdnn_tensor_desc          * pre_tfm_desc,
-                                                  zdnn_tensor_desc          * tfm_desc,
-                                                  zdnn_ztensor              * ztensor,
-                                                  int32_t                     slot) {
+// the ztensor uses the memory of scratch, so scratch must outlive it
+static inline void ggml_zdnn_init_scratch_ztensor(ggml_backend_zdnn_context     * ctx,
+                                                  zdnn_tensor_desc              * pre_tfm_desc,
+                                                  zdnn_tensor_desc              * tfm_desc,
+                                                  zdnn_ztensor                  * ztensor,
+                                                  ggml_zdnn_pool_alloc<uint8_t> & scratch) {
 
     const uint64_t size = zdnn_getsize_ztensor(tfm_desc);
 
-    if (ctx->scratch_ztensor_size[slot] < size) {
-        free(ctx->scratch_ztensor[slot]);
-        ctx->scratch_ztensor[slot]      = std::aligned_alloc(4096, GGML_PAD(size, 4096));
-        ctx->scratch_ztensor_size[slot] = GGML_PAD(size, 4096);
-        GGML_ASSERT(ctx->scratch_ztensor[slot] != nullptr);
+    zdnn_init_quantized_ztensor(pre_tfm_desc, tfm_desc, 1.0f, 0.0f, ztensor);
+    ztensor->buffer      = scratch.alloc(ctx->pool, size);
+    ztensor->buffer_size = size;
+}
+
+// same as ggml_zdnn_init_scratch_ztensor, for a DLFLOAT16 ztensor
+static inline void ggml_zdnn_init_scratch_ztensor_dlf16(ggml_backend_zdnn_context     * ctx,
+                                                        zdnn_tensor_desc              * pre_tfm_desc,
+                                                        zdnn_tensor_desc              * tfm_desc,
+                                                        zdnn_ztensor                  * ztensor,
+                                                        ggml_zdnn_pool_alloc<uint8_t> & scratch) {
+
+    const uint64_t size = zdnn_getsize_ztensor(tfm_desc);
+
+    zdnn_init_ztensor(pre_tfm_desc, tfm_desc, ztensor);
+    ztensor->buffer      = scratch.alloc(ctx->pool, size);
+    ztensor->buffer_size = size;
+}
+
+//
+// Stacked ztensors
+//
+
+// view of the stacks [s0, s0 + n) of a stacked ztensor, the stacks are the outermost dim of its buffer
+// the view descs must outlive the view
+static inline zdnn_ztensor ggml_zdnn_stack_view(zdnn_tensor_desc * pre_tfm_desc,
+                                                zdnn_tensor_desc * tfm_desc,
+                                          const zdnn_ztensor     & src,
+                                                int64_t            n_stacks,
+                                                int64_t            s0,
+                                                int64_t            n) {
+
+    GGML_ASSERT(src.buffer_size % n_stacks == 0);
+    const uint64_t stack_bytes = src.buffer_size / n_stacks;
+
+    *pre_tfm_desc = *src.pre_transformed_desc;
+    *tfm_desc     = *src.transformed_desc;
+
+    switch (pre_tfm_desc->layout) {
+        case ZDNN_3DS:
+            {
+                pre_tfm_desc->dim3 = n;
+            } break;
+        case ZDNN_2DS:
+            {
+                pre_tfm_desc->dim2 = n;
+            } break;
+        default:
+            GGML_ABORT("%s: unsupported stacked layout", __func__);
     }
 
-    zdnn_init_quantized_ztensor(pre_tfm_desc, tfm_desc, 1.0f, 0.0f, ztensor);
-    ztensor->buffer      = ctx->scratch_ztensor[slot];
-    ztensor->buffer_size = size;
+    tfm_desc->dim4 = n;
+    GGML_ASSERT(zdnn_getsize_ztensor(tfm_desc) == n*stack_bytes);
+
+    zdnn_ztensor view = src;
+    view.pre_transformed_desc = pre_tfm_desc;
+    view.transformed_desc     = tfm_desc;
+    view.buffer               = (char *) src.buffer + s0*stack_bytes;
+    view.buffer_size          = n * stack_bytes;
+    return view;
+}
+
+// zero bias of n_stacks stacks, created once per shape as it never changes
+static inline const zdnn_ztensor & ggml_zdnn_zero_bias(ggml_backend_zdnn_context * ctx, int64_t n_stacks, int64_t ne0) {
+    std::unique_ptr<ggml_backend_zdnn_buffer> & bias = ctx->zero_bias[{ n_stacks, ne0 }];
+
+    if (!bias) {
+        bias = std::make_unique<ggml_backend_zdnn_buffer>();
+
+        std::vector<float> zeros(n_stacks * ne0, 0.0f);
+
+        zdnn_init_pre_transformed_desc(ZDNN_2DS, FP32, &bias->pre_tfm_desc, n_stacks, ne0);
+        ZDNN_CHECK(zdnn_generate_transformed_desc(&bias->pre_tfm_desc, &bias->tfm_desc));
+        ZDNN_CHECK(zdnn_init_ztensor_with_malloc(&bias->pre_tfm_desc, &bias->tfm_desc, &bias->ztensor));
+        ZDNN_CHECK(zdnn_transform_ztensor(&bias->ztensor, zeros.data()));
+    }
+
+    return bias->ztensor;
 }
 
 //

ggml/src/ggml-zdnn/fattn.cpp

@@ -0,0 +1,255 @@
+#include "ggml.h"
+#include "fattn.hpp"
+
+#include <algorithm>
+#include <cstring>
+
+// DLFLOAT16 has no infinity, so masked scores use the lowest FP16 value, exp of which is still 0
+static inline float ggml_zdnn_fattn_mask_value(ggml_fp16_t m) {
+    return std::max(GGML_FP16_TO_FP32(m), -65504.0f);
+}
+
+bool ggml_zdnn_supports_flash_attn_ext(const ggml_backend_zdnn_device_context * ctx_dev, const ggml_tensor * op) {
+    const ggml_tensor * q    = op->src[0];
+    const ggml_tensor * k    = op->src[1];
+    const ggml_tensor * v    = op->src[2];
+    const ggml_tensor * mask = op->src[3];
+
+    float max_bias      = 0.0f;
+    float logit_softcap = 0.0f;
+    memcpy(&max_bias,      (const float *) op->op_params + 1, sizeof(float));
+    memcpy(&logit_softcap, (const float *) op->op_params + 2, sizeof(float));
+
+    // sinks, alibi, softcap and sparse k/v sets are not implemented
+    if (op->src[4] != nullptr || max_bias != 0.0f || logit_softcap != 0.0f || ggml_get_op_params_i32(op, 4) != 0) {
+        return false;
+    }
+
+    if (op->type != GGML_TYPE_F32 || q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16) {
+        return false;
+    }
+
+    if (mask != nullptr && (mask->type != GGML_TYPE_F16 || mask->ne[2] != 1 || mask->ne[3] != 1)) {
+        return false;
+    }
+
+    if (q->ne[3] != 1 || k->ne[3] != 1 || v->ne[3] != 1 || k->ne[2] != v->ne[2] || q->ne[2] % k->ne[2] != 0) {
+        return false;
+    }
+
+    // rows are gathered with their strides, but each row has to be contiguous
+    if (q->nb[0] != sizeof(float) || k->nb[0] != sizeof(ggml_fp16_t) || v->nb[0] != sizeof(ggml_fp16_t)) {
+        return false;
+    }
+
+    const int64_t max_dim = ctx_dev->max_size;
+    const int64_t group   = q->ne[2] / k->ne[2];
+
+    return q->ne[0] <= max_dim && v->ne[0] <= max_dim && k->ne[1] <= max_dim &&
+           q->ne[2] <= max_dim && group * q->ne[1] <= max_dim;
+}
+
+void ggml_zdnn_flash_attn_ext(ggml_backend_zdnn_context * ctx, ggml_tensor * dst) {
+    const ggml_tensor * q    = dst->src[0];
+    const ggml_tensor * k    = dst->src[1];
+    const ggml_tensor * v    = dst->src[2];
+    const ggml_tensor * mask = dst->src[3];
+
+    const int64_t d_k       = q->ne[0];
+    const int64_t n_q       = q->ne[1];
+    const int64_t n_head    = q->ne[2];
+    const int64_t n_kv      = k->ne[1];
+    const int64_t n_head_kv = k->ne[2];
+    const int64_t d_v       = v->ne[0];
+    const int64_t group     = n_head / n_head_kv;
+
+    float scale = 1.0f;
+    memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float));
+
+    const int32_t n_threads = n_q*n_kv >= 65536 ? ctx->n_threads : 1;
+
+    // one query runs all heads in one call with the mask row as the K*Q bias, more queries run per head and add the mask
+    const bool single = n_q == 1;
+
+    const int64_t n_mask = mask ? (single ? n_head_kv*n_kv : n_q*n_kv) : 0;
+
+    ggml_zdnn_pool_alloc<float>       q_host_alloc(ctx->pool, n_head*n_q*d_k);
+    ggml_zdnn_pool_alloc<float>       o_host_alloc(ctx->pool, n_head*n_q*d_v);
+    ggml_zdnn_pool_alloc<float>       mask_host_alloc(ctx->pool, n_mask);
+    ggml_zdnn_pool_alloc<ggml_fp16_t> k_host_alloc(ctx->pool, n_head_kv*n_kv*d_k);
+    ggml_zdnn_pool_alloc<ggml_fp16_t> v_host_alloc(ctx->pool, n_head_kv*n_kv*d_v);
+
+    float       * q_host    = q_host_alloc.get();
+    float       * o_host    = o_host_alloc.get();
+    float       * mask_host = mask_host_alloc.get();
+    ggml_fp16_t * k_host    = k_host_alloc.get();
+    ggml_fp16_t * v_host    = v_host_alloc.get();
+
+    // q is [n_head][n_q][d_k] with the scale folded in, k and v are [n_head_kv][n_kv][d]
+    ggml_zdnn_parallel_for(n_threads, n_head, [&](int64_t h_start, int64_t h_end) {
+        for (int64_t h = h_start; h < h_end; h++) {
+            for (int64_t iq = 0; iq < n_q; iq++) {
+                const float * src = (const float *)((const char *) q->data + iq*q->nb[1] + h*q->nb[2]);
+                float       * row = q_host + (h*n_q + iq)*d_k;
+                for (int64_t i0 = 0; i0 < d_k; i0++) {
+                    row[i0] = src[i0] * scale;
+                }
+            }
+        }
+    });
+
+    ggml_zdnn_parallel_for(n_threads, n_head_kv*n_kv, [&](int64_t r_start, int64_t r_end) {
+        for (int64_t r = r_start; r < r_end; r++) {
+            const int64_t hk = r / n_kv;
+            const int64_t ik = r % n_kv;
+            memcpy(k_host + r*d_k, (const char *) k->data + ik*k->nb[1] + hk*k->nb[2], d_k*sizeof(ggml_fp16_t));
+            memcpy(v_host + r*d_v, (const char *) v->data + ik*v->nb[1] + hk*v->nb[2], d_v*sizeof(ggml_fp16_t));
+        }
+    });
+
+    zdnn_tensor_desc q_pre_tfm_desc, q_tfm_desc;
+    zdnn_tensor_desc k_pre_tfm_desc, k_tfm_desc;
+    zdnn_tensor_desc v_pre_tfm_desc, v_tfm_desc;
+    zdnn_tensor_desc o_pre_tfm_desc, o_tfm_desc;
+    zdnn_ztensor     q_ztensor, k_ztensor, v_ztensor, o_ztensor;
+
+    ggml_zdnn_pool_alloc<uint8_t> q_ztensor_alloc, k_ztensor_alloc, v_ztensor_alloc, o_ztensor_alloc;
+
+    // a single query keeps the heads of a kv head in one stack, which is the same memory order as one head per stack
+    if (single) {
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &q_pre_tfm_desc, n_head_kv, group, d_k);
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &o_pre_tfm_desc, n_head_kv, group, d_v);
+    } else {
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &q_pre_tfm_desc, n_head, n_q, d_k);
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &o_pre_tfm_desc, n_head, n_q, d_v);
+    }
+    zdnn_init_pre_transformed_desc(ZDNN_3DS, FP16, &k_pre_tfm_desc, n_head_kv, n_kv, d_k);
+    zdnn_init_pre_transformed_desc(ZDNN_3DS, FP16, &v_pre_tfm_desc, n_head_kv, n_kv, d_v);
+
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&q_pre_tfm_desc, &q_tfm_desc));
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&k_pre_tfm_desc, &k_tfm_desc));
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&v_pre_tfm_desc, &v_tfm_desc));
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&o_pre_tfm_desc, &o_tfm_desc));
+
+    ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &q_pre_tfm_desc, &q_tfm_desc, &q_ztensor, q_ztensor_alloc);
+    ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &k_pre_tfm_desc, &k_tfm_desc, &k_ztensor, k_ztensor_alloc);
+    ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &v_pre_tfm_desc, &v_tfm_desc, &v_ztensor, v_ztensor_alloc);
+    ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &o_pre_tfm_desc, &o_tfm_desc, &o_ztensor, o_ztensor_alloc);
+
+    ZDNN_CHECK(zdnn_transform_ztensor(&q_ztensor, q_host));
+    ZDNN_CHECK(zdnn_transform_ztensor(&k_ztensor, k_host));
+    ZDNN_CHECK(zdnn_transform_ztensor(&v_ztensor, v_host));
+
+    ggml_zdnn_pool_alloc<uint8_t> softmax_save_area(ctx->pool, ZDNN_SOFTMAX_SAVEAREA_SIZE);
+
+    zdnn_tensor_desc s_pre_tfm_desc, s_tfm_desc;
+    zdnn_ztensor     s_ztensor, s2_ztensor;
+
+    ggml_zdnn_pool_alloc<uint8_t> s_ztensor_alloc, s2_ztensor_alloc;
+
+    if (single) {
+        // the bias of K*Q is the mask row of the query, repeated for every kv head
+        zdnn_tensor_desc bias_pre_tfm_desc, bias_tfm_desc;
+        zdnn_ztensor     bias_ztensor;
+
+        ggml_zdnn_pool_alloc<uint8_t> bias_ztensor_alloc;
+
+        const zdnn_ztensor * bias = &ggml_zdnn_zero_bias(ctx, n_head_kv, n_kv);
+        if (mask != nullptr) {
+            const ggml_fp16_t * mask_row = (const ggml_fp16_t *) mask->data;
+            for (int64_t hk = 0; hk < n_head_kv; hk++) {
+                for (int64_t ik = 0; ik < n_kv; ik++) {
+                    mask_host[hk*n_kv + ik] = ggml_zdnn_fattn_mask_value(mask_row[ik]);
+                }
+            }
+
+            zdnn_init_pre_transformed_desc(ZDNN_2DS, FP32, &bias_pre_tfm_desc, n_head_kv, n_kv);
+            ZDNN_CHECK(zdnn_generate_transformed_desc(&bias_pre_tfm_desc, &bias_tfm_desc));
+            ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &bias_pre_tfm_desc, &bias_tfm_desc, &bias_ztensor, bias_ztensor_alloc);
+            ZDNN_CHECK(zdnn_transform_ztensor(&bias_ztensor, mask_host));
+            bias = &bias_ztensor;
+        }
+
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &s_pre_tfm_desc, n_head_kv, group, n_kv);
+        ZDNN_CHECK(zdnn_generate_transformed_desc(&s_pre_tfm_desc, &s_tfm_desc));
+        ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &s_pre_tfm_desc, &s_tfm_desc, &s_ztensor,  s_ztensor_alloc);
+        ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &s_pre_tfm_desc, &s_tfm_desc, &s2_ztensor, s2_ztensor_alloc);
+
+        ZDNN_CHECK(zdnn_matmul_transpose_op(&q_ztensor, &k_ztensor, bias, false, true, MATMUL_OP_ADDITION, &s_ztensor));
+        ZDNN_CHECK(zdnn_softmax(&s_ztensor, softmax_save_area.get(), SOFTMAX_ACT_NONE, &s2_ztensor));
+        ZDNN_CHECK(zdnn_matmul_op(&s2_ztensor, &v_ztensor, &ggml_zdnn_zero_bias(ctx, n_head_kv, d_v), MATMUL_OP_ADDITION, &o_ztensor));
+    } else {
+        zdnn_tensor_desc mask_pre_tfm_desc, mask_tfm_desc;
+        zdnn_ztensor     mask_ztensor;
+
+        ggml_zdnn_pool_alloc<uint8_t> mask_ztensor_alloc;
+
+        if (mask != nullptr) {
+            ggml_zdnn_parallel_for(n_threads, n_q, [&](int64_t iq_start, int64_t iq_end) {
+                for (int64_t iq = iq_start; iq < iq_end; iq++) {
+                    const ggml_fp16_t * mask_row = (const ggml_fp16_t *)((const char *) mask->data + iq*mask->nb[1]);
+                    for (int64_t ik = 0; ik < n_kv; ik++) {
+                        mask_host[iq*n_kv + ik] = ggml_zdnn_fattn_mask_value(mask_row[ik]);
+                    }
+                }
+            });
+
+            zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &mask_pre_tfm_desc, 1, n_q, n_kv);
+            ZDNN_CHECK(zdnn_generate_transformed_desc(&mask_pre_tfm_desc, &mask_tfm_desc));
+            ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &mask_pre_tfm_desc, &mask_tfm_desc, &mask_ztensor, mask_ztensor_alloc);
+            ZDNN_CHECK(zdnn_transform_ztensor(&mask_ztensor, mask_host));
+        }
+
+        zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &s_pre_tfm_desc, 1, n_q, n_kv);
+        ZDNN_CHECK(zdnn_generate_transformed_desc(&s_pre_tfm_desc, &s_tfm_desc));
+        ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &s_pre_tfm_desc, &s_tfm_desc, &s_ztensor,  s_ztensor_alloc);
+        ggml_zdnn_init_scratch_ztensor_dlf16(ctx, &s_pre_tfm_desc, &s_tfm_desc, &s2_ztensor, s2_ztensor_alloc);
+
+        const zdnn_ztensor & k_bias = ggml_zdnn_zero_bias(ctx, 1, n_kv);
+        const zdnn_ztensor & v_bias = ggml_zdnn_zero_bias(ctx, 1, d_v);
+
+        zdnn_tensor_desc view_pre_tfm_desc[4];
+        zdnn_tensor_desc view_tfm_desc[4];
+
+        for (int64_t h = 0; h < n_head; h++) {
+            const zdnn_ztensor q_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[0], &view_tfm_desc[0], q_ztensor, n_head,    h,         1);
+            const zdnn_ztensor k_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[1], &view_tfm_desc[1], k_ztensor, n_head_kv, h / group, 1);
+            const zdnn_ztensor v_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[2], &view_tfm_desc[2], v_ztensor, n_head_kv, h / group, 1);
+            zdnn_ztensor       o_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[3], &view_tfm_desc[3], o_ztensor, n_head,    h,         1);
+
+            s_ztensor.is_transformed  = false;
+            s2_ztensor.is_transformed = false;
+
+            ZDNN_CHECK(zdnn_matmul_transpose_op(&q_view, &k_view, &k_bias, false, true, MATMUL_OP_ADDITION, &s_ztensor));
+
+            zdnn_ztensor * scores = &s_ztensor;
+            zdnn_ztensor * probs  = &s2_ztensor;
+            if (mask != nullptr) {
+                ZDNN_CHECK(zdnn_add(&s_ztensor, &mask_ztensor, &s2_ztensor));
+                scores = &s2_ztensor;
+                probs  = &s_ztensor;
+                probs->is_transformed = false;
+            }
+
+            ZDNN_CHECK(zdnn_softmax(scores, softmax_save_area.get(), SOFTMAX_ACT_NONE, probs));
+            ZDNN_CHECK(zdnn_matmul_op(probs, &v_view, &v_bias, MATMUL_OP_ADDITION, &o_view));
+        }
+
+        o_ztensor.is_transformed = true;
+    }
+
+    ZDNN_CHECK(zdnn_transform_origtensor(&o_ztensor, o_host));
+
+    // the result is [d_v][n_head][n_q]
+    for (int64_t iq = 0; iq < n_q; iq++) {
+        for (int64_t h = 0; h < n_head; h++) {
+            memcpy((char *) dst->data + h*dst->nb[1] + iq*dst->nb[2], o_host + (h*n_q + iq)*d_v, d_v*sizeof(float));
+        }
+    }
+
+    // the result is written to dst->data directly, so the dst ztensor is stale
+    ggml_backend_zdnn_buffer * dst_extra = (ggml_backend_zdnn_buffer *) dst->extra;
+    if (dst_extra != nullptr) {
+        dst_extra->ztensor.is_transformed = false;
+    }
+}

ggml/src/ggml-zdnn/fattn.hpp

@@ -0,0 +1,10 @@
+#ifndef GGML_ZDNN_FATTN_HPP
+#define GGML_ZDNN_FATTN_HPP
+
+#include "common.hpp"
+
+bool ggml_zdnn_supports_flash_attn_ext(const ggml_backend_zdnn_device_context * ctx_dev, const ggml_tensor * op);
+
+void ggml_zdnn_flash_attn_ext(ggml_backend_zdnn_context * ctx, ggml_tensor * dst);
+
+#endif  // GGML_ZDNN_FATTN_HPP

ggml/src/ggml-zdnn/ggml-zdnn.cpp

@@ -3,6 +3,7 @@
 #include "ggml-backend-impl.h"
 
 #include "ggml-zdnn/common.hpp"
+#include "ggml-zdnn/fattn.hpp"
 #include "ggml-zdnn/mmf.hpp"
 #include "ggml-zdnn/mmq.hpp"
 #include "ggml-zdnn/utils.hpp"
@@ -37,6 +38,11 @@ static bool ggml_zdnn_compute_forward(
                 ggml_zdnn_compute_forward_mul_mat(ctx, dst);
             } break;
 
+        case GGML_OP_FLASH_ATTN_EXT:
+            {
+                ggml_zdnn_flash_attn_ext(ctx, dst);
+            } break;
+
         default:
             return false;
     }
@@ -120,6 +126,9 @@ static bool ggml_zdnn_supports_op(const ggml_backend_zdnn_device_context * ctx_d
                 }
             } break;
 
+        case GGML_OP_FLASH_ATTN_EXT:
+            return ggml_zdnn_supports_flash_attn_ext(ctx_dev, op);
+
         default:
             return false;
     }
@@ -205,9 +214,6 @@ static void ggml_zdnn_free(ggml_backend_zdnn_context * ctx) {
     for (const auto & it : ctx->zero_bias) {
         ZDNN_CHECK(zdnn_free_ztensor_buffer(&it.second->ztensor));
     }
-    for (void * scratch : ctx->scratch_ztensor) {
-        free(scratch);
-    }
 
     delete ctx;
 }
@@ -309,7 +315,19 @@ static void ggml_backend_zdnn_buffer_memset_tensor(ggml_backend_buffer_t buffer,
 static void ggml_backend_zdnn_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
     memcpy((char *)tensor->data + offset, data, size);
 
+    // views share the ztensor of their source, which is now stale, and empty tensors have none
+    if (tensor->view_src != nullptr) {
+        ggml_backend_zdnn_buffer * src_extra = (ggml_backend_zdnn_buffer *)tensor->view_src->extra;
+        if (src_extra != nullptr && src_extra->ztensor.is_transformed) {
+            zdnn_reset_ztensor(&src_extra->ztensor);
+        }
+        return;
+    }
+
     ggml_backend_zdnn_buffer * extra = (ggml_backend_zdnn_buffer *)tensor->extra;
+    if (extra == nullptr) {
+        return;
+    }
 
     if (ggml_is_quantized(tensor->type)) {
         GGML_ASSERT(offset == 0 && size == ggml_nbytes(tensor));
@@ -319,8 +337,8 @@ static void ggml_backend_zdnn_buffer_set_tensor(ggml_backend_buffer_t buffer, gg
 
     // Fixes the LLAMA_SET_ROWS bug
     // see: https://github.com/ggml-org/llama.cpp/issues/15414
-    // compute tensors are transformed on first use instead, as not every op reads the ztensor
-    if (tensor->buffer->usage == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+    // only weights are transformed here, others on first use, as not every op reads the ztensor and some data (masks) is not finite
+    if (tensor->buffer->usage != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
         zdnn_reset_ztensor(&extra->ztensor);
         return;
     }

ggml/src/ggml-zdnn/mmq.cpp

@@ -6,63 +6,6 @@
 #include <future>
 #include <vector>
 
-// view of the stacks [s0, s0 + n) of a stacked ztensor, the stacks are the outermost dim of its buffer
-// the view descs must outlive the view
-static zdnn_ztensor ggml_zdnn_stack_view(zdnn_tensor_desc * pre_tfm_desc,
-                                         zdnn_tensor_desc * tfm_desc,
-                                   const zdnn_ztensor     & src,
-                                         int64_t            n_stacks,
-                                         int64_t            s0,
-                                         int64_t            n) {
-
-    GGML_ASSERT(src.buffer_size % n_stacks == 0);
-    const uint64_t stack_bytes = src.buffer_size / n_stacks;
-
-    *pre_tfm_desc = *src.pre_transformed_desc;
-    *tfm_desc     = *src.transformed_desc;
-
-    switch (pre_tfm_desc->layout) {
-        case ZDNN_3DS:
-            {
-                pre_tfm_desc->dim3 = n;
-            } break;
-        case ZDNN_2DS:
-            {
-                pre_tfm_desc->dim2 = n;
-            } break;
-        default:
-            GGML_ABORT("%s: unsupported stacked layout", __func__);
-    }
-
-    tfm_desc->dim4 = n;
-    GGML_ASSERT(zdnn_getsize_ztensor(tfm_desc) == n*stack_bytes);
-
-    zdnn_ztensor view = src;
-    view.pre_transformed_desc = pre_tfm_desc;
-    view.transformed_desc     = tfm_desc;
-    view.buffer               = (char *) src.buffer + s0*stack_bytes;
-    view.buffer_size          = n * stack_bytes;
-    return view;
-}
-
-// zero bias of n_stacks stacks, created once per shape as it never changes
-static const zdnn_ztensor & ggml_zdnn_zero_bias(ggml_backend_zdnn_context * ctx, int64_t n_stacks, int64_t ne0) {
-    std::unique_ptr<ggml_backend_zdnn_buffer> & bias = ctx->zero_bias[{ n_stacks, ne0 }];
-
-    if (!bias) {
-        bias = std::make_unique<ggml_backend_zdnn_buffer>();
-
-        std::vector<float> zeros(n_stacks * ne0, 0.0f);
-
-        zdnn_init_pre_transformed_desc(ZDNN_2DS, FP32, &bias->pre_tfm_desc, n_stacks, ne0);
-        ZDNN_CHECK(zdnn_generate_transformed_desc(&bias->pre_tfm_desc, &bias->tfm_desc));
-        ZDNN_CHECK(zdnn_init_ztensor_with_malloc(&bias->pre_tfm_desc, &bias->tfm_desc, &bias->ztensor));
-        ZDNN_CHECK(zdnn_transform_ztensor(&bias->ztensor, zeros.data()));
-    }
-
-    return bias->ztensor;
-}
-
 // the weights are stacked into groups along ne00, and each group of each row of the inputs and weights has its own scale
 // zDNN multiplies the groups in stacked calls, then the partial results of each group are scaled and summed here
 static void ggml_zdnn_mul_mat_q_grouped(ggml_backend_zdnn_context * ctx,
@@ -81,18 +24,20 @@ static void ggml_zdnn_mul_mat_q_grouped(ggml_backend_zdnn_context * ctx,
     GGML_ASSERT(n_groups * gs >= ne10 && (n_groups - 1) * gs < ne10);
 
     // the rounding errors of the inputs are stacked after the inputs, see ggml_zdnn_quantize_inputs
-    int8_t * inputs_q = ggml_zdnn_resize_scratch_max(ctx->scratch_q, 2 * n_groups * gs * ne11);
+    ggml_zdnn_pool_alloc<int8_t> inputs_q_alloc(ctx->pool, 2 * n_groups * gs * ne11);
+    int8_t * inputs_q = inputs_q_alloc.get();
     std::vector<float> inputs_scales(n_groups * ne11);
     ggml_zdnn_quantize_inputs(n_threads, inputs, gs, weights_extra->channel_scales.data(),
                             inputs_q, inputs_q + n_groups * gs * ne11, inputs_scales.data());
 
-    zdnn_tensor_desc inputs_pre_tfm_desc;
-    zdnn_tensor_desc inputs_tfm_desc;
-    zdnn_ztensor     inputs_ztensor;
+    zdnn_tensor_desc              inputs_pre_tfm_desc;
+    zdnn_tensor_desc              inputs_tfm_desc;
+    zdnn_ztensor                  inputs_ztensor;
+    ggml_zdnn_pool_alloc<uint8_t> inputs_ztensor_alloc;
 
     zdnn_init_pre_transformed_desc(ZDNN_3DS, INT8, &inputs_pre_tfm_desc, 2 * n_groups, ne11, gs);
     ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&inputs_pre_tfm_desc, QUANTIZED_INT8, &inputs_tfm_desc));
-    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, 0);
+    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, inputs_ztensor_alloc);
     ZDNN_CHECK(zdnn_transform_quantized_ztensor(&inputs_ztensor, false, -127, 127, inputs_q));
 
     // each group of the partial results is read back as its own 2D ztensor
@@ -116,19 +61,21 @@ static void ggml_zdnn_mul_mat_q_grouped(ggml_backend_zdnn_context * ctx,
 
     // two partial buffers, so the next chunk runs on the NNPA while this one is summed
     // each holds the results of a chunk of inputs, followed by the results of their rounding errors
-    zdnn_tensor_desc partial_pre_tfm_desc;
-    zdnn_tensor_desc partial_tfm_desc;
-    zdnn_ztensor     partial_ztensor[2];
+    zdnn_tensor_desc              partial_pre_tfm_desc;
+    zdnn_tensor_desc              partial_tfm_desc;
+    zdnn_ztensor                  partial_ztensor[2];
+    ggml_zdnn_pool_alloc<uint8_t> partial_ztensor_alloc[2];
 
     zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &partial_pre_tfm_desc, 2 * n_chunk, ne11, ne0);
     ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&partial_pre_tfm_desc, QUANTIZED_DLFLOAT16, &partial_tfm_desc));
     GGML_ASSERT(zdnn_getsize_ztensor(&partial_tfm_desc) == (uint64_t)(2 * n_chunk) * group_bytes);
     for (int32_t i = 0; i < 2; i++) {
-        ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor[i], 1 + i);
+        ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor[i], partial_ztensor_alloc[i]);
     }
 
     const float * weights_scales = weights_extra->scales.data();
-    float * partial     = ggml_zdnn_resize_scratch_max(ctx->scratch_f32, 2 * ne11 * ne0);
+    ggml_zdnn_pool_alloc<float> partial_alloc(ctx->pool, 2 * ne11 * ne0);
+    float * partial     = partial_alloc.get();
     float * partial_res = partial + ne11 * ne0;
 
     const auto sum_chunk = [&](const zdnn_ztensor & chunk, int64_t g0, int64_t n) {
@@ -262,27 +209,30 @@ void ggml_zdnn_mul_mat_q(
 
     // zDNN takes one scale per zTensor, so quantize each row of the inputs with its own scale
     // the rounding errors of the inputs are stacked after the inputs, see ggml_zdnn_quantize_inputs
-    int8_t * inputs_q = ggml_zdnn_resize_scratch_max(ctx->scratch_q, 2 * ne10 * ne11);
+    ggml_zdnn_pool_alloc<int8_t> inputs_q_alloc(ctx->pool, 2 * ne10 * ne11);
+    int8_t * inputs_q = inputs_q_alloc.get();
     std::vector<float> inputs_scales(ne11);
     ggml_zdnn_quantize_inputs(n_threads, inputs, ne10, weights_extra->channel_scales.data(),
                             inputs_q, inputs_q + ne10 * ne11, inputs_scales.data());
 
-    zdnn_tensor_desc inputs_pre_tfm_desc;
-    zdnn_tensor_desc inputs_tfm_desc;
-    zdnn_ztensor     inputs_ztensor;
+    zdnn_tensor_desc              inputs_pre_tfm_desc;
+    zdnn_tensor_desc              inputs_tfm_desc;
+    zdnn_ztensor                  inputs_ztensor;
+    ggml_zdnn_pool_alloc<uint8_t> inputs_ztensor_alloc;
 
     zdnn_init_pre_transformed_desc(ZDNN_3DS, INT8, &inputs_pre_tfm_desc, 2, ne11, ne10);
     ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&inputs_pre_tfm_desc, QUANTIZED_INT8, &inputs_tfm_desc));
-    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, 0);
+    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, inputs_ztensor_alloc);
     ZDNN_CHECK(zdnn_transform_quantized_ztensor(&inputs_ztensor, false, -127, 127, inputs_q));
 
-    zdnn_tensor_desc partial_pre_tfm_desc;
-    zdnn_tensor_desc partial_tfm_desc;
-    zdnn_ztensor     partial_ztensor;
+    zdnn_tensor_desc              partial_pre_tfm_desc;
+    zdnn_tensor_desc              partial_tfm_desc;
+    zdnn_ztensor                  partial_ztensor;
+    ggml_zdnn_pool_alloc<uint8_t> partial_ztensor_alloc;
 
     zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &partial_pre_tfm_desc, 2, ne11, ne0);
     ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&partial_pre_tfm_desc, QUANTIZED_DLFLOAT16, &partial_tfm_desc));
-    ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor, 1);
+    ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor, partial_ztensor_alloc);
 
     // both stacks are multiplied by the same unstacked weights in one broadcast call
     // the bias is all zeros, so it is already the pre-computed bias
@@ -300,7 +250,8 @@ void ggml_zdnn_mul_mat_q(
     const uint64_t stack_bytes = zdnn_getsize_ztensor(&stack_tfm_desc);
     GGML_ASSERT(partial_ztensor.buffer_size == 2*stack_bytes);
 
-    float * partial[2] = { (float *)output->data, ggml_zdnn_resize_scratch_max(ctx->scratch_f32, ne11 * ne0) };
+    ggml_zdnn_pool_alloc<float> partial_res_alloc(ctx->pool, ne11 * ne0);
+    float * partial[2] = { (float *)output->data, partial_res_alloc.get() };
 
     ggml_zdnn_parallel_for(n_threads, 2, [&](int64_t s_start, int64_t s_end) {
         zdnn_ztensor stack_ztensor;

ggml/src/ggml-zdnn/scratch.cpp

@@ -0,0 +1,45 @@
+#include "ggml.h"
+#include "scratch.hpp"
+
+#include <algorithm>
+#include <cstdint>
+#include <cstdlib>
+
+ggml_zdnn_pool::~ggml_zdnn_pool() {
+    for (const buffer & b : cached) {
+        std::free(b.ptr);
+    }
+}
+
+void * ggml_zdnn_pool::alloc(size_t size, size_t * actual_size) {
+    int64_t ibest = -1;
+    for (int64_t i = 0; i < (int64_t)cached.size(); i++) {
+        if (cached[i].size >= size && (ibest < 0 || cached[i].size < cached[ibest].size)) {
+            ibest = i;
+        }
+    }
+
+    if (ibest >= 0) {
+        void * ptr   = cached[ibest].ptr;
+        *actual_size = cached[ibest].size;
+        cached.erase(cached.begin() + ibest);
+        return ptr;
+    }
+
+    // every cached buffer is too small, free the largest so the pool does not grow while sizes grow with the context
+    if (!cached.empty()) {
+        const auto largest = std::max_element(cached.begin(), cached.end(), [](const buffer & a, const buffer & b) { return a.size < b.size; });
+        std::free(largest->ptr);
+        cached.erase(largest);
+    }
+
+    // 5% more, so slowly growing sizes do not allocate every time
+    *actual_size = GGML_PAD(std::max<size_t>(1, (size_t)(1.05 * size)), 4096);
+    void * ptr   = std::aligned_alloc(4096, *actual_size);
+    GGML_ASSERT(ptr != nullptr);
+    return ptr;
+}
+
+void ggml_zdnn_pool::free(void * ptr, size_t size) {
+    cached.push_back({ ptr, size });
+}

ggml/src/ggml-zdnn/scratch.hpp

@@ -0,0 +1,61 @@
+#ifndef GGML_ZDNN_SCRATCH_HPP
+#define GGML_ZDNN_SCRATCH_HPP
+
+#include "ggml.h"
+
+#include <cstddef>
+#include <vector>
+
+// 4K aligned host buffers that are cached for reuse, see ggml_cuda_pool_leg
+struct ggml_zdnn_pool {
+    struct buffer {
+        void * ptr  = nullptr;
+        size_t size = 0;
+    };
+
+    std::vector<buffer> cached;
+
+    ~ggml_zdnn_pool();
+
+    void * alloc(size_t size, size_t * actual_size);
+    void   free(void * ptr, size_t size);
+};
+
+// a buffer of the pool that goes back to it at the end of its scope, see ggml_cuda_pool_alloc
+template <typename T>
+struct ggml_zdnn_pool_alloc {
+    ggml_zdnn_pool * pool = nullptr;
+    T * ptr = nullptr;
+    size_t actual_size = 0;
+
+    ggml_zdnn_pool_alloc() = default;
+
+    ggml_zdnn_pool_alloc(ggml_zdnn_pool & pool, size_t size) {
+        alloc(pool, size);
+    }
+
+    ~ggml_zdnn_pool_alloc() {
+        if (ptr != nullptr) {
+            pool->free(ptr, actual_size);
+        }
+    }
+
+    // size is in number of elements
+    T * alloc(ggml_zdnn_pool & pool, size_t size) {
+        GGML_ASSERT(ptr == nullptr);
+        this->pool = &pool;
+        ptr = (T *) pool.alloc(size * sizeof(T), &actual_size);
+        return ptr;
+    }
+
+    T * get() {
+        return ptr;
+    }
+
+    ggml_zdnn_pool_alloc(const ggml_zdnn_pool_alloc &) = delete;
+    ggml_zdnn_pool_alloc(ggml_zdnn_pool_alloc &&) = delete;
+    ggml_zdnn_pool_alloc & operator=(const ggml_zdnn_pool_alloc &) = delete;
+    ggml_zdnn_pool_alloc & operator=(ggml_zdnn_pool_alloc &&) = delete;
+};
+
+#endif  // GGML_ZDNN_SCRATCH_HPP

gguf-py/gguf/constants.py

@@ -419,6 +419,8 @@ class Keys:
         PREPROC_MIN_TILES     = "clip.vision.preproc_min_tiles"
         PREPROC_MAX_TILES     = "clip.vision.preproc_max_tiles"
         PREPROC_IMAGE_SIZE    = "clip.vision.preproc_image_size"
+        PREPROC_TILE_SIZE     = "clip.vision.preproc_tile_size"
+        PREPROC_MAX_PIXELS_TOLERANCE = "clip.vision.preproc_max_pixels_tolerance"
         PATCH_SIZE            = "clip.vision.patch_size"
         EMBEDDING_LENGTH      = "clip.vision.embedding_length"
         FEED_FORWARD_LENGTH   = "clip.vision.feed_forward_length"

gguf-py/gguf/gguf_writer.py

@@ -1470,6 +1470,12 @@ class GGUFWriter:
     def add_vision_preproc_image_size(self, value: int) -> None:
         self.add_uint32(Keys.ClipVision.PREPROC_IMAGE_SIZE, value)
 
+    def add_vision_preproc_tile_size(self, value: int) -> None:
+        self.add_uint32(Keys.ClipVision.PREPROC_TILE_SIZE, value)
+
+    def add_vision_preproc_max_pixels_tolerance(self, value: float) -> None:
+        self.add_float32(Keys.ClipVision.PREPROC_MAX_PIXELS_TOLERANCE, value)
+
     def add_vision_projector_query_side(self, value: int) -> None:
         self.add_uint32(Keys.ClipVision.Projector.QUERY_SIDE, value)
 

scripts/check-apiabi-compat.sh

@@ -244,8 +244,10 @@ elif [ "$DO_CHECK" -eq 1 ]; then
     # From here on, we don't want to exit on first error
     set +e
 
+    ACC_STRICT=
+    [ "$STRICT" -eq 0 ] || ACC_STRICT="-strict"
     abi-compliance-checker \
-        -strict \
+        $ACC_STRICT \
         -source \
         -library "$PROJECT_NAME" \
         -old "$BUILD_DIR/apiabi/api.dump" \

scripts/snapdragon/build.py

@@ -11,6 +11,9 @@ import platform
 import shutil
 import logging
 import glob
+import shlex
+import tarfile
+import tempfile
 
 from sdk import validate_windows_sdks
 
@@ -19,19 +22,70 @@ logger = logging.getLogger("build")
 
 def parse_target(target_str):
     if not target_str:
-        return None, None
-    if target_str.startswith("adb") or target_str.startswith("android"):
-        parts = target_str.split(":", 1)
-        serial = parts[1] if len(parts) > 1 else None
-        return "android", serial
-    elif target_str.startswith("lnx") or target_str.startswith("linux") or target_str.startswith("ubuntu"):
-        parts = target_str.split(":", 1)
-        host = parts[1] if len(parts) > 1 else None
-        return "linux", host
+        return None, None, None
+
+    parts = target_str.split(":")
+    target_name = parts[0]
+
+    if target_name in ("adb", "android"):
+        serial = ":".join(parts[1:]) if len(parts) > 1 else None
+        return "android", serial, "adb"
+    elif target_name in ("lnx", "linux", "ubuntu"):
+        if len(parts) > 1 and parts[1] == "adb":
+            serial = ":".join(parts[2:]) if len(parts) > 2 else None
+            return "linux", serial, "adb"
+        host = ":".join(parts[1:]) if len(parts) > 1 else None
+        return "linux", host, "ssh"
     elif target_str in ("wos", "windows"):
-        return "windows", None
+        return "windows", None, None
     else:
-        return None, None
+        return None, None, None
+
+
+def default_target_dir(target_type, target_transport, home_dir):
+    if target_type == "android":
+        return "/data/local/tmp/llama.cpp"
+    if target_type == "linux":
+        if home_dir:
+            return home_dir.rstrip("/") + "/llama.cpp"
+        if target_transport == "adb":
+            return "/home/llama.cpp"
+        return "~/llama.cpp"
+    return "~/llama.cpp"
+
+
+def normalize_target_dir(target_dir):
+    target_dir = target_dir.rstrip("/")
+    if not target_dir or target_dir == "/":
+        raise RuntimeError("target directory must not be the filesystem root")
+    return target_dir
+
+
+def adb_root(adb_cmd):
+    root_cmd = adb_cmd + ["root"]
+    logger.info(f"+ {' '.join(root_cmd)}")
+    res = subprocess.run(root_cmd)
+    if res.returncode != 0:
+        logger.error("ADB root failed.")
+        sys.exit(res.returncode)
+
+    wait_cmd = adb_cmd + ["wait-for-device"]
+    logger.info(f"+ {' '.join(wait_cmd)}")
+    res = subprocess.run(wait_cmd)
+    if res.returncode != 0:
+        logger.error("ADB wait-for-device failed.")
+        sys.exit(res.returncode)
+
+
+def make_package_archive(src_path):
+    fd, archive_path = tempfile.mkstemp(prefix="llama-cpp-", suffix=".tar.gz")
+    os.close(fd)
+    with tarfile.open(archive_path, "w:gz", dereference=False) as tar:
+        for item in os.listdir(src_path):
+            if item.startswith("."):
+                continue
+            tar.add(os.path.join(src_path, item), arcname=item, recursive=True)
+    return archive_path
 
 
 def get_uid_gid():
@@ -45,7 +99,7 @@ def main():
     parser = argparse.ArgumentParser(
         description="Build llama.cpp for Snapdragon using cross-compilation docker containers or natively."
     )
-    parser.add_argument("--target", default="android", help="Compilation target and deployment definition (e.g. android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, windows/wos) (default: android)")
+    parser.add_argument("--target", default="android", help="Compilation target and deployment definition (e.g. android[:serial]/adb[:serial], ubuntu:adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, windows/wos) (default: android)")
     parser.add_argument("--build-dir", help="Build directory name (defaults to build-TARGET[-dbg], e.g. build-android)")
     parser.add_argument("--install-dir", help="Install directory name (defaults to pkg-TARGET[-dbg], e.g. pkg-android)")
     parser.add_argument("--jobs", "-j", type=int, help="Number of build jobs (defaults to CPU thread count)")
@@ -53,21 +107,22 @@ def main():
     parser.add_argument("--preset", help="Override the CMake preset to use")
     parser.add_argument("--debug", action="store_true", help="Build in debug mode (uses -debug presets instead of -release)")
     parser.add_argument("--hex-only", "--hexagon-only", action="store_true", help="Only build and push ggml-hexagon and libggml-htp libraries")
-    parser.add_argument("--hex-arch", choices=["v73", "v75", "v79", "v81", "73", "75", "79", "81"], help="Specific Hexagon DSP architecture to build with --hex-only (default: all)")
+    parser.add_argument("--hex-arch", choices=["v73", "v75", "v79", "v81", "v85", "73", "75", "79", "81", "85"], help="Specific Hexagon DSP architecture to build with --hex-only (default: all)")
 
     # Push options
     parser.add_argument("--push", action="store_true", help="Push built package to the target device via ADB or SSH/SCP")
-    parser.add_argument("--target-dir", help="Target directory on the device (default: /data/local/tmp/llama.cpp for Android, ~/llama.cpp for Linux)")
+    parser.add_argument("--target-dir", help="Target directory on the device (default: /data/local/tmp/llama.cpp for Android, ~/llama.cpp for Linux, /home/llama.cpp for Linux via ADB)")
+    parser.add_argument("--home-dir", help="Home directory for Linux targets when deriving the default target directory")
 
     # Toolchain options
-    parser.add_argument("--toolchain-version", default="v0.7", help="Docker toolchain image version/tag (default: v0.7)")
+    parser.add_argument("--toolchain-version", default="v0.8", help="Docker toolchain image version/tag (default: v0.8)")
     parser.add_argument("--toolchain-url", default="ghcr.io/snapdragon-toolchain", help="Docker toolchain registry URL/namespace (default: ghcr.io/snapdragon-toolchain)")
 
     args = parser.parse_args()
 
-    target_type, target_val = parse_target(args.target)
+    target_type, target_val, target_transport = parse_target(args.target)
     if not target_type:
-        logger.error(f"Error: Invalid target format '{args.target}'. Must be android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, or windows/wos.")
+        logger.error(f"Error: Invalid target format '{args.target}'. Must be android[:serial]/adb[:serial], ubuntu:adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, or windows/wos.")
         sys.exit(1)
 
     if target_type == "windows":
@@ -140,7 +195,7 @@ def main():
         if hex_arch:
             hex_targets = ["ggml-hexagon", f"htp-{hex_arch}"]
         else:
-            hex_targets = ["ggml-hexagon", "htp-v73", "htp-v75", "htp-v79", "htp-v81"]
+            hex_targets = ["ggml-hexagon", "htp-v73", "htp-v75", "htp-v79", "htp-v81", "htp-v85"]
 
     if args.no_docker:
         # Native/local host build
@@ -218,8 +273,8 @@ def main():
         # Resolve target directory on device
         target_dir = args.target_dir
         if not target_dir:
-            target_dir = "/data/local/tmp/llama.cpp" if target_type == "android" else "~/llama.cpp"
-        target_dir = target_dir.rstrip("/")
+            target_dir = default_target_dir(target_type, target_transport, args.home_dir)
+        target_dir = normalize_target_dir(target_dir)
 
         sub_items = [item for item in os.listdir(src_path) if not item.startswith(".")]
 
@@ -235,14 +290,23 @@ def main():
                 logger.error(f"Error: no hexagon libraries found in {lib_dir} to push.")
                 sys.exit(1)
 
-        if target_type == "android":
+        if target_transport == "adb":
             adb_cmd = ["adb"]
             if target_val: # serial
                 adb_cmd += ["-s", target_val]
 
+            if target_type == "linux":
+                adb_root(adb_cmd)
+
             if args.hex_only:
-                logger.info("\nPushing Hexagon libraries to Android device via ADB...")
+                logger.info("\nPushing Hexagon libraries to target device via ADB...")
                 target_lib_dir = f"{target_dir}/lib"
+                mkdir_cmd = adb_cmd + ["shell", f"mkdir -p {shlex.quote(target_lib_dir)}"]
+                logger.info(f"+ {' '.join(mkdir_cmd)}")
+                res = subprocess.run(mkdir_cmd)
+                if res.returncode != 0:
+                    logger.error("ADB mkdir failed.")
+                    sys.exit(res.returncode)
                 for lib in hex_libs:
                     push_cmd = adb_cmd + ["push", lib, f"{target_lib_dir}/"]
                     logger.info(f"+ {' '.join(push_cmd)}")
@@ -251,6 +315,44 @@ def main():
                         logger.error("ADB push failed.")
                         sys.exit(res.returncode)
                 logger.info("Hexagon libraries deployed successfully!")
+            elif target_type == "linux":
+                logger.info("\nDeploying Linux package to target device via ADB tar...")
+                remote_archive = f"/tmp/llama.cpp-package-{os.getpid()}.tar.gz"
+                archive_path = make_package_archive(src_path)
+                try:
+                    mkdir_cmd = adb_cmd + ["shell", f"mkdir -p {shlex.quote(target_dir)}"]
+                    logger.info(f"+ {' '.join(mkdir_cmd)}")
+                    res = subprocess.run(mkdir_cmd)
+                    if res.returncode != 0:
+                        logger.error("ADB mkdir failed.")
+                        sys.exit(res.returncode)
+
+                    if sub_items:
+                        clean_paths = " ".join(shlex.quote(f"{target_dir}/{item}") for item in sub_items)
+                        clean_cmd = adb_cmd + ["shell", f"rm -rf {clean_paths}"]
+                        logger.info(f"+ {' '.join(clean_cmd)}")
+                        subprocess.run(clean_cmd)
+
+                    push_cmd = adb_cmd + ["push", archive_path, remote_archive]
+                    logger.info(f"+ {' '.join(push_cmd)}")
+                    res = subprocess.run(push_cmd)
+                    if res.returncode != 0:
+                        logger.error("ADB push failed.")
+                        sys.exit(res.returncode)
+
+                    extract_cmd = adb_cmd + ["shell", f"tar -xzf {shlex.quote(remote_archive)} -C {shlex.quote(target_dir)} && rm -f {shlex.quote(remote_archive)}"]
+                    logger.info(f"+ {' '.join(extract_cmd)}")
+                    res = subprocess.run(extract_cmd)
+                    if res.returncode != 0:
+                        logger.error("ADB untar failed.")
+                        sys.exit(res.returncode)
+
+                    chmod_cmd = adb_cmd + ["shell", f"chmod -R 755 {shlex.quote(target_dir)}/bin 2>/dev/null || true"]
+                    logger.info(f"+ {' '.join(chmod_cmd)}")
+                    subprocess.run(chmod_cmd)
+                finally:
+                    os.unlink(archive_path)
+                logger.info("ADB tar deploy completed successfully!")
             else:
                 # Clean stale package files on device
                 if sub_items:

scripts/snapdragon/ggml-hexagon-inspect.py

@@ -255,7 +255,7 @@ class HexagonToolchain:
         use_docker: bool = False,
         image_url: str = "ghcr.io/snapdragon-toolchain",
         image_name: str = "arm64-android",
-        image_ver: str = "v0.7",
+        image_ver: str = "v0.8",
     ):
         self.repo_root = repo_root
         self.image = f"{image_url}/{image_name}:{image_ver}"
@@ -322,15 +322,14 @@ class HexagonToolchain:
         if not self.docker_bin:
             raise RuntimeError("Docker is required but not installed or found on PATH.")
 
-        container_tools_dir = "/opt/hexagon/6.6.0.0/tools/HEXAGON_Tools/19.0.07/Tools/bin"
         if "objdump" in tool_name:
-            container_tool = f"{container_tools_dir}/hexagon-llvm-objdump"
+            tool_bin = "hexagon-llvm-objdump"
         elif "addr2line" in tool_name:
-            container_tool = f"{container_tools_dir}/hexagon-addr2line"
+            tool_bin = "hexagon-addr2line"
         elif "nm" in tool_name:
-            container_tool = f"{container_tools_dir}/llvm-nm"
+            tool_bin = "llvm-nm"
         else:
-            container_tool = f"{container_tools_dir}/{tool_name}"
+            tool_bin = tool_name
 
         # Translate file paths from host to /workspace
         translated_args = [translate_container_arg(arg, self.repo_root) for arg in args]
@@ -350,7 +349,9 @@ class HexagonToolchain:
         if platform.system() != "Windows":
             docker_cmd += ["-u", f"{os.getuid()}:{os.getgid()}"]
 
-        docker_cmd += [self.image, container_tool] + translated_args
+        default_tools_root = "/opt/hexagon/6.7.0.0/tools/HEXAGON_Tools/21.0.04"
+        sh_cmd = f'exec "${{HEXAGON_TOOLS_ROOT:-{default_tools_root}}}/Tools/bin/{tool_bin}" "$@"'
+        docker_cmd += [self.image, "sh", "-c", sh_cmd, "--"] + translated_args
 
         res = subprocess.run(docker_cmd, capture_output=True, text=True, input=stdin_data)
         if res.returncode != 0:
@@ -1457,8 +1458,8 @@ def main():
     )
     parser.add_argument(
         "--toolchain-version",
-        default="v0.7",
-        help="Docker toolchain tag (default: v0.7).",
+        default="v0.8",
+        help="Docker toolchain tag (default: v0.8).",
     )
     parser.add_argument(
         "--toolchain-url",

scripts/snapdragon/run.py

@@ -49,27 +49,46 @@ MANAGED_ENV_NAMES = (
     "GGML_OPENCL_ADRENO_XMEM_GEMM",
     "GGML_OPENCL_ADRENO_USE_LARGE_BUFFER",
     "GGML_SCHED_DEBUG",
+    "ADSP_LIBRARY_PATH",
+    "CDSP_LIBRARY_PATH",
+    "CDSP1_LIBRARY_PATH",
+    "CDSP2_LIBRARY_PATH",
+    "CDSP3_LIBRARY_PATH",
     "MTMD_BACKEND_DEVICE",
     "D",
     "DEVICE",
 )
 
 
+DSP_LIBRARY_ENV_NAMES = (
+    "ADSP_LIBRARY_PATH",
+    "CDSP_LIBRARY_PATH",
+    "CDSP1_LIBRARY_PATH",
+    "CDSP2_LIBRARY_PATH",
+    "CDSP3_LIBRARY_PATH",
+)
+
+
 def parse_target(target_str):
     if not target_str:
-        return None, None
-    if target_str.startswith("adb") or target_str.startswith("android"):
-        parts = target_str.split(":", 1)
-        serial = parts[1] if len(parts) > 1 else None
-        return "android", serial
-    elif target_str.startswith("lnx") or target_str.startswith("linux") or target_str.startswith("ubuntu"):
-        parts = target_str.split(":", 1)
-        host = parts[1] if len(parts) > 1 else None
-        return "linux", host
+        return None, None, None
+
+    parts = target_str.split(":")
+    target_name = parts[0]
+
+    if target_name in ("adb", "android"):
+        serial = ":".join(parts[1:]) if len(parts) > 1 else None
+        return "android", serial, "adb"
+    elif target_name in ("lnx", "linux", "ubuntu"):
+        if len(parts) > 1 and parts[1] == "adb":
+            serial = ":".join(parts[2:]) if len(parts) > 2 else None
+            return "linux", serial, "adb"
+        host = ":".join(parts[1:]) if len(parts) > 1 else None
+        return "linux", host, "ssh"
     elif target_str in ("wos", "windows"):
-        return "windows", None
+        return "windows", None, None
     else:
-        return None, None
+        return None, None, None
 
 
 def shlex_join(args_list):
@@ -79,6 +98,41 @@ def shlex_join(args_list):
     return " ".join(pipes.quote(x) for x in args_list)
 
 
+def default_target_dir(target_type, target_transport, home_dir):
+    if target_type == "android":
+        return "/data/local/tmp/llama.cpp"
+    if target_type == "linux":
+        if home_dir:
+            return home_dir.rstrip("/") + "/llama.cpp"
+        if target_transport == "adb":
+            return "/home/llama.cpp"
+        return "~/llama.cpp"
+    return "~/llama.cpp"
+
+
+def library_env_parts(lib_dir):
+    env_parts = [f"LD_LIBRARY_PATH={lib_dir}"]
+    for name in DSP_LIBRARY_ENV_NAMES:
+        env_parts.append(f"{name}={lib_dir}")
+    return env_parts
+
+
+def adb_root(adb_base):
+    root_cmd = adb_base + ["root"]
+    logger.info(f"+ {' '.join(root_cmd)}")
+    res = subprocess.run(root_cmd)
+    if res.returncode != 0:
+        logger.error("ADB root failed.")
+        sys.exit(res.returncode)
+
+    wait_cmd = adb_base + ["wait-for-device"]
+    logger.info(f"+ {' '.join(wait_cmd)}")
+    res = subprocess.run(wait_cmd)
+    if res.returncode != 0:
+        logger.error("ADB wait-for-device failed.")
+        sys.exit(res.returncode)
+
+
 def split_device_list(devices):
     parts = []
     curr = []
@@ -144,8 +198,9 @@ def main():
     parser = argparse.ArgumentParser(
         description="Unified runner for llama.cpp tools on Snapdragon (natively, via ADB, or via SSH)."
     )
-    parser.add_argument("--target", help="Execution target (e.g. android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, windows/wos) (default: local run)")
-    parser.add_argument("--target-dir", help="Target directory on the device (default: /data/local/tmp/llama.cpp for Android, ~/llama.cpp for Linux)")
+    parser.add_argument("--target", help="Execution target (e.g. android[:serial]/adb[:serial], ubuntu:adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, windows/wos) (default: local run)")
+    parser.add_argument("--target-dir", help="Target directory on the device (default: /data/local/tmp/llama.cpp for Android, ~/llama.cpp for Linux, /home/llama.cpp for Linux via ADB)")
+    parser.add_argument("--home-dir", help="Home directory for Linux targets when deriving the default target directory")
     parser.add_argument("--install-dir", help="Install directory name (defaults to pkg-TARGET or pkg-TARGET-dbg prefix based on target)")
     parser.add_argument("--debug", action="store_true", help="Use debug build (defaults to pkg-TARGET-dbg folder)")
     parser.add_argument("--devices", "--device", "-d", help="Select execution devices (split into NPU and OpenCL GPUs automatically, default: HTP0)")
@@ -201,11 +256,12 @@ def main():
 
     target_type = None
     target_val = None
+    target_transport = None
     target_prefix = None
     if args.target:
-        target_type, target_val = parse_target(args.target)
+        target_type, target_val, target_transport = parse_target(args.target)
         if not target_type:
-            logger.error(f"Error: Invalid target format '{args.target}'. Must be android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, or windows/wos.")
+            logger.error(f"Error: Invalid target format '{args.target}'. Must be android[:serial]/adb[:serial], ubuntu:adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, or windows/wos.")
             sys.exit(1)
         target_prefix = args.target.split(":", 1)[0]
 
@@ -435,18 +491,19 @@ def main():
     # Resolve target directory on device
     target_dir = args.target_dir
     if not target_dir:
-        target_dir = "/data/local/tmp/llama.cpp" if target_type == "android" else "~/llama.cpp"
+        target_dir = default_target_dir(target_type, target_transport, args.home_dir)
 
-    if target_type == "android":
-        # Run via ADB
+    if target_transport == "adb":
         adb_base = ["adb"]
         if target_val: # serial
             adb_base += ["-s", target_val]
 
-        env_parts = [
-            "LD_LIBRARY_PATH=./lib",
-            "ADSP_LIBRARY_PATH=./lib"
-        ]
+        if target_type == "linux":
+            adb_root(adb_base)
+
+        env_parts = library_env_parts("./lib")
+        if target_type == "linux":
+            env_parts.insert(0, f"HOME={args.home_dir or '/home'}")
         for k, v in env_vars.items():
             env_parts.append(f"{k}={v}")
         env_str = " ".join(env_parts)
@@ -466,10 +523,7 @@ def main():
             sys.exit(1)
 
         # Linux remote run via SSH
-        env_parts = [
-            "LD_LIBRARY_PATH=./lib",
-            "ADSP_LIBRARY_PATH=./lib"
-        ]
+        env_parts = library_env_parts("./lib")
         for k, v in env_vars.items():
             env_parts.append(f"{k}={v}")
         env_str = " ".join(env_parts)
@@ -490,7 +544,6 @@ def main():
         # Run locally
         local_env = os.environ.copy()
         lib_dir = os.path.normpath(os.path.join(package_path, "lib"))
-        local_env["ADSP_LIBRARY_PATH"] = lib_dir
         if platform.system() == "Windows":
             local_env["PATH"] = lib_dir + os.path.pathsep + local_env.get("PATH", "")
         else:
@@ -498,6 +551,8 @@ def main():
 
         for k in MANAGED_ENV_NAMES:
             local_env.pop(k, None)
+        for name in DSP_LIBRARY_ENV_NAMES:
+            local_env[name] = lib_dir
         for k, v in env_vars.items():
             local_env[k] = v
 

scripts/snapdragon/sdk.py

@@ -6,7 +6,7 @@ SDK_CONFIGS = (
     {
         "name": "Hexagon SDK",
         "repo": "snapdragon-toolchain/hexagon-sdk",
-        "default_version": "6.6.0.0",
+        "default_version": "6.7.0.0",
         "parent_dir": "Hexagon_SDK",
         "archive_prefix": "hexagon-sdk-v",
         "markers": ("hexagon_sdk.json",),

src/llama-vocab.cpp

@@ -418,7 +418,7 @@ struct llm_tokenizer_bpe : llm_tokenizer {
                 // original regex from tokenizer.json
                 // "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
                 regex_exprs = {
-                    "[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}])([^a-z]))*((?=[\\p{L}])([^A-Z]))+|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}])([^a-z]))+((?=[\\p{L}])([^A-Z]))*|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+                    "[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))*((?=[\\p{L}\\p{M}])([^A-Z]))+|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))+((?=[\\p{L}\\p{M}])([^A-Z]))*|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
                 };
                 break;
             case LLAMA_VOCAB_PRE_TYPE_CHAMELEON:
@@ -440,7 +440,7 @@ struct llm_tokenizer_bpe : llm_tokenizer {
                 regex_exprs = {
                     // original regex from tokenizer.json
                     // "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
-                    "[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}])([^a-z]))*((?=[\\p{L}])([^A-Z]))+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}])([^a-z]))+((?=[\\p{L}])([^A-Z]))*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+                    "[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))*((?=[\\p{L}\\p{M}])([^A-Z]))+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))+((?=[\\p{L}\\p{M}])([^A-Z]))*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
                 };
                 break;
             case LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI:

src/models/kolibri1.cpp

@@ -131,11 +131,8 @@ llama_model_kolibri1::graph::graph(const llama_model & model, const llm_graph_pa
         cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
         cb(cur, "ffn_norm", il);
 
-        // The router runs in f32: the selection adds the bias to the logits, the weights
-        // are the sigmoid of the unbiased logits, so the same logits feed both.
+        // The selection adds the bias to the logits, the weights are the sigmoid of the unbiased logits, so the same logits feed both.
         ggml_tensor * logits = build_lora_mm(layer.ffn_gate_inp, cur);
-        ggml_prec_set_acc(logits, GGML_PREC_F32);
-        ggml_prec_set_src(logits, GGML_PREC_F32, 1);
         cb(logits, "ffn_moe_logits", il);
 
         ggml_tensor * selection = ggml_add(ctx0, logits, layer.ffn_exp_probs_b);

tests/test-backend-ops.cpp

@@ -10659,7 +10659,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     // ne[2] > 65535 exceeds the CUDA gridDim.y limit (#27901)
     test_cases.emplace_back(new test_rms_norm_mul_rope({4, 1, 65536, 1}, 1e-6f, false, false, false, GGML_ROPE_TYPE_NORMAL));
     test_cases.emplace_back(new test_rms_norm_mul_rope({4, 1, 65536, 1}, 1e-6f, false, true,  false, GGML_ROPE_TYPE_NEOX));
-    for (int64_t d_conv : {3, 4, 9}) {
+    for (int64_t d_conv : {3, 4, 7, 9}) {
         for (int64_t d_inner: {1024, 1536, 2048}) {
             test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_conv, d_inner, 1, 1}, {d_conv, d_inner, 1, 1}));
             test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {2 * d_conv, d_inner, 1, 1}, {d_conv, d_inner, 1, 1}));
@@ -10672,7 +10672,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
 
     // fused ssm_conv + (optional) bias_add + silu. The bias-only graph (no silu) is intentionally
     // not tested since there's no fusion for that pattern in ggml_cuda_can_fuse.
-    for (int64_t d_conv : {3, 4, 9}) {
+    for (int64_t d_conv : {3, 4, 7, 9}) {
         for (int64_t d_inner : {1024, 1536, 2048}) {
             for (bool fuse_bias : {false, true}) {
                 // short token path (n_t <= 32)
@@ -12199,6 +12199,17 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
         }
     }
 
+    // POOL_1D max over MoE down-projection amax shapes:
+    // amax is [1, n_expert, n_tokens] after reducing the whole intermediate dim.
+    for (int64_t n_ff : { 2048, 4096, 8192 }) {
+        for (int64_t n_expert : { 8, 16 }) {
+            for (int64_t n_tokens : { 1, 512 }) {
+                test_cases.emplace_back(new test_pool1d(
+                    GGML_OP_POOL_MAX, GGML_TYPE_F32, {n_ff, n_expert, n_tokens, 1}, n_ff, n_ff, 0));
+            }
+        }
+    }
+
     // LEAKY_RELU at FFN activation width, for direct comparison with RELU
     for (int64_t n_tokens : {512, 2048}) {
         test_cases.emplace_back(new test_leaky_relu(GGML_TYPE_F32, { 17408, n_tokens, 1, 1 }, 0.1f));

tests/test-llama-archs.cpp

@@ -85,10 +85,18 @@ static void fill_resid_gate(std::vector<float> & dst, const char * name, const t
     }
 }
 
+// a random router bias can tie within the noise between the backends, a ramp keeps the top-k selection the same
+static void fill_exp_probs_b(std::vector<float> & dst, const char *, const tensor_data_params &) {
+    for (size_t i = 0; i < dst.size(); i++) {
+        dst[i] = 2.0f*i;
+    }
+}
+
 // a tensor whose values hold a property of their own, matched on a substring of its name, first match wins
 static const struct { const char * match; tensor_filler fill; } TENSOR_FILLERS[] = {
-    { "ssm_a",   fill_ssm_a      },
-    { "_resid_", fill_resid_gate },
+    { "ssm_a",       fill_ssm_a       },
+    { "_resid_",     fill_resid_gate  },
+    { "exp_probs_b", fill_exp_probs_b },
 };
 
 static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) {
@@ -934,6 +942,18 @@ static int save_models(const std::string & arch_filter, const size_t seed, const
     return 0;
 }
 
+// the CPU and the GPU of an Apple SoC carry the same description, the name tells them apart (CPU, MTL0)
+static std::string get_device_label(ggml_backend_dev_t dev) {
+    const char * desc = ggml_backend_dev_description(dev);
+    const char * name = ggml_backend_dev_name(dev);
+
+    if (strcmp(desc, name) == 0) {
+        return desc;
+    }
+
+    return std::string(desc) + " (" + name + ")";
+}
+
 static int test_backends(const std::string & arch_filter, const size_t seed, const float stdev, const int verbosity, const char * target_backend) {
     struct user_data_t {
         struct {
@@ -986,7 +1006,7 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con
                 if (target_backend != nullptr && strcmp(target_backend, ggml_backend_dev_name(dev)) != 0) {
                     continue;
                 }
-                dev_configs.emplace_back(std::vector<ggml_backend_dev_t>{dev}, ggml_backend_dev_description(dev), LLAMA_SPLIT_MODE_LAYER);
+                dev_configs.emplace_back(std::vector<ggml_backend_dev_t>{dev}, get_device_label(dev), LLAMA_SPLIT_MODE_LAYER);
                 max_device_label_length = std::max(max_device_label_length, dev_configs.back().label.length());
 
                 // cpu-based devices cannot be used in tensor split mode

tests/test-mtmd-impl.cpp

@@ -96,6 +96,10 @@ MAKE_TEST(test_image_preprocessor_lfm2) {
     hparams.patch_size = 16;
     hparams.n_merge = 2;
     hparams.set_limit_image_tokens(64, 256);
+    hparams.preproc_min_tiles            = 2;
+    hparams.preproc_max_tiles            = 10;
+    hparams.preproc_tile_size            = 512;
+    hparams.preproc_max_pixels_tolerance = 2.0f;
 
     // { image size, expected tiling }
     const std::vector<std::pair<clip_image_size, bool>> cases = {
@@ -116,6 +120,12 @@ MAKE_TEST(test_image_preprocessor_lfm2) {
             std::string(expected ? "tiled" : "single"),
             std::string(actual   ? "tiled" : "single"));
     }
+
+    // do_image_splitting=false
+    hparams.preproc_min_tiles = 1;
+    hparams.preproc_max_tiles = 1;
+    t.assert_equal("tiling for 1024x977 with max_tiles=1", std::string("single"),
+            std::string(mtmd_image_preprocessor_lfm2::should_tile(hparams, { 1024, 977 }) ? "tiled" : "single"));
 }
 
 //

tools/mtmd/clip-impl.h

@@ -56,6 +56,8 @@
 #define KEY_PREPROC_MIN_TILES       "clip.vision.preproc_min_tiles"
 #define KEY_PREPROC_MAX_TILES       "clip.vision.preproc_max_tiles"
 #define KEY_PREPROC_IMAGE_SIZE      "clip.vision.preproc_image_size"
+#define KEY_PREPROC_TILE_SIZE       "clip.vision.preproc_tile_size"
+#define KEY_PREPROC_MAX_PIXELS_TOLERANCE "clip.vision.preproc_max_pixels_tolerance"
 #define KEY_PATCH_SIZE              "clip.vision.patch_size"
 #define KEY_IMAGE_MEAN              "clip.vision.image_mean"
 #define KEY_IMAGE_STD               "clip.vision.image_std"

tools/mtmd/clip-model.h

@@ -73,7 +73,8 @@ struct clip_hparams {
     int32_t preproc_min_tiles = 0;
     int32_t preproc_max_tiles = 0;
     int32_t max_slice_nums = 9; // llava-uhd slice cap; per-model, carried in the GGUF
-    int32_t preproc_tile_size = 0; // local tile size (deepseek-ocr)
+    int32_t preproc_tile_size = 0; // local tile size (deepseek-ocr, lfm2)
+    float preproc_max_pixels_tolerance = 1.0f;
     resize_algo image_resize_algo_rf = RESIZE_ALGO_BICUBIC;
     resize_algo image_resize_algo_ov = RESIZE_ALGO_BICUBIC;
     pad_style image_pad_rf = PAD_CEIL;  // padding style for the refined image (e.g. llava-1.6)

tools/mtmd/clip.cpp

@@ -1598,7 +1598,17 @@ struct clip_model_loader {
                             hparams.image_pad_rf = PAD_NONE;
                         }
                         get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false);
+                        // defaults for GGUFs converted before tiling params were written
                         // ref: https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B/blob/main/processor_config.json
+                        hparams.preproc_min_tiles            = 2;
+                        hparams.preproc_max_tiles            = 10;
+                        hparams.preproc_tile_size            = 512;
+                        hparams.preproc_max_pixels_tolerance = 2.0f;
+                        get_u32(KEY_PREPROC_MIN_TILES, hparams.preproc_min_tiles, false);
+                        get_u32(KEY_PREPROC_MAX_TILES, hparams.preproc_max_tiles, false);
+                        get_u32(KEY_PREPROC_TILE_SIZE, hparams.preproc_tile_size, false);
+                        get_f32(KEY_PREPROC_MAX_PIXELS_TOLERANCE, hparams.preproc_max_pixels_tolerance, false);
+                        GGML_ASSERT(hparams.preproc_min_tiles >= 1 && hparams.preproc_min_tiles <= hparams.preproc_max_tiles);
                         hparams.set_limit_image_tokens(64, 256);
                     } break;
                 case PROJECTOR_TYPE_PHI4:
@@ -1725,6 +1735,8 @@ struct clip_model_loader {
                 case PROJECTOR_TYPE_GEMMA4UV:
                     {
                         hparams.rope_theta = 100.0f;
+                        hparams.ffn_op = FFN_GELU;
+                        log_ffn_op = "gelu";
                         hparams.n_merge = 3; // pooling_kernel_size
                         hparams.image_resize_algo = RESIZE_ALGO_BICUBIC;
                         get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false);

tools/mtmd/mtmd-image.cpp

@@ -947,8 +947,9 @@ bool mtmd_image_preprocessor_lfm2::should_tile(
     const int h_bar = std::max(hparams.patch_size, round_by_factor(original_size.height));
     const int w_bar = std::max(hparams.patch_size, round_by_factor(original_size.width));
 
-    return static_cast<double>(h_bar) * static_cast<double>(w_bar) >
-           static_cast<double>(hparams.image_max_pixels) * max_pixels_tolerance;
+    return hparams.preproc_max_tiles > 1 &&
+           static_cast<double>(h_bar) * static_cast<double>(w_bar) >
+           static_cast<double>(hparams.image_max_pixels) * hparams.preproc_max_pixels_tolerance;
 }
 
 mtmd_image_preprocessor_llava_uhd::slice_instructions mtmd_image_preprocessor_lfm2::get_slice_instructions(const clip_image_size & original_size) const {
@@ -959,6 +960,7 @@ mtmd_image_preprocessor_llava_uhd::slice_instructions mtmd_image_preprocessor_lf
                             { align_size, hparams.image_min_pixels, hparams.image_max_pixels, 0 });
 
     const bool needs_tiling = should_tile(hparams, original_size);
+    const int  tile_size    = hparams.preproc_tile_size;
 
     if (!needs_tiling) {
         inst.refined_size = clip_image_size{0, 0};
@@ -1009,6 +1011,7 @@ clip_image_size mtmd_image_preprocessor_lfm2::find_closest_aspect_ratio(
             best_ratio_diff = ratio_diff;
             best_ratio = ratio;
         } else if (ratio_diff == best_ratio_diff) {
+            const int   tile_size   = hparams.preproc_tile_size;
             const float target_area = static_cast<float>(tile_size * tile_size * ratio.width * ratio.height);
             if (area > 0.5f * target_area) {
                 best_ratio = ratio;
@@ -1019,6 +1022,8 @@ clip_image_size mtmd_image_preprocessor_lfm2::find_closest_aspect_ratio(
 }
 
 std::vector<clip_image_size> mtmd_image_preprocessor_lfm2::get_target_ratios() const {
+    const int min_tiles = hparams.preproc_min_tiles;
+    const int max_tiles = hparams.preproc_max_tiles;
     std::vector<clip_image_size> ratios;
     for (int n = min_tiles; n <= max_tiles; n++) {
         for (int w = 1; w <= n; w++) {

tools/mtmd/mtmd-image.h

@@ -191,14 +191,9 @@ protected:
 };
 
 // custom llava-uhd slicing logic for LFM2
+// tiling params (min/max tiles, tile size, pixels tolerance) come from hparams
 // ref: https://github.com/huggingface/transformers/blob/v5.1.0/src/transformers/models/lfm2_vl/image_processing_lfm2_vl_fast.py
 struct mtmd_image_preprocessor_lfm2 : mtmd_image_preprocessor_llava_uhd {
-    // ref: https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B/blob/main/processor_config.json
-    static constexpr int   min_tiles            = 2;
-    static constexpr int   max_tiles            = 10;
-    static constexpr float max_pixels_tolerance = 2.0f;
-    static constexpr int   tile_size            = 512;
-
     using mtmd_image_preprocessor_llava_uhd::mtmd_image_preprocessor_llava_uhd;
     mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const override;
     slice_instructions get_slice_instructions(const clip_image_size & original_size) const override;

tools/server/server-chat.cpp

@@ -590,7 +590,7 @@ json server_chat_convert_anthropic_to_oai(const json & body) {
     }
 
     // Pass through common params
-    for (const auto & key : {"temperature", "top_p", "top_k", "stream", "chat_template_kwargs"}) {
+    for (const auto & key : {"temperature", "top_p", "top_k", "stream", "chat_template_kwargs", "id_slot"}) {
         if (body.contains(key)) {
             oai_body[key] = body.at(key);
         }

tools/server/server-context.cpp

@@ -1180,7 +1180,8 @@ private:
         }
 
         // optionally fit mmproj memory usage
-        if (has_mmproj && params_base.fit_params) {
+        // on resume params_base already holds the margin from the first load
+        if (has_mmproj && params_base.fit_params && !is_resume) {
             const auto & mmproj_mem = mmproj_usage.backend_mem_usage;
             const int64_t t_elapsed = mmproj_usage_t_us;
             if (!mmproj_mem.empty()) {

tools/server/server-models.cpp

@@ -2836,9 +2836,10 @@ 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_pipes = [pipe, cli]() {
         pipe->close_read();
         pipe->close_write();
+        cli->stop();
     };
 
     // wire up the receive end of the pipe

tools/server/server-queue.cpp

@@ -448,6 +448,7 @@ server_task_result_ptr server_response::recv(const std::unordered_set<int> & id_
 }
 
 server_task_result_ptr server_response::recv_with_timeout(const std::unordered_set<int> & id_tasks, int timeout) {
+    const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout);
     while (true) {
         std::unique_lock<std::mutex> lock(mutex_results);
 
@@ -459,7 +460,7 @@ server_task_result_ptr server_response::recv_with_timeout(const std::unordered_s
             }
         }
 
-        std::cv_status cr_res = condition_results.wait_for(lock, std::chrono::seconds(timeout));
+        std::cv_status cr_res = condition_results.wait_until(lock, deadline);
         if (!running) {
             RES_DBG("%s : queue result stop\n", __func__);
             std::terminate(); // we cannot return here since the caller is HTTP code

tools/server/server-task.cpp

@@ -591,8 +591,10 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp() {
 json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() {
     std::vector<json> server_sent_events;
     std::vector<json> output;
+    size_t output_index = 0;
 
     if (oaicompat_msg.reasoning_content != "") {
+        const size_t item_output_index = output_index++;
         const json output_item = json {
             {"id",      oai_resp_reasoning_id},
             {"summary", json::array()},
@@ -608,6 +610,7 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() {
             {"event", "response.output_item.done"},
             {"data", json {
                 {"type", "response.output_item.done"},
+                {"output_index", item_output_index},
                 {"item", output_item}
             }}
         });
@@ -615,11 +618,14 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() {
     }
 
     if (oaicompat_msg.content != "") {
+        const size_t item_output_index = output_index++;
         server_sent_events.push_back(json {
             {"event", "response.output_text.done"},
             {"data", json {
                 {"type",    "response.output_text.done"},
                 {"item_id", oai_resp_message_id},
+                {"output_index", item_output_index},
+                {"content_index", 0},
                 {"text",    oaicompat_msg.content}
             }}
         });
@@ -636,6 +642,8 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() {
             {"data", json {
                 {"type",    "response.content_part.done"},
                 {"item_id", oai_resp_message_id},
+                {"output_index", item_output_index},
+                {"content_index", 0},
                 {"part",    content_part}
             }}
         });
@@ -651,6 +659,7 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() {
             {"event", "response.output_item.done"},
             {"data", json {
                 {"type", "response.output_item.done"},
+                {"output_index", item_output_index},
                 {"item", output_item}
             }}
         });
@@ -658,6 +667,7 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() {
     }
 
     for (const common_chat_tool_call & tool_call : oaicompat_msg.tool_calls) {
+        const size_t item_output_index = output_index++;
         const json output_item = {
             {"id",        "fc_" + tool_call.id},
             {"type",      "function_call"},
@@ -670,6 +680,7 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() {
             {"event", "response.output_item.done"},
             {"data", json {
                 {"type", "response.output_item.done"},
+                {"output_index", item_output_index},
                 {"item", output_item}
             }}
         });
@@ -1202,6 +1213,7 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
                     {"event", "response.output_item.added"},
                     {"data", json {
                         {"type", "response.output_item.added"},
+                        {"output_index", 0},
                         {"item", json {
                             {"id",                oai_resp_reasoning_id},
                             {"summary",           json::array()},
@@ -1218,6 +1230,8 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
                 {"event", "response.reasoning_text.delta"},
                 {"data", json {
                     {"type",    "response.reasoning_text.delta"},
+                    {"output_index", 0},
+                    {"content_index", 0},
                     {"delta",   diff.reasoning_content_delta},
                     {"item_id", oai_resp_reasoning_id},
                 }},
@@ -1230,6 +1244,7 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
                     {"event", "response.output_item.added"},
                     {"data", json {
                         {"type", "response.output_item.added"},
+                        {"output_index", anthropic_has_reasoning ? 1 : 0},
                         {"item", json {
                             {"content", json::array()},
                             {"id",      oai_resp_message_id},
@@ -1244,6 +1259,8 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
                     {"data", json {
                         {"type",    "response.content_part.added"},
                         {"item_id", oai_resp_message_id},
+                        {"output_index", anthropic_has_reasoning ? 1 : 0},
+                        {"content_index", 0},
                         {"part", json {
                             {"type", "output_text"},
                             {"text", ""},
@@ -1257,6 +1274,8 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
                 {"data", json {
                     {"type",    "response.output_text.delta"},
                     {"item_id", oai_resp_message_id},
+                    {"output_index", anthropic_has_reasoning ? 1 : 0},
+                    {"content_index", 0},
                     {"delta",   diff.content_delta},
                 }},
             });
@@ -1267,6 +1286,7 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
                 {"event", "response.output_item.added"},
                 {"data", json {
                     {"type",  "response.output_item.added"},
+                    {"output_index", (anthropic_has_reasoning ? 1 : 0) + (text_block_started ? 1 : 0) + diff.tool_call_index},
                     {"item", json {
                         {"id",        "fc_" + diff.tool_call_delta.id},
                         {"arguments", ""},
@@ -1287,6 +1307,7 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() {
                     {"type",    "response.function_call_arguments.delta"},
                     {"delta",   diff.tool_call_delta.arguments},
                     {"item_id", "fc_" + oai_resp_fc_id},
+                    {"output_index", (anthropic_has_reasoning ? 1 : 0) + (text_block_started ? 1 : 0) + diff.tool_call_index},
                 }},
             });
         }