llama : reuse the tensor staging buffer and skip the direct-io bounce buffer (#31321 )

server : restore the spec checkpoint for drafts longer than n_rs_seq (#31350 )
Assisted-by: Claude Co-authored-by: Zeshawn Zahid <zzshawnn@gmail.com>
2026-10-30 16:37:36 -05:00 · 2026-10-29 19:22:43 +08:00 · 2026-10-29 11:28:03 +02:00 · 2026-10-29 10:18:13 +02:00 · 2026-10-29 10:15:54 +02:00 · 2026-10-29 10:15:31 +02:00
361 changed files with 52266 additions and 22145 deletions
@@ -107,3 +107,52 @@ jobs:
          source .venv/bin/activate
          cd build
          ctest -L behavior -LE 'model|backend' --verbose --timeout 900
+
+  windows:
+    runs-on: windows-2025
+
+    cache-mode: write
+    permissions:
+      actions: write
+      contents: read
+
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.24
+        with:
+          key: behavior-check-windows-2025
+          evict-old-files: 1d
+          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+
+      - name: Python Dependencies
+        id: python_depends
+        shell: bash
+        run: |
+          python3 -m venv .venv
+          source .venv/Scripts/activate
+          pip install jinja2==3.1.6
+
+      - name: Build
+        id: cmake_build
+        shell: cmd
+        run: |
+          call "C:\Program Files\Microsoft Visual Studio\18\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" x64
+          cmake -B build -G "Ninja Multi-Config" ^
+            -DCMAKE_TOOLCHAIN_FILE=cmake/x64-windows-llvm.cmake ^
+            -DCMAKE_BUILD_TYPE=Release ^
+            -DGGML_NATIVE=OFF ^
+            -DGGML_BACKEND_DL=OFF
+          set /A NINJA_JOBS=%NUMBER_OF_PROCESSORS%-1
+          cmake --build build --config Release -j %NINJA_JOBS%
+
+      - name: Test
+        id: cmake_test
+        shell: bash
+        run: |
+          source .venv/Scripts/activate
+          cd build
+          ctest -C Release -L behavior -LE 'model|backend' --verbose --timeout 900
@@ -43,9 +43,9 @@ env:
  LLAMA_ARG_LOG_TIMESTAMPS: 1

 jobs:
-  ubuntu-24-s390x:
-    name: ubuntu-24-s390x (VXE ${{ matrix.vxe }})
-    runs-on: ubuntu-24.04-s390x
+  ubuntu-26-s390x:
+    name: ubuntu-26-s390x (VXE ${{ matrix.vxe }})
+    runs-on: ubuntu-26.04-s390x

    strategy:
      fail-fast: false
@@ -66,12 +66,6 @@ jobs:
            libjpeg-dev build-essential libssl-dev \
            git-lfs

-      - name: Toolchain workaround (GCC 14)
-        run: |
-          sudo apt-get install -y gcc-14 g++-14
-          echo "CC=gcc-14" >> "$GITHUB_ENV"
-          echo "CXX=g++-14" >> "$GITHUB_ENV"
-
      - name: Python Dependencies
        id: python_depends
        run: |
@@ -116,30 +110,19 @@ jobs:

  ubuntu-26-zdnn-s390x:
    name: ubuntu-26-zdnn-s390x
-    runs-on: ubuntu-24.04-s390x
-    container: ubuntu:26.04  # required to get GCC 15.1 and binutils 2.44
-    defaults:
-      run:
-        shell: bash
+    runs-on: ubuntu-26.04-s390x

    steps:
-      - name: Build Dependencies
-        id: build_depends
-        run: |
-          apt-get update
-          apt-get install -y --no-install-recommends \
-            build-essential cmake git ca-certificates \
-            libssl-dev libzdnn-dev
-
      - name: Clone
        id: checkout
        uses: actions/checkout@v6

-      - name: Toolchain workaround (GCC 15)
+      - name: Build Dependencies
+        id: build_depends
        run: |
-          apt-get install -y gcc-15 g++-15
-          echo "CC=gcc-15" >> "$GITHUB_ENV"
-          echo "CXX=g++-15" >> "$GITHUB_ENV"
+          sudo apt-get update
+          sudo apt-get install -y --no-install-recommends \
+            build-essential libssl-dev libzdnn-dev

      - name: Build with zDNN Backend
        id: cmake_build
@@ -58,7 +58,12 @@ jobs:
      - name: Install dependencies
        run: |
          sudo apt update
-          sudo apt install -y build-essential cmake libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libvulkan1 mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 libssl-dev time unzip wget python3 python3-venv python3-pip
+          sudo apt install -y software-properties-common
+          sudo add-apt-repository ppa:deadsnakes/ppa
+          sudo apt update
+          sudo apt install -y build-essential cmake libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libvulkan1 mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 libssl-dev time unzip wget python3.12 python3.12-venv
+          sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1
+          sudo update-alternatives --config python3

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.24
@@ -172,6 +172,8 @@ jobs:
            - [Ubuntu x64 (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-x64.tar.gz)
            - [Ubuntu arm64 (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-arm64.tar.gz)
            - [Ubuntu s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz)
+            - [Ubuntu 26.04 s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-26.04-s390x.tar.gz)
+            - [Ubuntu 26.04 s390x (zDNN)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-zdnn-s390x.tar.gz)
            - [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz)
            - [Ubuntu arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-arm64.tar.gz)
            - [Ubuntu x64 (CUDA 12)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-12.8-x64.tar.gz) - [CUDA 12.8 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-12.8-x64.tar.gz)
@@ -180,6 +180,8 @@ jobs:
            os: ubuntu-24.04-arm
          - build: 's390x'
            os: ubuntu-24.04-s390x
+          - build: '26.04-s390x'
+            os: ubuntu-26.04-s390x

    runs-on: ${{ matrix.os }}

@@ -200,7 +202,7 @@ jobs:
        id: depends
        run: |
          sudo apt-get update
-          sudo apt-get install build-essential libssl-dev
+          sudo apt-get install -y build-essential libssl-dev

      - name: Toolchain workaround (GCC 14)
        if: ${{ contains(matrix.os, 'ubuntu-24.04') }}
@@ -210,7 +212,7 @@ jobs:
          echo "CXX=g++-14" >> "$GITHUB_ENV"

      - name: ccache
-        if: ${{ matrix.build != 's390x' }}
+        if: ${{ !contains(matrix.os, 's390x') }}
        uses: ggml-org/ccache-action@v1.2.24
        with:
          key: release-${{ matrix.os }}-cpu
@@ -242,7 +244,7 @@ jobs:
          archive: false

      - name: ccache-clear
-        if: ${{ matrix.build != 's390x' }}
+        if: ${{ !contains(matrix.os, 's390x') }}
        uses: ./.github/actions/ccache-clear
        with:
          key: release-${{ matrix.os }}-cpu
@@ -968,6 +970,11 @@ jobs:
          evict-old-files: 1d
          max-size: "1G"

+      - name: Tune ccache for ROCm
+        run: |
+          # since ROCm 10.1, clang -E output has no GPU code, so the preprocessor fallback can reuse objects with old GPU code
+          ccache --set-config=depend_mode=true
+
      # - name: Cache ROCm Installation
      #   id: cache-rocm
      #   uses: actions/cache@v5
@@ -1545,6 +1552,8 @@ jobs:
          # builds hit.
          ccache --set-config=compiler_check=content
          ccache --set-config=sloppiness=time_macros,include_file_mtime,include_file_ctime
+          # since ROCm 10.1, clang -E output has no GPU code, so the preprocessor fallback can reuse objects with old GPU code
+          ccache --set-config=depend_mode=true

      - name: Dependencies
        id: depends
@@ -1619,6 +1628,60 @@ jobs:
        with:
          key: release-ubuntu-24.04-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}

+  ubuntu-26-zdnn:
+    needs: [check-release, ui-build]
+    if: ${{ needs.check-release.outputs.should_release == 'true' }}
+
+    runs-on: ubuntu-26.04-s390x
+
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+
+      - name: Download UI build
+        uses: actions/download-artifact@v8
+        with:
+          name: llama-ui.zip
+          path: tools/ui/dist
+
+      - name: Dependencies
+        id: depends
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y --no-install-recommends \
+            build-essential libssl-dev libzdnn-dev
+
+      - name: Build
+        id: cmake_build
+        run: |
+          cmake -B build \
+            -DCMAKE_INSTALL_RPATH='$ORIGIN' \
+            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
+            -DGGML_BACKEND_DL=ON \
+            -DGGML_NATIVE=OFF \
+            -DGGML_VXE=ON \
+            -DGGML_ZDNN=ON \
+            -DCMAKE_C_FLAGS="-march=arch15" \
+            -DCMAKE_CXX_FLAGS="-march=arch15" \
+            -DLLAMA_FATAL_WARNINGS=ON \
+            ${{ env.CMAKE_ARGS }}
+          cmake --build build --config Release -j $(nproc)
+
+      - name: Pack artifacts
+        id: pack_artifacts
+        run: |
+          cp LICENSE ./build/bin/
+          tar -czvf llama-${{ needs.check-release.outputs.tag_name }}-bin-ubuntu-zdnn-s390x.tar.gz --transform "s,^\.,llama-${{ needs.check-release.outputs.tag_name }}," -C ./build/bin .
+
+      - name: Upload artifacts
+        uses: actions/upload-artifact@v7
+        with:
+          path: llama-${{ needs.check-release.outputs.tag_name }}-bin-ubuntu-zdnn-s390x.tar.gz
+          archive: false
+
  ios-xcode:
    needs: [check-release]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
@@ -4,7 +4,7 @@ include(CheckIncludeFileCXX)

 ### llama.cpp version
 set(LLAMA_VERSION_MAJOR 0)
-set(LLAMA_VERSION_MINOR 7)
+set(LLAMA_VERSION_MINOR 8)
 set(LLAMA_VERSION_PATCH 0)
 set(LLAMA_VERSION_BASE "${LLAMA_VERSION_MAJOR}.${LLAMA_VERSION_MINOR}.${LLAMA_VERSION_PATCH}")

@@ -103,7 +103,7 @@ endif()

 # subprocess spawning isn't a supported/sandbox-friendly operation on mobile OSes or in WASM
 if (CMAKE_SYSTEM_NAME STREQUAL "iOS" OR CMAKE_SYSTEM_NAME STREQUAL "Android" OR ANDROID
-        OR CMAKE_SYSTEM_NAME STREQUAL "Emscripten" OR EMSCRIPTEN)
+        OR CMAKE_SYSTEM_NAME STREQUAL "Emscripten" OR EMSCRIPTEN OR CMAKE_SYSTEM_NAME STREQUAL "WASI")
    set(LLAMA_SUBPROCESS_DEFAULT OFF)
 else()
    set(LLAMA_SUBPROCESS_DEFAULT ON)
@@ -2,20 +2,29 @@
 # multiple collaborators per item can be specified
 #
 # ggml-org/ci               : CISC, danbev, ggerganov, netrunnereve, ngxson, taronaeo
+# ggml-org/ggml-base        : aendk, fairydreaming, ggerganov
 # ggml-org/ggml-cann        : hipudding
-# ggml-org/ggml-cuda        : JohannesGaessler, am17an, IMbackK, ORippler
+# ggml-org/ggml-cpu         : angt, chaxu01, ggerganov, max-krasnyansky, taronaeo
+# ggml-org/ggml-cuda        : am17an, gaugarg-nv, IMbackK, JohannesGaessler, ORippler, thelittlefireman, ynankani
+# ggml-org/ggml-et          : marty1885
 # ggml-org/ggml-hexagon     : lhez, max-krasnyansky
-# ggml-org/ggml-metal       : ggerganov
+# ggml-org/ggml-hip         : IMbackK, thelittlefireman
+# ggml-org/ggml-metal       : forforever73, ggerganov, jhen0409, masterFoad
 # ggml-org/ggml-opencl      : lhez, max-krasnyansky
-# ggml-org/ggml-rpc         : rgerganov
-# ggml-org/ggml-sycl        : arthw
+# ggml-org/ggml-openvino    : ravi9, wine99
+# ggml-org/ggml-riscv       : xctan
+# ggml-org/ggml-rpc         : am17an, rgerganov
+# ggml-org/ggml-sycl        : arthw, Titaniumtown
 # ggml-org/ggml-vulkan      : 0cc4m, jeffbolznv
 # ggml-org/ggml-webgpu      : reeselevine, yomaytk
-# ggml-org/ggml-zdnn        : taronaeo
-# ggml-org/llama-common     : ggerganov, aldehir, angt, danbev, ngxson, pwilkin
+# ggml-org/ggml-zdnn        : AlekseiNikiforovIBM, Andreas-Krebbel, taronaeo
+# ggml-org/llama-base       : am17an, CISC, ggerganov, ngxson, sanmai, ServeurpersoCom
+# ggml-org/llama-common     : aldehir, angt, bartowski1182, danbev, ggerganov, ngxson, nikwen, pwilkin, ruixiang63, sanmai
+# ggml-org/llama-examples   : ggerganov
+# ggml-org/llama-models     : CISC, danbev, ggerganov, ngxson
 # ggml-org/llama-mtmd       : ngxson
-# ggml-org/llama-server     : ggerganov, ngxson, allozaur, angt, ServeurpersoCom
-# ggml-org/llama-ui           : allozaur
+# ggml-org/llama-server     : allozaur, angt, ggerganov, ngxson, ServeurpersoCom
+# ggml-org/llama-ui         : allozaur

 /.devops/*.Dockerfile                   @ngxson
 /.github/actions/                       @ggml-org/ci
@@ -29,28 +38,12 @@
 /conversion/                            @CISC
 /convert_*.py                           @CISC
 /docs/backend/snapdragon/               @ggml-org/ggml-hexagon
-/examples/batched.swift/                @ggerganov
-/examples/batched/                      @ggerganov
-/examples/convert-llama2c-to-ggml/      @ggerganov
+/examples/                              @ggml-org/llama-examples
 /examples/debug/                        @danbev @pwilkin
-/examples/deprecation-warning/          @ggerganov
 /examples/diffusion/                    @am17an
-/examples/embedding/                    @ggerganov
-/examples/eval-callback/                @ggerganov
-/examples/export-docs/                  @ggerganov
-/examples/gen-docs/                     @ggerganov
-/examples/gguf/                         @ggerganov
 /examples/llama.android/                @ggerganov @hanyin-arm @naco-siren
-/examples/llama.swiftui/                @ggerganov
-/examples/llama.vim                     @ggerganov
-/examples/lookahead/                    @ggerganov
 /examples/lookup/                       @JohannesGaessler
 /examples/model-conversion/             @danbev
-/examples/parallel/                     @ggerganov
-/examples/passkey/                      @ggerganov
-/examples/retrieval/                    @ggerganov
-/examples/speculative-simple/           @ggerganov
-/examples/speculative/                  @ggerganov
 /ggml/cmake/                            @ggerganov
 /ggml/include/                          @ggerganov
 /ggml/src/ggml-backend-meta.cpp         @JohannesGaessler
@@ -62,7 +55,7 @@
 /ggml/src/ggml-cuda/                    @ggml-org/ggml-cuda
 /ggml/src/ggml-cuda/vendors/hip.h       @IMbackK
 /ggml/src/ggml-hexagon/                 @ggml-org/ggml-hexagon
-/ggml/src/ggml-hip/                     @IMbackK
+/ggml/src/ggml-hip/                     @ggml-org/ggml-hip
 /ggml/src/ggml-et/                      @marty1885
 /ggml/src/ggml-impl.h                   @ggerganov
 /ggml/src/ggml-metal/                   @ggml-org/ggml-metal
@@ -76,7 +69,7 @@
 /ggml/src/ggml-virtgpu/                 @kpouget
 /ggml/src/ggml-vulkan/                  @ggml-org/ggml-vulkan
 /ggml/src/ggml-webgpu/                  @ggml-org/ggml-webgpu
-/ggml/src/ggml-zdnn/                    @ggml-org/ggml-zdnn @Andreas-Krebbel @AlekseiNikiforovIBM
+/ggml/src/ggml-zdnn/                    @ggml-org/ggml-zdnn
 /ggml/src/ggml-zendnn/                  @avinashcpandey @Jiten1parmar
 /ggml/src/ggml.c                        @ggerganov
 /ggml/src/ggml.cpp                      @ggerganov
@@ -87,14 +80,14 @@
 /scripts/get*                           @ggerganov
 /scripts/sync*                          @ggerganov
 /scripts/snapdragon/                    @ggml-org/ggml-hexagon
-/src/                                   @ggerganov
+/src/                                   @ggml-org/llama-base
 /src/llama-adapter.*                    @CISC
 /src/llama-arch.*                       @CISC
 /src/llama-chat.*                       @ngxson
 /src/llama-graph.*                      @CISC
 /src/llama-model.*                      @CISC
 /src/llama-vocab.*                      @CISC
-/src/models/                            @CISC
+/src/models/                            @ggml-org/llama-models
 /tests/                                 @ggerganov
 /tests/test-chat.*                      @pwilkin
 /tests/test-server-component.cpp        @ggml-org/llama-server
@@ -13,7 +13,7 @@
 [![Docker](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/docker.yml?label=Docker)](https://github.com/ggml-org/llama.cpp/actions/workflows/docker.yml)
 [![Winget](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/winget.yml?label=Winget)](https://github.com/ggml-org/llama.cpp/actions/workflows/winget.yml)

-[ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/issues?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Ajhen0409%20OR%20author%3Abartowski1182%20OR%20author%3Anikwen%20OR%20author%3Ahipudding%20OR%20author%3Aravi9%20OR%20author%3AServeurpersoCom%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3Amarty1885%20OR%20author%3A0cc4m%20OR%20author%3ATitaniumtown%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3Aallozaur%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Awine99%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Alhez%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming)%20sort%3Aupdated-desc) / [dev stats](https://github.com/ggml-org/llama.cpp-dev) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291)
+[ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/issues?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Asanmai%20OR%20author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Ajhen0409%20OR%20author%3Abartowski1182%20OR%20author%3Anikwen%20OR%20author%3Ahipudding%20OR%20author%3Aravi9%20OR%20author%3AServeurpersoCom%20OR%20author%3Athelittlefireman%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3Amarty1885%20OR%20author%3A0cc4m%20OR%20author%3ATitaniumtown%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3AmasterFoad%20OR%20author%3Aallozaur%20OR%20author%3AAndreas-Krebbel%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Awine99%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Achaxu01%20OR%20author%3Alhez%20OR%20author%3AAlekseiNikiforovIBM%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming%20OR%20author%3Aynankani)%20sort%3Aupdated-desc) / [dev stats](https://github.com/ggml-org/llama.cpp-dev) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291)

 </div>

@@ -924,7 +924,7 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
        // TODO @ngxson : maybe show a list of available models in CLI in this case
        bool can_skip_model = params.usage || params.completion || !params.server_base.empty();
        if (!can_skip_model && params.model.path.empty()) {
-            throw std::invalid_argument("error: --model is required\n");
+            throw std::invalid_argument("error: --model or -hf is required\n");
        }
    }

@@ -3708,6 +3708,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            params.models_max = value;
        }
    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODELS_MAX"));
+    add_opt(common_arg(
+        {"--models-memory-margin"}, "N",
+        string_format("for router server, MiB of device memory to keep free per device, unload LRU models until a new one fits (default: %d, 0 = disabled)", params.models_memory_margin),
+        [](common_params & params, int value) {
+            params.models_memory_margin = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODELS_MEMORY_MARGIN"));
    add_opt(common_arg(
        {"--models-autoload"},
        {"--no-models-autoload"},
@@ -4505,6 +4512,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            params.tts_lang = value;
        }
    ).set_examples({LLAMA_EXAMPLE_TTS}));
+    add_opt(common_arg(
+        {"--tts-voice"}, "VOICE",
+        "voice for audio generation (model-dependent)\n"
+        "see tts/README.md for per-model usage notes",
+        [](common_params & params, const std::string & value) {
+            params.tts_voice = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_TTS}));
    add_opt(common_arg(
        {"--tts-speaker-file"}, "FNAME",
        "speaker file path for audio generation",
@@ -190,10 +190,17 @@ json common_chat_msg::to_json_oaicompat(bool concat_typed_text) const {
        } else {
            auto & parts = jmsg["content"] = json::array();
            for (const auto & part : content_parts) {
-                parts.push_back({
-                    {"type", part.type},
-                    {"text", part.text},
-                });
+                if (part.type == "tool_reference") {
+                    parts.push_back({
+                        {"type", part.type},
+                        {"name", part.name},
+                    });
+                } else {
+                    parts.push_back({
+                        {"type", part.type},
+                        {"text", part.text},
+                    });
+                }
            }
        }
    } else {
@@ -309,6 +316,7 @@ struct common_chat_templates {
    bool has_explicit_template;  // Model had builtin template or template overridden was specified.
    std::unique_ptr<common_chat_template> template_default;  // always set (defaults to chatml)
    std::unique_ptr<common_chat_template> template_tool_use;
+    std::unique_ptr<common_chat_template> template_tts;  // the prompt of a text to speak
 };

 common_chat_tool_choice common_chat_tool_choice_parse_oaicompat(const std::string & tool_choice) {
@@ -370,12 +378,16 @@ std::vector<common_chat_msg> common_chat_msgs_parse_oaicompat(const json & messa
                            throw std::invalid_argument("Missing content part type: " + part.dump());
                        }
                        const auto & type = part.at("type");
-                        if (type != "text" && type != "media_marker") {
-                            throw std::invalid_argument("Unsupported content part type: " + type.dump());
-                        }
                        common_chat_msg_content_part msg_part;
                        msg_part.type = type;
-                        msg_part.text = part.at("text");
+                        if (type == "text" || type == "media_marker") {
+                            msg_part.text = part.at("text");
+                        } else if (type == "tool_reference") {
+                            // names a request tool, the template renders its definition
+                            msg_part.name = part.at("name");
+                        } else {
+                            throw std::invalid_argument("Unsupported content part type: " + type.dump());
+                        }
                        msg.content_parts.push_back(msg_part);
                    }
                } else if (!content.is_null()) {
@@ -610,13 +622,17 @@ json common_chat_tools_to_json_oaicompat(const std::vector<common_chat_tool> & t

    auto result = json::array();
    for (const auto & tool : tools) {
+        json function {
+            { "name",        tool.name },
+            { "description", tool.description },
+            { "parameters",  json::parse(tool.parameters) },
+        };
+        if (tool.defer_loading) {
+            function["defer_loading"] = true;
+        }
        result.push_back({
            { "type",     "function" },
-            { "function", {
-                { "name", tool.name },
-                { "description", tool.description },
-                { "parameters", json::parse(tool.parameters) },
-            }},
+            { "function", function },
        });
    }
    return result;
@@ -653,10 +669,19 @@ std::vector<common_chat_tool> common_chat_tools_parse_oaicompat(const json & too
                }

                const auto & function = tool.at("function");
+
+                // "defer_loading" sits on the function, or on the tool next to "type"; the function wins
+                const json & defer = function.contains("defer_loading") ? function.at("defer_loading")
+                                                                        : tool.value("defer_loading", json(false));
+                if (!defer.is_boolean()) {
+                    throw std::invalid_argument("Expected 'defer_loading' to be a boolean, got " + defer.dump());
+                }
+
                result.push_back({
                    /* .name = */ function.at("name"),
                    /* .description = */ function.value("description", ""),
                    /* .parameters = */ function.value("parameters", json::object()).dump(),
+                    /* .defer_loading = */ defer.get<bool>(),
                });
            }
        }
@@ -779,19 +804,70 @@ bool common_chat_templates_was_explicit(const struct common_chat_templates * tmp
    return tmpls->has_explicit_template;
 }

-common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates) {
+// renders a task template (e.g. "tts") with the given variables
+static std::string common_chat_template_render(const common_chat_template & tmpl, const json & vars) {
+    jinja::context ctx(tmpl.source());
+    jinja::global_from_json(ctx, vars, false);
+    ctx.set_val("bos_token", jinja::mk_val<jinja::value_string>(tmpl.bos_token()));
+    ctx.set_val("eos_token", jinja::mk_val<jinja::value_string>(tmpl.eos_token()));
+    jinja::runtime runtime(ctx);
+    return jinja::runtime::gather_string_parts(runtime.execute(tmpl.prog))->as_string().str();
+}
+
+// LFM3 format detection: template renders deferred tools inside <additional_tools>...</additional_tools>
+static bool is_lfm3_template(const std::string & src) {
+    return src.find("<additional_tools>") != std::string::npos;
+}
+
+// LFM3-Audio: a known language code maps to its name, no language lets the model detect it
+static std::string lfm3_asr_system_prompt(const std::string & language) {
+    static const std::map<std::string, std::string> names = {
+        { "ar", "Arabic" }, { "da", "Danish" }, { "de", "German" }, { "el", "Greek" }, { "en", "English" },
+        { "es", "Spanish" }, { "fi", "Finnish" }, { "fr", "French" }, { "he", "Hebrew" }, { "it", "Italian" },
+        { "ja", "Japanese" }, { "ko", "Korean" }, { "nl", "Dutch" }, { "pl", "Polish" }, { "pt", "Portuguese" },
+        { "sv", "Swedish" }, { "tr", "Turkish" }, { "zh", "Chinese" },
+    };
+    if (language.empty()) {
+        return "Perform ASR.";
+    }
+    std::string code = language;
+    std::transform(code.begin(), code.end(), code.begin(), [](unsigned char c) { return (char) std::tolower(c); });
+    const auto it = names.find(code);
+    return "Perform ASR in " + (it != names.end() ? it->second : language) + ".";
+}
+
+common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates, const std::string & language) {
    common_chat_prompt_preset asr_preset;
    asr_preset.system = "";
    asr_preset.user   = "Transcribe audio to text";

-    if (chat_templates && chat_templates->template_default && is_lfm2_template(chat_templates->template_default->source())) {
-        asr_preset.system = "Perform ASR.";
-        asr_preset.user   = "";
+    if (chat_templates && chat_templates->template_default) {
+        const std::string src = chat_templates->template_default->source();
+        if (is_lfm2_template(src)) {
+            asr_preset.system = "Perform ASR.";
+            asr_preset.user   = "";
+        } else if (is_lfm3_template(src)) {
+            asr_preset.system = lfm3_asr_system_prompt(language);
+            asr_preset.user   = "";
+            asr_preset.has_language = true;
+        }
    }

    return asr_preset;
 }

+std::string common_chat_get_tts_prompt(const common_chat_templates * chat_templates, const std::string & text,
+                                       const std::string & language, const std::string & voice) {
+    if (!chat_templates || !chat_templates->template_tts) {
+        return "";
+    }
+    return common_chat_template_render(*chat_templates->template_tts, {
+        {"messages", json::array({{{"role", "user"}, {"content", text}}})},
+        {"language", language},
+        {"voice",    voice},
+    });
+}
+
 std::string common_chat_templates_source(const struct common_chat_templates * tmpls, const std::string & variant) {
    if (!variant.empty()) {
        if (variant == "tool_use") {
@@ -902,6 +978,16 @@ common_chat_templates_ptr common_chat_templates_init(const struct llama_model *
            LOG_ERR("%s: failed to parse tool use chat template (ignoring it): %s\n", __func__, e.what());
        }
    }
+    // task templates of the model, kept with a chat template override
+    if (model) {
+        if (const char * src = llama_model_chat_template(model, "tts")) {
+            try {
+                tmpls->template_tts = std::make_unique<common_chat_template>(src, token_bos, token_eos);
+            } catch (const std::exception & e) {
+                LOG_ERR("%s: failed to parse the tts chat template (ignoring it): %s\n", __func__, e.what());
+            }
+        }
+    }
    return tmpls;
 }

@@ -1717,6 +1803,9 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena &          src_pars
    //LOG_DBG("Parsing PEG input with format %s: %s\n", common_chat_format_name(params.format), effective_input.text.c_str());

    common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_LENIENT;
+    if (!is_partial) {
+        flags |= COMMON_PEG_PARSE_FLAG_COMPLETE;
+    }
    if (params.debug) {
        flags |= COMMON_PEG_PARSE_FLAG_DEBUG;
    }
@@ -1854,12 +1943,13 @@ const common_chat_msg & common_chat_session::finish(const common_chat_input & ch
    return result;
 }

-std::map<std::string, bool> common_chat_templates_get_caps(const common_chat_templates * chat_templates) {
+common_json common_chat_templates_get_caps(const common_chat_templates * chat_templates) {
    GGML_ASSERT(chat_templates != nullptr);
    GGML_ASSERT(chat_templates->template_default != nullptr);
-    if (chat_templates->template_tool_use != nullptr) {
-        // take the more expressive template when available
-        return chat_templates->template_tool_use->caps.to_map();
-    }
-    return chat_templates->template_default->caps.to_map();
+    // take the more expressive template when available
+    const auto & caps = chat_templates->template_tool_use != nullptr ? chat_templates->template_tool_use->caps
+                                                                      : chat_templates->template_default->caps;
+    common_json result = caps.to_map();
+    result["reasoning_effort_levels"] = caps.reasoning_effort_levels;
+    return result;
 }
@@ -45,6 +45,7 @@ struct common_chat_tool_call {
 struct common_chat_msg_content_part {
    std::string type;
    std::string text;
+    std::string name; // tool_reference: name of a tool from the request tools

    // TODO @ngxson : no known chat templates support reasoning_content in content parts yet
    //                this can be useful for models with interleaved thinking (like Kimi-K2)
@@ -52,7 +53,7 @@ struct common_chat_msg_content_part {
    // std::string reasoning_content;

    bool operator==(const common_chat_msg_content_part & other) const {
-        return type == other.type && text == other.text;
+        return type == other.type && text == other.text && name == other.name;
    }
 };

@@ -226,6 +227,9 @@ struct common_chat_tool {
    std::string name;
    std::string description;
    std::string parameters;
+
+    // withheld from the system prompt, a tool_reference reveals it later in the conversation
+    bool defer_loading = false;
 };

 enum common_chat_tool_choice {
@@ -468,7 +472,7 @@ common_json common_chat_tools_to_json_oaicompat(const std::vector<common_chat_to
 common_json common_chat_tool_parameters(const common_json & function);

 // get template caps, useful for reporting to server /props endpoint
-std::map<std::string, bool> common_chat_templates_get_caps(const common_chat_templates * chat_templates);
+common_json common_chat_templates_get_caps(const common_chat_templates * chat_templates);

 common_chat_prompt common_chat_template_direct_apply(
    const common_chat_template & tmpl,
@@ -487,6 +491,12 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
 struct common_chat_prompt_preset {
    std::string system;
    std::string user;
+    bool has_language = false; // the preset already names the requested language
 };

-common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates);
+common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates, const std::string & language = "");
+
+// prompt of the model's "tts" chat template for the text to speak, empty if the model has none
+// language and voice are optional, the template defines their defaults
+std::string common_chat_get_tts_prompt(const common_chat_templates * chat_templates, const std::string & text,
+                                       const std::string & language = "", const std::string & voice = "");
@@ -219,6 +219,23 @@ int32_t common_cpu_get_num_math() {
            }
        }
    }
+#elif defined(__riscv) && defined(__linux__)
+    std::ifstream cpuinfo("/proc/cpuinfo");
+    std::string line;
+    int32_t n_a100 = 0;
+    while (std::getline(cpuinfo, line)) {
+        std::istringstream fields(line);
+        std::string key;
+        std::string separator;
+        uint64_t marchid;
+        if (fields >> key >> separator >> std::hex >> marchid &&
+            key == "marchid" && separator == ":" && marchid == UINT64_C(0x8000000041000002)) {
+            ++n_a100;
+        }
+    }
+    if (n_a100 > 0) {
+        return n_a100;
+    }
 #elif defined(__powerpc64__) || defined(__powerpc__)
    int32_t smt_factor = 1;
    int phy_cpus = common_cpu_get_num_physical_cores();
@@ -1187,6 +1204,9 @@ static const std::map<common_decision_type, std::string> COMMON_DECISION_TYPE_NA
    { COMMON_DECISION_TYPE_LFM2_D1_OMNI,   "lfm2-d1-omni"  },
    { COMMON_DECISION_TYPE_GRANITE_SWITCH, "granite-switch"},
    { COMMON_DECISION_TYPE_H2O_LIGHTNING,  "h2o-lightning" },
+    { COMMON_DECISION_TYPE_MINICPM_OCR_JEV, "openbmb-minicpm-ocr-jev" },
+    { COMMON_DECISION_TYPE_LFM3_D1,        "lfm3-d1"       },
+    { COMMON_DECISION_TYPE_LFM3_D1_OMNI,   "lfm3-d1-omni"  },
    // for future models, ALWAYS add org name prefix to avoid name collisions, e.g. "orgname-modelname"
 };

@@ -1252,44 +1272,50 @@ common_gguf_info common_get_gguf_info(const std::string & fname) {
    return info;
 }

+void common_fit_apply(common_params & params, llama_model_params & mparams, llama_context_params & cparams) {
+    if (!params.fit_params) {
+        return;
+    }
+
+    COM_TRC("%s", "fitting params to device memory ...\n");
+    COM_TRC("%s", "(for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n");
+
+    // the draft context is created from the same base params and follows the main context, fit both together
+    const bool has_draft = params.speculative.has_dft();
+    const bool spec_mtp  = std::find(params.speculative.types.begin(), params.speculative.types.end(),
+        COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end();
+
+    common_params params_dft = common_base_params_to_speculative(params);
+
+    auto mparams_dft = common_model_params_to_llama(params_dft);
+    auto cparams_dft = common_context_params_to_llama(params_dft);
+    if (spec_mtp) {
+        cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
+    }
+    cparams_dft.n_rs_seq = 0;
+
+    const common_fit_extra_model extra = {
+        /*.path_model   =*/ params_dft.model.path.c_str(),
+        /*.mparams      =*/ &mparams_dft,
+        /*.cparams      =*/ &cparams_dft,
+        /*.shares_model =*/ !has_draft, // an MTP context runs on the weights of the main model
+    };
+
+    common_fit_params(params.model.path.c_str(), &mparams, &cparams,
+        params.tensor_split,
+        params.tensor_buft_overrides.data(),
+        params.fit_params_target.data(),
+        params.fit_params_min_ctx,
+        has_draft || spec_mtp ? &extra : nullptr,
+        params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
+}
+
 common_init_result::common_init_result(common_params & params, bool model_only) :
    pimpl(new impl{}) {
    auto mparams = common_model_params_to_llama(params);
    auto cparams = common_context_params_to_llama(params);

-    if (params.fit_params) {
-        COM_TRC("%s", "fitting params to device memory ...\n");
-        COM_TRC("%s", "(for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n");
-
-        // the draft context is created from the same base params and follows the main context, fit both together
-        const bool has_draft = params.speculative.has_dft();
-        const bool spec_mtp  = std::find(params.speculative.types.begin(), params.speculative.types.end(),
-            COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end();
-
-        common_params params_dft = common_base_params_to_speculative(params);
-
-        auto mparams_dft = common_model_params_to_llama(params_dft);
-        auto cparams_dft = common_context_params_to_llama(params_dft);
-        if (spec_mtp) {
-            cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
-        }
-        cparams_dft.n_rs_seq = 0;
-
-        const common_fit_extra_model extra = {
-            /*.path_model   =*/ params_dft.model.path.c_str(),
-            /*.mparams      =*/ &mparams_dft,
-            /*.cparams      =*/ &cparams_dft,
-            /*.shares_model =*/ !has_draft, // an MTP context runs on the weights of the main model
-        };
-
-        common_fit_params(params.model.path.c_str(), &mparams, &cparams,
-            params.tensor_split,
-            params.tensor_buft_overrides.data(),
-            params.fit_params_target.data(),
-            params.fit_params_min_ctx,
-            has_draft || spec_mtp ? &extra : nullptr,
-            params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
-    }
+    common_fit_apply(params, mparams, cparams);

    llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams);
    if (model == NULL) {
@@ -1307,13 +1333,17 @@ common_init_result::common_init_result(common_params & params, bool model_only)
    // these decision models return a score for each token via the embeddings output
    // TODO: maybe improve this in the future
    const auto decision_type = common_get_decision_type(model);
+    // minicpm-ocr-jev reads its class scores from the pooled output, the pooling type is the one of the GGUF
    if (decision_type == COMMON_DECISION_TYPE_LAYA || decision_type == COMMON_DECISION_TYPE_KEV || decision_type == COMMON_DECISION_TYPE_CLEF ||
-        decision_type == COMMON_DECISION_TYPE_LFM2_D1_OMNI || decision_type == COMMON_DECISION_TYPE_VLLM_DECISION2) {
+        decision_type == COMMON_DECISION_TYPE_LFM2_D1_OMNI || decision_type == COMMON_DECISION_TYPE_VLLM_DECISION2 ||
+        decision_type == COMMON_DECISION_TYPE_MINICPM_OCR_JEV) {
+        const auto pooling_type = decision_type == COMMON_DECISION_TYPE_MINICPM_OCR_JEV ? LLAMA_POOLING_TYPE_UNSPECIFIED : LLAMA_POOLING_TYPE_NONE;
+
        params.embedding    = true;
-        params.pooling_type = LLAMA_POOLING_TYPE_NONE;
+        params.pooling_type = pooling_type;

        cparams.embeddings            = true;
-        cparams.pooling_type          = LLAMA_POOLING_TYPE_NONE;
+        cparams.pooling_type          = pooling_type;
        cparams.n_outputs_max         = cparams.n_batch;
        cparams.n_outputs_max_per_seq = 1;

@@ -1460,7 +1490,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
        if (params.control_vector_layer_start <= 0) params.control_vector_layer_start = 1;
        if (params.control_vector_layer_end   <= 0) params.control_vector_layer_end   = llama_model_n_layer(model);

-        const auto cvec = common_control_vector_load(params.control_vectors);
+        const auto cvec = common_control_vector_load(params.control_vectors, llama_model_n_layer(model));
        if (cvec.n_embd == -1) {
            return res;
        }
@@ -1982,7 +2012,7 @@ float common_embd_similarity_cos(const float * embd1, const float * embd2, int n
 // Control vector utils
 //

-static common_control_vector_data common_control_vector_load_one(const common_control_vector_load_info & load_info) {
+static common_control_vector_data common_control_vector_load_one(const common_control_vector_load_info & load_info, int32_t n_layer) {
    common_control_vector_data result = { -1, {} };

    ggml_context * ctx = nullptr;
@@ -2023,6 +2053,10 @@ static common_control_vector_data common_control_vector_load_one(const common_co
            COM_ERR("invalid (zero) direction tensor layer index in %s\n", load_info.fname.c_str());
            result.n_embd = -1;
            break;
+        } else if (layer_idx > n_layer) {
+            COM_ERR("direction tensor layer index %d exceeds the model layer count %d in %s\n", layer_idx, n_layer, load_info.fname.c_str());
+            result.n_embd = -1;
+            break;
        }

        struct ggml_tensor * tensor = ggml_get_tensor(ctx, name.c_str());
@@ -2046,10 +2080,10 @@ static common_control_vector_data common_control_vector_load_one(const common_co
        }

        // extend if necessary - do not store data for layer 0 (it's not used)
-        result.data.resize(std::max(result.data.size(), static_cast<size_t>(result.n_embd * layer_idx)), 0.0f);
+        result.data.resize(std::max(result.data.size(), (size_t) result.n_embd * layer_idx), 0.0f);

        const float * src = (const float *) tensor->data;
-        float * dst = result.data.data() + result.n_embd * (layer_idx - 1);  // layer 1 at [0]
+        float * dst = result.data.data() + (size_t) result.n_embd * (layer_idx - 1);  // layer 1 at [0]
        for (int j = 0; j < result.n_embd; j++) {
            dst[j] += src[j] * load_info.strength;  // allows multiple directions for same layer in same file
        }
@@ -2067,11 +2101,11 @@ static common_control_vector_data common_control_vector_load_one(const common_co
    return result;
 }

-common_control_vector_data common_control_vector_load(const std::vector<common_control_vector_load_info> & load_infos) {
+common_control_vector_data common_control_vector_load(const std::vector<common_control_vector_load_info> & load_infos, int32_t n_layer) {
    common_control_vector_data result = { -1, {} };

    for (const auto & info : load_infos) {
-        auto cur = common_control_vector_load_one(info);
+        auto cur = common_control_vector_load_one(info, n_layer);

        if (cur.n_embd == -1) {
            result.n_embd = -1;
@@ -701,6 +701,7 @@ struct common_params {
    std::string models_dir    = "";     // directory containing models for the router server
    std::string models_preset = "";     // directory containing model presets for the router server
    int models_max = 4;                 // maximum number of models to load simultaneously
+    int models_memory_margin = 0;       // MiB of device memory to keep free when loading models (0 = no memory accounting)
    bool models_autoload = true;        // automatically load models when requested via the router server
    std::string models_preset_hf = "";  // show a warning about remote presets on router loaded (if not empty)

@@ -771,6 +772,7 @@ struct common_params {

    // TTS params
    std::string tts_lang = "";
+    std::string tts_voice = "";
    std::string tts_speaker_file = "";

    bool is_gen_docs = false; // whether we are running inside llama-gen-docs
@@ -978,6 +980,9 @@ enum common_decision_type {
    COMMON_DECISION_TYPE_LFM2_D1_OMNI, // same as laya, other prompt layout
    COMMON_DECISION_TYPE_GRANITE_SWITCH, // same as openjev, run by the decision aLoRA of granite switch
    COMMON_DECISION_TYPE_H2O_LIGHTNING, // same as openjev, the labels are read after an answer prefill
+    COMMON_DECISION_TYPE_MINICPM_OCR_JEV, // one score per class, the pooled embeddings output of the last token
+    COMMON_DECISION_TYPE_LFM3_D1,       // same as lfm2-d1, non-causal, label codes A..Z then AA, AB..
+    COMMON_DECISION_TYPE_LFM3_D1_OMNI,  // same as lfm3-d1
    COMMON_DECISION_TYPE_UNKNOWN, // a decision model of a type that is not supported
 };

@@ -1017,6 +1022,9 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
 struct llama_model_params   common_model_params_to_llama  (      common_params & params);
 struct llama_context_params common_context_params_to_llama(const common_params & params);

+// apply --fit to mparams and cparams, no-op if params.fit_params is false
+void common_fit_apply(common_params & params, struct llama_model_params & mparams, struct llama_context_params & cparams);
+
 // clear LoRA adapters from context, then apply new list of adapters
 void common_set_adapter_lora(struct llama_context * ctx, std::vector<common_adapter_lora_info> & lora);

@@ -1229,8 +1237,9 @@ struct common_control_vector_load_info {
 };

 // Load control vectors, scale each by strength, and add them together.
+// Direction layer indices must be in [1, n_layer].
 // On error, returns {-1, empty}
-common_control_vector_data common_control_vector_load(const std::vector<common_control_vector_load_info> & load_infos);
+common_control_vector_data common_control_vector_load(const std::vector<common_control_vector_load_info> & load_infos, int32_t n_layer);

 //
 // Split utils
@@ -1,6 +1,8 @@
 #include "fit.h"

+#include "common.h"
 #include "json.h"
+#include "llama-cpp.h"
 #include "log.h"

 #include "../src/llama-ext.h"
@@ -27,6 +29,32 @@ class common_params_fit_exception : public std::runtime_error {
    using std::runtime_error::runtime_error;
 };

+// RAII: send llama logs below min_level to the debug log
+struct common_fit_log_scope {
+    common_fit_log_scope(ggml_log_level min_level) : min_level(min_level) {
+        llama_log_get(&original_logger.callback, &original_logger.user_data);
+        llama_log_set([](ggml_log_level level, const char * text, void * user_data) {
+            const common_fit_log_scope * self = (const common_fit_log_scope *) user_data;
+            const ggml_log_level level_eff = level >= self->min_level ? level : GGML_LOG_LEVEL_DEBUG;
+            self->original_logger.callback(level_eff, text, self->original_logger.user_data);
+        }, this);
+    }
+
+    ~common_fit_log_scope() {
+        llama_log_set(original_logger.callback, original_logger.user_data);
+    }
+
+    common_fit_log_scope(const common_fit_log_scope &) = delete;
+    common_fit_log_scope & operator=(const common_fit_log_scope &) = delete;
+
+private:
+    struct {
+        ggml_log_callback callback;
+        void * user_data;
+    } original_logger;
+    ggml_log_level min_level;
+};
+
 static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(
        const char * path_model,
        const llama_model_params * mparams,
@@ -36,22 +64,7 @@ static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(
        uint32_t & hp_n_ctx_train,
        uint32_t & hp_n_expert,
        ggml_log_level log_level) {
-    struct user_data_t {
-        struct {
-            ggml_log_callback callback;
-            void * user_data;
-        } original_logger;
-        ggml_log_level min_level; // prints below this log level go to debug log
-    };
-    user_data_t ud;
-    llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data);
-    ud.min_level = log_level;
-
-    llama_log_set([](ggml_log_level level, const char * text, void * user_data) {
-        const user_data_t * ud = (const user_data_t *) user_data;
-        const ggml_log_level level_eff = level >= ud->min_level ? level : GGML_LOG_LEVEL_DEBUG;
-        ud->original_logger.callback(level_eff, text, ud->original_logger.user_data);
-    }, &ud);
+    common_fit_log_scope log_scope(log_level);

    llama_model_params mparams_copy = *mparams;
    mparams_copy.no_alloc  = true;
@@ -59,14 +72,12 @@ static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(

    llama_model * model = llama_model_load_from_file(path_model, mparams_copy);
    if (model == nullptr) {
-        llama_log_set(ud.original_logger.callback, ud.original_logger.user_data);
        throw std::runtime_error("failed to load model");
    }

    llama_context * ctx = llama_init_from_model(model, *cparams);
    if (ctx == nullptr) {
        llama_model_free(model);
-        llama_log_set(ud.original_logger.callback, ud.original_logger.user_data);
        throw std::runtime_error("failed to create llama_context from model");
    }

@@ -148,11 +159,81 @@ static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(

    llama_free(ctx);
    llama_model_free(model);
-    llama_log_set(ud.original_logger.callback, ud.original_logger.user_data);

    return ret;
 }

+std::vector<common_memory_breakdown_row> common_fit_measure(
+        const char * path_model,
+        const llama_model_params * mparams,
+        const llama_context_params * cparams,
+        uint32_t * n_ctx_train,
+        uint32_t * n_ctx,
+        ggml_log_level log_level) {
+    common_fit_log_scope log_scope(log_level);
+
+    llama_model_params mparams_copy = *mparams;
+    mparams_copy.no_alloc  = true;
+    mparams_copy.load_mode = LLAMA_LOAD_MODE_NONE;
+
+    llama_model_ptr model(llama_model_load_from_file(path_model, mparams_copy));
+    if (!model) {
+        throw std::runtime_error("failed to load model");
+    }
+
+    llama_context_ptr ctx(llama_init_from_model(model.get(), *cparams));
+    if (!ctx) {
+        throw std::runtime_error("failed to create llama_context from model");
+    }
+
+    if (n_ctx_train) {
+        *n_ctx_train = llama_model_n_ctx_train(model.get());
+    }
+    if (n_ctx) {
+        *n_ctx = llama_n_ctx(ctx.get());
+    }
+
+    return common_memory_breakdown_rows(ctx.get());
+}
+
+common_memory_reports common_measure_memory(common_params & params) {
+    const ggml_log_level log_level = params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR;
+    const char * path_model = params.model.path.c_str();
+
+    common_memory_reports res;
+
+    auto mparams = common_model_params_to_llama(params);
+    auto cparams = common_context_params_to_llama(params);
+
+    // max: what --fit picks now, no-op with -fit off
+    {
+        auto mparams_fit = mparams;
+        auto cparams_fit = cparams;
+        common_fit_apply(params, mparams_fit, cparams_fit);
+        res.info_max = common_fit_measure(path_model, &mparams_fit, &cparams_fit, &res.n_ctx_train, &res.n_ctx_max, log_level);
+    }
+
+    // min: the smallest context --fit may pick (--fit only changes n_ctx when it is 0)
+    uint32_t n_ctx_min = cparams.n_ctx;
+    if (n_ctx_min == 0 && params.fit_params) {
+        const uint32_t n_streams = cparams.kv_unified ? 1 : std::max<uint32_t>(1, cparams.n_seq_max);
+        // -c 0 sets fit_params_min_ctx to UINT32_MAX, cap at n_ctx_train like --fit
+        n_ctx_min = (uint32_t) std::min<uint64_t>(
+            uint64_t(params.fit_params_min_ctx) * n_streams,
+            uint64_t(res.n_ctx_train) * n_streams);
+    }
+    if (n_ctx_min == 0 || n_ctx_min >= res.n_ctx_max) {
+        // nothing to reduce, or -fit off with the trained context
+        res.n_ctx_min = res.n_ctx_max;
+        res.info_min  = res.info_max;
+    } else {
+        cparams.n_ctx = n_ctx_min;
+        res.info_min  = common_fit_measure(path_model, &mparams, &cparams, nullptr, &res.n_ctx_min, log_level);
+    }
+
+    return res;
+}
+
 common_device_memory_data_vec common_get_device_memory_data(
        const char * path_model,
        const llama_model_params * mparams,
@@ -903,8 +984,7 @@ enum common_params_fit_status common_fit_params(
    return status;
 }

-void common_memory_breakdown_print(const struct llama_context * ctx) {
-    //const auto & devices = ctx->get_model().devices;
+std::vector<common_memory_breakdown_row> common_memory_breakdown_rows(const struct llama_context * ctx) {
    const auto * model = llama_get_model(ctx);

    std::vector<ggml_backend_dev_t> devices;
@@ -914,18 +994,6 @@ void common_memory_breakdown_print(const struct llama_context * ctx) {

    llama_memory_breakdown memory_breakdown = llama_get_memory_breakdown(ctx);

-    std::vector<std::array<std::string, 9>> table_data;
-    table_data.reserve(devices.size());
-
-    // same data as the table below, for --log-jsonl consumers
-    common_json rows = common_json::array();
-    const std::string template_header = "%s: | %s | %s   %s    %s   %s   %s   %s    %s |\n";
-    const std::string template_gpu    = "%s: | %s | %s = %s + (%s = %s + %s + %s) + %s |\n";
-    const std::string template_other  = "%s: | %s | %s   %s    %s = %s + %s + %s    %s |\n";
-
-    table_data.push_back({template_header, "memory breakdown [MiB]", "total", "free", "self", "model", "context", "compute", "unaccounted"});
-
-    constexpr size_t MiB = 1024 * 1024;
    const std::vector<std::string> desc_prefixes_strip = {"NVIDIA ", "GeForce ", "Tesla ", "AMD ", "Radeon ", "Instinct "};

    // track seen buffer types to avoid double counting:
@@ -964,104 +1032,130 @@ void common_memory_breakdown_print(const struct llama_context * ctx) {
        }
    }

-    // print memory breakdown for each device:
+    std::vector<common_memory_breakdown_row> rows;
+    rows.reserve(devices.size() + 1 + memory_breakdown.size());
+
    for (size_t i = 0; i < devices.size(); i++) {
        ggml_backend_dev_t dev = devices[i];
-        llama_memory_breakdown_data mb = mb_dev[i];

-        const std::string name = ggml_backend_dev_name(dev);
-        std::string desc = ggml_backend_dev_description(dev);
+        common_memory_breakdown_row row;
+        row.kind        = common_memory_breakdown_row::DEVICE;
+        row.name        = ggml_backend_dev_name(dev);
+        row.description = ggml_backend_dev_description(dev);
        for (const std::string & prefix : desc_prefixes_strip) {
-            if (desc.length() >= prefix.length() && desc.substr(0, prefix.length()) == prefix) {
-                desc = desc.substr(prefix.length());
+            if (row.description.length() >= prefix.length() && row.description.substr(0, prefix.length()) == prefix) {
+                row.description = row.description.substr(prefix.length());
            }
        }

        size_t free, total;
        ggml_backend_dev_memory(dev, &free, &total);
-
-        const size_t self = mb.model + mb.context + mb.compute;
-        const int64_t unaccounted = static_cast<int64_t>(total) - static_cast<int64_t>(free) - static_cast<int64_t>(self);
-
-        table_data.push_back({
-            template_gpu,
-            "  - " + name + " (" + desc + ")",
-            std::to_string(total / MiB),
-            std::to_string(free / MiB),
-            std::to_string(self / MiB),
-            std::to_string(mb.model / MiB),
-            std::to_string(mb.context / MiB),
-            std::to_string(mb.compute / MiB),
-            std::to_string(unaccounted / static_cast<int64_t>(MiB))});
-
-        rows.push_back({
-            {"kind",        "device"},
-            {"name",        name},
-            {"description", desc},
-            {"total",       total / MiB},
-            {"free",        free / MiB},
-            {"self",        self / MiB},
-            {"model",       mb.model / MiB},
-            {"context",     mb.context / MiB},
-            {"compute",     mb.compute / MiB},
-            {"unaccounted", unaccounted / static_cast<int64_t>(MiB)},
-        });
+        row.total   = total;
+        row.free    = free;
+        row.model   = mb_dev[i].model;
+        row.context = mb_dev[i].context;
+        row.compute = mb_dev[i].compute;
+        rows.push_back(std::move(row));
    }

-    // print memory breakdown for host:
    {
-        const size_t self = mb_host.model + mb_host.context + mb_host.compute;
-        table_data.push_back({
-            template_other,
-            "  - Host",
-            "", // total
-            "", // free
-            std::to_string(self / MiB),
-            std::to_string(mb_host.model / MiB),
-            std::to_string(mb_host.context / MiB),
-            std::to_string(mb_host.compute / MiB),
-            ""}); // unaccounted
-
-        rows.push_back({
-            {"kind",    "host"},
-            {"name",    "Host"},
-            {"self",    self / MiB},
-            {"model",   mb_host.model / MiB},
-            {"context", mb_host.context / MiB},
-            {"compute", mb_host.compute / MiB},
-        });
+        common_memory_breakdown_row row;
+        row.kind    = common_memory_breakdown_row::HOST;
+        row.name    = "Host";
+        row.model   = mb_host.model;
+        row.context = mb_host.context;
+        row.compute = mb_host.compute;
+        rows.push_back(std::move(row));
    }

-    // print memory breakdown for all remaining buffer types:
+    // all remaining buffer types:
    for (const auto & buft_mb : memory_breakdown) {
        ggml_backend_buffer_type_t          buft = buft_mb.first;
        const llama_memory_breakdown_data & mb   = buft_mb.second;
        if (seen_buffer_types.count(buft) == 1) {
            continue;
        }
-        const std::string name = ggml_backend_buft_name(buft);
-        const size_t self = mb.model + mb.context + mb.compute;
+        common_memory_breakdown_row row;
+        row.kind    = common_memory_breakdown_row::BUFFER_TYPE;
+        row.name    = ggml_backend_buft_name(buft);
+        row.model   = mb.model;
+        row.context = mb.context;
+        row.compute = mb.compute;
+        rows.push_back(std::move(row));
+
+        seen_buffer_types.insert(buft);
+    }
+
+    return rows;
+}
+
+void common_memory_breakdown_print(const struct llama_context * ctx) {
+    const std::vector<common_memory_breakdown_row> rows = common_memory_breakdown_rows(ctx);
+
+    std::vector<std::array<std::string, 9>> table_data;
+    table_data.reserve(rows.size() + 1);
+
+    // same data as the table below, for --log-jsonl consumers
+    common_json json_rows = common_json::array();
+    const std::string template_header = "%s: | %s | %s   %s    %s   %s   %s   %s    %s |\n";
+    const std::string template_gpu    = "%s: | %s | %s = %s + (%s = %s + %s + %s) + %s |\n";
+    const std::string template_other  = "%s: | %s | %s   %s    %s = %s + %s + %s    %s |\n";
+
+    table_data.push_back({template_header, "memory breakdown [MiB]", "total", "free", "self", "model", "context", "compute", "unaccounted"});
+
+    constexpr size_t MiB = 1024 * 1024;
+
+    for (const auto & row : rows) {
+        const size_t self = row.model + row.context + row.compute;
+
+        if (row.kind == common_memory_breakdown_row::DEVICE) {
+            const int64_t unaccounted = row.total - row.free - static_cast<int64_t>(self);
+
+            table_data.push_back({
+                template_gpu,
+                "  - " + row.name + " (" + row.description + ")",
+                std::to_string(row.total / MiB),
+                std::to_string(row.free / MiB),
+                std::to_string(self / MiB),
+                std::to_string(row.model / MiB),
+                std::to_string(row.context / MiB),
+                std::to_string(row.compute / MiB),
+                std::to_string(unaccounted / static_cast<int64_t>(MiB))});
+
+            json_rows.push_back({
+                {"kind",        "device"},
+                {"name",        row.name},
+                {"description", row.description},
+                {"total",       row.total / MiB},
+                {"free",        row.free / MiB},
+                {"self",        self / MiB},
+                {"model",       row.model / MiB},
+                {"context",     row.context / MiB},
+                {"compute",     row.compute / MiB},
+                {"unaccounted", unaccounted / static_cast<int64_t>(MiB)},
+            });
+            continue;
+        }
+
        table_data.push_back({
            template_other,
-            "  - " + name,
+            "  - " + row.name,
            "", // total
            "", // free
            std::to_string(self / MiB),
-            std::to_string(mb.model / MiB),
-            std::to_string(mb.context / MiB),
-            std::to_string(mb.compute / MiB),
+            std::to_string(row.model / MiB),
+            std::to_string(row.context / MiB),
+            std::to_string(row.compute / MiB),
            ""}); // unaccounted

-        rows.push_back({
-            {"kind",    "buffer_type"},
-            {"name",    name},
+        json_rows.push_back({
+            {"kind",    row.kind == common_memory_breakdown_row::HOST ? "host" : "buffer_type"},
+            {"name",    row.name},
            {"self",    self / MiB},
-            {"model",   mb.model / MiB},
-            {"context", mb.context / MiB},
-            {"compute", mb.compute / MiB},
+            {"model",   row.model / MiB},
+            {"context", row.context / MiB},
+            {"compute", row.compute / MiB},
        });
-
-        seen_buffer_types.insert(buft);
    }

    for (size_t j = 1; j < table_data[0].size(); j++) {
@@ -1081,7 +1175,7 @@ void common_memory_breakdown_print(const struct llama_context * ctx) {

    LOG_JSON("fit_memory_breakdown", common_json({
        {"unit", "MiB"},
-        {"rows", rows},
+        {"rows", json_rows},
    }));
 }

@@ -3,6 +3,7 @@
 #include "ggml.h"
 #include "llama.h"

+#include <string>
 #include <vector>

 enum common_params_fit_status {
@@ -45,6 +46,47 @@ void common_fit_print(

 void common_memory_breakdown_print(const llama_context * ctx);

+// one row of the memory breakdown: a device, the host or another buffer type
+struct common_memory_breakdown_row {
+    enum kind_t { DEVICE, HOST, BUFFER_TYPE };
+
+    kind_t      kind = DEVICE;
+    std::string name;        // device or buffer type name, "Host" for the host row
+    std::string description; // device description, empty otherwise
+    int64_t     total = 0;   // device memory in bytes, 0 unless kind == DEVICE
+    int64_t     free  = 0;   // free device memory in bytes, 0 unless kind == DEVICE
+    size_t      model   = 0; // bytes held by this context, by kind of buffer
+    size_t      context = 0;
+    size_t      compute = 0;
+};
+
+// memory a model needs at the min and max context --fit may pick
+struct common_memory_reports {
+    uint32_t n_ctx_train = 0;
+    uint32_t n_ctx_max   = 0; // the largest context that fits the free device memory right now
+    uint32_t n_ctx_min   = 0; // --fit-ctx (the smallest context --fit may pick), with the placement set in params
+    std::vector<common_memory_breakdown_row> info_max; // breakdown at n_ctx_max
+    std::vector<common_memory_breakdown_row> info_min; // breakdown at n_ctx_min
+};
+
+// rows of common_memory_breakdown_print: model devices first, then host
+std::vector<common_memory_breakdown_row> common_memory_breakdown_rows(const llama_context * ctx);
+
+// load model + context with no_alloc and return the memory breakdown, params are used as-is
+// n_ctx_train and n_ctx are optional outputs; throws on failure
+std::vector<common_memory_breakdown_row> common_fit_measure(
+                         const char * path_model,
+           const llama_model_params * mparams,
+         const llama_context_params * cparams,
+                           uint32_t * n_ctx_train,
+                           uint32_t * n_ctx,
+                     ggml_log_level   log_level);
+
+struct common_params;
+
+// measure min and max without allocating, params may be modified as with --fit; throws on failure
+common_memory_reports common_measure_memory(common_params & params);
+
 struct common_device_memory_data {
    int64_t total;
    int64_t free;
@@ -6,7 +6,9 @@
 // we can remove it in the future when we figure out a better way to define inputs using jinja::value
 #include "json.h"

+#include <algorithm>
 #include <functional>
+#include <optional>
 #include <sstream>

 #define FILENAME "jinja-caps"
@@ -184,6 +186,10 @@ caps caps_get(jinja::program & prog) {
                    {
                        {"role", "user"},
                        {"content", json::array({
+                            {
+                                {"type", "text"},
+                                {"text", content_marker}
+                            }
                        })}
                    }
                });
@@ -584,6 +590,58 @@ caps caps_get(jinja::program & prog) {
        }
    );

+    if (result.supports_reasoning_effort) {
+        JJ_DEBUG("%s\n", ">>> Running capability check: reasoning effort levels");
+
+        static const std::vector<std::string> candidates = { "minimal", "low", "medium", "high", "xhigh", "max" };
+
+        // renders a single user turn with the given effort, nothing when the template raises
+        const auto render = [&](const std::string & effort) {
+            std::optional<std::string> rendered;
+            caps_try_execute(
+                prog,
+                [&]() {
+                    // messages
+                    return json::array({
+                        {
+                            {"role", "user"},
+                            {"content", "User message"}
+                        },
+                    });
+                },
+                [&](context & ctx) {
+                    ctx.set_val("enable_thinking", mk_val<value_bool>(true));
+                    caps_apply_reasoning_effort(ctx, effort);
+                },
+                nullptr, // tools_fn
+                [&](context &, bool success, value &, value &, const std::string & output) {
+                    if (success) {
+                        rendered = output;
+                    }
+                }
+            );
+            return rendered;
+        };
+
+        // an unknown level shows the render of the template fallback, a level that renders the same is not honored
+        // levels sharing a render are aliases and the highest one is kept
+        const auto unknown = render("REASONING_EFFORT_UNKNOWN");
+        std::vector<std::string> renders;
+        for (const auto & level : candidates) {
+            const auto rendered = render(level);
+            if (!rendered || rendered == unknown) {
+                continue;
+            }
+            const auto it = std::find(renders.begin(), renders.end(), *rendered);
+            if (it != renders.end()) {
+                result.reasoning_effort_levels[it - renders.begin()] = level;
+            } else {
+                renders.push_back(*rendered);
+                result.reasoning_effort_levels.push_back(level);
+            }
+        }
+    }
+
    JJ_DEBUG("%s\n", result.to_string().c_str());

    return result;
@@ -4,6 +4,7 @@

 #include <string>
 #include <map>
+#include <vector>

 namespace jinja {

@@ -19,6 +20,9 @@ struct caps {
    // supports reasoning effort levels
    bool supports_reasoning_effort = false;

+    // reasoning effort levels the template renders distinctly, from lowest to highest
+    std::vector<std::string> reasoning_effort_levels;
+
    // one of the 2 content capabilities must be true
    bool supports_string_content = true;
    bool supports_typed_content = false;
@@ -76,9 +76,10 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat

        auto reasoning = p.eps();
        if (supports_reasoning && extract_reasoning) {
+            // The template strips the newline before </think>, so leave it out of the reasoning
            reasoning = p.optional("<think>" + p.space() +
-                                   p.reasoning(p.until_one_of({ "</think>", "<tool_call>" })) +
-                                   (p.literal("</think>") | p.peek(p.literal("<tool_call>"))));
+                                   p.reasoning(p.until_one_of({ "\n</think>", "</think>", "<tool_call>" })) +
+                                   (p.literal("\n</think>") | p.literal("</think>") | p.peek(p.literal("<tool_call>"))));
        }

        // Response format parser
@@ -166,8 +167,15 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
            auto calls      = inputs.parallel_tool_calls ? tool_call_first + p.zero_or_more(tool_call) : tool_call_first;
            auto tool_calls = p.trigger_rule("tool-call-root", p.repeat(calls, min_calls, 1));

+            // The templates separate content from the tool calls with one or two newlines,
+            // stop the content before them
+            std::vector<std::string> content_ends;
+            for (const auto & start : tool_call_starts) {
+                content_ends.insert(content_ends.end(), { "\n\n" + start, "\n" + start, start });
+            }
+
            return generation_prompt +
-                   (reasoning << p.content(p.until_one_of(tool_call_starts)) << tool_calls);
+                   (reasoning << p.content(p.until_one_of(content_ends)) << tool_calls);
        }

        // Content only parser
@@ -721,8 +721,9 @@ struct parser_executor {
                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos, {}, std::move(invalid_utf8));
            }

-            if (match == common_trie::PARTIAL_MATCH) {
-                // Found a partial match extending to end of input, return everything before it
+            if (match == common_trie::PARTIAL_MATCH && !ctx.is_complete()) {
+                // Found a partial match extending to end of input, return everything before it.
+                // On a complete input the delimiter can no longer form, so the bytes belong to the text
                return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start_pos, pos, {}, std::move(invalid_utf8));
            }

@@ -162,9 +162,10 @@ struct common_peg_parse_result {
 };

 enum common_peg_parse_flags {
-    COMMON_PEG_PARSE_FLAG_NONE    = 0,
-    COMMON_PEG_PARSE_FLAG_LENIENT = 1 << 0,
-    COMMON_PEG_PARSE_FLAG_DEBUG   = 1 << 1,
+    COMMON_PEG_PARSE_FLAG_NONE     = 0,
+    COMMON_PEG_PARSE_FLAG_LENIENT  = 1 << 0,
+    COMMON_PEG_PARSE_FLAG_DEBUG    = 1 << 1,
+    COMMON_PEG_PARSE_FLAG_COMPLETE = 1 << 2,
 };

 inline common_peg_parse_flags operator|(common_peg_parse_flags a, common_peg_parse_flags b) {
@@ -209,6 +210,7 @@ struct common_peg_parse_context {

    bool is_lenient() const { return flags & COMMON_PEG_PARSE_FLAG_LENIENT; }
    bool is_debug() const { return flags & COMMON_PEG_PARSE_FLAG_DEBUG; }
+    bool is_complete() const { return flags & COMMON_PEG_PARSE_FLAG_COMPLETE; }
 };

 class common_peg_arena;
@@ -172,6 +172,11 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "D1OmniModel": "lfm2",
    "Lfm2Model": "lfm2",
    "Lfm2MoeForCausalLM": "lfm2",
+    "Lfm3D1Model": "lfm3",
+    "Lfm3ForCausalLM": "lfm3",
+    "Lfm3OmniForConditionalGeneration": "lfm3",
+    "Lfm3VlForConditionalGeneration": "lfm3",
+    "LimiteForCausalLM": "limite",
    "Llama4ForCausalLM": "llama",
    "Llama4ForConditionalGeneration": "llama",
    "LlamaBidirectionalModel": "llama",
@@ -197,6 +202,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "MiniCPMForCausalLM": "minicpm",
    "MiniCPMV4_6ForConditionalGeneration": "minicpm",
    "MiniCPMV4_7ForConditionalGeneration": "minicpm",
+    "MiniCPMOCRForSequenceClassification": "qwen",
    "MiniCPMOcrForConditionalGeneration": "minicpm",
    "MiniMaxText01ForCausalLM": "minimax",
    "MiniMaxM1ForCausalLM": "minimax",
@@ -357,6 +363,9 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "Lfm2AudioForConditionalGeneration": "lfm2",
    "D1OmniModel": "lfm2",
    "Lfm2VlForConditionalGeneration": "lfm2",
+    "Lfm3AudioForConditionalGeneration": "lfm3",
+    "Lfm3OmniForConditionalGeneration": "lfm3",
+    "Lfm3VlForConditionalGeneration": "lfm3",
    "LightOnOCRForConditionalGeneration": "lighton_ocr",
    "Llama4ForConditionalGeneration": "llama4",
    "LlavaForConditionalGeneration": "llava",
@@ -365,6 +374,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "MiniMaxM3SparseForConditionalGeneration": "minimax",
    "MiniCPMV4_6ForConditionalGeneration": "minicpm",
    "MiniCPMV4_7ForConditionalGeneration": "minicpm",
+    "MiniCPMOCRForSequenceClassification": "qwen3vl",
    "MiniCPMOcrForConditionalGeneration": "minicpm",
    "Mistral3ForConditionalGeneration": "llava",
    "Mistral3Model": "llava",
@@ -15,25 +15,30 @@ from .base import MmprojModel, ModelBase, TextModel, gguf, jinja_str_or_json, lo
 from .gemma import ConformerAudioModel


+def lfm2_feed_forward_length(hparams: dict[str, Any]) -> int:
+    ff_dim = hparams.get("block_ff_dim", hparams.get("intermediate_size"))
+    assert ff_dim is not None
+    auto_adjust_ff_dim = hparams["block_auto_adjust_ff_dim"]
+    ffn_dim_multiplier = hparams["block_ffn_dim_multiplier"]
+    multiple_of = hparams["block_multiple_of"]
+
+    if auto_adjust_ff_dim:
+        ff_dim = int(2 * ff_dim / 3)
+        # custom dim factor multiplier
+        if ffn_dim_multiplier is not None:
+            ff_dim = int(ffn_dim_multiplier * ff_dim)
+        ff_dim = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of)
+
+    return ff_dim
+
+
@ModelBase.register("Lfm2ForCausalLM", "LFM2ForCausalLM")
@ModelBase.example("LiquidAI/LFM2-1.2B", "LiquidAI/LFM2.5-350M")
 class LFM2Model(TextModel):
    model_arch = gguf.MODEL_ARCH.LFM2

    def _add_feed_forward_length(self):
-        ff_dim = self.find_hparam(["block_ff_dim", "intermediate_size"])
-        auto_adjust_ff_dim = self.hparams["block_auto_adjust_ff_dim"]
-        ffn_dim_multiplier = self.hparams["block_ffn_dim_multiplier"]
-        multiple_of = self.hparams["block_multiple_of"]
-
-        if auto_adjust_ff_dim:
-            ff_dim = int(2 * ff_dim / 3)
-            # custom dim factor multiplier
-            if ffn_dim_multiplier is not None:
-                ff_dim = int(ffn_dim_multiplier * ff_dim)
-            ff_dim = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of)
-
-        self.gguf_writer.add_feed_forward_length(ff_dim)
+        self.gguf_writer.add_feed_forward_length(lfm2_feed_forward_length(self.hparams))

    def set_gguf_parameters(self):
        # set num_key_value_heads only for attention layers
@@ -72,7 +77,9 @@ def _is_d1_checkpoint(dir_model: Path) -> bool:
    if not (dir_model / "config.json").is_file():
        return False
    with open(dir_model / "config.json", encoding="utf-8") as f:
-        return json.load(f).get("auto_map", {}).get("AutoModel", "").endswith(".D1Model")
+        config = json.load(f)
+    # d1 on LFM3 is converted by conversion/lfm3.py
+    return config.get("auto_map", {}).get("AutoModel", "").endswith(".D1Model") and not config.get("model_type", "").startswith("lfm3")


@ModelBase.register_hparams_loader(_is_d1_checkpoint)
@@ -0,0 +1,632 @@
+from __future__ import annotations
+
+import json
+import math
+
+from pathlib import Path
+from typing import Any, Callable, Iterable, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import MmprojModel, ModelBase, TextModel, gguf, logger
+from .lfm2 import D1Model, lfm2_feed_forward_length
+
+
+@ModelBase.register("Lfm3ForCausalLM", "Lfm3VlForConditionalGeneration", "Lfm3OmniForConditionalGeneration")
+@ModelBase.example("LiquidAI/LFM3-1.7B", "LiquidAI/LFM3-1.7B-Base")
+class LFM3Model(TextModel):
+    model_arch = gguf.MODEL_ARCH.LFM3
+
+    # LFM3-Audio/Omni: encoder, projector and audio-output heads live next to the text model
+    _media_prefixes = ("encoder.", "mlp.", "codebook_embedding.", "depth_embeddings.", "depth_linear.", "depthformer.",
+                       "mask_embedding", "model.codebook_embedding.", "model.depth_decoder.")
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+
+        if name.startswith(cls._media_prefixes):
+            return None
+
+        name = name.removeprefix("lfm.")
+
+        return super().filter_tensors((name, gen))
+
+    # LFM3-Audio live ASR prompt of the model card, the audio follows in the assistant turn
+    _asr_template = (
+        "{{- bos_token -}}"
+        "{{- '<|im_start|>system\\nPerform live ASR.<|im_end|>\\n' -}}"
+        "{%- if add_generation_prompt -%}{{- '<|im_start|>assistant\\n' -}}{%- endif -%}"
+    )
+
+    # TTS prompt of the model card: a language code is looked up, anything else is taken as the language name
+    _tts_langs = {
+        "ar": "Arabic", "da": "Danish", "de": "German", "el": "Greek", "en": "American English", "en-us": "American English",
+        "en-gb": "British English", "fi": "Finnish", "fr": "French", "he": "Hebrew", "it": "Italian", "ja": "Japanese",
+        "ko": "Korean", "nl": "Dutch", "pl": "Polish", "pt": "Portuguese", "es": "Spanish", "sv": "Swedish", "tr": "Turkish",
+        "zh": "Chinese",
+    }
+
+    # "language" and "voice" are optional, the text to speak is the first message
+    def _tts_template(self, think: bool) -> str:
+        return (
+            "{%- set names = " + json.dumps(self._tts_langs) + " -%}"
+            "{%- set code = (language | lower) if language is defined and language else 'en' -%}"
+            "{%- set lang = names[code] if code in names else language -%}"
+            "{%- set v = (voice | lower) if voice is defined and voice and (voice | lower) in ['female', 'male'] else 'female' -%}"
+            "{{- bos_token -}}"
+            "{{- '<|im_start|>system\\nPerform TTS in ' + lang + ' using a ' + v + ' voice.<|im_end|>\\n' -}}"
+            "{{- '<|im_start|>user\\n' + messages[0]['content'] + '<|im_end|>\\n' -}}"
+            "{{- '<|im_start|>assistant\\n" + ("<think></think>" if think else "") + "<|audio_start|>' -}}"
+        )
+
+    def set_vocab(self):
+        super().set_vocab()
+        # hparams hold the merged text config, the checkpoint's own model_type is in config.json
+        with open(self.dir_model / "config.json", encoding="utf-8") as f:
+            model_type = json.load(f).get("model_type")
+        is_audio = model_type == "lfm3_audio"
+        is_omni  = model_type == "lfm3_omni"
+        templates = []
+        if is_audio:
+            templates.append({"name": "asr", "template": self._asr_template})
+        if is_audio or is_omni:
+            # LFM3-Omni opens every assistant turn with an empty think block
+            templates.append({"name": "tts", "template": self._tts_template(think=is_omni)})
+        if templates:
+            self.gguf_writer.add_chat_template(templates)
+
+    def set_gguf_parameters(self):
+        layer_types = self.hparams["layer_types"]
+        n_kv_heads = self.hparams["num_key_value_heads"]
+
+        # set num_key_value_heads only for attention layers
+        self.hparams["num_key_value_heads"] = [
+            n_kv_heads if lt != "conv" else 0
+            for lt in layer_types
+        ]
+
+        super().set_gguf_parameters()
+
+        self.gguf_writer.add_rope_dimension_count(self.hparams["head_dim"])
+
+        self.gguf_writer.add_sliding_window_pattern([lt == "sliding_attention" for lt in layer_types])
+        self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])
+
+        self.gguf_writer.add_shortconv_l_cache(self.hparams["conv_L_cache"])
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # Partial RoPE is represented by masking trailing frequency pairs.
+        # rope_parameters is keyed by layer type; only full_attention layers use partial RoPE.
+        rope_parameters = self.hparams["rope_parameters"]["full_attention"]
+        partial = rope_parameters.get("partial_rotary_factor", 1.0)
+        head_dim = self.hparams["head_dim"]
+        rope_angles = int(partial * head_dim) // 2
+        n_pairs = head_dim // 2
+        assert 0 < rope_angles <= n_pairs
+        values = [1.0] * rope_angles + [1e30] * (n_pairs - rope_angles)
+        rope_freqs = torch.tensor(values, dtype=torch.float32)
+        yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), rope_freqs)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # conv op requires 2d tensor
+        if "conv.conv" in name:
+            data_torch = data_torch.squeeze(1)
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
+# d1 on LFM3-VL / LFM3-Omni: the stock weights, the d1 prompt read at its last token
+_D1_TYPES = {"lfm3_vl": gguf.DecisionType.LFM3_D1, "lfm3_omni": gguf.DecisionType.LFM3_D1_OMNI}
+
+
+def _is_lfm3_d1_checkpoint(dir_model: Path) -> bool:
+    if not (dir_model / "config.json").is_file():
+        return False
+    with open(dir_model / "config.json", encoding="utf-8") as f:
+        config = json.load(f)
+    return config.get("architectures") == ["D1Model"] and config.get("model_type") in _D1_TYPES
+
+
+@ModelBase.register_hparams_loader(_is_lfm3_d1_checkpoint)
+def _load_lfm3_d1_hparams(dir_model: Path) -> dict[str, Any]:
+    logger.info("gguf: detected LFM3 d1 checkpoint")
+    hparams = ModelBase.load_hparams(dir_model, False, guess=False)
+    hparams["d1_family"] = hparams["model_type"]
+    # the mmproj is the base model's, the text model writes the decision
+    hparams["architectures"] = ["Lfm3VlForConditionalGeneration" if hparams["model_type"] == "lfm3_vl" else "Lfm3OmniForConditionalGeneration"]
+    hparams["text_config"]["architectures"] = ["Lfm3D1Model"]
+    return hparams
+
+
+@ModelBase.register("Lfm3D1Model")
+@ModelBase.example("LiquidAI/d1-2B", "LiquidAI/d1-omni-2B")
+class LFM3D1Model(LFM3Model):
+    model_arch = gguf.MODEL_ARCH.LFM3
+
+    # how the checkpoint serves, hard-coded in the server: assert the checkpoint records the same
+    _served = {"d1_codes": 2, "d1_attention": "set"}
+
+    def set_vocab(self):
+        # the vocab without the speech task templates of LFM3-Audio/Omni, d1 only decides
+        TextModel.set_vocab(self)
+        # the d1 prompt of LFM2, the LFM3 template writes an empty thought before the answer
+        self.gguf_writer.add_chat_template([{"name": "systemone", "template": D1Model._systemone_template() + "<think></think>"}])
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+
+        for key, value in self._served.items():
+            if self.hparams.get(key) != value:
+                raise ValueError(f"unsupported d1 checkpoint: {key} = {self.hparams.get(key)!r}, expected {value!r}")
+
+        # a lone question reads the whole prompt both ways on the full attention layers
+        self.gguf_writer.add_causal_attention(False)
+        self.gguf_writer.add_decision_type(_D1_TYPES[self.hparams["d1_family"]])
+
+        # the option answers share one temperature
+        with open(self.dir_model / "calibration.json", encoding="utf-8") as f:
+            calibration = json.load(f)
+        if calibration.get("temperatures"):
+            raise ValueError("unsupported d1 calibration: per-bucket temperatures")
+        for name in ("noul", "choice", "score"):
+            self.gguf_writer.add_decision_temperature(name, calibration["default"])
+
+
+class LFM3VLVisionMixin(MmprojModel):
+    """LFM3-VL vision tower and projector, shared with LFM3-Omni."""
+
+    def _set_lfm3vl_gguf_parameters(self):
+        assert self.hparams_vision is not None
+        # clip.cpp hardcodes the GELU projector and its input norm eps
+        if (self.global_config["projector_hidden_act"] != "gelu" or not self.global_config["projector_bias"]
+                or self.global_config["projector_layer_norm_eps"] != 1e-6):
+            raise ValueError("unsupported projector_hidden_act / projector_bias / projector_layer_norm_eps")
+        # clip.cpp hardcodes the vision rope base, the RGB patchifier and the 3x3 spatial conv
+        if (self.hparams_vision["rope_parameters"]["rope_theta"] != 1e6 or self.hparams_vision["in_channels"] != 3
+                or self.hparams_vision["conv_kernel_size"] != 3):
+            raise ValueError("unsupported vision rope_theta / in_channels / conv_kernel_size")
+        self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.LFM3VL)
+        self.gguf_writer.add_vision_head_count_kv(self.hparams_vision["num_key_value_heads"])
+        self.gguf_writer.add_vision_spatial_merge_size(self.global_config["downsample_factor"])
+        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision["norm_eps"])
+        self.gguf_writer.add_vision_use_silu(True)
+        size = self.preprocessor_config["size"]
+        self.gguf_writer.add_vision_min_pixels(size["shortest_edge"])
+        self.gguf_writer.add_vision_max_pixels(size["longest_edge"])
+        # PIL resample enum, the torchvision backend has no lanczos
+        resample = self.preprocessor_config["resample"]
+        resize_algo = {2: "bilinear", 3: "bicubic"}.get(resample)
+        if resize_algo is None:
+            raise ValueError(f"unsupported resample: {resample}")
+        self.gguf_writer.add_vision_image_resize_algo(resize_algo)
+        self._check_video_config()
+
+    # mtmd_get_video_ts_style and the lfm3vl video frame sizing hardcode these, conversion fails on a mismatch
+    _video = {"fps": 2.0, "min_frames": 4, "temporal_patch_size": 2, "size": {"shortest_edge": 4096, "longest_edge": 25165824}}
+    _video_max_frames = 768
+
+    def _check_video_config(self):
+        import json
+        video = self.preprocessor_config.get("video_processor")
+        path = self.dir_model / "video_preprocessor_config.json"
+        if video is None and path.is_file():
+            with open(path, encoding="utf-8") as f:
+                video = json.load(f)
+        if video is None:
+            return
+        for k, v in {**self._video, "max_frames": self._video_max_frames}.items():
+            if video.get(k) != v:
+                raise ValueError(f"unsupported video {k}={video.get(k)!r}, expected {v!r}")
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if "patch_embd" in new_name:
+            return gguf.GGMLQuantizationType.F16
+        if "conv2d_spatial" in new_name:
+            # CUDA CONV_2D_DW requires an f32 kernel; the tensor is tiny
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    def _lfm3vl_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name == "model.vision_tower.patchifier.weight":
+            # qwen2vl-style layout [channel, temporal, ph, pw], one conv kernel per temporal slice
+            assert self.hparams_vision is not None
+            t = self.hparams_vision["temporal_patch_size"]
+            p = self.hparams_vision["patch_size"]
+            n_embd = data_torch.shape[0]
+            w = data_torch.view(n_embd, 3, t, p, p)
+            yield ("v.patch_embd.weight", w[:, :, 0].contiguous())
+            if t == 2:
+                yield ("v.patch_embd.weight.1", w[:, :, 1].contiguous())
+            else:
+                assert t == 1, f"unsupported vision_temporal_patch_size {t}"
+            return
+        if name == "model.vision_tower.embedding_norm.weight":
+            yield ("v.embedding_norm.weight", data_torch)
+            return
+        if "conv.in_proj.weight" in name:
+            yield (f"v.blk.{bid}.conv2d_in_proj.weight", data_torch)
+            return
+        if "conv.spatial_conv.weight" in name:
+            yield (f"v.blk.{bid}.conv2d_spatial.weight", data_torch)
+            return
+        if "conv.out_proj.weight" in name:
+            yield (f"v.blk.{bid}.conv2d_out_proj.weight", data_torch)
+            return
+
+        # projector: input LayerNorm -> linear_1 -> GELU -> linear_2
+        if name.startswith("model.multi_modal_projector."):
+            layer, suffix = name.removeprefix("model.multi_modal_projector.").split(".")
+            yield ({"layer_norm": "mm.input_norm", "linear_1": "mm.0", "linear_2": "mm.2"}[layer] + "." + suffix, data_torch)
+            return
+
+        yield (self.map_tensor_name(name), data_torch)
+
+
+@ModelBase.register("Lfm3VlForConditionalGeneration")
+@ModelBase.example("LiquidAI/LFM3-VL-2B")
+class LFM3VLMmprojModel(LFM3VLVisionMixin):
+    has_vision_encoder = True
+    has_audio_encoder = False
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.hparams_vision is not None
+        # dynamic resolution, image_size is only set for MmprojModel compatibility
+        self.hparams_vision["image_size"] = 512
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self._set_lfm3vl_gguf_parameters()
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.startswith(("model.vision_tower.", "model.multi_modal_projector.")):
+            yield from self._lfm3vl_tensors(data_torch, name, bid)
+
+
+def _is_lfm3_audio_checkpoint(dir_model: Path) -> bool:
+    if not (dir_model / "config.json").is_file():
+        return False
+    with open(dir_model / "config.json", encoding="utf-8") as f:
+        config = json.load(f)
+    return config.get("model_type") == "lfm3_audio" and "lfm" in config
+
+
+# read the release config.json as is: a transformers that knows lfm3_audio rewrites it into another layout
+@ModelBase.register_hparams_loader(_is_lfm3_audio_checkpoint)
+def _load_lfm3_audio_hparams(dir_model: Path) -> dict[str, Any]:
+    with open(dir_model / "config.json", encoding="utf-8") as f:
+        config = json.load(f)
+    config["text_config"] = config["lfm"]
+    return config
+
+
+@ModelBase.register("Lfm3AudioForConditionalGeneration")
+@ModelBase.example("LiquidAI/LFM3-Audio-2B")
+class LFM3AudioMmprojModel(MmprojModel):
+    has_vision_encoder = False
+    has_audio_encoder = True
+
+    # the downsampler has the LFM2-Audio pre-encoder layout: conv, dw conv, pw conv, dw conv, pw conv
+    _downsample_idx = {"conv_in": 0, "conv_mid": 2, "lin_mid": 3, "conv_end": 5, "lin_end": 6}
+
+    # audio output: MaskGIT depthformer (a.gen.code.*) and the codec decoder (a.gen.wav.*)
+    _depth_layer_map = {
+        "operator_norm":      gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_NORM,
+        "self_attn.q_proj":   gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_Q,
+        "self_attn.q_norm":   gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_Q_NORM,
+        "self_attn.k_proj":   gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_K,
+        "self_attn.k_norm":   gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_K_NORM,
+        "self_attn.v_proj":   gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_V,
+        "self_attn.o_proj":   gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_OUT,
+        "self_attn.gate_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_GATE,
+        "ffn_norm":           gguf.MODEL_TENSOR.A_GEN_CODE_FFN_NORM,
+        "feed_forward.w1":    gguf.MODEL_TENSOR.A_GEN_CODE_FFN_GATE,
+        "feed_forward.w3":    gguf.MODEL_TENSOR.A_GEN_CODE_FFN_UP,
+        "feed_forward.w2":    gguf.MODEL_TENSOR.A_GEN_CODE_FFN_DOWN,
+    }
+    _codec_layer_map = {
+        "operator_norm":         gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM,
+        "post_operator_norm":    gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_POST_NORM,
+        "self_attn.q_proj":      gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q,
+        "self_attn.q_norm":      gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q_NORM,
+        "self_attn.k_proj":      gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K,
+        "self_attn.k_norm":      gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K_NORM,
+        "self_attn.v_proj":      gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V,
+        "self_attn.o_proj":      gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT,
+        "self_attn.gate_proj":   gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_GATE,
+        "conv.in_proj":          gguf.MODEL_TENSOR.A_GEN_WAV_TFM_CONV_PW1,
+        "conv.conv":             gguf.MODEL_TENSOR.A_GEN_WAV_TFM_CONV_DW,
+        "conv.out_proj":         gguf.MODEL_TENSOR.A_GEN_WAV_TFM_CONV_PW2,
+        "ffn_norm":              gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM,
+        "post_ffn_norm":         gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_POST_NORM,
+        "feed_forward.w1":       gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_GATE,
+        "feed_forward.w3":       gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP,
+        "feed_forward.w2":       gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN,
+    }
+    # the codec graph in clip.cpp hardcodes the liquid-codec FSQ front end and decoder shape, conversion fails on a mismatch
+    _codec = {
+        "model_type": "liquid_codec", "fsq_levels": [8, 7, 6, 6], "n_q": 12, "quantizer_dim": 256,
+        "encoder_dim": 512, "stft_n_fft": 1280, "stft_hop_length": 320, "stft_win_length": 1280, "stft_frame_stack": 6,
+        "decoder_window": 30, "sampling_rate": 24000,
+    }
+    # the lfm3a pipeline in mtmd-helper-gen.cpp hardcodes these
+    _generation = {"codebooks": 12, "maskgit_schedule": [4, 4, 4], "interleaved_n_text": 6, "interleaved_n_audio": 12}
+
+    # speech-code sampling of the model card, generation_config.json overrides it
+    _audio_sampling = {
+        "audio_temperature": 1.0, "audio_top_k": 128, "audio_repetition_penalty": 1.1, "audio_repetition_window": 4096,
+    }
+
+    # mtmd_audio_preprocessor_lfm3a hardcodes this front end, conversion fails on a mismatch
+    _feature_extractor = {
+        "sampling_rate": 16000, "feature_size": 128, "n_fft": 400, "win_length": 400, "hop_length": 160,
+        "log_floor": -12.0, "normalization_offset": 4.0, "normalization_divisor": 4.0, "encoder_stride": 8,
+    }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # audio output is exported when the checkpoint ships the codec
+        self._audio_output = (self.dir_model / "audio_tokenizer").is_dir()
+        gen = self.get_generation_config()
+        self._n_codebooks: int = gen["codebooks"]
+        self._audio_vocab: int = gen["audio_vocab_size"]
+        self._depth_embd: dict[int, Tensor] = {}
+        self._depth_norm: dict[int, Tensor] = {}
+
+    def get_audio_config(self) -> dict[str, Any] | None:
+        enc = self.global_config["encoder"]
+        cfg = dict(enc["lfm"])
+        cfg["num_mel_bins"] = enc["dim_in"]
+        cfg["sliding_window"] = enc["sliding_window"]
+        return cfg
+
+    def get_generation_config(self) -> dict[str, Any]:
+        c = self.global_config
+        steps = c["maskgit_steps"]
+        return {
+            "codebooks": c["codebooks"], "audio_vocab_size": c["audio_vocab_size"],
+            "maskgit_schedule": [c["codebooks"] // steps] * steps,
+            "interleaved_n_text": c["interleaved_n_text"], "interleaved_n_audio": c["interleaved_n_audio"],
+        }
+
+    def get_depth_config(self) -> dict[str, Any]:
+        return self.global_config["depthformer"]
+
+    def set_gguf_parameters(self):
+        assert self.hparams_audio is not None
+        fe = self.preprocessor_config["feature_extractor"]
+        if fe["feature_extractor_type"] not in ("Lfm3AudioFeatureExtractor", "Lfm3OmniFeatureExtractor"):
+            raise ValueError(f"unsupported feature_extractor_type {fe['feature_extractor_type']!r}")
+        for k, v in self._feature_extractor.items():
+            if fe[k] != v:
+                raise ValueError(f"unsupported feature_extractor.{k}={fe[k]!r}, expected {v!r}")
+
+        # clip.cpp hardcodes the encoder rope base
+        if self.hparams_audio["rope_parameters"]["rope_theta"] != 1e6:
+            raise ValueError("unsupported audio encoder rope_theta")
+
+        self.hparams_audio["intermediate_size"] = lfm2_feed_forward_length(self.hparams_audio)
+        super().set_gguf_parameters()
+        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.LFM3A)
+        self.gguf_writer.add_audio_head_count_kv(self.hparams_audio["num_key_value_heads"])
+        self.gguf_writer.add_audio_attention_layernorm_eps(self.hparams_audio["norm_eps"])
+        self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"])
+        self.gguf_writer.add_audio_window_size(self.hparams_audio["sliding_window"])
+
+        if not self._audio_output:
+            return
+
+        # audio generation: depthformer hparams, the codec ones are constant and live in clip.cpp
+        gen = self.get_generation_config()
+        for k, v in self._generation.items():
+            if gen[k] != v:
+                raise ValueError(f"unsupported {k}={gen[k]!r}, expected {v!r}")
+        depth = self.get_depth_config()
+        if depth.get("is_causal", False) or depth["use_sandwich_norm"] or set(depth["layer_types"]) != {"full_attention"}:
+            raise ValueError("unsupported depthformer, expected a bidirectional all-attention stack")
+        if depth["rope_parameters"]["full_attention"]["rope_theta"] != 1e6:
+            raise ValueError("unsupported depthformer rope_theta")
+        self.gguf_writer.add_clip_has_gen_audio_encoder(True)
+        self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.LFM3A_GEN)
+        self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text)
+        self.gguf_writer.add_gen_audio_embedding_length(depth["hidden_size"])
+        self.gguf_writer.add_gen_audio_feed_forward_length(depth["intermediate_size"])
+        self.gguf_writer.add_gen_audio_block_count(depth["num_hidden_layers"])
+        self.gguf_writer.add_gen_audio_head_count(depth["num_attention_heads"])
+        self.gguf_writer.add_gen_audio_head_count_kv(depth["num_key_value_heads"])
+        self.gguf_writer.add_gen_audio_attention_layernorm_eps(depth["norm_eps"])
+
+        # sampling defaults of the speech codes, next to the text ones that become general.sampling.*
+        gen = gguf.Metadata.load_generation_config(self.dir_model)
+        for key, add in (("audio_temperature",        self.gguf_writer.add_gen_audio_sampling_temp),
+                         ("audio_top_k",              self.gguf_writer.add_gen_audio_sampling_top_k),
+                         ("audio_repetition_penalty", self.gguf_writer.add_gen_audio_sampling_penalty_repeat),
+                         ("audio_repetition_window",  self.gguf_writer.add_gen_audio_sampling_penalty_last_n)):
+            value = gen.get(key, self._audio_sampling.get(key))
+            if value is not None:
+                add(value)
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if new_name.startswith("a.conv1d.") or "conv_dw" in new_name or new_name.startswith("a.gen.wav.quant."):
+            return gguf.GGMLQuantizationType.F32
+        # the codec decoder is small and its log-magnitude/phase output does not survive quantization
+        if new_name.startswith("a.gen.wav."):
+            return gguf.GGMLQuantizationType.F16
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        if self._audio_output:
+            yield from self._codec_tensors()
+
+    def _codec_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        import json
+        from safetensors.torch import load_file
+
+        T = gguf.MODEL_TENSOR
+        codec_dir = self.dir_model / "audio_tokenizer"
+        with open(codec_dir / "config.json", encoding="utf-8") as f:
+            cfg = json.load(f)
+        for k, v in self._codec.items():
+            if cfg.get(k) != v:
+                raise ValueError(f"unsupported audio_tokenizer {k}={cfg.get(k)!r}, expected {v!r}")
+        if math.prod(cfg["fsq_levels"]) != self._audio_vocab - 1:
+            raise ValueError(f"audio_vocab_size {self._audio_vocab} does not match the codec, expected codebook size + 1")
+
+        state_dict = load_file(codec_dir / "model.safetensors")
+
+        # FSQ code table: index -> the len(levels) quantized values, see FiniteScalarQuantization._indices_to_codes
+        levels = torch.tensor(cfg["fsq_levels"], dtype=torch.int64)
+        basis = torch.cumprod(torch.cat([torch.ones(1, dtype=torch.int64), levels[:-1]]), dim=0)
+        half = levels // 2
+        idx = torch.arange(int(levels.prod()), dtype=torch.int64)[:, None]
+        codes = ((idx // basis) % levels - half) / half
+        yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_FIRST_CB), codes.to(torch.float32))
+
+        # one projection over the concatenated code vectors of all codebooks, under quantizer.fsq. in newer exports
+        quantizer = "quantizer.fsq." if "quantizer.fsq.project_out.weight" in state_dict else "quantizer."
+        for suffix in ("weight", "bias"):
+            yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_FIRST_OUT, suffix=f".{suffix}"), state_dict[f"{quantizer}project_out.{suffix}"])
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_IN_PROJ, suffix=f".{suffix}"), state_dict[f"from_quantizer.{suffix}"])
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_OUT_PROJ, suffix=f".{suffix}"), state_dict[f"decoder.lin.{suffix}"])
+        yield (self.format_tensor_name(T.A_GEN_WAV_TFM_OUTPUT_NORM), state_dict["decoder.lfm.embedding_norm.weight"])
+
+        for name, data_torch in state_dict.items():
+            if not name.startswith("decoder.lfm.layers."):
+                continue
+            bid_str, key = name.removeprefix("decoder.lfm.layers.").rsplit(".weight", 1)[0].split(".", 1)
+            if key == "conv.conv":
+                data_torch = data_torch.squeeze(1)
+            yield (self.format_tensor_name(self._codec_layer_map[key], int(bid_str)), data_torch)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.startswith("a.gen.wav."):
+            yield (name, data_torch)
+            return
+
+        if name.startswith(("depthformer.", "depth_", "codebook_embedding.", "mask_embedding")):
+            if self._audio_output:
+                yield from self._depthformer_tensor(data_torch, name, bid)
+            return
+
+        if not name.startswith(("encoder.", "mlp.")):
+            return
+
+        if name.startswith("encoder.downsample."):
+            layer, suffix = name.removeprefix("encoder.downsample.").split(".")
+            if layer == "lin_out":
+                if suffix == "weight":
+                    data_torch = data_torch.squeeze(-1)
+                yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_ENC_OUT, suffix=f".{suffix}"), data_torch)
+                return
+            if suffix == "bias":
+                data_torch = data_torch[:, None, None]
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_ENC_CONV1D, self._downsample_idx[layer], f".{suffix}"), data_torch)
+            return
+
+        # conv op requires 2d tensor
+        if name.endswith("conv.conv.weight"):
+            data_torch = data_torch.squeeze(1)
+
+        yield (self.map_tensor_name(name), data_torch)
+
+    def _depthformer_tensor(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        T = gguf.MODEL_TENSOR
+
+        if name.startswith("depthformer.model.layers."):
+            key = name.removeprefix("depthformer.model.layers.").split(".", 1)[1].rsplit(".weight", 1)[0]
+            yield (self.format_tensor_name(self._depth_layer_map[key], bid), data_torch)
+            return
+        if name == "depthformer.model.embedding_norm.weight":
+            yield (self.format_tensor_name(T.A_GEN_CODE_OUTPUT_NORM), data_torch)
+            return
+        if name.startswith("depth_linear."):
+            yield (self.format_tensor_name(T.A_GEN_CODE_PROJ_IN, suffix="." + name.rsplit(".", 1)[1]), data_torch)
+            return
+        if name == "mask_embedding":
+            yield (self.format_tensor_name(T.A_GEN_CODE_MASK_EMBD, suffix=""), data_torch)
+            return
+        if name == "codebook_embedding.weight":
+            assert data_torch.shape[0] == self._n_codebooks * self._audio_vocab
+            yield (self.format_tensor_name(T.A_GEN_CODE_OUT_EMBD), data_torch)
+            return
+
+        # per-codebook heads: the embedding table is tied to to_logits, merge the 12 tables into one tensor
+        assert name.startswith("depth_embeddings.")
+        idx = int(name.split(".")[1])
+        key = name.split(".", 2)[2]
+        if key == "embedding.weight":
+            self._depth_embd[idx] = data_torch
+            if len(self._depth_embd) == self._n_codebooks:
+                yield (self.format_tensor_name(T.A_GEN_CODE_EMBD), torch.stack([self._depth_embd[i] for i in range(self._n_codebooks)], dim=0))
+        elif key == "embedding_norm.weight":
+            self._depth_norm[idx] = data_torch
+            if len(self._depth_norm) == self._n_codebooks:
+                yield (self.format_tensor_name(T.A_GEN_CODE_HEAD_NORM), torch.stack([self._depth_norm[i] for i in range(self._n_codebooks)], dim=0))
+        elif key == "to_logits.weight":
+            if not torch.equal(data_torch, self._depth_embd[idx]):
+                raise ValueError(f"depth_embeddings.{idx} is not tied")
+        else:
+            raise ValueError(f"unexpected tensor {name}")
+
+
+@ModelBase.register("Lfm3OmniForConditionalGeneration")
+@ModelBase.example("LiquidAI/LFM3-Omni-2B")
+class LFM3OmniMmprojModel(LFM3VLVisionMixin, LFM3AudioMmprojModel):
+    has_vision_encoder = True
+    has_audio_encoder = True
+
+    _audio_sampling = {"audio_temperature": 0.4, "audio_repetition_penalty": 1.2, "audio_repetition_window": 64}
+    _video_max_frames = 384
+
+    # the audio side is the LFM3-Audio model under the HF-native names
+    _audio_renames = (
+        ("model.audio_encoder.lfm.", "encoder.lfm.lfm."),
+        ("model.audio_encoder.", "encoder."),
+        ("model.audio_projector.layer_norm.", "mlp.model.0."),
+        ("model.audio_projector.linear_1.", "mlp.model.1."),
+        ("model.audio_projector.linear_2.", "mlp.model.3."),
+        ("model.depth_decoder.model.", "depthformer.model."),
+        ("model.depth_decoder.", ""),
+        ("model.codebook_embedding.", "codebook_embedding."),
+    )
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.hparams_vision is not None
+        # dynamic resolution, image_size is only set for MmprojModel compatibility
+        self.hparams_vision["image_size"] = 512
+
+    def get_audio_config(self) -> dict[str, Any] | None:
+        enc = self.global_config["audio_encoder_config"]
+        cfg = dict(enc["text_config"])
+        cfg["num_mel_bins"] = enc["num_mel_bins"]
+        cfg["sliding_window"] = enc["sliding_window"]
+        return cfg
+
+    def get_generation_config(self) -> dict[str, Any]:
+        c = self.global_config
+        return {
+            "codebooks": c["num_codebooks"], "audio_vocab_size": c["codebook_size"], "maskgit_schedule": c["maskgit_schedule"],
+            "interleaved_n_text": c["interleaved_n_text"], "interleaved_n_audio": c["interleaved_n_audio"],
+        }
+
+    def get_depth_config(self) -> dict[str, Any]:
+        return self.global_config["depth_config"]
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self._set_lfm3vl_gguf_parameters()
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.startswith(("model.vision_tower.", "model.multi_modal_projector.")):
+            yield from self._lfm3vl_tensors(data_torch, name, bid)
+            return
+        for old, new in self._audio_renames:
+            if name.startswith(old):
+                yield from super().modify_tensors(data_torch, new + name.removeprefix(old), bid)
+                return
+        yield from super().modify_tensors(data_torch, name, bid)
@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+from typing import Iterable, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import ModelBase, TextModel, gguf
+
+
+@ModelBase.register("LimiteForCausalLM")
+@ModelBase.example("paradigma-inc/limite-1b-violetto")
+class LimiteModel(TextModel):
+    model_arch = gguf.MODEL_ARCH.LIMITE
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        hp = self.hparams
+
+        expected = {
+            "pos_mode": "rope",
+            "qk_norm": "rms_pre_rope",
+            "rope_style": "interleaved_pairs_odd_lane_sign_flip",
+            "rope_per_layer": False,
+            "rms_norm_has_weight": False,
+            "rms_norm_eps_mode": "torch_finfo_default",
+            "mlp_type": "swiglu",
+            "attn_gate_applied": "per_head_before_o_proj",
+            "ve_head_slice": "first_num_key_value_heads",
+            "ve_applied_before_qk_norm": True,
+            "global_nope": True,
+            "global_window": -1,
+            "mudd_r_site": "resid",
+            "tie_word_embeddings": True,
+            "final_softcap": 0.0,
+        }
+        for key, value in expected.items():
+            if hp.get(key) != value:
+                raise NotImplementedError(f"Limite: unsupported {key}={hp.get(key)!r}, expected {value!r}")
+        if hp["softcap_logits"]["kind"] != "sigmoid":
+            raise NotImplementedError(f"Limite: unsupported softcap_logits kind {hp['softcap_logits']['kind']!r}")
+
+        self.n_head    = hp["num_attention_heads"]
+        self.n_head_kv = hp["num_key_value_heads"]
+        self.head_dim  = hp["head_dim"]
+        self.n_pairs   = hp["rope_n_pairs"]
+
+        self.mudd_taps: dict[int, list[int]] = {}
+        if hp.get("mudd"):
+            if not hp.get("mudd_mlp"):
+                raise NotImplementedError("Limite: mudd without mudd_mlp is not supported")
+            self.mudd_taps = {int(il): [int(t) for t in taps] for il, taps in hp["mudd_tap_idx"].items()}
+            assert all(len(taps) == hp["mudd_taps"] for taps in self.mudd_taps.values())
+
+    def set_vocab(self):
+        self._set_vocab_gpt2()
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        hp = self.hparams
+
+        self.gguf_writer.add_key_length(self.head_dim)
+        self.gguf_writer.add_value_length(self.head_dim)
+
+        # torch rms_norm with eps=None uses the fp32 epsilon, even for bf16 inputs
+        self.gguf_writer.add_layer_norm_rms_eps(torch.finfo(torch.float32).eps)
+        self.gguf_writer.add_attention_scale(hp["attention_softmax_scale"])
+        self.gguf_writer.add_xsa_eps(hp["xsa_normalize_eps"] if hp.get("xsa") else 0.0)
+
+        # HF freq_i = base^(-i/(n_pairs-1)), ggml uses freq_base^(-i/n_pairs)
+        self.gguf_writer.add_rope_dimension_count(2 * self.n_pairs)
+        self.gguf_writer.add_rope_freq_base(hp["rope_base_local"] ** (self.n_pairs / (self.n_pairs - 1)))
+
+        # HF keeps keys k >= q - window, so a window of W spans W + 1 tokens
+        global_layers = set(hp["global_layers"])
+        self.gguf_writer.add_sliding_window(hp["sliding_window"] + 1)
+        self.gguf_writer.add_sliding_window_pattern([il not in global_layers for il in range(self.block_count)])
+
+        cap = hp["softcap_logits"]
+        self.gguf_writer.add_final_logit_sigmoid_capping(cap["a"], cap["b"], cap["c"])
+
+        if self.mudd_taps:
+            n_taps = hp["mudd_taps"]
+            tap_indices = [-1] * (self.block_count * n_taps)
+            for il, taps in self.mudd_taps.items():
+                tap_indices[il * n_taps:(il + 1) * n_taps] = taps
+            self.gguf_writer.add_mudd_feed_forward_length(hp["mudd_inter"])
+            self.gguf_writer.add_mudd_tap_count(n_taps)
+            self.gguf_writer.add_mudd_tap_indices(tap_indices)
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        # tiny mixing weights, kept at full precision as in the reference
+        # the mudd tensors come from a global source tensor, so take the block id from the new name
+        if new_name.startswith("blk."):
+            bid = int(new_name.split(".")[1])
+        for key in (gguf.MODEL_TENSOR.ATTN_GATE, gguf.MODEL_TENSOR.ATTN_VE_GATE, gguf.MODEL_TENSOR.MUDD_DOWN, gguf.MODEL_TENSOR.MUDD_UP):
+            if self.match_model_tensor_name(new_name, key, bid):
+                return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    def _swap_rope_pairs(self, w: Tensor, n_head: int) -> Tensor:
+        # HF rotates each pair by -theta; swapping the pair lanes gives the ggml +theta rotation.
+        # q and k are both swapped, so the attention scores do not change.
+        idx = torch.arange(self.head_dim)
+        idx[:2 * self.n_pairs] = idx[:2 * self.n_pairs].view(-1, 2).flip(-1).reshape(-1)
+        w = w.reshape(n_head, self.head_dim, -1)
+        return w[:, idx, :].reshape(n_head * self.head_dim, -1)
+
+    def _pad_cols(self, w: Tensor) -> Tensor:
+        # gates read only the first channels of the input, pad them to n_embd
+        return torch.nn.functional.pad(w, (0, self.hparams["hidden_size"] - w.shape[-1]))
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        hp = self.hparams
+        T = gguf.MODEL_TENSOR
+
+        def load(n: str) -> Tensor:
+            return self.model_tensors[n]().float()
+
+        if name == "model.value_embeds.weight":
+            ve = data_torch.reshape(data_torch.shape[0], -1, hp["ve_dim"])[:, :self.n_head_kv]
+            ve = torch.nn.functional.pad(ve, (0, self.head_dim - hp["ve_dim"]))
+            yield self.format_tensor_name(T.VALUE_EMBD), ve.reshape(ve.shape[0], -1) * hp["ve_gate_scale"]
+            return
+
+        if name == "model.mudd.dense1":
+            for il in self.mudd_taps:
+                yield self.format_tensor_name(T.MUDD_DOWN, il), data_torch
+            return
+        if name == "model.mudd.dense2":
+            # rows [0, n) mix the attention input, rows [n, 2n) mix the residual
+            dense2_mlp = load("model.mudd.dense2_mlp")
+            bias       = load("model.mudd.bias")
+            bias_mlp   = load("model.mudd.bias_mlp")
+            for il, taps in self.mudd_taps.items():
+                n = len(taps)
+                w = torch.cat([data_torch[il, :n], dense2_mlp[il, :n]], 0)
+                b = torch.cat([bias[il, :n], bias_mlp[il, :n]], 0)
+                yield self.format_tensor_name(T.MUDD_UP, il), w
+                yield self.format_tensor_name(T.MUDD_UP, il, ".bias"), b
+            return
+        if name.startswith("model.mudd."):
+            return
+
+        if bid is None:
+            yield from super().modify_tensors(data_torch, name, bid)
+            return
+
+        prefix = f"model.layers.{bid}."
+        name = name.removeprefix(prefix)
+
+        if name in ("self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"):
+            data_torch = data_torch * load(prefix + "self_attn.qkv_scale")
+            if name == "self_attn.q_proj.weight":
+                data_torch = self._swap_rope_pairs(data_torch, self.n_head)
+            elif name == "self_attn.k_proj.weight":
+                data_torch = self._swap_rope_pairs(data_torch, self.n_head_kv)
+        elif name == "self_attn.o_proj.weight":
+            data_torch = data_torch * load(prefix + "self_attn.o_scale") * load(prefix + "post_lambda_attn") * hp["attn_gate_scale"]
+        elif name == "mlp.down_proj.weight":
+            data_torch = data_torch * load(prefix + "post_lambda_mlp")
+        elif name == "self_attn.attn_gate":
+            yield self.format_tensor_name(T.ATTN_GATE, bid), self._pad_cols(data_torch)
+            return
+        elif name == "self_attn.ve_gate":
+            yield self.format_tensor_name(T.ATTN_VE_GATE, bid), self._pad_cols(data_torch[:self.n_head_kv])
+            return
+        elif name == "self_attn.xsa_alpha":
+            if hp.get("xsa") and bid in hp["xsa_layers"]:
+                yield self.format_tensor_name(T.ATTN_XSA_ALPHA, bid), torch.tanh(data_torch)
+            return
+        elif name == "resid_lambda_attn":
+            yield self.format_tensor_name(T.ATTN_RESID_SCALE, bid), data_torch.reshape(1)
+            return
+        elif name == "resid_lambda_mlp":
+            yield self.format_tensor_name(T.FFN_RESID_SCALE, bid), data_torch.reshape(1)
+            return
+        elif name in ("self_attn.qkv_scale", "self_attn.o_scale", "post_lambda_attn", "post_lambda_mlp"):
+            # folded into the projections above
+            return
+
+        yield from super().modify_tensors(data_torch, prefix + name, bid)
@@ -690,6 +690,52 @@ class Qwen3_5TextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
        self._try_set_pooling_type()


+@ModelBase.register("MiniCPMOCRForSequenceClassification")
+class MiniCPMOCRSeqClsModel(Qwen3_5TextModel):
+    # OCR consistency judge: Qwen3.5 VL trunk with a value head read at the last token
+    # the class scores are [inconsistent, consistent], the softmax temperature is 1
+    model_arch = gguf.MODEL_ARCH.QWEN35
+    no_mtp = True  # the checkpoint has no MTP head
+
+    CLASS_LABELS = ["inconsistent", "consistent"]
+
+    def set_vocab(self):
+        super().set_vocab()
+        self.gguf_writer.add_chat_template([{"name": "systemone", "template": self._systemone_template()}])
+
+    def _systemone_template(self) -> str:
+        # the prompt of the training data, the two images are the original page and the one
+        # rendered from the OCR text; the task is fixed, there is nothing to fill in
+        # the trailing newlines are in an expression, so the template does not drop them
+        return (
+            "<|im_start|>user\n{{ images[0] }}\nOCR文本：\n{{ state }}\n{{ images[1] }}\n"
+            "第一张图是原图，第二张图是上述 OCR 文本渲染得到的图片。\n"
+            "请判断该 OCR 文本及其对应渲染图是否与原图内容一致。\n"
+            "这是一个二分类任务：一致标为 1，不一致标为 0。"
+            "{{ '<|im_end|>\\n<|im_start|>assistant\\n<think>\\n\\n</think>\\n\\n' }}"
+        )
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+
+        # the config.json of the checkpoint keeps the one of the base model, it has no num_labels / id2label
+        if "score.weight" not in self.model_tensors:
+            raise ValueError("the checkpoint has no score.weight classification head")
+
+        self.gguf_writer.add_pooling_type(gguf.PoolingType.LAST)
+        self.gguf_writer.add_embedding_length_out(len(self.CLASS_LABELS))
+        self.gguf_writer.add_classifier_output_labels(self.CLASS_LABELS)
+        self.gguf_writer.add_decision_type(gguf.DecisionType.MINICPM_OCR_JEV)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name == "score.weight":
+            if data_torch.shape[0] != len(self.CLASS_LABELS):
+                raise ValueError(f"score.weight must have {len(self.CLASS_LABELS)} rows, got {data_torch.shape[0]}")
+            yield (gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.CLS_OUT] + ".weight", data_torch)
+            return
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
 def _is_openjev_checkpoint(dir_model: Path) -> bool:
    return (dir_model / "helper" / "shim.py").is_file() and (dir_model / "config.json").is_file()

@@ -144,6 +144,16 @@ class Qwen3VLVisionModel(MmprojModel):
        yield from MmprojModel.modify_tensors(self, data_torch, name, bid)


+@ModelBase.register("MiniCPMOCRForSequenceClassification")
+class MiniCPMOCRSeqClsVisionModel(Qwen3VLVisionModel):
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        # trained with the image budget of its processor, wider than the mtmd default of Qwen-VL
+        size = self.preprocessor_config["size"]
+        self.gguf_writer.add_vision_min_pixels(int(size["shortest_edge"]))
+        self.gguf_writer.add_vision_max_pixels(int(size["longest_edge"]))
+
+
@ModelBase.register("Qwen3OmniMoeForConditionalGeneration")
@ModelBase.example("Qwen/Qwen3-Omni-30B-A3B-Instruct")
 class Qwen3OmniMmprojModel(Qwen3VLVisionModel, Qwen25AudioModel):
@@ -72,6 +72,8 @@ However, A6x GPUs in phones are likely not supported due to the outdated driver
 | MXFP4                  | Support                    |
 | IQ4_NL                 | Support                    |

+`MUL_MAT_ID` accepts an expert id of `-1` and writes zeros to that slot. `ADD_ID` adds no bias for `-1`, including in fused MoE operations.
+
 ## Model Preparation

 Since common quantizations are supported now, it is recommanded to download GGUF models directly from Huggingface.
@@ -784,6 +784,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_TARGET   | INTEL *(default)*                     | Set the SYCL target device type.            |
 | GGML_SYCL_DEVICE_ARCH | Optional                           | Set the SYCL device architecture. Setting the device architecture can improve the performance. See the table [--offload-arch](https://github.com/intel/llvm/blob/sycl/sycl/doc/design/OffloadDesign.md#--offload-arch) for a list of valid architectures. |
 | GGML_SYCL_F16      | OFF *(default)* \|ON *(optional)*     | Enable FP16 build with SYCL code path. [1.](#sycl-f16-footnote) |
+| GGML_SYCL_SUPPORT_XMX | ON *(default)* \|OFF *(Optional)* | Build the code paths that use XMX (DPAS) instructions. Set OFF for AOT builds targeting iGPUs without XMX, which otherwise fail to compile. |
 | GGML_SYCL_GRAPH    | ON *(default)* \|OFF *(Optional)*     | Enable build with [SYCL Graph extension](https://github.com/intel/llvm/blob/sycl/sycl/doc/extensions/experimental/sycl_ext_oneapi_graph.asciidoc). |
 | GGML_SYCL_DNN      | ON *(default)* \|OFF *(Optional)*     | Enable build with oneDNN.                   |
 | GGML_SYCL_HOST_MEM_FALLBACK | ON *(default)* \|OFF *(Optional)* | Allow host memory fallback when device memory is full during quantized weight reorder. Enables inference to continue at reduced speed (reading over PCIe) instead of failing. Requires Linux kernel 6.8+. |
@@ -812,6 +813,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_ENABLE_DNN | 0 or 1 (default)| Enable running computations through oneDNN and always use oneMKL. |
 | GGML_SYCL_FA_ONEDNN | 1 (default) or 0 | Enable the oneDNN fused SDPA (flash-attention) path on supported GPUs. Set to 0 to always use the native SYCL flash-attention kernel. |
 | GGML_SYCL_FA_ONEDNN_MAX_KV | 0 (default, disabled) or positive integer | By default (0), all sequences are handled by the oneDNN fused SDPA path, regardless of KV length; a positive value caps that length, past which sequences fall back to the native kernel. If GPU driver watchdog resets (DEVICE_LOST) occur during long-context inference, set this near the context depth where they start, e.g. 24576. |
+| GGML_SYCL_FA_DPAS | 1 (default) or 0 | Enable the XMX flash-attention kernel for a q4_0 or q8_0 KV cache (`-ctk q4_0 -ctv q4_0` or `-ctk q8_0 -ctv q8_0`) on Xe2 or newer GPUs (Arc B-series, Lunar Lake, Panther Lake, Wildcat Lake). It reads the quantized cache directly for decode and batches of up to 31 tokens, head sizes 64, 96, 128, 256 and 512, up to 32 query heads per KV head. Set to 0 to use the native SYCL flash-attention kernels. |
 | GGML_SYCL_ENABLE_VMM | 0 or 1 (default) | Enable the virtual-memory device pool. |
 | GGML_SYCL_ENABLE_MKL_FA | 1 (default) or 0 | Enable oneMKL GEMM flash attention for XMX-accelerated prompt processing with quantized KV cache. Automatically activates during prefill (prompt processing) when all conditions are met: (1) flash-attn enabled (`-fa` or `--flash-attn on`), (2) KV cache quantized (`--cache-type-k q8_0 --cache-type-v q8_0` or other `*_0/*_1` types), (3) batch size ≥ 1024 (`--batch-size 1024`), (4) prompt length ≥ 1024 tokens. Set to 0 to force the TILE kernel for A/B testing. Example minimum command: `llama-cli -m model.gguf -fa -ngl 99 --cache-type-k q8_0 --cache-type-v q8_0 --batch-size 1024 -p "your prompt"` |
 | GGML_SYCL_MKL_FA_DEBUG | 0 (default) or 1 | Enable per-call diagnostic logging for MKL flash attention: GEMM/softmax timings, interleaved-head detection, and buffer memory usage. |
@@ -846,6 +848,18 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
   Then benchmark with your preferred tool, such as `llama-bench`.
   For graph activity details (including reuse), set `GGML_SYCL_DEBUG=1` and look for `[SYCL-GRAPH]` messages.

+### Flash-attention kernels
+
+Attention dispatches on shape and hardware.
+
+DPAS is the quantized path. q4_0 or q8_0 KV cache, decode and batches up to 31 tokens, Xe2 and newer only. It issues dpas straight on the quantized blocks and never materialises an f16 copy. vISA gets used here because joint_matrix has no s4 type.
+
+oneDNN does XMX through SDPA but converts KV to f16 first, so it ends up on prefill: 32+ tokens and 1K+ keys. MKL is the prompt processing equivalent through oneMKL GEMM, also f16.
+
+VEC/TILE are plain kernels with no XMX in them. They run when DPAS is off or the hardware doesn't support it.
+
+joint_matrix belongs to the matmul GEMMs. Attention doesn't use it.
+
 ## Compile-time Flags

 Pass these via `CXXFLAGS` or add a one-off `#define` to enable a flag on the spot.
@@ -324,6 +324,9 @@ on 4 physical NPUs, or `--devices 'HTP0[0-1:0],HTP1[0-1:1]'` on 2 physical NPUs
 - `GGML_HEXAGON_NHVX=0`
  Controls the number of HVX hardware threads to use. The default is all (actual number varies depending on the hardware version).

+- `GGML_HEXAGON_VTCM_SIZE=0`
+  Overrides the VTCM size used by each session. Accepts `K`, `M` and `G` suffixes; a plain number is in MB. The default (0) uses the size the hardware reports. Values above the hardware size will fail to start the session.
+
 - `GGML_HEXAGON_HOSTBUF=1` (default: 0, disabled)
  Enables allocating host buffers for debugging. By default, host buffers are disabled.

@@ -14,22 +14,22 @@ Legend:

 | Operation | BLAS | CANN | CPU | CUDA | ET | HTP | MTL | OpenCL | SYCL | Vulkan | WebGPU | ZenDNN | zDNN |
 |-----------|------|------|------|------|------|------|------|------|------|------|------|------|------|
-|                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              ACC | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             ADD1 | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                           ADD_ID | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            CLAMP | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        COL2IM_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                       CONV_1D_DW | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ |
+|                       CONV_1D_DW | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ |
 |                          CONV_2D | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
-|                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          CONV_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                CONV_TRANSPOSE_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                CONV_TRANSPOSE_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -38,49 +38,49 @@ Legend:
 |                              CPY | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
 |               CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                           CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                              DIV | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              DIV | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                     DSV4_HC_COMB | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                     DSV4_HC_POST | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                      DSV4_HC_PRE | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                              DUP | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                              ELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              EXP | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                            EXPM1 | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             FILL | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              ELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              EXP | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            EXPM1 | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             FILL | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                   FLASH_ATTN_EXT | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
-|                            FLOOR | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                  GATED_DELTA_NET | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            FLOOR | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                  GATED_DELTA_NET | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                GATED_LINEAR_ATTN | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            GEGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        GEGLU_ERF | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                      GEGLU_QUICK | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            GEGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                        GEGLU_ERF | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                      GEGLU_QUICK | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             GELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         GELU_ERF | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                       GELU_QUICK | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         GET_ROWS | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ❌ | ❌ |
 |                    GET_ROWS_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ |
 |                       GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | 🟡 | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ❌ |
+|                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | 🟡 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                LIGHTNING_INDEXER | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                LIGHTNING_INDEXER | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ | ❌ |
-|                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
-|                 MUL_MAT_HADAMARD | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 |
+|                 MUL_MAT_HADAMARD | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ✅ | ❌ | 🟡 |
 |                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | 🟡 | ❌ |
-|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                     MUL_MAT_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                     MUL_MAT_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                     OPT_STEP_SGD | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -89,45 +89,46 @@ Legend:
 |                   PAD_REFLECT_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                          POOL_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                          POOL_2D | ❌ | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             RELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                      REPEAT_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                         RMS_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         RMS_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    RMS_NORM_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                             ROLL | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                             ROPE | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        ROPE_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            ROUND | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            ROUND | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        RWKV_WKV6 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                        RWKV_WKV7 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            SCALE | ❌ | 🟡 | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              SET | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ |
+|                            SCALE | ❌ | 🟡 | ✅ | ✅ | ✅ | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              SET | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ❌ |
 |                         SET_ROWS | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
-|                              SGN | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              SGN | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          SIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             SILU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        SILU_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
 |                              SIN | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                         SOFTPLUS | ❌ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         SOFTPLUS | ❌ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         SOFT_MAX | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    SOFT_MAX_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ |
-|                        SOLVE_TRI | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ |
+|                        SOLVE_TRI | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ❌ |
 |                              SQR | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             SQRT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         SSM_CONV | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                         SSM_SCAN | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ❌ | ❌ |
-|                             STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              SUB | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         SSM_SCAN | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ❌ |
+|                             STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              SUB | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              SUM | ❌ | 🟡 | ✅ | 🟡 | ❌ | 🟡 | 🟡 | ❌ | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
 |                         SUM_ROWS | ❌ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ |
-|                           SWIGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                     SWIGLU_CLAMP | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                           SWIGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                     SWIGLU_CLAMP | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                       SWIGLU_OAI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |               TIMESTEP_EMBEDDING | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | 🟡 | ✅ | ❌ | ❌ |
-|                              TRI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                            TRUNC | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
+|                       TOP_K_MASK | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                              TRI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            TRUNC | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          UPSCALE | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            XIELU | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -100,8 +100,9 @@ llama-server -m Qwen3-4B.gguf -md Qwen3-4B-DSpark.gguf \

 `--spec-draft-n-max` is clamped to the draft model's trained block size.

-`--spec-draft-conf-min P` truncates each drafted block at the first position whose predicted
-acceptance (from the draft's confidence head, if present) falls below `P` (default 0 = disabled).
+`--spec-draft-p-min P` truncates each drafted block at the first position whose predicted
+acceptance (from the draft's confidence head) falls below `P` (default 0 = disabled).
+Values above 0 require a draft with a confidence head.

 Currently only drafts with a Qwen3 backbone are supported; support for other backbones
 (e.g. Gemma4) is planned.
@@ -0,0 +1,211 @@
+#!/usr/bin/env python3
+"""Example client for the llama-server /v1/realtime websocket, with the openai SDK.
+
+Start a server with an audio model, e.g. LFM3-Audio:
+
+    llama-server -hf LiquidAI/LFM3-Audio-2B-GGUF
+
+Spoken conversation, the audio of each wav file is one user turn. It is sent in chunks as it would come from a
+microphone, the server prefills it as it arrives so the response starts right after the last chunk:
+
+    python examples/realtime_client.py --instructions "Respond with interleaved text and audio." question.wav
+
+Output modalities: left out, the model picks; "text", "audio" or "text,audio" (interleaved) force the response.
+Without --instructions, a forced mode comes with the system prompt of the LFM3-Audio or LFM3-Omni model card:
+
+    python examples/realtime_client.py --output-modalities text,audio --out answer.wav question.wav
+
+Tool calls: --tools gives the conversation a demo tool, the client runs its calls and sends the results back:
+
+    python examples/realtime_client.py --tools --text "Set the volume to 30."
+
+Text turns go in with --text, transcription sessions transcribe each file:
+
+    python examples/realtime_client.py --text "What is the capital of France?"
+    python examples/realtime_client.py --type transcription --language en speech.wav
+
+Live transcription models (Voxtral Realtime, LFM3-Audio live ASR) transcribe while the audio is sent, a transcription
+session opens with ?intent=transcription:
+
+    python examples/realtime_client.py --type transcription speech.wav
+
+Requirements: pip install "openai[realtime]" numpy
+"""
+
+import argparse
+import asyncio
+import base64
+import json
+import time
+import wave
+
+import numpy as np
+from openai import AsyncOpenAI
+
+
+# system prompts of the LiquidAI/LFM3-Audio-2B and LiquidAI/LFM3-Omni-2B model cards for a forced output mode
+MODE_INSTRUCTIONS = {
+    "audio": {"text,audio": "Respond with interleaved text and audio."},
+    "omni":  {"text": "Respond in text.", "text,audio": "Respond in interleaved text and audio."},
+}
+# and with tools: calls, or an answer in interleaved text and speech
+TOOLS_INSTRUCTIONS = "Respond with function calls, and interleaved text and audio.\n\nList of tools:"
+
+# the demo tool of --tools
+DEMO_TOOLS = [{
+    "type": "function",
+    "name": "set_volume",
+    "description": "Set the device speaker volume.",
+    "parameters": {
+        "type": "object",
+        "properties": {"volume": {"type": "integer", "description": "Volume from 0 to 100."}},
+        "required": ["volume"],
+    },
+}]
+
+
+def run_tool(name, arguments):
+    if name == "set_volume":
+        return {"volume": json.loads(arguments)["volume"]}
+    return {"error": f"unknown tool {name}"}
+
+
+def read_wav(path, rate):
+    # PCM16 wav -> mono int16 at the given rate
+    with wave.open(path, "rb") as f:
+        assert f.getsampwidth() == 2, "the example reads PCM16 wav files"
+        x = np.frombuffer(f.readframes(f.getnframes()), dtype=np.int16).astype(np.float32)
+        x = x.reshape(-1, f.getnchannels()).mean(axis=1)
+        sr = f.getframerate()
+    if sr != rate:
+        n = int(round(len(x) * rate / sr))
+        x = np.interp(np.arange(n) * sr / rate, np.arange(len(x)), x)
+    return np.clip(x, -32768, 32767).astype(np.int16)
+
+
+def write_wav(path, pcm, rate):
+    with wave.open(path, "wb") as f:
+        f.setnchannels(1)
+        f.setsampwidth(2)
+        f.setframerate(rate)
+        f.writeframes(pcm.tobytes())
+
+
+async def main(args):
+    # the SDK derives a wss:// URL from base_url, a plain http server needs ws://
+    client = AsyncOpenAI(base_url=args.url, api_key=args.api_key, websocket_base_url=args.url.replace("http", "ws", 1))
+    query = {"intent": "transcription"} if args.type == "transcription" else {}
+    async with client.realtime.connect(model=args.model, extra_query=query) as conn:
+        ev = await conn.recv()
+        if ev.type == "error":
+            print("error:", ev.error.message)
+            return
+        # the audio goes in at the rate of the model, PCM16 mono
+        session_info = ev.model_dump(warnings=False)["session"]
+        audio_fmt = session_info["audio"]
+        card = "omni" if "omni" in str(session_info.get("model", "")).lower() else "audio"
+        in_rate = audio_fmt["input"]["format"]["rate"]
+        out_rate = (audio_fmt.get("output") or {}).get("format", {}).get("rate", 24000)
+
+        session = {"type": args.type}
+        if args.type == "realtime":
+            instructions = args.instructions
+            if instructions is None:
+                instructions = TOOLS_INSTRUCTIONS if args.tools else MODE_INSTRUCTIONS[card].get(args.output_modalities)
+            if args.tools:
+                session["tools"] = DEMO_TOOLS
+            if instructions is not None:
+                session["instructions"] = instructions
+            if args.output_modalities:
+                session["output_modalities"] = args.output_modalities.split(",")
+        elif args.language:
+            session["audio"] = {"input": {"transcription": {"language": args.language}}}
+        await conn.session.update(session=session)  # ty: ignore[invalid-argument-type]
+        ev = await conn.recv()
+        if ev.type == "error":
+            print("error:", ev.error.message)
+            return
+
+        turns = [("audio", path) for path in args.wav] + [("text", text) for text in args.text]
+        for i, (kind, src) in enumerate(turns):
+            t_last = 0.0
+
+            # the input is sent while the events are read: live transcription arrives while the audio is sent
+            async def send_turn():
+                nonlocal t_last
+                if kind == "audio":
+                    pcm = read_wav(src, in_rate)
+                    step = in_rate * args.chunk_ms // 1000
+                    for j in range(0, len(pcm), step):
+                        chunk = pcm[j:j + step]
+                        await conn.input_audio_buffer.append(audio=base64.b64encode(chunk.tobytes()).decode())
+                        if args.realtime:
+                            await asyncio.sleep(len(chunk) / in_rate)
+                    if args.type == "transcription":
+                        # "final" ends the input, models that transcribe with a delay then complete the last item
+                        await conn.send({"type": "input_audio_buffer.commit", "final": i == len(args.wav) - 1})  # ty: ignore[invalid-argument-type]
+                else:
+                    await conn.conversation.item.create(item={
+                        "type": "message", "role": "user", "content": [{"type": "input_text", "text": src}]})
+                if args.type == "realtime":
+                    # the audio not committed yet ends the user turn
+                    await conn.response.create()
+                t_last = time.perf_counter()
+
+            print(f"--- turn {i} ({kind}: {src})")
+            sender = asyncio.create_task(send_turn())
+            t_first = None
+            audio = []
+            async for ev in conn:
+                if ev.type == "response.output_text.delta" or ev.type == "conversation.item.input_audio_transcription.delta":
+                    t_first = t_first or time.perf_counter()
+                    print(ev.delta, end="", flush=True)
+                elif ev.type == "response.output_audio.delta":
+                    t_first = t_first or time.perf_counter()
+                    audio.append(np.frombuffer(base64.b64decode(ev.delta), dtype=np.int16))
+                elif ev.type == "response.done":
+                    # the calls of the response are run, their results go back and the model answers them
+                    calls = [o for o in ev.model_dump()["response"]["output"] if o["type"] == "function_call"]
+                    for call in calls:
+                        result = json.dumps(run_tool(call["name"], call["arguments"]))
+                        print(f"\n[call {call['name']}({call['arguments']}) -> {result}]", flush=True)
+                        await conn.conversation.item.create(item={
+                            "type": "function_call_output", "call_id": call["call_id"], "output": result})
+                    if not calls:
+                        break
+                    await conn.response.create()
+                elif ev.type == "conversation.item.input_audio_transcription.completed":
+                    break
+                elif ev.type == "error":
+                    print("\nerror:", ev.error.message)
+                    return
+            await sender
+            print()
+            if t_first is not None and t_first > t_last:
+                print(f"first output {1000 * (t_first - t_last):.0f} ms after the end of the input")
+            elif t_first is not None:
+                print(f"output started {1000 * (t_last - t_first):.0f} ms before the end of the input, "
+                      f"done {1000 * (time.perf_counter() - t_last):.0f} ms after it")
+            if audio:
+                pcm_out = np.concatenate(audio)
+                path = f"{args.out.removesuffix('.wav')}_{i}.wav" if len(turns) > 1 else args.out
+                write_wav(path, pcm_out, out_rate)
+                print(f"audio: {len(pcm_out) / out_rate:.2f} s -> {path}")
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("wav", nargs="*", help="PCM16 wav files, one user turn each")
+    p.add_argument("--text", action="append", default=[], help="a user turn as text, after the audio turns, repeatable")
+    p.add_argument("--type", choices=["realtime", "transcription"], default="realtime", help="session type")
+    p.add_argument("--instructions", help="system prompt of the conversation")
+    p.add_argument("--tools", action="store_true", help="give the conversation a demo tool (set_volume), its calls are run")
+    p.add_argument("--output-modalities", choices=["text", "audio", "text,audio"], help="default: the model picks")
+    p.add_argument("--language", help="transcription language, e.g. en")
+    p.add_argument("--chunk-ms", type=int, default=100, help="audio chunk size")
+    p.add_argument("--realtime", action=argparse.BooleanOptionalAction, default=True, help="send the audio in real time")
+    p.add_argument("--out", default="response.wav", help="where the audio of the responses goes")
+    p.add_argument("--url", default="http://127.0.0.1:9931/v1")
+    p.add_argument("--api-key", default="none")
+    p.add_argument("--model", default="default", help="model name, used by the router")
+    asyncio.run(main(p.parse_args()))
@@ -4,7 +4,7 @@ project("ggml" C CXX ASM)

 ### GGML Version
 set(GGML_VERSION_MAJOR 0)
-set(GGML_VERSION_MINOR 27)
+set(GGML_VERSION_MINOR 28)
 set(GGML_VERSION_PATCH 0)
 set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")

@@ -250,6 +250,7 @@ option(GGML_OPENMP_FETCH                    "ggml: fetch LLVM OpenMP"
 option(GGML_RPC                             "ggml: use RPC"                                   OFF)
 option(GGML_SYCL                            "ggml: use SYCL"                                  OFF)
 option(GGML_SYCL_F16                        "ggml: use 16 bit floats for sycl calculations"   OFF)
+option(GGML_SYCL_SUPPORT_XMX                "ggml: build the XMX (DPAS) code paths in the SYCL backend" ON)
 option(GGML_SYCL_GRAPH                      "ggml: enable graphs in the SYCL backend"         ON)
 option(GGML_SYCL_HOST_MEM_FALLBACK          "ggml: allow host memory fallback in SYCL reorder (requires kernel 6.8+)" ON)
 option(GGML_SYCL_SUPPORT_LEVEL_ZERO_API     "ggml: use Level Zero API in SYCL backend"  ON)
@@ -535,6 +535,60 @@ void ggml_gemv_mxfp4_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const v
    UNUSED(ncols_interleaved);
    UNUSED(blocklen);

+#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE)
+    if (svcntb() * 8 == 256) {
+        const svbool_t  pg16 = svptrue_pat_b8(SV_VL16);
+        const svbool_t  pg4  = svptrue_pat_b32(SV_VL4);
+        const svbool_t  pgb  = svptrue_b8();
+        const svbool_t  pgw  = svptrue_b32();
+        const svbool_t  pg_hi4 = svnot_b_z(pgw, pg4);
+        const svuint8_t m4b = svdup_n_u8(0x0f);
+        const svuint8_t values = svld1_u8(pg16, (const uint8_t *) kvalues_mxfp4);
+
+        // Reorder Q8 four-byte groups as [0,2,4,6,1,3,5,7]
+        const svuint32_t q8_group_index = svsub_n_u32_m(pg_hi4, svindex_u32(0, 2), 7);
+
+        const block_q8_0 * a_ptr = (const block_q8_0 *) vy;
+        float * res_ptr = s;
+
+        for (int x = 0; x < nc / ncols_interleaved; x++) {
+            const block_mxfp4x4 * b_ptr = (const block_mxfp4x4 *) vx + (x * nb);
+            svfloat32_t sumf = svdup_n_f32(0.0f);
+
+            for (int l = 0; l < nb; l++) {
+                const svint8_t q8 = svld1_s8(pgb, a_ptr[l].qs);
+                const svint8_t q8_even_odd = svreinterpret_s8_u32(svtbl_u32(svreinterpret_u32_s8(q8), q8_group_index));
+                const svuint8_t qmx01 = svld1_u8(pgb, b_ptr[l].qs + 0);
+                const svuint8_t qmx_lo01 = svtbl_u8(values, svand_u8_x(pgb, qmx01, m4b));
+                const svuint8_t qmx_hi01 = svtbl_u8(values, svlsr_n_u8_x(pgb, qmx01, 4));
+                const svuint8_t qmx23 = svld1_u8(pgb, b_ptr[l].qs + 32);
+                const svuint8_t qmx_lo23 = svtbl_u8(values, svand_u8_x(pgb, qmx23, m4b));
+                const svuint8_t qmx_hi23 = svtbl_u8(values, svlsr_n_u8_x(pgb, qmx23, 4));
+
+                svint32_t sumi = svdup_n_s32(0);
+                sumi = svdot_lane_s32(sumi, svreinterpret_s8_u8(qmx_lo01), q8_even_odd, 0);
+                sumi = svdot_lane_s32(sumi, svreinterpret_s8_u8(qmx_lo23), q8_even_odd, 1);
+                sumi = svdot_lane_s32(sumi, svreinterpret_s8_u8(qmx_hi01), q8_even_odd, 2);
+                sumi = svdot_lane_s32(sumi, svreinterpret_s8_u8(qmx_hi23), q8_even_odd, 3);
+                const svint32_t sumi_cols = svadd_s32_z(pg4, sumi, svext_s32(sumi, sumi, 4));
+
+                // Vector form of GGML_CPU_E8M0_TO_FP32_HALF
+                const svuint32_t e_u32 = svld1ub_u32(pg4, b_ptr[l].e);
+                const svbool_t e_is_denormal = svcmplt_n_u32(pg4, e_u32, 2);
+                const svuint32_t e_denormal_bits = svlsl_u32_x(pg4, svdup_n_u32(0x00200000), e_u32);
+                const svuint32_t e_normal_bits = svlsl_n_u32_x(pg4, svsub_n_u32_x(pg4, e_u32, 1), 23);
+                const svuint32_t e_bits = svsel_u32(e_is_denormal, e_denormal_bits, e_normal_bits);
+                const svfloat16_t d_f16 = svreinterpret_f16_u16(svdup_n_u16(a_ptr[l].d));
+                const svfloat32_t d_f32 = svcvt_f32_f16_x(pg4, d_f16);
+                const svfloat32_t scale = svmul_f32_x(pg4, svreinterpret_f32_u32(e_bits), d_f32);
+                sumf = svmla_f32_m(pg4, sumf,svcvt_f32_s32_m(svdup_n_f32(0.0f), pg4, sumi_cols), scale);
+            }
+            svst1_f32(pg4, res_ptr + x * 4, sumf);
+        }
+        return;
+    }
+#endif // defined(__aarch64__) && defined(__ARM_FEATURE_SVE)
+
 #if ! ((defined(_MSC_VER)) && ! defined(__clang__)) && defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD)
    const int8x16_t kvalues = vld1q_s8(kvalues_mxfp4);
    const block_q8_0 * a_ptr = (const block_q8_0 *) vy;
@@ -1791,6 +1845,71 @@ void ggml_gemv_q8_0_4x8_q8_0(int                        n,
    UNUSED(ncols_interleaved);
    UNUSED(blocklen);

+#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE)
+    // 256-bit SVE: 32 x int8 -> 8 x int32
+    if (ggml_cpu_get_sve_cnt() == QK8_0) {
+        const block_q8_0x4 * GGML_RESTRICT b_ptr = (const block_q8_0x4 *) vx;
+
+        const svbool_t pg   = svptrue_b8();
+        const svbool_t pg4w = svptrue_pat_b32(SV_VL4);
+
+        // after svtbl + svcvt the scale lanes are [d0,d0,d1,d1,d2,d2,d3,d3]: lane pair j = row j
+        const svuint16_t d_idx = svreinterpret_u16_u32(svlsr_n_u32_x(pg, svindex_u32(0, 1), 1));
+
+        for (int c = 0; c < nc; c += ncols_interleaved) {
+            const block_q8_0 * GGML_RESTRICT a_ptr = (const block_q8_0 *) vy;
+
+            svfloat32_t acc = svdup_n_f32(0.0f);
+
+            for (int b = 0; b < nb; ++b) {
+                const int8_t * GGML_RESTRICT bq = (const int8_t *) b_ptr->qs;
+
+                uint64_t a_chunk0;
+                uint64_t a_chunk1;
+                uint64_t a_chunk2;
+                uint64_t a_chunk3;
+                uint64_t b_d;
+
+                memcpy(&a_chunk0, a_ptr->qs +  0, sizeof(a_chunk0));
+                memcpy(&a_chunk1, a_ptr->qs +  8, sizeof(a_chunk1));
+                memcpy(&a_chunk2, a_ptr->qs + 16, sizeof(a_chunk2));
+                memcpy(&a_chunk3, a_ptr->qs + 24, sizeof(a_chunk3));
+                memcpy(&b_d, b_ptr->d, sizeof(b_d));
+
+                const svint8_t aq0 = svreinterpret_s8_u64(svdup_n_u64(a_chunk0));
+                const svint8_t aq1 = svreinterpret_s8_u64(svdup_n_u64(a_chunk1));
+                const svint8_t aq2 = svreinterpret_s8_u64(svdup_n_u64(a_chunk2));
+                const svint8_t aq3 = svreinterpret_s8_u64(svdup_n_u64(a_chunk3));
+
+                const svint8_t bq0 = svld1_s8(pg, bq +  0);
+                const svint8_t bq1 = svld1_s8(pg, bq + 32);
+                const svint8_t bq2 = svld1_s8(pg, bq + 64);
+                const svint8_t bq3 = svld1_s8(pg, bq + 96);
+
+                const svint32_t iacc0 = svdot_s32(svdup_n_s32(0), bq0, aq0);
+                const svint32_t iacc1 = svdot_s32(svdup_n_s32(0), bq1, aq1);
+                const svint32_t iacc2 = svdot_s32(svdup_n_s32(0), bq2, aq2);
+                const svint32_t iacc3 = svdot_s32(svdup_n_s32(0), bq3, aq3);
+
+                const svint32_t isum = svadd_s32_x(pg, svadd_s32_x(pg, iacc0, iacc1), svadd_s32_x(pg, iacc2, iacc3));
+
+                const svfloat16_t bd_h  = svtbl_f16(svreinterpret_f16_u64(svdup_n_u64(b_d)), d_idx);
+                const svfloat32_t scale = svmul_n_f32_x(pg, svcvt_f32_f16_x(pg, bd_h), GGML_CPU_FP16_TO_FP32(a_ptr->d));
+
+                acc = svmla_f32_x(pg, acc, svcvt_f32_s32_x(pg, isum), scale);
+
+                ++a_ptr;
+                ++b_ptr;
+            }
+
+            svst1_f32(pg4w, s, svadd_f32_x(pg, svuzp1_f32(acc, acc), svuzp2_f32(acc, acc)));
+            s += ncols_interleaved;
+        }
+
+        return;
+    }
+#endif  // defined(__aarch64__) && defined(__ARM_FEATURE_SVE)
+
 #if defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_DOTPROD)
    const block_q8_0x4 * b_ptr = (const block_q8_0x4 *) vx;

@@ -4874,6 +4993,7 @@ void ggml_gemm_q6_K_8x8_q8_K(int                        n,

    constexpr int ncols_interleaved = 8;
    constexpr int blocklen          = 8;
+    constexpr int q8_k_blocklen     = 4;

    assert(n % qk == 0);
    assert(nr % 4 == 0);
@@ -4882,9 +5002,195 @@ void ggml_gemm_q6_K_8x8_q8_K(int                        n,
    UNUSED(nb);
    UNUSED(ncols_interleaved);
    UNUSED(blocklen);
+    UNUSED(q8_k_blocklen);
+
+#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(__ARM_FEATURE_MATMUL_INT8)
+    if (svcntb() * 8 == 256) {
+        const svbool_t pg8  = svptrue_b8();
+        const svbool_t pg32 = svptrue_b32();
+
+        const svuint8_t low_nibble = svdup_n_u8(0x0f);
+        const svuint8_t high_lo2   = svdup_n_u8(0x03);
+        const svuint8_t high_hi2   = svdup_n_u8(0x30);
+        const svint8_t  q6_bias    = svdup_n_s8(32);
+        const svint32_t zero_s32   = svdup_n_s32(0);
+
+        for (int y = 0; y < nr / q8_k_blocklen; ++y) {
+            const block_q8_Kx4 * GGML_RESTRICT q8_ptr = (const block_q8_Kx4 *) vy + y * nb;
+
+            for (int x = 0; x < nc / ncols_interleaved; ++x) {
+                const block_q6_Kx8 * GGML_RESTRICT q6_ptr = (const block_q6_Kx8 *) vx + x * nb;
+
+                // one accumlator per output row
+                svfloat32_t acc_f32_r0 = svdup_n_f32(0);
+                svfloat32_t acc_f32_r1 = svdup_n_f32(0);
+                svfloat32_t acc_f32_r2 = svdup_n_f32(0);
+                svfloat32_t acc_f32_r3 = svdup_n_f32(0);
+
+                for (int b = 0; b < nb; ++b) {
+                    svint32_t acc_c03_r01 = zero_s32;
+                    svint32_t acc_c03_r23 = zero_s32;
+                    svint32_t acc_c47_r01 = zero_s32;
+                    svint32_t acc_c47_r23 = zero_s32;
+
+                    for (int half = 0; half < 2; ++half) {
+                        const uint8_t * ql_base = q6_ptr[b].ql + half * 512;
+                        const uint8_t * qh_base = q6_ptr[b].qh + half * 256;
+
+                        for (int sb = 0; sb < QK_K / 64; ++sb) {
+                            const int8_t * q8_base_l = q8_ptr[b].qs + half * 512 + sb * 64;
+                            const int8_t * q8_base_h = q8_ptr[b].qs + half * 512 + 256 + sb * 64;
+
+                            // load 16 bytes (8B of R0, 8B of R1) and repeat in upper half
+                            // low lanes will multiply with 16 bytes of C0 and high lanes will multiply with 16 bytes of C1
+                            const svint8_t q8_l_r01_k0 = svld1rq_s8(pg8, q8_base_l + 0);
+                            const svint8_t q8_l_r01_k1 = svld1rq_s8(pg8, q8_base_l + 32);
+                            const svint8_t q8_l_r23_k0 = svld1rq_s8(pg8, q8_base_l + 16);
+                            const svint8_t q8_l_r23_k1 = svld1rq_s8(pg8, q8_base_l + 48);
+
+                            const svint8_t q8_h_r01_k0 = svld1rq_s8(pg8, q8_base_h + 0);
+                            const svint8_t q8_h_r01_k1 = svld1rq_s8(pg8, q8_base_h + 32);
+                            const svint8_t q8_h_r23_k0 = svld1rq_s8(pg8, q8_base_h + 16);
+                            const svint8_t q8_h_r23_k1 = svld1rq_s8(pg8, q8_base_h + 48);
+
+                            const int ql_off = sb * QK_K / 2;
+                            const int qh_off = ql_off & 255;
+
+                            const int scale_idx_l = half * 8 + sb;
+                            const int scale_idx_h = scale_idx_l + 4;
+
+                            const svint32_t scales_l_all = svld1sb_s32(pg32, q6_ptr[b].scales + scale_idx_l * 8);
+                            const svint32_t scales_h_all = svld1sb_s32(pg32, q6_ptr[b].scales + scale_idx_h * 8);
+
+                            const svint32_t scales_l_c03 = svzip1_s32(scales_l_all, scales_l_all);
+                            const svint32_t scales_h_c03 = svzip1_s32(scales_h_all, scales_h_all);
+
+                            {
+                                svuint8_t q6_ql0_01 = svld1_u8(pg8, ql_base + ql_off);
+                                svuint8_t q6_ql1_01 = svld1_u8(pg8, ql_base + ql_off + 64);
+                                svuint8_t q6_qh0_01 = svld1_u8(pg8, qh_base + qh_off);
+                                svuint8_t q6_qh1_01 = svld1_u8(pg8, qh_base + qh_off + 64);
+
+                                if (sb > 1) {
+                                    q6_qh0_01 = svlsr_n_u8_x(pg8, q6_qh0_01, 2);
+                                    q6_qh1_01 = svlsr_n_u8_x(pg8, q6_qh1_01, 2);
+                                }
+
+                                const svuint8_t q6_qh0_01_hi = svand_u8_x(pg8, q6_qh0_01, high_hi2);
+                                const svuint8_t q6_qh1_01_hi = svand_u8_x(pg8, q6_qh1_01, high_hi2);
+
+                                const svuint8_t q6_l0_01_bits = svorr_u8_x(pg8, svand_u8_x(pg8, q6_ql0_01, low_nibble), svlsl_n_u8_x(pg8, svand_u8_x(pg8, q6_qh0_01, high_lo2), 4));
+                                const svuint8_t q6_l1_01_bits = svorr_u8_x(pg8, svand_u8_x(pg8, q6_ql1_01, low_nibble), svlsl_n_u8_x(pg8, svand_u8_x(pg8, q6_qh1_01, high_lo2), 4));
+                                const svuint8_t q6_h0_01_bits = svorr_u8_x(pg8, svlsr_n_u8_x(pg8, q6_ql0_01, 4), q6_qh0_01_hi);
+                                const svuint8_t q6_h1_01_bits = svorr_u8_x(pg8, svlsr_n_u8_x(pg8, q6_ql1_01, 4), q6_qh1_01_hi);
+
+                                const svint8_t q6_l0_01 = svsub_s8_x(pg8, svreinterpret_s8_u8(q6_l0_01_bits), q6_bias);
+                                const svint8_t q6_l1_01 = svsub_s8_x(pg8, svreinterpret_s8_u8(q6_l1_01_bits), q6_bias);
+                                const svint8_t q6_h0_01 = svsub_s8_x(pg8, svreinterpret_s8_u8(q6_h0_01_bits), q6_bias);
+                                const svint8_t q6_h1_01 = svsub_s8_x(pg8, svreinterpret_s8_u8(q6_h1_01_bits), q6_bias);
+
+                                svint32_t dot_l_r01 = svmmla_s32(zero_s32, q6_l0_01, q8_l_r01_k0);
+                                dot_l_r01 = svmmla_s32(dot_l_r01, q6_l1_01, q8_l_r01_k1);
+
+                                svint32_t dot_h_r01 = svmmla_s32(zero_s32, q6_h0_01, q8_h_r01_k0);
+                                dot_h_r01 = svmmla_s32(dot_h_r01, q6_h1_01, q8_h_r01_k1);
+
+                                svint32_t dot_l_r23 = svmmla_s32(zero_s32, q6_l0_01, q8_l_r23_k0);
+                                dot_l_r23 = svmmla_s32(dot_l_r23, q6_l1_01, q8_l_r23_k1);
+
+                                svint32_t dot_h_r23 = svmmla_s32(zero_s32, q6_h0_01, q8_h_r23_k0);
+                                dot_h_r23 = svmmla_s32(dot_h_r23, q6_h1_01, q8_h_r23_k1);
+
+                                acc_c03_r01 = svmla_s32_x(pg32, acc_c03_r01, dot_l_r01, scales_l_c03);
+                                acc_c03_r01 = svmla_s32_x(pg32, acc_c03_r01, dot_h_r01, scales_h_c03);
+                                acc_c03_r23 = svmla_s32_x(pg32, acc_c03_r23, dot_l_r23, scales_l_c03);
+                                acc_c03_r23 = svmla_s32_x(pg32, acc_c03_r23, dot_h_r23, scales_h_c03);
+                            }
+
+                            const svint32_t scales_l_c47 = svzip2_s32(scales_l_all, scales_l_all);
+                            const svint32_t scales_h_c47 = svzip2_s32(scales_h_all, scales_h_all);
+
+                            {
+                                svuint8_t q6_ql0_23 = svld1_u8(pg8, ql_base + ql_off + 32);
+                                svuint8_t q6_ql1_23 = svld1_u8(pg8, ql_base + ql_off + 96);
+                                svuint8_t q6_qh0_23 = svld1_u8(pg8, qh_base + qh_off + 32);
+                                svuint8_t q6_qh1_23 = svld1_u8(pg8, qh_base + qh_off + 96);
+
+                                if (sb > 1) {
+                                    q6_qh0_23 = svlsr_n_u8_x(pg8, q6_qh0_23, 2);
+                                    q6_qh1_23 = svlsr_n_u8_x(pg8, q6_qh1_23, 2);
+                                }
+
+                                const svuint8_t q6_qh0_23_hi = svand_u8_x(pg8, q6_qh0_23, high_hi2);
+                                const svuint8_t q6_qh1_23_hi = svand_u8_x(pg8, q6_qh1_23, high_hi2);
+
+                                const svuint8_t q6_l0_23_bits = svorr_u8_x(pg8, svand_u8_x(pg8, q6_ql0_23, low_nibble), svlsl_n_u8_x(pg8, svand_u8_x(pg8, q6_qh0_23, high_lo2), 4));
+                                const svuint8_t q6_l1_23_bits = svorr_u8_x(pg8, svand_u8_x(pg8, q6_ql1_23, low_nibble), svlsl_n_u8_x(pg8, svand_u8_x(pg8, q6_qh1_23, high_lo2), 4));
+                                const svuint8_t q6_h0_23_bits = svorr_u8_x(pg8, svlsr_n_u8_x(pg8, q6_ql0_23, 4), q6_qh0_23_hi);
+                                const svuint8_t q6_h1_23_bits = svorr_u8_x(pg8, svlsr_n_u8_x(pg8, q6_ql1_23, 4), q6_qh1_23_hi);
+
+                                const svint8_t q6_l0_23 = svsub_s8_x(pg8, svreinterpret_s8_u8(q6_l0_23_bits), q6_bias);
+                                const svint8_t q6_l1_23 = svsub_s8_x(pg8, svreinterpret_s8_u8(q6_l1_23_bits), q6_bias);
+                                const svint8_t q6_h0_23 = svsub_s8_x(pg8, svreinterpret_s8_u8(q6_h0_23_bits), q6_bias);
+                                const svint8_t q6_h1_23 = svsub_s8_x(pg8, svreinterpret_s8_u8(q6_h1_23_bits), q6_bias);
+
+                                svint32_t dot_l_r01 = svmmla_s32(zero_s32, q6_l0_23, q8_l_r01_k0);
+                                dot_l_r01 = svmmla_s32(dot_l_r01, q6_l1_23, q8_l_r01_k1);
+
+                                svint32_t dot_h_r01 = svmmla_s32(zero_s32, q6_h0_23, q8_h_r01_k0);
+                                dot_h_r01 = svmmla_s32(dot_h_r01, q6_h1_23, q8_h_r01_k1);
+
+                                svint32_t dot_l_r23 = svmmla_s32(zero_s32, q6_l0_23, q8_l_r23_k0);
+                                dot_l_r23 = svmmla_s32(dot_l_r23, q6_l1_23, q8_l_r23_k1);
+
+                                svint32_t dot_h_r23 = svmmla_s32(zero_s32, q6_h0_23, q8_h_r23_k0);
+                                dot_h_r23 = svmmla_s32(dot_h_r23, q6_h1_23, q8_h_r23_k1);
+
+                                acc_c47_r01 = svmla_s32_x(pg32, acc_c47_r01, dot_l_r01, scales_l_c47);
+                                acc_c47_r01 = svmla_s32_x(pg32, acc_c47_r01, dot_h_r01, scales_h_c47);
+                                acc_c47_r23 = svmla_s32_x(pg32, acc_c47_r23, dot_l_r23, scales_l_c47);
+                                acc_c47_r23 = svmla_s32_x(pg32, acc_c47_r23, dot_h_r23, scales_h_c47);
+                            }
+                        }
+                    }
+
+                    // acc_c03_r01 -> c0-r0,c0-r1,c1-r0,c1-r1,c2-r0,c2-r1,c3-r0,c3-r1
+                    // acc_c47_r01 -> c4-r0,c4-r1,c5-r0,c5-r1,c6-r0,c6-r1,c7-r0,c7-r1
+                    // svuzp1_s32 picks even lanes -> c0-r0,c1-r0,c2-r0,c3-r0,c4-r0,c5-r0,c6-r0,c7-r0
+                    // similarly svuzp2_s32 picks remaining values from odd lanes
+                    const svint32_t dot_r0 = svuzp1_s32(acc_c03_r01, acc_c47_r01);
+                    const svint32_t dot_r1 = svuzp2_s32(acc_c03_r01, acc_c47_r01);
+                    const svint32_t dot_r2 = svuzp1_s32(acc_c03_r23, acc_c47_r23);
+                    const svint32_t dot_r3 = svuzp2_s32(acc_c03_r23, acc_c47_r23);
+
+                    const svuint32_t q6_d_bits = svld1uh_u32(pg32, (const uint16_t *) q6_ptr[b].d);
+                    const svfloat32_t q6_d_f32 = svcvt_f32_f16_x(pg32, svreinterpret_f16_u32(q6_d_bits));
+
+                    const svfloat32_t scale_r0 = svmul_n_f32_x(pg32, q6_d_f32, q8_ptr[b].d[0]);
+                    const svfloat32_t scale_r1 = svmul_n_f32_x(pg32, q6_d_f32, q8_ptr[b].d[1]);
+                    const svfloat32_t scale_r2 = svmul_n_f32_x(pg32, q6_d_f32, q8_ptr[b].d[2]);
+                    const svfloat32_t scale_r3 = svmul_n_f32_x(pg32, q6_d_f32, q8_ptr[b].d[3]);
+
+                    acc_f32_r0 = svmla_f32_x(pg32, acc_f32_r0, svcvt_f32_s32_x(pg32, dot_r0), scale_r0);
+                    acc_f32_r1 = svmla_f32_x(pg32, acc_f32_r1, svcvt_f32_s32_x(pg32, dot_r1), scale_r1);
+                    acc_f32_r2 = svmla_f32_x(pg32, acc_f32_r2, svcvt_f32_s32_x(pg32, dot_r2), scale_r2);
+                    acc_f32_r3 = svmla_f32_x(pg32, acc_f32_r3, svcvt_f32_s32_x(pg32, dot_r3), scale_r3);
+                }
+
+                const int row = y * q8_k_blocklen;
+                const int col = x * ncols_interleaved;
+
+                svst1_f32(pg32, s + (row + 0) * bs + col, acc_f32_r0);
+                svst1_f32(pg32, s + (row + 1) * bs + col, acc_f32_r1);
+                svst1_f32(pg32, s + (row + 2) * bs + col, acc_f32_r2);
+                svst1_f32(pg32, s + (row + 3) * bs + col, acc_f32_r3);
+            }
+        }
+        return;
+    }
+#endif  // SVE compile-time end

 #if defined(__aarch64__) && defined(__ARM_NEON) && defined(__ARM_FEATURE_MATMUL_INT8)
-    constexpr int    q8_k_blocklen = 4;
    const uint8x16_t m4b           = vdupq_n_u8(0x0f);
    const uint8x16_t mask_lo       = vdupq_n_u8(0x03);
    const uint8x16_t mask_hi       = vdupq_n_u8(0x30);
@@ -1914,213 +1914,87 @@ void ggml_compute_forward_repeat_back(
 }

 // ggml_compute_forward_concat
-static void ggml_compute_forward_concat_any(
-    const ggml_compute_params * params,
-    ggml_tensor * dst) {

-    const ggml_tensor * src0 = dst->src[0];
-    const ggml_tensor * src1 = dst->src[1];
+// rows of dst copied together, so that a strided source is read along neighboring rows
+static constexpr int64_t GGML_CONCAT_TILE = 32;

-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    GGML_TENSOR_BINARY_OP_LOCALS
-
-    const int32_t dim = ggml_get_op_params_i32(dst, 0);
-
-    GGML_ASSERT(dim >= 0 && dim < 4);
-    GGML_ASSERT(ggml_is_contiguous_rows(src0));
-    GGML_ASSERT(ggml_is_contiguous_rows(src1));
-
-    int64_t o[4] = {0, 0, 0, 0};
-
-    if (dim == 0) {
-        GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
-        GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
-
-        o[dim] = src0->ne[dim]/ggml_blck_size(src0->type);
-    } else {
-        o[dim] = src0->ne[dim];
-    }
-
-    // Region 1: copy rows from src0
-    for (int i3 = 0; i3 < ne03; i3++) {
-        for (int i2 = ith; i2 < ne02; i2 += nth) {
-            for (int i1 = 0; i1 < ne01; i1++) {
-                const char * x = (const char *) src0->data + i1*nb01 + i2*nb02 + i3*nb03;
-                      char * y = (      char *) dst->data  + i1*nb1  + i2*nb2  + i3*nb3;
-                memcpy(y, x, ggml_row_size(src0->type, ne00));
-            }
-        }
-    }
-
-    // Region 2: copy rows from src1, offset into dst by o[]
-    for (int i3 = 0; i3 < ne13; i3++) {
-        for (int i2 = ith; i2 < ne12; i2 += nth) {
-            for (int i1 = 0; i1 < ne11; i1++) {
-                const char * x = (const char *) src1->data + i1*nb11         + i2*nb12         + i3*nb13;
-                      char * y = (      char *)  dst->data + (i1 + o[1])*nb1 + (i2 + o[2])*nb2 + (i3 + o[3])*nb3 + o[0]*nb0;
-                memcpy(y, x, ggml_row_size(src1->type, ne10));
-            }
+// copy n_rows rows of n elements of es bytes from a source with element stride nb0, element by element
+template <size_t es>
+static void ggml_concat_copy_strided(char * y, size_t nby, const char * x, size_t nbx, size_t nb0, int64_t n, int64_t n_rows) {
+    for (int64_t i0 = 0; i0 < n; i0++) {
+        for (int64_t r = 0; r < n_rows; r++) {
+            memcpy(y + r*nby + i0*es, x + r*nbx + i0*nb0, es);
        }
    }
 }

-static void ggml_compute_forward_concat_i8(
-    const ggml_compute_params * params,
-    ggml_tensor * dst) {
+// copy n_rows rows of n elements of src into dst, rows are nbx bytes apart in src and nby bytes apart in dst
+static void ggml_concat_copy_rows(char * y, size_t nby, const ggml_tensor * src, const char * x, size_t nbx, int64_t n, int64_t n_rows) {
+    const size_t es = ggml_type_size(src->type);

-    const ggml_tensor * src0 = dst->src[0];
-    const ggml_tensor * src1 = dst->src[1];
-
-    GGML_ASSERT(ggml_type_size(src0->type) == sizeof(int8_t));
-
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    GGML_TENSOR_BINARY_OP_LOCALS
-
-    const int32_t dim = ggml_get_op_params_i32(dst, 0);
-
-    GGML_ASSERT(dim >= 0 && dim < 4);
-
-    int64_t o[4] = {0, 0, 0, 0};
-    o[dim] = src0->ne[dim];
-
-    const int8_t * x;
-
-    // TODO: smarter multi-theading
-    for (int i3 = 0; i3 < ne3; i3++) {
-        for (int i2 = ith; i2 < ne2; i2 += nth) {
-            for (int i1 = 0; i1 < ne1; i1++) {
-                for (int i0 = 0; i0 < ne0; i0++) {
-                    if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) {
-                        x = (const int8_t *) ((const char *)src0->data + (i0       )*nb00 + (i1       )*nb01 + (i2       )*nb02 + (i3       )*nb03);
-                    } else {
-                        x = (const int8_t *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13);
-                    }
-
-                    int8_t * y = (int8_t *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3);
-
-                    *y = *x;
-                }
-            }
-        }
-    }
-}
-
-static void ggml_compute_forward_concat_f16(
-    const ggml_compute_params * params,
-    ggml_tensor * dst) {
-
-    const ggml_tensor * src0 = dst->src[0];
-    const ggml_tensor * src1 = dst->src[1];
-
-    GGML_ASSERT(ggml_type_size(src0->type) == sizeof(ggml_fp16_t));
-
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    GGML_TENSOR_BINARY_OP_LOCALS
-
-    const int32_t dim = ggml_get_op_params_i32(dst, 0);
-
-    GGML_ASSERT(dim >= 0 && dim < 4);
-
-    int64_t o[4] = {0, 0, 0, 0};
-    o[dim] = src0->ne[dim];
-
-    const ggml_fp16_t * x;
-
-    // TODO: smarter multi-theading
-    for (int i3 = 0; i3 < ne3; i3++) {
-        for (int i2 = ith; i2 < ne2; i2 += nth) {
-            for (int i1 = 0; i1 < ne1; i1++) {
-                for (int i0 = 0; i0 < ne0; i0++) {
-                    if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) {
-                        x = (const ggml_fp16_t *) ((const char *)src0->data + (i0       )*nb00 + (i1       )*nb01 + (i2       )*nb02 + (i3       )*nb03);
-                    } else {
-                        x = (const ggml_fp16_t *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13);
-                    }
-
-                    ggml_fp16_t * y = (ggml_fp16_t *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3);
-
-                    *y = *x;
-                }
-            }
-        }
-    }
-}
-
-static void ggml_compute_forward_concat_f32(
-    const ggml_compute_params * params,
-    ggml_tensor * dst) {
-
-    const ggml_tensor * src0 = dst->src[0];
-    const ggml_tensor * src1 = dst->src[1];
-
-    GGML_ASSERT(ggml_type_size(src0->type) == sizeof(float));
-
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    GGML_TENSOR_BINARY_OP_LOCALS
-
-    const int32_t dim = ggml_get_op_params_i32(dst, 0);
-
-    GGML_ASSERT(dim >= 0 && dim < 4);
-
-    int64_t o[4] = {0, 0, 0, 0};
-    o[dim] = src0->ne[dim];
-
-    const float * x;
-
-    // TODO: smarter multi-theading
-    for (int i3 = 0; i3 < ne3; i3++) {
-        for (int i2 = ith; i2 < ne2; i2 += nth) {
-            for (int i1 = 0; i1 < ne1; i1++) {
-                for (int i0 = 0; i0 < ne0; i0++) {
-                    if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) {
-                        x = (const float *) ((const char *)src0->data + (i0       )*nb00 + (i1       )*nb01 + (i2       )*nb02 + (i3       )*nb03);
-                    } else {
-                        x = (const float *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13);
-                    }
-
-                    float * y = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3);
-
-                    *y = *x;
-                }
-            }
+    if (src->nb[0] == es) {
+        const size_t row_size = ggml_row_size(src->type, n);
+        for (int64_t r = 0; r < n_rows; r++) {
+            memcpy(y + r*nby, x + r*nbx, row_size);
        }
+        return;
+    }
+
+    switch (es) {
+        case 1: ggml_concat_copy_strided<1>(y, nby, x, nbx, src->nb[0], n, n_rows); break;
+        case 2: ggml_concat_copy_strided<2>(y, nby, x, nbx, src->nb[0], n, n_rows); break;
+        case 4: ggml_concat_copy_strided<4>(y, nby, x, nbx, src->nb[0], n, n_rows); break;
+        case 8: ggml_concat_copy_strided<8>(y, nby, x, nbx, src->nb[0], n, n_rows); break;
+        default: GGML_ABORT("unsupported element size");
    }
 }

+// each dst row comes from src0 or src1, or for dim 0 from the row of src0 followed by the row of src1
 void ggml_compute_forward_concat(
    const ggml_compute_params * params,
    ggml_tensor * dst) {

    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];

-    switch (src0->type) {
-        case GGML_TYPE_F16:
-        case GGML_TYPE_BF16:
-        case GGML_TYPE_I16:
-            {
-                ggml_compute_forward_concat_f16(params, dst);
-            } break;
-        case GGML_TYPE_I8:
-            {
-                ggml_compute_forward_concat_i8(params, dst);
-            } break;
-        case GGML_TYPE_F32:
-        case GGML_TYPE_I32:
-            {
-                ggml_compute_forward_concat_f32(params, dst);
-            } break;
-        default:
-            {
-                ggml_compute_forward_concat_any(params, dst);
-            }
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    const int32_t dim = ggml_get_op_params_i32(dst, 0);
+
+    GGML_ASSERT(dim >= 0 && dim < 4);
+    GGML_ASSERT(nb0 == ggml_type_size(dst->type));
+    GGML_ASSERT(!ggml_is_quantized(dst->type) || (nb00 == nb0 && nb10 == nb0));
+    GGML_ASSERT(dim != 0 || (ne00 % ggml_blck_size(dst->type) == 0 && ne10 % ggml_blck_size(dst->type) == 0));
+
+    const size_t y1 = dim == 0 ? ggml_row_size(dst->type, ne00) : 0;
+
+    const int64_t nr = ne1*ne2*ne3;
+    const int64_t dr = (nr + params->nth - 1)/params->nth;
+
+    const int64_t ir0 = dr*params->ith;
+    const int64_t ir1 = MIN(ir0 + dr, nr);
+
+    for (int64_t ir = ir0; ir < ir1; ) {
+        int64_t i[4] = { 0, ir % ne1, (ir/ne1) % ne2, ir/(ne1*ne2) };
+
+        // a tile of consecutive rows sharing i2, i3 and their source
+        int64_t nt = MIN(MIN(ir1 - ir, ne1 - i[1]), GGML_CONCAT_TILE);
+        if (dim == 1 && i[1] < ne01) {
+            nt = MIN(nt, ne01 - i[1]);
+        }
+
+        char * y = (char *) dst->data + i[1]*nb1 + i[2]*nb2 + i[3]*nb3;
+
+        if (dim == 0 || i[dim] < src0->ne[dim]) {
+            ggml_concat_copy_rows(y, nb1, src0, (const char *) src0->data + i[1]*nb01 + i[2]*nb02 + i[3]*nb03, nb01, ne00, nt);
+        }
+
+        if (dim == 0 || i[dim] >= src0->ne[dim]) {
+            i[dim] -= src0->ne[dim]; // i[0] does not address rows
+            ggml_concat_copy_rows(y + y1, nb1, src1, (const char *) src1->data + i[1]*nb11 + i[2]*nb12 + i[3]*nb13, nb11, ne10, nt);
+        }
+
+        ir += nt;
    }
 }

@@ -4972,8 +4846,8 @@ static void ggml_compute_forward_scale_f32(

    const ggml_tensor * src0 = dst->src[0];

-    GGML_ASSERT(ggml_is_contiguous(src0));
-    GGML_ASSERT(ggml_is_contiguous(dst));
+    GGML_ASSERT(ggml_is_contiguous_rows(src0));
+    GGML_ASSERT(ggml_is_contiguous_rows(dst));
    GGML_ASSERT(ggml_are_same_shape(src0, dst));

    float s; // scale factor
@@ -4988,6 +4862,8 @@ static void ggml_compute_forward_scale_f32(
    const int nc = src0->ne[0];
    const int nr = ggml_nrows(src0);

+    GGML_TENSOR_UNARY_OP_LOCALS
+
    // rows per thread
    const int dr = (nr + nth - 1)/nth;

@@ -4995,25 +4871,23 @@ static void ggml_compute_forward_scale_f32(
    const int ir0 = dr*ith;
    const int ir1 = MIN(ir0 + dr, nr);

-    const size_t nb01 = src0->nb[1];
+    for (int i = ir0; i < ir1; i++) {
+        const int64_t i1 = i % ne01;
+        const int64_t i2 = (i / ne01) % ne02;
+        const int64_t i3 = i / (ne01*ne02);

-    const size_t nb1 = dst->nb[1];
+        float *       y = (float *)       ((char *) dst->data  + i1*nb1  + i2*nb2  + i3*nb3);
+        const float * x = (const float *) ((char *) src0->data + i1*nb01 + i2*nb02 + i3*nb03);

-    if (b == 0.0f) {
-        for (int i1 = ir0; i1 < ir1; i1++) {
+        if (b == 0.0f) {
            if (dst->data != src0->data) {
                // src0 is same shape as dst => same indices
                // TODO: add x parameter to ggml_vec_scale_f32 and remove this memcpy
-                memcpy((char *)dst->data + i1*nb1, (char *)src0->data + i1*nb01, nc * sizeof(float));
+                memcpy(y, x, nc * sizeof(float));
            }
-            ggml_vec_scale_f32(nc, (float *) ((char *) dst->data + i1*nb1), s);
-        }
-    } else {
-        for (int i1 = ir0; i1 < ir1; i1++) {
-            ggml_vec_mad1_f32(nc,
-                (float *) ((char *) dst->data  + i1*nb1),
-                (float *) ((char *) src0->data + i1*nb1),
-                s, b);
+            ggml_vec_scale_f32(nc, y, s);
+        } else {
+            ggml_vec_mad1_f32(nc, y, x, s, b);
        }
    }
 }
@@ -5024,8 +4898,8 @@ static void ggml_compute_forward_scale_bf16(

    const ggml_tensor * src0 = dst->src[0];

-    GGML_ASSERT(ggml_is_contiguous(src0));
-    GGML_ASSERT(ggml_is_contiguous(dst));
+    GGML_ASSERT(ggml_is_contiguous_rows(src0));
+    GGML_ASSERT(ggml_is_contiguous_rows(dst));
    GGML_ASSERT(ggml_are_same_shape(src0, dst));

    float s; // scale factor
@@ -5040,6 +4914,8 @@ static void ggml_compute_forward_scale_bf16(
    const int nc = src0->ne[0];
    const int nr = ggml_nrows(src0);

+    GGML_TENSOR_UNARY_OP_LOCALS
+
    // rows per thread
    const int dr = (nr + nth - 1)/nth;

@@ -5047,9 +4923,13 @@ static void ggml_compute_forward_scale_bf16(
    const int ir0 = dr*ith;
    const int ir1 = MIN(ir0 + dr, nr);

-    for (int i1 = ir0; i1 < ir1; i1++) {
-        const ggml_bf16_t * x = (const ggml_bf16_t *) ((const char *) src0->data + i1*src0->nb[1]);
-        ggml_bf16_t       * y = (ggml_bf16_t *)       ((char *)       dst->data  + i1*dst->nb[1]);
+    for (int i = ir0; i < ir1; i++) {
+        const int64_t i1 = i % ne01;
+        const int64_t i2 = (i / ne01) % ne02;
+        const int64_t i3 = i / (ne01*ne02);
+
+        const ggml_bf16_t * x = (const ggml_bf16_t *) ((const char *) src0->data + i1*nb01 + i2*nb02 + i3*nb03);
+        ggml_bf16_t       * y = (ggml_bf16_t *)       ((char *)       dst->data  + i1*nb1  + i2*nb2  + i3*nb3);
        for (int k = 0; k < nc; k++) {
            y[k] = GGML_FP32_TO_BF16(s*GGML_BF16_TO_FP32(x[k]) + b);
        }
@@ -180,7 +180,6 @@ static __global__ void conv2d_implicit_gemm_f16(const half * __restrict__ input,
    const int                  end            = int(int64_t(tiles) * (split + 1) / split_k) * BK;
    for (int k0 = begin; k0 < end; k0 += BK) {
        if (k_total % 8 == 0 && uintptr_t(weight) % 16 == 0) {
-#pragma unroll
            for (int i = tid; i < BM * BK / 8; i += nthreads) {
                const int  row = i / (BK / 8), col = 8 * (i % (BK / 8));
                const int4 v                 = m0 + row < oc && k0 + col < k_total ?
@@ -189,7 +188,6 @@ static __global__ void conv2d_implicit_gemm_f16(const half * __restrict__ input,
                *(int4 *) &a_s[row][col / 2] = v;
            }
        } else {
-#pragma unroll
            for (int i = tid; i < BM * BK / 2; i += nthreads) {
                const int row = i / (BK / 2), col = 2 * (i % (BK / 2));
                half      lo = __float2half(0.0f), hi = lo;
@@ -202,7 +200,6 @@ static __global__ void conv2d_implicit_gemm_f16(const half * __restrict__ input,
                a_s[row][col / 2] = __halves2half2(lo, hi);
            }
        }
-#pragma unroll
        for (int k = load_row; k < BK; k += nthreads / (BN / 2)) {
            const int ki = k0 + k;
            const int ci = ki / (kw * kh), ky = ki / kw % kh, kx = ki % kw;
@@ -1866,7 +1866,7 @@ static __global__ void flash_attn_ext_f16(
 #endif // defined(AMD_WMMA_AVAILABLE)

 #if defined(AMD_MFMA_AVAILABLE)
-    if (ncols1*ncols2 < 16 || (DKQ > 256 && ncols1*ncols2 < 32)) {
+    if (ncols1*ncols2 < 16 || (DKQ > 256 && ncols1*ncols2 < 32 && !use_sparse)) {
        NO_DEVICE_CODE;
        return;
    }
@@ -2054,7 +2054,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
    bool use_sparse = false;
    if (logit_softcap == 0.0f) {
        constexpr bool use_logit_softcap = false;
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if !defined(GGML_USE_MUSA)
        if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2)) {
            if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, ncols1, ncols2)) {
                constexpr bool use_sparse_kernel = true;
@@ -2077,7 +2077,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
                }
            }
        } else
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#endif // !defined(GGML_USE_MUSA)
        {
            constexpr bool use_sparse_kernel = false;
            fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;
@@ -5,7 +5,8 @@
 #include "fattn-vec.cuh"
 #include "fattn.cuh"

-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if !defined(GGML_USE_MUSA)
+
 // one list per group of ncols1 queries: a column is selected if any query of the group can see it
 template <int ncols1, bool oob>
 __launch_bounds__(256, 1)
@@ -14,10 +15,20 @@ static __global__ void flash_attn_mask_to_sparse_indices(
        const int n_kv_max, const int64_t s31, const int64_t s33) {
    ggml_cuda_pdl_sync();

+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
    constexpr int values_per_lane = 8;
+    using lane_mask_t = std::conditional_t<warp_size == 64, uint64_t, uint32_t>;
+    constexpr uint64_t ballot_mask = ~(lane_mask_t)0;
+    const auto popc = [](const lane_mask_t mask) {
+        if constexpr (warp_size == 64) {
+            return __popcll(mask);
+        } else {
+            return __popc(mask);
+        }
+    };
    const int tid      = threadIdx.x;
-    const int warp     = tid / WARP_SIZE;
-    const int lane     = tid % WARP_SIZE;
+    const int warp     = tid / warp_size;
+    const int lane     = tid % warp_size;
    const int sequence = blockIdx.y;
    const int group    = blockIdx.x;

@@ -27,7 +38,7 @@ static __global__ void flash_attn_mask_to_sparse_indices(
    const half * mask = mask_ptr + sequence*s33 + q0*s31;
    int32_t * indices = indices_ptr + (int64_t(sequence)*gridDim.x + group)*n_kv_max;

-    __shared__ int warp_offsets[256/WARP_SIZE];
+    __shared__ int warp_offsets[256/warp_size];
    __shared__ int row_count;
    __shared__ int chunk_count;

@@ -37,11 +48,11 @@ static __global__ void flash_attn_mask_to_sparse_indices(
    __syncthreads();

    for (int i0 = 0; i0 < ne30; i0 += blockDim.x*values_per_lane) {
-        uint32_t selected_warp[values_per_lane];
+        lane_mask_t selected_warp[values_per_lane];
        int warp_count = 0;
 #pragma unroll
        for (int item = 0; item < values_per_lane; ++item) {
-            const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane;
+            const int i = i0 + (warp*values_per_lane + item)*warp_size + lane;
            bool selected = false;
            if (i < ne30) {
 #pragma unroll
@@ -49,8 +60,8 @@ static __global__ void flash_attn_mask_to_sparse_indices(
                    selected |= (!oob || q < q1 - q0) && isfinite(__half2float(mask[q*s31 + i]));
                }
            }
-            selected_warp[item] = __ballot_sync(0xFFFFFFFF, selected);
-            warp_count += __popc(selected_warp[item]);
+            selected_warp[item] = (lane_mask_t)__ballot_sync(ballot_mask, (int) selected);
+            warp_count += popc(selected_warp[item]);
        }

        if (lane == 0) {
@@ -61,7 +72,7 @@ static __global__ void flash_attn_mask_to_sparse_indices(
        if (tid == 0) {
            int offset = 0;
 #pragma unroll
-            for (int iw = 0; iw < 256/WARP_SIZE; ++iw) {
+            for (int iw = 0; iw < 256/warp_size; ++iw) {
                const int count = warp_offsets[iw];
                warp_offsets[iw] = offset;
                offset += count;
@@ -70,16 +81,16 @@ static __global__ void flash_attn_mask_to_sparse_indices(
        }
        __syncthreads();

-        const uint32_t lane_mask = lane == 0 ? 0 : (1u << lane) - 1;
+        const lane_mask_t lane_mask = lane == 0 ? 0 : ((lane_mask_t)1 << lane) - 1;
        int warp_item_offset = 0;
 #pragma unroll
        for (int item = 0; item < values_per_lane; ++item) {
-            const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane;
-            const int dst = row_count + warp_offsets[warp] + warp_item_offset + __popc(selected_warp[item] & lane_mask);
-            if ((selected_warp[item] & (uint32_t(1) << lane)) && dst < n_kv_max) {
+            const int i = i0 + (warp*values_per_lane + item)*warp_size + lane;
+            const int dst = row_count + warp_offsets[warp] + warp_item_offset + popc(selected_warp[item] & lane_mask);
+            if ((selected_warp[item] & ((lane_mask_t)1 << lane)) && dst < n_kv_max) {
                indices[dst] = i;
            }
-            warp_item_offset += __popc(selected_warp[item]);
+            warp_item_offset += popc(selected_warp[item]);
        }
        __syncthreads();

@@ -101,13 +112,13 @@ static __global__ void flash_attn_mask_to_sparse_indices(
    // the dependent grid reads indices, signal once the row is complete
    ggml_cuda_pdl_lc();
 }
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#endif // !defined(GGML_USE_MUSA)

 void ggml_cuda_flash_attn_ext_compact_mask(
        const ggml_tensor * mask, int32_t * indices, int32_t * counts, int32_t n_queries, int32_t ncols1, int32_t n_kv_max, cudaStream_t stream) {
-#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+#ifdef GGML_USE_MUSA
    GGML_UNUSED_VARS(mask, indices, counts, n_queries, ncols1, n_kv_max, stream);
-    GGML_ABORT("sparse flash attention is only supported on NVIDIA CUDA");
+    GGML_ABORT("sparse flash attention is not supported on MUSA");
 #else
    const int64_t s31 = mask->nb[1] / sizeof(half);
    const int64_t s33 = mask->nb[3] / sizeof(half);
@@ -122,16 +133,17 @@ void ggml_cuda_flash_attn_ext_compact_mask(
    ggml_cuda_kernel_launch(kernel, launch_params,
        (const half *) mask->data, indices, counts, int(mask->ne[0]), n_queries, n_kv_max, s31, s33);
    CUDA_CHECK(cudaGetLastError());
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#endif // !defined(GGML_USE_MUSA)
 }

 bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_tensor * dst, const int ncols1, const int ncols2) {
-#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+#ifdef GGML_USE_MUSA
    GGML_UNUSED_VARS(cc, dst, ncols1, ncols2);
    return false;
 #else
    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];

    float max_bias = 0.0f;
@@ -139,16 +151,21 @@ bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_
    memcpy(&max_bias,      (const float *) dst->op_params + 1, sizeof(float));
    memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float));

+    // On CDNA only the DKQ=576 kernel variant with ncols2=16 is compiled with sparse support.
+    if (amd_mfma_available(cc) && !ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(int(Q->ne[0]), int(V->ne[0]), 1, 16)) {
+        return false;
+    }
+
    const int32_t n_kv_max = ggml_get_op_params_i32(dst, 4);

    // the dense kernel handles up to 64/ncols2 queries per K/V pass, the single-query gather has to beat that
    const int64_t n_gather = (ncols1 == 1 ? std::min<int64_t>(Q->ne[1], 64/ncols2) : ncols1) * (int64_t) n_kv_max;

-    return GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) &&
+    return ((GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc)) || amd_mfma_available(cc)) &&
        mask != nullptr && n_kv_max > 0 && max_bias == 0.0f && logit_softcap == 0.0f &&
        mask->ne[0] == K->ne[1] && mask->ne[1] >= Q->ne[1] && mask->ne[2] == 1 &&
-        K->ne[1] >= std::max<int64_t>(4096, 2*n_gather);
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+        K->ne[1] >= std::max<int64_t>(4096, GGML_CUDA_CC_IS_NVIDIA(cc) ? 2*n_gather : 2LL*n_kv_max);
+#endif // !defined(GGML_USE_MUSA)
 }

 template <int DKQ, int DV, int ncols2>
@@ -156,7 +173,7 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con
    const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
    const ggml_tensor * Q = dst->src[0];

-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if !defined(GGML_USE_MUSA)
    if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, 1, ncols2)) {
        // a sparse variant at the full tile width gathers the union of its queries once, prefer it for large batches
        constexpr bool has_wide_sparse = ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, 64/ncols2, ncols2);
@@ -165,7 +182,7 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con
            return;
        }
    }
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#endif // !defined(GGML_USE_MUSA)

    if constexpr (ncols2 <= 8) {
        if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) {
@@ -365,7 +382,9 @@ static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, gg

            GGML_ASSERT(Q->ne[2] % K->ne[2] == 0);
            const int gqa_ratio = Q->ne[2] / K->ne[2];
-            if (gqa_ratio == 20 && GGML_CUDA_CC_IS_NVIDIA(cc)) { // GLM 4.7 Flash
+            if (amd_mfma_available(cc) && ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1, 16)) {
+                ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 16>(ctx, dst);
+            } else if (gqa_ratio == 20 && GGML_CUDA_CC_IS_NVIDIA(cc)) { // GLM 4.7 Flash
                if (cc >= GGML_CUDA_CC_DGX_SPARK) {
                    if (Q->ne[1] <= 8) {
                        ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 16>(ctx, dst);
@@ -679,6 +698,11 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
        return BEST_FATTN_KERNEL_MMA_F16;
    }

+    // On CDNA, the sparse MMA kernel is used for DKQ=576 regardless of the head size heuristics.
+    if (amd_mfma_available(cc) && ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1, 16)) {
+        return BEST_FATTN_KERNEL_MMA_F16;
+    }
+
    // AMD MFMA needs a certain minimum batch size to outscale the tile kernel for large head sizes.
    if (amd_mfma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) {
        if ((Q->ne[0] <= 64 && Q->ne[1] * gqa_ratio_eff > 8)) {
@@ -1962,7 +1962,7 @@ static bool ggml_cuda_mul_mat_q_fusion_matches(const ggml_tensor * up, const ggm
    const ggml_tensor * x_up = up->src[0];
    const ggml_tensor * y    = up->src[1];

-    if (up->op != GGML_OP_MUL_MAT || !ggml_cuda_mmq_fusion_supported(x_up->type) || y->type != GGML_TYPE_F32) {
+    if ((up->op != GGML_OP_MUL_MAT && up->op != GGML_OP_MUL_MAT_ID) || !ggml_cuda_mmq_fusion_supported(x_up->type) || y->type != GGML_TYPE_F32) {
        return false;
    }

@@ -2004,12 +2004,18 @@ static bool ggml_cuda_should_fuse_mul_mat_q(const ggml_tensor * up, const ggml_t
    const ggml_tensor * x_up = up->src[0];
    const ggml_tensor * y    = up->src[1];
    const int           cc   = ggml_cuda_info().devices[device].cc;
-    if (ggml_cuda_should_use_mmvq(x_up->type, cc, y->ne[1]) || !ggml_cuda_should_use_mmq(x_up->type, cc, y->ne[1], 0)) {
+
+    // For MUL_MAT_ID: same kernel choice as ggml_cuda_mul_mat_id, with the nominal tiles of all experts (the per-expert token counts are only known on the GPU).
+    const bool    is_id     = up->op == GGML_OP_MUL_MAT_ID;
+    const int64_t ncols     = is_id ? y->ne[2]    : y->ne[1];
+    const int64_t nchannels = is_id ? x_up->ne[2] : y->ne[2]*y->ne[3];
+    const bool    use_mmvq  = is_id ? ncols <= get_mmvq_mmid_max_batch(x_up->type, cc) : ggml_cuda_should_use_mmvq(x_up->type, cc, ncols);
+    if (use_mmvq || !ggml_cuda_should_use_mmq(x_up->type, cc, ncols, /*n_experts =*/ is_id ? x_up->ne[2] : 0)) {
        return false;
    }

    // The fused kernel needs one block per tile (no stream-k split), so only fuse if the tiling is efficient.
-    return ggml_cuda_mmq_fusion_is_efficient(x_up->type, x_up->ne[1], y->ne[1], y->ne[2]*y->ne[3], device);
+    return ggml_cuda_mmq_fusion_is_efficient(x_up->type, x_up->ne[1], ncols, nchannels, device);
 }

 static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
@@ -3822,6 +3828,11 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
            return false;
        }

+        // TODO: make the fused softcap kernel stride aware, it reads src and writes dst as dense
+        if (!ggml_is_contiguous(scale->src[0]) || !ggml_is_contiguous(scale2)) {
+            return false;
+        }
+
        if (ggml_get_unary_op(tanh) != GGML_UNARY_OP_TANH) {
            return false;
        }
@@ -5972,7 +5983,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
        case GGML_OP_LOG:
            return true;
        case GGML_OP_SCALE:
-            return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_BF16) && op->type == op->src[0]->type;
+            return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_BF16) && op->type == op->src[0]->type &&
+                   ggml_is_contiguous_rows(op->src[0]);
        case GGML_OP_CLAMP:
            return ggml_is_contiguous_rows(op->src[0]);
        case GGML_OP_ADD:
@@ -170,10 +170,13 @@ void ggml_cuda_mul_mat_q(

    // Fused gate/up/GLU: src0 is the up weight, dst the GLU output.
    const ggml_tensor * gate = fusion ? fusion->gate : nullptr;
+    ggml_cuda_mm_fusion_args_device fusion_device = {};
    if (fusion) {
-        // TODO MoE support
-        GGML_ASSERT(gate && !ids && !fusion->x_bias && !fusion->gate_bias && !fusion->x_scale && !fusion->gate_scale);
+        GGML_ASSERT(gate && !fusion->x_bias && !fusion->gate_bias && !fusion->x_scale && !fusion->gate_scale);
        GGML_ASSERT(gate->type == src0->type && ggml_are_same_shape(gate, src0) && ggml_are_same_stride(gate, src0));
+        fusion_device.gate      = gate->data;
+        fusion_device.glu_op    = fusion->glu_op;
+        fusion_device.glu_limit = fusion->glu_limit;
    }

    GGML_TENSOR_BINARY_OP_LOCALS;
@@ -286,13 +289,6 @@ void ggml_cuda_mul_mat_q(
                                ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
        const int64_t s13 = ne12*s12;

-        ggml_cuda_mm_fusion_args_device fusion_device = {};
-        if (gate) {
-            fusion_device.gate      = gate->data;
-            fusion_device.glu_op    = fusion->glu_op;
-            fusion_device.glu_limit = fusion->glu_limit;
-        }
-
        const mmq_args args = {
            src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d,
            src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr,
@@ -378,7 +374,7 @@ void ggml_cuda_mul_mat_q(
        ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
        ne02, ne02, s02, s12, s2,
        ne03, ne13, s03, s13, s3,
-        ne12, J_best};
+        ne12, J_best, fusion_device};

    ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
 }
@@ -3,23 +3,45 @@

 #define MAX_GRIDDIM_X 0x7FFFFFFF

+// src and dst may be views: rows are contiguous, dims 1..3 follow the strides (in elements).
 template <typename T>
-static __global__ void scale_f32(const T * x, T * dst, const float scale, const float bias, const int64_t nelements) {
+static __global__ void scale_f32(const T * x, T * dst, const float scale, const float bias, const uint32_t k,
+                                 const uint3 ne0, const uint3 ne1, const uint3 ne2,
+                                 const uint32_t s01, const uint32_t s02, const uint32_t s03,
+                                 const uint32_t s1,  const uint32_t s2,  const uint32_t s3) {
    ggml_cuda_pdl_lc();
-    int64_t tid = (int64_t)blockIdx.x * (int64_t)blockDim.x + (int64_t)threadIdx.x;
-    int64_t stride = (int64_t)blockDim.x * (int64_t)gridDim.x;
+    const uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;

    ggml_cuda_pdl_sync();
-    for (int64_t i = tid; i < nelements; i += stride) {
-        dst[i] = ggml_cuda_cast<T>(scale * ggml_cuda_cast<float>(x[i]) + bias);
+    if (i >= k) {
+        return;
    }
+
+    const uint2 d0 = fast_div_modulo(i,    ne0); // <i / ne0, i0>
+    const uint2 d1 = fast_div_modulo(d0.x, ne1); // <i / (ne0*ne1), i1>
+    const uint2 d2 = fast_div_modulo(d1.x, ne2); // <i3, i2>
+
+    const size_t i_src = d0.y + size_t(d1.y)*s01 + size_t(d2.y)*s02 + size_t(d2.x)*s03;
+    const size_t i_dst = d0.y + size_t(d1.y)*s1  + size_t(d2.y)*s2  + size_t(d2.x)*s3;
+
+    dst[i_dst] = ggml_cuda_cast<T>(scale * ggml_cuda_cast<float>(x[i_src]) + bias);
 }

 template <typename T>
-static void scale_f32_cuda(const T * x, T * dst, const float scale, const float bias, const int64_t nelements, cudaStream_t stream) {
-    const int64_t num_blocks = (nelements + CUDA_SCALE_BLOCK_SIZE - 1) / CUDA_SCALE_BLOCK_SIZE;
+static void scale_f32_cuda(const T * x, T * dst, const float scale, const float bias, const ggml_tensor * src0, const ggml_tensor * t, cudaStream_t stream) {
+    const int64_t k  = ggml_nelements(src0);
+    const size_t  ts = sizeof(T);
+    GGML_ASSERT(k <= std::numeric_limits<uint32_t>::max());
+
+    const uint3 ne0 = init_fastdiv_values(src0->ne[0]);
+    const uint3 ne1 = init_fastdiv_values(src0->ne[1]);
+    const uint3 ne2 = init_fastdiv_values(src0->ne[2]);
+
+    const int64_t num_blocks = (k + CUDA_SCALE_BLOCK_SIZE - 1) / CUDA_SCALE_BLOCK_SIZE;
    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(MIN(MAX_GRIDDIM_X, num_blocks), CUDA_SCALE_BLOCK_SIZE, 0, stream);
-    ggml_cuda_kernel_launch(scale_f32<T>, launch_params, x, dst, scale, bias, nelements);
+    ggml_cuda_kernel_launch(scale_f32<T>, launch_params, x, dst, scale, bias, (uint32_t) k, ne0, ne1, ne2,
+        src0->nb[1]/ts, src0->nb[2]/ts, src0->nb[3]/ts,
+        t->nb[1]/ts,    t->nb[2]/ts,    t->nb[3]/ts);
 }

 void ggml_cuda_op_scale(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
@@ -28,6 +50,7 @@ void ggml_cuda_op_scale(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {

    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_BF16);
    GGML_ASSERT(dst->type == src0->type);
+    GGML_ASSERT(ggml_is_contiguous_rows(src0) && ggml_is_contiguous_rows(dst));

    float scale;
    float bias;
@@ -35,8 +58,8 @@ void ggml_cuda_op_scale(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
    memcpy(&bias,  (float *) dst->op_params + 1, sizeof(float));

    if (src0->type == GGML_TYPE_BF16) {
-        scale_f32_cuda((const nv_bfloat16 *) src0->data, (nv_bfloat16 *) dst->data, scale, bias, ggml_nelements(src0), stream);
+        scale_f32_cuda((const nv_bfloat16 *) src0->data, (nv_bfloat16 *) dst->data, scale, bias, src0, dst, stream);
    } else {
-        scale_f32_cuda((const float *) src0->data, (float *) dst->data, scale, bias, ggml_nelements(src0), stream);
+        scale_f32_cuda((const float *) src0->data, (float *) dst->data, scale, bias, src0, dst, stream);
    }
 }
@@ -402,6 +402,10 @@ static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx,
            ggml_cuda_kernel_launch(topk_moe_cuda<288, has_bias>, launch_params,
                logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config);
            break;
+        case 384: // DeepSeek V4.1
+            ggml_cuda_kernel_launch(topk_moe_cuda<384, has_bias>, launch_params,
+                logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config);
+            break;
        case 512:
            ggml_cuda_kernel_launch(topk_moe_cuda<512, has_bias>, launch_params,
                logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config);
@@ -472,7 +476,7 @@ bool ggml_cuda_should_use_topk_moe(const ggml_tensor * gating_op,
    // must match an instantiation of launch_topk_moe_cuda: a power of 2 up to 512,
    // or one of the non-power-of-2 expert counts of supported models
    const int n_expert = ids->nb[1] / ids->nb[0];
-    if (((n_expert & (n_expert - 1)) != 0 || n_expert > 512) && n_expert != 288 && n_expert != 576) {
+    if (((n_expert & (n_expert - 1)) != 0 || n_expert > 512) && n_expert != 288 && n_expert != 384 && n_expert != 576) {
        return false;
    }

@@ -254,6 +254,10 @@ typedef __hip_fp8_e4m3 __nv_fp8_e4m3;
 #define FP8_AVAILABLE
 #endif // HIP_VERSION >= 60300000

+#if HIP_VERSION < 70000000
+#define __ballot_sync(mask, var) __ballot(var)
+#endif //HIP_VERSION < 70000000
+
 typedef int8_t int8x4_t __attribute__((ext_vector_type(4)));
 typedef uint8_t uint8x4_t __attribute__((ext_vector_type(4)));
 static __device__ __forceinline__ int __vsubss4(const int a, const int b) {
@@ -85,6 +85,7 @@ int    opt_arch       = 0; // autodetect
 size_t opt_ndev       = 1;
 size_t opt_nhvx       = 0; // use all
 int    opt_nhmx       = 1; // when set, enable HMX; when 0, use HVX only
+size_t opt_vtcm_size  = 0; // VTCM size override in bytes; 0 = auto-detect
 size_t opt_vmem       = HTP_OP_MAX_VMEM_DEFAULT;  // max available va space for buffer mappings
 int    opt_etm        = 0;
 int    opt_verbose    = 0;
@@ -2152,6 +2153,12 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
            this->n_hvx     = this->n_threads;
            this->n_hmx     = (opt_nhmx != 0) ? (uint32_t)hw_n_hmx : 0;
            this->vtcm_size = (uint64_t)hw_vtcm_size;
+            if (opt_vtcm_size > 0) {
+                if (opt_vtcm_size > hw_vtcm_size) {
+                    GGML_LOG_WARN("ggml-hex: %s VTCM override %zu exceeds HW size %llu\n", this->c_name(), opt_vtcm_size, hw_vtcm_size);
+                }
+                this->vtcm_size = opt_vtcm_size;
+            }
            GGML_LOG_INFO("ggml-hex: %s hwinfo: threads %u, hvx %u, hmx %u, vtcm %llu MB\n",
                          this->c_name(), this->n_threads, this->n_hvx, this->n_hmx,
                          (unsigned long long)(this->vtcm_size / (1024 * 1024)));
@@ -2161,7 +2168,7 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
            this->n_threads = opt_nhvx > 0 ? (uint32_t) (std::min)(opt_nhvx, (size_t) HTP_MAX_NTHREADS) : default_n_threads;
            this->n_hvx     = this->n_threads;
            this->n_hmx     = (opt_nhmx != 0) ? 1 : 0;
-            this->vtcm_size = 8 * 1024 * 1024;
+            this->vtcm_size = opt_vtcm_size > 0 ? opt_vtcm_size : 8 * 1024 * 1024;
        }
    }

@@ -2230,7 +2237,8 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
    this->op_batch = new ggml_hexagon_opbatch(this, opt_opbatch, this->max_vmem);

    // Start dspqueue/opbatch processing
-    err = htp_iface_start(this->handle, this->session_id, this->queue_id, this->n_threads, opt_nhmx, this->max_vmem, (uint32_t) opt_power_mode);
+    err = htp_iface_start(this->handle, this->session_id, this->queue_id, this->n_threads, opt_nhmx, opt_vtcm_size, this->max_vmem,
+                          (uint32_t) opt_power_mode);
    if (err != 0) {
        GGML_LOG_ERROR("ggml-hex: %s failed to start session: 0x%08x\n", this->c_name(), (unsigned) err);
        throw std::runtime_error("ggml-hex: iface start failed (see log for details)");
@@ -3391,6 +3399,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
    const char * str_etm      = getenv("GGML_HEXAGON_ETM");
    const char * str_nhvx     = getenv("GGML_HEXAGON_NHVX");
    const char * str_nhmx     = getenv("GGML_HEXAGON_NHMX");
+    const char * str_vtcm_size = getenv("GGML_HEXAGON_VTCM_SIZE");
    const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT");
    const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT");
    const char * str_fa_split  = getenv("GGML_HEXAGON_FA_SPLIT");
@@ -3456,6 +3465,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
    }
    opt_nhvx      = str_nhvx     ? strtoul(str_nhvx, NULL, 0)             : opt_nhvx;
    opt_nhmx      = str_nhmx     ? atoi(str_nhmx)                         : opt_nhmx;
+    opt_vtcm_size = str_vtcm_size ? parse_size(str_vtcm_size)                 : opt_vtcm_size;
    opt_mm_select = str_mm_select ? atoi(str_mm_select)                   : opt_mm_select;
    opt_fa_select = str_fa_select ? atoi(str_fa_select)                   : opt_fa_select;
    opt_fa_split  = str_fa_split  ? atoi(str_fa_split)                   : opt_fa_split;
@@ -57,7 +57,10 @@ htp_op_code htp_op_remap(const struct ggml_tensor * t) {
        case GGML_OP_PAD:             return HTP_OP_PAD;
        case GGML_OP_IM2COL:          return HTP_OP_IM2COL;
        case GGML_OP_IM2COL_3D:       return HTP_OP_IM2COL_3D;
+        case GGML_OP_CONV_2D:         return HTP_OP_CONV_2D;
+        case GGML_OP_CONV_2D_DW:      return HTP_OP_CONV_2D_DW;
        case GGML_OP_ROLL:            return HTP_OP_ROLL;
+        case GGML_OP_UPSCALE:         return HTP_OP_UPSCALE;

        case GGML_OP_UNARY:
            switch (ggml_get_unary_op(t)) {
@@ -1790,11 +1793,6 @@ static bool htp_kparams_solve_softmax(
        return false;
    }

-    const int64_t ne0 = src0->ne[0];
-    if ((ne0 > 32 && (ne0 & (32 - 1)) != 0) || ne0 > 131072) {
-        return false;
-    }
-
    const uint32_t src0_nrows          = src0->ne[1] * src0->ne[2] * src0->ne[3];
    const size_t   dst_data_row_size   = op->ne[0] * ggml_type_size(op->type);
    const uint32_t min_rows_per_thread = (uint32_t) (std::max<size_t>(1, (128 + dst_data_row_size - 1) / dst_data_row_size));
@@ -2188,7 +2186,8 @@ static bool htp_kparams_solve_ssm_conv(
        target_max_tile          = (std::max)(target_max_tile, 32u);
        target_max_tile          = (std::min)(target_max_tile, 128u);

-        uint32_t d_inner_tile = (uint32_t) (avail_for_tiles / (2 * (ncs + n_t) * sizeof(float)));
+        const size_t raw_stride_bytes = ncs * sizeof(float) + sizeof(float);
+        uint32_t d_inner_tile = (uint32_t) (avail_for_tiles / (2 * (raw_stride_bytes + n_t * sizeof(float))));
        d_inner_tile          = (d_inner_tile / 32) * 32;
        if (d_inner_tile == 0) {
            d_inner_tile = 32;
@@ -2375,6 +2374,183 @@ static bool htp_kparams_solve_pool_2d(
    return true;
 }

+bool htp_kparams_solve_upscale(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * dst,
+    struct htp_upscale_kernel_params * kparams
+) {
+    memset(kparams, 0, sizeof(*kparams));
+
+    if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    const uint32_t src_x = (uint32_t) src0->ne[0];
+    const uint32_t src_y = (uint32_t) src0->ne[1];
+    const uint32_t src_z = (uint32_t) src0->ne[2];
+    const uint32_t src_w = (uint32_t) src0->ne[3];
+
+    const uint32_t dst_x = (uint32_t) dst->ne[0];
+    const uint32_t dst_y = (uint32_t) dst->ne[1];
+    const uint32_t dst_z = (uint32_t) dst->ne[2];
+    const uint32_t dst_w = (uint32_t) dst->ne[3];
+
+    if (src_x == 0 || src_y == 0 || src_z == 0 || src_w == 0 ||
+        dst_x == 0 || dst_y == 0 || dst_z == 0 || dst_w == 0) {
+        return false;
+    }
+
+    const bool is_src_contig_row = (src0->nb[0] == sizeof(float));
+    const bool is_src_strided_2d = (src0->nb[0] > sizeof(float)) &&
+                                   (src0->nb[0] % sizeof(float) == 0) &&
+                                   (src0->nb[1] == (size_t) src_x * src0->nb[0]) &&
+                                   ((size_t) src_x * src_y <= 0xFFFFu);
+
+    if ((!is_src_contig_row && !is_src_strided_2d) || dst->nb[0] != sizeof(float)) {
+        return false;
+    }
+
+    const int32_t mode_flags = dst->op_params[0];
+    const uint32_t mode_raw  = (uint32_t) (mode_flags & 0xFF);
+    if (mode_raw >= GGML_SCALE_MODE_COUNT) {
+        return false;
+    }
+
+    const bool align_corners = (mode_flags & GGML_SCALE_FLAG_ALIGN_CORNERS) != 0;
+    const bool antialias     = (mode_flags & GGML_SCALE_FLAG_ANTIALIAS) != 0;
+
+    float sf0 = (float) dst_x / (float) src_x;
+    float sf1 = (float) dst_y / (float) src_y;
+    float sf2 = (float) dst_z / (float) src_z;
+    float sf3 = (float) dst_w / (float) src_w;
+    float pixel_offset = 0.5f;
+
+    if (align_corners) {
+        pixel_offset = 0.0f;
+        sf0 = (dst_x > 1 && src_x > 1) ? (float) (dst_x - 1) / (float) (src_x - 1) : sf0;
+        sf1 = (dst_y > 1 && src_y > 1) ? (float) (dst_y - 1) / (float) (src_y - 1) : sf1;
+    }
+
+    const float inv_sf0 = 1.0f / sf0;
+    const float inv_sf1 = 1.0f / sf1;
+    const float inv_sf2 = 1.0f / sf2;
+    const float inv_sf3 = 1.0f / sf3;
+
+    const float filter_radius0 = antialias ? (std::max)(1.0f, inv_sf0) : 0.0f;
+    const float filter_radius1 = antialias ? (std::max)(1.0f, inv_sf1) : 0.0f;
+    const float invscale0      = antialias ? (1.0f / filter_radius0) : 0.0f;
+    const float invscale1      = antialias ? (1.0f / filter_radius1) : 0.0f;
+
+    const uint32_t total_slices = dst_z * dst_w;
+    uint32_t       n_threads    = htp_session_n_threads(sess);
+
+    struct htp_upscale_vtcm_layout layout;
+    uint32_t dst_nrows_per_chunk = 0;
+    while (n_threads > 0) {
+        if (htp_upscale_solve_layout(&layout, src_x, src_y, dst_x, dst_y, inv_sf1, filter_radius1, n_threads,
+                                     htp_session_vtcm_size(sess), &dst_nrows_per_chunk)) {
+            break;
+        }
+        n_threads--;
+    }
+    if (n_threads == 0) {
+        return false;
+    }
+
+    size_t lut_x_size = 0;
+    size_t elem_y_sz  = sizeof(uint32_t);
+    if (mode_raw == HTP_UPSCALE_MODE_BICUBIC) {
+        lut_x_size = hex_round_up(dst_x * sizeof(struct upscale_bicubic_x), 128);
+        elem_y_sz  = sizeof(struct upscale_bicubic_y);
+    } else if (antialias) {
+        lut_x_size = hex_round_up(dst_x * sizeof(struct upscale_antialias_x), 128);
+        elem_y_sz  = sizeof(struct upscale_antialias_y);
+    } else if (mode_raw == HTP_UPSCALE_MODE_BILINEAR) {
+        const uint32_t n_chunks = (dst_x + 31) / 32;
+        lut_x_size = hex_round_up(n_chunks * sizeof(struct upscale_bilinear_x_chunk), 128);
+        elem_y_sz  = sizeof(struct upscale_bilinear_y);
+    } else {
+        lut_x_size = hex_round_up(dst_x * sizeof(uint32_t), 128);
+    }
+    const size_t total_lut_size = lut_x_size + hex_round_up(dst_y * elem_y_sz, 128);
+    if (total_lut_size > HTP_DDR_SPAD_SIZE) {
+        return false;
+    }
+
+    const uint32_t chunks_per_slice = (dst_y + dst_nrows_per_chunk - 1) / dst_nrows_per_chunk;
+    const uint32_t total_chunks     = total_slices * chunks_per_slice;
+
+    n_threads = (std::min)(n_threads, total_chunks > 0 ? total_chunks : 1);
+    htp_upscale_vtcm_layout_build(&layout, src_x, src_y, dst_x, dst_y, dst_nrows_per_chunk, inv_sf1, filter_radius1, n_threads);
+
+    const uint32_t slice_dst_bytes = dst_x * dst_y * sizeof(float);
+    const uint32_t chunk_dst_bytes = dst_x * dst_nrows_per_chunk * sizeof(float);
+    const bool     dst_contig      = (dst->nb[1] == (size_t) dst_x * sizeof(float)) &&
+                                     (dst_z <= 1 || dst->nb[2] == (size_t) dst_y * dst->nb[1]) &&
+                                     (dst_w <= 1 || dst->nb[3] == (size_t) dst_z * dst->nb[2]);
+
+    uint32_t chunks_per_group = 0;
+    if (dst_contig) {
+        const uint32_t slices_per_group = (slice_dst_bytes > 0)
+            ? (HEX_L2_LINE_SIZE / hex_gcd_u32(slice_dst_bytes, HEX_L2_LINE_SIZE))
+            : 1;
+        if (slices_per_group == 1 && (chunk_dst_bytes % HEX_L2_LINE_SIZE == 0)) {
+            chunks_per_group = 1;
+        } else {
+            chunks_per_group = slices_per_group * chunks_per_slice;
+        }
+    }
+
+    uint8_t kernel_type = HTP_UPSCALE_KERNEL_BILINEAR;
+    if (mode_raw == HTP_UPSCALE_MODE_NEAREST) {
+        if (dst_x == 2 * src_x && dst_y == 2 * src_y && !align_corners) {
+            kernel_type = HTP_UPSCALE_KERNEL_NEAREST_2X;
+        } else if (dst_x == 4 * src_x && dst_y == 4 * src_y && !align_corners) {
+            kernel_type = HTP_UPSCALE_KERNEL_NEAREST_4X;
+        } else {
+            kernel_type = HTP_UPSCALE_KERNEL_NEAREST;
+        }
+    } else if (mode_raw == HTP_UPSCALE_MODE_BICUBIC) {
+        kernel_type = HTP_UPSCALE_KERNEL_BICUBIC;
+    } else if (antialias) {
+        kernel_type = HTP_UPSCALE_KERNEL_BILINEAR_AA;
+    } else if (dst_x == 2 * src_x && dst_y == 2 * src_y && !align_corners && (src_x % 32 == 0)) {
+        kernel_type = HTP_UPSCALE_KERNEL_BILINEAR_2X;
+    } else {
+        kernel_type = HTP_UPSCALE_KERNEL_BILINEAR;
+    }
+
+    kparams->kernel_type          = kernel_type;
+    kparams->mode                 = (uint8_t) mode_raw;
+    kparams->align_corners        = align_corners ? 1 : 0;
+    kparams->antialias            = antialias ? 1 : 0;
+
+    kparams->inv_sf0              = inv_sf0;
+    kparams->inv_sf1              = inv_sf1;
+    kparams->inv_sf2              = inv_sf2;
+    kparams->inv_sf3              = inv_sf3;
+    kparams->pixel_offset         = pixel_offset;
+
+    kparams->filter_radius0       = filter_radius0;
+    kparams->filter_radius1       = filter_radius1;
+    kparams->invscale0            = invscale0;
+    kparams->invscale1            = invscale1;
+
+    kparams->n_threads            = n_threads;
+    kparams->chunks_per_group     = chunks_per_group;
+
+    kparams->dst_nrows_per_chunk  = dst_nrows_per_chunk;
+    kparams->chunks_per_slice     = chunks_per_slice;
+    kparams->total_chunks         = total_chunks;
+    kparams->vtcm_size            = (uint32_t) layout.total_bytes;
+
+    kparams->div_chunks_per_slice = init_fastdiv_values(chunks_per_slice > 0 ? chunks_per_slice : 1);
+    kparams->div_dst_z            = init_fastdiv_values(dst_z > 0 ? dst_z : 1);
+
+    return true;
+}
+
 static bool htp_kparams_solve_concat(
    const struct ggml_hexagon_session * sess,
    const struct ggml_tensor * op,
@@ -2729,7 +2905,8 @@ static bool htp_kparams_solve_im2col(
            const uint32_t dst_row_bytes = hex_round_up((uint32_t) dst_row_bytes_raw, 256);

            struct htp_im2col_vtcm_layout layout;
-            htp_im2col_vtcm_layout_build(&layout, src_row_bytes, dst_row_bytes, n_threads);
+            htp_im2col_vtcm_layout_build(&layout, src_row_bytes, dst_row_bytes, n_threads,
+                                         htp_im2col_scratch_bytes(HTP_IM2COL_KERNEL_PATCHEMBED_DMA, KH, patch_stride));

            if (vtcm_size == 0 || layout.total_bytes <= vtcm_size) {
                kparams->kernel_type          = HTP_IM2COL_KERNEL_PATCHEMBED_DMA;
@@ -2778,18 +2955,24 @@ static bool htp_kparams_solve_im2col(
        const uint32_t src_row_bytes = hex_round_up((uint32_t) src_row_bytes_raw, 256);
        const uint32_t dst_row_bytes = hex_round_up((uint32_t) dst_row_bytes_raw, 256);

+        // threads work on output-col blocks, so a single output row (1D) can use them all
+        const uint32_t nbpr        = (OW + owb - 1) / owb;
+        const uint64_t nblocks     = (uint64_t) total_rows * nbpr;
+        const uint32_t blk_threads = (uint32_t) (nblocks < max_threads ? nblocks : max_threads);
+
        struct htp_im2col_vtcm_layout layout;
-        htp_im2col_vtcm_layout_build(&layout, src_row_bytes, dst_row_bytes, n_threads);
+        htp_im2col_vtcm_layout_build(&layout, src_row_bytes, dst_row_bytes, blk_threads,
+                                     htp_im2col_scratch_bytes(HTP_IM2COL_KERNEL_BLOCKED_DMA, KH, patch_stride));

        if (vtcm_size == 0 || layout.total_bytes <= vtcm_size) {
            kparams->kernel_type          = HTP_IM2COL_KERNEL_BLOCKED_DMA;
-            kparams->n_threads            = (uint8_t) n_threads;
+            kparams->n_threads            = (uint8_t) blk_threads;
            kparams->vtcm_size            = (uint32_t) layout.total_bytes;
            kparams->src_row_bytes        = src_row_bytes;
            kparams->dst_row_bytes        = dst_row_bytes;
            kparams->owb                  = owb;
            kparams->wb                   = (uint32_t) Wb;
-            kparams->nbpr                 = (OW + owb - 1) / owb;
+            kparams->nbpr                 = nbpr;
            kparams->N                    = N;
            kparams->IC                   = IC;
            kparams->ID                   = 1;
@@ -2806,6 +2989,7 @@ static bool htp_kparams_solve_im2col(
            kparams->rows_per_chunk       = rows_per_chunk;
            kparams->div_OH               = init_fastdiv_values(OH);
            kparams->div_OD               = init_fastdiv_values(1);
+            kparams->div_nbpr             = init_fastdiv_values(nbpr);
            return true;
        }
        if (owb == 1) break;
@@ -2907,7 +3091,7 @@ static bool htp_kparams_solve_im2col_3d(
            const uint32_t dst_row_bytes = hex_round_up((uint32_t) dst_row_bytes_raw, 256);

            struct htp_im2col_vtcm_layout layout;
-            htp_im2col_vtcm_layout_build(&layout, src_row_bytes, dst_row_bytes, n_threads);
+            htp_im2col_vtcm_layout_build(&layout, src_row_bytes, dst_row_bytes, n_threads, 256);

            if (vtcm_size == 0 || layout.total_bytes <= vtcm_size) {
                kparams->kernel_type          = HTP_IM2COL_KERNEL_PATCHEMBED_DMA;
@@ -2957,7 +3141,7 @@ static bool htp_kparams_solve_im2col_3d(
        const uint32_t dst_row_bytes = hex_round_up((uint32_t) dst_row_bytes_raw, 256);

        struct htp_im2col_vtcm_layout layout;
-        htp_im2col_vtcm_layout_build(&layout, src_row_bytes, dst_row_bytes, n_threads);
+        htp_im2col_vtcm_layout_build(&layout, src_row_bytes, dst_row_bytes, n_threads, 256);

        if (vtcm_size == 0 || layout.total_bytes <= vtcm_size) {
            kparams->kernel_type          = HTP_IM2COL_KERNEL_BLOCKED_DMA;
@@ -3286,6 +3470,545 @@ static bool htp_kparams_solve_solve_tri(
    return true;
 }

+// Implicit GEMM over sliding windows of fp16 input croutons (HMX activation:single), no im2col
+static bool htp_kparams_conv2d_win(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * dst,
+    struct htp_conv2d_kernel_params * kp
+) {
+    if (htp_session_n_hmx(sess) == 0) {
+        return false;
+    }
+
+    const struct ggml_tensor * src0 = dst->src[0];
+    const struct ggml_tensor * src1 = dst->src[1];
+    const int32_t s0 = ggml_get_op_params_i32(dst, 0);
+    const int32_t s1 = ggml_get_op_params_i32(dst, 1);
+    const int32_t p0 = ggml_get_op_params_i32(dst, 2);
+    const int32_t p1 = ggml_get_op_params_i32(dst, 3);
+    const int32_t d0 = ggml_get_op_params_i32(dst, 4);
+    const int32_t d1 = ggml_get_op_params_i32(dst, 5);
+
+    const uint32_t KW   = src0->ne[0];
+    const uint32_t KH   = src0->ne[1];
+    const uint32_t Cin  = src0->ne[2];
+    const uint32_t Cout = src0->ne[3];
+    const uint32_t IW   = src1->ne[0];
+    const uint32_t IH   = src1->ne[1];
+    const uint32_t OW   = dst->ne[0];
+    const uint32_t OH   = dst->ne[1];
+    const uint32_t N    = dst->ne[3];
+    const size_t   es   = ggml_type_size(src0->type);
+
+    if (s0 != 1 || s1 != 1 || p0 < 0 || p1 < 0 || dst->nb[0] != sizeof(float)) {
+        return false;
+    }
+    if (dst->nb[1] < OW * sizeof(float) || (Cout > 1 && dst->nb[2] < OW * sizeof(float))) {
+        return false;
+    }
+
+    const bool whcn = src1->nb[0] == sizeof(float) && src1->nb[1] >= IW * sizeof(float) && (Cin == 1 || src1->nb[2] >= IH * src1->nb[1]);
+    const bool cwhn = !whcn && src1->nb[2] == sizeof(float) && src1->nb[0] == Cin * sizeof(float) && src1->nb[1] >= IW * src1->nb[0];
+    if (!whcn && !cwhn) {
+        return false;
+    }
+
+    const bool w_rows = src0->nb[0] == es && (KH == 1 || src0->nb[1] == KW * es) && (Cin == 1 || src0->nb[2] == KW * KH * es);
+    const bool co_c   = (Cout == 1 || src0->nb[3] == es) && (Cin == 1 || src0->nb[2] >= (std::min)(Cout, 32u) * es);
+    if (!w_rows && !co_c) {
+        return false;
+    }
+
+    // short output rows: fold the kw taps into channels so the virtual row is OW wide and output rows are contiguous
+    const bool fold = (whcn || (cwhn && d0 == 1)) && KW > 1 && KW * Cin <= 32 && OW < 64 && KW * KH <= 16 && Cout >= 16;
+
+    const uint32_t T     = KW * KH;
+    const uint32_t n_ch  = fold ? KW * Cin : Cin;
+    const uint32_t n_cc  = (n_ch + 31) / 32;
+    const uint32_t rem   = n_ch - (n_cc - 1) * 32;
+    const uint32_t stop  = (rem + 7) / 8 * 8 - 1;
+    const uint32_t SR    = fold ? OW : IW + 2 * p0;
+    const uint32_t IHS   = IH + 2 * p1;
+    const uint32_t n_wt  = fold ? KH : T;
+    const uint32_t moff  = (KH - 1) * d1 * SR + (fold ? 0 : (KW - 1) * d0);
+    const uint32_t halo  = moff / 32 + 1;
+    const uint32_t n_rt  = (Cout + 31) / 32;
+    const uint64_t n_q   = (uint64_t) (N * IHS - (IHS - OH) - 1) * SR + OW;
+    const uint32_t nth   = kp->n_threads;
+    const size_t   tile  = 2048;
+
+    if ((uint64_t) N * IHS * SR + moff + 64 >= (1ull << 31)) {
+        return false;
+    }
+    const uint32_t n_tiles = (uint32_t) ((n_q + 31) / 32);
+
+    // weight packing scratch
+    uint32_t w_cs  = 0;
+    size_t   wraw  = 0;
+    size_t   wscr  = 0;
+    if (w_rows) {
+        const uint32_t kpw = es == 2 ? 64 : 32;
+        for (uint32_t cs = 32; cs >= 2; cs /= 2) {
+            w_cs = cs;
+            wraw = (size_t) 32 * ((cs * T + kpw - 1) / kpw) * 128;
+            if (3 * wraw + 64 * 128 <= 256 * 1024 || (fold && cs == 32)) {
+                break;
+            }
+        }
+        // enough packing units for all threads
+        const uint32_t kpw2 = es == 2 ? 64 : 32;
+        while (!fold && w_cs > 2 && n_rt * ((Cin + 31) / 32) * (32 / w_cs) < kp->n_threads) {
+            w_cs /= 2;
+        }
+        wraw = (std::min)(wraw, (size_t) 32 * ((w_cs * T + kpw2 - 1) / kpw2) * 128);
+        wscr = 3 * wraw + 64 * 128 + 128;
+    } else {
+        w_cs = fold ? KW : (std::min)(T, 16u);
+        while (!fold && w_cs > 1 && n_rt * ((Cin + 31) / 32) * ((T + w_cs - 1) / w_cs) < kp->n_threads) {
+            w_cs = (w_cs + 1) / 2;
+        }
+        wscr = 2 * (size_t) w_cs * 32 * 128;
+    }
+
+    uint32_t rt_block = n_rt;
+    uint32_t ct       = (std::max)(4 * halo, (n_tiles + 3) / 4);
+    ct = (std::min)((std::max)(ct, 8u), 512u);
+    ct = (std::min)(ct, n_tiles);
+    uint32_t seg_b = 8;
+    // longer output rows make the output DMA more efficient
+    uint32_t seg_t = (std::min)(32u, (std::max)(8u, (OW + 31) / 32));
+
+    while (rt_block > 1 && (size_t) rt_block * n_wt * n_cc * tile > htp_session_vtcm_size(sess) / 2) {
+        rt_block = (rt_block + 1) / 2;
+    }
+
+    size_t stage = 0, thr = 0, act = 0, out = 0, w_size = 0, total = 0;
+    uint32_t n_bc = 0;
+    while (true) {
+        n_bc   = ct + halo;
+        stage  = cwhn ? (size_t) seg_b * 32 * Cin * sizeof(float) + 256 : (size_t) n_ch * seg_b * 32 * sizeof(float);
+        if (fold) {
+            // raw input rows of the staged pixels with 128 bytes of read slack on both sides
+            stage = (size_t) Cin * ((seg_b * 32 + OW - 1) / OW + 1) * IW * sizeof(float) + 256;
+        }
+        stage  = hex_round_up(stage, 128);
+        thr    = hex_round_up((std::max)(2 * stage + (size_t) HTP_CONV2D_WIN_OBUF * 32 * seg_t * 32 * sizeof(float), wscr), tile);
+        w_size = (size_t) rt_block * n_wt * n_cc * tile;
+        act    = (size_t) n_cc * n_bc * tile;
+        out    = (size_t) rt_block * ct * tile;
+        struct htp_conv2d_vtcm_layout L;
+        htp_conv2d_vtcm_layout_build(&L, rt_block, w_size, act, out, thr, nth);
+        total  = L.total_bytes;
+        if (total <= htp_session_vtcm_size(sess)) {
+            break;
+        }
+        if (ct > 64) {
+            ct /= 2;
+        } else if (rt_block > 1) {
+            rt_block = (rt_block + 1) / 2;
+        } else if (seg_b > 1 || seg_t > 1) {
+            seg_b = (std::max)(seg_b / 2, 1u);
+            seg_t = (std::max)(seg_t / 2, 1u);
+        } else if (ct > 1) {
+            ct /= 2;
+        } else {
+            return false;
+        }
+    }
+
+    if (hex_round_up((size_t) T * sizeof(uint32_t), 128) + 2 * hex_round_up((size_t) ct, 128) > HTP_DDR_SPAD_SIZE) {
+        return false;
+    }
+
+    const uint32_t TG   = w_cs;
+    const uint32_t CS   = w_cs;
+    const uint32_t n_tg = (T + TG - 1) / TG;
+    const uint32_t n_cs = (32 + CS - 1) / CS;
+
+    kp->kernel_type        = (uint8_t) HTP_CONV2D_KERNEL_HMX_WINDOW;
+    kp->n_threads          = (uint8_t) nth;
+    kp->cwhn               = cwhn ? 1 : 0;
+    kp->w_kmajor           = !w_rows ? 1 : 0;
+    kp->n_row_tiles        = n_rt;
+    kp->ct_chunk           = ct;
+    kp->rt_block           = rt_block;
+    kp->stage_size         = stage;
+    kp->thr_size           = thr;
+    kp->vtcm_size          = total;
+    kp->window.fold        = fold;
+    kp->window.n_cc        = n_cc;
+    kp->window.act_rt_last = 0x780 | 0x2 | (stop << 2);
+    kp->window.sr          = SR;
+    kp->window.ihs         = IHS;
+    kp->window.max_off     = moff;
+    kp->window.n_bc        = n_bc;
+    kp->window.seg_b       = seg_b;
+    kp->window.seg_t       = seg_t;
+    kp->window.w_cs        = w_cs;
+    kp->window.div_sr      = init_fastdiv_values(SR);
+    kp->window.div_ihs     = init_fastdiv_values(IHS);
+    kp->window.div_kw      = init_fastdiv_values(KW);
+    kp->window.div_tg      = init_fastdiv_values(n_tg);
+    kp->window.div_cc      = init_fastdiv_values(n_cc);
+    kp->window.div_cs      = init_fastdiv_values(n_cs);
+    return true;
+}
+
+// Tiled implicit GEMM, MMA on HMX if available, else HVX
+static bool htp_kparams_conv2d_tiled(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * dst,
+    struct htp_conv2d_kernel_params * kp
+) {
+    if (htp_session_n_hmx(sess) == 0) {
+        return false;
+    }
+
+    const struct ggml_tensor * src0 = dst->src[0];
+    const struct ggml_tensor * src1 = dst->src[1];
+    const int32_t s0 = ggml_get_op_params_i32(dst, 0);
+    const int32_t d0 = ggml_get_op_params_i32(dst, 4);
+
+    const uint32_t KW   = src0->ne[0];
+    const uint32_t KH   = src0->ne[1];
+    const uint32_t Cin  = src0->ne[2];
+    const uint32_t Cout = src0->ne[3];
+    const uint32_t OW   = dst->ne[0];
+    const uint32_t OH   = dst->ne[1];
+    const uint32_t N    = dst->ne[3];
+    const size_t   es   = ggml_type_size(src0->type);
+
+    if (s0 != 1 || dst->nb[0] != sizeof(float)) {
+        return false;
+    }
+
+    // kernel k order must match the staged input: (cin, kh, kw) for WHCN input, (kh, kw, cin) for CWHN input.
+    // Kernel is either k-contiguous rows per cout (WHCN only) or cout-contiguous with a uniform k stride.
+    const bool   whcn   = src1->nb[0] == sizeof(float);
+    const bool   cwhn   = !whcn && src1->nb[2] == sizeof(float) && src1->nb[0] == Cin * sizeof(float);
+    const bool   co_c   = Cout == 1 || src0->nb[3] == es;
+    const bool   w_rows = whcn && src0->nb[0] == es && (KH == 1 || src0->nb[1] == KW * es) && (Cin == 1 || src0->nb[2] == KW * KH * es);
+    const size_t ks     = whcn ? src0->nb[0] : src0->nb[2];
+    bool         kmajor = false;
+    if (whcn) {
+        kmajor = co_c && ks >= Cout * es && (KH == 1 || src0->nb[1] == KW * ks) && (Cin == 1 || src0->nb[2] == KW * KH * ks);
+    } else if (cwhn) {
+        kmajor = co_c && ks >= Cout * es && src0->nb[0] == Cin * ks && (KH == 1 || src0->nb[1] == KW * Cin * ks);
+    }
+    if (!w_rows && !kmajor) {
+        return false;
+    }
+    kmajor = kmajor && !w_rows;
+
+    const uint32_t K     = KW * KH * Cin;
+    const uint32_t kg    = !cwhn ? K : d0 == 1 ? KW * Cin : Cin;
+    const uint32_t n_grp = K / kg;
+    const uint32_t kg_p  = hex_round_up(kg, 32);
+    const uint32_t k_pad = hex_round_up(n_grp * kg_p, 64);
+    const uint32_t n_kt  = k_pad / 32;
+    const uint32_t n_rt  = (Cout + 31) / 32;
+    const uint32_t n_owt = (OW + 31) / 32;
+    const uint32_t n_ct  = N * OH * n_owt;
+    const uint32_t wk    = (std::min)(k_pad, 512u);
+    const uint32_t nth   = htp_session_n_threads(sess);
+    const size_t   tile  = 2048;
+
+    uint32_t rt_block = n_rt;
+    uint32_t ct       = (std::min)(n_ct, 64u);
+    uint32_t seg      = (std::min)(8u, n_owt);
+    uint32_t n_stage  = 1;
+    uint32_t sr       = 0;
+    size_t   stage = 0, ostage = 0, prow = 0, wpack = 0;
+    size_t   w_size = 0, act_size = 0, out_size = 0, thr_size = 0, total = 0;
+
+    while (true) {
+        const uint32_t wb = 32 * seg + (KW - 1) * d0;
+        sr     = cwhn ? hex_round_up(wb * Cin, 64) : hex_round_up(wb, 32);
+        stage  = cwhn ? (size_t) KH * sr * sizeof(float) : (size_t) Cin * KH * sr * sizeof(float);
+        ostage = 2 * 32 * (size_t) seg * 32 * sizeof(float);
+        prow   = cwhn ? stage / 2 + 256 + 64 * 128 : 0;
+        wpack  = kmajor ? 2 * (size_t) wk * 64 * sizeof(float) + 96 * 128 : 2 * (size_t) 32 * wk * (sizeof(float) + sizeof(uint16_t));
+        thr_size = hex_round_up((std::max)(stage + ostage + prow, wpack), tile);
+        w_size   = (size_t) rt_block * n_kt * tile;
+        act_size = (size_t) ct * n_kt * tile;
+        out_size = (size_t) rt_block * ct * tile;
+        struct htp_conv2d_vtcm_layout L;
+        htp_conv2d_vtcm_layout_build(&L, rt_block, w_size, act_size, out_size, thr_size, nth);
+        total    = L.total_bytes;
+        if (total <= htp_session_vtcm_size(sess)) {
+            break;
+        }
+        if (ct > 8) {
+            ct /= 2;
+        } else if (seg > 1) {
+            seg--;
+        } else if (rt_block > 1) {
+            rt_block = (rt_block + 1) / 2;
+        } else if (ct > 1) {
+            ct /= 2;
+        } else {
+            return false;
+        }
+    }
+
+    // Second stage buffer: DMA of the next segment overlaps HVX pack of this one.
+    // Skip large stages. They already fill VTCM, and the extra DMA slows HMX.
+    {
+        const size_t thr2 = hex_round_up((std::max)(2 * stage + ostage + prow, wpack), tile);
+        struct htp_conv2d_vtcm_layout L2;
+        htp_conv2d_vtcm_layout_build(&L2, rt_block, w_size, act_size, out_size, thr2, nth);
+        if (stage <= 128 * 1024 && L2.total_bytes <= htp_session_vtcm_size(sess)) {
+            n_stage  = 2;
+            thr_size = thr2;
+            total    = L2.total_bytes;
+        }
+    }
+
+    kp->kernel_type        = (uint8_t) HTP_CONV2D_KERNEL_HMX_TILED;
+    kp->n_threads          = (uint8_t) nth;
+    kp->cwhn               = cwhn ? 1 : 0;
+    kp->w_kmajor           = kmajor ? 1 : 0;
+    kp->n_row_tiles        = n_rt;
+    kp->ct_chunk           = ct;
+    kp->rt_block           = rt_block;
+    kp->stage_size         = stage;
+    kp->thr_size           = thr_size;
+    kp->vtcm_size          = total;
+    kp->tiled.w_kstride    = ks;
+    kp->tiled.k_size       = K;
+    kp->tiled.k_group      = kg;
+    kp->tiled.k_group_pad  = kg_p;
+    kp->tiled.n_groups     = n_grp;
+    kp->tiled.n_k_tiles    = n_kt;
+    kp->tiled.n_ow_tiles   = n_owt;
+    kp->tiled.seg_tiles    = seg;
+    kp->tiled.stage_row    = sr;
+    kp->tiled.n_stage      = n_stage;
+    kp->tiled.wk_chunk     = wk;
+    kp->tiled.div_ow_tiles = init_fastdiv_values(n_owt);
+    kp->tiled.div_oh       = init_fastdiv_values(OH);
+    return true;
+}
+
+static bool htp_kparams_solve_conv2d_unchecked(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * dst,
+    struct htp_conv2d_kernel_params * kparams
+) {
+    memset(kparams, 0, sizeof(*kparams));
+
+    const uint32_t total_rows = (uint32_t) (dst->ne[1] * dst->ne[2] * dst->ne[3]);
+    kparams->total_rows = total_rows;
+    kparams->n_threads  = total_rows == 0 ? 1 : (std::min)((uint32_t) htp_session_n_threads(sess), total_rows);
+
+    if (htp_kparams_conv2d_win(sess, dst, kparams)) {
+        return true;
+    }
+
+    return htp_kparams_conv2d_tiled(sess, dst, kparams);
+}
+
+static bool htp_kparams_solve_conv2d(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    struct htp_conv2d_kernel_params * kparams
+) {
+    memset(kparams, 0, sizeof(*kparams));
+
+    const struct ggml_tensor * src0 = op->src[0];
+    const struct ggml_tensor * src1 = op->src[1];
+    const struct ggml_tensor * dst  = op;
+
+    if ((src0->type != GGML_TYPE_F16 && src0->type != GGML_TYPE_F32) ||
+        src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    if (src0->ne[2] != src1->ne[2]) {
+        return false;
+    }
+
+    const int32_t s0 = ggml_get_op_params_i32(op, 0);
+    const int32_t s1 = ggml_get_op_params_i32(op, 1);
+    const int32_t p0 = ggml_get_op_params_i32(op, 2);
+    const int32_t p1 = ggml_get_op_params_i32(op, 3);
+    const int32_t d0 = ggml_get_op_params_i32(op, 4);
+    const int32_t d1 = ggml_get_op_params_i32(op, 5);
+
+    const int64_t ow = (src1->ne[0] + 2 * p0 - d0 * (src0->ne[0] - 1) - 1) / s0 + 1;
+    const int64_t oh = (src1->ne[1] + 2 * p1 - d1 * (src0->ne[1] - 1) - 1) / s1 + 1;
+    if (ow != dst->ne[0] || oh != dst->ne[1] || src0->ne[3] != dst->ne[2] || src1->ne[3] != dst->ne[3]) {
+        return false;
+    }
+
+    return htp_kparams_solve_conv2d_unchecked(sess, op, kparams);
+}
+
+static bool htp_kparams_solve_conv2d_dw(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    struct htp_conv2d_dw_kernel_params * kparams
+) {
+    memset(kparams, 0, sizeof(*kparams));
+
+    const struct ggml_tensor * src0 = op->src[0];
+    const struct ggml_tensor * src1 = op->src[1];
+    const struct ggml_tensor * dst  = op;
+
+    if ((src0->type != GGML_TYPE_F16 && src0->type != GGML_TYPE_F32) ||
+        src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    if (src0->ne[2] != 1 || src0->ne[3] != src1->ne[2]) {
+        return false;
+    }
+
+    if (src0->ne[0] * src0->ne[1] > 64) {
+        return false;
+    }
+
+    const int32_t s0 = ggml_get_op_params_i32(op, 0);
+    const int32_t s1 = ggml_get_op_params_i32(op, 1);
+    const int32_t p0 = ggml_get_op_params_i32(op, 2);
+    const int32_t p1 = ggml_get_op_params_i32(op, 3);
+    const int32_t d0 = ggml_get_op_params_i32(op, 4);
+    const int32_t d1 = ggml_get_op_params_i32(op, 5);
+
+    if (s0 <= 0 || s1 <= 0 || d0 <= 0 || d1 <= 0) {
+        return false;
+    }
+
+    const int64_t ow = (src1->ne[0] + 2 * p0 - d0 * (src0->ne[0] - 1) - 1) / s0 + 1;
+    const int64_t oh = (src1->ne[1] + 2 * p1 - d1 * (src0->ne[1] - 1) - 1) / s1 + 1;
+    if (ow != dst->ne[0] || oh != dst->ne[1] || dst->ne[2] != src1->ne[2] || dst->ne[3] != src1->ne[3]) {
+        return false;
+    }
+
+    const bool cwhn = ggml_is_contiguous_channels(src1);
+    if (!cwhn && !ggml_is_contiguous(src1)) {
+        return false;
+    }
+
+    if (!cwhn && (d0 != 1 || (s0 != 1 && s0 != 2))) {
+        return false;
+    }
+
+    const uint32_t c  = (uint32_t) src1->ne[2];
+    const uint32_t n  = (uint32_t) src1->ne[3];
+    const uint32_t w  = (uint32_t) src1->ne[0];
+    const uint32_t h  = (uint32_t) src1->ne[1];
+    const uint32_t kw = (uint32_t) src0->ne[0];
+    const uint32_t kh = (uint32_t) src0->ne[1];
+    const uint32_t ow_u = (uint32_t) ow;
+    const uint32_t oh_u = (uint32_t) oh;
+
+    const uint32_t row_src_bytes = cwhn
+        ? (w * c * (uint32_t) sizeof(float))
+        : (w * (uint32_t) sizeof(float));
+    const uint32_t row_dst_bytes = cwhn
+        ? (ow_u * c * (uint32_t) sizeof(float))
+        : (ow_u * (uint32_t) sizeof(float));
+
+    uint8_t kernel_type = HTP_CONV2D_DW_KERNEL_UNSUPPORTED;
+    if (cwhn) {
+        kernel_type = HTP_CONV2D_DW_KERNEL_CWHN;
+    } else if ((w & 31) == 0 && kw == 3 && kh == 3 && p0 == 1 && p1 == 1 && d0 == 1 && d1 == 1 && s0 == 1 && s1 == 1) {
+        kernel_type = HTP_CONV2D_DW_KERNEL_WHCN_S1_K3X3;
+    } else if ((w & 31) == 0 && kw == 3 && kh == 3 && p0 == 1 && p1 == 1 && d0 == 1 && d1 == 1 && s0 == 2 && s1 == 2) {
+        if (w == 32 && ow_u == 16) {
+            kernel_type = HTP_CONV2D_DW_KERNEL_WHCN_S2_K3X3_W32;
+        } else if (w == 64 && ow_u == 32) {
+            kernel_type = HTP_CONV2D_DW_KERNEL_WHCN_S2_K3X3_W64;
+        } else {
+            kernel_type = HTP_CONV2D_DW_KERNEL_WHCN_S2_K3X3;
+        }
+    } else if (s0 == 1) {
+        kernel_type = HTP_CONV2D_DW_KERNEL_WHCN_S1;
+    } else {
+        kernel_type = HTP_CONV2D_DW_KERNEL_WHCN_S2;
+    }
+
+    const bool     fast_k3x3      = kernel_type >= HTP_CONV2D_DW_KERNEL_WHCN_S1_K3X3;
+    const uint32_t pad_rows       = fast_k3x3 ? 4 : 0;
+    const uint32_t pad_top_stride = fast_k3x3 ? w : 0;
+
+    const size_t vtcm_budget = (sess && htp_session_vtcm_size(sess) > 0) ? htp_session_vtcm_size(sess) : 0;
+    if (vtcm_budget == 0) {
+        return false;
+    }
+
+    const uint32_t w_size = (uint32_t) (kw * kh * c * ggml_type_size(src0->type));
+
+    uint32_t max_threads_sess = (uint32_t) htp_session_n_threads(sess);
+    uint32_t oh_chunk = oh_u;
+
+    uint32_t n_threads = max_threads_sess;
+    struct htp_conv2d_dw_vtcm_layout layout;
+
+    while (true) {
+        const uint32_t max_ih_chunk = (std::min)(h, (oh_chunk - 1) * (uint32_t) s1 + (kh - 1) * (uint32_t) d1 + 1);
+        const uint32_t max_src_bytes = (max_ih_chunk + pad_rows) * row_src_bytes + 128;
+        const uint32_t max_dst_bytes = oh_chunk * row_dst_bytes + 128;
+
+        htp_conv2d_dw_vtcm_layout_build(&layout, max_src_bytes, max_dst_bytes, w_size, n_threads);
+
+        if (layout.total_bytes <= vtcm_budget && layout.total_bytes > 0) {
+            break;
+        }
+
+        if (oh_chunk > 1) {
+            oh_chunk = (oh_chunk + 1) / 2;
+        } else if (n_threads > 1) {
+            n_threads--;
+        } else {
+            return false;
+        }
+    }
+
+    const uint32_t n_chunks = (oh_u + oh_chunk - 1) / oh_chunk;
+    const uint32_t total_tasks = cwhn ? (n * n_chunks) : (c * n * n_chunks);
+    n_threads = (std::min)(n_threads, total_tasks);
+    if (n_threads == 0) {
+        n_threads = 1;
+    }
+
+    const uint32_t max_ih_chunk = (std::min)(h, (oh_chunk - 1) * (uint32_t) s1 + (kh - 1) * (uint32_t) d1 + 1);
+    const uint32_t max_src_bytes = (max_ih_chunk + pad_rows) * row_src_bytes + 128;
+    const uint32_t max_dst_bytes = oh_chunk * row_dst_bytes + 128;
+    htp_conv2d_dw_vtcm_layout_build(&layout, max_src_bytes, max_dst_bytes, w_size, n_threads);
+
+    kparams->kernel_type    = kernel_type;
+    kparams->n_threads      = (uint8_t) n_threads;
+    kparams->cwhn           = cwhn ? 1 : 0;
+    kparams->is_f16         = (src0->type == GGML_TYPE_F16) ? 1 : 0;
+    kparams->total_tasks    = total_tasks;
+    kparams->vtcm_size      = (uint32_t) layout.total_bytes;
+    kparams->w              = w;
+    kparams->h              = h;
+    kparams->c              = c;
+    kparams->n              = n;
+    kparams->kw             = kw;
+    kparams->kh             = kh;
+    kparams->ow             = ow_u;
+    kparams->oh             = oh_u;
+    kparams->s0             = s0;
+    kparams->s1             = s1;
+    kparams->p0             = p0;
+    kparams->p1             = p1;
+    kparams->d0             = d0;
+    kparams->d1             = d1;
+    kparams->oh_chunk       = oh_chunk;
+    kparams->n_chunks       = n_chunks;
+    kparams->row_src_bytes  = row_src_bytes;
+    kparams->row_dst_bytes  = row_dst_bytes;
+    kparams->pad_top_stride = pad_top_stride;
+    kparams->div_c          = init_fastdiv_values(c);
+    kparams->div_n_chunks   = init_fastdiv_values(n_chunks > 0 ? n_chunks : 1);
+
+    return true;
+}
+
 template <auto SolveFn>
 static inline bool htp_kparams_invoke(
    const struct ggml_hexagon_session * sess,
@@ -3372,6 +4095,31 @@ void htp_opnode_collapse(const struct ggml_hexagon_session * sess, htp_opnode &
        }
        return;
    }
+
+    if (node.opcode == HTP_OP_UPSCALE) {
+        if (node.inputs.empty() || !node.inputs[0] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        const struct ggml_tensor * src0 = node.inputs[0];
+        const struct ggml_tensor * dst  = node.outputs[0];
+
+        if ((src0->ne[3] > 1 || dst->ne[3] > 1) &&
+            src0->ne[2] == dst->ne[2] && src0->ne[3] == dst->ne[3] &&
+            src0->nb[3] == (size_t) src0->ne[2] * src0->nb[2] &&
+            dst->nb[3]  == (size_t) dst->ne[2]  * dst->nb[2]) {
+            struct ggml_tensor src0_c = *src0;
+            struct ggml_tensor dst_c  = *dst;
+            src0_c.ne[2] *= src0_c.ne[3];
+            src0_c.ne[3] = 1;
+            src0_c.nb[3] = (size_t) src0_c.ne[2] * src0_c.nb[2];
+            dst_c.ne[2]  *= dst_c.ne[3];
+            dst_c.ne[3]  = 1;
+            dst_c.nb[3]  = (size_t) dst_c.ne[2] * dst_c.nb[2];
+            node.inputs[0]  = node.add_dummy(src0_c);
+            node.outputs[0] = node.add_dummy(dst_c);
+        }
+        return;
+    }
 }

 bool htp_kparams_solve(
@@ -3401,6 +4149,9 @@ bool htp_kparams_solve(
        const size_t src2_size = src2 ? hex_round_up(ggml_nbytes(src2), 128) : 0;
        return htp_kparams_solve_matmul_impl(sess, src0, src1, dst, src2 ? src2->nb[1] : 0, src2_size, node.kparams<htp_mm_kernel_params>());
    }
+    if (node.opcode == HTP_OP_UPSCALE) {
+        return htp_kparams_solve_upscale(sess, src0, dst, node.kparams<htp_upscale_kernel_params>());
+    }

    return htp_kparams_solve(sess, node.node, node.kernel_params);
 }
@@ -3411,7 +4162,7 @@ bool htp_kparams_solve(
    void * kparams
 ) {
    const htp_op_code opcode = htp_op_remap(op);
-    if (htp_op_is_unary(opcode) || htp_op_is_glu(opcode)) {
+    if (htp_op_is_unary(opcode) || htp_op_is_glu(opcode) || opcode == HTP_OP_UPSCALE) {
        htp_opnode tmp(opcode, op);
        if (!htp_kparams_solve(sess, tmp)) {
            return false;
@@ -3476,6 +4227,12 @@ bool htp_kparams_solve(
        case HTP_OP_IM2COL_3D:
            return htp_kparams_invoke<htp_kparams_solve_im2col_3d>(sess, op, kparams);

+        case HTP_OP_CONV_2D:
+            return htp_kparams_invoke<htp_kparams_solve_conv2d>(sess, op, kparams);
+
+        case HTP_OP_CONV_2D_DW:
+            return htp_kparams_invoke<htp_kparams_solve_conv2d_dw>(sess, op, kparams);
+
        case HTP_OP_SUM_ROWS:
        case HTP_OP_SUM:
        case HTP_OP_ARGMAX:
@@ -7,6 +7,7 @@
 #include "htp/argsort-ops.h"
 #include "htp/binary-ops.h"
 #include "htp/concat-ops.h"
+#include "htp/conv2d-ops.h"
 #include "htp/cpy-ops.h"
 #include "htp/cumsum-ops.h"
 #include "htp/diag-ops.h"
@@ -28,6 +29,7 @@
 #include "htp/ssm-conv.h"
 #include "htp/sum-rows-ops.h"
 #include "htp/unary-ops.h"
+#include "htp/upscale-ops.h"

 #include <type_traits>
 #include <utility>
@@ -133,4 +135,19 @@ bool htp_kparams_solve_unary_mul_glu(
    struct htp_glu_kernel_params * kparams
 );

+bool htp_kparams_solve_upscale(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * dst,
+    struct htp_upscale_kernel_params * kparams
+);
+
+static inline bool htp_kparams_solve_upscale(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    struct htp_upscale_kernel_params * kparams
+) {
+    return htp_kparams_solve_upscale(sess, op->src[0], op, kparams);
+}
+
 #endif // HTP_KPARAMS_H
@@ -19,6 +19,7 @@
 #include "htp/ssm-conv.h"
 #include "htp/sum-rows-ops.h"
 #include "htp/unary-ops.h"
+#include "htp/upscale-ops.h"

 #include <stdio.h>

@@ -224,6 +225,19 @@ static std::string format_kernel_params(const htp_opnode & node) {
    } else if (node.opcode == HTP_OP_SOLVE_TRI) {
        const auto * kparams = node.kparams<htp_solve_tri_kernel_params>();
        snprintf(str, max_size, "%s nth %d", kparams->batched ? "batched" : "chunked", (int) kparams->n_threads);
+    } else if (node.opcode == HTP_OP_UPSCALE) {
+        const auto * kparams = node.kparams<htp_upscale_kernel_params>();
+        const char * ktype_str = "unknown";
+        switch (kparams->kernel_type) {
+            case HTP_UPSCALE_KERNEL_NEAREST:     ktype_str = "nearest";     break;
+            case HTP_UPSCALE_KERNEL_NEAREST_2X:  ktype_str = "nearest-2x";  break;
+            case HTP_UPSCALE_KERNEL_NEAREST_4X:  ktype_str = "nearest-4x";  break;
+            case HTP_UPSCALE_KERNEL_BILINEAR:    ktype_str = kparams->align_corners ? "bilinear-ac" : "bilinear"; break;
+            case HTP_UPSCALE_KERNEL_BILINEAR_2X: ktype_str = "bilinear-2x"; break;
+            case HTP_UPSCALE_KERNEL_BILINEAR_AA: ktype_str = "bilinear-aa"; break;
+            case HTP_UPSCALE_KERNEL_BICUBIC:     ktype_str = "bicubic";     break;
+        }
+        snprintf(str, max_size, "%s nth %d chk %d vtcm %u", ktype_str, (int) kparams->n_threads, (int) kparams->dst_nrows_per_chunk, (unsigned int) kparams->vtcm_size);
    } else {
        snprintf(str, max_size, "----");
    }
@@ -44,8 +44,10 @@ add_library(${HTP_LIB} SHARED
    pad-ops.c
    argsort-ops.c
    im2col-ops.c
+    conv2d-ops.c
    roll-ops.c
    pool-ops.c
+    upscale-ops.c
    allreduce-ops.c
 )

@@ -47,87 +47,108 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
    const uint32_t src0_ne0 = src0->ne[0];
    const uint32_t src1_ne0 = src1->ne[0];

-    const uint32_t row_end = cctx->row_start + cctx->nrows;
-    const uint32_t start_i = cctx->row_start + ith * cctx->nrows_per_thread;
-    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < row_end) ? (start_i + cctx->nrows_per_thread) : row_end;
-    if (start_i >= end_i) return;
+    // dst rows [row_begin, row_end) of this thread; a dst row is a src0 row followed by a src1 column
+    const uint32_t row_begin = cctx->row_start + ith * cctx->nrows_per_thread;
+    const uint32_t row_end   = MIN(row_begin + cctx->nrows_per_thread, cctx->row_start + cctx->nrows);
+    if (row_begin >= row_end) return;

    dma_queue * dma_q = octx->ctx->dma[ith];
-
-    uint8_t * spad0_base = cctx->spad0_base + ith * cctx->spad0_size_per_thread;
-    uint8_t * spad1_base = cctx->spad1_base + ith * cctx->spad1_size_per_thread;
-
-    const uint32_t block_i = 32;
-    const uint32_t spad1_stride = block_i * sizeof(float);
-
-    const HVX_Vector offsets = hvx_vec_gather_offsets_w(spad1_stride);
-    const uint32_t src1_ne0_padded = hex_round_up(src1_ne0, 32);
-    const uint32_t src0_row_bytes  = src0_ne0 * sizeof(float);
-    const uint32_t src0_row_padded = hex_round_up(src0_row_bytes, VLEN);
-    const uint32_t src0_pre        = src0_row_padded - src0_row_bytes;
-    const uint32_t spad0_row_bytes = src0_row_padded + src1_ne0_padded * sizeof(float);
-
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];

    const struct fastdiv_values * div_ne2 = &cctx->div_ne2;
    const uint32_t ne2 = dst->ne[2];

-    uint32_t p = 0;
-    uint32_t i = start_i;
+    // tile: block_rows dst rows x block_cols src1 columns
+    const uint32_t block_rows = 128;
+    const uint32_t block_cols = 256;
+    // +1 element spreads the gather lanes across VTCM, so they do not serialize on one path
+    const uint32_t spad1_stride = block_rows * sizeof(float) + sizeof(float);

-    const dma_addr_t src1_addr = concat_plane_addr(src1, p, div_ne2, ne2) + i * src1->nb[1];
-    dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], MIN(end_i - i, block_i) * sizeof(float), src1_ne0);
+    const HVX_Vector offsets = hvx_vec_gather_offsets_w(spad1_stride);
+    const uint32_t block_cols_padded = MIN(hex_round_up(src1_ne0, 32), block_cols);
+    const uint32_t src0_row_bytes    = src0_ne0 * sizeof(float);
+    const uint32_t src0_row_padded   = hex_round_up(src0_row_bytes, VLEN);
+    const uint32_t src0_pre          = src0_row_padded - src0_row_bytes;
+    const uint32_t spad0_stride      = src0_row_padded + block_cols_padded * sizeof(float);

-    const dma_addr_t src0_addr = concat_plane_addr(src0, p, div_ne2, ne2) + i * src0->nb[1];
-    dma_queue_push(dma_q, dma_make_data(spad0_base + src0_pre, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, MIN(end_i - i, block_i));
+    // two slots per scratchpad: tile b gathers in slot b&1 while dst[b-1] and src1[b+1] DMA in the other
+    uint8_t * spad0_base = cctx->spad0_base + ith * cctx->spad0_size_per_thread;
+    uint8_t * spad1_base = cctx->spad1_base + ith * cctx->spad1_size_per_thread;
+    const uint32_t spad0_slot = block_rows * spad0_stride;
+    const uint32_t spad1_slot = block_cols_padded * spad1_stride;

-    dma_queue_pop(dma_q); // src1
+    // tiles walk (plane, row, col) with col fastest; src0 goes out with the col == 0 tile only
+    uint32_t plane = 0;
+    uint32_t row   = row_begin;
+    uint32_t col   = 0;
+    uint32_t slot  = 0;
+    bool dst_pending = false;

-    while (p < cctx->nplanes) {
-        const uint32_t current_block_i = MIN(end_i - i, block_i);
+    const dma_addr_t src1_addr = concat_plane_addr(src1, plane, div_ne2, ne2) + row * src1->nb[1];
+    dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], MIN(row_end - row, block_rows) * sizeof(float), MIN(src1_ne0, block_cols));

-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
-        for (uint32_t j = 0; j < src1_ne0; j += 32) {
-            const uint8_t * src_ptr = spad1_base + j * spad1_stride;
-            uint8_t * dst_ptr = spad0_base + src0_row_padded + j * sizeof(float);
-            hvx_transpose_32x32_w_gather(dst_ptr, spad0_row_bytes, src_ptr, spad1_stride, offsets, current_block_i, MIN(src1_ne0 - j, 32));
+    const dma_addr_t src0_addr = concat_plane_addr(src0, plane, div_ne2, ne2) + row * src0->nb[1];
+    dma_queue_push(dma_q, dma_make_data(spad0_base + src0_pre, src0_addr), spad0_stride, src0->nb[1], src0_row_bytes, MIN(row_end - row, block_rows));
+
+    while (plane < cctx->nplanes) {
+        const uint32_t nrows = MIN(row_end - row, block_rows);
+        const uint32_t ncols = MIN(src1_ne0 - col, block_cols);
+
+        uint8_t * spad0 = spad0_base + slot * spad0_slot;
+        uint8_t * spad1 = spad1_base + slot * spad1_slot;
+
+        dma_queue_pop(dma_q); // src1
+        if (col == 0) {
+            dma_queue_pop(dma_q); // src0
        }
-        hvx_gather_sync(spad0_base + src0_row_padded);
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);

-        uint32_t np = p;
-        uint32_t ni = i + block_i;
-        if (ni >= end_i) {
-            ni = start_i;
-            np++;
+        uint32_t next_plane = plane;
+        uint32_t next_row   = row;
+        uint32_t next_col   = col + block_cols;
+        if (next_col >= src1_ne0) {
+            next_col = 0;
+            next_row = row + block_rows;
+            if (next_row >= row_end) {
+                next_row = row_begin;
+                next_plane++;
+            }
        }
-        const bool has_next = np < cctx->nplanes;
-        const uint32_t next_block_i = MIN(end_i - ni, block_i);
+        const bool has_next = next_plane < cctx->nplanes;
+        const uint32_t next_nrows = MIN(row_end - next_row, block_rows);
+        const uint32_t next_ncols = MIN(src1_ne0 - next_col, block_cols);

-        // spad1 is free after the gather sync, prefetch next src1 ahead of the dst write
+        // other spad1 slot is free since the previous gather sync
        if (has_next) {
-            const dma_addr_t nsrc1_addr = concat_plane_addr(src1, np, div_ne2, ne2) + ni * src1->nb[1];
-            dma_queue_push(dma_q, dma_make_data(spad1_base, nsrc1_addr), spad1_stride, src1->nb[0], next_block_i * sizeof(float), src1_ne0);
+            const dma_addr_t addr = concat_plane_addr(src1, next_plane, div_ne2, ne2) + next_row * src1->nb[1] + next_col * src1->nb[0];
+            dma_queue_push(dma_q, dma_make_data(spad1_base + (slot ^ 1) * spad1_slot, addr), spad1_stride, src1->nb[0], next_nrows * sizeof(float), next_ncols);
        }

-        dma_queue_pop(dma_q); // src0
-
-        const dma_addr_t dst_addr = concat_plane_addr(dst, p, div_ne2, ne2) + i * dst->nb[1];
-        dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base + src0_pre), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(float), current_block_i);
-
-        if (has_next) {
-            dma_queue_pop(dma_q); // next src1
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row);
+        for (uint32_t c = 0; c < ncols; c += 32) {
+            hvx_transpose_32x32_w_gather(spad0 + src0_row_padded + c * sizeof(float), spad0_stride, spad1 + c * spad1_stride, spad1_stride, offsets, nrows, MIN(ncols - c, 32));
        }
-        dma_queue_pop(dma_q); // dst
+        hvx_gather_sync(spad0 + src0_row_padded);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row);

-        // spad0 is free after the dst write
-        if (has_next) {
-            const dma_addr_t nsrc0_addr = concat_plane_addr(src0, np, div_ne2, ne2) + ni * src0->nb[1];
-            dma_queue_push(dma_q, dma_make_data(spad0_base + src0_pre, nsrc0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, next_block_i);
+        // other spad0 slot is free once dst[b-1] is written
+        if (dst_pending) {
+            dma_queue_pop(dma_q); // dst
+        }
+        if (has_next && next_col == 0) {
+            const dma_addr_t addr = concat_plane_addr(src0, next_plane, div_ne2, ne2) + next_row * src0->nb[1];
+            dma_queue_push(dma_q, dma_make_data(spad0_base + (slot ^ 1) * spad0_slot + src0_pre, addr), spad0_stride, src0->nb[1], src0_row_bytes, next_nrows);
        }

-        p = np;
-        i = ni;
+        // the col == 0 tile also carries the src0 columns in front
+        const uint32_t lead = (col == 0) ? src0_row_bytes : 0;
+        const dma_addr_t dst_addr = concat_plane_addr(dst, plane, div_ne2, ne2) + row * dst->nb[1] + src0_row_bytes + col * sizeof(float) - lead;
+        dma_queue_push(dma_q, dma_make_data(dst_addr, spad0 + src0_row_padded - lead), dst->nb[1], spad0_stride, lead + ncols * sizeof(float), nrows);
+        dst_pending = true;
+
+        plane = next_plane;
+        row   = next_row;
+        col   = next_col;
+        slot ^= 1;
    }
    dma_queue_flush(dma_q);
 }
@@ -143,87 +164,108 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
    const uint32_t src0_ne0 = src0->ne[0];
    const uint32_t src1_ne0 = src1->ne[0];

-    const uint32_t row_end = cctx->row_start + cctx->nrows;
-    const uint32_t start_i = cctx->row_start + ith * cctx->nrows_per_thread;
-    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < row_end) ? (start_i + cctx->nrows_per_thread) : row_end;
-    if (start_i >= end_i) return;
+    // dst rows [row_begin, row_end) of this thread; a dst row is a src0 row followed by a src1 column
+    const uint32_t row_begin = cctx->row_start + ith * cctx->nrows_per_thread;
+    const uint32_t row_end   = MIN(row_begin + cctx->nrows_per_thread, cctx->row_start + cctx->nrows);
+    if (row_begin >= row_end) return;

    dma_queue * dma_q = octx->ctx->dma[ith];
-
-    uint8_t * spad0_base = cctx->spad0_base + ith * cctx->spad0_size_per_thread;
-    uint8_t * spad1_base = cctx->spad1_base + ith * cctx->spad1_size_per_thread;
-
-    const uint32_t block_i = 64;
-    const uint32_t spad1_stride = block_i * sizeof(__fp16);
-
-    const HVX_Vector offsets = hvx_vec_gather_offsets_h(spad1_stride);
-    const uint32_t src1_ne0_padded = hex_round_up(src1_ne0, 64);
-    const uint32_t src0_row_bytes  = src0_ne0 * sizeof(__fp16);
-    const uint32_t src0_row_padded = hex_round_up(src0_row_bytes, VLEN);
-    const uint32_t src0_pre        = src0_row_padded - src0_row_bytes;
-    const uint32_t spad0_row_bytes = src0_row_padded + src1_ne0_padded * sizeof(__fp16);
-
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];

    const struct fastdiv_values * div_ne2 = &cctx->div_ne2;
    const uint32_t ne2 = dst->ne[2];

-    uint32_t p = 0;
-    uint32_t i = start_i;
+    // tile: block_rows dst rows x block_cols src1 columns
+    const uint32_t block_rows = 128;
+    const uint32_t block_cols = 512;
+    // +1 element spreads the gather lanes across VTCM, so they do not serialize on one path
+    const uint32_t spad1_stride = block_rows * sizeof(__fp16) + sizeof(__fp16);

-    const dma_addr_t src1_addr = concat_plane_addr(src1, p, div_ne2, ne2) + i * src1->nb[1];
-    dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], MIN(end_i - i, block_i) * sizeof(__fp16), src1_ne0);
+    const HVX_Vector offsets = hvx_vec_gather_offsets_h(spad1_stride);
+    const uint32_t block_cols_padded = MIN(hex_round_up(src1_ne0, 64), block_cols);
+    const uint32_t src0_row_bytes    = src0_ne0 * sizeof(__fp16);
+    const uint32_t src0_row_padded   = hex_round_up(src0_row_bytes, VLEN);
+    const uint32_t src0_pre          = src0_row_padded - src0_row_bytes;
+    const uint32_t spad0_stride      = src0_row_padded + block_cols_padded * sizeof(__fp16);

-    const dma_addr_t src0_addr = concat_plane_addr(src0, p, div_ne2, ne2) + i * src0->nb[1];
-    dma_queue_push(dma_q, dma_make_data(spad0_base + src0_pre, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, MIN(end_i - i, block_i));
+    // two slots per scratchpad: tile b gathers in slot b&1 while dst[b-1] and src1[b+1] DMA in the other
+    uint8_t * spad0_base = cctx->spad0_base + ith * cctx->spad0_size_per_thread;
+    uint8_t * spad1_base = cctx->spad1_base + ith * cctx->spad1_size_per_thread;
+    const uint32_t spad0_slot = block_rows * spad0_stride;
+    const uint32_t spad1_slot = block_cols_padded * spad1_stride;

-    dma_queue_pop(dma_q); // src1
+    // tiles walk (plane, row, col) with col fastest; src0 goes out with the col == 0 tile only
+    uint32_t plane = 0;
+    uint32_t row   = row_begin;
+    uint32_t col   = 0;
+    uint32_t slot  = 0;
+    bool dst_pending = false;

-    while (p < cctx->nplanes) {
-        const uint32_t current_block_i = MIN(end_i - i, block_i);
+    const dma_addr_t src1_addr = concat_plane_addr(src1, plane, div_ne2, ne2) + row * src1->nb[1];
+    dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], MIN(row_end - row, block_rows) * sizeof(__fp16), MIN(src1_ne0, block_cols));

-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
-        for (uint32_t j = 0; j < src1_ne0; j += 64) {
-            const uint8_t * src_ptr = spad1_base + j * spad1_stride;
-            uint8_t * dst_ptr = spad0_base + src0_row_padded + j * sizeof(__fp16);
-            hvx_transpose_64x64_h_gather(dst_ptr, spad0_row_bytes, src_ptr, spad1_stride, offsets, current_block_i, MIN(src1_ne0 - j, 64));
+    const dma_addr_t src0_addr = concat_plane_addr(src0, plane, div_ne2, ne2) + row * src0->nb[1];
+    dma_queue_push(dma_q, dma_make_data(spad0_base + src0_pre, src0_addr), spad0_stride, src0->nb[1], src0_row_bytes, MIN(row_end - row, block_rows));
+
+    while (plane < cctx->nplanes) {
+        const uint32_t nrows = MIN(row_end - row, block_rows);
+        const uint32_t ncols = MIN(src1_ne0 - col, block_cols);
+
+        uint8_t * spad0 = spad0_base + slot * spad0_slot;
+        uint8_t * spad1 = spad1_base + slot * spad1_slot;
+
+        dma_queue_pop(dma_q); // src1
+        if (col == 0) {
+            dma_queue_pop(dma_q); // src0
        }
-        hvx_gather_sync(spad0_base + src0_row_padded);
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);

-        uint32_t np = p;
-        uint32_t ni = i + block_i;
-        if (ni >= end_i) {
-            ni = start_i;
-            np++;
+        uint32_t next_plane = plane;
+        uint32_t next_row   = row;
+        uint32_t next_col   = col + block_cols;
+        if (next_col >= src1_ne0) {
+            next_col = 0;
+            next_row = row + block_rows;
+            if (next_row >= row_end) {
+                next_row = row_begin;
+                next_plane++;
+            }
        }
-        const bool has_next = np < cctx->nplanes;
-        const uint32_t next_block_i = MIN(end_i - ni, block_i);
+        const bool has_next = next_plane < cctx->nplanes;
+        const uint32_t next_nrows = MIN(row_end - next_row, block_rows);
+        const uint32_t next_ncols = MIN(src1_ne0 - next_col, block_cols);

-        // spad1 is free after the gather sync, prefetch next src1 ahead of the dst write
+        // other spad1 slot is free since the previous gather sync
        if (has_next) {
-            const dma_addr_t nsrc1_addr = concat_plane_addr(src1, np, div_ne2, ne2) + ni * src1->nb[1];
-            dma_queue_push(dma_q, dma_make_data(spad1_base, nsrc1_addr), spad1_stride, src1->nb[0], next_block_i * sizeof(__fp16), src1_ne0);
+            const dma_addr_t addr = concat_plane_addr(src1, next_plane, div_ne2, ne2) + next_row * src1->nb[1] + next_col * src1->nb[0];
+            dma_queue_push(dma_q, dma_make_data(spad1_base + (slot ^ 1) * spad1_slot, addr), spad1_stride, src1->nb[0], next_nrows * sizeof(__fp16), next_ncols);
        }

-        dma_queue_pop(dma_q); // src0
-
-        const dma_addr_t dst_addr = concat_plane_addr(dst, p, div_ne2, ne2) + i * dst->nb[1];
-        dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base + src0_pre), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(__fp16), current_block_i);
-
-        if (has_next) {
-            dma_queue_pop(dma_q); // next src1
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row);
+        for (uint32_t c = 0; c < ncols; c += 64) {
+            hvx_transpose_64x64_h_gather(spad0 + src0_row_padded + c * sizeof(__fp16), spad0_stride, spad1 + c * spad1_stride, spad1_stride, offsets, nrows, MIN(ncols - c, 64));
        }
-        dma_queue_pop(dma_q); // dst
+        hvx_gather_sync(spad0 + src0_row_padded);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row);

-        // spad0 is free after the dst write
-        if (has_next) {
-            const dma_addr_t nsrc0_addr = concat_plane_addr(src0, np, div_ne2, ne2) + ni * src0->nb[1];
-            dma_queue_push(dma_q, dma_make_data(spad0_base + src0_pre, nsrc0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, next_block_i);
+        // other spad0 slot is free once dst[b-1] is written
+        if (dst_pending) {
+            dma_queue_pop(dma_q); // dst
+        }
+        if (has_next && next_col == 0) {
+            const dma_addr_t addr = concat_plane_addr(src0, next_plane, div_ne2, ne2) + next_row * src0->nb[1];
+            dma_queue_push(dma_q, dma_make_data(spad0_base + (slot ^ 1) * spad0_slot + src0_pre, addr), spad0_stride, src0->nb[1], src0_row_bytes, next_nrows);
        }

-        p = np;
-        i = ni;
+        // the col == 0 tile also carries the src0 columns in front
+        const uint32_t lead = (col == 0) ? src0_row_bytes : 0;
+        const dma_addr_t dst_addr = concat_plane_addr(dst, plane, div_ne2, ne2) + row * dst->nb[1] + src0_row_bytes + col * sizeof(__fp16) - lead;
+        dma_queue_push(dma_q, dma_make_data(dst_addr, spad0 + src0_row_padded - lead), dst->nb[1], spad0_stride, lead + ncols * sizeof(__fp16), nrows);
+        dst_pending = true;
+
+        plane = next_plane;
+        row   = next_row;
+        col   = next_col;
+        slot ^= 1;
    }
    dma_queue_flush(dma_q);
 }
@@ -53,13 +53,17 @@ static inline void htp_concat_transposed_vtcm_layout_build(
    uint32_t type_size,
    uint32_t n_threads) {

-    uint32_t block_i = (type_size == 4) ? 32 : 64;
-    uint32_t spad1_stride = block_i * type_size;
-    uint32_t src1_ne0_padded = hex_round_up(src1_ne0, block_i);
-    uint32_t spad0_row_bytes = hex_round_up(src0_ne0 * type_size, 128) + src1_ne0_padded * type_size;
+    // src1_ne0 is processed in chunks of block_cols tokens, so VTCM does not grow with the batch
+    uint32_t block_rows = 128;
+    uint32_t block_cols = (type_size == 4) ? 256 : 512;
+    // +1 element spreads the gather lanes across VTCM, so they do not serialize on one path
+    uint32_t spad1_stride = block_rows * type_size + type_size;
+    uint32_t block_cols_padded = MIN(hex_round_up(src1_ne0, (type_size == 4) ? 32 : 64), block_cols);
+    uint32_t spad0_row_bytes = hex_round_up(src0_ne0 * type_size, 128) + block_cols_padded * type_size;

-    L->src0_spad_size_per_thread = block_i * spad0_row_bytes;
-    L->src1_spad_size_per_thread = src1_ne0_padded * spad1_stride;
+    // double-buffered: two slots per thread
+    L->src0_spad_size_per_thread = 2 * block_rows * spad0_row_bytes;
+    L->src1_spad_size_per_thread = 2 * block_cols_padded * spad1_stride;

    size_t off = 0;
    VTCM_LAYOUT_ALLOC(off, off_spad0, (size_t) n_threads * L->src0_spad_size_per_thread);
@@ -0,0 +1,197 @@
+#ifndef HTP_CONV2D_OPS_H
+#define HTP_CONV2D_OPS_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-ops.h"
+#include "htp-vtcm.h"
+
+enum htp_conv2d_kernel_type {
+    HTP_CONV2D_KERNEL_UNSUPPORTED = 0,
+    HTP_CONV2D_KERNEL_HMX_WINDOW,       // HMX convolution with sliding-window staging
+    HTP_CONV2D_KERNEL_HMX_TILED,        // tiled implicit-GEMM convolution on HMX
+};
+
+// HMX_WINDOW output staging buffers per thread
+#define HTP_CONV2D_WIN_OBUF 2
+
+// HMX scales of one row tile at the start of VTCM: 32 (scale, bias) fp16 pairs and 128 zero bytes
+#define HTP_CONV2D_HMX_SCALES_SIZE 256
+
+struct htp_conv2d_tiled_params {
+    uint32_t w_kstride;
+    uint32_t k_size;
+    uint32_t k_group;
+    uint32_t k_group_pad;
+    uint32_t n_groups;
+    uint32_t n_k_tiles;
+    uint32_t n_ow_tiles;
+    uint32_t seg_tiles;
+    uint32_t stage_row;
+    uint32_t n_stage;
+    uint32_t wk_chunk;
+    struct fastdiv_values div_ow_tiles;
+    struct fastdiv_values div_oh;
+};
+
+struct htp_conv2d_window_params {
+    uint32_t fold;
+    uint32_t n_cc;
+    uint32_t act_rt_last;
+    uint32_t sr;
+    uint32_t ihs;
+    uint32_t max_off;
+    uint32_t n_bc;
+    uint32_t seg_b;
+    uint32_t seg_t;
+    uint32_t w_cs;
+    struct fastdiv_values div_sr;
+    struct fastdiv_values div_ihs;
+    struct fastdiv_values div_kw;
+    struct fastdiv_values div_tg;
+    struct fastdiv_values div_cc;
+    struct fastdiv_values div_cs;
+};
+
+struct htp_conv2d_kernel_params {
+    uint8_t  kernel_type;
+    uint8_t  n_threads;
+    uint8_t  cwhn;
+    uint8_t  w_kmajor;
+
+    uint32_t total_rows;
+    uint32_t n_row_tiles;
+    uint32_t ct_chunk;
+    uint32_t rt_block;
+    uint32_t stage_size;
+    uint32_t thr_size;
+    uint32_t vtcm_size;
+
+    union {
+        struct htp_conv2d_tiled_params  tiled;
+        struct htp_conv2d_window_params window;
+    };
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_conv2d_kernel_params) <= 128, "htp_conv2d_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_conv2d_kernel_params) <= 128, "htp_conv2d_kernel_params is too large for kernel_params blob");
+#endif
+
+enum htp_conv2d_dw_kernel_type {
+    HTP_CONV2D_DW_KERNEL_UNSUPPORTED = 0,
+    HTP_CONV2D_DW_KERNEL_CWHN,             // CWHN channel-contiguous layout
+    HTP_CONV2D_DW_KERNEL_WHCN_S1,          // WHCN general stride-1
+    HTP_CONV2D_DW_KERNEL_WHCN_S2,          // WHCN general stride-2
+    HTP_CONV2D_DW_KERNEL_WHCN_S1_K3X3,     // WHCN aligned w%32==0, 3x3 s=1 p=1 d=1
+    HTP_CONV2D_DW_KERNEL_WHCN_S2_K3X3_W32, // WHCN w=32->ow=16, 3x3 s=2 p=1 d=1 (2 rows per vector)
+    HTP_CONV2D_DW_KERNEL_WHCN_S2_K3X3_W64, // WHCN w=64->ow=32, 3x3 s=2 p=1 d=1 (1 vector rolling kh)
+    HTP_CONV2D_DW_KERNEL_WHCN_S2_K3X3,     // WHCN aligned w%32==0 (w>=96), 3x3 s=2 p=1 d=1
+    HTP_CONV2D_DW_KERNEL_COUNT,
+};
+
+struct htp_conv2d_dw_kernel_params {
+    uint8_t  kernel_type;
+    uint8_t  n_threads;
+    uint8_t  cwhn;
+    uint8_t  is_f16;
+
+    uint32_t total_tasks;
+    uint32_t vtcm_size;
+    uint32_t w;
+    uint32_t h;
+    uint32_t c;
+    uint32_t n;
+    uint32_t kw;
+    uint32_t kh;
+    uint32_t ow;
+    uint32_t oh;
+    int32_t  s0;
+    int32_t  s1;
+    int32_t  p0;
+    int32_t  p1;
+    int32_t  d0;
+    int32_t  d1;
+    uint32_t oh_chunk;
+    uint32_t n_chunks;
+    uint32_t row_src_bytes;
+    uint32_t row_dst_bytes;
+    uint32_t pad_top_stride;
+    struct fastdiv_values div_c;
+    struct fastdiv_values div_n_chunks;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_conv2d_dw_kernel_params) <= 128, "htp_conv2d_dw_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_conv2d_dw_kernel_params) <= 128, "htp_conv2d_dw_kernel_params is too large for kernel_params blob");
+#endif
+
+struct htp_conv2d_dw_vtcm_layout {
+    size_t   total_bytes;
+    size_t   off_w;
+    size_t   off_src;
+    size_t   off_dst;
+    size_t   src_spad_half_size;
+    size_t   dst_spad_half_size;
+    size_t   src_bytes_per_thread;
+    size_t   dst_bytes_per_thread;
+};
+
+static inline void htp_conv2d_dw_vtcm_layout_build(
+    struct htp_conv2d_dw_vtcm_layout * L,
+    uint32_t max_src_bytes,
+    uint32_t max_dst_bytes,
+    uint32_t w_size,
+    uint32_t n_threads
+) {
+    L->src_spad_half_size   = hex_round_up(max_src_bytes + 256, 128);
+    L->dst_spad_half_size   = hex_round_up(max_dst_bytes + 128, 128);
+    L->src_bytes_per_thread = 2 * L->src_spad_half_size;
+    L->dst_bytes_per_thread = 2 * L->dst_spad_half_size;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_w, hex_round_up(w_size, 128));
+    VTCM_LAYOUT_ALLOC(off, off_src, L->src_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst, L->dst_bytes_per_thread * n_threads);
+    L->total_bytes = off;
+}
+
+struct htp_conv2d_vtcm_layout {
+    size_t total_bytes;
+    size_t off_scales;
+    size_t off_w;
+    size_t off_act[2];
+    size_t off_out[2];
+    size_t off_thr;
+
+    size_t act_size;
+};
+
+static inline void htp_conv2d_vtcm_layout_build(
+    struct htp_conv2d_vtcm_layout * L,
+    uint32_t rt_block,
+    size_t   w_size,
+    size_t   act_size,
+    size_t   out_size,
+    size_t   thr_size,
+    uint32_t n_threads
+) {
+    const size_t sc = hex_round_up((size_t) rt_block * HTP_CONV2D_HMX_SCALES_SIZE, 2048);
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_scales, sc);
+    VTCM_LAYOUT_ALLOC(off, off_w,      w_size);
+    VTCM_LAYOUT_ALLOC(off, off_act[0], act_size);
+    VTCM_LAYOUT_ALLOC(off, off_act[1], act_size);
+    VTCM_LAYOUT_ALLOC(off, off_out[0], out_size);
+    VTCM_LAYOUT_ALLOC(off, off_out[1], out_size);
+    VTCM_LAYOUT_ALLOC(off, off_thr,    (size_t) n_threads * thr_size);
+    L->act_size    = act_size;
+    L->total_bytes = off;
+}
+
+#endif // HTP_CONV2D_OPS_H
@@ -1456,9 +1456,10 @@ static inline void gdn_dma_push_chunk_gb(
    const dma_addr_t g_dma    = g->data + (uint64_t) iv3 * g->nb[3] + (uint64_t) t_chunk * g->nb[2] + (uint64_t) iv1 * g->nb[1];
    const dma_addr_t beta_dma = beta->data + (uint64_t) iv3 * beta->nb[3] + (uint64_t) t_chunk * beta->nb[2] + (uint64_t) iv1 * beta->nb[1];
    const uint32_t row_bytes  = n_batch * sizeof(float);
+    const uint32_t dst_stride = (n_batch > 1) ? row_bytes + sizeof(float) : row_bytes;

-    dma_queue_push(dma_q, dma_make_data(vtcm_g_raw, g_dma), row_bytes, g->nb[2], row_bytes, chunk_size);
-    dma_queue_push(dma_q, dma_make_data(vtcm_b_raw, beta_dma), row_bytes, beta->nb[2], row_bytes, chunk_size);
+    dma_queue_push(dma_q, dma_make_data(vtcm_g_raw, g_dma), dst_stride, g->nb[2], row_bytes, chunk_size);
+    dma_queue_push(dma_q, dma_make_data(vtcm_b_raw, beta_dma), dst_stride, beta->nb[2], row_bytes, chunk_size);
 }

 static inline void gdn_pack_s_col_tiles(
@@ -1707,19 +1708,16 @@ static void gdn_hvx_phase1a_worker(unsigned int n, unsigned int i, void * data)
        hvx_vmem(head->b_f32[curr_buf] + 0)  = hvx_vmem(bctx->vtcm_b_raw + 0);
        hvx_vmem(head->b_f32[curr_buf] + 32) = hvx_vmem(bctx->vtcm_b_raw + 32);
    } else {
-        int32_t offsets[32] __attribute__((aligned(128)));
-        for (int k = 0; k < 32; ++k) {
-            offsets[k] = k * n_batch * sizeof(float);
-        }
-        HVX_Vector vv = *(const HVX_Vector *) offsets;
+        const uint32_t vtcm_stride = (n_batch + 1) * sizeof(float);
+        const HVX_Vector vv = hvx_vec_gather_offsets_w(vtcm_stride);
        const size_t rt_g = (size_t) ((const uint8_t *) bctx->vtcm_g_raw + i * sizeof(float));
        const size_t rt_b = (size_t) ((const uint8_t *) bctx->vtcm_b_raw + i * sizeof(float));
-        const size_t mu   = 64 * n_batch * sizeof(float);
+        const size_t mu   = 64 * vtcm_stride;

        Q6_vgather_ARMVw((HVX_Vector *) (head->g_f32[curr_buf] + 0),  rt_g, mu, vv);
-        Q6_vgather_ARMVw((HVX_Vector *) (head->g_f32[curr_buf] + 32), rt_g + 32 * n_batch * sizeof(float), mu, vv);
+        Q6_vgather_ARMVw((HVX_Vector *) (head->g_f32[curr_buf] + 32), rt_g + 32 * vtcm_stride, mu, vv);
        Q6_vgather_ARMVw((HVX_Vector *) (head->b_f32[curr_buf] + 0),  rt_b, mu, vv);
-        Q6_vgather_ARMVw((HVX_Vector *) (head->b_f32[curr_buf] + 32), rt_b + 32 * n_batch * sizeof(float), mu, vv);
+        Q6_vgather_ARMVw((HVX_Vector *) (head->b_f32[curr_buf] + 32), rt_b + 32 * vtcm_stride, mu, vv);
    }

    const uint32_t t_chunk = bctx->c * 64;
@@ -204,7 +204,7 @@ static inline void htp_gdn_hmx_vtcm_layout_build(
    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_g_f32[1], bh * dma_scalar_sz, pipeline);
    VTCM_LAYOUT_ALLOC(off, off_b_f32[0], bh * dma_scalar_sz);
    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_b_f32[1], bh * dma_scalar_sz, pipeline);
-    const size_t raw_gb_sz = hex_round_up(bh * chunk_size * sizeof(float), 128);
+    const size_t raw_gb_sz = hex_round_up((bh + 1) * chunk_size * sizeof(float), 128);
    VTCM_LAYOUT_ALLOC(off, off_g_raw[0], raw_gb_sz);
    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_g_raw[1], raw_gb_sz, pipeline);
    VTCM_LAYOUT_ALLOC(off, off_b_raw[0], raw_gb_sz);
@@ -182,9 +182,12 @@ int op_gated_delta_net(struct htp_ops_context * octx);
 int op_pad(struct htp_ops_context * octx);
 int op_im2col(struct htp_ops_context * octx);
 int op_im2col_3d(struct htp_ops_context * octx);
+int op_conv_2d(struct htp_ops_context * octx);
+int op_conv_2d_dw(struct htp_ops_context * octx);
 int op_pool_2d(struct htp_ops_context * octx);
 int op_pool_1d(struct htp_ops_context * octx);
 int op_allreduce(struct htp_ops_context * octx);
 int op_roll(struct htp_ops_context * octx);
+int op_upscale(struct htp_ops_context * octx);

 #endif /* HTP_CTX_H */
@@ -96,6 +96,8 @@ enum htp_op_code {
    HTP_OP_CLAMP,
    HTP_OP_LEAKY_RELU,
    HTP_OP_IM2COL,
+    HTP_OP_CONV_2D,
+    HTP_OP_CONV_2D_DW,
    HTP_OP_FENCE,
    HTP_OP_ALLREDUCE,
    HTP_OP_ALLREDUCE_ADD,
@@ -114,6 +116,7 @@ enum htp_op_code {
    HTP_OP_GLU_SIGMOID,
    HTP_OP_UNARY_CEIL,
    HTP_OP_UNARY_SGN,
+    HTP_OP_UPSCALE,

    HTP_OP_INVALID
 };
@@ -134,6 +137,8 @@ enum htp_op_code {

 #define HTP_MMAP_MAX_VMEM  (2147483648u)

+#define HTP_DDR_SPAD_SIZE  (512 * 1024)
+
 enum htp_tensor_flags {
    HTP_TENSOR_WEIGHT   = (1U << 0), // Tensor buffer model weight data (not compute)
    HTP_TENSOR_TILED    = (1U << 1), // Tensor is in tiled format
@@ -11,7 +11,7 @@ struct htp_iface_pmu_conf {
 };

 interface htp_iface : remote_handle64 {
-    AEEResult start(in uint32 sess_id, in uint64 dsp_queue_id, in uint32 n_hvx, in uint32 n_hmx, in uint64 max_vmem, in uint32 power_mode);
+    AEEResult start(in uint32 sess_id, in uint64 dsp_queue_id, in uint32 n_hvx, in uint32 n_hmx, in uint64 vtcm_size, in uint64 max_vmem, in uint32 power_mode);
    AEEResult stop();
    AEEResult mmap(in uint32 fd, in uint64 size);
    AEEResult munmap(in uint32 fd);
@@ -26,6 +26,7 @@ struct htp_im2col_context {
    uint32_t pe_row_base;                          // first output row index assigned to this dev (DMA path)
    uint32_t pe_nrows;                             // number of output rows assigned to this dev (DMA path)
    uint32_t pe_rows_per_thread;                   // output rows per worker
+    uint32_t pe_blocks_per_thread;                 // output-col blocks per worker (blocked DMA path)
    uint32_t pe_src_row_bytes;                     // one output row's source, rounded 256
    uint32_t pe_dst_row_bytes;                     // one output row's dst: OW*patch_stride*2, rounded 256

@@ -38,6 +39,7 @@ struct htp_im2col_context {
    uint32_t pe_owb;                               // output-col block size
    uint32_t pe_wb;                                // staged source window width
    uint32_t pe_nbpr;                              // blocks per row
+    uint32_t pe_gather_nvec;                       // vectors per gathered patch (KH == 1 blocked path)

    // Precomputed dimensions for worker threads
    uint32_t N;
@@ -51,6 +53,7 @@ struct htp_im2col_context {
    uint32_t thread_in[HTP_MAX_NTHREADS];
    uint32_t thread_iod[HTP_MAX_NTHREADS];
    uint32_t thread_ioh[HTP_MAX_NTHREADS];
+    uint32_t thread_c0[HTP_MAX_NTHREADS];

    void * compute;
    uint32_t dst_elem;
@@ -215,6 +218,109 @@ static void FNAME(
 IM2COL_BLOCKED_COMPUTE_BODY(im2col_blocked_compute_f16, uint16_t, hvx_copy_f16_f32_uu, hvx_splat_u16_u)
 IM2COL_BLOCKED_COMPUTE_BODY(im2col_blocked_compute_f32, float,    hvx_copy_f32_uu,     hvx_splat_f32_u)

+// Offsets of one output patch in the staged window, for KH == 1: lane k reads channel k / KW at tap k % KW.
+// Lanes past the patch read the first element and are not stored.
+// The first vector is built lane by lane; each next one is 32 lanes further: the tap advances by 32 % KW
+// and the channel by 32 / KW, plus one more when the tap wraps past KW.
+static __attribute__((noinline)) void im2col_gather_table_build(const struct htp_im2col_context * ictx,
+                                                                HVX_Vector * restrict table) {
+    const int32_t  d0 = ictx->octx->op_params[4];
+    const uint32_t KW = ictx->KW, Wb = ictx->pe_wb;
+
+    int32_t  tap0[32] __attribute__((aligned(128)));
+    int32_t  off0[32] __attribute__((aligned(128)));
+    uint32_t iic = 0, ikw = 0;
+    for (uint32_t k = 0; k < 32; k++) {
+        tap0[k] = (int32_t) ikw;
+        off0[k] = (int32_t) ((iic * Wb + ikw * d0) * sizeof(float));
+        if (++ikw == KW) {
+            ikw = 0;
+            iic++;
+        }
+    }
+    // (iic, ikw) is now (32 / KW, 32 % KW)
+    const HVX_Vector v_kw       = Q6_V_vsplat_R(KW);
+    const HVX_Vector v_kw_m1    = Q6_V_vsplat_R(KW - 1);
+    const HVX_Vector v_tap_step = Q6_V_vsplat_R(ikw);
+    const HVX_Vector v_off_step = Q6_V_vsplat_R((iic * Wb + ikw * d0) * sizeof(float));
+    const HVX_Vector v_off_wrap = Q6_V_vsplat_R((Wb - KW * d0) * sizeof(float));
+    const HVX_Vector v_32       = Q6_V_vsplat_R(32);
+    const HVX_Vector v_n        = Q6_V_vsplat_R(ictx->patch_stride);
+
+    HVX_Vector lane = hvx_vec_gather_offsets_w(1);
+    HVX_Vector tap  = *(HVX_Vector *) tap0;
+    HVX_Vector off  = *(HVX_Vector *) off0;
+    for (uint32_t g = 0; g < ictx->pe_gather_nvec; g++) {
+        table[g] = Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VwVw(v_n, lane), off, Q6_V_vzero());
+
+        lane = Q6_Vw_vadd_VwVw(lane, v_32);
+        tap  = Q6_Vw_vadd_VwVw(tap, v_tap_step);
+        off  = Q6_Vw_vadd_VwVw(off, v_off_step);
+
+        const HVX_VectorPred wrap = Q6_Q_vcmp_gt_VwVw(tap, v_kw_m1);
+        tap = Q6_V_vmux_QVV(wrap, Q6_Vw_vsub_VwVw(tap, v_kw), tap);
+        off = Q6_V_vmux_QVV(wrap, Q6_Vw_vadd_VwVw(off, v_off_wrap), off);
+    }
+}
+
+// KH == 1: each output row is one patch, gathered from the staged [IC][Wb] window with a fixed offset
+// table shifted by s0 per row, then converted with a single copy.
+// scratch holds the offset table followed by the gathered row (both VTCM).
+#define IM2COL_GATHER_COMPUTE_BODY(FNAME, DST_CTYPE, COPY_FN)                                    \
+static void FNAME(                                                                               \
+    const struct htp_im2col_context * ictx,                                                      \
+    void * restrict dstb_raw, const float * restrict srcb_c,                                     \
+    float * restrict scratch,                                                                    \
+    uint32_t nb, uint32_t c0, uint32_t ioh) {                                                    \
+    DST_CTYPE * restrict dstb = (DST_CTYPE *) dstb_raw;                                          \
+    float *     restrict srcb = (float *) srcb_c;                                                \
+    const int32_t  s0 = ictx->octx->op_params[0], p0 = ictx->octx->op_params[2];                 \
+    const int32_t  s1 = ictx->octx->op_params[1], p1 = ictx->octx->op_params[3];                 \
+    const uint32_t IC = ictx->IC, IH = ictx->IH, IW = ictx->IW, Wb = ictx->pe_wb;                \
+    const uint32_t patch_stride = ictx->patch_stride;                                            \
+    const uint32_t ng   = ictx->pe_gather_nvec;                                                  \
+    const int32_t  win0 = (int32_t) c0 * s0 - p0;                                                \
+    const int32_t  iih  = (int32_t) ioh * s1 - p1;                                               \
+                                                                                                 \
+    /* the DMA stage fills only the columns inside [0, IW): zero the rest of the window */       \
+    const int32_t lo = win0 < 0 ? -win0 : 0;                                                     \
+    const int32_t hi = MIN((int32_t) Wb, (int32_t) IW - win0);                                   \
+    if (iih < 0 || iih >= (int32_t) IH || hi <= lo) {                                            \
+        for (uint32_t iic = 0; iic < IC; iic++) {                                                \
+            hvx_splat_f32_u(srcb + (size_t) iic * Wb, 0, Wb);                                    \
+        }                                                                                        \
+    } else if (lo > 0 || hi < (int32_t) Wb) {                                                    \
+        for (uint32_t iic = 0; iic < IC; iic++) {                                                \
+            float * row = srcb + (size_t) iic * Wb;                                              \
+            if (lo > 0) {                                                                        \
+                hvx_splat_f32_u(row, 0, (uint32_t) lo);                                          \
+            }                                                                                    \
+            if (hi < (int32_t) Wb) {                                                             \
+                hvx_splat_f32_u(row + hi, 0, Wb - (uint32_t) hi);                                \
+            }                                                                                    \
+        }                                                                                        \
+    }                                                                                            \
+                                                                                                 \
+    const HVX_Vector * table  = (const HVX_Vector *) scratch;                                    \
+    float *            row    = scratch + (size_t) ng * 32;                                      \
+    const uint32_t     region = IC * Wb * sizeof(float);                                         \
+                                                                                                 \
+    for (uint32_t j = 0; j < nb; j++) {                                                          \
+        const uint32_t shift = j * (uint32_t) s0 * sizeof(float);                                \
+        const size_t   base  = (size_t) ((const uint8_t *) srcb + shift);                        \
+        const uint32_t mu    = region - shift - 1;                                               \
+        for (uint32_t g = 0; g < ng; g++) {                                                      \
+            Q6_vgather_ARMVw((HVX_Vector *) (row + g * 32), base, mu, table[g]);                 \
+        }                                                                                        \
+        hvx_gather_sync(row + (ng - 1) * 32);                                                    \
+        DST_CTYPE * out = dstb + (size_t) j * patch_stride;                                      \
+        COPY_FN((uint8_t *) out, (const uint8_t *) row, patch_stride);                           \
+    }                                                                                            \
+}
+
+IM2COL_GATHER_COMPUTE_BODY(im2col_gather_compute_f16, uint16_t, hvx_copy_f16_f32_uu)
+IM2COL_GATHER_COMPUTE_BODY(im2col_gather_compute_f32, float,    hvx_copy_f32_uu)
+
 static void im2col_blocked_dma_thread(unsigned int nth, unsigned int ith, void * data) {
    (void) nth;
    struct htp_im2col_context * ictx      = (struct htp_im2col_context *) data;
@@ -234,24 +340,24 @@ static void im2col_blocked_dma_thread(unsigned int nth, unsigned int ith, void *
    float *        srcb2[2]               = { (float *) srcb_base, (float *) (srcb_base + ictx->pe_src_row_bytes) };
    void *         dstb2[2]               = { dstb_base, dstb_base + ictx->pe_dst_row_bytes };
    float *        scratch                = (float *) (srcb_base + 2 * ictx->pe_src_row_bytes);
-    const uint32_t row_end_max            = ictx->pe_row_base + ictx->pe_nrows;
-    const uint32_t row_start              = ictx->pe_row_base + ictx->pe_rows_per_thread * ith;
-    const uint32_t row_end                = MIN(row_start + ictx->pe_rows_per_thread, row_end_max);
-    if (row_start >= row_end) {
-        return;
-    }
-    const uint32_t nrows_local  = row_end - row_start;
-    const uint32_t total_blocks = nrows_local * nbpr;
-    if (total_blocks == 0) {
+    // work is split over output-col blocks, so a single output row (1D) still uses every thread
+    const uint32_t blk_end_max  = ictx->pe_nrows * nbpr;
+    const uint32_t blk_start    = ictx->pe_blocks_per_thread * ith;
+    const uint32_t blk_end      = MIN(blk_start + ictx->pe_blocks_per_thread, blk_end_max);
+    if (blk_start >= blk_end) {
        return;
    }
+    const uint32_t total_blocks = blk_end - blk_start;
    struct im2col_cursor cur = {
-        .r   = row_start,
+        .r   = ictx->thread_in[ith] * ictx->OH + ictx->thread_ioh[ith],
        .in  = ictx->thread_in[ith],
        .iod = 0,
        .ioh = ictx->thread_ioh[ith],
-        .c0  = 0,
+        .c0  = ictx->thread_c0[ith],
    };
+    if (KH == 1) {
+        im2col_gather_table_build(ictx, (HVX_Vector *) scratch);
+    }
    struct im2col_cursor stage = cur;
    for (uint32_t i = 0; i < total_blocks && i < 2; i++) {
        dma_queue_push(dma_q, dma_make_data(dst_data, dstb2[i]), 0, 0, 0, 0);
@@ -307,6 +413,37 @@ static void im2col_patchembed_compute_f16(
    const uint32_t KH = ictx->KH, KW = ictx->KW;
    const uint32_t OW = ictx->OW, patch_stride = ictx->patch_stride;

+    if (KW == 16 && (KH % 4 == 0)) {
+        const HVX_VectorPred q_lo = Q6_Q_vsetq_R(64);
+        for (uint32_t iow = 0; iow < OW; iow++) {
+            uint16_t * dst_patch = dstb + (size_t) iow * patch_stride;
+            for (uint32_t iic = 0; iic < IC; iic++) {
+                for (uint32_t ikh = 0; ikh < KH; ikh += 4) {
+                    const int32_t iih0 = (int32_t) ioh * (int32_t) KH + (int32_t) (ikh + 0);
+                    const int32_t iih1 = (int32_t) ioh * (int32_t) KH + (int32_t) (ikh + 1);
+                    const int32_t iih2 = (int32_t) ioh * (int32_t) KH + (int32_t) (ikh + 2);
+                    const int32_t iih3 = (int32_t) ioh * (int32_t) KH + (int32_t) (ikh + 3);
+
+                    const float * s0 = srcb + (size_t) (iic * KH + ikh + 0) * IW + (size_t) iow * 16;
+                    const float * s1 = srcb + (size_t) (iic * KH + ikh + 1) * IW + (size_t) iow * 16;
+                    const float * s2 = srcb + (size_t) (iic * KH + ikh + 2) * IW + (size_t) iow * 16;
+                    const float * s3 = srcb + (size_t) (iic * KH + ikh + 3) * IW + (size_t) iow * 16;
+
+                    const HVX_Vector v0 = (iih0 >= 0 && iih0 < (int32_t) IH) ? *((const HVX_UVector *) s0) : Q6_V_vzero();
+                    const HVX_Vector v1 = (iih1 >= 0 && iih1 < (int32_t) IH) ? *((const HVX_UVector *) s1) : Q6_V_vzero();
+                    const HVX_Vector v2 = (iih2 >= 0 && iih2 < (int32_t) IH) ? *((const HVX_UVector *) s2) : Q6_V_vzero();
+                    const HVX_Vector v3 = (iih3 >= 0 && iih3 < (int32_t) IH) ? *((const HVX_UVector *) s3) : Q6_V_vzero();
+
+                    const HVX_Vector w01 = Q6_V_vmux_QVV(q_lo, v0, Q6_V_vror_VR(v1, 64));
+                    const HVX_Vector w23 = Q6_V_vmux_QVV(q_lo, v2, Q6_V_vror_VR(v3, 64));
+
+                    *((HVX_Vector *) (dst_patch + (size_t) (iic * KH + ikh) * 16)) = hvx_vec_f32_to_f16(w01, w23);
+                }
+            }
+        }
+        return;
+    }
+
    for (uint32_t iow = 0; iow < OW; iow++) {
        uint16_t * dst_patch = dstb + (size_t) iow * patch_stride;
        for (uint32_t ikh = 0; ikh < KH; ikh++) {
@@ -332,6 +469,30 @@ static void im2col_patchembed_compute_f32(
    const uint32_t KH = ictx->KH, KW = ictx->KW;
    const uint32_t OW = ictx->OW, patch_stride = ictx->patch_stride;

+    if (KW == 16 && (KH % 2 == 0)) {
+        const HVX_VectorPred q_lo = Q6_Q_vsetq_R(64);
+        for (uint32_t iow = 0; iow < OW; iow++) {
+            float * dst_patch = dstb + (size_t) iow * patch_stride;
+            for (uint32_t iic = 0; iic < IC; iic++) {
+                for (uint32_t ikh = 0; ikh < KH; ikh += 2) {
+                    const int32_t iih0 = (int32_t) ioh * (int32_t) KH + (int32_t) (ikh + 0);
+                    const int32_t iih1 = (int32_t) ioh * (int32_t) KH + (int32_t) (ikh + 1);
+
+                    const float * s0 = srcb + (size_t) (iic * KH + ikh + 0) * IW + (size_t) iow * 16;
+                    const float * s1 = srcb + (size_t) (iic * KH + ikh + 1) * IW + (size_t) iow * 16;
+
+                    const HVX_Vector v0 = (iih0 >= 0 && iih0 < (int32_t) IH) ? *((const HVX_UVector *) s0) : Q6_V_vzero();
+                    const HVX_Vector v1 = (iih1 >= 0 && iih1 < (int32_t) IH) ? *((const HVX_UVector *) s1) : Q6_V_vzero();
+
+                    const HVX_Vector w01 = Q6_V_vmux_QVV(q_lo, v0, Q6_V_vror_VR(v1, 64));
+
+                    *((HVX_Vector *) (dst_patch + (size_t) (iic * KH + ikh) * 16)) = w01;
+                }
+            }
+        }
+        return;
+    }
+
    for (uint32_t iow = 0; iow < OW; iow++) {
        float * dst_patch = dstb + (size_t) iow * patch_stride;
        for (uint32_t ikh = 0; ikh < KH; ikh++) {
@@ -784,12 +945,16 @@ int op_im2col(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

-    if (!htp_ops_context_set_n_threads(octx, MIN((uint32_t) kparams->n_threads, nrows))) {
+    // the blocked path splits work over output-col blocks, so a single output row (1D) can use every thread
+    const bool     blocked = kparams->kernel_type == HTP_IM2COL_KERNEL_BLOCKED_DMA;
+    const uint32_t nwork   = blocked ? nrows * kparams->nbpr : nrows;
+    if (!htp_ops_context_set_n_threads(octx, MIN((uint32_t) kparams->n_threads, nwork))) {
        return HTP_STATUS_INVAL_PARAMS;
    }

    struct htp_im2col_vtcm_layout layout;
-    htp_im2col_vtcm_layout_build(&layout, kparams->src_row_bytes, kparams->dst_row_bytes, kparams->n_threads);
+    htp_im2col_vtcm_layout_build(&layout, kparams->src_row_bytes, kparams->dst_row_bytes, kparams->n_threads,
+                                 htp_im2col_scratch_bytes(kparams->kernel_type, kparams->KH, kparams->patch_stride));

    struct htp_im2col_context ictx = { 0 };
    uint8_t * const base = octx->ctx->vtcm_base;
@@ -799,6 +964,7 @@ int op_im2col(struct htp_ops_context * octx) {
    ictx.pe_row_base            = row_base;
    ictx.pe_nrows               = nrows;
    ictx.pe_rows_per_thread     = fastdiv(nrows + pth - 1, &octx->n_threads_div);
+    ictx.pe_blocks_per_thread   = fastdiv(nwork + pth - 1, &octx->n_threads_div);
    ictx.pe_src_row_bytes       = kparams->src_row_bytes;
    ictx.pe_dst_row_bytes       = kparams->dst_row_bytes;
    ictx.pe_vtcm_src            = VTCM_LAYOUT_PTR(uint8_t, base, layout.off_src);
@@ -820,12 +986,22 @@ int op_im2col(struct htp_ops_context * octx) {
    ictx.OH                     = kparams->OH;
    ictx.OW                     = kparams->OW;
    ictx.patch_stride           = kparams->patch_stride;
+    ictx.pe_gather_nvec         = hex_round_up(kparams->patch_stride, 32) / 32;

+    // a thread starts at a row, or at an output-col block (possibly mid-row) on the blocked path
    for (uint32_t ith = 0; ith < pth; ith++) {
-        const uint32_t r  = row_base + ictx.pe_rows_per_thread * ith;
+        uint32_t r  = row_base + ictx.pe_rows_per_thread * ith;
+        uint32_t c0 = 0;
+        if (blocked) {
+            const uint32_t blk = ictx.pe_blocks_per_thread * ith;
+            const uint32_t q   = fastdiv(blk, &kparams->div_nbpr);
+            r  = row_base + q;
+            c0 = (blk - q * kparams->nbpr) * kparams->owb;
+        }
        const uint32_t in = fastdiv(r, &kparams->div_OH);
        ictx.thread_ioh[ith] = r - in * kparams->OH;
        ictx.thread_in[ith]  = in;
+        ictx.thread_c0[ith]  = c0;
    }

    ictx.dst_elem = dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float);
@@ -835,9 +1011,14 @@ int op_im2col(struct htp_ops_context * octx) {
                                                 : (void *) im2col_patchembed_compute_f32;
        work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth);
        return HTP_STATUS_OK;
-    } else if (kparams->kernel_type == HTP_IM2COL_KERNEL_BLOCKED_DMA) {
-        ictx.compute = dst->type == HTP_TYPE_F16 ? (void *) im2col_blocked_compute_f16
-                                                 : (void *) im2col_blocked_compute_f32;
+    } else if (blocked) {
+        if (kparams->KH == 1) {
+            ictx.compute = dst->type == HTP_TYPE_F16 ? (void *) im2col_gather_compute_f16
+                                                     : (void *) im2col_gather_compute_f32;
+        } else {
+            ictx.compute = dst->type == HTP_TYPE_F16 ? (void *) im2col_blocked_compute_f16
+                                                     : (void *) im2col_blocked_compute_f32;
+        }
        work_queue_run(octx->ctx->work_queue, im2col_blocked_dma_thread, &ictx, pth);
        return HTP_STATUS_OK;
    }
@@ -873,7 +1054,7 @@ int op_im2col_3d(struct htp_ops_context * octx) {
    }

    struct htp_im2col_vtcm_layout layout;
-    htp_im2col_vtcm_layout_build(&layout, kparams->src_row_bytes, kparams->dst_row_bytes, kparams->n_threads);
+    htp_im2col_vtcm_layout_build(&layout, kparams->src_row_bytes, kparams->dst_row_bytes, kparams->n_threads, 256);

    struct htp_im2col_context ictx = { 0 };
    uint8_t * const base = octx->ctx->vtcm_base;
@@ -20,13 +20,15 @@ struct htp_im2col_vtcm_layout {
    uint32_t dst_bytes_per_thread;
 };

+// scratch_bytes: per-thread VTCM after the two src buffers (at least 256)
 static inline void htp_im2col_vtcm_layout_build(
    struct htp_im2col_vtcm_layout * L,
    uint32_t src_row_bytes,
    uint32_t dst_row_bytes,
-    uint32_t n_threads) {
+    uint32_t n_threads,
+    uint32_t scratch_bytes) {

-    L->src_bytes_per_thread = 2 * src_row_bytes + 256;
+    L->src_bytes_per_thread = 2 * src_row_bytes + (scratch_bytes > 256 ? scratch_bytes : 256);
    L->dst_bytes_per_thread = 2 * dst_row_bytes;

    size_t off = 0;
@@ -35,6 +37,15 @@ static inline void htp_im2col_vtcm_layout_build(
    L->total_bytes = off;
 }

+// per-thread scratch: the KH == 1 blocked path keeps its offset table and one gathered row there,
+// both padded to whole vectors
+static inline uint32_t htp_im2col_scratch_bytes(uint32_t kernel_type, uint32_t KH, uint32_t patch_stride) {
+    if (kernel_type == HTP_IM2COL_KERNEL_BLOCKED_DMA && KH == 1) {
+        return 2 * hex_round_up(patch_stride, 32) * (uint32_t) sizeof(float);
+    }
+    return 256;
+}
+
 struct htp_im2col_kernel_params {
    uint8_t  kernel_type;
    uint8_t  n_threads;
@@ -58,6 +69,7 @@ struct htp_im2col_kernel_params {

    struct fastdiv_values div_OH;
    struct fastdiv_values div_OD;
+    struct fastdiv_values div_nbpr;
 };

 #if defined(__cplusplus)
@@ -325,8 +325,11 @@ static unsigned int vtcm_query_size(void) {
    return vtcm_size;
 }

-static int vtcm_alloc(struct htp_context * ctx) {
+static int vtcm_alloc(struct htp_context * ctx, uint64_t size_override) {
    unsigned int vtcm_size = vtcm_query_size();
+    if (size_override) {
+        vtcm_size = (unsigned int) size_override;
+    }

    compute_res_attr_t attr;
    HAP_compute_res_attr_init(&attr);
@@ -438,7 +441,7 @@ static const struct htp_power_profile htp_power_profiles[HTP_POWER_MODE_COUNT] =

 #undef HTP_POWER_PROFILE_ENTRY

-AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp_queue_id, uint32_t n_hvx, uint32_t n_hmx, uint64_t max_vmem, uint32_t power_mode) {
+AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp_queue_id, uint32_t n_hvx, uint32_t n_hmx, uint64_t vtcm_size, uint64_t max_vmem, uint32_t power_mode) {
    struct htp_handle * h = (struct htp_handle *) handle;
    if (!h) {
        return AEE_EBADPARM;
@@ -557,7 +560,7 @@ AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp
    ctx->max_vmem    = max_vmem;
    ctx->dsp_queue   = dsp_queue;

-    err = vtcm_alloc(ctx);
+    err = vtcm_alloc(ctx, vtcm_size);
    if (err != AEE_SUCCESS) {
        FARF(ERROR, "Unable to allocate VTCM");
        htp_iface_stop(handle);
@@ -674,7 +677,7 @@ AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp
        dma_ptr_curr += size_dma_q;
    }

-    ctx->ddr_spad_size = 512 * 1024; // 512 KB
+    ctx->ddr_spad_size = HTP_DDR_SPAD_SIZE;
    ctx->ddr_spad_base = memalign(128, ctx->ddr_spad_size);

    void * wq_ptr = (void *) ((uintptr_t) block + offset_wq);
@@ -1022,6 +1025,12 @@ static int execute_op(struct htp_ops_context * octx) {
        case HTP_OP_IM2COL_3D:
            return op_im2col_3d(octx);

+        case HTP_OP_CONV_2D:
+            return op_conv_2d(octx);
+
+        case HTP_OP_CONV_2D_DW:
+            return op_conv_2d_dw(octx);
+
        case HTP_OP_ROLL:
            return op_roll(octx);

@@ -1040,6 +1049,9 @@ static int execute_op(struct htp_ops_context * octx) {
        case HTP_OP_TRI:
            return op_unary(octx);

+        case HTP_OP_UPSCALE:
+            return op_upscale(octx);
+
        case HTP_OP_INVALID:
            break;
    }
@@ -16,6 +16,7 @@
 #include "hvx-dump.h"
 #include "hvx-arith.h"
 #include "hvx-reduce.h"
+#include "hvx-copy.h"

 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
@@ -1413,7 +1414,7 @@ static void hvx_mv_id_tiled(unsigned int nth, unsigned int ith, void * data) {

    for (uint32_t ie1 = 0; ie1 < n_aids; ++ie1) {  // for each expert
        const int32_t eid = *(const int32_t *) ((const uint8_t *) src2->data + ie1 * src2->nb[0]);
-        if (eid == -1) {
+        if (eid < 0) {
            continue;
        }
        assert(eid >= 0 && eid < (int32_t) n_ids);
@@ -4111,7 +4112,7 @@ static int hvx_mm_matmul_id_nx(
    return HTP_STATUS_OK;
 }

-static inline void scan_expert_ids_n(
+static __attribute__((noinline)) uint32_t scan_expert_ids_n(
    const struct htp_tensor * ids,
    const uint32_t n_ids,
    uint32_t n_as,
@@ -4121,15 +4122,17 @@ static inline void scan_expert_ids_n(
 ) {
    const size_t ids_nb1 = ids->nb[1];
    const uint8_t * ids_data = (const uint8_t *) ids->data;
+    uint32_t n_skipped = 0;

    for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
        const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
        for (uint32_t id = 0; id < n_ids; ++id) {
            const int32_t i02 = row_ptr[id];
-            if (i02 == -1) {
+            if (i02 < 0) {
+                n_skipped++;
                continue;
            }
-            assert(i02 >= 0 && i02 < n_as);
+            assert(i02 >= 0 && i02 < (int32_t) n_as);

            if (matrix_rows) {
                matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
@@ -4137,31 +4140,10 @@ static inline void scan_expert_ids_n(
            counts[i02] += 1;
        }
    }
+    return n_skipped;
 }

-// the matmuls only write the rows that an expert owns, so the rows of the skipped slots are zeroed here
-static inline void zero_skipped_dst_rows(
-    const struct htp_tensor * ids,
-    uint32_t n_ids,
-    const struct htp_tensor * dst,
-    uint32_t ne0,
-    size_t nb1,
-    size_t nb2
-) {
-    const uint8_t * ids_data = (const uint8_t *) ids->data;
-
-    for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
-        for (uint32_t id = 0; id < n_ids; ++id) {
-            const int32_t i02 = *(const int32_t *) (ids_data + iid1 * ids->nb[1] + id * ids->nb[0]);
-            if (i02 != -1) {
-                continue;
-            }
-            memset((uint8_t *) (dst->data + id * nb1 + iid1 * nb2), 0, ne0 * sizeof(float));
-        }
-    }
-}
-
-static inline void scan_expert_ids(
+static __attribute__((noinline)) uint32_t scan_expert_ids(
    const struct htp_tensor * ids,
    uint32_t n_ids,
    uint32_t n_as,
@@ -4173,23 +4155,25 @@ static inline void scan_expert_ids(

    if (ids_nb0 == 4) {
        switch (n_ids) {
-            case 8:  scan_expert_ids_n(ids, 8,     n_as, counts, matrix_rows, mapping_stride); break;
-            case 4:  scan_expert_ids_n(ids, 4,     n_as, counts, matrix_rows, mapping_stride); break;
-            case 2:  scan_expert_ids_n(ids, 2,     n_as, counts, matrix_rows, mapping_stride); break;
-            default: scan_expert_ids_n(ids, n_ids, n_as, counts, matrix_rows, mapping_stride); break;
+            case 8:  return scan_expert_ids_n(ids, 8,     n_as, counts, matrix_rows, mapping_stride);
+            case 4:  return scan_expert_ids_n(ids, 4,     n_as, counts, matrix_rows, mapping_stride);
+            case 2:  return scan_expert_ids_n(ids, 2,     n_as, counts, matrix_rows, mapping_stride);
+            default: return scan_expert_ids_n(ids, n_ids, n_as, counts, matrix_rows, mapping_stride);
        }
    } else {
        // Strided fallback
        const size_t ids_nb1 = ids->nb[1];
        const uint8_t * ids_data = (const uint8_t *) ids->data;
+        uint32_t n_skipped = 0;
        for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
            const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
            for (uint32_t id = 0; id < n_ids; ++id) {
                const int32_t i02 = *(const int32_t *) ((const uint8_t *) row_ptr + id * ids_nb0);
-                if (i02 == -1) {
+                if (i02 < 0) {
+                    n_skipped++;
                    continue;
                }
-                assert(i02 >= 0 && i02 < n_as);
+                assert(i02 >= 0 && i02 < (int32_t) n_as);

                if (matrix_rows) {
                    matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
@@ -4197,9 +4181,124 @@ static inline void scan_expert_ids(
                counts[i02] += 1;
            }
        }
+        return n_skipped;
    }
 }

+static __attribute__((noinline)) uint32_t count_skipped_ids_single(const struct htp_tensor * ids, uint32_t n_ids) {
+    const uint8_t * ids_data = (const uint8_t *) ids->data;
+    const size_t ids_nb0 = ids->nb[0];
+    uint32_t n_skipped = 0;
+    for (uint32_t id = 0; id < n_ids; ++id) {
+        const int32_t i02 = *(const int32_t *) (ids_data + id * ids_nb0);
+        if (i02 < 0) {
+            n_skipped++;
+        }
+    }
+    return n_skipped;
+}
+
+// Zero rows for skipped slots (ids < 0). Matmuls only write rows that an expert owns.
+static void zero_skipped_dst_rows(
+    struct htp_ops_context * octx,
+    const struct htp_tensor * ids,
+    uint32_t n_ids,
+    const struct htp_tensor * dst
+) {
+    const uint32_t elem_size = (dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float);
+    uint32_t col_start = 0;
+    uint32_t col_count = dst->ne[0];
+
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_can_row_partition(dst, elem_size);
+        const uint32_t chunk_elems = (dst->type == HTP_TYPE_F16) ? 64 : 32;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+            dst->ne[0], can_split ? chunk_elems : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        col_start = range.start;
+        col_count = range.count;
+    }
+
+    if (col_count == 0) {
+        return;
+    }
+
+    const uint8_t * ids_data = (const uint8_t *) ids->data;
+    const size_t ids_nb0 = ids->nb[0];
+    const size_t ids_nb1 = ids->nb[1];
+    const size_t dst_nb1 = dst->nb[1];
+    const size_t dst_nb2 = dst->nb[2];
+    const uint32_t col_offset_bytes = col_start * elem_size;
+
+    for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
+        const uint8_t * ids_row = ids_data + iid1 * ids_nb1;
+        uint8_t * dst_row_group = (uint8_t *) dst->data + iid1 * dst_nb2 + col_offset_bytes;
+        for (uint32_t id = 0; id < n_ids; ++id) {
+            const int32_t i02 = *(const int32_t *) (ids_row + id * ids_nb0);
+            if (i02 < 0) {
+                hvx_splat_u(dst_row_group + id * dst_nb1, Q6_V_vzero(), col_count, elem_size);
+            }
+        }
+    }
+}
+
+static int htp_mmid_prepare_mapping(
+    struct htp_ops_context * octx,
+    const struct htp_tensor * ids,
+    uint32_t n_ids,
+    uint32_t n_as,
+    uint32_t act_nrows,
+    uint8_t ** out_mapping_buf,
+    uint32_t ** out_matrix_row_counts,
+    struct mmid_row_mapping ** out_matrix_rows,
+    uint32_t * out_mapping_stride,
+    uint32_t * out_n_skipped
+) {
+    uint8_t  * mapping_buf       = octx->ctx->ddr_spad_base;
+    uint32_t   mapping_stride    = 1;
+    uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
+    struct mmid_row_mapping * matrix_rows = NULL;
+    uint32_t n_skipped = 0;
+
+    if (act_nrows > 1) {
+        const size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
+        assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size);
+
+        hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]);
+
+        hvx_splat_u(matrix_row_counts, Q6_V_vzero(), n_as, sizeof(uint32_t));
+        n_skipped = scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);
+
+        uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as);
+        mapping_stride = max_count > 0 ? max_count : 1;
+
+        size_t matrix_row_map_size  = n_as * mapping_stride * sizeof(struct mmid_row_mapping);
+        const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
+
+        if (total_map_size > octx->ctx->ddr_spad_size) {
+            mapping_buf = memalign(128, total_map_size);
+            if (!mapping_buf) {
+                return HTP_STATUS_INTERNAL_ERR;
+            }
+        }
+
+        matrix_row_counts = (uint32_t *) mapping_buf;
+        matrix_rows       = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size);
+
+        hvx_splat_u(matrix_row_counts, Q6_V_vzero(), n_as, sizeof(uint32_t));
+        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
+    } else {
+        n_skipped = count_skipped_ids_single(ids, n_ids);
+    }
+
+    *out_mapping_buf       = mapping_buf;
+    *out_matrix_row_counts = matrix_row_counts;
+    *out_matrix_rows       = matrix_rows;
+    *out_mapping_stride    = mapping_stride;
+    *out_n_skipped         = n_skipped;
+
+    return HTP_STATUS_OK;
+}
+
 int op_matmul_id(struct htp_ops_context * octx) {
    htp_matmul_tensors_preamble;

@@ -4235,41 +4334,22 @@ int op_matmul_id(struct htp_ops_context * octx) {
    const int n_ids = ids->ne[0];  // n_expert_used
    const int n_as  = ne02;        // n_expert

-    uint8_t  * mapping_buf       = octx->ctx->ddr_spad_base;
-    uint32_t   mapping_stride    = 1;
-    uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
+    uint8_t  * mapping_buf       = NULL;
+    uint32_t * matrix_row_counts = NULL;
    struct mmid_row_mapping * matrix_rows = NULL;
+    uint32_t   mapping_stride    = 1;
+    uint32_t   n_skipped         = 0;

-    if (act_nrows > 1) {
-        const size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
-        assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size);
-
-        hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]);
-
-        memset(matrix_row_counts, 0, matrix_row_counts_size);
-        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);
-
-        uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as);
-        mapping_stride = max_count > 0 ? max_count : 1;
-
-        size_t matrix_row_map_size  = n_as * mapping_stride * sizeof(struct mmid_row_mapping);
-        const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
-
-        if (total_map_size > octx->ctx->ddr_spad_size) {
-            mapping_buf = memalign(128, total_map_size);
-            if (!mapping_buf) {
-                return HTP_STATUS_INTERNAL_ERR;
-            }
-        }
-
-        matrix_row_counts = (uint32_t *) mapping_buf;
-        matrix_rows       = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size);
-
-        memset(matrix_row_counts, 0, n_as * sizeof(uint32_t));
-        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
+    int map_status = htp_mmid_prepare_mapping(octx, ids, n_ids, n_as, act_nrows,
+                                              &mapping_buf, &matrix_row_counts, &matrix_rows,
+                                              &mapping_stride, &n_skipped);
+    if (map_status != HTP_STATUS_OK) {
+        return map_status;
    }

-    zero_skipped_dst_rows(ids, n_ids, dst, ne0, nb1, nb2);
+    if (n_skipped > 0) {
+        zero_skipped_dst_rows(octx, ids, n_ids, dst);
+    }

    mmctx->matrix_row_counts    = matrix_row_counts;
    mmctx->matrix_rows          = matrix_rows;
@@ -4360,38 +4440,26 @@ int op_matmul_id_nx(struct htp_ops_context * octx) {
    const int n_ids = ids->ne[0];
    const int n_as  = src0->ne[2];

-    uint8_t  * mapping_buf       = octx->ctx->ddr_spad_base;
-    uint32_t   mapping_stride    = 1;
-    uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
+    uint8_t  * mapping_buf       = NULL;
+    uint32_t * matrix_row_counts = NULL;
    struct mmid_row_mapping * matrix_rows = NULL;
+    uint32_t   mapping_stride    = 1;
+    uint32_t   n_skipped         = 0;

-    if (act_nrows > 1) {
-        const size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
-        assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size);
+    int map_status = htp_mmid_prepare_mapping(octx, ids, n_ids, n_as, act_nrows,
+                                              &mapping_buf, &matrix_row_counts, &matrix_rows,
+                                              &mapping_stride, &n_skipped);
+    if (map_status != HTP_STATUS_OK) {
+        return map_status;
+    }

-        hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]);
-
-        memset(matrix_row_counts, 0, matrix_row_counts_size);
-        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);
-
-        uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as);
-        mapping_stride = max_count > 0 ? max_count : 1;
-
-        size_t matrix_row_map_size  = n_as * mapping_stride * sizeof(struct mmid_row_mapping);
-        const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
-
-        if (total_map_size > octx->ctx->ddr_spad_size) {
-            mapping_buf = memalign(128, total_map_size);
-            if (!mapping_buf) {
-                return HTP_STATUS_INTERNAL_ERR;
+    if (n_skipped > 0) {
+        for (uint32_t p = 0; p < n_weights; ++p) {
+            const struct htp_tensor * restrict d = octx->dsts[p];
+            if (d) {
+                zero_skipped_dst_rows(octx, ids, n_ids, d);
            }
        }
-
-        matrix_row_counts = (uint32_t *) mapping_buf;
-        matrix_rows       = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size);
-
-        memset(matrix_row_counts, 0, n_as * sizeof(uint32_t));
-        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
    }

    mmctx->matrix_row_counts    = matrix_row_counts;
@@ -173,11 +173,12 @@ static inline void hvx_ssm_conv_decode_4(const float * x, const float * w, float
 // the lanes past the channel tail.
 static inline void hvx_ssm_conv_transpose_block(const float * raw_block,
                                                float *       T,
+                                                uint32_t      raw_stride_bytes,
                                                uint32_t      ncs,
                                                uint32_t      cb_n,
                                                HVX_Vector    vv) {
    const size_t   base = (size_t) raw_block;
-    const uint32_t mu   = cb_n * ncs * sizeof(float) - 1;
+    const uint32_t mu   = cb_n * raw_stride_bytes - 1;

    for (uint32_t t = 0; t < ncs; ++t) {
        Q6_vgather_ARMVw((HVX_Vector *) (T + (size_t) t * VLEN_FP32), base + t * sizeof(float), mu, vv);
@@ -393,7 +394,8 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void

    // src0 spad holds two raw tiles (fetch of tile n+1 overlaps compute of tile n) plus
    // the transposed block. dst spad holds two tiles so a writeback can stay in flight.
-    const size_t src0_tile_raw_bytes = hex_round_up(d_inner_tile * ncs * sizeof(float), 128);
+    const size_t raw_stride_bytes    = ncs * sizeof(float) + sizeof(float);
+    const size_t src0_tile_raw_bytes = hex_round_up(d_inner_tile * raw_stride_bytes, 128);
    const size_t dst_tile_bytes      = hex_round_up(d_inner_tile * n_t * sizeof(float), 128);

    float * src0_tile_raw[2] = { (float *) src0_spad_base, (float *) (src0_spad_base + src0_tile_raw_bytes) };
@@ -419,7 +421,7 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
                                          (ir0 + tile_off_fetch) * src0_stride_inner_bytes;       \
            dma_queue_push(dma_q,                                                                 \
                           dma_make_data((uint8_t *) src0_tile_raw[(c) & 1], fetch_ddr),          \
-                           row_bytes, src0_stride_inner_bytes, row_bytes, cur_tile_n);            \
+                           raw_stride_bytes, src0_stride_inner_bytes, row_bytes, cur_tile_n);     \
            tile_off_fetch += d_inner_tile;                                                       \
            if (tile_off_fetch >= d_inner_per_thread) {                                           \
                tile_off_fetch = 0;                                                               \
@@ -446,11 +448,7 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
    const uint32_t C_TILE = VLEN_FP32;

    // gather offsets: lane r reads channel r of the raw tile block
-    uint32_t __attribute__((aligned(VLEN))) gather_off[VLEN_FP32];
-    for (uint32_t r = 0; r < VLEN_FP32; ++r) {
-        gather_off[r] = r * ncs * sizeof(float);
-    }
-    const HVX_Vector vv = *(const HVX_Vector *) gather_off;
+    const HVX_Vector vv = hvx_vec_gather_offsets_w(raw_stride_bytes);

    uint32_t i3       = 0;
    uint32_t tile_off = 0;
@@ -474,7 +472,8 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
            const uint32_t cb_n = MIN(C_TILE, tile_n - cb);

            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, (uint16_t) tile_off);
-            hvx_ssm_conv_transpose_block(raw + (size_t) cb * ncs, src0_T, ncs, cb_n, vv);
+            const float * raw_ch = (const float *) ((const uint8_t *) raw + (size_t) cb * raw_stride_bytes);
+            hvx_ssm_conv_transpose_block(raw_ch, src0_T, raw_stride_bytes, ncs, cb_n, vv);
            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, (uint16_t) tile_off);

            const float * restrict wp = src1_T + tile_off + cb;
@@ -54,11 +54,12 @@ static inline void htp_ssm_conv_vtcm_layout_build(
        L->src0_bytes_per_thread         = 2 * src0_tile_raw_bytes + src0_T_bytes;
        L->dst_bytes_per_thread          = 2 * hex_round_up((size_t) d_inner_tile * sizeof(float), 128);
    } else {
-        const uint32_t ncs          = d_conv - 1 + n_t;
-        const size_t   src0_block_T = hex_round_up((size_t) ncs * 32 * sizeof(float), 128);
-        const size_t   src0_tile_raw = hex_round_up((size_t) d_inner_tile * ncs * sizeof(float), 128);
-        L->src0_bytes_per_thread    = 2 * src0_tile_raw + src0_block_T;
-        L->dst_bytes_per_thread     = 2 * hex_round_up((size_t) d_inner_tile * n_t * sizeof(float), 128);
+        const uint32_t ncs              = d_conv - 1 + n_t;
+        const size_t   raw_stride_bytes = ncs * sizeof(float) + sizeof(float);
+        const size_t   src0_block_T     = hex_round_up((size_t) ncs * 32 * sizeof(float), 128);
+        const size_t   src0_tile_raw    = hex_round_up((size_t) d_inner_tile * raw_stride_bytes, 128);
+        L->src0_bytes_per_thread        = 2 * src0_tile_raw + src0_block_T;
+        L->dst_bytes_per_thread         = 2 * hex_round_up((size_t) d_inner_tile * n_t * sizeof(float), 128);
    }

    L->src0_bytes = L->src0_bytes_per_thread * n_threads;
@@ -0,0 +1,958 @@
+#pragma clang diagnostic ignored "-Wunused-variable"
+
+#include <float.h>
+#include <math.h>
+#include <HAP_farf.h>
+
+#include "hex-common.h"
+#include "dma-queue.h"
+#include "hex-profile.h"
+#include "htp-ctx.h"
+#include "htp-ops.h"
+#include "htp-tensor.h"
+#include "hvx-types.h"
+#include "hvx-base.h"
+#include "hvx-utils.h"
+#include "upscale-ops.h"
+
+static inline int32_t clamp_i32(int32_t v, int32_t lo, int32_t hi) {
+    return v < lo ? lo : (v > hi ? hi : v);
+}
+
+static inline float clamp_f32(float v, float lo, float hi) {
+    return v < lo ? lo : (v > hi ? hi : v);
+}
+
+static inline float bicubic_weight1(float x) {
+    const float a = -0.75f;
+    return ((a + 2.0f) * x - (a + 3.0f)) * x * x + 1.0f;
+}
+
+static inline float bicubic_weight2(float x) {
+    const float a = -0.75f;
+    return ((a * x - 5.0f * a) * x + 8.0f * a) * x - 4.0f * a;
+}
+
+static inline void bicubic_weights(float x, float w[4]) {
+    w[0] = bicubic_weight2(x + 1.0f);
+    w[1] = bicubic_weight1(x);
+    w[2] = bicubic_weight1(1.0f - x);
+    w[3] = bicubic_weight2(2.0f - x);
+}
+
+static void upscale_init_luts_nearest(
+    uint32_t * lut_x,
+    uint32_t * lut_y,
+    const struct htp_tensor * src0,
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p
+) {
+    const uint32_t src_x = src0->ne[0];
+    const uint32_t src_y = src0->ne[1];
+    const uint32_t dst_x = dst->ne[0];
+    const uint32_t dst_y = dst->ne[1];
+    const uint32_t dst_x_aligned = (dst_x + VLEN_FP32 - 1) & ~(VLEN_FP32 - 1);
+
+    for (uint32_t ox = 0; ox < dst_x_aligned; ++ox) {
+        if (ox < dst_x) {
+            uint32_t x0 = (uint32_t) ((float) ox * p->inv_sf0);
+            lut_x[ox] = MIN(x0, src_x - 1) * sizeof(float);
+        } else {
+            lut_x[ox] = 0;
+        }
+    }
+    for (uint32_t oy = 0; oy < dst_y; ++oy) {
+        uint32_t y0 = (uint32_t) ((float) oy * p->inv_sf1);
+        lut_y[oy] = MIN(y0, src_y - 1);
+    }
+}
+
+static void upscale_init_luts_bilinear(
+    struct upscale_bilinear_x_chunk * lut_x,
+    struct upscale_bilinear_y * lut_y,
+    const struct htp_tensor * src0,
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p
+) {
+    const uint32_t src_x = src0->ne[0];
+    const uint32_t src_y = src0->ne[1];
+    const uint32_t dst_x = dst->ne[0];
+    const uint32_t dst_y = dst->ne[1];
+    const uint32_t n_chunks = (dst_x + VLEN_FP32 - 1) / VLEN_FP32;
+
+    for (uint32_t c = 0; c < n_chunks; ++c) {
+        struct upscale_bilinear_x_chunk * chunk = &lut_x[c];
+        const uint32_t ox_start = c * VLEN_FP32;
+        for (uint32_t l = 0; l < VLEN_FP32; ++l) {
+            const uint32_t ox = ox_start + l;
+            if (ox < dst_x) {
+                float xf = ((float) ox + p->pixel_offset) * p->inv_sf0 - p->pixel_offset;
+                int32_t x0 = (int32_t) floorf(xf);
+                int32_t x1 = x0 + 1;
+                x0 = clamp_i32(x0, 0, (int32_t) src_x - 1);
+                x1 = clamp_i32(x1, 0, (int32_t) src_x - 1);
+                float dx = clamp_f32(xf - (float) x0, 0.0f, 1.0f);
+
+                chunk->off0[l] = x0 * (int32_t) sizeof(float);
+                chunk->off1[l] = x1 * (int32_t) sizeof(float);
+                chunk->w0[l]   = 1.0f - dx;
+                chunk->dx[l]   = dx;
+            } else {
+                chunk->off0[l] = 0;
+                chunk->off1[l] = 0;
+                chunk->w0[l]   = 0.0f;
+                chunk->dx[l]   = 0.0f;
+            }
+        }
+    }
+    for (uint32_t oy = 0; oy < dst_y; ++oy) {
+        float yf = ((float) oy + p->pixel_offset) * p->inv_sf1 - p->pixel_offset;
+        int32_t y0 = (int32_t) floorf(yf);
+        int32_t y1 = y0 + 1;
+        y0 = clamp_i32(y0, 0, (int32_t) src_y - 1);
+        y1 = clamp_i32(y1, 0, (int32_t) src_y - 1);
+        float dy = clamp_f32(yf - (float) y0, 0.0f, 1.0f);
+
+        lut_y[oy].y0  = (uint32_t) y0;
+        lut_y[oy].y1  = (uint32_t) y1;
+        lut_y[oy].wy1 = dy;
+        lut_y[oy].wy0 = 1.0f - dy;
+    }
+}
+
+static void upscale_init_luts_bicubic(
+    struct upscale_bicubic_x * lut_x,
+    struct upscale_bicubic_y * lut_y,
+    const struct htp_tensor * src0,
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p
+) {
+    const uint32_t src_x = src0->ne[0];
+    const uint32_t src_y = src0->ne[1];
+    const uint32_t dst_x = dst->ne[0];
+    const uint32_t dst_y = dst->ne[1];
+
+    for (uint32_t ox = 0; ox < dst_x; ++ox) {
+        float xf = ((float) ox + p->pixel_offset) * p->inv_sf0 - p->pixel_offset;
+        int32_t x0 = (int32_t) floorf(xf);
+        float dx = xf - (float) x0;
+        bicubic_weights(dx, lut_x[ox].w);
+
+        lut_x[ox].x[0] = clamp_i32(x0 - 1, 0, (int32_t) src_x - 1);
+        lut_x[ox].x[1] = clamp_i32(x0 + 0, 0, (int32_t) src_x - 1);
+        lut_x[ox].x[2] = clamp_i32(x0 + 1, 0, (int32_t) src_x - 1);
+        lut_x[ox].x[3] = clamp_i32(x0 + 2, 0, (int32_t) src_x - 1);
+    }
+    for (uint32_t oy = 0; oy < dst_y; ++oy) {
+        float yf = ((float) oy + p->pixel_offset) * p->inv_sf1 - p->pixel_offset;
+        int32_t y0 = (int32_t) floorf(yf);
+        float dy = yf - (float) y0;
+        bicubic_weights(dy, lut_y[oy].w);
+
+        lut_y[oy].y[0] = clamp_i32(y0 - 1, 0, (int32_t) src_y - 1);
+        lut_y[oy].y[1] = clamp_i32(y0 + 0, 0, (int32_t) src_y - 1);
+        lut_y[oy].y[2] = clamp_i32(y0 + 1, 0, (int32_t) src_y - 1);
+        lut_y[oy].y[3] = clamp_i32(y0 + 2, 0, (int32_t) src_y - 1);
+    }
+}
+
+static void upscale_init_luts_antialias(
+    struct upscale_antialias_x * lut_x,
+    struct upscale_antialias_y * lut_y,
+    const struct htp_tensor * src0,
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p
+) {
+    const uint32_t src_x = src0->ne[0];
+    const uint32_t src_y = src0->ne[1];
+    const uint32_t dst_x = dst->ne[0];
+    const uint32_t dst_y = dst->ne[1];
+
+    const float filter_radius0 = p->filter_radius0;
+    const float filter_radius1 = p->filter_radius1;
+
+    for (uint32_t ox = 0; ox < dst_x; ++ox) {
+        const float xf = ((float) ox + p->pixel_offset) * p->inv_sf0;
+        int32_t x_min = (int32_t) MAX(xf - filter_radius0 + p->pixel_offset, 0.0f);
+        int32_t x_max = (int32_t) MIN(xf + filter_radius0 + p->pixel_offset, (float) src_x);
+        lut_x[ox].x_min = clamp_i32(x_min, 0, (int32_t) src_x);
+        lut_x[ox].x_max = clamp_i32(x_max, 0, (int32_t) src_x);
+        lut_x[ox].xf    = xf;
+    }
+    for (uint32_t oy = 0; oy < dst_y; ++oy) {
+        const float yf = ((float) oy + p->pixel_offset) * p->inv_sf1;
+        int32_t y_min = (int32_t) MAX(yf - filter_radius1 + p->pixel_offset, 0.0f);
+        int32_t y_max = (int32_t) MIN(yf + filter_radius1 + p->pixel_offset, (float) src_y);
+        lut_y[oy].y_min = clamp_i32(y_min, 0, (int32_t) src_y);
+        lut_y[oy].y_max = clamp_i32(y_max, 0, (int32_t) src_y);
+        lut_y[oy].yf    = yf;
+    }
+}
+
+static inline void get_src_y_window(
+    const struct htp_tensor * src0,
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p,
+    const void * lut_y,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t * out_src_y_start,
+    uint32_t * out_src_y_count
+) {
+    const uint32_t src_y = src0->ne[1];
+    const uint32_t dst_y = dst->ne[1];
+
+    if (dst_y_count == dst_y) {
+        *out_src_y_start = 0;
+        *out_src_y_count = src_y;
+        return;
+    }
+
+    const uint32_t last_oy = dst_y_start + dst_y_count - 1;
+
+    if (p->mode == HTP_UPSCALE_MODE_NEAREST) {
+        const uint32_t * ly = (const uint32_t *) lut_y;
+        const uint32_t y0 = ly[dst_y_start];
+        const uint32_t y1 = ly[last_oy];
+        *out_src_y_start = y0;
+        *out_src_y_count = y1 - y0 + 1;
+    } else if (p->mode == HTP_UPSCALE_MODE_BICUBIC) {
+        const struct upscale_bicubic_y * ly = (const struct upscale_bicubic_y *) lut_y;
+        const int32_t y0 = ly[dst_y_start].y[0];
+        const int32_t y1 = ly[last_oy].y[3];
+        *out_src_y_start = (uint32_t) y0;
+        *out_src_y_count = (uint32_t) (y1 - y0 + 1);
+    } else if (p->antialias) {
+        const struct upscale_antialias_y * ly = (const struct upscale_antialias_y *) lut_y;
+        const int32_t y0 = ly[dst_y_start].y_min;
+        const int32_t y1 = ly[last_oy].y_max;
+        *out_src_y_start = (uint32_t) y0;
+        *out_src_y_count = (uint32_t) (y1 - y0);
+    } else {
+        const struct upscale_bilinear_y * ly = (const struct upscale_bilinear_y *) lut_y;
+        const uint32_t y0 = ly[dst_y_start].y0;
+        const uint32_t y1 = ly[last_oy].y1;
+        *out_src_y_start = y0;
+        *out_src_y_count = y1 - y0 + 1;
+    }
+}
+
+static inline void resolve_dst_slice(
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p,
+    uint32_t chunk_idx,
+    dma_addr_t * out_slice_dst_off,
+    uint32_t * out_dst_ys,
+    uint32_t * out_dst_yc
+) {
+    const uint32_t s_idx = fastdiv(chunk_idx, &p->div_chunks_per_slice);
+    const uint32_t c_idx = chunk_idx - s_idx * p->chunks_per_slice;
+    const uint32_t d_ys  = c_idx * p->dst_nrows_per_chunk;
+    const uint32_t d_yc  = MIN(p->dst_nrows_per_chunk, dst->ne[1] - d_ys);
+    const uint32_t i3    = fastdiv(s_idx, &p->div_dst_z);
+    const uint32_t i2    = s_idx - i3 * dst->ne[2];
+
+    *out_slice_dst_off = (dma_addr_t) i2 * dst->nb[2] + (dma_addr_t) i3 * dst->nb[3];
+    *out_dst_ys        = d_ys;
+    *out_dst_yc        = d_yc;
+}
+
+static inline void resolve_src_slice(
+    const struct htp_tensor * src0,
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p,
+    const void * lut_y,
+    uint32_t chunk_idx,
+    dma_addr_t * out_slice_src_off,
+    uint32_t * out_src_ys,
+    uint32_t * out_src_yc
+) {
+    const uint32_t s_idx = fastdiv(chunk_idx, &p->div_chunks_per_slice);
+    const uint32_t c_idx = chunk_idx - s_idx * p->chunks_per_slice;
+    const uint32_t d_ys  = c_idx * p->dst_nrows_per_chunk;
+    const uint32_t d_yc  = MIN(p->dst_nrows_per_chunk, dst->ne[1] - d_ys);
+    uint32_t s_ys, s_yc;
+    get_src_y_window(src0, dst, p, lut_y, d_ys, d_yc, &s_ys, &s_yc);
+
+    const uint32_t i3  = fastdiv(s_idx, &p->div_dst_z);
+    const uint32_t i2  = s_idx - i3 * dst->ne[2];
+    const uint32_t i03 = (uint32_t) ((float) i3 * p->inv_sf3);
+    const uint32_t i02 = (uint32_t) ((float) i2 * p->inv_sf2);
+
+    *out_slice_src_off = (dma_addr_t) i02 * src0->nb[2] + (dma_addr_t) i03 * src0->nb[3];
+    *out_src_ys        = s_ys;
+    *out_src_yc        = s_yc;
+}
+
+struct upscale_context;
+
+typedef void (*upscale_chunk_fn_t)(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+);
+
+struct upscale_context {
+    struct htp_ops_context * octx;
+    const struct htp_upscale_kernel_params * kparams;
+    struct htp_upscale_vtcm_layout layout;
+    upscale_chunk_fn_t compute_chunk;
+    const void * lut_x;
+    const void * lut_y;
+    uint32_t chunk_start;
+    uint32_t chunk_count;
+    uint32_t chunks_per_thread;
+};
+
+static void compute_nearest_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+    const uint32_t * lut_x = (const uint32_t *) ctx->lut_x;
+    const uint32_t * lut_y = (const uint32_t *) ctx->lut_y;
+    const uint32_t region_size = hex_round_up(src_x * sizeof(float), 128);
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const uint32_t y0 = lut_y[oy];
+        const float * r0  = srcb + (y0 - src_y_start) * src_x;
+        float * dst_row   = dstb + (oy - dst_y_start) * dst_x;
+
+        for (uint32_t ix = 0; ix < src_x; ix += VLEN_FP32) {
+            *(HVX_Vector *) (R_buf + ix) = hvx_vmemu(r0 + ix);
+        }
+
+        for (uint32_t ox = 0; ox < dst_x; ox += VLEN_FP32) {
+            const uint32_t rem = dst_x - ox;
+            const uint32_t nbytes = (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN;
+            const HVX_Vector voff = hvx_vmem(lut_x + ox);
+            Q6_vgather_ARMVw((HVX_Vector *) gather_tmp, (size_t) R_buf, region_size, voff);
+            hvx_gather_sync(gather_tmp);
+            hvx_vec_store_u(dst_row + ox, nbytes, hvx_vmem(gather_tmp));
+        }
+    }
+}
+
+static void compute_nearest_2x_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    (void) R_buf;
+    (void) gather_tmp;
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const uint32_t sy = oy >> 1;
+        const float * r0  = srcb + (sy - src_y_start) * src_x;
+        float * dst_row   = dstb + (oy - dst_y_start) * dst_x;
+
+        for (uint32_t ix = 0; ix < src_x; ix += VLEN_FP32) {
+            const HVX_Vector vin = hvx_vmemu(r0 + ix);
+
+            const HVX_VectorPair p = Q6_W_vshuff_VVR(vin, vin, -4);
+            const HVX_Vector v_lo = Q6_V_lo_W(p);
+            const HVX_Vector v_hi = Q6_V_hi_W(p);
+
+            const uint32_t ox0 = 2 * ix;
+            const uint32_t rem_out0 = (dst_x > ox0) ? (dst_x - ox0) : 0;
+            const uint32_t nbytes_out0 = (rem_out0 < VLEN_FP32) ? (rem_out0 * sizeof(float)) : VLEN;
+            hvx_vec_store_u(dst_row + ox0, nbytes_out0, v_lo);
+
+            const uint32_t ox1 = ox0 + VLEN_FP32;
+            if (dst_x > ox1) {
+                const uint32_t rem_out1 = dst_x - ox1;
+                const uint32_t nbytes_out1 = (rem_out1 < VLEN_FP32) ? (rem_out1 * sizeof(float)) : VLEN;
+                hvx_vec_store_u(dst_row + ox1, nbytes_out1, v_hi);
+            }
+        }
+    }
+}
+
+static void compute_nearest_4x_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    (void) R_buf;
+    (void) gather_tmp;
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const uint32_t sy = oy >> 2;
+        const float * r0  = srcb + (sy - src_y_start) * src_x;
+        float * dst_row   = dstb + (oy - dst_y_start) * dst_x;
+
+        for (uint32_t ix = 0; ix < src_x; ix += VLEN_FP32) {
+            const HVX_Vector vin = hvx_vmemu(r0 + ix);
+
+            const HVX_VectorPair p = Q6_W_vshuff_VVR(vin, vin, -4);
+            const HVX_Vector v_lo = Q6_V_lo_W(p);
+            const HVX_Vector v_hi = Q6_V_hi_W(p);
+
+            const HVX_VectorPair p0 = Q6_W_vshuff_VVR(v_lo, v_lo, -4);
+            const HVX_VectorPair p1 = Q6_W_vshuff_VVR(v_hi, v_hi, -4);
+
+            const HVX_Vector out0 = Q6_V_lo_W(p0);
+            const HVX_Vector out1 = Q6_V_hi_W(p0);
+            const HVX_Vector out2 = Q6_V_lo_W(p1);
+            const HVX_Vector out3 = Q6_V_hi_W(p1);
+
+            const uint32_t ox = 4 * ix;
+            if (dst_x > ox) {
+                const uint32_t rem = dst_x - ox;
+                hvx_vec_store_u(dst_row + ox, (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN, out0);
+            }
+            if (dst_x > ox + VLEN_FP32) {
+                const uint32_t rem = dst_x - (ox + VLEN_FP32);
+                hvx_vec_store_u(dst_row + ox + VLEN_FP32, (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN, out1);
+            }
+            if (dst_x > ox + 2 * VLEN_FP32) {
+                const uint32_t rem = dst_x - (ox + 2 * VLEN_FP32);
+                hvx_vec_store_u(dst_row + ox + 2 * VLEN_FP32, (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN, out2);
+            }
+            if (dst_x > ox + 3 * VLEN_FP32) {
+                const uint32_t rem = dst_x - (ox + 3 * VLEN_FP32);
+                hvx_vec_store_u(dst_row + ox + 3 * VLEN_FP32, (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN, out3);
+            }
+        }
+    }
+}
+
+static inline void compute_bilinear_chunk_32(
+    float * restrict dst_chunk,
+    const float * restrict R_buf,
+    float * restrict tmp_a,
+    float * restrict tmp_b,
+    uint32_t region_size,
+    const struct upscale_bilinear_x_chunk * restrict chunk,
+    uint32_t nbytes
+) {
+    const HVX_Vector voff0 = hvx_vmem(chunk->off0);
+    const HVX_Vector voff1 = hvx_vmem(chunk->off1);
+    const HVX_Vector vw0   = hvx_vmem(chunk->w0);
+    const HVX_Vector vdx   = hvx_vmem(chunk->dx);
+
+    Q6_vgather_ARMVw((HVX_Vector *) tmp_a, (size_t) R_buf, region_size, voff0);
+    Q6_vgather_ARMVw((HVX_Vector *) tmp_b, (size_t) R_buf, region_size, voff1);
+    hvx_gather_sync(tmp_b);
+
+    const HVX_Vector va = hvx_vmem(tmp_a);
+    const HVX_Vector vb = hvx_vmem(tmp_b);
+
+    const HVX_Vector val = hvx_vec_add_f32_f32(
+        hvx_vec_mul_f32_f32(vw0, va),
+        hvx_vec_mul_f32_f32(vdx, vb)
+    );
+
+    hvx_vec_store_u(dst_chunk, nbytes, val);
+}
+
+static void compute_bilinear_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+    const struct upscale_bilinear_x_chunk * lut_x = (const struct upscale_bilinear_x_chunk *) ctx->lut_x;
+    const struct upscale_bilinear_y * lut_y = (const struct upscale_bilinear_y *) ctx->lut_y;
+    const uint32_t n_chunks = (dst_x + VLEN_FP32 - 1) / VLEN_FP32;
+    const uint32_t region_size = hex_round_up(src_x * sizeof(float), 128);
+    float * tmp_a = gather_tmp;
+    float * tmp_b = gather_tmp + VLEN_FP32;
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const struct upscale_bilinear_y * ey = &lut_y[oy];
+        const float * r0 = srcb + (ey->y0 - src_y_start) * src_x;
+        const float * r1 = srcb + (ey->y1 - src_y_start) * src_x;
+        float * dst_row  = dstb + (oy - dst_y_start) * dst_x;
+
+        if (ey->y0 == ey->y1) {
+            for (uint32_t ix = 0; ix < src_x; ix += VLEN_FP32) {
+                *(HVX_Vector *) (R_buf + ix) = hvx_vmemu(r0 + ix);
+            }
+        } else {
+            const HVX_Vector vwy0 = hvx_vec_splat_f32(ey->wy0);
+            const HVX_Vector vwy1 = hvx_vec_splat_f32(ey->wy1);
+            for (uint32_t ix = 0; ix < src_x; ix += VLEN_FP32) {
+                const HVX_Vector vr0 = hvx_vmemu(r0 + ix);
+                const HVX_Vector vr1 = hvx_vmemu(r1 + ix);
+                *(HVX_Vector *) (R_buf + ix) = hvx_vec_add_f32_f32(
+                    hvx_vec_mul_f32_f32(vwy0, vr0),
+                    hvx_vec_mul_f32_f32(vwy1, vr1)
+                );
+            }
+        }
+
+        for (uint32_t c = 0; c < n_chunks; ++c) {
+            const uint32_t ox = c * VLEN_FP32;
+            const uint32_t rem = dst_x - ox;
+            const uint32_t nbytes = (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN;
+            compute_bilinear_chunk_32(dst_row + ox, R_buf, tmp_a, tmp_b, region_size, &lut_x[c], nbytes);
+        }
+    }
+}
+
+static void compute_bilinear_2x_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    (void) R_buf;
+    (void) gather_tmp;
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+    const struct upscale_bilinear_y * lut_y = (const struct upscale_bilinear_y *) ctx->lut_y;
+    const HVX_Vector c_075 = hvx_vec_splat_f32(0.75f);
+    const HVX_Vector c_025 = hvx_vec_splat_f32(0.25f);
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const struct upscale_bilinear_y * ey = &lut_y[oy];
+        const float * r0 = srcb + (ey->y0 - src_y_start) * src_x;
+        const float * r1 = srcb + (ey->y1 - src_y_start) * src_x;
+        float * dst_row  = dstb + (oy - dst_y_start) * dst_x;
+
+        const bool same_row = (ey->y0 == ey->y1);
+        const HVX_Vector vwy0 = hvx_vec_splat_f32(ey->wy0);
+        const HVX_Vector vwy1 = hvx_vec_splat_f32(ey->wy1);
+
+        HVX_Vector v_curr = hvx_vmem(r0);
+        if (!same_row) {
+            const HVX_Vector v1 = hvx_vmem(r1);
+            v_curr = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(vwy0, v_curr), hvx_vec_mul_f32_f32(vwy1, v1));
+        }
+
+        HVX_Vector v_prev = Q6_V_vror_VR(v_curr, 4);
+
+        for (uint32_t ix = 0; ix < src_x; ix += VLEN_FP32) {
+            HVX_Vector v_next_chunk;
+            const uint32_t next_ix = ix + VLEN_FP32;
+            if (next_ix < src_x) {
+                v_next_chunk = hvx_vmem(r0 + next_ix);
+                if (!same_row) {
+                    const HVX_Vector v1 = hvx_vmem(r1 + next_ix);
+                    v_next_chunk = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(vwy0, v_next_chunk), hvx_vec_mul_f32_f32(vwy1, v1));
+                }
+            } else {
+                v_next_chunk = Q6_V_vror_VR(v_curr, 124);
+            }
+
+            const HVX_Vector v_shl = Q6_V_valign_VVR(v_curr, v_prev, 124);
+            const HVX_Vector v_shr = Q6_V_valign_VVR(v_next_chunk, v_curr, 4);
+
+            const HVX_Vector ve = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(c_075, v_curr), hvx_vec_mul_f32_f32(c_025, v_shl));
+            const HVX_Vector vo = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(c_075, v_curr), hvx_vec_mul_f32_f32(c_025, v_shr));
+
+            const HVX_VectorPair p = Q6_W_vshuff_VVR(vo, ve, -4);
+            const HVX_Vector out0 = Q6_V_lo_W(p);
+            const HVX_Vector out1 = Q6_V_hi_W(p);
+
+            const uint32_t ox0 = 2 * ix;
+            hvx_vmem(dst_row + ox0)             = out0;
+            hvx_vmem(dst_row + ox0 + VLEN_FP32) = out1;
+
+            v_prev = v_curr;
+            v_curr = v_next_chunk;
+        }
+    }
+}
+
+static void compute_bilinear_antialias_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    (void) R_buf;
+    (void) gather_tmp;
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+    const struct htp_upscale_kernel_params * p = ctx->kparams;
+    const struct upscale_antialias_x * lut_x = (const struct upscale_antialias_x *) ctx->lut_x;
+    const struct upscale_antialias_y * lut_y = (const struct upscale_antialias_y *) ctx->lut_y;
+
+    const float invscale1 = p->invscale1;
+    const float invscale0 = p->invscale0;
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const struct upscale_antialias_y * ey = &lut_y[oy];
+        const float yf = ey->yf;
+        const int32_t y_min = ey->y_min;
+        const int32_t y_max = ey->y_max;
+
+        float * dst_row = dstb + (oy - dst_y_start) * dst_x;
+
+        for (uint32_t ox = 0; ox < dst_x; ox += VLEN_FP32) {
+            const uint32_t rem = dst_x - ox;
+            const uint32_t nbytes = (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN;
+            float tmp[VLEN_FP32] __attribute__((aligned(128)));
+            const uint32_t cnt = rem < VLEN_FP32 ? rem : VLEN_FP32;
+
+            for (uint32_t l = 0; l < cnt; ++l) {
+                const uint32_t cur_ox = ox + l;
+                const struct upscale_antialias_x * ex = &lut_x[cur_ox];
+                const float xf = ex->xf;
+                const int32_t x_min = ex->x_min;
+                const int32_t x_max = ex->x_max;
+
+                float val = 0.0f;
+                float total_weight = 0.0f;
+
+                for (int32_t sy = y_min; sy < y_max; ++sy) {
+                    const float diff_y = fabsf((float) sy - yf + p->pixel_offset) * invscale1;
+                    const float weight_y = MAX(1.0f - diff_y, 0.0f);
+                    if (weight_y <= 0.0f) {
+                        continue;
+                    }
+                    const float * r = srcb + (uint32_t) (sy - (int32_t) src_y_start) * src_x;
+
+                    for (int32_t sx = x_min; sx < x_max; ++sx) {
+                        const float diff_x = fabsf((float) sx - xf + p->pixel_offset) * invscale0;
+                        const float weight_x = MAX(1.0f - diff_x, 0.0f);
+                        const float weight = weight_x * weight_y;
+                        if (weight <= 0.0f) {
+                            continue;
+                        }
+                        val += r[sx] * weight;
+                        total_weight += weight;
+                    }
+                }
+
+                tmp[l] = (total_weight > 0.0f) ? (val / total_weight) : 0.0f;
+            }
+
+            hvx_vec_store_u(dst_row + ox, nbytes, *(const HVX_Vector *) tmp);
+        }
+    }
+}
+
+static void compute_bicubic_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    (void) R_buf;
+    (void) gather_tmp;
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+    const struct upscale_bicubic_x * lut_x = (const struct upscale_bicubic_x *) ctx->lut_x;
+    const struct upscale_bicubic_y * lut_y = (const struct upscale_bicubic_y *) ctx->lut_y;
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const struct upscale_bicubic_y * ey = &lut_y[oy];
+        const float * r0 = srcb + (uint32_t) (ey->y[0] - (int32_t) src_y_start) * src_x;
+        const float * r1 = srcb + (uint32_t) (ey->y[1] - (int32_t) src_y_start) * src_x;
+        const float * r2 = srcb + (uint32_t) (ey->y[2] - (int32_t) src_y_start) * src_x;
+        const float * r3 = srcb + (uint32_t) (ey->y[3] - (int32_t) src_y_start) * src_x;
+
+        float * dst_row = dstb + (oy - dst_y_start) * dst_x;
+
+        const HVX_Vector vwy0 = hvx_vec_splat_f32(ey->w[0]);
+        const HVX_Vector vwy1 = hvx_vec_splat_f32(ey->w[1]);
+        const HVX_Vector vwy2 = hvx_vec_splat_f32(ey->w[2]);
+        const HVX_Vector vwy3 = hvx_vec_splat_f32(ey->w[3]);
+
+        for (uint32_t ox = 0; ox < dst_x; ox += VLEN_FP32) {
+            const uint32_t rem = dst_x - ox;
+            const uint32_t nbytes = (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN;
+
+            float v0[VLEN_FP32] __attribute__((aligned(128)));
+            float v1[VLEN_FP32] __attribute__((aligned(128)));
+            float v2[VLEN_FP32] __attribute__((aligned(128)));
+            float v3[VLEN_FP32] __attribute__((aligned(128)));
+
+            const uint32_t cnt = rem < VLEN_FP32 ? rem : VLEN_FP32;
+            for (uint32_t l = 0; l < cnt; ++l) {
+                const struct upscale_bicubic_x * e = &lut_x[ox + l];
+                v0[l] = r0[e->x[0]] * e->w[0] + r0[e->x[1]] * e->w[1] + r0[e->x[2]] * e->w[2] + r0[e->x[3]] * e->w[3];
+                v1[l] = r1[e->x[0]] * e->w[0] + r1[e->x[1]] * e->w[1] + r1[e->x[2]] * e->w[2] + r1[e->x[3]] * e->w[3];
+                v2[l] = r2[e->x[0]] * e->w[0] + r2[e->x[1]] * e->w[1] + r2[e->x[2]] * e->w[2] + r2[e->x[3]] * e->w[3];
+                v3[l] = r3[e->x[0]] * e->w[0] + r3[e->x[1]] * e->w[1] + r3[e->x[2]] * e->w[2] + r3[e->x[3]] * e->w[3];
+            }
+            for (uint32_t l = cnt; l < VLEN_FP32; ++l) {
+                v0[l] = v1[l] = v2[l] = v3[l] = 0.0f;
+            }
+
+            const HVX_Vector vv0 = *(const HVX_Vector *) v0;
+            const HVX_Vector vv1 = *(const HVX_Vector *) v1;
+            const HVX_Vector vv2 = *(const HVX_Vector *) v2;
+            const HVX_Vector vv3 = *(const HVX_Vector *) v3;
+
+            const HVX_Vector acc0 = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(vwy0, vv0), hvx_vec_mul_f32_f32(vwy1, vv1));
+            const HVX_Vector acc1 = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(vwy2, vv2), hvx_vec_mul_f32_f32(vwy3, vv3));
+            const HVX_Vector val  = hvx_vec_add_f32_f32(acc0, acc1);
+
+            hvx_vec_store_u(dst_row + ox, nbytes, val);
+        }
+    }
+}
+
+static inline void dma_push_src_chunk(
+    dma_queue * dma_q,
+    float * spad,
+    const struct htp_tensor * src0,
+    dma_addr_t slice_src_off,
+    uint32_t src_ys,
+    uint32_t src_yc
+) {
+    const uint32_t src_x = src0->ne[0];
+    const dma_addr_t src_addr = src0->data + slice_src_off + (dma_addr_t) src_ys * src0->nb[1];
+
+    if (src0->nb[0] == sizeof(float)) {
+        const uint32_t row_size   = src_x * sizeof(float);
+        const uint32_t dst_stride = row_size;
+        const uint32_t src_stride = src0->nb[1];
+        dma_queue_push(dma_q, dma_make_data(spad, src_addr), dst_stride, src_stride, row_size, src_yc);
+    } else {
+        const uint32_t row_size   = sizeof(float);
+        const uint32_t dst_stride = sizeof(float);
+        const uint32_t src_stride = src0->nb[0];
+        const uint32_t nrows      = src_x * src_yc;
+        dma_queue_push(dma_q, dma_make_data(spad, src_addr), dst_stride, src_stride, row_size, nrows);
+    }
+}
+
+static void upscale_worker(unsigned int nth, unsigned int ith, void * data) {
+    struct upscale_context * ctx = (struct upscale_context *) data;
+    const struct htp_upscale_kernel_params * p = ctx->kparams;
+    const struct htp_upscale_vtcm_layout * L = &ctx->layout;
+    const struct htp_tensor * src0 = ctx->octx->src[0];
+    const struct htp_tensor * dst  = ctx->octx->dst;
+
+    const uint32_t first = ctx->chunk_start + ith * ctx->chunks_per_thread;
+    const uint32_t last  = MIN(first + ctx->chunks_per_thread, ctx->chunk_start + ctx->chunk_count);
+
+    if (first >= last) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &ctx->octx->ctx->trace[ith];
+    dma_queue * dma_q            = ctx->octx->ctx->dma[ith];
+
+    const uint32_t src_spad_half = (uint32_t) L->src_spad_half_size;
+    const uint32_t dst_spad_half = (uint32_t) L->dst_spad_half_size;
+
+    uint8_t * vtcm_base  = (uint8_t *) ctx->octx->ctx->vtcm_base;
+    uint8_t * src_spad   = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L->off_src)    + ith * L->src_bytes_per_thread;
+    uint8_t * dst_spad   = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L->off_dst)    + ith * L->dst_bytes_per_thread;
+    float   * R_buf      = VTCM_LAYOUT_PTR(float,   vtcm_base, L->off_r)      + ith * (L->r_bytes_per_thread / sizeof(float));
+    float   * gather_tmp = VTCM_LAYOUT_PTR(float,   vtcm_base, L->off_gather) + ith * (L->gather_bytes_per_thread / sizeof(float));
+    const void * lut_y   = ctx->lut_y;
+
+    float * srcb2[2] = { (float *) src_spad, (float *) (src_spad + src_spad_half) };
+    float * dstb2[2] = { (float *) dst_spad, (float *) (dst_spad + dst_spad_half) };
+
+    const uint32_t dst_row_size = dst->ne[0] * sizeof(float);
+    const uint32_t dst_stride   = dst->nb[1];
+
+    const uint32_t total = last - first;
+
+    // Warmup: prime pipeline with initial transfers
+    for (uint32_t i = 0; i < total && i < 2; ++i) {
+        dma_queue_push(dma_q, dma_make_data(dst->data, dstb2[i]), dst_stride, dst_row_size, dst_row_size, 0);
+
+        dma_addr_t slice_src_off;
+        uint32_t src_ys, src_yc;
+        resolve_src_slice(src0, dst, p, lut_y, first + i, &slice_src_off, &src_ys, &src_yc);
+
+        dma_push_src_chunk(dma_q, srcb2[i], src0, slice_src_off, src_ys, src_yc);
+    }
+
+    // Steady state loop
+    for (uint32_t i = 0; i < total; ++i) {
+        const uint32_t buf = i & 1u;
+        float * srcb = srcb2[buf];
+        float * dstb = dstb2[buf];
+
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+
+        dma_addr_t slice_dst_off;
+        uint32_t dst_ys, dst_yc;
+        resolve_dst_slice(dst, p, first + i, &slice_dst_off, &dst_ys, &dst_yc);
+
+        uint32_t src_ys, src_yc;
+        get_src_y_window(src0, dst, p, lut_y, dst_ys, dst_yc, &src_ys, &src_yc);
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (first + i));
+        ctx->compute_chunk(ctx, srcb, dstb, R_buf, gather_tmp, dst_ys, dst_yc, src_ys);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (first + i));
+
+        const dma_addr_t dst_addr = dst->data + slice_dst_off + (dma_addr_t) dst_ys * dst_stride;
+        dma_queue_push(dma_q, dma_make_data(dst_addr, dstb), dst_stride, dst_row_size, dst_row_size, dst_yc);
+
+        const uint32_t next = i + 2;
+        if (next < total) {
+            dma_addr_t next_src_off;
+            uint32_t next_src_ys, next_src_yc;
+            resolve_src_slice(src0, dst, p, lut_y, first + next, &next_src_off, &next_src_ys, &next_src_yc);
+
+            dma_push_src_chunk(dma_q, srcb, src0, next_src_off, next_src_ys, next_src_yc);
+        }
+    }
+
+    dma_queue_flush(dma_q);
+
+    (void) nth;
+}
+
+int op_upscale(struct htp_ops_context * octx) {
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;
+    const struct htp_upscale_kernel_params * p =
+        (const struct htp_upscale_kernel_params *) octx->kernel_params;
+
+    const uint32_t src_x = src0->ne[0];
+    const uint32_t src_y = src0->ne[1];
+    const uint32_t dst_x = dst->ne[0];
+    const uint32_t dst_y = dst->ne[1];
+
+    uint32_t chunk_start = 0;
+    uint32_t chunk_count = p->total_chunks;
+
+    if (octx->ctx->mdev.count > 1) {
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+            chunk_count, htp_tensor_mdev_data_aligned(dst) ? p->chunks_per_group : 0,
+            octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        chunk_start = range.start;
+        chunk_count = range.count;
+    }
+
+    if (chunk_count == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = MIN(p->n_threads, chunk_count);
+    if (!htp_ops_context_set_n_threads(octx, n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    const uint32_t chunks_per_thread = fastdiv(chunk_count + n_threads - 1, &octx->n_threads_div);
+
+    size_t lut_x_size = 0;
+    size_t elem_y_sz  = sizeof(uint32_t);
+
+    if (p->mode == HTP_UPSCALE_MODE_BICUBIC) {
+        lut_x_size = hex_round_up(dst_x * sizeof(struct upscale_bicubic_x), 128);
+        elem_y_sz  = sizeof(struct upscale_bicubic_y);
+    } else if (p->antialias) {
+        lut_x_size = hex_round_up(dst_x * sizeof(struct upscale_antialias_x), 128);
+        elem_y_sz  = sizeof(struct upscale_antialias_y);
+    } else if (p->mode == HTP_UPSCALE_MODE_BILINEAR) {
+        const uint32_t n_chunks = (dst_x + VLEN_FP32 - 1) / VLEN_FP32;
+        lut_x_size = hex_round_up(n_chunks * sizeof(struct upscale_bilinear_x_chunk), 128);
+        elem_y_sz  = sizeof(struct upscale_bilinear_y);
+    } else {
+        lut_x_size = hex_round_up(dst_x * sizeof(uint32_t), 128);
+    }
+
+    const size_t lut_y_size = hex_round_up(dst_y * elem_y_sz, 128);
+    const size_t total_lut_size = lut_x_size + lut_y_size;
+
+    if (total_lut_size > octx->ctx->ddr_spad_size) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    void * base_scratch = octx->ctx->ddr_spad_base;
+    void * lut_x        = base_scratch;
+    void * lut_y        = (uint8_t *) base_scratch + lut_x_size;
+
+    if (p->mode == HTP_UPSCALE_MODE_NEAREST) {
+        upscale_init_luts_nearest((uint32_t *) lut_x, (uint32_t *) lut_y, src0, dst, p);
+    } else if (p->mode == HTP_UPSCALE_MODE_BICUBIC) {
+        upscale_init_luts_bicubic((struct upscale_bicubic_x *) lut_x, (struct upscale_bicubic_y *) lut_y, src0, dst, p);
+    } else if (p->antialias) {
+        upscale_init_luts_antialias((struct upscale_antialias_x *) lut_x, (struct upscale_antialias_y *) lut_y, src0, dst, p);
+    } else {
+        upscale_init_luts_bilinear((struct upscale_bilinear_x_chunk *) lut_x, (struct upscale_bilinear_y *) lut_y, src0, dst, p);
+    }
+
+    upscale_chunk_fn_t compute_chunk;
+    switch (p->kernel_type) {
+        case HTP_UPSCALE_KERNEL_NEAREST:
+            compute_chunk = compute_nearest_chunk;
+            break;
+        case HTP_UPSCALE_KERNEL_NEAREST_2X:
+            compute_chunk = compute_nearest_2x_chunk;
+            break;
+        case HTP_UPSCALE_KERNEL_NEAREST_4X:
+            compute_chunk = compute_nearest_4x_chunk;
+            break;
+        case HTP_UPSCALE_KERNEL_BILINEAR:
+            compute_chunk = compute_bilinear_chunk;
+            break;
+        case HTP_UPSCALE_KERNEL_BILINEAR_2X:
+            compute_chunk = compute_bilinear_2x_chunk;
+            break;
+        case HTP_UPSCALE_KERNEL_BILINEAR_AA:
+            compute_chunk = compute_bilinear_antialias_chunk;
+            break;
+        case HTP_UPSCALE_KERNEL_BICUBIC:
+            compute_chunk = compute_bicubic_chunk;
+            break;
+        default:
+            return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    struct upscale_context ctx = {
+        .octx               = octx,
+        .kparams            = p,
+        .compute_chunk      = compute_chunk,
+        .lut_x              = lut_x,
+        .lut_y              = lut_y,
+        .chunk_start        = chunk_start,
+        .chunk_count        = chunk_count,
+        .chunks_per_thread  = chunks_per_thread,
+    };
+
+    htp_upscale_vtcm_layout_build(&ctx.layout, src_x, src_y, dst_x, dst_y,
+                                  p->dst_nrows_per_chunk, p->inv_sf1, p->filter_radius1, n_threads);
+
+    work_queue_run(octx->ctx->work_queue, upscale_worker, &ctx, n_threads);
+
+    return HTP_STATUS_OK;
+}
@@ -0,0 +1,204 @@
+#ifndef HTP_UPSCALE_OPS_H
+#define HTP_UPSCALE_OPS_H
+
+#include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
+#include <math.h>
+
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
+
+enum htp_upscale_mode {
+    HTP_UPSCALE_MODE_NEAREST  = 0,
+    HTP_UPSCALE_MODE_BILINEAR = 1,
+    HTP_UPSCALE_MODE_BICUBIC  = 2,
+};
+
+enum htp_upscale_kernel_type {
+    HTP_UPSCALE_KERNEL_NEAREST     = 0,
+    HTP_UPSCALE_KERNEL_NEAREST_2X  = 1,
+    HTP_UPSCALE_KERNEL_NEAREST_4X  = 2,
+    HTP_UPSCALE_KERNEL_BILINEAR    = 3,
+    HTP_UPSCALE_KERNEL_BILINEAR_2X = 4,
+    HTP_UPSCALE_KERNEL_BILINEAR_AA = 5,
+    HTP_UPSCALE_KERNEL_BICUBIC     = 6,
+};
+
+struct htp_upscale_kernel_params {
+    uint8_t  kernel_type;
+    uint8_t  mode;
+    uint8_t  align_corners;
+    uint8_t  antialias;
+
+    float    inv_sf0;
+    float    inv_sf1;
+    float    inv_sf2;
+    float    inv_sf3;
+    float    pixel_offset;
+
+    float    filter_radius0;
+    float    filter_radius1;
+    float    invscale0;
+    float    invscale1;
+
+    uint32_t n_threads;
+    uint32_t chunks_per_group;
+
+    uint32_t dst_nrows_per_chunk;
+    uint32_t chunks_per_slice;
+    uint32_t total_chunks;
+    uint32_t vtcm_size;
+
+    struct fastdiv_values div_chunks_per_slice;
+    struct fastdiv_values div_dst_z;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_upscale_kernel_params) <= 128, "htp_upscale_kernel_params is too large");
+#else
+_Static_assert(sizeof(struct htp_upscale_kernel_params) <= 128, "htp_upscale_kernel_params is too large");
+#endif
+
+struct __attribute__((aligned(128))) upscale_bilinear_x_chunk {
+    int32_t  off0[32];
+    int32_t  off1[32];
+    float    w0[32];
+    float    dx[32];
+};
+
+struct upscale_bilinear_y {
+    uint32_t y0;
+    uint32_t y1;
+    float    wy0;
+    float    wy1;
+};
+
+struct upscale_bicubic_x {
+    int32_t x[4];
+    float   w[4];
+};
+
+struct upscale_bicubic_y {
+    int32_t y[4];
+    float   w[4];
+};
+
+struct upscale_antialias_x {
+    int32_t x_min;
+    int32_t x_max;
+    float   xf;
+};
+
+struct upscale_antialias_y {
+    int32_t y_min;
+    int32_t y_max;
+    float   yf;
+};
+
+struct htp_upscale_vtcm_layout {
+    size_t   total_bytes;
+    size_t   off_src;
+    size_t   off_dst;
+    size_t   off_r;
+    size_t   off_gather;
+    size_t   src_bytes_per_thread;
+    size_t   dst_bytes_per_thread;
+    size_t   r_bytes_per_thread;
+    size_t   gather_bytes_per_thread;
+    size_t   src_spad_half_size;
+    size_t   dst_spad_half_size;
+    uint32_t max_src_chunk_rows;
+    uint32_t dst_chunk_rows;
+};
+
+static inline void htp_upscale_vtcm_layout_build(
+    struct htp_upscale_vtcm_layout * L,
+    uint32_t src_x,
+    uint32_t src_y,
+    uint32_t dst_x,
+    uint32_t dst_y,
+    uint32_t dst_chunk_rows,
+    float    inv_sf1,
+    float    filter_radius1,
+    uint32_t n_threads
+) {
+    uint32_t max_src_rows;
+    if (dst_chunk_rows >= dst_y) {
+        max_src_rows = src_y;
+    } else {
+        const float est = (float) dst_chunk_rows * inv_sf1 + 2.0f * filter_radius1;
+        max_src_rows = (uint32_t) ceilf(est) + 4;
+        if (max_src_rows > src_y) {
+            max_src_rows = src_y;
+        }
+    }
+
+    const uint32_t src_chunk_bytes = src_x * max_src_rows * sizeof(float);
+    const uint32_t dst_chunk_bytes = dst_x * dst_chunk_rows * sizeof(float);
+
+    L->max_src_chunk_rows      = max_src_rows;
+    L->dst_chunk_rows          = dst_chunk_rows;
+    L->src_spad_half_size      = hex_round_up(src_chunk_bytes + 256, 128);
+    L->dst_spad_half_size      = hex_round_up(dst_chunk_bytes, 128);
+    L->src_bytes_per_thread    = 2 * L->src_spad_half_size;
+    L->dst_bytes_per_thread    = 2 * L->dst_spad_half_size;
+    L->r_bytes_per_thread      = hex_round_up(src_x * sizeof(float), 128);
+    L->gather_bytes_per_thread = 256;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src,    L->src_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,    L->dst_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_r,      L->r_bytes_per_thread   * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_gather, L->gather_bytes_per_thread * n_threads);
+    L->total_bytes = off;
+}
+
+static inline bool htp_upscale_solve_layout(
+    struct htp_upscale_vtcm_layout * layout,
+    uint32_t src_x,
+    uint32_t src_y,
+    uint32_t dst_x,
+    uint32_t dst_y,
+    float    inv_sf1,
+    float    filter_radius1,
+    uint32_t n_threads,
+    size_t   vtcm_budget,
+    uint32_t * out_dst_chunk_rows
+) {
+    // 1. Try full plane
+    uint32_t chunk_rows = dst_y;
+    htp_upscale_vtcm_layout_build(layout, src_x, src_y, dst_x, dst_y, chunk_rows, inv_sf1, filter_radius1, n_threads);
+    if (layout->total_bytes <= vtcm_budget) {
+        *out_dst_chunk_rows = chunk_rows;
+        return true;
+    }
+
+    // 2. Try smaller chunk bands: 64, 32, 16, 8, 4, 2, 1
+    chunk_rows = 64;
+    while (chunk_rows > 1) {
+        if (chunk_rows < dst_y) {
+            htp_upscale_vtcm_layout_build(layout, src_x, src_y, dst_x, dst_y, chunk_rows, inv_sf1, filter_radius1, n_threads);
+            if (layout->total_bytes <= vtcm_budget) {
+                *out_dst_chunk_rows = chunk_rows;
+                return true;
+            }
+        }
+        chunk_rows /= 2;
+    }
+
+    // 3. Fallback to 1 row
+    chunk_rows = 1;
+    htp_upscale_vtcm_layout_build(layout, src_x, src_y, dst_x, dst_y, chunk_rows, inv_sf1, filter_radius1, n_threads);
+    if (layout->total_bytes <= vtcm_budget) {
+        *out_dst_chunk_rows = chunk_rows;
+        return true;
+    }
+
+    return false;
+}
+
+int op_upscale(struct htp_ops_context * octx);
+
+#endif // HTP_UPSCALE_OPS_H
@@ -507,6 +507,47 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexe
    return res;
 }

+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_mask(
+        ggml_metal_library_t lib,
+        const ggml_tensor * op,
+        bool    fill,
+        bool    has_score,
+        bool    has_blocks,
+        bool    has_extra,
+        bool    vec) {
+    GGML_ASSERT(op->op == GGML_OP_TOP_K_MASK);
+
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_top_k_mask%s_%s", fill ? "_fill" : "", ggml_type_name(op->type));
+
+    ggml_metal_pipeline_with_params res;
+
+    ggml_metal_cv_t cv = ggml_metal_cv_init();
+
+    if (fill) {
+        snprintf(name, 256, "%s_vec=%d", base, vec ? 1 : 0);
+
+        ggml_metal_cv_set_bool(cv, vec, FC_TOP_K_MASK + 3);
+    } else {
+        snprintf(name, 256, "%s_sc=%d_bl=%d_ex=%d", base, has_score ? 1 : 0, has_blocks ? 1 : 0, has_extra ? 1 : 0);
+
+        ggml_metal_cv_set_bool(cv, has_score,  FC_TOP_K_MASK + 0);
+        ggml_metal_cv_set_bool(cv, has_blocks, FC_TOP_K_MASK + 1);
+        ggml_metal_cv_set_bool(cv, has_extra,  FC_TOP_K_MASK + 2);
+    }
+
+    res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
+    }
+
+    ggml_metal_cv_free(cv);
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc(ggml_metal_library_t lib, const ggml_tensor * op) {
    char name[256];
    const char * base = nullptr;
@@ -1062,7 +1103,6 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
            {
                nsg = N_SG_MXFP4;
                nr0 = N_R0_MXFP4;
-                smem = 32*sizeof(float);
            } break;
        case GGML_TYPE_Q2_K:
            {
@@ -1358,7 +1398,6 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m
            {
                nsg = N_SG_MXFP4;
                nr0 = N_R0_MXFP4;
-                smem = 32*sizeof(float);
            } break;
        case GGML_TYPE_Q2_K:
            {
@@ -126,6 +126,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_ad
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_tri               (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max          (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexer (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_mask        (ggml_metal_library_t lib, const struct ggml_tensor * op, bool fill, bool has_score, bool has_blocks, bool has_extra, bool vec);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc           (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv          (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t nc, bool use_silu);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched  (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs, int32_t nc, bool use_silu);
@@ -1555,6 +1555,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te

    switch (op->op) {
        case GGML_OP_SCALE:
+            return ggml_is_contiguous_rows(op->src[0]) && (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_BF16);
        case GGML_OP_FILL:
        case GGML_OP_CLAMP:
        case GGML_OP_SQR:
@@ -1834,6 +1835,15 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                ggml_is_contiguous_rows(op->src[1]) &&
                ggml_is_contiguous_rows(op->src[2]) &&
                (op->src[3] == NULL || ggml_is_contiguous_rows(op->src[3]));
+        case GGML_OP_TOP_K_MASK:
+            return (op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32) &&
+                op->src[0]->type == op->type &&
+                op->src[0]->nb[0] == ggml_type_size(op->src[0]->type) &&
+                op->nb[0] == ggml_type_size(op->type) &&
+                op->src[1]->type == GGML_TYPE_I32 &&
+                (op->src[2] == NULL || op->src[2]->type == GGML_TYPE_F32) &&
+                (op->src[3] == NULL || op->src[3]->type == GGML_TYPE_I32) &&
+                (op->src[4] == NULL || op->src[4]->type == GGML_TYPE_I32);
        case GGML_OP_SSM_SCAN:
            return has_simdgroup_reduction;
        case GGML_OP_SSM_CONV:
@@ -125,6 +125,7 @@
 #define FC_FLASH_ATTN_EXT_TENSOR       2200
 #define FC_LIGHTNING_INDEXER           2200
 #define FC_MUL_MV_MMA                  2300
+#define FC_TOP_K_MASK                  2400

 // op-specific constants
 #define OP_FLASH_ATTN_EXT_NQPSG 8
@@ -1295,6 +1296,38 @@ typedef struct {
    uint64_t nbm3;
 } ggml_metal_kargs_lightning_indexer;

+typedef struct {
+    int32_t  n_kv;
+    int32_t  m1;
+    int32_t  m2;
+    int32_t  m3;
+    int32_t  n_top;
+    int32_t  n_pick;
+    int32_t  n_blk;
+    int32_t  n_extra;
+    int64_t  n_elems;
+    uint64_t nb1;
+    uint64_t nb2;
+    uint64_t nb3;
+    uint64_t nbm1;
+    uint64_t nbm2;
+    uint64_t nbm3;
+    uint64_t nbk0;
+    uint64_t nbk1;
+    uint64_t nbk2;
+    uint64_t nbk3;
+    uint64_t nbs0;
+    uint64_t nbs1;
+    uint64_t nbs2;
+    uint64_t nbs3;
+    uint64_t nbb0;
+    uint64_t nbb1;
+    uint64_t nbe0;
+    uint64_t nbe1;
+    uint64_t nbe2;
+    uint64_t nbe3;
+} ggml_metal_kargs_top_k_mask;
+
 typedef struct {
    int32_t  n_tokens;
    int32_t  n_iter;
@@ -345,6 +345,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
            {
                n_fuse = ggml_metal_op_lightning_indexer(ctx, idx);
            } break;
+        case GGML_OP_TOP_K_MASK:
+            {
+                n_fuse = ggml_metal_op_top_k_mask(ctx, idx);
+            } break;
        case GGML_OP_DSV4_HC_COMB:
        case GGML_OP_DSV4_HC_PRE:
        case GGML_OP_DSV4_HC_POST:
@@ -895,7 +899,9 @@ int ggml_metal_op_unary(ggml_metal_op_t ctx, int idx) {
    if (pipeline.cnt) {
        const int n = pipeline.c4 ? ggml_nelements(op)/4 : ggml_nelements(op);

-        ggml_metal_encoder_dispatch_threadgroups(enc, n, 1, 1, 1, 1, 1);
+        const int nth = MIN(256, n);
+
+        ggml_metal_encoder_dispatch_threadgroups(enc, (n + nth - 1)/nth, 1, 1, nth, 1, 1);
    } else {
        const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
        const int nth = MIN(args.ne00, nth_max);
@@ -948,7 +954,9 @@ int ggml_metal_op_glu(ggml_metal_op_t ctx, int idx) {

    const int64_t nrows = ggml_nrows(op->src[0]);

-    const int32_t nth = std::max(1, std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne00/2));
+    // split each row over several threadgroups: with a single row (decode) one threadgroup would run on one GPU core
+    const int32_t nth = std::max(1, std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), std::min(256, ne0)));
+    const int32_t nk  = (ne0 + nth - 1)/nth;

    ggml_metal_encoder_set_pipeline(enc, pipeline);
    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
@@ -960,7 +968,7 @@ int ggml_metal_op_glu(ggml_metal_op_t ctx, int idx) {
    }
    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         3);

-    ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, nth, 1, 1);
+    ggml_metal_encoder_dispatch_threadgroups(enc, nk, nrows, 1, nth, 1, 1);

    return 1;
 }
@@ -1424,6 +1432,124 @@ int ggml_metal_op_lightning_indexer(ggml_metal_op_t ctx, int idx) {
    return 1;
 }

+int ggml_metal_op_top_k_mask(ggml_metal_op_t ctx, int idx) {
+    ggml_tensor * op = ctx->node(idx);
+
+    ggml_metal_library_t lib = ctx->lib;
+    ggml_metal_encoder_t enc = ctx->enc;
+
+    const ggml_tensor * mask   = op->src[0];
+    const ggml_tensor * top_k  = op->src[1];
+    const ggml_tensor * score  = op->src[2];
+    const ggml_tensor * blocks = op->src[3];
+    const ggml_tensor * extra  = op->src[4];
+
+    GGML_ASSERT(op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32);
+    GGML_ASSERT(mask->type == op->type);
+    GGML_ASSERT(top_k->type == GGML_TYPE_I32);
+    GGML_ASSERT(!score  || score->type  == GGML_TYPE_F32);
+    GGML_ASSERT(!blocks || blocks->type == GGML_TYPE_I32);
+    GGML_ASSERT(!extra  || extra->type  == GGML_TYPE_I32);
+
+    const int32_t n_kv    = mask->ne[0];
+    const int32_t n_top   = (int32_t) top_k->ne[0];
+    const int32_t n_pick  = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv;
+    const int32_t n_blk   = blocks ? blocks->ne[0] : 0;
+    const int32_t n_extra = extra  ? extra->ne[0]  : 0;
+
+    const ggml_metal_buffer_id bid = ggml_metal_get_buffer_id(op);
+
+    // the vector fill stores 4 elements at a time
+    const int32_t ts      = ggml_type_size(op->type);
+    const int64_t n_elems = ggml_nelements(op);
+
+    const bool has_score  = score  != NULL;
+    const bool has_blocks = blocks != NULL;
+    const bool has_extra  = extra  != NULL;
+
+    const bool vec = n_elems % 4 == 0 && bid.offs % (4*ts) == 0;
+
+    ggml_metal_kargs_top_k_mask args = {
+        /*.n_kv    =*/ n_kv,
+        /*.m1      =*/ (int32_t) mask->ne[1],
+        /*.m2      =*/ (int32_t) mask->ne[2],
+        /*.m3      =*/ (int32_t) mask->ne[3],
+        /*.n_top   =*/ n_top,
+        /*.n_pick  =*/ n_pick,
+        /*.n_blk   =*/ n_blk,
+        /*.n_extra =*/ n_extra,
+        /*.n_elems =*/ n_elems,
+        /*.nb1     =*/ op->nb[1],
+        /*.nb2     =*/ op->nb[2],
+        /*.nb3     =*/ op->nb[3],
+        /*.nbm1    =*/ mask->nb[1],
+        /*.nbm2    =*/ mask->nb[2],
+        /*.nbm3    =*/ mask->nb[3],
+        /*.nbk0    =*/ top_k->nb[0],
+        /*.nbk1    =*/ top_k->nb[1],
+        /*.nbk2    =*/ top_k->nb[2],
+        /*.nbk3    =*/ top_k->nb[3],
+        /*.nbs0    =*/ score  ? score->nb[0]  : 0,
+        /*.nbs1    =*/ score  ? score->nb[1]  : 0,
+        /*.nbs2    =*/ score  ? score->nb[2]  : 0,
+        /*.nbs3    =*/ score  ? score->nb[3]  : 0,
+        /*.nbb0    =*/ blocks ? blocks->nb[0] : 0,
+        /*.nbb1    =*/ blocks ? blocks->nb[1] : 0,
+        /*.nbe0    =*/ extra  ? extra->nb[0]  : 0,
+        /*.nbe1    =*/ extra  ? extra->nb[1]  : 0,
+        /*.nbe2    =*/ extra  ? extra->nb[2]  : 0,
+        /*.nbe3    =*/ extra  ? extra->nb[3]  : 0,
+    };
+
+    // the fill covers the whole dst, the scatter overwrites the kept cells with the mask values
+    {
+        auto pipeline_fill = ggml_metal_library_get_pipeline_top_k_mask(lib, op, true, false, false, false, vec);
+
+        const int64_t n_fill  = n_elems/(vec ? 4 : 1);
+        const int32_t nth_max = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline_fill);
+
+        int nth = 32; // SIMD width
+
+        while (nth < nth_max && nth < n_fill) {
+            nth *= 2;
+        }
+
+        ggml_metal_encoder_set_pipeline(enc, pipeline_fill);
+        ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0);
+        ggml_metal_encoder_set_buffer(enc, bid, 1);
+
+        ggml_metal_encoder_dispatch_threadgroups(enc, (n_fill + nth - 1)/nth, 1, 1, nth, 1, 1);
+
+        ggml_metal_op_concurrency_reset(ctx);
+    }
+
+    {
+        auto pipeline = ggml_metal_library_get_pipeline_top_k_mask(lib, op, false, has_score, has_blocks, has_extra, false);
+
+        const int32_t n_max   = n_top + n_extra;
+        const int32_t nth_max = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline);
+
+        int nth = 32; // SIMD width
+
+        while (nth < nth_max && nth < n_max) {
+            nth *= 2;
+        }
+
+        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(mask),  1);
+        ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(top_k), 2);
+        ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(score  ? score  : mask), 3);
+        ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(blocks ? blocks : mask), 4);
+        ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(extra  ? extra  : mask), 5);
+        ggml_metal_encoder_set_buffer(enc, bid, 6);
+
+        ggml_metal_encoder_dispatch_threadgroups(enc, op->ne[1], op->ne[2], op->ne[3], nth, 1, 1);
+    }
+
+    return 1;
+}
+
 int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) {
    ggml_tensor * op = ctx->node(idx);

@@ -4206,7 +4332,9 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
    ggml_metal_encoder_set_buffer  (enc, bid_dst,  3);

    if (pipeline.cnt) {
-        ggml_metal_encoder_dispatch_threadgroups(enc, args.ne0, ggml_nrows(op), 1, 1, 1, 1);
+        const int nth = MIN(256, args.ne0);
+
+        ggml_metal_encoder_dispatch_threadgroups(enc, (args.ne0 + nth - 1)/nth, ggml_nrows(op), 1, nth, 1, 1);
    } else {
        const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));

@@ -60,6 +60,7 @@ int ggml_metal_op_get_rows          (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_set_rows          (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_diag              (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_lightning_indexer (ggml_metal_op_t ctx, int idx);
+int ggml_metal_op_top_k_mask        (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_dsv4_hc           (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_soft_max          (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_ssm_conv          (ggml_metal_op_t ctx, int idx);
@@ -22,8 +22,13 @@ kernel void kernel_bin_fuse_impl(

    if (FC_RB) {
        // row broadcast
-        const uint i0 = tgpig.y*args.ne00 + tgpig.x;
-        const uint i1 = FC_CB ? tgpig.x%args.ne10 : tgpig.x;
+        const uint i = tgpig.x*ntg.x + tpitg.x;
+        if (i >= (uint) args.ne0) {
+            return;
+        }
+
+        const uint i0 = tgpig.y*args.ne00 + i;
+        const uint i1 = FC_CB ? i%args.ne10 : i;

        device const T0 * src0_row = (device const T0 *) (src0);
        device       T  * dst_row  = (device       T  *) (dst);
@@ -906,3 +906,99 @@ kernel void kernel_dsv4_hc_post_f32(
        *(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst];
    }
 }
+
+constant bool FC_top_k_mask_has_score  [[function_constant(FC_TOP_K_MASK + 0)]];
+constant bool FC_top_k_mask_has_blocks [[function_constant(FC_TOP_K_MASK + 1)]];
+constant bool FC_top_k_mask_has_extra  [[function_constant(FC_TOP_K_MASK + 2)]];
+constant bool FC_top_k_mask_vec        [[function_constant(FC_TOP_K_MASK + 3)]];
+
+// the fill runs before the scatter, both are fully parallel over the whole dst and over the picks of one row
+template <typename T, typename T4>
+kernel void kernel_top_k_mask_fill(
+        constant ggml_metal_kargs_top_k_mask & args,
+        device char * dst,
+        uint  tid[[thread_position_in_grid]],
+        uint ntot[[threads_per_grid]]) {
+    const T ninf = T(-INFINITY);
+
+    if (FC_top_k_mask_vec) {
+        device T4 * dst_4 = (device T4 *) dst;
+        for (int64_t i = tid; i < args.n_elems/4; i += ntot) {
+            dst_4[i] = T4(ninf);
+        }
+    } else {
+        device T * dst_t = (device T *) dst;
+        for (int64_t i = tid; i < args.n_elems; i += ntot) {
+            dst_t[i] = ninf;
+        }
+    }
+}
+
+typedef decltype(kernel_top_k_mask_fill<half,  half4>)  kernel_top_k_mask_fill_h_t;
+typedef decltype(kernel_top_k_mask_fill<float, float4>) kernel_top_k_mask_fill_f_t;
+
+template [[host_name("kernel_top_k_mask_fill_f16")]] kernel kernel_top_k_mask_fill_h_t kernel_top_k_mask_fill<half,  half4>;
+template [[host_name("kernel_top_k_mask_fill_f32")]] kernel kernel_top_k_mask_fill_f_t kernel_top_k_mask_fill<float, float4>;
+
+// one threadgroup per row: copy the mask cells of the live picks and of the extra cells over the fill
+template <typename T, typename T4>
+kernel void kernel_top_k_mask(
+        constant ggml_metal_kargs_top_k_mask & args,
+        device const  char * mask,
+        device const  char * top_k,
+        device const  char * score,
+        device const  char * blocks,
+        device const  char * extra,
+        device        char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        uint3   tpitg[[thread_position_in_threadgroup]],
+        uint3 ntg_max[[threads_per_threadgroup]]) {
+    const int i1 = tgpig.x;
+    const int i2 = tgpig.y;
+    const int i3 = tgpig.z;
+
+    const int   tiitg = tpitg.x;
+    const int     ntg = ntg_max.x;
+
+    device const char * mask_row  = mask  + (i1 % args.m1)*args.nbm1 + (i2 % args.m2)*args.nbm2 + (i3 % args.m3)*args.nbm3;
+    device       char * dst_row   = dst   + i1*args.nb1 + i2*args.nb2 + i3*args.nb3;
+    device const char * top_k_row = top_k + i1*args.nbk1 + i2*args.nbk2 + i3*args.nbk3;
+
+    device const T * mask_row_t = (device const T *) mask_row;
+    device       T * dst_row_t  = (device       T *) dst_row;
+
+    for (int i = tiitg; i < args.n_top; i += ntg) {
+        const int32_t p = *(device const int32_t *) (top_k_row + i*args.nbk0);
+        if (p < 0 || p >= args.n_pick) {
+            continue;
+        }
+        if (FC_top_k_mask_has_score && !(*(device const float *) (score + p*args.nbs0 + i1*args.nbs1 + i2*args.nbs2 + i3*args.nbs3) > -INFINITY)) {
+            continue;
+        }
+        if (FC_top_k_mask_has_blocks) {
+            FOR_UNROLL (int k = 0; k < args.n_blk; ++k) {
+                const int32_t c = *(device const int32_t *) (blocks + k*args.nbb0 + p*args.nbb1);
+                if (c >= 0 && c < args.n_kv) {
+                    dst_row_t[c] = mask_row_t[c];
+                }
+            }
+        } else if (p < args.n_kv) {
+            dst_row_t[p] = mask_row_t[p];
+        }
+    }
+
+    if (FC_top_k_mask_has_extra) {
+        for (int i = tiitg; i < args.n_extra; i += ntg) {
+            const int32_t c = *(device const int32_t *) (extra + i*args.nbe0 + i1*args.nbe1 + i2*args.nbe2 + i3*args.nbe3);
+            if (c >= 0 && c < args.n_kv) {
+                dst_row_t[c] = mask_row_t[c];
+            }
+        }
+    }
+}
+
+typedef decltype(kernel_top_k_mask<half,  half4>)  kernel_top_k_mask_h_t;
+typedef decltype(kernel_top_k_mask<float, float4>) kernel_top_k_mask_f_t;
+
+template [[host_name("kernel_top_k_mask_f16")]] kernel kernel_top_k_mask_h_t kernel_top_k_mask<half,  half4>;
+template [[host_name("kernel_top_k_mask_f32")]] kernel kernel_top_k_mask_f_t kernel_top_k_mask<float, float4>;
@@ -3030,6 +3030,22 @@ kernel void kernel_mul_mv_iq4_xs_f32(
    kernel_mul_mv_iq4_xs_f32_impl<N_R0_IQ4_XS, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
 }

+// decode the E2M1 value in the low nibble of n, scaled by 2^-14, without a lookup table: the 3 magnitude bits
+// land in the fp16 exponent/mantissa (e=0 gives the 0.5 subnormal). same decode as MLX's fp4_e2m1
+static inline half mxfp4_dq1(uint32_t n) {
+    const half c = as_type<half>(ushort((n & 7) << 9));
+    return (n & 8) ? -c : c;
+}
+
+// low / high nibbles of the 4 bytes of w, in byte order, scaled by 2^-14
+static inline float4 mxfp4_lo4(uint32_t w) {
+    return float4(mxfp4_dq1(w), mxfp4_dq1(w >> 8), mxfp4_dq1(w >> 16), mxfp4_dq1(w >> 24));
+}
+
+static inline float4 mxfp4_hi4(uint32_t w) {
+    return float4(mxfp4_dq1(w >> 4), mxfp4_dq1(w >> 12), mxfp4_dq1(w >> 20), mxfp4_dq1(w >> 28));
+}
+
 template<int NR0, typename args_t>
 void kernel_mul_mv_mxfp4_f32_impl(
        args_t args,
@@ -3042,8 +3058,6 @@ void kernel_mul_mv_mxfp4_f32_impl(
        ushort sgitg) {
    const short NSG = FC_mul_mv_nsg;

-    threadgroup float * shmem_f32 = (threadgroup float *) shmem;
-
    const int r0 = tgpig.x;
    const int r1 = tgpig.y;
    const int im = tgpig.z;
@@ -3065,9 +3079,6 @@ void kernel_mul_mv_mxfp4_f32_impl(
    const short ix = tiisg/2;  // 0...15
    const short it = tiisg%2;  // 0 or 1

-    shmem_f32[tiisg] = kvalues_mxfp4_f[tiisg%16];
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-
    float4 yl[4];
    float sumf[NR0]={0.f};

@@ -3087,14 +3098,25 @@ void kernel_mul_mv_mxfp4_f32_impl(
            device const block_mxfp4 & xb = x[row*ns01 + ib];
            device const uint8_t     * q2 = (device const uint8_t *)(xb.qs + 8*it);

-            float4 acc1 = yl[0]*float4(shmem_f32[q2[0] &  0x0F], shmem_f32[q2[1] &  0x0F], shmem_f32[q2[2] &  0x0F], shmem_f32[q2[3] &  0x0F]);
-            float4 acc2 = yl[1]*float4(shmem_f32[q2[0] >> 4   ], shmem_f32[q2[1] >> 4   ], shmem_f32[q2[2] >> 4   ], shmem_f32[q2[3] >> 4   ]);
-            float4 acc3 = yl[2]*float4(shmem_f32[q2[4] &  0x0F], shmem_f32[q2[5] &  0x0F], shmem_f32[q2[6] &  0x0F], shmem_f32[q2[7] &  0x0F]);
-            float4 acc4 = yl[3]*float4(shmem_f32[q2[4] >> 4   ], shmem_f32[q2[5] >> 4   ], shmem_f32[q2[6] >> 4   ], shmem_f32[q2[7] >> 4   ]);
+            // blocks are 17 bytes, so the 8 quant bytes are only byte aligned: read the aligned words that cover them
+            // and funnel-shift. the third word is read only when needed so no word past the block end is touched
+            device const uint32_t * pw = (device const uint32_t *)((ulong) q2 & ~3ul);
+            const uint sh = ((uint)(ulong) q2 & 3u)*8;
+            const uint32_t a = pw[0];
+            const uint32_t b = pw[1];
+            const uint32_t c = pw[sh ? 2 : 1];
+            const uint32_t w0 = (a >> sh) | ((b << 1) << (31 - sh));
+            const uint32_t w1 = (b >> sh) | ((c << 1) << (31 - sh));
+
+            float4 acc1 = yl[0]*mxfp4_lo4(w0);
+            float4 acc2 = yl[1]*mxfp4_hi4(w0);
+            float4 acc3 = yl[2]*mxfp4_lo4(w1);
+            float4 acc4 = yl[3]*mxfp4_hi4(w1);

            acc1 = (acc1 + acc3) + (acc2 + acc4);

-            sumf[row] += e8m0_to_fp32(xb.e) * ((acc1[0] + acc1[1]) + (acc1[2] + acc1[3]));
+            // 16384 undoes the 2^-14 of the decode
+            sumf[row] += (e8m0_to_fp32(xb.e)*16384.0f) * ((acc1[0] + acc1[1]) + (acc1[2] + acc1[3]));
        }

        yb += 16 * QK_MXFP4;
@@ -3,6 +3,24 @@
 constant short FC_unary_op [[function_constant(FC_UNARY + 0)]];
 constant bool  FC_unary_cnt[[function_constant(FC_UNARY + 1)]];

+template <typename T, typename U>
+inline constexpr T ggml_metal_unary_cast(U x) {
+    return (T) x;
+}
+
+#if defined(GGML_METAL_HAS_BF16)
+// MSL has no C-style cast from float vectors to bfloat vectors
+template <>
+inline bfloat4 ggml_metal_unary_cast<bfloat4, float4>(float4 x) {
+    return bfloat4(bfloat(x[0]), bfloat(x[1]), bfloat(x[2]), bfloat(x[3]));
+}
+
+template <>
+inline bfloat4 ggml_metal_unary_cast<bfloat4, float>(float x) {
+    return bfloat4(bfloat(x), bfloat(x), bfloat(x), bfloat(x));
+}
+#endif
+
 template <typename T0, typename T, typename TC>
 kernel void kernel_unary_impl(
        constant ggml_metal_kargs_unary & args,
@@ -20,7 +38,7 @@ kernel void kernel_unary_impl(
    int i0;

    if (FC_CNT) {
-        i0 = tgpig.x;
+        i0 = tgpig.x*ntg.x + tpitg.x;

        src0_ptr = (device const T0 *) (src0);
        dst_ptr  = (device       T  *) (dst);
@@ -43,84 +61,88 @@ kernel void kernel_unary_impl(
            if (i0 >= args.ne0) {
                return;
            }
+        } else {
+            if (i0 >= args.ne00*args.ne01*args.ne02*args.ne03) {
+                return;
+            }
        }

        const TC x = (TC) src0_ptr[i0];

        if (FC_OP == OP_UNARY_NUM_SCALE) {
-            dst_ptr[i0] = (T) (args.scale * x + args.bias);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(args.scale * x + args.bias);
        }

        if (FC_OP == OP_UNARY_NUM_FILL) {
-            dst_ptr[i0] = (T) args.val;
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(args.val);
        }

        if (FC_OP == OP_UNARY_NUM_CLAMP) {
-            dst_ptr[i0] = (T) clamp(x, args.min, args.max);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(clamp(x, args.min, args.max));
        }

        if (FC_OP == OP_UNARY_NUM_SQR) {
-            dst_ptr[i0] = (T) (x * x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(x * x);
        }

        if (FC_OP == OP_UNARY_NUM_SQRT) {
-            dst_ptr[i0] = (T) sqrt(x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(sqrt(x));
        }

        if (FC_OP == OP_UNARY_NUM_SIN) {
-            dst_ptr[i0] = (T) sin(x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(sin(x));
        }

        if (FC_OP == OP_UNARY_NUM_COS) {
-            dst_ptr[i0] = (T) cos(x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(cos(x));
        }

        if (FC_OP == OP_UNARY_NUM_LOG) {
-            dst_ptr[i0] = (T) log(x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(log(x));
        }

        if (FC_OP == OP_UNARY_NUM_LEAKY_RELU) {
-            dst_ptr[i0] = (T) (TC(x > 0)*x + TC(x <= 0)*(x * args.slope));
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(TC(x > 0)*x + TC(x <= 0)*(x * args.slope));
        }

        if (FC_OP == OP_UNARY_NUM_TANH) {
-            dst_ptr[i0] = (T) precise::tanh(x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(precise::tanh(x));
        }

        if (FC_OP == OP_UNARY_NUM_RELU) {
-            dst_ptr[i0] = (T) fmax(0, x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(fmax(0, x));
        }

        if (FC_OP == OP_UNARY_NUM_SIGMOID) {
-            dst_ptr[i0] = (T) (1 / (1 + exp(-x)));
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(1 / (1 + exp(-x)));
        }

        if (FC_OP == OP_UNARY_NUM_GELU) {
-            dst_ptr[i0] = (T) (0.5*x*(1 + precise::tanh(SQRT_2_OVER_PI*x*(1 + GELU_COEF_A*x*x))));
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(0.5*x*(1 + precise::tanh(SQRT_2_OVER_PI*x*(1 + GELU_COEF_A*x*x))));
        }

        if (FC_OP == OP_UNARY_NUM_GELU_ERF) {
-            dst_ptr[i0] = (T) (0.5*x*(1 + erf_approx(SQRT_2_INV*x)));
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(0.5*x*(1 + erf_approx(SQRT_2_INV*x)));
        }

        if (FC_OP == OP_UNARY_NUM_GELU_QUICK) {
-            dst_ptr[i0] = (T) (x * (1/(1 + exp(GELU_QUICK_COEF*x))));
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(x * (1/(1 + exp(GELU_QUICK_COEF*x))));
        }

        if (FC_OP == OP_UNARY_NUM_SILU) {
-            dst_ptr[i0] = (T) (x / (1 + exp(-x)));
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(x / (1 + exp(-x)));
        }

        if (FC_OP == OP_UNARY_NUM_ELU) {
-            dst_ptr[i0] = (T) elu_approx(x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(elu_approx(x));
        }

        if (FC_OP == OP_UNARY_NUM_NEG) {
-            dst_ptr[i0] = (T) -x;
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(-x);
        }

        if (FC_OP == OP_UNARY_NUM_ABS) {
-            dst_ptr[i0] = (T) fabs(x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(fabs(x));
        }

        if (FC_OP == OP_UNARY_NUM_SGN) {
@@ -132,40 +154,40 @@ kernel void kernel_unary_impl(
        }

        if (FC_OP == OP_UNARY_NUM_HARDSWISH) {
-            dst_ptr[i0] = (T) (x * fmax(0, fmin(1, x/6 + 0.5)));
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(x * fmax(0, fmin(1, x/6 + 0.5)));
        }

        if (FC_OP == OP_UNARY_NUM_HARDSIGMOID) {
-            dst_ptr[i0] = (T) fmax(0, fmin(1, x/6 + 0.5));
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(fmax(0, fmin(1, x/6 + 0.5)));
        }

        if (FC_OP == OP_UNARY_NUM_EXP) {
-            dst_ptr[i0] = (T) exp(x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(exp(x));
        }

        if (FC_OP == OP_UNARY_NUM_SOFTPLUS) {
-            dst_ptr[i0] = (T) select(log(1 + exp(x)), x, x > 20);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(select(log(1 + exp(x)), x, x > 20));
        }

        if (FC_OP == OP_UNARY_NUM_EXPM1) {
            // TODO: precise implementation
-            dst_ptr[i0] = (T) (exp(x) - 1);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(exp(x) - 1);
        }

        if (FC_OP == OP_UNARY_NUM_FLOOR) {
-            dst_ptr[i0] = (T) floor(x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(floor(x));
        }

        if (FC_OP == OP_UNARY_NUM_CEIL) {
-            dst_ptr[i0] = (T) ceil(x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(ceil(x));
        }

        if (FC_OP == OP_UNARY_NUM_ROUND) {
-            dst_ptr[i0] = (T) round(x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(round(x));
        }

        if (FC_OP == OP_UNARY_NUM_TRUNC) {
-            dst_ptr[i0] = (T) trunc(x);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(trunc(x));
        }

        if (FC_OP == OP_UNARY_NUM_XIELU) {
@@ -174,7 +196,7 @@ kernel void kernel_unary_impl(
            const TC clamped = fmin(xi, TC(args.val));
            const TC y_pos   = TC(args.scale) * xi * xi + TC(args.bias) * xi;
            const TC y_neg   = (exp(clamped) - TC(1.0f) - xi) * TC(args.slope) + TC(args.bias) * xi;
-            dst_ptr[i0] = (T) (gate * y_pos + (TC(1.0f) - gate) * y_neg);
+            dst_ptr[i0] = ggml_metal_unary_cast<T>(gate * y_pos + (TC(1.0f) - gate) * y_neg);
        }
    }

@@ -188,6 +210,10 @@ template [[host_name("kernel_unary_f32_f32")]]   kernel kernel_unary_t kernel_un
 template [[host_name("kernel_unary_f32_f32_4")]] kernel kernel_unary_t kernel_unary_impl<float4, float4, float4>;
 template [[host_name("kernel_unary_f16_f16")]]   kernel kernel_unary_t kernel_unary_impl<half,   half,   float>;
 template [[host_name("kernel_unary_f16_f16_4")]] kernel kernel_unary_t kernel_unary_impl<half4,  half4,  float4>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_unary_bf16_bf16")]]   kernel kernel_unary_t kernel_unary_impl<bfloat,  bfloat,  float>;
+template [[host_name("kernel_unary_bf16_bf16_4")]] kernel kernel_unary_t kernel_unary_impl<bfloat4, bfloat4, float4>;
+#endif

 kernel void kernel_silu_back_f32(
        constant ggml_metal_kargs_silu_back & args,
@@ -209,14 +235,15 @@ kernel void kernel_reglu(
        device const char * src0,
        device const char * src1,
        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+        uint2 tgpig[[threadgroup_position_in_grid]],
+        uint2 tpitg[[thread_position_in_threadgroup]],
+        uint2   ntg[[threads_per_threadgroup]],
+        uint2  ntgs[[threadgroups_per_grid]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig.y*args.nb1);

-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+    for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
        const float x0 = src0_row[i0];
        const float x1 = src1_row[i0];

@@ -235,14 +262,15 @@ kernel void kernel_geglu(
        device const char * src0,
        device const char * src1,
        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+        uint2 tgpig[[threadgroup_position_in_grid]],
+        uint2 tpitg[[thread_position_in_threadgroup]],
+        uint2   ntg[[threads_per_threadgroup]],
+        uint2  ntgs[[threadgroups_per_grid]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig.y*args.nb1);

-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+    for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
        const float x0 = src0_row[i0];
        const float x1 = src1_row[i0];

@@ -263,14 +291,15 @@ kernel void kernel_swiglu(
        device const char * src0,
        device const char * src1,
        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+        uint2 tgpig[[threadgroup_position_in_grid]],
+        uint2 tpitg[[thread_position_in_threadgroup]],
+        uint2   ntg[[threads_per_threadgroup]],
+        uint2  ntgs[[threadgroups_per_grid]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig.y*args.nb1);

-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+    for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
        const float x0 = src0_row[i0];
        const float x1 = src1_row[i0];

@@ -291,14 +320,15 @@ kernel void kernel_swiglu_oai(
        device const char * src0,
        device const char * src1,
        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+        uint2 tgpig[[threadgroup_position_in_grid]],
+        uint2 tpitg[[thread_position_in_threadgroup]],
+        uint2   ntg[[threads_per_threadgroup]],
+        uint2  ntgs[[threadgroups_per_grid]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig.y*args.nb1);

-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+    for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
        float x0 = src0_row[i0];
        float x1 = src1_row[i0];

@@ -323,14 +353,15 @@ kernel void kernel_swiglu_clamp(
        device const char * src0,
        device const char * src1,
        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+        uint2 tgpig[[threadgroup_position_in_grid]],
+        uint2 tpitg[[thread_position_in_threadgroup]],
+        uint2   ntg[[threads_per_threadgroup]],
+        uint2  ntgs[[threadgroups_per_grid]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig.y*args.nb1);

-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+    for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
        const float gate = min((float) src0_row[i0], args.limit);
        const float up = clamp((float) src1_row[i0], -args.limit, args.limit);

@@ -349,14 +380,15 @@ kernel void kernel_geglu_erf(
        device const char * src0,
        device const char * src1,
        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+        uint2 tgpig[[threadgroup_position_in_grid]],
+        uint2 tpitg[[thread_position_in_threadgroup]],
+        uint2   ntg[[threads_per_threadgroup]],
+        uint2  ntgs[[threadgroups_per_grid]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig.y*args.nb1);

-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+    for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
        const float x0 = src0_row[i0];
        const float x1 = src1_row[i0];

@@ -377,14 +409,15 @@ kernel void kernel_geglu_quick(
        device const char * src0,
        device const char * src1,
        device       char * dst,
-        uint tgpig[[threadgroup_position_in_grid]],
-        uint tpitg[[thread_position_in_threadgroup]],
-        uint   ntg[[threads_per_threadgroup]]) {
-    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig*args.nb01) + args.i00;
-    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig*args.nb11) + args.i10;
-    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig*args.nb1);
+        uint2 tgpig[[threadgroup_position_in_grid]],
+        uint2 tpitg[[thread_position_in_threadgroup]],
+        uint2   ntg[[threads_per_threadgroup]],
+        uint2  ntgs[[threadgroups_per_grid]]) {
+    device const T * src0_row = (device const T *) ((device const char *) src0 + tgpig.y*args.nb01) + args.i00;
+    device const T * src1_row = (device const T *) ((device const char *) src1 + tgpig.y*args.nb11) + args.i10;
+    device       T * dst_row  = (device       T *) ((device       char *) dst  + tgpig.y*args.nb1);

-    for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) {
+    for (int i0 = tgpig.x*ntg.x + tpitg.x; i0 < args.ne0; i0 += ntg.x*ntgs.x) {
        const float x0 = src0_row[i0];
        const float x1 = src1_row[i0];

@@ -104,7 +104,34 @@ set(GGML_OPENCL_KERNELS
    gemm_noshuffle_q2_k_q8_1_dp4a
    gemv_noshuffle_q2_k_f32
    gemm_noshuffle_q2_k_f32
+    gemm_noshuffle_q3_k_f32
+    gemm_noshuffle_iq2_s_f32
+    gemm_noshuffle_iq2_xs_f32
+    gemm_noshuffle_iq2_xxs_f32
+    gemm_noshuffle_iq2_xxs_q8_1_dp4a
+    gemv_noshuffle_iq2_xxs_f32
+    gemm_noshuffle_iq2_xs_q8_1_dp4a
+    gemv_noshuffle_iq2_xs_f32
+    gemm_noshuffle_iq2_s_q8_1_dp4a
+    gemv_noshuffle_iq2_s_f32
+    gemm_noshuffle_q3_k_q8_1_dp4a
+    gemv_noshuffle_q3_k_f32
    mul_mv_q2_k_f32
+    mul_mv_q3_k_f32
+    iq2_s_grid
+    mul_mv_iq2_s_f32
+    iq2_xs_grid
+    mul_mv_iq2_xs_f32
+    iq2_xxs_grid
+    mul_mv_iq2_xxs_f32
+    mul_mv_iq2_xxs_f32_flat
+    mul_mm_iq2_xxs_f32_l4_lm
+    mul_mv_iq2_xs_f32_flat
+    mul_mm_iq2_xs_f32_l4_lm
+    mul_mv_iq2_s_f32_flat
+    mul_mm_iq2_s_f32_l4_lm
+    mul_mv_q3_k_f32_flat
+    mul_mm_q3_k_f32_l4_lm
    mul_mv_q2_k_f32_flat
    mul_mm_q2_k_f32_l4_lm
    mul_mv_q4_k_f32
@@ -128,6 +155,7 @@ set(GGML_OPENCL_KERNELS
    mul_mv_id_q8_0_f32_flat
    mul_mv_id_mxfp4_f32
    mul_mv_id_mxfp4_f32_flat
+    mul_mat_id_utils
    gemm_moe_q4_0_f32_ns
    gemm_moe_q4_0_q8_1_dp4a
    gemv_moe_q4_0_f32_ns
@@ -199,6 +227,7 @@ set(GGML_OPENCL_KERNELS
    gemv_noshuffle_q4_k_f32
    gemv_noshuffle_q4_k_f32_o4
    gemv_noshuffle_q4_k_f32_tiled
+    gemm_noshuffle_q4_K_f32_tiled_dp4a
    gemm_noshuffle_q4_k_f32
    gemv_noshuffle_q4_k_f32_32b_trans
    gemv_noshuffle_q6_k_f32
@@ -2773,3 +2773,285 @@ kernel void kernel_restore_block_q2_k_noshuffle(
        }
    }
 }
+
+struct block_q3_K {
+    uchar hmask[QK_K/8];
+    uchar qs[QK_K/4];
+    uchar scales[12];
+    half d;
+};
+
+kernel void kernel_convert_block_q3_K(
+    global struct block_q3_K * src0,
+    global uchar * dst_qs,
+    global uchar * dst_hm,
+    global uchar * dst_sc,
+    global half * dst_d,
+    ulong n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_q3_K * b = src0 + i;
+    for (int j = 0; j < 64; ++j) {
+        dst_qs[64*i + j] = b->qs[j];
+    }
+    for (int j = 0; j < 32; ++j) {
+        dst_hm[32*i + j] = b->hmask[j];
+    }
+    for (int j = 0; j < 12; ++j) {
+        dst_sc[12*i + j] = b->scales[j];
+    }
+    dst_d[i] = b->d;
+}
+
+kernel void kernel_restore_block_q3_K(
+    global uchar * src_qs,
+    global uchar * src_hm,
+    global uchar * src_sc,
+    global half * src_d,
+    global struct block_q3_K * dst,
+    ulong n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_q3_K * b = dst + i;
+    for (int j = 0; j < 64; ++j) {
+        b->qs[j] = src_qs[64*i + j];
+    }
+    for (int j = 0; j < 32; ++j) {
+        b->hmask[j] = src_hm[32*i + j];
+    }
+    for (int j = 0; j < 12; ++j) {
+        b->scales[j] = src_sc[12*i + j];
+    }
+    b->d = src_d[i];
+}
+
+kernel void kernel_convert_block_q3_k_noshuffle(
+    global struct block_q3_K * src0,
+    global uchar * dst_qs,      // QK_K/4 per block
+    global uchar * dst_hm,      // QK_K/8 per block
+    global uint  * dst_sc,      // 3 per block
+    global half  * dst_d,       // 1
+    ulong          n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_q3_K * b = (global struct block_q3_K *) src0 + i;
+
+    dst_d[i] = b->d;
+
+    for (int w = 0; w < 3; ++w) {
+        global const uchar * p = b->scales + 4*w;
+        dst_sc[3*i + w] = (uint)p[0] | ((uint)p[1] << 8) | ((uint)p[2] << 16) | ((uint)p[3] << 24);
+    }
+
+    for (int g = 0; g < QK_K/4; ++g) {
+        uint pk = 0;
+        for (int t = 0; t < 4; ++t) {
+            const int e  = 4*g + t;
+            const int sb = e >> 5;
+            const int ee = e & 31;
+            pk |= (uint)(((b->qs[32*(sb >> 2) + ee] >> (2*(sb & 3))) & 3) << (2*t));
+        }
+        dst_qs[(QK_K/4) * i + g] = (uchar)pk;
+    }
+
+    // two groups per hmask byte, so this loop owns the byte and cannot race
+    for (int t = 0; t < QK_K/8; ++t) {
+        uint v = 0;
+        for (int s = 0; s < 2; ++s) {
+            const int g = 2*t + s;
+            for (int u = 0; u < 4; ++u) {
+                const int e  = 4*g + u;
+                const int sb = e >> 5;
+                const int ee = e & 31;
+                v |= (uint)(((b->hmask[ee] >> sb) & 1) << (4*s + u));
+            }
+        }
+        dst_hm[(QK_K/8) * i + t] = (uchar)v;
+    }
+}
+
+kernel void kernel_restore_block_q3_k_noshuffle(
+    global uchar * src_qs,
+    global uchar * src_hm,
+    global uint  * src_sc,
+    global half  * src_d,
+    global struct block_q3_K * dst,
+    ulong          n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_q3_K * b = (global struct block_q3_K *) dst + i;
+
+    b->d = src_d[i];
+
+    for (int w = 0; w < 3; ++w) {
+        const uint x = src_sc[3*i + w];
+        global uchar * p = b->scales + 4*w;
+        p[0] = (uchar)( x        & 0xFF);
+        p[1] = (uchar)((x >>  8) & 0xFF);
+        p[2] = (uchar)((x >> 16) & 0xFF);
+        p[3] = (uchar)((x >> 24) & 0xFF);
+    }
+
+    for (int n2 = 0; n2 < 2; ++n2) {
+        for (int ee = 0; ee < 32; ++ee) {
+            uint v = 0;
+            for (int j = 0; j < 4; ++j) {
+                const int sb = 4*n2 + j;
+                const int e  = 32*sb + ee;
+                v |= ((uint)(src_qs[(QK_K/4) * i + (e >> 2)] >> (2*(e & 3))) & 3u) << (2*j);
+            }
+            b->qs[32*n2 + ee] = (uchar)v;
+        }
+    }
+
+    for (int ee = 0; ee < 32; ++ee) {
+        uint v = 0;
+        for (int sb = 0; sb < 8; ++sb) {
+            const int e = 32*sb + ee;
+            const int g = e >> 2;
+            v |= ((uint)(src_hm[(QK_K/8) * i + (g >> 1)] >> (4*(g & 1) + (e & 3))) & 1u) << sb;
+        }
+        b->hmask[ee] = (uchar)v;
+    }
+}
+
+struct block_iq2_s {
+    half d;
+    uchar qs[QK_K/4];
+    uchar qh[QK_K/32];
+    uchar scales[QK_K/32];
+};
+
+kernel void kernel_convert_block_iq2_s(
+    global struct block_iq2_s * src0,
+    global uchar * dst_qs,      // QK_K/8 per block
+    global uchar * dst_sg,      // QK_K/8
+    global uchar * dst_qh,      // QK_K/32
+    global uchar * dst_sc,      // QK_K/32
+    global half  * dst_d,       // 1
+    ulong          n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_iq2_s * b = src0 + i;
+
+    dst_d[i] = b->d;
+    for (int j = 0; j < QK_K/8;  ++j) { dst_qs[(QK_K/8)  * i + j] = b->qs[j];            }
+    for (int j = 0; j < QK_K/8;  ++j) { dst_sg[(QK_K/8)  * i + j] = b->qs[QK_K/8 + j];   }
+    for (int j = 0; j < QK_K/32; ++j) { dst_qh[(QK_K/32) * i + j] = b->qh[j];            }
+    for (int j = 0; j < QK_K/32; ++j) { dst_sc[(QK_K/32) * i + j] = b->scales[j];        }
+}
+
+kernel void kernel_restore_block_iq2_s(
+    global uchar * src_qs,
+    global uchar * src_sg,
+    global uchar * src_qh,
+    global uchar * src_sc,
+    global half  * src_d,
+    global struct block_iq2_s * dst,
+    ulong          n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_iq2_s * b = dst + i;
+
+    b->d = src_d[i];
+    for (int j = 0; j < QK_K/8;  ++j) { b->qs[j]          = src_qs[(QK_K/8)  * i + j]; }
+    for (int j = 0; j < QK_K/8;  ++j) { b->qs[QK_K/8 + j] = src_sg[(QK_K/8)  * i + j]; }
+    for (int j = 0; j < QK_K/32; ++j) { b->qh[j]          = src_qh[(QK_K/32) * i + j]; }
+    for (int j = 0; j < QK_K/32; ++j) { b->scales[j]      = src_sc[(QK_K/32) * i + j]; }
+}
+
+struct block_iq2_xs {
+    half d;
+    ushort qs[QK_K/8];
+    uchar scales[QK_K/32];
+};
+
+kernel void kernel_convert_block_iq2_xs(
+    global struct block_iq2_xs * src0,
+    global ushort * dst_qs,
+    global uchar * dst_sc,
+    global half * dst_d,
+    ulong n_blocks
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blocks) { return; }
+    dst_d[i] = src0[i].d;
+    for (int j = 0; j < 32; ++j) { dst_qs[32*i + j] = src0[i].qs[j]; }
+    for (int j = 0; j < 8; ++j) { dst_sc[8*i + j] = src0[i].scales[j]; }
+}
+
+kernel void kernel_restore_block_iq2_xs(
+    global ushort * src_qs,
+    global uchar * src_sc,
+    global half * src_d,
+    global struct block_iq2_xs * dst,
+    ulong n_blocks
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blocks) { return; }
+    dst[i].d = src_d[i];
+    for (int j = 0; j < 32; ++j) { dst[i].qs[j] = src_qs[32*i + j]; }
+    for (int j = 0; j < 8; ++j) { dst[i].scales[j] = src_sc[8*i + j]; }
+}
+
+struct block_iq2_xxs {
+    half d;
+    ushort qs[QK_K/8];
+};
+
+kernel void kernel_convert_block_iq2_xxs(
+    global struct block_iq2_xxs * src0,
+    global uchar * dst_qs,
+    global uint * dst_sas,
+    global half * dst_d,
+    ulong n_blocks
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blocks) { return; }
+    dst_d[i] = src0[i].d;
+    for (int j = 0; j < 8; ++j) {
+        global ushort * q = src0[i].qs + 4*j;
+        uint indices = (uint)q[0] | ((uint)q[1] << 16);
+        vstore4(as_uchar4(indices), 8*i + j, dst_qs);
+        dst_sas[8*i + j] = (uint)q[2] | ((uint)q[3] << 16);
+    }
+}
+
+kernel void kernel_restore_block_iq2_xxs(
+    global uchar * src_qs,
+    global uint * src_sas,
+    global half * src_d,
+    global struct block_iq2_xxs * dst,
+    ulong n_blocks
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blocks) { return; }
+    dst[i].d = src_d[i];
+    for (int j = 0; j < 8; ++j) {
+        global ushort * q = dst[i].qs + 4*j;
+        uint indices = as_uint(vload4(8*i + j, src_qs));
+        uint sas = src_sas[8*i + j];
+        q[0] = (ushort)indices;
+        q[1] = (ushort)(indices >> 16);
+        q[2] = (ushort)sas;
+        q[3] = (ushort)(sas >> 16);
+    }
+}
@@ -0,0 +1,93 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_qcom_reqd_sub_group_size
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+__attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+kernel void kernel_gemm_noshuffle_iq2_s_f32(
+    global const uchar * src0_qs,
+    global const uchar * src0_sg,
+    global const uchar * src0_qh,
+    global const uchar * src0_sc,
+    global const half  * src0_d,
+    read_only image1d_buffer_t src1,
+    global float * dst,
+    ulong offsetd,
+    int m,
+    int n,
+    int k,
+    int n_no_padding
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+    const int gy = get_global_id(0);
+    const int row = get_global_id(1) * 4;
+    const int n_4 = n / 4;
+    if (row >= m) {
+        return;
+    }
+
+    half8 c0 = 0, c1 = 0, c2 = 0, c3 = 0;
+    for (int i = 0; i < k; i += 16) {
+        const uchar4 sc = vload4(0, src0_sc + row + (i / 32) * m);
+        const half4 d = vload4(0, src0_d + row + (i / 256) * m);
+        const float4 scale_f = convert_float4(d) * 0.25f * (0.5f + convert_float4((sc >> (uchar4)(4*((i / 16) & 1))) & (uchar4)(15)));
+        const half4 scale = convert_half4(scale_f);
+
+        for (int l = 0; l < 16; l += 4) {
+            const int ki = i + l;
+            const uint4 qs = convert_uint4(vload4(0, src0_qs + row + (ki / 8) * m));
+            const uint4 qh = convert_uint4(vload4(0, src0_qh + row + (ki / 32) * m));
+            const uchar4 sg = vload4(0, src0_sg + row + (ki / 8) * m);
+            const uint4 gi = qs | (((qh >> (uint4)(2*((ki / 8) & 3))) & (uint4)(3)) << (uint4)(8));
+            const uint half_grid = (ki & 7) / 4;
+            const uint4 grid = (uint4)(iq2_s_grid[2*gi.s0 + half_grid], iq2_s_grid[2*gi.s1 + half_grid],
+                                      iq2_s_grid[2*gi.s2 + half_grid], iq2_s_grid[2*gi.s3 + half_grid]);
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                half8 b;
+                b.s0123 = read_imageh(src1, gy * 2     + (ki + j) * n_4);
+                b.s4567 = read_imageh(src1, gy * 2 + 1 + (ki + j) * n_4);
+                const int4 magnitude = convert_int4((grid >> (uint4)(8*j)) & (uint4)(255));
+                const int4 signs = convert_int4((sg >> (uchar4)((ki & 7) + j)) & (uchar4)(1));
+                const half4 w = convert_half4(magnitude * ((int4)(1) - 2*signs)) * scale;
+                c0 += b * w.s0;
+                c1 += b * w.s1;
+                c2 += b * w.s2;
+                c3 += b * w.s3;
+            }
+        }
+    }
+
+    int idx = gy * 8 * m + row;
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s0, c1.s0, c2.s0, c3.s0), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s1, c1.s1, c2.s1, c3.s1), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s2, c1.s2, c2.s2, c3.s2), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s3, c1.s3, c2.s3, c3.s3), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s4, c1.s4, c2.s4, c3.s4), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s5, c1.s5, c2.s5, c3.s5), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s6, c1.s6, c2.s6, c3.s6), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s7, c1.s7, c2.s7, c3.s7), 0, dst + idx);
+    }
+}
@@ -0,0 +1,149 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#ifdef cl_khr_integer_dot_product
+#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
+#endif
+
+#define QK_K 256
+
+#ifndef TILESIZE_N
+#define TILESIZE_N 32
+#endif
+
+inline uint iq2_s_pack(uint gv, uint sg, uint base) {
+    int v0 = (int)((gv >>  0) & 0xFF); if (sg & (1u << (base + 0))) { v0 = -v0; }
+    int v1 = (int)((gv >>  8) & 0xFF); if (sg & (1u << (base + 1))) { v1 = -v1; }
+    int v2 = (int)((gv >> 16) & 0xFF); if (sg & (1u << (base + 2))) { v2 = -v2; }
+    int v3 = (int)((gv >> 24) & 0xFF); if (sg & (1u << (base + 3))) { v3 = -v3; }
+    return ((uint)v0 & 0xFFu) | (((uint)v1 & 0xFFu) <<  8)
+         | (((uint)v2 & 0xFFu) << 16) | (((uint)v3 & 0xFFu) << 24);
+}
+
+// Pass staged operands by value; local dp4a operands miscompile on X2.
+inline int dot4_q8a_v(uint4 qw, uint4 a) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_iq2_s_q8_1_dp4a(
+        __global const uchar  * src0_qs,
+        __global const uchar  * src0_sg,
+        __global const uchar  * src0_qh,
+        __global const uchar  * src0_sc,
+        __global const half   * src0_d,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+
+#define IQ2_S_GRID(i) iq2_s_grid[(i)]
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+        const uint ib  = sub >> 3;
+
+        const uint  qhv = (uint)src0_qh[rrow + sub * (uint)m];
+        const uint  scv = (uint)src0_sc[rrow + sub * (uint)m];
+        const float dv  = (float)src0_d[rrow + ib * (uint)m];
+        const float dl0 = dv * (0.5f + (float)( scv       & 0xFu)) * 0.25f;
+        const float dl1 = dv * (0.5f + (float)( scv >> 4)        ) * 0.25f;
+
+        const uint gb = rrow + (step >> 3) * (uint)m;
+
+        uint4 qlo, qhi;
+        {
+            const uint gi0 = (uint)src0_qs[gb + 0u * (uint)m] | ((( qhv       & 3u) << 8));
+            const uint gi1 = (uint)src0_qs[gb + 1u * (uint)m] | ((((qhv >> 2) & 3u) << 8));
+            const uint gi2 = (uint)src0_qs[gb + 2u * (uint)m] | ((((qhv >> 4) & 3u) << 8));
+            const uint gi3 = (uint)src0_qs[gb + 3u * (uint)m] | ((((qhv >> 6) & 3u) << 8));
+            const uint s0  = (uint)src0_sg[gb + 0u * (uint)m];
+            const uint s1  = (uint)src0_sg[gb + 1u * (uint)m];
+            const uint s2  = (uint)src0_sg[gb + 2u * (uint)m];
+            const uint s3  = (uint)src0_sg[gb + 3u * (uint)m];
+            qlo.s0 = iq2_s_pack(IQ2_S_GRID(2u*gi0 + 0u), s0, 0u);
+            qlo.s1 = iq2_s_pack(IQ2_S_GRID(2u*gi0 + 1u), s0, 4u);
+            qlo.s2 = iq2_s_pack(IQ2_S_GRID(2u*gi1 + 0u), s1, 0u);
+            qlo.s3 = iq2_s_pack(IQ2_S_GRID(2u*gi1 + 1u), s1, 4u);
+            qhi.s0 = iq2_s_pack(IQ2_S_GRID(2u*gi2 + 0u), s2, 0u);
+            qhi.s1 = iq2_s_pack(IQ2_S_GRID(2u*gi2 + 1u), s2, 4u);
+            qhi.s2 = iq2_s_pack(IQ2_S_GRID(2u*gi3 + 0u), s3, 0u);
+            qhi.s3 = iq2_s_pack(IQ2_S_GRID(2u*gi3 + 1u), s3, 4u);
+        }
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+#define IQ2_S_COL(b) (dl0 * (float)dot4_q8a_v(qlo, (uint4)(sh_qa4[b][0]))  \
+                   + dl1 * (float)dot4_q8a_v(qhi, (uint4)(sh_qa4[b][1])))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+            rf.s0 = IQ2_S_COL(b+0);  rf.s1 = IQ2_S_COL(b+1);
+            rf.s2 = IQ2_S_COL(b+2);  rf.s3 = IQ2_S_COL(b+3);
+            acc[g] += LD4(sh_d, b) * rf;
+        }
+#undef IQ2_S_COL
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
+        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
+        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
+        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
+    }
+#undef NGROUPS
+#undef IQ2_S_GRID
+}
@@ -0,0 +1,91 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_qcom_reqd_sub_group_size
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+__attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+kernel void kernel_gemm_noshuffle_iq2_xs_f32(
+    global const ushort * src0_qs,
+    global const uchar * src0_sc,
+    global const half  * src0_d,
+    read_only image1d_buffer_t src1,
+    global float * dst,
+    ulong offsetd,
+    int m,
+    int n,
+    int k,
+    int n_no_padding
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+    const int gy = get_global_id(0);
+    const int row = get_global_id(1) * 4;
+    const int n_4 = n / 4;
+    if (row >= m) {
+        return;
+    }
+
+    half8 c0 = 0, c1 = 0, c2 = 0, c3 = 0;
+    for (int i = 0; i < k; i += 16) {
+        const uchar4 sc = vload4(0, src0_sc + row + (i / 32) * m);
+        const half4 d = vload4(0, src0_d + row + (i / 256) * m);
+        const float4 scale_f = convert_float4(d) * 0.25f * (0.5f + convert_float4((sc >> (uchar4)(4*((i / 16) & 1))) & (uchar4)(15)));
+        const half4 scale = convert_half4(scale_f);
+
+        for (int l = 0; l < 16; l += 4) {
+            const int ki = i + l;
+            const uint4 qs = convert_uint4(vload4(0, src0_qs + row + (ki / 8) * m));
+            const uint4 code = qs >> (uint4)(9);
+            const uint4 sg = code | ((popcount(code) & (uint4)(1)) << (uint4)(7));
+            const uint4 gi = qs & (uint4)(511);
+            const uint half_grid = (ki & 7) / 4;
+            const uint4 grid = (uint4)(iq2_xs_grid[2*gi.s0 + half_grid], iq2_xs_grid[2*gi.s1 + half_grid],
+                                      iq2_xs_grid[2*gi.s2 + half_grid], iq2_xs_grid[2*gi.s3 + half_grid]);
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                half8 b;
+                b.s0123 = read_imageh(src1, gy * 2     + (ki + j) * n_4);
+                b.s4567 = read_imageh(src1, gy * 2 + 1 + (ki + j) * n_4);
+                const int4 magnitude = convert_int4((grid >> (uint4)(8*j)) & (uint4)(255));
+                const int4 signs = convert_int4((sg >> (uint4)((ki & 7) + j)) & (uint4)(1));
+                const half4 w = convert_half4(magnitude * ((int4)(1) - 2*signs)) * scale;
+                c0 += b * w.s0;
+                c1 += b * w.s1;
+                c2 += b * w.s2;
+                c3 += b * w.s3;
+            }
+        }
+    }
+
+    int idx = gy * 8 * m + row;
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s0, c1.s0, c2.s0, c3.s0), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s1, c1.s1, c2.s1, c3.s1), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s2, c1.s2, c2.s2, c3.s2), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s3, c1.s3, c2.s3, c3.s3), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s4, c1.s4, c2.s4, c3.s4), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s5, c1.s5, c2.s5, c3.s5), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s6, c1.s6, c2.s6, c3.s6), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s7, c1.s7, c2.s7, c3.s7), 0, dst + idx);
+    }
+}
@@ -0,0 +1,150 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#ifdef cl_khr_integer_dot_product
+#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
+#endif
+
+#define QK_K 256
+
+#ifndef TILESIZE_N
+#define TILESIZE_N 32
+#endif
+
+inline uint iq2_xs_pack(uint gv, uint sg, uint base) {
+    int v0 = (int)((gv >>  0) & 0xFF); if (sg & (1u << (base + 0))) { v0 = -v0; }
+    int v1 = (int)((gv >>  8) & 0xFF); if (sg & (1u << (base + 1))) { v1 = -v1; }
+    int v2 = (int)((gv >> 16) & 0xFF); if (sg & (1u << (base + 2))) { v2 = -v2; }
+    int v3 = (int)((gv >> 24) & 0xFF); if (sg & (1u << (base + 3))) { v3 = -v3; }
+    return ((uint)v0 & 0xFFu) | (((uint)v1 & 0xFFu) <<  8)
+         | (((uint)v2 & 0xFFu) << 16) | (((uint)v3 & 0xFFu) << 24);
+}
+
+// Pass staged operands by value; local dp4a operands miscompile on X2.
+inline int dot4_q8a_v(uint4 qw, uint4 a) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_iq2_xs_q8_1_dp4a(
+        __global const ushort * src0_qs,
+        __global const uchar  * src0_sc,
+        __global const half   * src0_d,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+
+#define IQ2_XS_GRID(i) iq2_xs_grid[(i)]
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+        const uint ib  = sub >> 3;
+
+        const uint  scv = (uint)src0_sc[rrow + sub * (uint)m];
+        const float dv  = (float)src0_d[rrow + ib * (uint)m];
+        const float dl0 = dv * (0.5f + (float)( scv       & 0xFu)) * 0.25f;
+        const float dl1 = dv * (0.5f + (float)( scv >> 4)        ) * 0.25f;
+
+        const uint gb = rrow + (step >> 3) * (uint)m;
+
+        uint4 qlo, qhi;
+        {
+            const uint q0 = (uint)src0_qs[gb + 0u * (uint)m];
+            const uint q1 = (uint)src0_qs[gb + 1u * (uint)m];
+            const uint q2 = (uint)src0_qs[gb + 2u * (uint)m];
+            const uint q3 = (uint)src0_qs[gb + 3u * (uint)m];
+            const uint gi0 = (q0 & 511u) << 1;
+            const uint gi1 = (q1 & 511u) << 1;
+            const uint gi2 = (q2 & 511u) << 1;
+            const uint gi3 = (q3 & 511u) << 1;
+            const uint s0  = iq2_xs_signs(q0 >> 9);
+            const uint s1  = iq2_xs_signs(q1 >> 9);
+            const uint s2  = iq2_xs_signs(q2 >> 9);
+            const uint s3  = iq2_xs_signs(q3 >> 9);
+            qlo.s0 = iq2_xs_pack(IQ2_XS_GRID(gi0 + 0u), s0, 0u);
+            qlo.s1 = iq2_xs_pack(IQ2_XS_GRID(gi0 + 1u), s0, 4u);
+            qlo.s2 = iq2_xs_pack(IQ2_XS_GRID(gi1 + 0u), s1, 0u);
+            qlo.s3 = iq2_xs_pack(IQ2_XS_GRID(gi1 + 1u), s1, 4u);
+            qhi.s0 = iq2_xs_pack(IQ2_XS_GRID(gi2 + 0u), s2, 0u);
+            qhi.s1 = iq2_xs_pack(IQ2_XS_GRID(gi2 + 1u), s2, 4u);
+            qhi.s2 = iq2_xs_pack(IQ2_XS_GRID(gi3 + 0u), s3, 0u);
+            qhi.s3 = iq2_xs_pack(IQ2_XS_GRID(gi3 + 1u), s3, 4u);
+        }
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+#define IQ2_XS_COL(b) (dl0 * (float)dot4_q8a_v(qlo, (uint4)(sh_qa4[b][0]))  \
+                    + dl1 * (float)dot4_q8a_v(qhi, (uint4)(sh_qa4[b][1])))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+            rf.s0 = IQ2_XS_COL(b+0);  rf.s1 = IQ2_XS_COL(b+1);
+            rf.s2 = IQ2_XS_COL(b+2);  rf.s3 = IQ2_XS_COL(b+3);
+            acc[g] += LD4(sh_d, b) * rf;
+        }
+#undef IQ2_XS_COL
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
+        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
+        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
+        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
+    }
+#undef NGROUPS
+#undef IQ2_XS_GRID
+}
@@ -0,0 +1,90 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_qcom_reqd_sub_group_size
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+__attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+kernel void kernel_gemm_noshuffle_iq2_xxs_f32(
+    global const uchar * src0_qs,
+    global const uint * src0_sas,
+    global const half  * src0_d,
+    read_only image1d_buffer_t src1,
+    global float * dst,
+    ulong offsetd,
+    int m,
+    int n,
+    int k,
+    int n_no_padding
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+    const int gy = get_global_id(0);
+    const int row = get_global_id(1) * 4;
+    const int n_4 = n / 4;
+    if (row >= m) {
+        return;
+    }
+
+    half8 c0 = 0, c1 = 0, c2 = 0, c3 = 0;
+    for (int i = 0; i < k; i += 16) {
+        const uint4 sas = vload4(0, src0_sas + row + (i / 32) * m);
+        const half4 d = vload4(0, src0_d + row + (i / 256) * m);
+        const float4 scale_f = convert_float4(d) * 0.25f * (0.5f + convert_float4(sas >> (uint4)(28)));
+        const half4 scale = convert_half4(scale_f);
+
+        for (int l = 0; l < 16; l += 4) {
+            const int ki = i + l;
+            const uint4 gi = convert_uint4(vload4(0, src0_qs + row + (ki / 8) * m));
+            const uint4 code = (sas >> (uint4)(7*((ki / 8) & 3))) & (uint4)(127);
+            const uint4 sg = code | ((popcount(code) & (uint4)(1)) << (uint4)(7));
+            const uint half_grid = (ki & 7) / 4;
+            const uint4 grid = (uint4)(iq2_xxs_grid[2*gi.s0 + half_grid], iq2_xxs_grid[2*gi.s1 + half_grid],
+                                      iq2_xxs_grid[2*gi.s2 + half_grid], iq2_xxs_grid[2*gi.s3 + half_grid]);
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                half8 b;
+                b.s0123 = read_imageh(src1, gy * 2     + (ki + j) * n_4);
+                b.s4567 = read_imageh(src1, gy * 2 + 1 + (ki + j) * n_4);
+                const int4 magnitude = convert_int4((grid >> (uint4)(8*j)) & (uint4)(255));
+                const int4 signs = convert_int4((sg >> (uint4)((ki & 7) + j)) & (uint4)(1));
+                const half4 w = convert_half4(magnitude * ((int4)(1) - 2*signs)) * scale;
+                c0 += b * w.s0;
+                c1 += b * w.s1;
+                c2 += b * w.s2;
+                c3 += b * w.s3;
+            }
+        }
+    }
+
+    int idx = gy * 8 * m + row;
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s0, c1.s0, c2.s0, c3.s0), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s1, c1.s1, c2.s1, c3.s1), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s2, c1.s2, c2.s2, c3.s2), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s3, c1.s3, c2.s3, c3.s3), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s4, c1.s4, c2.s4, c3.s4), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s5, c1.s5, c2.s5, c3.s5), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s6, c1.s6, c2.s6, c3.s6), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s7, c1.s7, c2.s7, c3.s7), 0, dst + idx);
+    }
+}
@@ -0,0 +1,145 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#ifdef cl_khr_integer_dot_product
+#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
+#endif
+
+#define QK_K 256
+
+#ifndef TILESIZE_N
+#define TILESIZE_N 32
+#endif
+
+inline uint iq2_xxs_pack(uint gv, uint sg, uint base) {
+    int v0 = (int)((gv >>  0) & 0xFF); if (sg & (1u << (base + 0))) { v0 = -v0; }
+    int v1 = (int)((gv >>  8) & 0xFF); if (sg & (1u << (base + 1))) { v1 = -v1; }
+    int v2 = (int)((gv >> 16) & 0xFF); if (sg & (1u << (base + 2))) { v2 = -v2; }
+    int v3 = (int)((gv >> 24) & 0xFF); if (sg & (1u << (base + 3))) { v3 = -v3; }
+    return ((uint)v0 & 0xFFu) | (((uint)v1 & 0xFFu) <<  8)
+         | (((uint)v2 & 0xFFu) << 16) | (((uint)v3 & 0xFFu) << 24);
+}
+
+// Pass staged operands by value; local dp4a operands miscompile on X2.
+inline int dot4_q8a_v(uint4 qw, uint4 a) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_iq2_xxs_q8_1_dp4a(
+        __global const uchar  * src0_qs,
+        __global const uint   * src0_sas,
+        __global const half   * src0_d,
+        __global const uint   * src1_qa,
+        __global const half   * src1_da,
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+
+#define IQ2_XXS_GRID(i) iq2_xxs_grid[(i)]
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+        const uint ib  = sub >> 3;
+
+        const uint  aux = src0_sas[rrow + sub * (uint)m];
+        const float dv  = (float)src0_d[rrow + ib * (uint)m];
+        const float dl  = dv * (0.5f + (float)(aux >> 28)) * 0.25f;
+
+        const uint gb = rrow + (step >> 3) * (uint)m;
+
+        uint4 qlo, qhi;
+        {
+            const uint gi0 = (uint)src0_qs[gb + 0u * (uint)m] << 1;
+            const uint gi1 = (uint)src0_qs[gb + 1u * (uint)m] << 1;
+            const uint gi2 = (uint)src0_qs[gb + 2u * (uint)m] << 1;
+            const uint gi3 = (uint)src0_qs[gb + 3u * (uint)m] << 1;
+            const uint s0  = iq2_xxs_signs((aux >>  0) & 127u);
+            const uint s1  = iq2_xxs_signs((aux >>  7) & 127u);
+            const uint s2  = iq2_xxs_signs((aux >> 14) & 127u);
+            const uint s3  = iq2_xxs_signs((aux >> 21) & 127u);
+            qlo.s0 = iq2_xxs_pack(IQ2_XXS_GRID(gi0 + 0u), s0, 0u);
+            qlo.s1 = iq2_xxs_pack(IQ2_XXS_GRID(gi0 + 1u), s0, 4u);
+            qlo.s2 = iq2_xxs_pack(IQ2_XXS_GRID(gi1 + 0u), s1, 0u);
+            qlo.s3 = iq2_xxs_pack(IQ2_XXS_GRID(gi1 + 1u), s1, 4u);
+            qhi.s0 = iq2_xxs_pack(IQ2_XXS_GRID(gi2 + 0u), s2, 0u);
+            qhi.s1 = iq2_xxs_pack(IQ2_XXS_GRID(gi2 + 1u), s2, 4u);
+            qhi.s2 = iq2_xxs_pack(IQ2_XXS_GRID(gi3 + 0u), s3, 0u);
+            qhi.s3 = iq2_xxs_pack(IQ2_XXS_GRID(gi3 + 1u), s3, 4u);
+        }
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? vload4(0, src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+#define IQ2_XXS_COL(b) (dl * (float)(dot4_q8a_v(qlo, (uint4)(sh_qa4[b][0]))   \
+                                  + dot4_q8a_v(qhi, (uint4)(sh_qa4[b][1]))))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+            rf.s0 = IQ2_XXS_COL(b+0);  rf.s1 = IQ2_XXS_COL(b+1);
+            rf.s2 = IQ2_XXS_COL(b+2);  rf.s3 = IQ2_XXS_COL(b+3);
+            acc[g] += LD4(sh_d, b) * rf;
+        }
+#undef IQ2_XXS_COL
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
+        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
+        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
+        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
+    }
+#undef NGROUPS
+#undef IQ2_XXS_GRID
+}
@@ -0,0 +1,92 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_qcom_reqd_sub_group_size
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+__attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+kernel void kernel_gemm_noshuffle_q3_k_f32(
+    global const uchar * src0_qs,
+    global const uchar * src0_hm,
+    global const uint  * src0_sc,
+    global const half  * src0_d,
+    read_only image1d_buffer_t src1,
+    global float * dst,
+    ulong offsetd,
+    int m,
+    int n,
+    int k,
+    int n_no_padding
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+    const int gy = get_global_id(0);
+    const int row = get_global_id(1) * 4;
+    const int n_4 = n / 4;
+    if (row >= m) {
+        return;
+    }
+
+    half8 c0 = 0, c1 = 0, c2 = 0, c3 = 0;
+    for (int i = 0; i < k; i += 16) {
+        const int is = (i / 16) & 15;
+        const int word = is >> 2;
+        const int shift = 8 * (is & 3);
+        const int sc_base = row + 3 * (i / 256) * m;
+        const uint4 low = vload4(0, src0_sc + sc_base + (word & 1) * m) >> (uint4)(4 * (word >> 1));
+        const uint4 high = vload4(0, src0_sc + sc_base + 2 * m) >> (uint4)(2 * word);
+        const int4 sc = convert_int4(((low >> (uint4)(shift)) & (uint4)(15)) | (((high >> (uint4)(shift)) & (uint4)(3)) << (uint4)(4))) - (int4)(32);
+        const half4 d = vload4(0, src0_d + row + (i / 256) * m);
+        const half4 scale = convert_half4(convert_float4(d) * convert_float4(sc));
+
+        for (int l = 0; l < 16; l += 4) {
+            const int ki = i + l;
+            const uchar4 bits = vload4(0, src0_qs + row + (ki / 4) * m);
+            const uchar4 hbits = vload4(0, src0_hm + row + (ki / 8) * m);
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                half8 b;
+                b.s0123 = read_imageh(src1, gy * 2     + (ki + j) * n_4);
+                b.s4567 = read_imageh(src1, gy * 2 + 1 + (ki + j) * n_4);
+                const int4 low = convert_int4((bits >> (uchar4)(2 * j)) & (uchar4)(3));
+                const int4 high = convert_int4((hbits >> (uchar4)((ki & 7) + j)) & (uchar4)(1));
+                const half4 w = convert_half4(low - (int4)(4) + (int4)(4) * high) * scale;
+                c0 += b * w.s0;
+                c1 += b * w.s1;
+                c2 += b * w.s2;
+                c3 += b * w.s3;
+            }
+        }
+    }
+
+    int idx = gy * 8 * m + row;
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s0, c1.s0, c2.s0, c3.s0), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s1, c1.s1, c2.s1, c3.s1), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s2, c1.s2, c2.s2, c3.s2), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s3, c1.s3, c2.s3, c3.s3), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s4, c1.s4, c2.s4, c3.s4), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s5, c1.s5, c2.s5, c3.s5), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s6, c1.s6, c2.s6, c3.s6), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s7, c1.s7, c2.s7, c3.s7), 0, dst + idx);
+    }
+}
@@ -0,0 +1,160 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#ifdef cl_khr_integer_dot_product
+#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
+#endif
+
+#define QK_K 256
+
+#ifndef TILESIZE_N
+#define TILESIZE_N 32
+#endif
+
+inline int q3_k_scale(uint sc0, uint sc1, uint sc2, uint is) {
+    const uint w  = is >> 2;
+    const uint sh = 8u * (is & 3u);
+    const uint tb = (sc2 >> sh) & 0xFFu;
+    uint v;
+    if      (w == 0u) { v = ( (sc0 >> sh)       & 0xFu) | (((tb >> 0) & 3u) << 4); }
+    else if (w == 1u) { v = ( (sc1 >> sh)       & 0xFu) | (((tb >> 2) & 3u) << 4); }
+    else if (w == 2u) { v = (((sc0 >> sh) >> 4) & 0xFu) | (((tb >> 4) & 3u) << 4); }
+    else              { v = (((sc1 >> sh) >> 4) & 0xFu) | (((tb >> 6) & 3u) << 4); }
+    return (int)v - 32;
+}
+
+inline uint q3_k_pack(uint pk, uint hb) {
+    int v0 = (int)((pk      ) & 3u) - 4 + 4*(int)((hb     ) & 1u);
+    int v1 = (int)((pk >> 2) & 3u) - 4 + 4*(int)((hb >> 1) & 1u);
+    int v2 = (int)((pk >> 4) & 3u) - 4 + 4*(int)((hb >> 2) & 1u);
+    int v3 = (int)((pk >> 6) & 3u) - 4 + 4*(int)((hb >> 3) & 1u);
+    return ((uint)v0 & 0xFFu) | (((uint)v1 & 0xFFu) <<  8)
+         | (((uint)v2 & 0xFFu) << 16) | (((uint)v3 & 0xFFu) << 24);
+}
+
+// Pass staged operands by value; local dp4a operands miscompile on X2.
+#define KQ_STAGE(p) vload4(0, (p))
+
+inline int dot4_q8a_v(uint4 qw, uint4 a) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q3_k_q8_1_dp4a(
+        __global const uchar  * src0_qs,
+        __global const uchar  * src0_hm,
+        __global const uint   * src0_sc,
+        __global const half   * src0_d,
+        __global const uint   * src1_qa,   // q8_1 activations int8 (as uint, 4/elem) [N, K]
+        __global const half   * src1_da,   // q8_1 per-block scale [N, K/32]
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,                          // output features (rows)
+        int    n_no_padding,               // tokens (cols)
+        int    k                           // K (== ne00)
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local half sh_d[TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+        const uint ib  = sub >> 3;                // super-block along K
+        const uint sb  = sub & 7u;
+
+        const uint scb = rrow + (3u * ib) * (uint)m;
+        const uint sc0 = src0_sc[scb + 0u * (uint)m];
+        const uint sc1 = src0_sc[scb + 1u * (uint)m];
+        const uint sc2 = src0_sc[scb + 2u * (uint)m];
+
+        const float d_w = (float)src0_d[rrow + ib * (uint)m];
+        const float dl0 = d_w * (float)q3_k_scale(sc0, sc1, sc2, 2u*sb + 0u);
+        const float dl1 = d_w * (float)q3_k_scale(sc0, sc1, sc2, 2u*sb + 1u);
+
+        const uint qsb = rrow + (step >> 2) * (uint)m;
+        const uint hmb = rrow + (step >> 3) * (uint)m;
+
+        uint4 qlo, qhi;
+        {
+            const uint h0 = (uint)src0_hm[hmb + 0u * (uint)m];
+            const uint h1 = (uint)src0_hm[hmb + 1u * (uint)m];
+            const uint h2 = (uint)src0_hm[hmb + 2u * (uint)m];
+            const uint h3 = (uint)src0_hm[hmb + 3u * (uint)m];
+            qlo.s0 = q3_k_pack((uint)src0_qs[qsb + 0u * (uint)m],  h0       & 0xFu);
+            qlo.s1 = q3_k_pack((uint)src0_qs[qsb + 1u * (uint)m], (h0 >> 4) & 0xFu);
+            qlo.s2 = q3_k_pack((uint)src0_qs[qsb + 2u * (uint)m],  h1       & 0xFu);
+            qlo.s3 = q3_k_pack((uint)src0_qs[qsb + 3u * (uint)m], (h1 >> 4) & 0xFu);
+            qhi.s0 = q3_k_pack((uint)src0_qs[qsb + 4u * (uint)m],  h2       & 0xFu);
+            qhi.s1 = q3_k_pack((uint)src0_qs[qsb + 5u * (uint)m], (h2 >> 4) & 0xFu);
+            qhi.s2 = q3_k_pack((uint)src0_qs[qsb + 6u * (uint)m],  h3       & 0xFu);
+            qhi.s3 = q3_k_pack((uint)src0_qs[qsb + 7u * (uint)m], (h3 >> 4) & 0xFu);
+        }
+
+        for (uint idx = lid; idx < TILESIZE_N * 2; idx += 64) {
+            const uint t = idx >> 1;
+            const uint v = idx & 1;
+            const uint c = col_base + t;
+            sh_qa4[t][v] = (c < (uint)n_no_padding)
+                         ? KQ_STAGE(src1_qa + c * k_u + (step >> 2) + (v << 2))
+                         : (uint4)(0u);
+        }
+        if (lid < TILESIZE_N) {
+            const uint c = col_base + lid;
+            sh_d[lid] = (c < (uint)n_no_padding) ? src1_da[c * k_b + sub] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+#define Q3_K_COL(T) (dl0 * (float)dot4_q8a_v(qlo, (uint4)(sh_qa4[T][0]))  \
+                  + dl1 * (float)dot4_q8a_v(qhi, (uint4)(sh_qa4[T][1])))
+            rf.s0 = Q3_K_COL(b+0);  rf.s1 = Q3_K_COL(b+1);
+            rf.s2 = Q3_K_COL(b+2);  rf.s3 = Q3_K_COL(b+3);
+#undef Q3_K_COL
+            acc[g] += LD4(sh_d, b) * rf;
+        }
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
+        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
+        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
+        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
+    }
+#undef NGROUPS
+}
@@ -0,0 +1,131 @@
+// DP4A GEMM for the canonical 64-row Q4_K layout.
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#ifdef cl_khr_integer_dot_product
+#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
+#endif
+
+#ifdef cl_qcom_reqd_sub_group_size
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define ADRENO_GPU 1
+#define REQD_SUBGROUP_SIZE_64 __attribute__((qcom_reqd_sub_group_size("half")))
+#endif
+
+#define QK_K       256
+#define TILE_ROWS  64
+#define NTILES     4
+#define BN_DP      8
+#define WG_THREADS (NTILES * TILE_ROWS)
+
+inline void q4k_scale_min(int j, __global const uchar * q, uchar * d, uchar * m) {
+    if (j < 4) {
+        *d = q[j]   & 0x3F;
+        *m = q[j+4] & 0x3F;
+    } else {
+        *d = (q[j+4] & 0x0F) | ((q[j-4] & 0xC0) >> 2);
+        *m = ((q[j+4] >> 4) & 0x0F) | ((q[j] & 0xC0) >> 2);
+    }
+}
+
+// dp4a GEMM over the tiled q4_K layout (ne1 > 1) with a q8_1 activation; per 32-element sub-block:
+//   acc += d*sd*qad*dp4a(q, qa) - dm*sm*qas
+#if defined(ADRENO_GPU)
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_gemm_noshuffle_q4_K_f32_tiled_dp4a(
+    __global uint4 * src0_q,
+    __global half  * src0_d,
+    __global half  * src0_dm,
+    __global uchar * src0_s,
+    __global uint  * qa,        // q8_1 int8 activation, 4/uint  [ne11][ne00/4]
+    __global half  * qad,       // q8_1 per-32-block scale       [ne11][ne00/32]
+    __global half  * qas,       // q8_1 per-32-block sum (qad*Sum)[ne11][ne00/32]
+    global float * dst,
+    ulong offsetd,
+    int ne00,
+    int ne01,
+    int ne11
+) {
+    int rit = get_local_id(0);
+    int sg  = get_local_id(1);
+    int lid = sg * TILE_ROWS + rit;
+    int row = get_group_id(0) * WG_THREADS + lid;
+    int rt  = row / TILE_ROWS;
+    int col0 = get_global_id(2) * BN_DP;
+    int nb  = ne00 / QK_K;
+    int qcs = ne00 / 4;
+    int dcs = ne00 / 32;
+
+    const bool row_ok = row < ne01;
+    __local uint sh_qa[BN_DP * 64];
+    __local half sh_qad[BN_DP * 8];
+    __local half sh_qas[BN_DP * 8];
+
+    float acc[BN_DP];
+    #pragma unroll
+    for (int j = 0; j < BN_DP; ++j) acc[j] = 0.0f;
+
+    for (int sb = 0; sb < nb; ++sb) {
+        for (int p = lid; p < BN_DP * 64; p += WG_THREADS) {
+            int j = p >> 6, u = p & 63, c = col0 + j;
+            sh_qa[p] = (c < ne11) ? qa[c * qcs + sb * 64 + u] : 0u;
+        }
+        for (int p = lid; p < BN_DP * 8; p += WG_THREADS) {
+            int j = p >> 3, b = p & 7, c = col0 + j;
+            sh_qad[p] = (c < ne11) ? qad[c * dcs + sb * 8 + b] : (half)0;
+            sh_qas[p] = (c < ne11) ? qas[c * dcs + sb * 8 + b] : (half)0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+        if (row_ok) {
+            int tile_blk = rt * nb + sb;
+            float dval  = (float)src0_d [tile_blk * TILE_ROWS + rit];
+            float dmval = (float)src0_dm[tile_blk * TILE_ROWS + rit];
+            __global uchar * sc = src0_s + (tile_blk * TILE_ROWS + rit) * 12;
+
+            uint q[32];
+            #pragma unroll
+            for (int g = 0; g < 8; ++g) {
+                uint4 v = src0_q[(tile_blk * 8 + g) * TILE_ROWS + rit];
+                q[g*4+0] = v.x; q[g*4+1] = v.y; q[g*4+2] = v.z; q[g*4+3] = v.w;
+            }
+
+            for (int is = 0; is < 8; ++is) {
+                uchar sd, sm;
+                q4k_scale_min(is, sc, &sd, &sm);
+                float scl = dval * (float)sd, mnv = dmval * (float)sm;
+
+                uint cw[8];
+                #pragma unroll
+                for (int t = 0; t < 8; ++t) {
+                    uint pk = 0;
+                    #pragma unroll
+                    for (int u = 0; u < 4; ++u) {
+                        int e = is * 32 + t * 4 + u;
+                        uint code = (q[e >> 3] >> ((e & 7) * 4)) & 0xF;
+                        pk |= (code & 0xFF) << (u * 8);
+                    }
+                    cw[t] = pk;
+                }
+                #pragma unroll
+                for (int j = 0; j < BN_DP; ++j) {
+                    __local uint * a = sh_qa + j * 64 + is * 8;
+                    int raw = 0;
+                    #pragma unroll
+                    for (int t = 0; t < 8; ++t) raw = dot_acc_sat_4x8packed_ss_int(cw[t], a[t], raw);
+                    acc[j] += scl * (float)sh_qad[j*8 + is] * (float)raw
+                            - mnv * (float)sh_qas[j*8 + is];
+                }
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (row_ok) {
+        dst = (global float*)((global char*)dst + offsetd);
+        #pragma unroll
+        for (int j = 0; j < BN_DP; ++j) {
+            int c = col0 + j;
+            if (c < ne11) dst[(ulong)c * ne01 + row] = acc[j];
+        }
+    }
+}
@@ -0,0 +1,122 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#define QK_K 256
+
+#ifndef IQ2_S_MV_NSG
+#define IQ2_S_MV_NSG 8
+#endif
+
+inline float4 iq2_s_vals(uint gv, uint sg, uint base) {
+
+    const uint  s   = sg >> base;
+    const uint4 sgn = (uint4)(s << 31, s << 30, s << 29, s << 28) & 0x80000000u;
+    return as_float4(as_uint4(convert_float4(as_uchar4(gv))) ^ sgn);
+}
+
+kernel void kernel_iq2_s_grid_export(global uint * out) {
+    out[get_global_id(0)] = iq2_s_grid[get_global_id(0)];
+}
+
+kernel void kernel_gemv_noshuffle_iq2_s_f32(
+        read_only image1d_buffer_t grid_img,
+        global const uchar * src0_qs,
+        global const uchar * src0_sg,
+        global const uchar * src0_qh,
+        global const uchar * src0_sc,
+        global const half  * src0_d,
+        global const float * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne10,
+        int ne0
+) {
+    src1 = (global const float *)((global const char *)src1 + offset1);
+    dst  = (global float       *)((global char       *)dst  + offsetd);
+
+    const uint m   = (uint)ne01;
+    const uint K   = (uint)ne00;
+    const uint nsb = K / QK_K;
+
+    const uint lid = get_local_id(0);
+    const uint sgi = get_local_id(1);
+    const uint col = get_group_id(1);
+
+    global const float * y = src1 + (ulong)col * (uint)ne10;
+
+#define IQ2_S_GRID(i) read_imageui(grid_img, (int)(i)).x
+
+    const uint mh  = m >> 1;
+    const uint j   = get_group_id(0) * 64u + lid;
+    const uint row = j << 1;
+
+    float sumf = 0.f, sumf1 = 0.f;
+
+    if (j < mh) {
+        global const ushort * qsu = (global const ushort *)src0_qs;
+        global const ushort * sgu = (global const ushort *)src0_sg;
+        global const ushort * qhu = (global const ushort *)src0_qh;
+        global const ushort * scu = (global const ushort *)src0_sc;
+
+        for (uint ib = sgi; ib < nsb; ib += IQ2_S_MV_NSG) {
+            const half2 dh = vload2(j + ib * mh, src0_d);
+
+            float acc0 = 0.f, acc1 = 0.f;
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint sub = ib * 8u + sb;
+                const uint qhv = (uint)qhu[j + sub * mh];
+                const uint scv = (uint)scu[j + sub * mh];
+                const uint gb  = j + (sub * 4u) * mh;
+
+                for (uint h = 0; h < 2u; ++h) {
+                    const uint n0 = (h == 0u) ? ( scv        & 0xFu) : (( scv        >> 4) & 0xFu);
+                    const uint n1 = (h == 0u) ? ((scv >> 8)  & 0xFu) : (((scv >> 8)  >> 4) & 0xFu);
+
+                    float a0 = 0.f, a1 = 0.f;
+                    for (uint t = 0; t < 2u; ++t) {
+                        const uint l   = 2u*h + t;
+                        const uint qsv = (uint)qsu[gb + l * mh];
+                        const uint sgv = (uint)sgu[gb + l * mh];
+                        const uint gi0 = ( qsv       & 0xFFu) | ((((qhv      ) >> (2u*l)) & 3u) << 8);
+                        const uint gi1 = ((qsv >> 8) & 0xFFu) | ((((qhv >> 8) >> (2u*l)) & 3u) << 8);
+                        const uint s0  =  sgv       & 0xFFu;
+                        const uint s1  = (sgv >> 8) & 0xFFu;
+
+                        const uint grp = (ib * 64u + sb * 8u) + l * 2u;
+                        const float4 y0 = vload4(grp + 0u, y);
+                        const float4 y1 = vload4(grp + 1u, y);
+                        a0 += dot(y0, iq2_s_vals(IQ2_S_GRID(2u*gi0 + 0u), s0, 0u));
+                        a0 += dot(y1, iq2_s_vals(IQ2_S_GRID(2u*gi0 + 1u), s0, 4u));
+                        a1 += dot(y0, iq2_s_vals(IQ2_S_GRID(2u*gi1 + 0u), s1, 0u));
+                        a1 += dot(y1, iq2_s_vals(IQ2_S_GRID(2u*gi1 + 1u), s1, 4u));
+                    }
+                    acc0 += (0.5f + (float)n0) * a0;
+                    acc1 += (0.5f + (float)n1) * a1;
+                }
+            }
+            sumf  += (float)dh.s0 * 0.25f * acc0;
+            sumf1 += (float)dh.s1 * 0.25f * acc1;
+        }
+    }
+
+#if IQ2_S_MV_NSG > 1
+    __local float2 part[IQ2_S_MV_NSG][64];
+    part[sgi][lid] = (float2)(sumf, sumf1);
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < IQ2_S_MV_NSG; ++s) {
+        const float2 p = part[s][lid];
+        sumf  += p.s0;
+        sumf1 += p.s1;
+    }
+#endif
+
+    if (j < mh) {
+        vstore2((float2)(sumf, sumf1), 0, dst + (ulong)col * (uint)ne0 + row);
+    }
+#undef IQ2_S_GRID
+}
@@ -0,0 +1,128 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#define QK_K 256
+
+#ifndef IQ2_XS_MV_NSG
+#define IQ2_XS_MV_NSG 4
+#endif
+
+inline float4 iq2_xs_vals(uint gv, uint sgv, uint base) {
+    const uint  s   = sgv >> base;
+    const uint4 sgn = (uint4)(s << 31, s << 30, s << 29, s << 28) & 0x80000000u;
+    return as_float4(as_uint4(convert_float4(as_uchar4(gv))) ^ sgn);
+}
+
+kernel void kernel_gemv_noshuffle_iq2_xs_f32(
+        __read_only image1d_buffer_t grid_img,
+        global const ushort * src0_qs,
+        global const uchar  * src0_sc,
+        global const half   * src0_d,
+        global const float  * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne10,
+        int ne0
+) {
+    src1 = (global const float *)((global const char *)src1 + offset1);
+    dst  = (global float       *)((global char       *)dst  + offsetd);
+
+    const uint m   = (uint)ne01;
+    const uint K   = (uint)ne00;
+    const uint nsb = K / QK_K;
+
+    const uint lid = get_local_id(0);
+    const uint sgi = get_local_id(1);
+    const uint col = get_group_id(1);
+
+    global const float * y = src1 + (ulong)col * (uint)ne10;
+
+#define IQ2_XS_GRID(i) (read_imageui(grid_img, (int)(i)).x)
+
+    const uint mh  = m >> 1;
+    const uint j   = get_group_id(0) * 64u + lid;
+    const uint row = j << 1;
+
+    float sumf  = 0.f;
+    float sumf1 = 0.f;
+
+    if (j < mh) {
+        global const uint   * qsu = (global const uint   *)src0_qs;
+        global const ushort * scu = (global const ushort *)src0_sc;
+
+        for (uint ib = sgi; ib < nsb; ib += IQ2_XS_MV_NSG) {
+            const half2 dh = vload2(j + ib * mh, src0_d);
+            const float d0 = (float)dh.s0;
+            const float d1 = (float)dh.s1;
+
+            float acc0 = 0.f, acc1 = 0.f;
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint sub = ib * 8u + sb;
+                const uint scv = (uint)scu[j + sub * mh];
+
+                const float lo0 = 0.5f + (float)( scv        & 0xFu);
+                const float hi0 = 0.5f + (float)((scv >>  4) & 0xFu);
+                const float lo1 = 0.5f + (float)((scv >>  8) & 0xFu);
+                const float hi1 = 0.5f + (float)((scv >> 12) & 0xFu);
+
+                const uint grp = sub * 4u;
+                const uint qsb = j + grp * mh;
+
+                float alo0 = 0.f, ahi0 = 0.f, alo1 = 0.f, ahi1 = 0.f;
+                for (uint u = 0; u < 4u; ++u) {
+                    const uint qsv = qsu[qsb + u * mh];
+                    const uint q0  = qsv & 0xFFFFu;
+                    const uint q1  = qsv >> 16;
+
+                    const uint g0  = (q0 & 511u) << 1;
+                    const uint g1  = (q1 & 511u) << 1;
+                    const uint sg0 = iq2_xs_signs(q0 >> 9);
+                    const uint sg1 = iq2_xs_signs(q1 >> 9);
+
+                    const float4 yl = vload4(sub * 8u + u * 2u + 0u, y);
+                    const float4 yh = vload4(sub * 8u + u * 2u + 1u, y);
+
+                    const float a0 = dot(yl, iq2_xs_vals(IQ2_XS_GRID(g0    ), sg0, 0u))
+                                   + dot(yh, iq2_xs_vals(IQ2_XS_GRID(g0 + 1), sg0, 4u));
+                    const float a1 = dot(yl, iq2_xs_vals(IQ2_XS_GRID(g1    ), sg1, 0u))
+                                   + dot(yh, iq2_xs_vals(IQ2_XS_GRID(g1 + 1), sg1, 4u));
+
+                    if (u < 2u) { alo0 += a0; alo1 += a1; }
+                    else        { ahi0 += a0; ahi1 += a1; }
+                }
+                acc0 += lo0 * alo0 + hi0 * ahi0;
+                acc1 += lo1 * alo1 + hi1 * ahi1;
+            }
+            sumf  += d0 * 0.25f * acc0;
+            sumf1 += d1 * 0.25f * acc1;
+        }
+    }
+
+#if IQ2_XS_MV_NSG > 1
+    __local float2 part[IQ2_XS_MV_NSG][64];
+    part[sgi][lid] = (float2)(sumf, sumf1);
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < IQ2_XS_MV_NSG; ++s) {
+        const float2 p = part[s][lid];
+        sumf  += p.s0;
+        sumf1 += p.s1;
+    }
+#endif
+
+    if (j < mh) {
+        vstore2((float2)(sumf, sumf1), 0, dst + (ulong)col * (uint)ne0 + row);
+    }
+#undef IQ2_XS_GRID
+}
+
+kernel void kernel_iq2_xs_grid_export(global uint * out) {
+    const uint i = get_global_id(0);
+    if (i < 1024u) {
+        out[i] = iq2_xs_grid[i];
+    }
+}
@@ -0,0 +1,117 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#define QK_K 256
+
+#ifndef IQ2_XXS_MV_NSG
+#define IQ2_XXS_MV_NSG 4
+#endif
+
+inline float4 iq2_xxs_vals(uint gv, uint sgv, uint base) {
+    const uint  s   = sgv >> base;
+    const uint4 sgn = (uint4)(s << 31, s << 30, s << 29, s << 28) & 0x80000000u;
+    return as_float4(as_uint4(convert_float4(as_uchar4(gv))) ^ sgn);
+}
+
+kernel void kernel_gemv_noshuffle_iq2_xxs_f32(
+        __read_only image1d_buffer_t grid_img,
+        global const uchar * src0_qs,
+        global const uint  * src0_sas,
+        global const half  * src0_d,
+        global const float * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne10,
+        int ne0
+) {
+    src1 = (global const float *)((global const char *)src1 + offset1);
+    dst  = (global float       *)((global char       *)dst  + offsetd);
+
+    const uint m   = (uint)ne01;
+    const uint K   = (uint)ne00;
+    const uint nsb = K / QK_K;
+
+    const uint lid = get_local_id(0);
+    const uint sgi = get_local_id(1);
+    const uint col = get_group_id(1);
+
+    global const float * y = src1 + (ulong)col * (uint)ne10;
+
+#define IQ2_XXS_GRID(i) (read_imageui(grid_img, (int)(i)).x)
+
+    const uint mh  = m >> 1;
+    const uint j   = get_group_id(0) * 64u + lid;
+    const uint row = j << 1;
+
+    float sumf  = 0.f;
+    float sumf1 = 0.f;
+
+    if (j < mh) {
+        global const ushort * qsu = (global const ushort *)src0_qs;
+
+        for (uint ib = sgi; ib < nsb; ib += IQ2_XXS_MV_NSG) {
+            const half2 dh = vload2(j + ib * mh, src0_d);
+            const float d0 = (float)dh.s0;
+            const float d1 = (float)dh.s1;
+
+            float acc0 = 0.f, acc1 = 0.f;
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint  sub = ib * 8u + sb;
+                const uint2 aux = vload2(j + sub * mh, src0_sas);
+
+                const uint grp = sub * 4u;
+                const uint qsb = j + grp * mh;
+
+                float a0 = 0.f, a1 = 0.f;
+                for (uint u = 0; u < 4u; ++u) {
+                    const uint qsv = (uint)qsu[qsb + u * mh];
+                    const uint sg0 = iq2_xxs_signs((aux.s0 >> (7u * u)) & 127u);
+                    const uint sg1 = iq2_xxs_signs((aux.s1 >> (7u * u)) & 127u);
+
+                    const uint g0 = ( qsv       & 0xFFu) << 1;
+                    const uint g1 = ((qsv >> 8) & 0xFFu) << 1;
+
+                    const float4 yl = vload4(sub * 8u + u * 2u + 0u, y);
+                    const float4 yh = vload4(sub * 8u + u * 2u + 1u, y);
+
+                    a0 += dot(yl, iq2_xxs_vals(IQ2_XXS_GRID(g0    ), sg0, 0u))
+                        + dot(yh, iq2_xxs_vals(IQ2_XXS_GRID(g0 + 1), sg0, 4u));
+                    a1 += dot(yl, iq2_xxs_vals(IQ2_XXS_GRID(g1    ), sg1, 0u))
+                        + dot(yh, iq2_xxs_vals(IQ2_XXS_GRID(g1 + 1), sg1, 4u));
+                }
+                acc0 += (0.5f + (float)(aux.s0 >> 28)) * a0;
+                acc1 += (0.5f + (float)(aux.s1 >> 28)) * a1;
+            }
+            sumf  += d0 * 0.25f * acc0;
+            sumf1 += d1 * 0.25f * acc1;
+        }
+    }
+
+#if IQ2_XXS_MV_NSG > 1
+    __local float2 part[IQ2_XXS_MV_NSG][64];
+    part[sgi][lid] = (float2)(sumf, sumf1);
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < IQ2_XXS_MV_NSG; ++s) {
+        const float2 p = part[s][lid];
+        sumf  += p.s0;
+        sumf1 += p.s1;
+    }
+#endif
+
+    if (j < mh) {
+        vstore2((float2)(sumf, sumf1), 0, dst + (ulong)col * (uint)ne0 + row);
+    }
+#undef IQ2_XXS_GRID
+}
+
+kernel void kernel_iq2_xxs_grid_export(global uint * out) {
+    const uint i = get_global_id(0);
+    if (i < 512u) {
+        out[i] = iq2_xxs_grid[i];
+    }
+}
@@ -0,0 +1,249 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#define QK_K 256
+
+#ifndef Q3_K_MV_NSG
+#define Q3_K_MV_NSG 8
+#endif
+
+#ifndef Q3_K_MV_R
+#define Q3_K_MV_R 2
+#endif
+
+inline int q3_k_scale(uint sc0, uint sc1, uint sc2, uint is) {
+    const uint w  = is >> 2;
+    const uint sh = 8u * (is & 3u);
+    const uint tb = (sc2 >> sh) & 0xFFu;
+    uint v;
+    if      (w == 0u) { v = ( (sc0 >> sh)       & 0xFu) | (((tb >> 0) & 3u) << 4); }
+    else if (w == 1u) { v = ( (sc1 >> sh)       & 0xFu) | (((tb >> 2) & 3u) << 4); }
+    else if (w == 2u) { v = (((sc0 >> sh) >> 4) & 0xFu) | (((tb >> 4) & 3u) << 4); }
+    else              { v = (((sc1 >> sh) >> 4) & 0xFu) | (((tb >> 6) & 3u) << 4); }
+    return (int)v - 32;
+}
+
+inline float4 q3_k_vals(uint pk, uint hb) {
+    float4 v;
+    v.s0 = (float)((int)((pk      ) & 3u) - 4 + 4*(int)((hb     ) & 1u));
+    v.s1 = (float)((int)((pk >> 2) & 3u) - 4 + 4*(int)((hb >> 1) & 1u));
+    v.s2 = (float)((int)((pk >> 4) & 3u) - 4 + 4*(int)((hb >> 2) & 1u));
+    v.s3 = (float)((int)((pk >> 6) & 3u) - 4 + 4*(int)((hb >> 3) & 1u));
+    return v;
+}
+
+kernel void kernel_gemv_noshuffle_q3_k_f32(
+        global const uchar * src0_qs,
+        global const uchar * src0_hm,
+        global const uint  * src0_sc,
+        global const half  * src0_d,
+        global const float * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,      // K
+        int ne01,      // M
+        int ne10,      // activation row stride
+        int ne0        // dst row stride
+) {
+    src1 = (global const float *)((global const char *)src1 + offset1);
+    dst  = (global float       *)((global char       *)dst  + offsetd);
+
+    const uint m   = (uint)ne01;
+    const uint K   = (uint)ne00;
+    const uint nsb = K / QK_K;
+
+    const uint lid = get_local_id(0);
+    const uint sgi = get_local_id(1);
+    const uint col = get_group_id(1);
+
+    global const float * y = src1 + (ulong)col * (uint)ne10;
+
+#if Q3_K_MV_R == 4
+    const uint mq  = m >> 2;
+    const uint j   = get_group_id(0) * 64u + lid;
+    const uint row = j << 2;
+
+    float sumf = 0.f, sumf1 = 0.f, sumf2 = 0.f, sumf3 = 0.f;
+
+    if (j < mq) {
+        global const uint * qsw = (global const uint *)src0_qs;
+        global const uint * hmw = (global const uint *)src0_hm;
+
+        for (uint ib = sgi; ib < nsb; ib += Q3_K_MV_NSG) {
+            const half4 dh = vload4(j + ib * mq, src0_d);
+            const uint4 s0 = vload4(j + (3u * ib + 0u) * mq, src0_sc);
+            const uint4 s1 = vload4(j + (3u * ib + 1u) * mq, src0_sc);
+            const uint4 s2 = vload4(j + (3u * ib + 2u) * mq, src0_sc);
+
+            float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f, acc3 = 0.f;
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint grp = ib * 64u + sb * 8u;
+                const uint qsb = j + grp * mq;
+                const uint hmb = j + (grp >> 1) * mq;
+
+                for (uint h = 0; h < 2u; ++h) {         // two 16-weight halves
+                    const int l0 = q3_k_scale(s0.s0, s1.s0, s2.s0, 2u*sb + h);
+                    const int l1 = q3_k_scale(s0.s1, s1.s1, s2.s1, 2u*sb + h);
+                    const int l2 = q3_k_scale(s0.s2, s1.s2, s2.s2, 2u*sb + h);
+                    const int l3 = q3_k_scale(s0.s3, s1.s3, s2.s3, 2u*sb + h);
+
+                    float a0 = 0.f, a1 = 0.f, a2 = 0.f, a3 = 0.f;
+                    for (uint u = 0; u < 4u; ++u) {
+                        const uint gg  = 4u*h + u;
+                        const uint qsv = qsw[qsb + gg * mq];        // four rows, one load
+                        const uint hmv = hmw[hmb + (gg >> 1) * mq];
+                        const uint hsh = 4u * (gg & 1u);
+                        const float4 yv = vload4(grp + gg, y);      // shared by all four
+                        a0 += dot(yv, q3_k_vals((qsv      ) & 0xFFu, ((hmv      ) >> hsh) & 0xFu));
+                        a1 += dot(yv, q3_k_vals((qsv >>  8) & 0xFFu, ((hmv >>  8) >> hsh) & 0xFu));
+                        a2 += dot(yv, q3_k_vals((qsv >> 16) & 0xFFu, ((hmv >> 16) >> hsh) & 0xFu));
+                        a3 += dot(yv, q3_k_vals((qsv >> 24) & 0xFFu, ((hmv >> 24) >> hsh) & 0xFu));
+                    }
+                    acc0 += (float)l0 * a0;
+                    acc1 += (float)l1 * a1;
+                    acc2 += (float)l2 * a2;
+                    acc3 += (float)l3 * a3;
+                }
+            }
+            sumf  += (float)dh.s0 * acc0;
+            sumf1 += (float)dh.s1 * acc1;
+            sumf2 += (float)dh.s2 * acc2;
+            sumf3 += (float)dh.s3 * acc3;
+        }
+    }
+#elif Q3_K_MV_R == 2
+    const uint mh  = m >> 1;
+    const uint j   = get_group_id(0) * 64u + lid;
+    const uint row = j << 1;
+
+    float sumf  = 0.f;
+    float sumf1 = 0.f;
+
+    if (j < mh) {
+        global const ushort * qsu = (global const ushort *)src0_qs;
+        global const ushort * hmu = (global const ushort *)src0_hm;
+
+        for (uint ib = sgi; ib < nsb; ib += Q3_K_MV_NSG) {
+            const half2 dh = vload2(j + ib * mh, src0_d);
+            const uint2 s0 = vload2(j + (3u * ib + 0u) * mh, src0_sc);
+            const uint2 s1 = vload2(j + (3u * ib + 1u) * mh, src0_sc);
+            const uint2 s2 = vload2(j + (3u * ib + 2u) * mh, src0_sc);
+
+            float acc0 = 0.f, acc1 = 0.f;
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint grp = ib * 64u + sb * 8u;
+                const uint qsb = j + grp * mh;
+                const uint hmb = j + (grp >> 1) * mh;
+
+                for (uint h = 0; h < 2u; ++h) {         // two 16-weight halves
+                    const int l0 = q3_k_scale(s0.s0, s1.s0, s2.s0, 2u*sb + h);
+                    const int l1 = q3_k_scale(s0.s1, s1.s1, s2.s1, 2u*sb + h);
+
+                    float a0 = 0.f, a1 = 0.f;
+                    for (uint u = 0; u < 4u; ++u) {
+                        const uint gg  = 4u*h + u;
+                        const uint qsv = (uint)qsu[qsb + gg * mh];
+                        const uint hmv = (uint)hmu[hmb + (gg >> 1) * mh];
+                        const uint hsh = 4u * (gg & 1u);
+                        const float4 yv = vload4(grp + gg, y);
+                        a0 += dot(yv, q3_k_vals( qsv       & 0xFFu, ( hmv        >> hsh) & 0xFu));
+                        a1 += dot(yv, q3_k_vals((qsv >> 8) & 0xFFu, ((hmv >> 8)  >> hsh) & 0xFu));
+                    }
+                    acc0 += (float)l0 * a0;
+                    acc1 += (float)l1 * a1;
+                }
+            }
+            sumf  += (float)dh.s0 * acc0;
+            sumf1 += (float)dh.s1 * acc1;
+        }
+    }
+#else
+    const uint row = get_group_id(0) * 64u + lid;
+
+    float sumf = 0.f;
+
+    if (row < m) {
+        for (uint ib = sgi; ib < nsb; ib += Q3_K_MV_NSG) {
+            const float d  = (float)src0_d[row + ib * m];
+            const uint  s0 = src0_sc[row + (3u * ib + 0u) * m];
+            const uint  s1 = src0_sc[row + (3u * ib + 1u) * m];
+            const uint  s2 = src0_sc[row + (3u * ib + 2u) * m];
+
+            float acc = 0.f;
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint grp = ib * 64u + sb * 8u;
+                const uint qsb = row + grp * m;
+                const uint hmb = row + (grp >> 1) * m;
+
+                for (uint h = 0; h < 2u; ++h) {
+                    const int l = q3_k_scale(s0, s1, s2, 2u*sb + h);
+                    float a = 0.f;
+                    for (uint u = 0; u < 4u; ++u) {
+                        const uint gg  = 4u*h + u;
+                        const uint qsv = (uint)src0_qs[qsb + gg * m];
+                        const uint hmv = (uint)src0_hm[hmb + (gg >> 1) * m];
+                        const float4 yv = vload4(grp + gg, y);
+                        a += dot(yv, q3_k_vals(qsv, (hmv >> (4u * (gg & 1u))) & 0xFu));
+                    }
+                    acc += (float)l * a;
+                }
+            }
+            sumf += d * acc;
+        }
+    }
+#endif
+
+#if Q3_K_MV_NSG > 1
+#if Q3_K_MV_R == 4
+    __local float4 part[Q3_K_MV_NSG][64];
+    part[sgi][lid] = (float4)(sumf, sumf1, sumf2, sumf3);
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < Q3_K_MV_NSG; ++s) {
+        const float4 p = part[s][lid];
+        sumf  += p.s0;
+        sumf1 += p.s1;
+        sumf2 += p.s2;
+        sumf3 += p.s3;
+    }
+#elif Q3_K_MV_R == 2
+    __local float2 part[Q3_K_MV_NSG][64];
+    part[sgi][lid] = (float2)(sumf, sumf1);
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < Q3_K_MV_NSG; ++s) {
+        const float2 p = part[s][lid];
+        sumf  += p.s0;
+        sumf1 += p.s1;
+    }
+#else
+    __local float part[Q3_K_MV_NSG][64];
+    part[sgi][lid] = sumf;
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < Q3_K_MV_NSG; ++s) {
+        sumf += part[s][lid];
+    }
+#endif
+#endif
+
+#if Q3_K_MV_R == 4
+    if (j < mq) {
+        vstore4((float4)(sumf, sumf1, sumf2, sumf3), 0, dst + (ulong)col * (uint)ne0 + row);
+    }
+#elif Q3_K_MV_R == 2
+    if (j < mh) {
+        vstore2((float2)(sumf, sumf1), 0, dst + (ulong)col * (uint)ne0 + row);
+    }
+#else
+    if (row < m) {
+        dst[(ulong)col * (uint)ne0 + row] = sumf;
+    }
+#endif
+}
@@ -0,0 +1,259 @@
+// IQ2_S codebook, stored as pairs of four packed magnitudes.
+constant uint iq2_s_grid[2048] = {
+    0x08080808, 0x08080808, 0x0808082b, 0x08080808, 0x08081919, 0x08080808, 0x08082b08, 0x08080808,
+    0x08082b2b, 0x08080808, 0x08190819, 0x08080808, 0x08191908, 0x08080808, 0x0819192b, 0x08080808,
+    0x08192b19, 0x08080808, 0x082b0808, 0x08080808, 0x082b082b, 0x08080808, 0x082b1919, 0x08080808,
+    0x082b2b08, 0x08080808, 0x19080819, 0x08080808, 0x19081908, 0x08080808, 0x1908192b, 0x08080808,
+    0x19082b19, 0x08080808, 0x19190808, 0x08080808, 0x1919082b, 0x08080808, 0x19191919, 0x08080808,
+    0x19192b08, 0x08080808, 0x192b0819, 0x08080808, 0x192b1908, 0x08080808, 0x192b192b, 0x08080808,
+    0x192b2b19, 0x08080808, 0x2b080808, 0x08080808, 0x2b08082b, 0x08080808, 0x2b081919, 0x08080808,
+    0x2b082b08, 0x08080808, 0x2b190819, 0x08080808, 0x2b191908, 0x08080808, 0x2b2b0808, 0x08080808,
+    0x2b2b1919, 0x08080808, 0x2b2b2b2b, 0x08080808, 0x08080819, 0x08080819, 0x08081908, 0x08080819,
+    0x0808192b, 0x08080819, 0x08082b19, 0x08080819, 0x08190808, 0x08080819, 0x0819082b, 0x08080819,
+    0x08191919, 0x08080819, 0x08192b08, 0x08080819, 0x082b0819, 0x08080819, 0x082b1908, 0x08080819,
+    0x19080808, 0x08080819, 0x1908082b, 0x08080819, 0x19081919, 0x08080819, 0x19082b08, 0x08080819,
+    0x19190819, 0x08080819, 0x19191908, 0x08080819, 0x1919192b, 0x08080819, 0x19192b19, 0x08080819,
+    0x192b0808, 0x08080819, 0x192b1919, 0x08080819, 0x192b2b08, 0x08080819, 0x2b080819, 0x08080819,
+    0x2b081908, 0x08080819, 0x2b190808, 0x08080819, 0x2b19082b, 0x08080819, 0x2b191919, 0x08080819,
+    0x2b2b0819, 0x08080819, 0x2b2b1908, 0x08080819, 0x08080808, 0x0808082b, 0x0808082b, 0x0808082b,
+    0x08081919, 0x0808082b, 0x08082b08, 0x0808082b, 0x08190819, 0x0808082b, 0x08191908, 0x0808082b,
+    0x082b0808, 0x0808082b, 0x082b2b2b, 0x0808082b, 0x19080819, 0x0808082b, 0x19081908, 0x0808082b,
+    0x1908192b, 0x0808082b, 0x19082b19, 0x0808082b, 0x19190808, 0x0808082b, 0x19191919, 0x0808082b,
+    0x2b080808, 0x0808082b, 0x2b081919, 0x0808082b, 0x2b082b2b, 0x0808082b, 0x2b191908, 0x0808082b,
+    0x2b2b082b, 0x0808082b, 0x08080819, 0x08081908, 0x08081908, 0x08081908, 0x0808192b, 0x08081908,
+    0x08082b19, 0x08081908, 0x08190808, 0x08081908, 0x0819082b, 0x08081908, 0x08191919, 0x08081908,
+    0x08192b08, 0x08081908, 0x082b0819, 0x08081908, 0x082b1908, 0x08081908, 0x082b192b, 0x08081908,
+    0x082b2b19, 0x08081908, 0x19080808, 0x08081908, 0x1908082b, 0x08081908, 0x19081919, 0x08081908,
+    0x19082b08, 0x08081908, 0x19082b2b, 0x08081908, 0x19190819, 0x08081908, 0x19191908, 0x08081908,
+    0x1919192b, 0x08081908, 0x19192b19, 0x08081908, 0x192b0808, 0x08081908, 0x192b082b, 0x08081908,
+    0x192b1919, 0x08081908, 0x2b080819, 0x08081908, 0x2b081908, 0x08081908, 0x2b08192b, 0x08081908,
+    0x2b082b19, 0x08081908, 0x2b190808, 0x08081908, 0x2b191919, 0x08081908, 0x2b192b08, 0x08081908,
+    0x2b2b0819, 0x08081908, 0x2b2b1908, 0x08081908, 0x08080808, 0x08081919, 0x0808082b, 0x08081919,
+    0x08081919, 0x08081919, 0x08082b08, 0x08081919, 0x08082b2b, 0x08081919, 0x08190819, 0x08081919,
+    0x08191908, 0x08081919, 0x0819192b, 0x08081919, 0x08192b19, 0x08081919, 0x082b0808, 0x08081919,
+    0x082b1919, 0x08081919, 0x082b2b08, 0x08081919, 0x19080819, 0x08081919, 0x19081908, 0x08081919,
+    0x1908192b, 0x08081919, 0x19082b19, 0x08081919, 0x19190808, 0x08081919, 0x1919082b, 0x08081919,
+    0x19191919, 0x08081919, 0x19192b08, 0x08081919, 0x192b0819, 0x08081919, 0x192b1908, 0x08081919,
+    0x2b080808, 0x08081919, 0x2b08082b, 0x08081919, 0x2b081919, 0x08081919, 0x2b082b08, 0x08081919,
+    0x2b190819, 0x08081919, 0x2b191908, 0x08081919, 0x2b2b0808, 0x08081919, 0x08080819, 0x0808192b,
+    0x08081908, 0x0808192b, 0x0808192b, 0x0808192b, 0x08082b19, 0x0808192b, 0x08190808, 0x0808192b,
+    0x08191919, 0x0808192b, 0x19080808, 0x0808192b, 0x19081919, 0x0808192b, 0x19082b08, 0x0808192b,
+    0x19190819, 0x0808192b, 0x19191908, 0x0808192b, 0x192b0808, 0x0808192b, 0x2b080819, 0x0808192b,
+    0x2b081908, 0x0808192b, 0x2b190808, 0x0808192b, 0x08080808, 0x08082b08, 0x0808082b, 0x08082b08,
+    0x08081919, 0x08082b08, 0x08082b08, 0x08082b08, 0x08190819, 0x08082b08, 0x08191908, 0x08082b08,
+    0x0819192b, 0x08082b08, 0x08192b19, 0x08082b08, 0x082b0808, 0x08082b08, 0x082b1919, 0x08082b08,
+    0x082b2b2b, 0x08082b08, 0x19080819, 0x08082b08, 0x19081908, 0x08082b08, 0x1908192b, 0x08082b08,
+    0x19082b19, 0x08082b08, 0x19190808, 0x08082b08, 0x1919082b, 0x08082b08, 0x19191919, 0x08082b08,
+    0x19192b08, 0x08082b08, 0x192b0819, 0x08082b08, 0x192b1908, 0x08082b08, 0x2b080808, 0x08082b08,
+    0x2b081919, 0x08082b08, 0x2b191908, 0x08082b08, 0x2b2b2b2b, 0x08082b08, 0x08080819, 0x08082b19,
+    0x08081908, 0x08082b19, 0x08190808, 0x08082b19, 0x0819082b, 0x08082b19, 0x08191919, 0x08082b19,
+    0x08192b08, 0x08082b19, 0x082b0819, 0x08082b19, 0x19080808, 0x08082b19, 0x19081919, 0x08082b19,
+    0x19082b08, 0x08082b19, 0x19190819, 0x08082b19, 0x19191908, 0x08082b19, 0x192b0808, 0x08082b19,
+    0x2b080819, 0x08082b19, 0x2b190808, 0x08082b19, 0x08080808, 0x08082b2b, 0x08190819, 0x08082b2b,
+    0x08191908, 0x08082b2b, 0x082b082b, 0x08082b2b, 0x082b2b08, 0x08082b2b, 0x082b2b2b, 0x08082b2b,
+    0x19190808, 0x08082b2b, 0x2b192b19, 0x08082b2b, 0x08080819, 0x08190808, 0x08081908, 0x08190808,
+    0x0808192b, 0x08190808, 0x08082b19, 0x08190808, 0x08190808, 0x08190808, 0x0819082b, 0x08190808,
+    0x08191919, 0x08190808, 0x08192b08, 0x08190808, 0x082b0819, 0x08190808, 0x082b1908, 0x08190808,
+    0x082b192b, 0x08190808, 0x19080808, 0x08190808, 0x1908082b, 0x08190808, 0x19081919, 0x08190808,
+    0x19082b08, 0x08190808, 0x19190819, 0x08190808, 0x19191908, 0x08190808, 0x1919192b, 0x08190808,
+    0x19192b19, 0x08190808, 0x192b0808, 0x08190808, 0x192b082b, 0x08190808, 0x192b1919, 0x08190808,
+    0x192b2b08, 0x08190808, 0x2b080819, 0x08190808, 0x2b081908, 0x08190808, 0x2b08192b, 0x08190808,
+    0x2b190808, 0x08190808, 0x2b191919, 0x08190808, 0x2b192b08, 0x08190808, 0x2b2b0819, 0x08190808,
+    0x2b2b1908, 0x08190808, 0x08080808, 0x08190819, 0x0808082b, 0x08190819, 0x08081919, 0x08190819,
+    0x08082b08, 0x08190819, 0x08082b2b, 0x08190819, 0x08190819, 0x08190819, 0x08191908, 0x08190819,
+    0x0819192b, 0x08190819, 0x08192b19, 0x08190819, 0x082b0808, 0x08190819, 0x082b082b, 0x08190819,
+    0x082b1919, 0x08190819, 0x082b2b08, 0x08190819, 0x19080819, 0x08190819, 0x19081908, 0x08190819,
+    0x1908192b, 0x08190819, 0x19082b19, 0x08190819, 0x19190808, 0x08190819, 0x1919082b, 0x08190819,
+    0x19191919, 0x08190819, 0x19192b08, 0x08190819, 0x192b0819, 0x08190819, 0x192b1908, 0x08190819,
+    0x2b080808, 0x08190819, 0x2b08082b, 0x08190819, 0x2b081919, 0x08190819, 0x2b082b08, 0x08190819,
+    0x2b190819, 0x08190819, 0x2b191908, 0x08190819, 0x08080819, 0x0819082b, 0x08081908, 0x0819082b,
+    0x08082b19, 0x0819082b, 0x08190808, 0x0819082b, 0x08191919, 0x0819082b, 0x082b0819, 0x0819082b,
+    0x082b1908, 0x0819082b, 0x19080808, 0x0819082b, 0x19081919, 0x0819082b, 0x19190819, 0x0819082b,
+    0x19191908, 0x0819082b, 0x2b080819, 0x0819082b, 0x2b081908, 0x0819082b, 0x2b190808, 0x0819082b,
+    0x08080808, 0x08191908, 0x0808082b, 0x08191908, 0x08081919, 0x08191908, 0x08082b08, 0x08191908,
+    0x08190819, 0x08191908, 0x08191908, 0x08191908, 0x0819192b, 0x08191908, 0x08192b19, 0x08191908,
+    0x082b0808, 0x08191908, 0x082b1919, 0x08191908, 0x082b2b08, 0x08191908, 0x19080819, 0x08191908,
+    0x19081908, 0x08191908, 0x1908192b, 0x08191908, 0x19082b19, 0x08191908, 0x19190808, 0x08191908,
+    0x1919082b, 0x08191908, 0x19191919, 0x08191908, 0x19192b08, 0x08191908, 0x192b0819, 0x08191908,
+    0x192b1908, 0x08191908, 0x2b080808, 0x08191908, 0x2b08082b, 0x08191908, 0x2b081919, 0x08191908,
+    0x2b082b08, 0x08191908, 0x2b190819, 0x08191908, 0x2b191908, 0x08191908, 0x2b2b0808, 0x08191908,
+    0x08080819, 0x08191919, 0x08081908, 0x08191919, 0x0808192b, 0x08191919, 0x08082b19, 0x08191919,
+    0x08190808, 0x08191919, 0x0819082b, 0x08191919, 0x08191919, 0x08191919, 0x08192b08, 0x08191919,
+    0x082b0819, 0x08191919, 0x082b1908, 0x08191919, 0x19080808, 0x08191919, 0x1908082b, 0x08191919,
+    0x19081919, 0x08191919, 0x19082b08, 0x08191919, 0x19190819, 0x08191919, 0x19191908, 0x08191919,
+    0x192b0808, 0x08191919, 0x2b080819, 0x08191919, 0x2b081908, 0x08191919, 0x2b190808, 0x08191919,
+    0x08080808, 0x0819192b, 0x08081919, 0x0819192b, 0x08082b08, 0x0819192b, 0x08190819, 0x0819192b,
+    0x08191908, 0x0819192b, 0x082b0808, 0x0819192b, 0x19080819, 0x0819192b, 0x19081908, 0x0819192b,
+    0x19190808, 0x0819192b, 0x2b080808, 0x0819192b, 0x2b2b2b2b, 0x0819192b, 0x08080819, 0x08192b08,
+    0x08081908, 0x08192b08, 0x0808192b, 0x08192b08, 0x08082b19, 0x08192b08, 0x08190808, 0x08192b08,
+    0x08191919, 0x08192b08, 0x08192b08, 0x08192b08, 0x082b0819, 0x08192b08, 0x19080808, 0x08192b08,
+    0x1908082b, 0x08192b08, 0x19081919, 0x08192b08, 0x19082b08, 0x08192b08, 0x19190819, 0x08192b08,
+    0x19191908, 0x08192b08, 0x192b0808, 0x08192b08, 0x2b080819, 0x08192b08, 0x2b081908, 0x08192b08,
+    0x08080808, 0x08192b19, 0x0808082b, 0x08192b19, 0x08081919, 0x08192b19, 0x08082b08, 0x08192b19,
+    0x08190819, 0x08192b19, 0x08191908, 0x08192b19, 0x082b0808, 0x08192b19, 0x19080819, 0x08192b19,
+    0x19081908, 0x08192b19, 0x19190808, 0x08192b19, 0x192b2b19, 0x08192b19, 0x2b2b082b, 0x08192b19,
+    0x08081908, 0x08192b2b, 0x08190808, 0x08192b2b, 0x19080808, 0x08192b2b, 0x1919192b, 0x08192b2b,
+    0x08080808, 0x082b0808, 0x0808082b, 0x082b0808, 0x08081919, 0x082b0808, 0x08082b08, 0x082b0808,
+    0x08190819, 0x082b0808, 0x08191908, 0x082b0808, 0x0819192b, 0x082b0808, 0x08192b19, 0x082b0808,
+    0x082b0808, 0x082b0808, 0x082b1919, 0x082b0808, 0x082b2b2b, 0x082b0808, 0x19080819, 0x082b0808,
+    0x19081908, 0x082b0808, 0x19190808, 0x082b0808, 0x1919082b, 0x082b0808, 0x19191919, 0x082b0808,
+    0x192b1908, 0x082b0808, 0x2b080808, 0x082b0808, 0x2b082b2b, 0x082b0808, 0x2b191908, 0x082b0808,
+    0x2b2b2b2b, 0x082b0808, 0x08080819, 0x082b0819, 0x08081908, 0x082b0819, 0x08190808, 0x082b0819,
+    0x0819082b, 0x082b0819, 0x08191919, 0x082b0819, 0x082b0819, 0x082b0819, 0x19080808, 0x082b0819,
+    0x1908082b, 0x082b0819, 0x19081919, 0x082b0819, 0x19190819, 0x082b0819, 0x19191908, 0x082b0819,
+    0x192b0808, 0x082b0819, 0x2b080819, 0x082b0819, 0x2b081908, 0x082b0819, 0x2b190808, 0x082b0819,
+    0x08080808, 0x082b082b, 0x08082b2b, 0x082b082b, 0x082b082b, 0x082b082b, 0x082b2b08, 0x082b082b,
+    0x082b2b2b, 0x082b082b, 0x19081908, 0x082b082b, 0x19190808, 0x082b082b, 0x2b082b08, 0x082b082b,
+    0x2b082b2b, 0x082b082b, 0x2b2b2b08, 0x082b082b, 0x08080819, 0x082b1908, 0x08081908, 0x082b1908,
+    0x0808192b, 0x082b1908, 0x08082b19, 0x082b1908, 0x08190808, 0x082b1908, 0x08191919, 0x082b1908,
+    0x08192b08, 0x082b1908, 0x082b0819, 0x082b1908, 0x082b1908, 0x082b1908, 0x19080808, 0x082b1908,
+    0x1908082b, 0x082b1908, 0x19081919, 0x082b1908, 0x19082b08, 0x082b1908, 0x19190819, 0x082b1908,
+    0x19191908, 0x082b1908, 0x192b0808, 0x082b1908, 0x2b080819, 0x082b1908, 0x2b081908, 0x082b1908,
+    0x2b190808, 0x082b1908, 0x08080808, 0x082b1919, 0x08081919, 0x082b1919, 0x08082b08, 0x082b1919,
+    0x08190819, 0x082b1919, 0x08191908, 0x082b1919, 0x082b0808, 0x082b1919, 0x19080819, 0x082b1919,
+    0x19081908, 0x082b1919, 0x19190808, 0x082b1919, 0x192b192b, 0x082b1919, 0x2b080808, 0x082b1919,
+    0x08080819, 0x082b192b, 0x08081908, 0x082b192b, 0x08190808, 0x082b192b, 0x19080808, 0x082b192b,
+    0x19192b19, 0x082b192b, 0x08080808, 0x082b2b08, 0x08081919, 0x082b2b08, 0x08190819, 0x082b2b08,
+    0x08191908, 0x082b2b08, 0x19080819, 0x082b2b08, 0x19081908, 0x082b2b08, 0x19190808, 0x082b2b08,
+    0x2b082b2b, 0x082b2b08, 0x2b2b2b2b, 0x082b2b08, 0x08080819, 0x082b2b19, 0x08081908, 0x082b2b19,
+    0x08190808, 0x082b2b19, 0x2b191919, 0x082b2b19, 0x08082b2b, 0x082b2b2b, 0x082b082b, 0x082b2b2b,
+    0x192b1908, 0x082b2b2b, 0x2b082b08, 0x082b2b2b, 0x2b082b2b, 0x082b2b2b, 0x08080819, 0x19080808,
+    0x08081908, 0x19080808, 0x0808192b, 0x19080808, 0x08082b19, 0x19080808, 0x08190808, 0x19080808,
+    0x0819082b, 0x19080808, 0x08191919, 0x19080808, 0x08192b08, 0x19080808, 0x08192b2b, 0x19080808,
+    0x082b0819, 0x19080808, 0x082b1908, 0x19080808, 0x082b192b, 0x19080808, 0x19080808, 0x19080808,
+    0x1908082b, 0x19080808, 0x19081919, 0x19080808, 0x19082b08, 0x19080808, 0x19082b2b, 0x19080808,
+    0x19190819, 0x19080808, 0x19191908, 0x19080808, 0x1919192b, 0x19080808, 0x19192b19, 0x19080808,
+    0x192b0808, 0x19080808, 0x192b082b, 0x19080808, 0x192b1919, 0x19080808, 0x2b080819, 0x19080808,
+    0x2b081908, 0x19080808, 0x2b190808, 0x19080808, 0x2b191919, 0x19080808, 0x2b192b08, 0x19080808,
+    0x2b2b0819, 0x19080808, 0x2b2b1908, 0x19080808, 0x08080808, 0x19080819, 0x0808082b, 0x19080819,
+    0x08081919, 0x19080819, 0x08082b08, 0x19080819, 0x08190819, 0x19080819, 0x08191908, 0x19080819,
+    0x0819192b, 0x19080819, 0x08192b19, 0x19080819, 0x082b0808, 0x19080819, 0x082b082b, 0x19080819,
+    0x082b1919, 0x19080819, 0x19080819, 0x19080819, 0x19081908, 0x19080819, 0x1908192b, 0x19080819,
+    0x19082b19, 0x19080819, 0x19190808, 0x19080819, 0x1919082b, 0x19080819, 0x19191919, 0x19080819,
+    0x19192b08, 0x19080819, 0x192b0819, 0x19080819, 0x192b1908, 0x19080819, 0x2b080808, 0x19080819,
+    0x2b08082b, 0x19080819, 0x2b081919, 0x19080819, 0x2b082b08, 0x19080819, 0x2b190819, 0x19080819,
+    0x2b191908, 0x19080819, 0x2b2b0808, 0x19080819, 0x08080819, 0x1908082b, 0x08081908, 0x1908082b,
+    0x08190808, 0x1908082b, 0x0819082b, 0x1908082b, 0x08191919, 0x1908082b, 0x08192b08, 0x1908082b,
+    0x082b1908, 0x1908082b, 0x19080808, 0x1908082b, 0x19081919, 0x1908082b, 0x19082b08, 0x1908082b,
+    0x19190819, 0x1908082b, 0x19191908, 0x1908082b, 0x192b0808, 0x1908082b, 0x2b080819, 0x1908082b,
+    0x2b081908, 0x1908082b, 0x08080808, 0x19081908, 0x0808082b, 0x19081908, 0x08081919, 0x19081908,
+    0x08082b08, 0x19081908, 0x08082b2b, 0x19081908, 0x08190819, 0x19081908, 0x08191908, 0x19081908,
+    0x0819192b, 0x19081908, 0x08192b19, 0x19081908, 0x082b0808, 0x19081908, 0x082b082b, 0x19081908,
+    0x082b1919, 0x19081908, 0x082b2b08, 0x19081908, 0x19080819, 0x19081908, 0x19081908, 0x19081908,
+    0x1908192b, 0x19081908, 0x19082b19, 0x19081908, 0x19190808, 0x19081908, 0x1919082b, 0x19081908,
+    0x19191919, 0x19081908, 0x19192b08, 0x19081908, 0x192b0819, 0x19081908, 0x192b1908, 0x19081908,
+    0x2b080808, 0x19081908, 0x2b08082b, 0x19081908, 0x2b081919, 0x19081908, 0x2b082b08, 0x19081908,
+    0x2b190819, 0x19081908, 0x2b191908, 0x19081908, 0x2b2b0808, 0x19081908, 0x08080819, 0x19081919,
+    0x08081908, 0x19081919, 0x0808192b, 0x19081919, 0x08082b19, 0x19081919, 0x08190808, 0x19081919,
+    0x0819082b, 0x19081919, 0x08191919, 0x19081919, 0x08192b08, 0x19081919, 0x082b0819, 0x19081919,
+    0x082b1908, 0x19081919, 0x19080808, 0x19081919, 0x1908082b, 0x19081919, 0x19081919, 0x19081919,
+    0x19082b08, 0x19081919, 0x19190819, 0x19081919, 0x19191908, 0x19081919, 0x192b0808, 0x19081919,
+    0x192b2b2b, 0x19081919, 0x2b080819, 0x19081919, 0x2b081908, 0x19081919, 0x2b190808, 0x19081919,
+    0x08080808, 0x1908192b, 0x0808082b, 0x1908192b, 0x08081919, 0x1908192b, 0x08082b08, 0x1908192b,
+    0x08190819, 0x1908192b, 0x08191908, 0x1908192b, 0x082b0808, 0x1908192b, 0x19080819, 0x1908192b,
+    0x19081908, 0x1908192b, 0x19190808, 0x1908192b, 0x2b080808, 0x1908192b, 0x2b2b1919, 0x1908192b,
+    0x08080819, 0x19082b08, 0x08081908, 0x19082b08, 0x08082b19, 0x19082b08, 0x08190808, 0x19082b08,
+    0x0819082b, 0x19082b08, 0x08191919, 0x19082b08, 0x08192b08, 0x19082b08, 0x082b0819, 0x19082b08,
+    0x082b1908, 0x19082b08, 0x19080808, 0x19082b08, 0x1908082b, 0x19082b08, 0x19081919, 0x19082b08,
+    0x19082b08, 0x19082b08, 0x19190819, 0x19082b08, 0x19191908, 0x19082b08, 0x192b0808, 0x19082b08,
+    0x2b081908, 0x19082b08, 0x2b190808, 0x19082b08, 0x08080808, 0x19082b19, 0x0808082b, 0x19082b19,
+    0x08081919, 0x19082b19, 0x08082b08, 0x19082b19, 0x08190819, 0x19082b19, 0x08191908, 0x19082b19,
+    0x082b0808, 0x19082b19, 0x19080819, 0x19082b19, 0x19081908, 0x19082b19, 0x19190808, 0x19082b19,
+    0x2b080808, 0x19082b19, 0x2b19192b, 0x19082b19, 0x08080819, 0x19082b2b, 0x08081908, 0x19082b2b,
+    0x08190808, 0x19082b2b, 0x19080808, 0x19082b2b, 0x08080808, 0x19190808, 0x0808082b, 0x19190808,
+    0x08081919, 0x19190808, 0x08082b08, 0x19190808, 0x08190819, 0x19190808, 0x08191908, 0x19190808,
+    0x0819192b, 0x19190808, 0x08192b19, 0x19190808, 0x082b0808, 0x19190808, 0x082b082b, 0x19190808,
+    0x082b1919, 0x19190808, 0x082b2b08, 0x19190808, 0x19080819, 0x19190808, 0x19081908, 0x19190808,
+    0x1908192b, 0x19190808, 0x19082b19, 0x19190808, 0x19190808, 0x19190808, 0x1919082b, 0x19190808,
+    0x19191919, 0x19190808, 0x19192b08, 0x19190808, 0x192b0819, 0x19190808, 0x192b1908, 0x19190808,
+    0x2b080808, 0x19190808, 0x2b08082b, 0x19190808, 0x2b081919, 0x19190808, 0x2b082b08, 0x19190808,
+    0x2b190819, 0x19190808, 0x2b191908, 0x19190808, 0x08080819, 0x19190819, 0x08081908, 0x19190819,
+    0x0808192b, 0x19190819, 0x08082b19, 0x19190819, 0x08190808, 0x19190819, 0x0819082b, 0x19190819,
+    0x08191919, 0x19190819, 0x08192b08, 0x19190819, 0x082b0819, 0x19190819, 0x082b1908, 0x19190819,
+    0x19080808, 0x19190819, 0x1908082b, 0x19190819, 0x19081919, 0x19190819, 0x19082b08, 0x19190819,
+    0x19190819, 0x19190819, 0x19191908, 0x19190819, 0x192b0808, 0x19190819, 0x2b080819, 0x19190819,
+    0x2b081908, 0x19190819, 0x2b190808, 0x19190819, 0x08080808, 0x1919082b, 0x08081919, 0x1919082b,
+    0x08082b08, 0x1919082b, 0x08190819, 0x1919082b, 0x08191908, 0x1919082b, 0x082b0808, 0x1919082b,
+    0x19080819, 0x1919082b, 0x19081908, 0x1919082b, 0x19190808, 0x1919082b, 0x192b2b19, 0x1919082b,
+    0x2b080808, 0x1919082b, 0x08080819, 0x19191908, 0x08081908, 0x19191908, 0x0808192b, 0x19191908,
+    0x08082b19, 0x19191908, 0x08190808, 0x19191908, 0x0819082b, 0x19191908, 0x08191919, 0x19191908,
+    0x08192b08, 0x19191908, 0x082b0819, 0x19191908, 0x082b1908, 0x19191908, 0x19080808, 0x19191908,
+    0x1908082b, 0x19191908, 0x19081919, 0x19191908, 0x19082b08, 0x19191908, 0x19190819, 0x19191908,
+    0x19191908, 0x19191908, 0x192b0808, 0x19191908, 0x2b080819, 0x19191908, 0x2b081908, 0x19191908,
+    0x2b190808, 0x19191908, 0x08080808, 0x19191919, 0x0808082b, 0x19191919, 0x08081919, 0x19191919,
+    0x08082b08, 0x19191919, 0x08190819, 0x19191919, 0x08191908, 0x19191919, 0x082b0808, 0x19191919,
+    0x19080819, 0x19191919, 0x19081908, 0x19191919, 0x19190808, 0x19191919, 0x2b080808, 0x19191919,
+    0x08080819, 0x1919192b, 0x08081908, 0x1919192b, 0x08190808, 0x1919192b, 0x082b192b, 0x1919192b,
+    0x19080808, 0x1919192b, 0x08080808, 0x19192b08, 0x0808082b, 0x19192b08, 0x08081919, 0x19192b08,
+    0x08082b08, 0x19192b08, 0x08190819, 0x19192b08, 0x08191908, 0x19192b08, 0x082b0808, 0x19192b08,
+    0x19080819, 0x19192b08, 0x19081908, 0x19192b08, 0x19190808, 0x19192b08, 0x19192b2b, 0x19192b08,
+    0x2b080808, 0x19192b08, 0x08080819, 0x19192b19, 0x08081908, 0x19192b19, 0x08190808, 0x19192b19,
+    0x19080808, 0x19192b19, 0x08080808, 0x19192b2b, 0x08192b19, 0x19192b2b, 0x2b081919, 0x19192b2b,
+    0x2b2b2b08, 0x19192b2b, 0x08080819, 0x192b0808, 0x08081908, 0x192b0808, 0x0808192b, 0x192b0808,
+    0x08190808, 0x192b0808, 0x0819082b, 0x192b0808, 0x08191919, 0x192b0808, 0x08192b08, 0x192b0808,
+    0x082b0819, 0x192b0808, 0x082b1908, 0x192b0808, 0x19080808, 0x192b0808, 0x19081919, 0x192b0808,
+    0x19082b08, 0x192b0808, 0x19190819, 0x192b0808, 0x19191908, 0x192b0808, 0x192b0808, 0x192b0808,
+    0x2b081908, 0x192b0808, 0x2b190808, 0x192b0808, 0x08080808, 0x192b0819, 0x0808082b, 0x192b0819,
+    0x08081919, 0x192b0819, 0x08082b08, 0x192b0819, 0x08190819, 0x192b0819, 0x08191908, 0x192b0819,
+    0x082b0808, 0x192b0819, 0x19080819, 0x192b0819, 0x19081908, 0x192b0819, 0x19190808, 0x192b0819,
+    0x2b080808, 0x192b0819, 0x2b192b19, 0x192b0819, 0x08081908, 0x192b082b, 0x08190808, 0x192b082b,
+    0x19080808, 0x192b082b, 0x1919192b, 0x192b082b, 0x2b2b0819, 0x192b082b, 0x08080808, 0x192b1908,
+    0x08081919, 0x192b1908, 0x08082b08, 0x192b1908, 0x08190819, 0x192b1908, 0x08191908, 0x192b1908,
+    0x082b0808, 0x192b1908, 0x19080819, 0x192b1908, 0x19081908, 0x192b1908, 0x19190808, 0x192b1908,
+    0x2b080808, 0x192b1908, 0x08080819, 0x192b1919, 0x08081908, 0x192b1919, 0x08190808, 0x192b1919,
+    0x19080808, 0x192b1919, 0x19082b2b, 0x192b1919, 0x192b2b08, 0x192b1919, 0x2b19082b, 0x192b1919,
+    0x08080808, 0x192b192b, 0x2b191908, 0x192b192b, 0x08080819, 0x192b2b08, 0x08081908, 0x192b2b08,
+    0x08190808, 0x192b2b08, 0x192b1919, 0x192b2b08, 0x2b192b08, 0x192b2b08, 0x08080808, 0x192b2b19,
+    0x082b2b2b, 0x192b2b19, 0x1908082b, 0x192b2b2b, 0x2b2b0819, 0x192b2b2b, 0x08080808, 0x2b080808,
+    0x0808082b, 0x2b080808, 0x08081919, 0x2b080808, 0x08082b08, 0x2b080808, 0x08190819, 0x2b080808,
+    0x08191908, 0x2b080808, 0x08192b19, 0x2b080808, 0x082b0808, 0x2b080808, 0x082b1919, 0x2b080808,
+    0x19080819, 0x2b080808, 0x19081908, 0x2b080808, 0x19190808, 0x2b080808, 0x1919082b, 0x2b080808,
+    0x19191919, 0x2b080808, 0x19192b08, 0x2b080808, 0x192b0819, 0x2b080808, 0x2b080808, 0x2b080808,
+    0x2b081919, 0x2b080808, 0x2b190819, 0x2b080808, 0x2b191908, 0x2b080808, 0x08080819, 0x2b080819,
+    0x08081908, 0x2b080819, 0x08082b19, 0x2b080819, 0x08190808, 0x2b080819, 0x0819082b, 0x2b080819,
+    0x08191919, 0x2b080819, 0x08192b08, 0x2b080819, 0x082b0819, 0x2b080819, 0x082b1908, 0x2b080819,
+    0x19080808, 0x2b080819, 0x1908082b, 0x2b080819, 0x19081919, 0x2b080819, 0x19082b08, 0x2b080819,
+    0x19190819, 0x2b080819, 0x19191908, 0x2b080819, 0x2b080819, 0x2b080819, 0x2b081908, 0x2b080819,
+    0x2b190808, 0x2b080819, 0x2b2b2b19, 0x2b080819, 0x08080808, 0x2b08082b, 0x08081919, 0x2b08082b,
+    0x08082b2b, 0x2b08082b, 0x08190819, 0x2b08082b, 0x08191908, 0x2b08082b, 0x19080819, 0x2b08082b,
+    0x19081908, 0x2b08082b, 0x19190808, 0x2b08082b, 0x08080819, 0x2b081908, 0x08081908, 0x2b081908,
+    0x0808192b, 0x2b081908, 0x08082b19, 0x2b081908, 0x08190808, 0x2b081908, 0x0819082b, 0x2b081908,
+    0x08191919, 0x2b081908, 0x08192b08, 0x2b081908, 0x082b0819, 0x2b081908, 0x19080808, 0x2b081908,
+    0x1908082b, 0x2b081908, 0x19081919, 0x2b081908, 0x19082b08, 0x2b081908, 0x19190819, 0x2b081908,
+    0x19191908, 0x2b081908, 0x192b0808, 0x2b081908, 0x2b080819, 0x2b081908, 0x2b081908, 0x2b081908,
+    0x2b190808, 0x2b081908, 0x08080808, 0x2b081919, 0x0808082b, 0x2b081919, 0x08081919, 0x2b081919,
+    0x08082b08, 0x2b081919, 0x08190819, 0x2b081919, 0x08191908, 0x2b081919, 0x082b0808, 0x2b081919,
+    0x19080819, 0x2b081919, 0x19081908, 0x2b081919, 0x19190808, 0x2b081919, 0x2b080808, 0x2b081919,
+    0x2b082b2b, 0x2b081919, 0x08080819, 0x2b08192b, 0x08081908, 0x2b08192b, 0x08190808, 0x2b08192b,
+    0x082b2b19, 0x2b08192b, 0x19080808, 0x2b08192b, 0x08080808, 0x2b082b08, 0x08081919, 0x2b082b08,
+    0x08190819, 0x2b082b08, 0x08191908, 0x2b082b08, 0x19080819, 0x2b082b08, 0x19081908, 0x2b082b08,
+    0x19190808, 0x2b082b08, 0x2b2b082b, 0x2b082b08, 0x08080819, 0x2b082b19, 0x08081908, 0x2b082b19,
+    0x19080808, 0x2b082b19, 0x192b1919, 0x2b082b19, 0x082b082b, 0x2b082b2b, 0x19192b08, 0x2b082b2b,
+    0x19192b2b, 0x2b082b2b, 0x2b08082b, 0x2b082b2b, 0x2b2b082b, 0x2b082b2b, 0x08080819, 0x2b190808,
+    0x08081908, 0x2b190808, 0x08082b19, 0x2b190808, 0x08190808, 0x2b190808, 0x0819082b, 0x2b190808,
+    0x08191919, 0x2b190808, 0x08192b08, 0x2b190808, 0x082b1908, 0x2b190808, 0x19080808, 0x2b190808,
+    0x1908082b, 0x2b190808, 0x19081919, 0x2b190808, 0x19082b08, 0x2b190808, 0x19190819, 0x2b190808,
+    0x19191908, 0x2b190808, 0x192b0808, 0x2b190808, 0x2b080819, 0x2b190808, 0x2b081908, 0x2b190808,
+    0x2b190808, 0x2b190808, 0x08080808, 0x2b190819, 0x08081919, 0x2b190819, 0x08190819, 0x2b190819,
+    0x08191908, 0x2b190819, 0x19080819, 0x2b190819, 0x19081908, 0x2b190819, 0x19190808, 0x2b190819,
+    0x19192b2b, 0x2b190819, 0x08080819, 0x2b19082b, 0x08081908, 0x2b19082b, 0x08190808, 0x2b19082b,
+    0x19080808, 0x2b19082b, 0x2b2b192b, 0x2b19082b, 0x08080808, 0x2b191908, 0x0808082b, 0x2b191908,
+    0x08081919, 0x2b191908, 0x08082b08, 0x2b191908, 0x08190819, 0x2b191908, 0x08191908, 0x2b191908,
+    0x082b0808, 0x2b191908, 0x19080819, 0x2b191908, 0x19081908, 0x2b191908, 0x19190808, 0x2b191908,
+    0x2b080808, 0x2b191908, 0x2b19192b, 0x2b191908, 0x08080819, 0x2b191919, 0x08081908, 0x2b191919,
+    0x08190808, 0x2b191919, 0x19080808, 0x2b191919, 0x2b192b08, 0x2b191919, 0x2b2b0819, 0x2b191919,
+    0x08080808, 0x2b19192b, 0x1908192b, 0x2b19192b, 0x192b1908, 0x2b19192b, 0x08080819, 0x2b192b08,
+    0x08081908, 0x2b192b08, 0x08190808, 0x2b192b08, 0x082b192b, 0x2b192b08, 0x19080808, 0x2b192b08,
+    0x2b2b2b19, 0x2b192b08, 0x08080808, 0x2b192b19, 0x19082b19, 0x2b192b19, 0x1919082b, 0x2b192b19,
+    0x2b190808, 0x2b192b2b, 0x08080808, 0x2b2b0808, 0x08081919, 0x2b2b0808, 0x08082b2b, 0x2b2b0808,
+    0x08191908, 0x2b2b0808, 0x082b082b, 0x2b2b0808, 0x082b2b2b, 0x2b2b0808, 0x19080819, 0x2b2b0808,
+    0x19081908, 0x2b2b0808, 0x19190808, 0x2b2b0808, 0x2b2b082b, 0x2b2b0808, 0x2b2b2b2b, 0x2b2b0808,
+    0x19080808, 0x2b2b0819, 0x192b1919, 0x2b2b0819, 0x0808082b, 0x2b2b082b, 0x08082b2b, 0x2b2b082b,
+    0x082b082b, 0x2b2b082b, 0x082b2b08, 0x2b2b082b, 0x082b2b2b, 0x2b2b082b, 0x2b08082b, 0x2b2b082b,
+    0x2b082b08, 0x2b2b082b, 0x2b082b2b, 0x2b2b082b, 0x2b2b2b08, 0x2b2b082b, 0x08080819, 0x2b2b1908,
+    0x08081908, 0x2b2b1908, 0x08190808, 0x2b2b1908, 0x19080808, 0x2b2b1908, 0x2b082b19, 0x2b2b1908,
+    0x2b2b1908, 0x2b2b1908, 0x08080808, 0x2b2b1919, 0x08192b19, 0x2b2b1919, 0x19190819, 0x2b2b192b,
+    0x08082b2b, 0x2b2b2b08, 0x082b2b08, 0x2b2b2b08, 0x2b2b082b, 0x2b2b2b08, 0x19191908, 0x2b2b2b19,
+    0x2b08192b, 0x2b2b2b19, 0x08082b08, 0x2b2b2b2b, 0x08082b2b, 0x2b2b2b2b, 0x082b0808, 0x2b2b2b2b,
+    0x082b082b, 0x2b2b2b2b, 0x082b2b08, 0x2b2b2b2b, 0x2b082b08, 0x2b2b2b2b, 0x2b2b2b2b, 0x2b2b2b2b
+};
@@ -0,0 +1,135 @@
+// IQ2_XS codebook, stored as pairs of four packed magnitudes.
+constant uint iq2_xs_grid[1024] = {
+    0x08080808, 0x08080808, 0x0808082b, 0x08080808, 0x08081919, 0x08080808, 0x08082b08, 0x08080808,
+    0x08082b2b, 0x08080808, 0x08190819, 0x08080808, 0x08191908, 0x08080808, 0x0819192b, 0x08080808,
+    0x08192b19, 0x08080808, 0x082b0808, 0x08080808, 0x082b082b, 0x08080808, 0x082b1919, 0x08080808,
+    0x082b2b08, 0x08080808, 0x19080819, 0x08080808, 0x19081908, 0x08080808, 0x1908192b, 0x08080808,
+    0x19082b19, 0x08080808, 0x19190808, 0x08080808, 0x1919082b, 0x08080808, 0x19191919, 0x08080808,
+    0x19192b08, 0x08080808, 0x192b0819, 0x08080808, 0x192b1908, 0x08080808, 0x2b080808, 0x08080808,
+    0x2b08082b, 0x08080808, 0x2b081919, 0x08080808, 0x2b082b08, 0x08080808, 0x2b190819, 0x08080808,
+    0x2b191908, 0x08080808, 0x2b192b19, 0x08080808, 0x2b2b0808, 0x08080808, 0x08080819, 0x08080819,
+    0x08081908, 0x08080819, 0x0808192b, 0x08080819, 0x08082b19, 0x08080819, 0x08190808, 0x08080819,
+    0x0819082b, 0x08080819, 0x08191919, 0x08080819, 0x08192b08, 0x08080819, 0x08192b2b, 0x08080819,
+    0x082b0819, 0x08080819, 0x082b1908, 0x08080819, 0x19080808, 0x08080819, 0x1908082b, 0x08080819,
+    0x19081919, 0x08080819, 0x19082b08, 0x08080819, 0x19190819, 0x08080819, 0x19191908, 0x08080819,
+    0x192b0808, 0x08080819, 0x192b2b08, 0x08080819, 0x2b080819, 0x08080819, 0x2b081908, 0x08080819,
+    0x2b190808, 0x08080819, 0x08080808, 0x0808082b, 0x0808082b, 0x0808082b, 0x08081919, 0x0808082b,
+    0x08082b08, 0x0808082b, 0x08190819, 0x0808082b, 0x08191908, 0x0808082b, 0x082b0808, 0x0808082b,
+    0x19080819, 0x0808082b, 0x19081908, 0x0808082b, 0x19190808, 0x0808082b, 0x19191919, 0x0808082b,
+    0x2b080808, 0x0808082b, 0x2b082b2b, 0x0808082b, 0x08080819, 0x08081908, 0x08081908, 0x08081908,
+    0x0808192b, 0x08081908, 0x08082b19, 0x08081908, 0x08190808, 0x08081908, 0x0819082b, 0x08081908,
+    0x08191919, 0x08081908, 0x08192b08, 0x08081908, 0x082b0819, 0x08081908, 0x082b1908, 0x08081908,
+    0x19080808, 0x08081908, 0x1908082b, 0x08081908, 0x19081919, 0x08081908, 0x19082b08, 0x08081908,
+    0x19190819, 0x08081908, 0x19191908, 0x08081908, 0x1919192b, 0x08081908, 0x192b0808, 0x08081908,
+    0x2b080819, 0x08081908, 0x2b081908, 0x08081908, 0x2b190808, 0x08081908, 0x08080808, 0x08081919,
+    0x0808082b, 0x08081919, 0x08081919, 0x08081919, 0x08082b08, 0x08081919, 0x08190819, 0x08081919,
+    0x08191908, 0x08081919, 0x082b0808, 0x08081919, 0x19080819, 0x08081919, 0x19081908, 0x08081919,
+    0x19190808, 0x08081919, 0x192b0819, 0x08081919, 0x2b080808, 0x08081919, 0x08080819, 0x0808192b,
+    0x08081908, 0x0808192b, 0x08190808, 0x0808192b, 0x082b192b, 0x0808192b, 0x19080808, 0x0808192b,
+    0x1908082b, 0x0808192b, 0x2b081908, 0x0808192b, 0x08080808, 0x08082b08, 0x0808082b, 0x08082b08,
+    0x08081919, 0x08082b08, 0x08082b08, 0x08082b08, 0x08082b2b, 0x08082b08, 0x08190819, 0x08082b08,
+    0x08191908, 0x08082b08, 0x082b0808, 0x08082b08, 0x082b1919, 0x08082b08, 0x19080819, 0x08082b08,
+    0x19081908, 0x08082b08, 0x19190808, 0x08082b08, 0x19192b08, 0x08082b08, 0x2b080808, 0x08082b08,
+    0x2b2b0808, 0x08082b08, 0x2b2b2b2b, 0x08082b08, 0x08080819, 0x08082b19, 0x08081908, 0x08082b19,
+    0x08190808, 0x08082b19, 0x19080808, 0x08082b19, 0x2b080819, 0x08082b19, 0x2b082b19, 0x08082b19,
+    0x08080808, 0x08082b2b, 0x082b0808, 0x08082b2b, 0x082b2b08, 0x08082b2b, 0x2b19192b, 0x08082b2b,
+    0x2b2b0808, 0x08082b2b, 0x08080819, 0x08190808, 0x08081908, 0x08190808, 0x0808192b, 0x08190808,
+    0x08082b19, 0x08190808, 0x08190808, 0x08190808, 0x0819082b, 0x08190808, 0x08191919, 0x08190808,
+    0x08192b08, 0x08190808, 0x082b0819, 0x08190808, 0x082b1908, 0x08190808, 0x19080808, 0x08190808,
+    0x1908082b, 0x08190808, 0x19081919, 0x08190808, 0x19082b08, 0x08190808, 0x19190819, 0x08190808,
+    0x19191908, 0x08190808, 0x192b0808, 0x08190808, 0x192b2b2b, 0x08190808, 0x2b080819, 0x08190808,
+    0x2b081908, 0x08190808, 0x2b190808, 0x08190808, 0x08080808, 0x08190819, 0x0808082b, 0x08190819,
+    0x08081919, 0x08190819, 0x08082b08, 0x08190819, 0x08190819, 0x08190819, 0x08191908, 0x08190819,
+    0x082b0808, 0x08190819, 0x19080819, 0x08190819, 0x19081908, 0x08190819, 0x19190808, 0x08190819,
+    0x2b080808, 0x08190819, 0x2b191908, 0x08190819, 0x2b19192b, 0x08190819, 0x08080819, 0x0819082b,
+    0x08081908, 0x0819082b, 0x0808192b, 0x0819082b, 0x08190808, 0x0819082b, 0x19080808, 0x0819082b,
+    0x192b0808, 0x0819082b, 0x08080808, 0x08191908, 0x0808082b, 0x08191908, 0x08081919, 0x08191908,
+    0x08082b08, 0x08191908, 0x08190819, 0x08191908, 0x08191908, 0x08191908, 0x082b0808, 0x08191908,
+    0x19080819, 0x08191908, 0x19081908, 0x08191908, 0x19082b19, 0x08191908, 0x19190808, 0x08191908,
+    0x192b1908, 0x08191908, 0x2b080808, 0x08191908, 0x08080819, 0x08191919, 0x08081908, 0x08191919,
+    0x08190808, 0x08191919, 0x19080808, 0x08191919, 0x08080808, 0x0819192b, 0x08191908, 0x0819192b,
+    0x19082b19, 0x0819192b, 0x08080819, 0x08192b08, 0x08081908, 0x08192b08, 0x08190808, 0x08192b08,
+    0x0819082b, 0x08192b08, 0x19080808, 0x08192b08, 0x19191908, 0x08192b08, 0x2b08192b, 0x08192b08,
+    0x08080808, 0x08192b19, 0x08081919, 0x08192b19, 0x192b192b, 0x08192b19, 0x19190819, 0x08192b2b,
+    0x2b2b2b19, 0x08192b2b, 0x08080808, 0x082b0808, 0x0808082b, 0x082b0808, 0x08081919, 0x082b0808,
+    0x08082b08, 0x082b0808, 0x08082b2b, 0x082b0808, 0x08190819, 0x082b0808, 0x08191908, 0x082b0808,
+    0x082b0808, 0x082b0808, 0x19080819, 0x082b0808, 0x19081908, 0x082b0808, 0x19190808, 0x082b0808,
+    0x2b080808, 0x082b0808, 0x2b2b0808, 0x082b0808, 0x08080819, 0x082b0819, 0x08081908, 0x082b0819,
+    0x08190808, 0x082b0819, 0x19080808, 0x082b0819, 0x19082b08, 0x082b0819, 0x192b1919, 0x082b0819,
+    0x08080808, 0x082b082b, 0x082b082b, 0x082b082b, 0x2b080808, 0x082b082b, 0x2b2b2b08, 0x082b082b,
+    0x08080819, 0x082b1908, 0x08081908, 0x082b1908, 0x08190808, 0x082b1908, 0x082b2b19, 0x082b1908,
+    0x19080808, 0x082b1908, 0x08080808, 0x082b1919, 0x19080819, 0x082b1919, 0x1919082b, 0x082b1919,
+    0x2b192b19, 0x082b1919, 0x08080819, 0x082b192b, 0x08192b2b, 0x082b192b, 0x2b2b192b, 0x082b192b,
+    0x08080808, 0x082b2b08, 0x08082b08, 0x082b2b08, 0x08082b2b, 0x082b2b08, 0x082b0808, 0x082b2b08,
+    0x19191919, 0x082b2b08, 0x2b082b08, 0x082b2b08, 0x2b2b082b, 0x082b2b08, 0x192b2b08, 0x082b2b19,
+    0x2b190808, 0x082b2b19, 0x08082b08, 0x082b2b2b, 0x082b0808, 0x082b2b2b, 0x2b08082b, 0x082b2b2b,
+    0x2b082b08, 0x082b2b2b, 0x2b082b2b, 0x082b2b2b, 0x08080819, 0x19080808, 0x08081908, 0x19080808,
+    0x0808192b, 0x19080808, 0x08082b19, 0x19080808, 0x08190808, 0x19080808, 0x0819082b, 0x19080808,
+    0x08191919, 0x19080808, 0x08192b08, 0x19080808, 0x082b0819, 0x19080808, 0x082b1908, 0x19080808,
+    0x19080808, 0x19080808, 0x1908082b, 0x19080808, 0x19081919, 0x19080808, 0x19082b08, 0x19080808,
+    0x19082b2b, 0x19080808, 0x19190819, 0x19080808, 0x19191908, 0x19080808, 0x192b0808, 0x19080808,
+    0x192b1919, 0x19080808, 0x2b080819, 0x19080808, 0x2b081908, 0x19080808, 0x2b190808, 0x19080808,
+    0x08080808, 0x19080819, 0x0808082b, 0x19080819, 0x08081919, 0x19080819, 0x08082b08, 0x19080819,
+    0x08190819, 0x19080819, 0x08191908, 0x19080819, 0x082b0808, 0x19080819, 0x19080819, 0x19080819,
+    0x19081908, 0x19080819, 0x19190808, 0x19080819, 0x2b080808, 0x19080819, 0x2b081919, 0x19080819,
+    0x2b2b082b, 0x19080819, 0x08080819, 0x1908082b, 0x08081908, 0x1908082b, 0x08190808, 0x1908082b,
+    0x0819082b, 0x1908082b, 0x082b2b19, 0x1908082b, 0x19080808, 0x1908082b, 0x08080808, 0x19081908,
+    0x0808082b, 0x19081908, 0x08081919, 0x19081908, 0x08082b08, 0x19081908, 0x08190819, 0x19081908,
+    0x08191908, 0x19081908, 0x08192b19, 0x19081908, 0x082b0808, 0x19081908, 0x19080819, 0x19081908,
+    0x19081908, 0x19081908, 0x19190808, 0x19081908, 0x2b080808, 0x19081908, 0x2b191908, 0x19081908,
+    0x08080819, 0x19081919, 0x08081908, 0x19081919, 0x08190808, 0x19081919, 0x082b1908, 0x19081919,
+    0x19080808, 0x19081919, 0x2b192b2b, 0x19081919, 0x08080808, 0x1908192b, 0x08082b2b, 0x1908192b,
+    0x19081908, 0x1908192b, 0x19190808, 0x1908192b, 0x08080819, 0x19082b08, 0x08081908, 0x19082b08,
+    0x08190808, 0x19082b08, 0x19080808, 0x19082b08, 0x19081919, 0x19082b08, 0x19191908, 0x19082b08,
+    0x192b082b, 0x19082b08, 0x08080808, 0x19082b19, 0x08190819, 0x19082b19, 0x19081908, 0x19082b19,
+    0x19190808, 0x19082b19, 0x192b2b19, 0x19082b19, 0x08081908, 0x19082b2b, 0x08080808, 0x19190808,
+    0x0808082b, 0x19190808, 0x08081919, 0x19190808, 0x08082b08, 0x19190808, 0x08190819, 0x19190808,
+    0x08191908, 0x19190808, 0x082b0808, 0x19190808, 0x082b2b08, 0x19190808, 0x19080819, 0x19190808,
+    0x19081908, 0x19190808, 0x19190808, 0x19190808, 0x2b080808, 0x19190808, 0x08080819, 0x19190819,
+    0x08081908, 0x19190819, 0x08190808, 0x19190819, 0x08191919, 0x19190819, 0x19080808, 0x19190819,
+    0x1908082b, 0x19190819, 0x08080808, 0x1919082b, 0x19081908, 0x1919082b, 0x2b2b2b2b, 0x1919082b,
+    0x08080819, 0x19191908, 0x08081908, 0x19191908, 0x08190808, 0x19191908, 0x082b0819, 0x19191908,
+    0x19080808, 0x19191908, 0x192b0808, 0x19191908, 0x2b080819, 0x19191908, 0x2b2b0819, 0x19191908,
+    0x08080808, 0x19191919, 0x08082b08, 0x19191919, 0x2b080808, 0x19191919, 0x2b082b08, 0x19191919,
+    0x082b0819, 0x1919192b, 0x192b2b08, 0x1919192b, 0x2b2b0819, 0x1919192b, 0x08080808, 0x19192b08,
+    0x08191908, 0x19192b08, 0x19080819, 0x19192b08, 0x19190808, 0x19192b08, 0x2b192b19, 0x19192b08,
+    0x08192b2b, 0x19192b19, 0x19080808, 0x19192b19, 0x1908082b, 0x19192b19, 0x2b081919, 0x19192b2b,
+    0x08080819, 0x192b0808, 0x08081908, 0x192b0808, 0x08190808, 0x192b0808, 0x19080808, 0x192b0808,
+    0x19191908, 0x192b0808, 0x192b082b, 0x192b0808, 0x2b08192b, 0x192b0808, 0x2b2b2b19, 0x192b0808,
+    0x08080808, 0x192b0819, 0x082b1908, 0x192b082b, 0x19082b2b, 0x192b082b, 0x2b19082b, 0x192b082b,
+    0x08080808, 0x192b1908, 0x0819192b, 0x192b1908, 0x08190808, 0x192b1919, 0x19080808, 0x192b1919,
+    0x19081919, 0x192b1919, 0x2b2b1908, 0x192b1919, 0x08080819, 0x192b2b08, 0x192b2b2b, 0x192b2b08,
+    0x082b1919, 0x192b2b19, 0x0808192b, 0x192b2b2b, 0x19191908, 0x192b2b2b, 0x192b082b, 0x192b2b2b,
+    0x08080808, 0x2b080808, 0x0808082b, 0x2b080808, 0x08081919, 0x2b080808, 0x08082b08, 0x2b080808,
+    0x08190819, 0x2b080808, 0x08191908, 0x2b080808, 0x082b0808, 0x2b080808, 0x082b2b2b, 0x2b080808,
+    0x19080819, 0x2b080808, 0x19081908, 0x2b080808, 0x19190808, 0x2b080808, 0x2b080808, 0x2b080808,
+    0x2b08082b, 0x2b080808, 0x2b2b2b08, 0x2b080808, 0x2b2b2b2b, 0x2b080808, 0x08080819, 0x2b080819,
+    0x08081908, 0x2b080819, 0x0808192b, 0x2b080819, 0x08190808, 0x2b080819, 0x19080808, 0x2b080819,
+    0x19190819, 0x2b080819, 0x19192b19, 0x2b080819, 0x08080808, 0x2b08082b, 0x082b0808, 0x2b08082b,
+    0x2b080808, 0x2b08082b, 0x2b08082b, 0x2b08082b, 0x2b2b0808, 0x2b08082b, 0x2b2b2b08, 0x2b08082b,
+    0x08080819, 0x2b081908, 0x08081908, 0x2b081908, 0x08190808, 0x2b081908, 0x0819082b, 0x2b081908,
+    0x08191919, 0x2b081908, 0x19080808, 0x2b081908, 0x192b0808, 0x2b081908, 0x2b082b19, 0x2b081908,
+    0x08080808, 0x2b081919, 0x19081908, 0x2b081919, 0x2b2b1919, 0x2b081919, 0x08192b08, 0x2b08192b,
+    0x192b2b2b, 0x2b08192b, 0x08080808, 0x2b082b08, 0x08082b08, 0x2b082b08, 0x082b1919, 0x2b082b08,
+    0x19192b2b, 0x2b082b08, 0x2b080808, 0x2b082b08, 0x2b08082b, 0x2b082b08, 0x2b2b2b08, 0x2b082b08,
+    0x0808192b, 0x2b082b19, 0x082b082b, 0x2b082b2b, 0x2b080808, 0x2b082b2b, 0x2b082b08, 0x2b082b2b,
+    0x2b19192b, 0x2b082b2b, 0x2b2b2b08, 0x2b082b2b, 0x08080819, 0x2b190808, 0x08081908, 0x2b190808,
+    0x08190808, 0x2b190808, 0x19080808, 0x2b190808, 0x1919192b, 0x2b190808, 0x2b081908, 0x2b190808,
+    0x08080808, 0x2b190819, 0x082b082b, 0x2b190819, 0x192b1908, 0x2b190819, 0x1919192b, 0x2b19082b,
+    0x2b082b19, 0x2b19082b, 0x08080808, 0x2b191908, 0x08081919, 0x2b191908, 0x19081908, 0x2b191908,
+    0x19190808, 0x2b191908, 0x19192b08, 0x2b191908, 0x082b2b19, 0x2b191919, 0x2b190808, 0x2b191919,
+    0x2b19082b, 0x2b191919, 0x19080819, 0x2b19192b, 0x19190819, 0x2b192b08, 0x2b2b192b, 0x2b192b08,
+    0x19082b19, 0x2b192b19, 0x08191919, 0x2b192b2b, 0x192b0808, 0x2b192b2b, 0x08080808, 0x2b2b0808,
+    0x0808082b, 0x2b2b0808, 0x08082b08, 0x2b2b0808, 0x08082b2b, 0x2b2b0808, 0x082b0808, 0x2b2b0808,
+    0x082b2b2b, 0x2b2b0808, 0x2b2b0808, 0x2b2b0808, 0x19190819, 0x2b2b0819, 0x19192b19, 0x2b2b0819,
+    0x2b2b192b, 0x2b2b0819, 0x08080808, 0x2b2b082b, 0x0808082b, 0x2b2b082b, 0x08082b08, 0x2b2b082b,
+    0x082b2b2b, 0x2b2b082b, 0x2b080808, 0x2b2b082b, 0x2b2b0808, 0x2b2b082b, 0x19080808, 0x2b2b1908,
+    0x2b191919, 0x2b2b1908, 0x192b1919, 0x2b2b192b, 0x2b192b08, 0x2b2b192b, 0x08082b2b, 0x2b2b2b08,
+    0x082b0808, 0x2b2b2b08, 0x082b082b, 0x2b2b2b08, 0x082b2b08, 0x2b2b2b08, 0x2b2b0808, 0x2b2b2b08,
+    0x2b2b2b08, 0x2b2b2b08, 0x08081908, 0x2b2b2b19, 0x2b081908, 0x2b2b2b19, 0x2b08192b, 0x2b2b2b19,
+    0x082b2b08, 0x2b2b2b2b, 0x082b2b2b, 0x2b2b2b2b, 0x2b190819, 0x2b2b2b2b, 0x2b2b2b2b, 0x2b2b2b2b
+};
+
+inline uint iq2_xs_signs(uint code) {
+    return code | ((popcount(code) & 1u) << 7);
+}
@@ -0,0 +1,71 @@
+// IQ2_XXS codebook, stored as pairs of four packed magnitudes.
+constant uint iq2_xxs_grid[512] = {
+    0x08080808, 0x08080808, 0x0808082b, 0x08080808, 0x08081919, 0x08080808, 0x08082b08, 0x08080808,
+    0x08082b2b, 0x08080808, 0x08190819, 0x08080808, 0x08191908, 0x08080808, 0x082b0808, 0x08080808,
+    0x082b082b, 0x08080808, 0x082b2b08, 0x08080808, 0x082b2b2b, 0x08080808, 0x19080819, 0x08080808,
+    0x19081908, 0x08080808, 0x19190808, 0x08080808, 0x19192b08, 0x08080808, 0x192b0819, 0x08080808,
+    0x192b1908, 0x08080808, 0x2b080808, 0x08080808, 0x2b08082b, 0x08080808, 0x2b082b2b, 0x08080808,
+    0x2b2b082b, 0x08080808, 0x08080819, 0x08080819, 0x08081908, 0x08080819, 0x08190808, 0x08080819,
+    0x08191919, 0x08080819, 0x19080808, 0x08080819, 0x2b081908, 0x08080819, 0x2b192b08, 0x08080819,
+    0x08080808, 0x0808082b, 0x0808082b, 0x0808082b, 0x082b082b, 0x0808082b, 0x2b08082b, 0x0808082b,
+    0x08080819, 0x08081908, 0x08081908, 0x08081908, 0x08190808, 0x08081908, 0x082b0819, 0x08081908,
+    0x082b1908, 0x08081908, 0x19080808, 0x08081908, 0x1908082b, 0x08081908, 0x19082b08, 0x08081908,
+    0x192b0808, 0x08081908, 0x2b080819, 0x08081908, 0x2b081908, 0x08081908, 0x2b190808, 0x08081908,
+    0x2b2b1908, 0x08081908, 0x08080808, 0x08081919, 0x0808082b, 0x08081919, 0x08082b08, 0x08081919,
+    0x082b0808, 0x08081919, 0x1908192b, 0x08081919, 0x192b2b19, 0x08081919, 0x2b080808, 0x08081919,
+    0x2b190819, 0x08081919, 0x08082b19, 0x0808192b, 0x08190808, 0x0808192b, 0x19080808, 0x0808192b,
+    0x2b081908, 0x0808192b, 0x2b2b1908, 0x0808192b, 0x08080808, 0x08082b08, 0x08081919, 0x08082b08,
+    0x08082b08, 0x08082b08, 0x08191908, 0x08082b08, 0x082b2b08, 0x08082b08, 0x19080819, 0x08082b08,
+    0x19081908, 0x08082b08, 0x19190808, 0x08082b08, 0x1919082b, 0x08082b08, 0x2b082b08, 0x08082b08,
+    0x08081908, 0x08082b19, 0x19080808, 0x08082b19, 0x0808082b, 0x08082b2b, 0x08191908, 0x08082b2b,
+    0x08080819, 0x08190808, 0x08081908, 0x08190808, 0x08190808, 0x08190808, 0x082b0819, 0x08190808,
+    0x19080808, 0x08190808, 0x192b0808, 0x08190808, 0x2b081908, 0x08190808, 0x2b190808, 0x08190808,
+    0x2b191919, 0x08190808, 0x08080808, 0x08190819, 0x08082b08, 0x08190819, 0x082b0808, 0x08190819,
+    0x19190808, 0x08190819, 0x19192b2b, 0x08190819, 0x2b080808, 0x08190819, 0x082b1908, 0x0819082b,
+    0x19081919, 0x0819082b, 0x08080808, 0x08191908, 0x08082b08, 0x08191908, 0x082b0808, 0x08191908,
+    0x082b1919, 0x08191908, 0x19082b19, 0x08191908, 0x2b080808, 0x08191908, 0x08192b08, 0x08191919,
+    0x192b082b, 0x08191919, 0x08080808, 0x0819192b, 0x0819192b, 0x0819192b, 0x08080819, 0x08192b08,
+    0x08081908, 0x08192b08, 0x08190808, 0x08192b08, 0x19080808, 0x08192b08, 0x2b080819, 0x08192b08,
+    0x08080808, 0x08192b19, 0x08081919, 0x08192b19, 0x2b2b0808, 0x08192b19, 0x19190819, 0x08192b2b,
+    0x08080808, 0x082b0808, 0x0808082b, 0x082b0808, 0x08082b2b, 0x082b0808, 0x19081908, 0x082b0808,
+    0x192b0819, 0x082b0808, 0x2b080808, 0x082b0808, 0x2b08082b, 0x082b0808, 0x082b2b19, 0x082b0819,
+    0x19082b08, 0x082b0819, 0x08080808, 0x082b082b, 0x0808082b, 0x082b082b, 0x08080819, 0x082b1908,
+    0x08081908, 0x082b1908, 0x08190808, 0x082b1908, 0x19080808, 0x082b1908, 0x1919192b, 0x082b1908,
+    0x08080808, 0x082b1919, 0x19080819, 0x082b1919, 0x192b1908, 0x082b1919, 0x2b190808, 0x082b192b,
+    0x08082b08, 0x082b2b08, 0x082b0808, 0x082b2b08, 0x2b191908, 0x082b2b08, 0x19081908, 0x082b2b2b,
+    0x08080819, 0x19080808, 0x08081908, 0x19080808, 0x08190808, 0x19080808, 0x08192b08, 0x19080808,
+    0x082b0819, 0x19080808, 0x082b1908, 0x19080808, 0x19080808, 0x19080808, 0x19082b08, 0x19080808,
+    0x1919192b, 0x19080808, 0x192b0808, 0x19080808, 0x2b080819, 0x19080808, 0x2b081908, 0x19080808,
+    0x2b190808, 0x19080808, 0x08080808, 0x19080819, 0x082b0808, 0x19080819, 0x192b0819, 0x19080819,
+    0x2b080808, 0x19080819, 0x2b081919, 0x19080819, 0x08080819, 0x1908082b, 0x08190808, 0x1908082b,
+    0x19082b08, 0x1908082b, 0x1919192b, 0x1908082b, 0x192b2b08, 0x1908082b, 0x08080808, 0x19081908,
+    0x08082b08, 0x19081908, 0x082b0808, 0x19081908, 0x2b080808, 0x19081908, 0x2b192b19, 0x19081908,
+    0x0819082b, 0x19081919, 0x082b1908, 0x19081919, 0x08080808, 0x1908192b, 0x08080819, 0x19082b08,
+    0x08081908, 0x19082b08, 0x08190808, 0x19082b08, 0x19080808, 0x19082b08, 0x19081919, 0x19082b08,
+    0x08080808, 0x19082b19, 0x19192b08, 0x19082b19, 0x192b0819, 0x19082b19, 0x2b08082b, 0x19082b19,
+    0x19081919, 0x19082b2b, 0x2b190808, 0x19082b2b, 0x08080808, 0x19190808, 0x08082b08, 0x19190808,
+    0x08190819, 0x19190808, 0x08192b19, 0x19190808, 0x082b0808, 0x19190808, 0x2b080808, 0x19190808,
+    0x2b082b08, 0x19190808, 0x08081908, 0x19190819, 0x1908082b, 0x19190819, 0x2b2b1908, 0x19190819,
+    0x2b190819, 0x1919082b, 0x2b190808, 0x19191908, 0x2b19082b, 0x19191908, 0x08082b2b, 0x19191919,
+    0x08080819, 0x1919192b, 0x19191908, 0x1919192b, 0x08080808, 0x19192b08, 0x08190819, 0x19192b08,
+    0x08192b19, 0x19192b08, 0x192b1908, 0x19192b08, 0x19080808, 0x19192b19, 0x08082b08, 0x19192b2b,
+    0x08081908, 0x192b0808, 0x08190808, 0x192b0808, 0x19080808, 0x192b0808, 0x192b2b08, 0x192b0808,
+    0x08080808, 0x192b0819, 0x19191919, 0x192b0819, 0x08192b08, 0x192b082b, 0x192b0808, 0x192b082b,
+    0x08080808, 0x192b1908, 0x08081919, 0x192b1908, 0x08190808, 0x192b1919, 0x0819082b, 0x192b1919,
+    0x2b081908, 0x192b1919, 0x1908082b, 0x192b2b08, 0x08080808, 0x2b080808, 0x0808082b, 0x2b080808,
+    0x08082b2b, 0x2b080808, 0x19080819, 0x2b080808, 0x2b08082b, 0x2b080808, 0x08081908, 0x2b080819,
+    0x08192b08, 0x2b080819, 0x19080808, 0x2b080819, 0x08190819, 0x2b08082b, 0x08080819, 0x2b081908,
+    0x08081908, 0x2b081908, 0x08190808, 0x2b081908, 0x08191919, 0x2b081908, 0x19080808, 0x2b081908,
+    0x192b0808, 0x2b081908, 0x08080808, 0x2b081919, 0x1908192b, 0x2b081919, 0x2b191908, 0x2b081919,
+    0x08082b19, 0x2b08192b, 0x19080808, 0x2b08192b, 0x192b0808, 0x2b08192b, 0x0808082b, 0x2b082b08,
+    0x08081908, 0x2b082b19, 0x08190819, 0x2b082b2b, 0x08081908, 0x2b190808, 0x08190808, 0x2b190808,
+    0x082b1908, 0x2b190808, 0x19080808, 0x2b190808, 0x2b2b0819, 0x2b190808, 0x0819192b, 0x2b190819,
+    0x2b080808, 0x2b190819, 0x19081919, 0x2b19082b, 0x08080808, 0x2b191908, 0x082b082b, 0x2b191908,
+    0x19081908, 0x2b191908, 0x19190819, 0x2b191919, 0x2b080819, 0x2b192b08, 0x082b0808, 0x2b192b19,
+    0x0808082b, 0x2b2b0808, 0x19190808, 0x2b2b0808, 0x2b081919, 0x2b2b0808, 0x08082b19, 0x2b2b0819,
+    0x08080808, 0x2b2b082b, 0x08192b08, 0x2b2b1908, 0x19190808, 0x2b2b2b08, 0x08081908, 0x2b2b2b19
+};
+
+inline uint iq2_xxs_signs(uint code) {
+    return code | ((popcount(code) & 1u) << 7);
+}
@@ -57,13 +57,17 @@ kernel void kernel_add_id_add_id_swiglu_oai(

    global const float * g_row  = (global const float *)(src_g  + i1*nb01_g + i2*nb02_g);
    global const float * u_row  = (global const float *)(src_u  + i1*nb01_u + i2*nb02_u);
-    global const float * gb_row = (global const float *)(src_gb + i11*nb11_g);
-    global const float * ub_row = (global const float *)(src_ub + i11*nb11_u);
    global       float * d_row  = (global       float *)(dst    + i1*nbd1   + i2*nbd2);

    for (int i0 = get_local_id(0); i0 < ne0; i0 += get_local_size(0)) {
-        float x0 = g_row[i0] + gb_row[i0];
-        float x1 = u_row[i0] + ub_row[i0];
+        float x0 = g_row[i0];
+        float x1 = u_row[i0];
+        if (i11 != -1) {
+            global const float * gb_row = (global const float *)(src_gb + i11*nb11_g);
+            global const float * ub_row = (global const float *)(src_ub + i11*nb11_u);
+            x0 += gb_row[i0];
+            x1 += ub_row[i0];
+        }

        x0 = min(x0, limit);
        x1 = max(min(x1, limit), -limit);
@@ -42,8 +42,11 @@ __kernel void kernel_moe_combine_bias_f32(
    float4 acc = (float4)(0.0f);
    for (int e = 0; e < k; ++e) {
        const int i11 = *((__global const int *)(I + (ulong)e*sizeof(int)));
-        __global const float * Brow = (__global const float *)(B + (ulong)i11*nb_b1) + r4*4u;
-        const float4 v = vload4(0, E + (uint)e*e1) + vload4(0, Brow);
+        float4 v = vload4(0, E + (uint)e*e1);
+        if (i11 != -1) {
+            __global const float * Brow = (__global const float *)(B + (ulong)i11*nb_b1) + r4*4u;
+            v += vload4(0, Brow);
+        }
        acc = mad(v, (float4)(W[(uint)e*w1]), acc);
    }

@@ -80,32 +80,6 @@ __kernel void kernel_moe_scatter(
    emap[tile_idx] = val;
 }

-// the gemm only writes the rows an expert owns, so the skipped rows (negative ids) are zeroed here
-__kernel void kernel_moe_zero_dst(
-    __global const int * input,
-    __global float * dst,
-    uint N,
-    uint topK,
-    uint n_experts,
-    uint ne0
-) {
-    uint n = get_global_id(0);
-    uint k = get_global_id(1);
-
-    if (n >= N || k >= topK) {
-        return;
-    }
-
-    if (input[n * n_experts + k] != -1) {
-        return;
-    }
-
-    __global float * dst_row = dst + ((ulong)n * topK + k) * ne0;
-    for (uint i = get_local_id(2); i < ne0; i += get_local_size(2)) {
-        dst_row[i] = 0.0f;
-    }
-}
-
 // Deterministic replacement for kernel_moe_scatter.
 //
 // kernel_moe_scatter takes each token's slot from atomic_inc(slot_counter[expert]),
@@ -0,0 +1,22 @@
+kernel void kernel_moe_zero_dst(
+    global const char * ids,
+    ulong offset_ids,
+    ulong nb_ids1,
+    global char * dst,
+    ulong offset_dst,
+    ulong nb1,
+    ulong nb2,
+    int ne0
+) {
+    const size_t i1 = get_group_id(0);
+    const size_t i2 = get_group_id(1);
+    const int expert = *((global const int *)(ids + offset_ids + i1*sizeof(int) + i2*nb_ids1));
+    if (expert != -1) {
+        return;
+    }
+
+    global float * row = (global float *)(dst + offset_dst + i1*nb1 + i2*nb2);
+    for (size_t i0 = get_local_id(2); i0 < (size_t)ne0; i0 += get_local_size(2)) {
+        row[i0] = 0.0f;
+    }
+}
--- a/Show More
+++ b/Show More