tests: tolerate a shared pool abort in test_completion_unified (#28759 )

The expected success table holds when the four requests enter the shared pool together. On a loaded runner they are admitted tens of milliseconds apart, the slot lifetimes overlap differently and the pool overflows while a short request is still resident. The decode failure aborts every slot, so a request the table marks as successful comes back with the context error instead of its generation. Such a request now passes on that error too, while any other status, a different error or a truncated generation still fails the test.
ggml: skip 0-sized ids tensor when offloading selected experts (#28739 )
2026-09-11 17:37:10 +02:00 · 2026-09-11 15:50:12 +02:00 · 2026-09-11 15:17:08 +02:00 · 2026-09-11 14:12:55 +03:00 · 2026-09-11 13:01:29 +02:00 · 2026-09-11 12:59:43 +02:00
116 changed files with 6425 additions and 3278 deletions
@@ -137,19 +137,19 @@ runs:
      run: |
          mkdir -p "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4"
          choco install unzip -y
-          curl -O "https://packages.nvidia.com/bin-archive/pool/windows-x86_64/5B515474-7E78-11F1-8656-C51E4F4B317F/cccl-windows-x86_64-13.3.4.1.2-archive.zip"
-          curl -O "https://packages.nvidia.com/bin-archive/pool/windows-x86_64/5B515474-7E78-11F1-8656-C51E4F4B317F/cuda_crt-windows-x86_64-13.4.46-archive.zip"
-          curl -O "https://packages.nvidia.com/bin-archive/pool/windows-x86_64/5B515474-7E78-11F1-8656-C51E4F4B317F/cuda_nvcc-windows-x86_64-13.4.46-archive.zip"
-          curl -O "https://packages.nvidia.com/bin-archive/pool/windows-x86_64/5B515474-7E78-11F1-8656-C51E4F4B317F/libnvvm-windows-x86_64-13.4.46-archive.zip"
-          curl -O "https://packages.nvidia.com/bin-archive/pool/windows-arm64/5B515474-7E78-11F1-8656-C51E4F4B317F/cuda_cudart-windows-arm64-13.4.46-archive.zip"
-          curl -O "https://packages.nvidia.com/bin-archive/pool/windows-arm64/5B515474-7E78-11F1-8656-C51E4F4B317F/libcublas-windows-arm64-13.7.0.10-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cccl/windows-x86_64/cccl-windows-x86_64-13.3.4.2.1-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_crt/windows-x86_64/cuda_crt-windows-x86_64-13.4.59-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/windows-x86_64/cuda_nvcc-windows-x86_64-13.4.59-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/libnvvm/windows-x86_64/libnvvm-windows-x86_64-13.4.59-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/windows-arm64/cuda_cudart-windows-arm64-13.4.49-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/libcublas/windows-arm64/libcublas-windows-arm64-13.7.0.27-archive.zip"
          unzip '*.zip' -d "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4"
-          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cccl-windows-x86_64-13.3.4.1.2-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
-          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cuda_crt-windows-x86_64-13.4.46-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
-          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cuda_nvcc-windows-x86_64-13.4.46-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
-          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\libnvvm-windows-x86_64-13.4.46-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
-          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cuda_cudart-windows-arm64-13.4.46-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
-          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\libcublas-windows-arm64-13.7.0.10-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cccl-windows-x86_64-13.3.4.2.1-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cuda_crt-windows-x86_64-13.4.59-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cuda_nvcc-windows-x86_64-13.4.59-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\libnvvm-windows-x86_64-13.4.59-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cuda_cudart-windows-arm64-13.4.49-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\libcublas-windows-arm64-13.7.0.27-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
          echo "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\bin" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append
          echo "CUDA_PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
          echo "CUDA_PATH_V13_4=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
@@ -106,83 +106,3 @@ jobs:
          wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories260K/stories260K.bin
          ./bin/llama-convert-llama2c-to-ggml --copy-vocab-from-model ./tok512.bin --llama2c-model stories260K.bin --llama2c-output-model stories260K.gguf
          ./bin/llama-completion -m stories260K.gguf -p "One day, Lily met a Shoggoth" -n 500 -c 256
-
-  ubuntu-riscv64-native-sanitizer:
-    runs-on: ubuntu-24.04-riscv
-
-    continue-on-error: true
-
-    strategy:
-      matrix:
-        sanitizer: [ADDRESS, THREAD, UNDEFINED]
-        build_type: [Debug]
-
-    steps:
-      - name: Install dependencies
-        run: |
-          # Set gcc-14 and g++-14 as the default compilers
-          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-14 100
-          sudo update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-14 100
-
-          git lfs install
-
-      - name: GCC version check
-        run: |
-          gcc --version
-          g++ --version
-
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      # note: sparing some ccache since these jobs run on dedicated runners that are not part of the organitzation
-      #- name: ccache
-      #  uses: ggml-org/ccache-action@v1.2.24
-      #  with:
-      #    key: riscv-ubuntu-native-sanitizer-${{ matrix.sanitizer }}-${{ matrix.build_type }}
-      #    evict-old-files: 1d
-      #    save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-
-      - name: Build
-        id: cmake_build
-        if: ${{ matrix.sanitizer != 'THREAD' }}
-        run: |
-          cmake -B build \
-            -DLLAMA_OPENSSL=OFF \
-            -DCMAKE_BUILD_TYPE=${{ matrix.build_type }} \
-            -DGGML_OPENMP=ON \
-            -DLLAMA_BUILD_EXAMPLES=ON \
-            -DLLAMA_BUILD_TOOLS=ON \
-            -DLLAMA_BUILD_TESTS=OFF \
-            -DCMAKE_C_COMPILER_LAUNCHER=ccache \
-            -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
-            -DLLAMA_SANITIZE_${{ matrix.sanitizer }}=ON \
-            -DCMAKE_C_COMPILER=riscv64-linux-gnu-gcc-14 \
-            -DCMAKE_CXX_COMPILER=riscv64-linux-gnu-g++-14
-
-          cmake --build build --config ${{ matrix.build_type }} -j $(nproc)
-
-      - name: Build (no OpenMP)
-        id: cmake_build_no_openmp
-        if: ${{ matrix.sanitizer == 'THREAD' }}
-        run: |
-          cmake -B build \
-            -DLLAMA_OPENSSL=OFF \
-            -DCMAKE_BUILD_TYPE=${{ matrix.build_type }} \
-            -DGGML_OPENMP=OFF \
-            -DLLAMA_BUILD_EXAMPLES=ON \
-            -DLLAMA_BUILD_TOOLS=ON \
-            -DLLAMA_BUILD_TESTS=OFF \
-            -DCMAKE_C_COMPILER_LAUNCHER=ccache \
-            -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
-            -DLLAMA_SANITIZE_${{ matrix.sanitizer }}=ON \
-            -DCMAKE_C_COMPILER=riscv64-linux-gnu-gcc-14 \
-            -DCMAKE_CXX_COMPILER=riscv64-linux-gnu-g++-14
-
-          cmake --build build --config ${{ matrix.build_type }} -j $(nproc)
-
-      - name: Test
-        id: cmake_test
-        run: |
-          cd build
-          ctest -L main --verbose --timeout 900
@@ -101,8 +101,6 @@ jobs:

      - name: Test
        id: cmake_test
-        # skip run in Debug - very slow
-        if: ${{ matrix.sanitizer != 'UNDEFINED' }}
        run: |
          cd build
          ctest -L main -E tokenizer --verbose --timeout 900
@@ -58,18 +58,48 @@ env:

 jobs:
  gpu-cuda:
-    runs-on: [self-hosted, Linux, NVIDIA]
+    runs-on: "hf-jobs-t4-small:cuda13"

    steps:
      - name: Clone
        id: checkout
        uses: actions/checkout@v6

+      - name: Install dependencies
+        run: |
+          sudo apt update
+          sudo apt install -y cmake libssl-dev time unzip wget python3 python3-venv python3-pip
+
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.24
+        with:
+          restore: false
+          save: false
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        with:
+          key: self-hosted-gpu-cuda
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
+
      - name: Test
        id: ggml-ci
        run: |
          nvidia-smi
-          GG_BUILD_CUDA=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
+          GG_BUILD_CUDA=1 CUDACXX=/usr/local/cuda/bin/nvcc bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
+
+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: self-hosted-gpu-cuda
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true

  gpu-rocm:
    runs-on: [self-hosted, Linux, AMD]
@@ -0,0 +1,67 @@
+name: Fusion
+
+on:
+  workflow_dispatch: # allows manual triggering
+  push:
+    branches:
+      - master
+    paths: [
+      '.github/workflows/fusion.yml',
+      'ggml/**',
+      'tests/fusion/**',
+      'tests/test-fusion.cpp'
+    ]
+
+  pull_request:
+    types: [opened, synchronize, reopened]
+    paths: [
+      '.github/workflows/fusion.yml',
+      'ggml/**',
+      'tests/fusion/**',
+      'tests/test-fusion.cpp'
+    ]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
+  cancel-in-progress: true
+
+env:
+  GGML_NLOOP: 3
+  GGML_N_THREADS: 1
+  LLAMA_ARG_LOG_COLORS: 1
+  LLAMA_ARG_LOG_PREFIX: 1
+  LLAMA_ARG_LOG_TIMESTAMPS: 1
+
+jobs:
+  # TODO: add jobs for other backends as they adopt the fusion debug API
+  metal:
+    runs-on: [self-hosted, macOS, ARM64]
+
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+
+      - name: Build
+        id: cmake_build
+        run: |
+          cmake -B build \
+            -DCMAKE_BUILD_TYPE=Release \
+            -DLLAMA_FATAL_WARNINGS=ON \
+            -DLLAMA_OPENSSL=OFF \
+            -DGGML_SCHED_NO_REALLOC=ON \
+            -DGGML_BLAS=OFF \
+            -DGGML_METAL=ON
+          time cmake --build build --config Release --target test-llama-archs -j $(sysctl -n hw.logicalcpu)
+          time cmake --build build --config Release --target test-fusion -j $(sysctl -n hw.logicalcpu)
+
+      - name: Generate models
+        id: generate_models
+        run: |
+          rm -rf build-ci-models && mkdir -p build-ci-models
+          ./build/bin/test-llama-archs -o build-ci-models
+
+      - name: Test fusion
+        id: test_fusion
+        run: |
+          ./build/bin/test-fusion --models build-ci-models --device MTL0 --check tests/fusion/MTL.csv
@@ -1717,7 +1717,7 @@ jobs:
            - [Windows arm64 (OpenCL Adreno)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-opencl-adreno-arm64.zip)
            - [Windows x64 (CUDA 12)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-cuda-12.4-x64.zip) - [CUDA 12.4 DLLs](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-bin-win-cuda-12.4-x64.zip)
            - [Windows x64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-cuda-13.3-x64.zip) - [CUDA 13.3 DLLs](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-bin-win-cuda-13.3-x64.zip)
-            - [Windows arm64 (CUDA 13) (preview)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-cuda-13.4-arm64.zip) - [CUDA 13.4 DLLs](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-bin-win-cuda-13.4-arm64.zip)
+            - [Windows arm64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-cuda-13.4-arm64.zip) - [CUDA 13.4 DLLs](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-bin-win-cuda-13.4-arm64.zip)
            - [Windows x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-vulkan-x64.zip)
            - [Windows x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-openvino-${{ needs.windows-openvino.outputs.openvino_version }}-x64.zip)
            - [Windows x64 (SYCL)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-sycl-x64.zip)
@@ -32,6 +32,8 @@ on:
    ]

 env:
+  # note: this is dud token to avoid rate limiting (https://github.com/ggml-org/llama.cpp/pull/25706#issuecomment-4979941302)
+  HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
  LLAMA_ARG_LOG_COLORS: 1
  LLAMA_ARG_LOG_PREFIX: 1
  LLAMA_ARG_LOG_TIMESTAMPS: 1
@@ -43,7 +45,7 @@ concurrency:

 jobs:
  server:
-    runs-on: [self-hosted, CPU, Linux, llama-server]
+    runs-on: hf-jobs-cpu-upgrade

    strategy:
      matrix:
@@ -52,20 +54,6 @@ jobs:
      fail-fast: false

    steps:
-      #- name: Dependencies
-      #  id: depends
-      #  run: |
-      #    sudo apt-get update
-      #    sudo apt-get -y install \
-      #      build-essential \
-      #      xxd \
-      #      git \
-      #      cmake \
-      #      curl \
-      #      wget \
-      #      language-pack-en \
-      #      libssl-dev
-
      - name: Clone
        id: checkout
        uses: actions/checkout@v6
@@ -73,6 +61,24 @@ jobs:
          fetch-depth: 0
          ref: ${{ github.event.inputs.sha || github.event.pull_request.head.sha || github.sha || github.head_ref || github.ref_name }}

+      - name: Install dependencies
+        run: |
+          sudo apt update
+          sudo apt install -y build-essential cmake python3-full
+
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.24
+        with:
+          restore: false
+          save: false
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        with:
+          key: server-sanitize-${{ matrix.sanitizer }}
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
+
      - name: Build
        id: cmake_build
        run: |
@@ -87,9 +93,17 @@ jobs:
            -DLLAMA_SANITIZE_UNDEFINED=${{ matrix.sanitizer == 'UNDEFINED' }}
          cmake --build build --config ${{ matrix.build_type }} -j $(nproc) --target llama-server

-      - name: Python setup
-        id: setup_python
-        uses: actions/setup-python@v7
+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: server-sanitize-${{ matrix.sanitizer }}
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true

      - name: Install Python dependencies
        run: |
@@ -334,6 +334,35 @@ function gg_sum_test_llama_archs_tensor_split {
    gg_printf '```\n'
 }

+# test_llama_archs_models
+
+function gg_run_test_llama_archs_models {
+    cd ${SRC}
+
+    set -e
+
+    # TODO: fix and re-enable `test-llama-archs` on OpenVINO
+    # TODO: the `test-llama-archs` currently does not build on Windows, so we check if the binary exists
+    if [ -z ${GG_BUILD_OPENVINO} ] && [ -f ./build-ci-release/bin/test-llama-archs ]; then
+        rm -rf build-ci-models && mkdir -p build-ci-models
+
+        # generate the dummy models used by the model-dependent tests
+        ./build-ci-release/bin/test-llama-archs -o build-ci-models 2>&1
+    fi
+
+    set +e
+}
+
+function gg_sum_test_llama_archs_models {
+    gg_printf '### %s\n\n' "${ci}"
+
+    gg_printf 'Generates the dummy models used by the model-dependent tests\n'
+    gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
+    gg_printf '```\n'
+    gg_printf '%s\n' "$(cat $OUT/${ci}.log)"
+    gg_printf '```\n'
+}
+
 # test_scripts

 function gg_run_test_scripts {
@@ -790,6 +819,7 @@ ret=0
 test $ret -eq 0 && gg_run ctest_debug
 test $ret -eq 0 && gg_run ctest_release

+test $ret -eq 0 && gg_run test_llama_archs_models
 test $ret -eq 0 && gg_run test_llama_archs_tensor_split

 if [ ! -z ${GG_BUILD_HIGH_PERF} ]; then
@@ -134,6 +134,8 @@ set_target_properties(${TARGET} PROPERTIES
 target_include_directories(${TARGET} PUBLIC .)
 target_link_libraries     (${TARGET} PUBLIC vendor::nlohmann vendor::sheredom)
 target_compile_features   (${TARGET} PUBLIC cxx_std_17)
+target_precompile_headers (${TARGET} PRIVATE common.h)
+target_precompile_headers (${TARGET} PRIVATE chat.h)

 if (LLAMA_SUBPROCESS)
    target_compile_definitions(${TARGET} PUBLIC LLAMA_SUBPROCESS)
@@ -296,7 +296,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
            drafting[seq_id] = true;
            common_sampler_reset(smpls[seq_id].get());

-            common_batch_add(batch, dp.id_last, dp.n_past, { seq_id }, true);
+            common_batch_add(batch, dp.id_last, dp.pos0, { seq_id }, true);
        }

        int ret = llama_decode(ctx_dft, batch);
@@ -355,7 +355,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
                    continue;
                }

-                common_batch_add(batch, id, dp.n_past + i + 1, { seq_id }, true);
+                common_batch_add(batch, id, dp.pos0 + i + 1, { seq_id }, true);
            }

            if (batch.n_tokens == 0) {
@@ -1094,8 +1094,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

        // Target prefill may contain token IDs or multimodal embeddings. Both
        // produce the target-layer features used to seed the draft KV cache, so
-        // skipping the embedding batches leaves a hole in the draft's cache and
-        // the next injection fails to initialize.
+        // embeddings are injected too, except the pinned ones skipped below.
        // TODO: revisit after https://github.com/ggml-org/llama.cpp/pull/24669 is merged
        const bool has_tokens     = batch_in.token != nullptr;
        const bool has_embeddings = batch_in.embd  != nullptr;
@@ -1131,6 +1130,13 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
            }
            const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1;

+            // an M-RoPE image pins all its rows to one position, so a windowed draft
+            // cache cannot free cells for it - skip it, the draft can jump over the gap
+            const bool pos_pinned = batch_in.pos[i_batch_beg[seq_id]] == batch_in.pos[i_batch_end[seq_id]];
+            if (has_embeddings && n_rows > 1 && pos_pinned) {
+                continue;
+            }
+
            for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) {
                const int32_t n_chunk = std::min(n_ubatch, n_rows - offset);

@@ -1191,7 +1197,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

            common_sampler_reset(smpls[seq_id].get());

-            const int32_t n = (int32_t) dp.n_past;
+            const int32_t n = (int32_t) dp.pos0;

            const int32_t n_draft = params.n_max;

@@ -1487,7 +1493,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {

        const int32_t n_tokens = batch_in.n_tokens;

-        // remember the frist and last batch index for each sequence
+        // remember the first and last batch index for each sequence
        std::fill(i_batch_beg.begin(), i_batch_beg.end(), -1);
        std::fill(i_batch_end.begin(), i_batch_end.end(), -1);

@@ -1615,7 +1621,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
            drafting[seq_id] = true;
            common_sampler_reset(smpls[seq_id].get());

-            common_batch_add(batch, dp.id_last, dp.n_past, { seq_id }, true);
+            common_batch_add(batch, dp.id_last, dp.pos0, { seq_id }, true);
            std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, pending_h[seq_id].data(), row_bytes);

            i_last[seq_id] = batch.n_tokens - 1;
@@ -1629,16 +1635,16 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {

        while (n_drafting > 0) {
            // each step decodes under a different head, i.e. a different decoder layer, and
-            // KV is per layer. process() filled this layer's KV only for positions < n_past
+            // KV is per layer. process() filled this layer's KV only for positions < pos0
            // (prompt + accepted prefix) — nothing in the draft region yet. so reset the
-            // draft region (the seq_rm lower bound is n_past, leaving the prompt KV intact)
+            // draft region (the seq_rm lower bound is pos0, leaving the prompt KV intact)
            // and select head i so it rebuilds its own layer's KV there; decoding just the
            // latest token would leave its attention reading cells only another head wrote.
            if (chain_heads) {
                auto * mem_dft = llama_get_memory(ctx_dft);
                for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
                    if (drafting[seq_id]) {
-                        llama_memory_seq_rm(mem_dft, seq_id, dparams[seq_id].n_past, -1);
+                        llama_memory_seq_rm(mem_dft, seq_id, dparams[seq_id].pos0, -1);
                    }
                }
                llama_set_nextn_layer_offset(ctx_dft, i);
@@ -1704,17 +1710,17 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
                    const int n_rows = (int) result.size() + 1; // id_last + tokens drafted so far
                    for (int t = 0; t < n_rows; ++t) {
                        const llama_token tok = (t == 0) ? dp.id_last : result[t - 1];
-                        common_batch_add(batch, tok, dp.n_past + t, { seq_id }, t == n_rows - 1);
+                        common_batch_add(batch, tok, dp.pos0 + t, { seq_id }, t == n_rows - 1);
                        std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd,
                                    chain_h[seq_id].data() + (size_t) t * n_embd, row_bytes);
                    }
                } else if (is_mem_shared) {
                    // note: with shared memory (e.g. Gemma4 assistants) we use the same position for all draft tokens
                    // ref: https://github.com/huggingface/transformers/blob/effde20942e3f82a1b97449f60b3a48c5ff96145/docs/source/en/model_doc/gemma4_assistant.md?plain=1#L36-L37
-                    common_batch_add(batch, id, dp.n_past, { seq_id }, true);
+                    common_batch_add(batch, id, dp.pos0, { seq_id }, true);
                    std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, h_row, row_bytes);
                } else {
-                    common_batch_add(batch, id, dp.n_past + i + 1, { seq_id }, true);
+                    common_batch_add(batch, id, dp.pos0 + i + 1, { seq_id }, true);
                    std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, h_row, row_bytes);
                }

@@ -61,7 +61,7 @@ struct common_speculative_draft_params {
    // can be used to constraint the max draft based on the remaining context size
    int32_t n_max = -1;

-    llama_pos   n_past;
+    llama_pos   pos0;
    llama_token id_last;

    // TODO: remove in the future by keeping track of the prompt from the _begin() call and the consecutive accept calls
@@ -216,14 +216,14 @@ class NemotronHModel(GraniteHybridModel):
        hparams = kwargs.pop("hparams", None)
        if hparams is None:
            hparams = ModelBase.load_hparams(args[0], self.is_mistral_format)
-        has_moe_params = (
-            "num_experts_per_tok" in hparams
-            or (isinstance(hparams.get("llm_config"), dict) and "num_experts_per_tok" in hparams["llm_config"])
-        )
+        llm_config = {**hparams, **(hparams.get("llm_config") or {})}
+
+        has_moe_params = "num_experts_per_tok" in llm_config
+        layers_block_type = llm_config.get("layers_block_type")
+
        if has_moe_params:
            self.model_arch = gguf.MODEL_ARCH.NEMOTRON_H_MOE
            self.is_moe = True
-        layers_block_type = hparams.get("layers_block_type")
        if layers_block_type is not None:
            hparams["num_hidden_layers"] = len(layers_block_type)

@@ -0,0 +1,122 @@
+## Build profiling
+This page is a working document for analyzing the current build and try to
+identify ways to improve the build time.
+
+### Requirements
+The profiling script requires clang to be used as the compiler tool chain and
+also requires that ClangBuildAnalyzer is installed.
+
+Mac:
+```console
+brew install clang-build-analyzer
+```
+
+Linux:
+```console
+git clone https://github.com/aras-p/ClangBuildAnalyzer.git
+cd ClangBuildAnalyzer
+cmake -B build -DCMAKE_BUILD_TYPE=Release
+cmake --build build -j$(nproc)
+sudo cp build/ClangBuildAnalyzer /usr/local/bin/
+```
+
+Windows: install LLVM/clang and Ninja (e.g. via the
+[LLVM releases page](https://github.com/llvm/llvm-project/releases) and
+`winget install Ninja-build.Ninja`), then build ClangBuildAnalyzer the same
+way as on Linux:
+```console
+git clone https://github.com/aras-p/ClangBuildAnalyzer.git
+cd ClangBuildAnalyzer
+cmake -B build -G Ninja -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_BUILD_TYPE=Release
+cmake --build build --config Release
+```
+Then add `ClangBuildAnalyzer\build` to `PATH`.
+
+### Usage
+Mac/Linux:
+```console
+$ ./scripts/build-profile.sh
+```
+
+Windows:
+```console
+> .\scripts\build-profile.ps1
+```
+
+Both accept `--full`/`-Full` (include Server, Tools, and Tests) and a jobs
+override (`-jN` / `-Jobs N`).
+
+Note: on Windows, `cmake` defaults to the Visual Studio generator, which
+ignores `CMAKE_C_COMPILER`/`CMAKE_CXX_COMPILER` and silently falls back to
+MSVC. `build-profile.ps1` passes `-G Ninja` so clang is actually used, this
+is required on ARM64.
+
+### Linux (Ubuntu 24.04)
+
+Environment:
+- Clang:     18.1.3 (Ubuntu clang version 18.1.3 (1ubuntu1))
+- libstdc++: GCC 13.3.0 (Ubuntu 13.3.0-6ubuntu2~24.04.1)
+- Target:    x86_64-pc-linux-gnu
+
+```console
+------------------------+-----+------------+------------+------------+
+| Build                  | TUs | Frontend   | Backend    | Total      |
+------------------------+-----+------------+------------+------------+
+| Minimal, master        | 249 |   468.2 s  |   270.3 s  |   738.5 s  |
+| Minimal, with PCH      | 253 |   177.1 s  |   265.8 s  |   442.9 s  |
+| Full,    master        | 396 |   811.0 s  |   692.2 s  | 1,503.2 s  |
+| Full,    with PCH      | 405 |   380.0 s  |   664.7 s  | 1,044.7 s  |
+| Full,    with PCH + UB | 264 |   357.7 s  |   635.7 s  |   993.4 s  |
+------------------------+-----+------------+------------+------------+
+
+PCH  = precompiled header.
+Full = includes building Server, Tools, and Tests.
+UB   = unity build for models
+```
+Note that the number of translation units (TUs) increases when using precompiled
+headers — each PCH target adds one extra TU for the precompilation step itself.
+
+### Mac (Apple M3)
+
+Environment:
+- Clang:  Apple clang version 17.0.0 (clang-1700.3.19.1)
+- libc++: ships with Apple clang 17.0.0 (Xcode toolchain)
+- Target: arm64-apple-macosx15.6
+
+```console
+------------------------+-----+------------+------------+------------+
+| Build                  | TUs | Frontend   | Backend    | Total      |
+------------------------+-----+------------+------------+------------+
+| Minimal, master        | 256 |   154.5 s  |    94.8 s  |   249.3 s  |
+| Minimal, with PCH      | 261 |    65.9 s  |    90.0 s  |   155.9 s  |
+| Full,    master        | 407 |   265.7 s  |   209.7 s  |   475.4 s  |
+| Full,    with PCH      | 414 |   154.6 s  |   197.5 s  |   352.1 s  |
+| Full,    with PCH + UB | 274 |   143.0 s  |   192.2 s  |   335.2 s  |
+------------------------+-----+------------+------------+------------+
+
+PCH = precompiled header.
+Full = includes building Server, Tools, and Tests.
+UB   = unity build for models
+```
+
+### Windows (ARM64)
+
+Environment:
+- Clang:  clang version 22.1.8 (LLVM, `C:\Program Files\LLVM`)
+- STL:    MSVC STL (Visual Studio 2022 Build Tools 14.44.35207)
+- Target: aarch64-pc-windows-msvc
+
+```console
+------------------------+-----+------------+------------+------------+
+| Build                  | TUs | Frontend   | Backend    | Total      |
+------------------------+-----+------------+------------+------------+
+| Minimal, master        | 249 |   159.4 s  |    82.2 s  |   241.6 s  |
+| Full,    master        | 373 |   337.2 s  |   167.4 s  |   504.6 s  |
+| Minimal, with PCH + UB | 113 |    62.3 s  |    82.4 s  |   144.7 s  |
+| Full,    with PCH + UB | 240 |   233.0 s  |   185.1 s  |   418.1 s  |
+------------------------+-----+------------+------------+------------+
+
+PCH = precompiled header.
+Full = includes building Server, Tools, and Tests.
+UB   = unity build for models
+```
@@ -243,6 +243,7 @@ IBM VXE/VXE2 SIMD acceleration depends on the BLAS implementation. It is strongl
 | FP32       | ✅           | ✅    | ❓     |
 | FP16       | ✅           | ✅    | ❓     |
 | BF16       | ✅           | ✅    | ❓     |
+| Q1_0       | ✅           | ❓    | ❓     |
 | Q4_0       | ✅           | ❓    | ❓     |
 | Q4_1       | ✅           | ❓    | ❓     |
 | MXFP4      | ✅           | ❓    | ❓     |
@@ -272,4 +273,4 @@ IBM VXE/VXE2 SIMD acceleration depends on the BLAS implementation. It is strongl
 -   🚫 - acceleration unavailable, will still run using scalar implementation
 -   ❓ - acceleration unknown, please contribute if you can test it yourself

-Last Updated by **Aaron Teo (aaron.teo1@ibm.com)** on Feb 15, 2026.
+Last Updated by **Aaron Teo (aaron.teo1@ibm.com)** on Sep 8, 2026.
@@ -300,7 +300,8 @@ The following compilation options are also available to tweak performance:
 |-------------------------------|------------------------|---------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | GGML_CUDA_FORCE_MMQ           | Boolean                | false   | Force the use of custom matrix multiplication kernels for quantized models instead of FP16 cuBLAS even if there is no int8 tensor core implementation available (affects V100, CDNA and RDNA3+). MMQ kernels are enabled by default on GPUs with int8 tensor core support. With MMQ force enabled, speed for large batch sizes will be worse but VRAM consumption will be lower. |
 | GGML_CUDA_FORCE_CUBLAS        | Boolean                | false   | Force the use of FP16 cuBLAS instead of custom matrix multiplication kernels for quantized models. There may be issues with numerical overflows (except for V100, CDNA and RDNA4 which use FP32 compute type by default) and memory use will be higher. Prompt processing may become faster on recent datacenter GPUs (the custom kernels were tuned primarily for RTX 3000/4000).   |
-| GGML_CUDA_FA_ALL_QUANTS       | Boolean                | false   | Compile support for all KV cache quantization type (combinations) for the FlashAttention CUDA kernels. More fine-grained control over KV cache size but compilation takes much longer.                                                                                                                                                                                           |
+| GGML_CUDA_FA_QUANTS           | `all` or `type_K-type_V` list | q4_0-q4_0;q8_0-q8_0;f16-f16;bf16-bf16 | Select which K/V type combinations to compile the FlashAttention CUDA kernels for. `all` compiles every combination, but compilation takes much longer. Otherwise a `;`-separated list of `type_K-type_V` pairs; f16-f16 is always compiled. Combinations that were not compiled fall back to f16-f16 kernel with a warning. Legal types: f16, bf16, q4_0, q4_1, q5_0, q5_1, q8_0. |
+| GGML_CUDA_FA_ALL_QUANTS       | Boolean                | false   | Deprecated alias for `GGML_CUDA_FA_QUANTS=all`.                                                                                                                                                                                                                                                                                                                               |

 ## MUSA

@@ -188,7 +188,7 @@ int main(int argc, char ** argv) {
            common_speculative_get_draft_params(spec, seq_id) = {
                /* .drafting   = */ true,
                /* .n_max      = */ n_draft_max,
-                /* .n_past     = */ n_past,
+                /* .pos0       = */ n_past,
                /* .id_last    = */ id_last,
                /* .prompt     = */ &prompt_tgt,
                /* .result     = */ &draft, // output
@@ -204,6 +204,8 @@ option(GGML_CUDA_NO_PEER_COPY               "ggml: do not use peer to peer copie
 option(GGML_CUDA_NO_VMM                     "ggml: do not try to use CUDA VMM"                OFF)
 option(GGML_CUDA_FA                         "ggml: compile ggml FlashAttention CUDA kernels"  ON)
 option(GGML_CUDA_FA_ALL_QUANTS              "ggml: compile all quants for FlashAttention"     OFF)
+set   (GGML_CUDA_FA_QUANTS "q4_0-q4_0;q8_0-q8_0;f16-f16;bf16-bf16" CACHE STRING
+                                            "ggml: FlashAttention K-V type combinations to compile, \"all\" or a list such as \"q8_0-q8_0;q8_0-q4_0\"")
 option(GGML_CUDA_GRAPHS                     "ggml: use CUDA graphs (llama.cpp only)"          ${GGML_CUDA_GRAPHS_DEFAULT})
 option(GGML_CUDA_NCCL                       "ggml: use NVIDIA Collective Comm. Library"       ON)
 set   (GGML_CUDA_COMPRESSION_MODE "size" CACHE STRING
@@ -48,3 +48,74 @@ function(ggml_get_system_arch)
        set(GGML_SYSTEM_ARCH "UNKNOWN" PARENT_SCOPE)
    endif()
 endfunction()
+
+# Determines which FlashAttention vector kernel template instances to compile, returns them in OUT_SRCS.
+function(ggml_cuda_fattn_vec_instances DIR OUT_SRCS)
+    set(FA_TYPES q4_0 q4_1 q5_0 q5_1 q8_0 bf16 f16)
+
+    string(TOLOWER "${GGML_CUDA_FA_QUANTS}" FA_QUANTS)
+    string(STRIP   "${FA_QUANTS}" FA_QUANTS)
+    if (GGML_CUDA_FA_ALL_QUANTS)
+        message(WARNING "GGML_CUDA_FA_ALL_QUANTS is deprecated, use GGML_CUDA_FA_QUANTS=all instead")
+        set(FA_QUANTS all)
+    endif()
+    if (NOT FA_QUANTS)
+        message(FATAL_ERROR "GGML_CUDA_FA_QUANTS must not be empty")
+    endif()
+
+    if (FA_QUANTS STREQUAL "all")
+        set(FA_COMBINATIONS "")
+        foreach (TYPE_V IN LISTS FA_TYPES)
+            foreach (TYPE_K IN LISTS FA_TYPES)
+                list(APPEND FA_COMBINATIONS ${TYPE_K}-${TYPE_V})
+            endforeach()
+        endforeach()
+    else()
+        set(FA_COMBINATIONS f16-f16)
+
+        string(REPLACE "," ";" FA_SELECTED "${FA_QUANTS}")
+        foreach (COMBINATION IN LISTS FA_SELECTED)
+            string(STRIP   "${COMBINATION}" COMBINATION)
+            if (NOT COMBINATION MATCHES "^([a-z0-9_]+)-([a-z0-9_]+)$")
+                message(FATAL_ERROR "GGML_CUDA_FA_QUANTS: \"${COMBINATION}\" is not \"all\" or a <type_K>-<type_V> combination")
+            endif()
+            set(TYPE_K ${CMAKE_MATCH_1})
+            set(TYPE_V ${CMAKE_MATCH_2})
+            foreach (TYPE ${TYPE_K} ${TYPE_V})
+                if (NOT TYPE IN_LIST FA_TYPES)
+                    message(FATAL_ERROR
+                        "GGML_CUDA_FA_QUANTS: unknown type \"${TYPE}\" in \"${COMBINATION}\", must be one of: ${FA_TYPES}")
+                endif()
+            endforeach()
+            list(APPEND FA_COMBINATIONS ${TYPE_K}-${TYPE_V})
+        endforeach()
+    endif()
+    list(REMOVE_DUPLICATES FA_COMBINATIONS)
+
+    string(REPLACE ";" "," FA_QUANTS_DEFINE "${FA_QUANTS}")
+    add_compile_definitions(GGML_CUDA_FA_QUANTS="${FA_QUANTS_DEFINE}")
+    foreach (TYPE_V IN LISTS FA_TYPES)
+        foreach (TYPE_K IN LISTS FA_TYPES)
+            if ("${TYPE_K}-${TYPE_V}" IN_LIST FA_COMBINATIONS)
+                set(COMPILED 1)
+            else()
+                set(COMPILED 0)
+            endif()
+            string(TOUPPER "GGML_CUDA_FA_${TYPE_K}_${TYPE_V}" COMBINATION_DEF)
+            add_compile_definitions(${COMBINATION_DEF}=${COMPILED})
+        endforeach()
+    endforeach()
+
+    message(STATUS "FlashAttention K-V type combinations: ${FA_COMBINATIONS}")
+
+    set(SRCS "")
+    foreach (COMBINATION IN LISTS FA_COMBINATIONS)
+        set(SRC "${DIR}/template-instances/fattn-vec-instance-${COMBINATION}.cu")
+        if (NOT EXISTS "${SRC}")
+            message(FATAL_ERROR "FlashAttention template instance \"${SRC}\" does not exist")
+        endif()
+        list(APPEND SRCS "${SRC}")
+    endforeach()
+
+    set(${OUT_SRCS} ${SRCS} PARENT_SCOPE)
+endfunction()
@@ -1705,6 +1705,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                    ggml_tensor * ids_tensor = node->src[2];
                    ggml_backend_t ids_backend = split_backend;

+                    if (ggml_nelements(ids_tensor) == 0) {
+                        continue;
+                    }
+
                    // if the ids tensor is also an input of the split, it may not have been copied yet to the split backend
                    // in that case, we use the original ids tensor
                    for (int i = input_id + 1; i < split->n_inputs; i++) {
@@ -520,7 +520,9 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
    elseif (GGML_SYSTEM_ARCH STREQUAL "s390x")
        message(STATUS "s390x detected")
        list(APPEND GGML_CPU_SOURCES
-            ggml-cpu/arch/s390/quants.c)
+            ggml-cpu/arch/s390/quants.c
+            ggml-cpu/arch/s390/repack.cpp
+            )

        # for native compilation
        if (GGML_NATIVE)
@@ -673,6 +675,12 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
    target_compile_options(${GGML_CPU_NAME} PRIVATE ${ARCH_FLAGS})
    target_compile_definitions(${GGML_CPU_NAME} PRIVATE ${ARCH_DEFINITIONS})

+    if (CMAKE_C_COMPILER_ID STREQUAL "GNU" AND NOT GGML_SYSTEM_ARCH STREQUAL "x86")
+        message(STATUS "Skipping PCH for ${GGML_CPU_NAME}: GCC PCH is only enabled for x86 (arch: ${GGML_SYSTEM_ARCH})")
+    else()
+        target_precompile_headers(${GGML_CPU_NAME} PRIVATE ggml-impl.h)
+    endif()
+
    if (EMSCRIPTEN)
        set_target_properties(${GGML_CPU_NAME} PROPERTIES COMPILE_FLAGS "-msimd128")
    endif()
@@ -247,7 +247,6 @@
 // quants.c
 #define quantize_row_q8_K_generic quantize_row_q8_K
 #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
-#define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
 #define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
 #define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
 #define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
@@ -260,11 +259,9 @@
 #define ggml_vec_dot_iq1_s_q8_K_generic ggml_vec_dot_iq1_s_q8_K
 #define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K
 // repack.cpp
-#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
 #define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8
-#define ggml_gemv_q4_0_4x4_q8_0_generic ggml_gemv_q4_0_4x4_q8_0
 #define ggml_gemv_q4_0_4x8_q8_0_generic ggml_gemv_q4_0_4x8_q8_0
 #define ggml_gemv_q4_0_8x8_q8_0_generic ggml_gemv_q4_0_8x8_q8_0
 #define ggml_gemv_q2_K_8x8_q8_K_generic ggml_gemv_q2_K_8x8_q8_K
@@ -280,7 +277,6 @@
 #define ggml_gemv_mxfp4_8x8_q8_0_generic ggml_gemv_mxfp4_8x8_q8_0
 #define ggml_gemv_q8_0_4x4_q8_0_generic ggml_gemv_q8_0_4x4_q8_0
 #define ggml_gemv_q8_0_4x8_q8_0_generic ggml_gemv_q8_0_4x8_q8_0
-#define ggml_gemm_q4_0_4x4_q8_0_generic ggml_gemm_q4_0_4x4_q8_0
 #define ggml_gemm_q4_0_4x8_q8_0_generic ggml_gemm_q4_0_4x8_q8_0
 #define ggml_gemm_q4_0_8x8_q8_0_generic ggml_gemm_q4_0_8x8_q8_0
 #define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
@@ -146,6 +146,74 @@ void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, i

 //===================================== Dot products =================================

+void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK1_0;  // 128
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q1_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+#if defined(__VXE__) || defined(__VXE2__)
+    float32x4_t v_sumf = vec_splats(0.0f);
+
+    const uint8x16_t v_zero = vec_splats((uint8_t)0x00);  // zero
+    const uint8x16_t v_bias = vec_splats((uint8_t)0x80);  // bias from signed to unsigned
+                                                          // v ^ 0x80 == v + 128
+
+    const uint8x16_t v_idx = (const uint8x16_t){ 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1 };
+    const uint8x16_t v_bit = (const uint8x16_t){ 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128 };
+
+    for (int i = 0; i < nb; ++i) {
+        const uint8x16_t  v_x  = vec_xl(0, (const uint8_t *)x[i].qs);
+        const float32x4_t v_xd = vec_splats(GGML_CPU_FP16_TO_FP32(x[i].d));
+
+        for (int k = 0; k < 4; ++k) {
+            // sub-block k holds elements 32k .. 32k+31
+            const block_q8_0 * GGML_RESTRICT yb = &y[i*4 + k];
+            const float32x4_t v_yd = vec_splats(GGML_CPU_FP16_TO_FP32(yb->d));
+
+            const uint8x16_t v_xrl = vec_perm(v_x, v_x, vec_add(v_idx, vec_splats((uint8_t)(k*4 + 0))));
+            const uint8x16_t v_xrh = vec_perm(v_x, v_x, vec_add(v_idx, vec_splats((uint8_t)(k*4 + 2))));
+
+            // isolate each lane's bit, then set all ones where that bit is clear, the -d case
+            const int8x16_t v_ml = (int8x16_t)vec_cmpeq(vec_and(v_xrl, v_bit), v_zero);
+            const int8x16_t v_mh = (int8x16_t)vec_cmpeq(vec_and(v_xrh, v_bit), v_zero);
+
+            const int8x16_t v_yl = vec_xl(0,       (const int8_t *)yb->qs);
+            const int8x16_t v_yh = vec_xl(QK8_0/2, (const int8_t *)yb->qs);
+
+            // weights are only +1 or -1, so negate y
+            const int8x16_t v_ysl = vec_sub(vec_xor(v_yl, v_ml), v_ml);
+            const int8x16_t v_ysh = vec_sub(vec_xor(v_yh, v_mh), v_mh);
+
+            // bias to unsigned, then vec_sum4 adds each group of 4 bytes into one word
+            const uint32x4_t v_p = vec_add(vec_sum4(vec_xor((uint8x16_t)v_ysl, v_bias), v_zero),
+                                           vec_sum4(vec_xor((uint8x16_t)v_ysh, v_bias), v_zero));
+
+            // each word summed 8 biased bytes, so take back 8 * 128
+            const int32x4_t v_xy = vec_sub((int32x4_t)v_p, vec_splats((int32_t)1024));
+
+            // apply both block scales and add into the running total
+            v_sumf = vec_madd(vec_float(v_xy), vec_mul(v_xd, v_yd), v_sumf);
+        }
+    }
+
+    *s = vec_hsum_f32x4(v_sumf);
+#else
+    UNUSED(nb);
+    UNUSED(x);
+    UNUSED(y);
+    ggml_vec_dot_q1_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
 void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
    const int qk = QK8_0;
    const int nb = n / qk;
@@ -0,0 +1,223 @@
+#define GGML_COMMON_IMPL_CPP
+#define GGML_COMMON_DECL_CPP
+#include "ggml-common.h"
+#include "ggml-backend-impl.h"
+
+#include "ggml-impl.h"
+#include "ggml-cpu.h"
+#include "ggml-cpu-impl.h"
+#include "simd-mappings.h"
+#include "traits.h"
+
+#include <cmath>
+#include <cstring>
+#include <cassert>
+
+#define GGML_CPU_CLANG_WORKAROUND
+#include "../../repack.h"
+
+#define UNUSED GGML_UNUSED
+
+void ggml_quantize_mat_q8_0_4x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
+    assert(QK8_0 == 32);
+    assert(k % QK8_0 == 0);
+    const int nb = k / QK8_0;
+
+    block_q8_0x4 * GGML_RESTRICT y = (block_q8_0x4 *) vy;
+
+#if defined(__VXE__) || defined(__VXE2__)
+    float32x4_t v_src[4][8];
+    float       id[4];
+
+    for (int i = 0; i < nb; i++) {
+        float32x4_t v_asrc[8];
+        float32x4_t v_amax[8];
+
+        for (int row_iter = 0; row_iter < 4; row_iter++) {
+            for (int j = 0; j < 8; j++) v_src[row_iter][j] = vec_xl(0, x + row_iter * k + i * 32 + 4 * j);
+            for (int j = 0; j < 8; j++) v_asrc[j] = vec_abs(v_src[row_iter][j]);
+
+            for (int j = 0; j < 4; j++) v_amax[2 * j] = vec_max(v_asrc[2 * j], v_asrc[2 * j + 1]);
+            for (int j = 0; j < 2; j++) v_amax[4 * j] = vec_max(v_amax[4 * j], v_amax[4 * j + 2]);
+            for (int j = 0; j < 1; j++) v_amax[8 * j] = vec_max(v_amax[8 * j], v_amax[8 * j + 4]);
+
+            const float amax = MAX(MAX(vec_extract(v_amax[0], 0), vec_extract(v_amax[0], 1)),
+                                   MAX(vec_extract(v_amax[0], 2), vec_extract(v_amax[0], 3)));
+
+            const float d = amax / ((1 << 7) - 1);
+            id[row_iter] = d ? 1.0f / d : 0.0f;
+
+            y[i].d[row_iter] = GGML_CPU_FP32_TO_FP16(d);
+        }
+
+        for (int j = 0; j < 8; j++) {
+            /* Uses non-default rounding for vec_signed or vec_round */
+            const int32x4_t v_qs0 = vec_signed(__builtin_s390_vfisb(vec_mul(v_src[0][j], id[0]), 4, 1));
+            const int32x4_t v_qs1 = vec_signed(__builtin_s390_vfisb(vec_mul(v_src[1][j], id[1]), 4, 1));
+            const int32x4_t v_qs2 = vec_signed(__builtin_s390_vfisb(vec_mul(v_src[2][j], id[2]), 4, 1));
+            const int32x4_t v_qs3 = vec_signed(__builtin_s390_vfisb(vec_mul(v_src[3][j], id[3]), 4, 1));
+
+            const int16x8_t v_qs01 = vec_packs(v_qs0, v_qs1);
+            const int16x8_t v_qs23 = vec_packs(v_qs2, v_qs3);
+
+            vec_xst(vec_packs(v_qs01, v_qs23), 0, y[i].qs + 16 * j);
+        }
+    }
+#else
+    UNUSED(nb);
+    UNUSED(y);
+    ggml_quantize_mat_q8_0_4x4_generic(x, vy, k);
+#endif
+}
+
+static inline int16x8_t vxe_dot_acc(const int8x16_t v_x, const int8x16_t v_y, const int16x8_t v_acc) {
+    return vec_meadd(v_x, v_y, vec_moadd(v_x, v_y, v_acc));
+}
+
+static inline int8x16_t vxe_splat_granule(const int8_t * qs) {
+    uint32_t g;
+    memcpy(&g, qs, sizeof(g));
+    return (int8x16_t)vec_splats(g);
+}
+
+static inline int32x4_t vxe_fold(const int16x8_t v_sumi) {
+    const int16x8_t v_ones = vec_splats((int16_t)1);
+    return vec_add(vec_mule(v_sumi, v_ones), vec_mulo(v_sumi, v_ones));
+}
+
+void ggml_gemv_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) {
+    const int qk = QK8_0;
+    const int nb = n / qk;
+    const int ncols_interleaved = 4;
+
+    assert(nr == 1);
+    assert(n % qk == 0);
+    assert(nc % ncols_interleaved == 0);
+
+    UNUSED(bs);
+    UNUSED(nr);
+
+#if defined(__VXE__) || defined(__VXE2__)
+    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_q4_0x4 * b_ptr = (const block_q4_0x4 *) vx + (x * nb);
+
+        float32x4_t v_sumf = vec_splats(0.0f);
+
+        for (int l = 0; l < nb; l++) {
+            const int8_t * x_qs = b_ptr[l].qs;
+
+            const int8x16_t v_x0 = vec_xl( 0, x_qs);
+            const int8x16_t v_x1 = vec_xl(16, x_qs);
+            const int8x16_t v_x2 = vec_xl(32, x_qs);
+            const int8x16_t v_x3 = vec_xl(48, x_qs);
+
+            const int8x16_t v_x0l = vec_sra(vec_sl(v_x0, 4), 4);
+            const int8x16_t v_x1l = vec_sra(vec_sl(v_x1, 4), 4);
+            const int8x16_t v_x2l = vec_sra(vec_sl(v_x2, 4), 4);
+            const int8x16_t v_x3l = vec_sra(vec_sl(v_x3, 4), 4);
+
+            const int8x16_t v_x0h = vec_sra(v_x0, 4);
+            const int8x16_t v_x1h = vec_sra(v_x1, 4);
+            const int8x16_t v_x2h = vec_sra(v_x2, 4);
+            const int8x16_t v_x3h = vec_sra(v_x3, 4);
+
+            const int8_t * y_lo = a_ptr[l].qs;
+            const int8_t * y_hi = y_lo + qk / 2;
+
+            int16x8_t v_sumi = vec_splats((int16_t)0);
+
+            v_sumi = vxe_dot_acc(v_x0l, vxe_splat_granule(y_lo +  0), v_sumi);
+            v_sumi = vxe_dot_acc(v_x1l, vxe_splat_granule(y_lo +  4), v_sumi);
+            v_sumi = vxe_dot_acc(v_x2l, vxe_splat_granule(y_lo +  8), v_sumi);
+            v_sumi = vxe_dot_acc(v_x3l, vxe_splat_granule(y_lo + 12), v_sumi);
+
+            v_sumi = vxe_dot_acc(v_x0h, vxe_splat_granule(y_hi +  0), v_sumi);
+            v_sumi = vxe_dot_acc(v_x1h, vxe_splat_granule(y_hi +  4), v_sumi);
+            v_sumi = vxe_dot_acc(v_x2h, vxe_splat_granule(y_hi +  8), v_sumi);
+            v_sumi = vxe_dot_acc(v_x3h, vxe_splat_granule(y_hi + 12), v_sumi);
+
+            const float32x4_t v_yd = vec_splats(GGML_CPU_FP16_TO_FP32(a_ptr[l].d));
+            const float32x4_t v_xd = __lzs_f16cx4_load(b_ptr[l].d);
+            const float32x4_t v_d  = vec_mul(v_yd, v_xd);
+
+            v_sumf = vec_madd(vec_float(vxe_fold(v_sumi)), v_d, v_sumf);
+        }
+
+        vec_xst(v_sumf, 0, res_ptr + x * ncols_interleaved);
+    }
+#else
+    UNUSED(nb);
+    UNUSED(ncols_interleaved);
+    ggml_gemv_q4_0_4x4_q8_0_generic(n, s, bs, vx, vy, nr, nc);
+#endif
+}
+
+void ggml_gemm_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) {
+    const int qk = QK8_0;
+    const int nb = n / qk;
+    const int ncols_interleaved = 4;
+
+    assert(nr % 4 == 0);
+    assert(n % qk == 0);
+    assert(nc % ncols_interleaved == 0);
+
+#if defined(__VXE__) || defined(__VXE2__)
+    for (int y = 0; y < nr / 4; y++) {
+        const block_q8_0x4 * a_ptr = (const block_q8_0x4 *) vy + (y * nb);
+
+        for (int x = 0; x < nc / ncols_interleaved; x++) {
+            const block_q4_0x4 * b_ptr = (const block_q4_0x4 *) vx + (x * nb);
+
+            float32x4_t v_sumf[4];
+            for (int m = 0; m < 4; m++) {
+                v_sumf[m] = vec_splats(0.0f);
+            }
+
+            for (int l = 0; l < nb; l++) {
+                int16x8_t v_sumi0 = vec_splats((int16_t)0);
+                int16x8_t v_sumi1 = vec_splats((int16_t)0);
+                int16x8_t v_sumi2 = vec_splats((int16_t)0);
+                int16x8_t v_sumi3 = vec_splats((int16_t)0);
+
+                for (int k = 0; k < 4; k++) {
+                    const int8x16_t v_x  = vec_xl(0, b_ptr[l].qs + 16 * k);
+                    const int8x16_t v_xl = vec_sra(vec_sl(v_x, 4), 4);
+                    const int8x16_t v_xh = vec_sra(v_x, 4);
+
+                    const int8_t * y_lo = a_ptr[l].qs + 16 * k;
+                    const int8_t * y_hi = y_lo + qk / 2 * 4;
+
+                    v_sumi0 = vxe_dot_acc(v_xl, vxe_splat_granule(y_lo +  0), v_sumi0);
+                    v_sumi1 = vxe_dot_acc(v_xl, vxe_splat_granule(y_lo +  4), v_sumi1);
+                    v_sumi2 = vxe_dot_acc(v_xl, vxe_splat_granule(y_lo +  8), v_sumi2);
+                    v_sumi3 = vxe_dot_acc(v_xl, vxe_splat_granule(y_lo + 12), v_sumi3);
+
+                    v_sumi0 = vxe_dot_acc(v_xh, vxe_splat_granule(y_hi +  0), v_sumi0);
+                    v_sumi1 = vxe_dot_acc(v_xh, vxe_splat_granule(y_hi +  4), v_sumi1);
+                    v_sumi2 = vxe_dot_acc(v_xh, vxe_splat_granule(y_hi +  8), v_sumi2);
+                    v_sumi3 = vxe_dot_acc(v_xh, vxe_splat_granule(y_hi + 12), v_sumi3);
+                }
+
+                const float32x4_t v_yd = __lzs_f16cx4_load(a_ptr[l].d);
+                const float32x4_t v_xd = __lzs_f16cx4_load(b_ptr[l].d);
+
+                v_sumf[0] = vec_madd(vec_float(vxe_fold(v_sumi0)), vec_mul(v_xd, vec_splat(v_yd, 0)), v_sumf[0]);
+                v_sumf[1] = vec_madd(vec_float(vxe_fold(v_sumi1)), vec_mul(v_xd, vec_splat(v_yd, 1)), v_sumf[1]);
+                v_sumf[2] = vec_madd(vec_float(vxe_fold(v_sumi2)), vec_mul(v_xd, vec_splat(v_yd, 2)), v_sumf[2]);
+                v_sumf[3] = vec_madd(vec_float(vxe_fold(v_sumi3)), vec_mul(v_xd, vec_splat(v_yd, 3)), v_sumf[3]);
+            }
+
+            for (int m = 0; m < 4; m++) {
+                vec_xst(v_sumf[m], 0, s + (y * 4 + m) * bs + x * ncols_interleaved);
+            }
+        }
+    }
+#else
+    UNUSED(nb);
+    UNUSED(ncols_interleaved);
+    ggml_gemm_q4_0_4x4_q8_0_generic(n, s, bs, vx, vy, nr, nc);
+#endif
+}
@@ -18,7 +18,15 @@
 #endif
 #endif

+// -Winterference-size was introduced in GCC 12
+#if defined(__cplusplus) && defined(__GNUC__) && !defined(__clang__) && __GNUC__ >= 12
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Winterference-size"
+#endif
 static const size_t CACHE_LINE_SIZE_F32 = CACHE_LINE_SIZE/sizeof(float);
+#if defined(__cplusplus) && defined(__GNUC__) && !defined(__clang__) && __GNUC__ >= 12
+#pragma GCC diagnostic pop
+#endif

 // Work buffer size for im2col operations in CONV2D
 #define GGML_IM2COL_WORK_SIZE (16 * 1024 * 1024)
@@ -4586,6 +4586,11 @@ static const ggml::cpu::tensor_traits * ggml_repack_get_optimal_repack_type(cons
                return &q4_0_4x4_q8_0;
            }
        }
+        if (ggml_cpu_has_vxe()) {
+            if (cur->ne[1] % 4 == 0) {
+                return &q4_0_4x4_q8_0;
+            }
+        }
        if (ggml_cpu_has_riscv_v()) {
            #if defined __riscv_zvfh
            switch (__riscv_vlenb() * 8) {
@@ -112,17 +112,8 @@ if (CUDAToolkit_FOUND)
    file(GLOB   SRCS "template-instances/mmf*.cu")
    list(APPEND GGML_SOURCES_CUDA ${SRCS})

-    if (GGML_CUDA_FA_ALL_QUANTS)
-        file(GLOB   SRCS "template-instances/fattn-vec*.cu")
-        list(APPEND GGML_SOURCES_CUDA ${SRCS})
-        add_compile_definitions(GGML_CUDA_FA_ALL_QUANTS)
-    else()
-        list(APPEND GGML_SOURCES_CUDA
-            template-instances/fattn-vec-instance-f16-f16.cu
-            template-instances/fattn-vec-instance-q4_0-q4_0.cu
-            template-instances/fattn-vec-instance-q8_0-q8_0.cu
-            template-instances/fattn-vec-instance-bf16-bf16.cu)
-    endif()
+    ggml_cuda_fattn_vec_instances(${CMAKE_CURRENT_SOURCE_DIR} SRCS)
+    list(APPEND GGML_SOURCES_CUDA ${SRCS})

    ggml_add_backend_library(ggml-cuda
                             ${GGML_HEADERS_CUDA}
@@ -1133,12 +1133,21 @@ void launch_fattn(

    dim3 blocks_num;
    if (stream_k) {
-        // For short contexts it can be faster to have the SMs work on whole tiles because this lets us skip the fixup.
-        const int max_blocks = max_blocks_per_sm*nsm;
-        const int tiles_nwaves = (ntiles_dst + max_blocks - 1) / max_blocks;
-        const int tiles_efficiency_percent = 100 * ntiles_dst / (max_blocks*tiles_nwaves);
+        auto should_use_stream_k = [](const int cc, const int ntiles_dst, const int max_blocks, const int DKQ) {
+            const int tiles_nwaves             = (ntiles_dst + max_blocks - 1) / max_blocks;
+            const int tiles_efficiency_percent = 100 * ntiles_dst / (max_blocks*tiles_nwaves);

-        const bool use_stream_k = cc >= GGML_CUDA_CC_ADA_LOVELACE || amd_wmma_available(cc) || tiles_efficiency_percent < 75;
+            if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_ADA_LOVELACE) {
+                return true;
+            }
+            if (amd_wmma_available(cc) && DKQ == 64) {
+                return true; // TODO better configuration
+            }
+            return tiles_efficiency_percent < 75;
+        };
+
+        const int  max_blocks   = max_blocks_per_sm*nsm;
+        const bool use_stream_k = should_use_stream_k(cc, ntiles_dst, max_blocks, Q->ne[0]);

        blocks_num.x = ntiles_dst;
        blocks_num.y = 1;
@@ -158,8 +158,8 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co

    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8,  64, 2,  32, 128, 128, 128, 1, true);
    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16,  64, 2,  32, 128, 128, 128, 1, true);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 128, 2,  64, 128, 128,  64, 1, true);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2,  64, 128, 128,  64, 1, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 256, 2,  64, 128, 128,  64, 1, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 256, 2,  64, 128, 128,  64, 1, true);

    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 128, 2,  32, 160, 128, 128, 1, true);
    GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 64, 128, 2,  32, 160, 128, 128, 1, true);
@@ -1826,7 +1826,7 @@ static __global__ void flash_attn_ext_f16(
 #endif // __CUDA_ARCH__ == GGML_CUDA_CC_TURING

 #if defined(AMD_WMMA_AVAILABLE)
-    if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 128) {
+    if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 256) {
        NO_DEVICE_CODE;
        return;
    }
@@ -221,6 +221,24 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con
        }
    }

+    // On RDNA it is preferable to minimize wasted compute vs. duplicate I/O for the mask.
+    if (amd_wmma_available(cc)) {
+        if (use_gqa_opt && gqa_ratio % 8 == 0) {
+            ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 8>(ctx, dst);
+            return;
+        }
+
+        if (use_gqa_opt && gqa_ratio % 4 == 0) {
+            ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 4>(ctx, dst);
+            return;
+        }
+
+        if (use_gqa_opt && gqa_ratio % 2 == 0) {
+            ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 2>(ctx, dst);
+            return;
+        }
+    }
+
    if (use_gqa_opt && gqa_ratio > 4) {
        ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 8>(ctx, dst);
        return;
@@ -374,90 +392,101 @@ static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, gg
    }
 }

-#define FATTN_VEC_CASE(D, type_K, type_V)                                                                        \
-    {                                                                                                            \
-        const bool type_K_okay = K->type == (type_K) || (K->type == GGML_TYPE_F32 && (type_K) == GGML_TYPE_F16); \
-        const bool type_V_okay = V->type == (type_V) || (V->type == GGML_TYPE_F32 && (type_V) == GGML_TYPE_F16); \
-        if (Q->ne[0] == (D) && type_K_okay && type_V_okay) {                                                     \
-            ggml_cuda_flash_attn_ext_vec_case<D, type_K, type_V>(ctx, dst);                                      \
-            return;                                                                                              \
-        }                                                                                                        \
-    }                                                                                                            \
+#define FATTN_VEC_CASE(D, type_K_case, type_V_case)                                                                                \
+    if constexpr (GGML_CUDA_FA_##type_K_case##_##type_V_case) {                                                                    \
+        const bool type_K_okay = type_K == GGML_TYPE_##type_K_case || (type_K == GGML_TYPE_F32 && GGML_TYPE_##type_K_case == GGML_TYPE_F16); \
+        const bool type_V_okay = type_V == GGML_TYPE_##type_V_case || (type_V == GGML_TYPE_F32 && GGML_TYPE_##type_V_case == GGML_TYPE_F16); \
+        if (head_size == (D) && type_K_okay && type_V_okay) {                                                                      \
+            return ggml_cuda_flash_attn_ext_vec_case<D, GGML_TYPE_##type_K_case, GGML_TYPE_##type_V_case>;                         \
+        }                                                                                                                          \
+    }                                                                                                                              \

-#define FATTN_VEC_CASES_ALL_D(type_K, type_V) \
-    FATTN_VEC_CASE( 64, type_K, type_V)       \
-    FATTN_VEC_CASE(128, type_K, type_V)       \
-    FATTN_VEC_CASE(256, type_K, type_V)       \
+#define FATTN_VEC_CASES_ALL_D(type_K_case, type_V_case) \
+    FATTN_VEC_CASE( 64, type_K_case, type_V_case)       \
+    FATTN_VEC_CASE(128, type_K_case, type_V_case)       \
+    FATTN_VEC_CASE(256, type_K_case, type_V_case)       \
+
+typedef void (* fattn_vec_case_t)(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+// Vector kernel for the given head size and K/V types, nullptr if its template instance was not compiled:
+static fattn_vec_case_t ggml_cuda_get_fattn_vec_case(const int64_t head_size, const ggml_type type_K, const ggml_type type_V) {
+    FATTN_VEC_CASES_ALL_D(F16,  F16)
+    FATTN_VEC_CASES_ALL_D(Q4_0, F16)
+    FATTN_VEC_CASES_ALL_D(Q4_1, F16)
+    FATTN_VEC_CASES_ALL_D(Q5_0, F16)
+    FATTN_VEC_CASES_ALL_D(Q5_1, F16)
+    FATTN_VEC_CASES_ALL_D(Q8_0, F16)
+    FATTN_VEC_CASES_ALL_D(BF16, F16)
+
+    FATTN_VEC_CASES_ALL_D(F16,  Q4_0)
+    FATTN_VEC_CASES_ALL_D(Q4_0, Q4_0)
+    FATTN_VEC_CASES_ALL_D(Q4_1, Q4_0)
+    FATTN_VEC_CASES_ALL_D(Q5_0, Q4_0)
+    FATTN_VEC_CASES_ALL_D(Q5_1, Q4_0)
+    FATTN_VEC_CASES_ALL_D(Q8_0, Q4_0)
+    FATTN_VEC_CASES_ALL_D(BF16, Q4_0)
+
+    FATTN_VEC_CASES_ALL_D(F16,  Q4_1)
+    FATTN_VEC_CASES_ALL_D(Q4_0, Q4_1)
+    FATTN_VEC_CASES_ALL_D(Q4_1, Q4_1)
+    FATTN_VEC_CASES_ALL_D(Q5_0, Q4_1)
+    FATTN_VEC_CASES_ALL_D(Q5_1, Q4_1)
+    FATTN_VEC_CASES_ALL_D(Q8_0, Q4_1)
+    FATTN_VEC_CASES_ALL_D(BF16, Q4_1)
+
+    FATTN_VEC_CASES_ALL_D(F16,  Q5_0)
+    FATTN_VEC_CASES_ALL_D(Q4_0, Q5_0)
+    FATTN_VEC_CASES_ALL_D(Q4_1, Q5_0)
+    FATTN_VEC_CASES_ALL_D(Q5_0, Q5_0)
+    FATTN_VEC_CASES_ALL_D(Q5_1, Q5_0)
+    FATTN_VEC_CASES_ALL_D(Q8_0, Q5_0)
+    FATTN_VEC_CASES_ALL_D(BF16, Q5_0)
+
+    FATTN_VEC_CASES_ALL_D(F16,  Q5_1)
+    FATTN_VEC_CASES_ALL_D(Q4_0, Q5_1)
+    FATTN_VEC_CASES_ALL_D(Q4_1, Q5_1)
+    FATTN_VEC_CASES_ALL_D(Q5_0, Q5_1)
+    FATTN_VEC_CASES_ALL_D(Q5_1, Q5_1)
+    FATTN_VEC_CASES_ALL_D(Q8_0, Q5_1)
+    FATTN_VEC_CASES_ALL_D(BF16, Q5_1)
+
+    FATTN_VEC_CASES_ALL_D(F16,  Q8_0)
+    FATTN_VEC_CASES_ALL_D(Q4_0, Q8_0)
+    FATTN_VEC_CASES_ALL_D(Q4_1, Q8_0)
+    FATTN_VEC_CASES_ALL_D(Q5_0, Q8_0)
+    FATTN_VEC_CASES_ALL_D(Q5_1, Q8_0)
+    FATTN_VEC_CASES_ALL_D(Q8_0, Q8_0)
+    FATTN_VEC_CASES_ALL_D(BF16, Q8_0)
+
+    FATTN_VEC_CASES_ALL_D(F16,  BF16)
+    FATTN_VEC_CASES_ALL_D(Q4_0, BF16)
+    FATTN_VEC_CASES_ALL_D(Q4_1, BF16)
+    FATTN_VEC_CASES_ALL_D(Q5_0, BF16)
+    FATTN_VEC_CASES_ALL_D(Q5_1, BF16)
+    FATTN_VEC_CASES_ALL_D(Q8_0, BF16)
+    FATTN_VEC_CASES_ALL_D(BF16, BF16)
+
+    return nullptr;
+}

 static void ggml_cuda_flash_attn_ext_vec(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
-    ggml_tensor * Q = dst->src[0];
-    ggml_tensor * K = dst->src[1];
-    ggml_tensor * V = dst->src[2];
+    const ggml_tensor * Q = dst->src[0];
+    const ggml_tensor * K = dst->src[1];
+    const ggml_tensor * V = dst->src[2];

-#ifdef GGML_CUDA_FA_ALL_QUANTS
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16,  GGML_TYPE_F16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_F16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_1, GGML_TYPE_F16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_0, GGML_TYPE_F16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_F16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_F16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_F16)
-
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16,  GGML_TYPE_Q4_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_Q4_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_1, GGML_TYPE_Q4_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_0, GGML_TYPE_Q4_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_Q4_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_Q4_0)
-
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16,  GGML_TYPE_Q4_1)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_Q4_1)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_1, GGML_TYPE_Q4_1)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_0, GGML_TYPE_Q4_1)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_Q4_1)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q4_1)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_Q4_1)
-
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16,  GGML_TYPE_Q5_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_Q5_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_1, GGML_TYPE_Q5_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_0, GGML_TYPE_Q5_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_Q5_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q5_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_Q5_0)
-
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16,  GGML_TYPE_Q5_1)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_Q5_1)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_1, GGML_TYPE_Q5_1)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_0, GGML_TYPE_Q5_1)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_Q5_1)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q5_1)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_Q5_1)
-
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16,  GGML_TYPE_Q8_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_Q8_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_1, GGML_TYPE_Q8_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_0, GGML_TYPE_Q8_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_Q8_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_Q8_0)
-
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16,  GGML_TYPE_BF16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_BF16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_1, GGML_TYPE_BF16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_0, GGML_TYPE_BF16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q5_1, GGML_TYPE_BF16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_BF16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_BF16)
-#else
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_F16,  GGML_TYPE_F16)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q4_0, GGML_TYPE_Q4_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)
-    FATTN_VEC_CASES_ALL_D(GGML_TYPE_BF16, GGML_TYPE_BF16)
-#endif // GGML_CUDA_FA_ALL_QUANTS
-
-    GGML_ABORT("fatal error");
+    fattn_vec_case_t vec_case = ggml_cuda_get_fattn_vec_case(Q->ne[0], K->type, V->type);
+    if (vec_case == nullptr) {
+        static bool warned = false;
+        if (!warned) {
+            GGML_LOG_WARN("%s: no FlashAttention vector kernel compiled for K/V types %s-%s, converting K and V to f16 instead (slow). "
+                "Add \"%s-%s\" to GGML_CUDA_FA_QUANTS to compile it.\n",
+                __func__, ggml_type_name(K->type), ggml_type_name(V->type), ggml_type_name(K->type), ggml_type_name(V->type));
+            warned = true;
+        }
+        vec_case = ggml_cuda_get_fattn_vec_case(Q->ne[0], GGML_TYPE_F16, GGML_TYPE_F16);
+    }
+    GGML_ASSERT(vec_case != nullptr);
+    vec_case(ctx, dst);
 }

 // Best FlashAttention kernel for a specific GPU:
@@ -468,20 +497,17 @@ enum best_fattn_kernel {
    BEST_FATTN_KERNEL_MMA_F16 = 400,
 };

-static bool ggml_cuda_fattn_kv_type_supported(ggml_type type) {
+// K/V types for which there is a vector kernel template instance, other kernels convert these to f16:
+static bool ggml_cuda_fattn_kv_type_supported(const ggml_type type) {
    switch (type) {
        case GGML_TYPE_F32:
        case GGML_TYPE_F16:
-            return true;
+        case GGML_TYPE_BF16:
+        case GGML_TYPE_Q4_0:
        case GGML_TYPE_Q4_1:
        case GGML_TYPE_Q5_0:
        case GGML_TYPE_Q5_1:
-#ifndef GGML_CUDA_FA_ALL_QUANTS
-            return false;
-#endif // GGML_CUDA_FA_ALL_QUANTS
-        case GGML_TYPE_Q4_0:
        case GGML_TYPE_Q8_0:
-        case GGML_TYPE_BF16:
            return true;
        default:
            return false;
@@ -572,12 +598,6 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
            return BEST_FATTN_KERNEL_NONE;
    }

-#ifndef GGML_CUDA_FA_ALL_QUANTS
-    if (K->type != V->type) {
-        return BEST_FATTN_KERNEL_NONE;
-    }
-#endif // GGML_CUDA_FA_ALL_QUANTS
-
    if (!ggml_cuda_fattn_kv_type_supported(K->type) || !ggml_cuda_fattn_kv_type_supported(V->type)) {
        return BEST_FATTN_KERNEL_NONE;
    }
@@ -644,8 +664,9 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
        }
    }

-    // AMD WMMA is always faster than the tile kernel if the full tile width of 16 can be utilized.
-    if ((amd_wmma_available(cc) && gqa_opt_applies && Q->ne[0] <= 128) && Q->ne[0] != 40 && Q->ne[0] != 72 && Q->ne[1] * gqa_ratio_eff > 8) {
+    // AMD WMMA is faster than the tile kernel if the wide tiles with high arithmetic intensity can be utilized.
+    if ((amd_wmma_available(cc) && gqa_opt_applies && Q->ne[0] <= 256) && Q->ne[0] != 40 && Q->ne[0] != 72 &&
+            Q->ne[1] * gqa_ratio_eff > (Q->ne[0] <= 128 ? 8 : 16)) {
        return BEST_FATTN_KERNEL_MMA_F16;
    }

@@ -669,6 +690,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
 size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * dst) {
    GGML_ASSERT(dst->op == GGML_OP_FLASH_ATTN_EXT);

+    const ggml_tensor * Q = dst->src[0];
    const ggml_tensor * K = dst->src[1];
    const ggml_tensor * V = dst->src[2];

@@ -686,10 +708,11 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d
            need_f16_K = true;
            need_f16_V = true;
            break;
-        case BEST_FATTN_KERNEL_VEC:
-            need_f16_K = K->type == GGML_TYPE_F32;
-            need_f16_V = V->type == GGML_TYPE_F32;
-            break;
+        case BEST_FATTN_KERNEL_VEC: {
+            const bool f16_fallback = ggml_cuda_get_fattn_vec_case(Q->ne[0], K->type, V->type) == nullptr;
+            need_f16_K = K->type == GGML_TYPE_F32 || f16_fallback;
+            need_f16_V = V->type == GGML_TYPE_F32 || f16_fallback;
+        } break;
        case BEST_FATTN_KERNEL_NONE:
            break;
    }
@@ -5640,8 +5640,8 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t
        features.push_back({ "USE_GRAPHS", "1" });
    #endif

-    #ifdef GGML_CUDA_FA_ALL_QUANTS
-        features.push_back({ "FA_ALL_QUANTS", "1" });
+    #ifdef GGML_CUDA_FA_QUANTS
+        features.push_back({ "FA_QUANTS", GGML_CUDA_FA_QUANTS });
    #endif

    {
@@ -171,7 +171,7 @@ void ggml_cuda_mul_mat_q(
            ne00, ne01, ne1, s01, ne11, s1,
            ne02, ne12, s02, s12, s2,
            ne03, ne13, s03, s13, s3,
-            ne1};
+            ne1, ne1};
        ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
        return;
    }
@@ -244,6 +244,13 @@ void ggml_cuda_mul_mat_q(
                                         ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
    const int64_t s13 = ne12*s12;

+    // Each expert only sees ne12*n_expert_used/ne02 tokens on average.
+    // On RDNA3 and RDNA4 it is faster to pick the tile size against this value instead of ne12.
+    int64_t ncols_opt = ne12;
+    if (GGML_CUDA_CC_IS_RDNA3_0(cc) || GGML_CUDA_CC_IS_RDNA4(cc)) {
+        ncols_opt = (ne12*n_expert_used + ne02 - 1) / ne02;
+    }
+
    // Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid.
    const mmq_args args = {
        src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), dst_d,
@@ -251,7 +258,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};
+        ne12, ncols_opt};

    ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
 }
@@ -1376,6 +1376,7 @@ struct mmq_args {
    int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst;
    int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
    int64_t ncols_max;
+    int64_t ncols_opt; // value to optimize the tile size against, launch grid still uses ncols_max
 };

 static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) {
@@ -1486,7 +1487,7 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
            continue;
        }

-        const int ntiles_x = (args.ncols_max + config.J - 1) / config.J;
+        const int ntiles_x = (args.ncols_opt + config.J - 1) / config.J;

        if (ntiles_x < ntiles_J_best) {
            J_best = J;
@@ -56,6 +56,7 @@
 #include "htp/unary-ops.h"
 #include "htp/get-rows-ops.h"
 #include "htp/set-rows-ops.h"
+#include "htp/rope-ops.h"
 #include "htp_iface.h"
 #include "htp-drv.h"

@@ -299,6 +300,12 @@ static void ggml_hexagon_precompute_set_rows_params(
    struct htp_set_rows_kernel_params * kparams
 );

+static void ggml_hexagon_precompute_rope_params(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    struct htp_rope_kernel_params * kparams
+);
+
 static void ggml_hexagon_precompute_fused_mmnx_params(
    const struct ggml_hexagon_session * sess,
    const struct ggml_tensor * src0,
@@ -4148,6 +4155,36 @@ static void ggml_hexagon_precompute_set_rows_params(
    kparams->vtcm_size = vtcm_layout.total_bytes;
 }

+static void ggml_hexagon_precompute_rope_params(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    struct htp_rope_kernel_params * kparams
+) {
+    memset(kparams, 0, sizeof(*kparams));
+
+    const struct ggml_tensor * src0 = op->src[0];
+    const struct ggml_tensor * dst  = op;
+
+    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    const uint32_t n_threads  = (std::min)((uint32_t) sess->n_threads, src0_nrows);
+
+    struct htp_rope_vtcm_layout layout;
+    htp_rope_vtcm_layout_build(&layout, src0->ne[0], n_threads);
+
+    kparams->n_threads              = n_threads;
+    kparams->src0_nrows             = src0_nrows;
+    kparams->src0_nrows_per_thread  = (src0_nrows + n_threads - 1) / n_threads;
+    kparams->vtcm_size              = (uint32_t) layout.total_bytes;
+    kparams->spad_per_thread        = (uint32_t) layout.bytes_per_thread;
+    kparams->theta_cache_offset     = (uint32_t) layout.theta_cache_size_aligned;
+    kparams->src0_row_size_aligned  = (uint32_t) layout.src0_row_size_aligned;
+
+    if (src0_nrows > 0) {
+        kparams->div_ne2_ne1 = init_fastdiv_values(dst->ne[2] * dst->ne[1]);
+        kparams->div_ne1     = init_fastdiv_values(dst->ne[1]);
+    }
+}
+
 static void ggml_hexagon_precompute_fused_mmnx_params(
    const struct ggml_hexagon_session * sess,
    const struct ggml_tensor * src0, // W0
@@ -4706,56 +4743,82 @@ static bool ggml_hexagon_supported_argsort(const struct ggml_hexagon_session * s
 }

 static bool ggml_hexagon_supported_rope(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
-    const int32_t * op_params = &op->op_params[0];
-
-    // ggml_rope_set_offset: HVX kernels need a VLEN-aligned window start (32 f32 elems)
-    if (op_params[15] % 32 != 0) {
-        return false;
-    }
-
-    int mode = op_params[2];
-
-    // n_dims == ne0/2, so the rotation spans the full row
-    if (mode == GGML_ROPE_TYPE_VISION) {
-        const int n_dims = op_params[1];
-        if (n_dims != (int) (op->src[0]->ne[0] / 2)) {
-            return false;
-        }
-    }
-    if (mode & 1) {
-        return false;
-    }
-
    const struct ggml_tensor * src0 = op->src[0];
    const struct ggml_tensor * src1 = op->src[1];
    const struct ggml_tensor * src2 = op->src[2];
    const struct ggml_tensor * dst  = op;

-    if (src0->type != GGML_TYPE_F32) {
-        return false;  // FIXME: add support for GGML_TYPE_F16 for src0
-    }
-    if (dst->type != GGML_TYPE_F32) {
+    if (!ggml_are_same_shape(src0, dst)) {
        return false;
    }
-    if (src1->type != GGML_TYPE_I32) {
+
+    if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_I32) {
        return false;
    }
-    if (src2) {
-        if (src2->type != GGML_TYPE_F32) {
-            return false;
-        }
-        int n_dims = op_params[1];
-        if (src2->ne[0] < (n_dims / 2)) {
+
+    if (src0->ne[0] <= 0) {
+        return false;
+    }
+
+    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    if (src0_nrows == 0) {
+        return false;
+    }
+
+    const int32_t * op_params = &op->op_params[0];
+    const int n_dims = op_params[1];
+    const int mode   = op_params[2];
+    const int n_offs = op_params[15];
+
+    if (n_dims <= 0 || n_dims % 2 != 0) {
+        return false;
+    }
+
+    // ggml_rope_set_offset: HVX kernels need a VLEN-aligned window start (32 f32 elems)
+    if (n_offs < 0 || (n_offs % 32 != 0) || (n_offs + n_dims > src0->ne[0])) {
+        return false;
+    }
+
+    float freq_base;
+    memcpy(&freq_base, op_params + 5, sizeof(float));
+    if (freq_base <= 0.0f) {
+        return false;
+    }
+
+    if (mode != GGML_ROPE_TYPE_NORMAL &&
+        mode != GGML_ROPE_TYPE_NEOX &&
+        mode != GGML_ROPE_TYPE_MROPE &&
+        mode != GGML_ROPE_TYPE_VISION &&
+        mode != GGML_ROPE_TYPE_IMROPE) {
+        return false;
+    }
+
+    const bool is_mrope = (mode & GGML_ROPE_TYPE_MROPE) != 0;
+
+    // n_dims == ne0/2, so the rotation spans the full row
+    if (mode == GGML_ROPE_TYPE_VISION) {
+        if (n_dims != (int) (src0->ne[0] / 2) || n_offs != 0) {
            return false;
        }
    }

-    if (src2) {
-        if (!ggml_is_contiguous(src1) || !ggml_is_contiguous(src2)) {
+    if (is_mrope) {
+        const int32_t * sections = op_params + 11;
+        if (sections[0] <= 0 && sections[1] <= 0 && sections[2] <= 0) {
            return false;
        }
-    } else {
-        if (!ggml_is_contiguous(src1)) {
+    }
+
+    const int64_t min_pos_len = (is_mrope || mode == GGML_ROPE_TYPE_VISION) ? src0->ne[2] * 4 : src0->ne[2];
+    if (src1->ne[0] < min_pos_len || !ggml_is_contiguous(src1)) {
+        return false;
+    }
+
+    if (src2) {
+        if (src2->type != GGML_TYPE_F32 || !ggml_is_contiguous(src2)) {
+            return false;
+        }
+        if (src2->ne[0] < (n_dims / 2)) {
            return false;
        }
    }
@@ -4768,9 +4831,16 @@ static bool ggml_hexagon_supported_rope(const struct ggml_hexagon_session * sess
    if (src0->nb[1] < src0->ne[0] * sizeof(float) || dst->nb[1] < dst->ne[0] * sizeof(float)) {
        return false;
    }
-    return true;

-    GGML_UNUSED(sess);
+    const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, src0_nrows);
+
+    struct htp_rope_vtcm_layout layout;
+    htp_rope_vtcm_layout_build(&layout, src0->ne[0], n_threads);
+    if (layout.total_bytes > sess->vtcm_size) {
+        return false;
+    }
+
+    return true;
 }

 static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -5206,6 +5276,11 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg
                    node.node->src[0], node.node->src[1], node.dst(),
                    (struct htp_set_rows_kernel_params *)node.kernel_params
                );
+            } else if (node.opcode == HTP_OP_ROPE) {
+                ggml_hexagon_precompute_rope_params(sess,
+                    node.node,
+                    (struct htp_rope_kernel_params *)node.kernel_params
+                );
            }
            computed_nodes.push_back(std::move(node));
        }
@@ -4,87 +4,75 @@
 #include "hvx-base.h"
 #include "hvx-floor.h"

-static inline HVX_Vector hvx_vec_cos_f32(HVX_Vector x) {
-    HVX_Vector const_inv_pi = hvx_vec_splat_f32(0.3183098861837907f);
-    HVX_Vector const_half   = hvx_vec_splat_f32(0.5f);
-    HVX_Vector const_pi     = hvx_vec_splat_f32(3.141592653589793f);
-    HVX_Vector const_one    = hvx_vec_splat_f32(1.0f);
+// Range-reduce x to y in [-pi/2, pi/2] and the quadrant sign (-1)^n.
+// Floor/truncate need IEEE bits, so convert qf32 back to sf before them.
+static inline void hvx_vec_sincos_reduce_f32(HVX_Vector x, HVX_Vector * y, HVX_Vector * sign) {
+    HVX_Vector const_inv_pi  = hvx_vec_splat_f32(0.3183098861837907f);
+    HVX_Vector const_half    = hvx_vec_splat_f32(0.5f);
+    HVX_Vector const_pi      = hvx_vec_splat_f32(3.141592653589793f);
+    HVX_Vector const_one     = hvx_vec_splat_f32(1.0f);
    HVX_Vector const_neg_one = hvx_vec_splat_f32(-1.0f);
+    HVX_Vector const_one_i   = Q6_V_vsplat_R(1);

-    // n = floor(x * (1/pi) + 0.5)
-    HVX_Vector n_float = hvx_vec_floor_f32(hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(x, const_inv_pi), const_half));
+    HVX_Vector x_over_pi = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(x, const_inv_pi));
+    x_over_pi = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(x_over_pi, const_half));

-    // y = x - n * pi
-    HVX_Vector y = hvx_vec_sub_f32_f32(x, hvx_vec_mul_f32_f32(n_float, const_pi));
+    HVX_Vector n_float = hvx_vec_floor_f32(x_over_pi);
+    HVX_Vector n_int   = hvx_vec_truncate_f32(n_float);

-    // Sign determination: if n is odd, sign is -1.0f, else 1.0f
-    // half_n = n * 0.5f
-    HVX_Vector half_n = hvx_vec_mul_f32_f32(n_float, const_half);
-    // floor_half_n = floor(half_n)
-    HVX_Vector floor_half_n = hvx_vec_floor_f32(half_n);
-    // is_odd = half_n > floor_half_n
-    HVX_VectorPred is_odd = Q6_Q_vcmp_gt_VsfVsf(half_n, floor_half_n);
-    // sign = vmux(is_odd, -1.0f, 1.0f)
-    HVX_Vector sign = Q6_V_vmux_QVV(is_odd, const_neg_one, const_one);
+    HVX_Vector n_pi = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(n_float, const_pi));
+    *y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(x, n_pi));

-    // z = y^2
-    HVX_Vector z = hvx_vec_mul_f32_f32(y, y);
+    HVX_VectorPred is_odd = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(n_int, const_one_i), const_one_i);
+    *sign = Q6_V_vmux_QVV(is_odd, const_neg_one, const_one);
+}
+
+static inline void hvx_vec_sincos_f32(HVX_Vector x, HVX_Vector * vcos, HVX_Vector * vsin) {
+    HVX_Vector y;
+    HVX_Vector sign;
+    hvx_vec_sincos_reduce_f32(x, &y, &sign);
+
+    HVX_Vector z = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(y, y));

-    // Chebyshev approximation for cos(y)
    HVX_Vector c4 = hvx_vec_splat_f32(2.3557242013849433e-05f);
    HVX_Vector c3 = hvx_vec_splat_f32(-0.0013871428263450528f);
    HVX_Vector c2 = hvx_vec_splat_f32(0.041665895266688284f);
    HVX_Vector c1 = hvx_vec_splat_f32(-0.4999999360426369f);
    HVX_Vector c0 = hvx_vec_splat_f32(0.9999999999071725f);

-    HVX_Vector cos_y = hvx_vec_add_f32_f32(c3, hvx_vec_mul_f32_f32(z, c4));
-    cos_y = hvx_vec_add_f32_f32(c2, hvx_vec_mul_f32_f32(z, cos_y));
-    cos_y = hvx_vec_add_f32_f32(c1, hvx_vec_mul_f32_f32(z, cos_y));
-    cos_y = hvx_vec_add_f32_f32(c0, hvx_vec_mul_f32_f32(z, cos_y));
+    HVX_Vector cos_y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(c3, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(z, c4))));
+    cos_y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(c2, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(z, cos_y))));
+    cos_y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(c1, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(z, cos_y))));
+    cos_y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(c0, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(z, cos_y))));

-    return hvx_vec_mul_f32_f32(cos_y, sign);
-}
-
-static inline HVX_Vector hvx_vec_sin_f32(HVX_Vector x) {
-    HVX_Vector const_inv_pi = hvx_vec_splat_f32(0.3183098861837907f);
-    HVX_Vector const_half   = hvx_vec_splat_f32(0.5f);
-    HVX_Vector const_pi     = hvx_vec_splat_f32(3.141592653589793f);
-    HVX_Vector const_one    = hvx_vec_splat_f32(1.0f);
-    HVX_Vector const_neg_one = hvx_vec_splat_f32(-1.0f);
-
-    // n = floor(x * (1/pi) + 0.5)
-    HVX_Vector n_float = hvx_vec_floor_f32(hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(x, const_inv_pi), const_half));
-
-    // y = x - n * pi
-    HVX_Vector y = hvx_vec_sub_f32_f32(x, hvx_vec_mul_f32_f32(n_float, const_pi));
-
-    // Sign determination: if n is odd, sign is -1.0f, else 1.0f
-    // half_n = n * 0.5f
-    HVX_Vector half_n = hvx_vec_mul_f32_f32(n_float, const_half);
-    // floor_half_n = floor(half_n)
-    HVX_Vector floor_half_n = hvx_vec_floor_f32(half_n);
-    // is_odd = half_n > floor_half_n
-    HVX_VectorPred is_odd = Q6_Q_vcmp_gt_VsfVsf(half_n, floor_half_n);
-    // sign = vmux(is_odd, -1.0f, 1.0f)
-    HVX_Vector sign = Q6_V_vmux_QVV(is_odd, const_neg_one, const_one);
-
-    // z = y^2
-    HVX_Vector z = hvx_vec_mul_f32_f32(y, y);
-
-    // Chebyshev approximation for sin(y)
    HVX_Vector s4 = hvx_vec_splat_f32(2.642186986152672e-06f);
    HVX_Vector s3 = hvx_vec_splat_f32(-0.00019825318964070864f);
    HVX_Vector s2 = hvx_vec_splat_f32(0.00833326283319605f);
    HVX_Vector s1 = hvx_vec_splat_f32(-0.16666666082087775f);
    HVX_Vector s0 = hvx_vec_splat_f32(0.999999999915155f);

-    HVX_Vector sin_y = hvx_vec_add_f32_f32(s3, hvx_vec_mul_f32_f32(z, s4));
-    sin_y = hvx_vec_add_f32_f32(s2, hvx_vec_mul_f32_f32(z, sin_y));
-    sin_y = hvx_vec_add_f32_f32(s1, hvx_vec_mul_f32_f32(z, sin_y));
-    sin_y = hvx_vec_add_f32_f32(s0, hvx_vec_mul_f32_f32(z, sin_y));
-    sin_y = hvx_vec_mul_f32_f32(y, sin_y);
+    HVX_Vector sin_y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(s3, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(z, s4))));
+    sin_y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(s2, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(z, sin_y))));
+    sin_y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(s1, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(z, sin_y))));
+    sin_y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(s0, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(z, sin_y))));
+    sin_y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(y, sin_y));

-    return hvx_vec_mul_f32_f32(sin_y, sign);
+    *vcos = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(cos_y, sign));
+    *vsin = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(sin_y, sign));
+}
+
+static inline HVX_Vector hvx_vec_cos_f32(HVX_Vector x) {
+    HVX_Vector vcos;
+    HVX_Vector vsin;
+    hvx_vec_sincos_f32(x, &vcos, &vsin);
+    return vcos;
+}
+
+static inline HVX_Vector hvx_vec_sin_f32(HVX_Vector x) {
+    HVX_Vector vcos;
+    HVX_Vector vsin;
+    hvx_vec_sincos_f32(x, &vcos, &vsin);
+    return vsin;
 }

 #endif /* HVX_SIN_COS_H */
@@ -17,8 +17,8 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
 #include "htp-tensor.h"
+#include "rope-ops.h"

 // Redefined the rope type constants as we can't include ggml.h
 #define HTP_ROPE_TYPE_NORMAL 0
@@ -27,9 +27,6 @@
 #define HTP_ROPE_TYPE_VISION 24
 #define HTP_ROPE_TYPE_IMROPE 40

-#define HTP_ROPE_SPAD_NROWS  16
-#define HTP_ROPE_SPAD_BLOCK  (HTP_ROPE_SPAD_NROWS/2)
-
 #define htp_rope_preamble              \
    const uint32_t ne00 = src0->ne[0]; \
    const uint32_t ne01 = src0->ne[1]; \
@@ -65,26 +62,27 @@ struct htp_rope_context {
    float beta_fast;
    float beta_slow;
    float theta_scale;
+    float theta_scale_32;
+    float theta_powers[32];
    float corr_dims[2];

    uint32_t src0_nrows_per_thread;
-    size_t spad_stride;

    struct htp_ops_context * octx;

+    uint8_t * vtcm_base;
+    size_t    spad_per_thread;
+    size_t    theta_cache_offset;
+
    size_t src0_row_size;
    size_t src0_row_stride;
    size_t dst_row_size;
    size_t dst_row_stride;
    size_t src0_row_size_aligned;
-    size_t dst_row_size_aligned;
-    size_t theta_cache_offset;
    uint32_t src0_nrows;

    struct fastdiv_values div_ne2_ne1;
    struct fastdiv_values div_ne1;
-
-    uint64_t t_start;
 };

 static float rope_yarn_ramp(const float low, const float high, const int i0) {
@@ -112,94 +110,80 @@ static inline void rope_yarn_one(float theta, float freq_scale, float * corr_dim
        mscale_final  *= 1.0f + 0.1f * logf(1.0f / freq_scale);
    }

-    cache[i0 + 0] = cosf(theta_final) * mscale_final;
-    cache[i0 + 1] = sinf(theta_final) * mscale_final;
+    const uint32_t b = i0 / 64;
+    const uint32_t k = (i0 % 64) / 2;
+    cache[b * 64 + k]      = cosf(theta_final) * mscale_final;
+    cache[b * 64 + 32 + k] = sinf(theta_final) * mscale_final;
+}
+
+// 32 thetas -> 32 deinterleaved pairs [cos[32] | sin[32]] at cache[i0].
+static inline void rope_cache_hvx_32(float * cache, uint32_t i0,
+                                     HVX_Vector v_theta,
+                                     const float * freq_factors,
+                                     HVX_Vector v_freq_scale,
+                                     HVX_Vector v_mscale) {
+    if (freq_factors) {
+        HVX_Vector v_ff = hvx_vmemu(freq_factors + i0 / 2);
+        v_theta = hvx_vec_mul_f32_f32(v_theta, hvx_vec_inverse_f32(v_ff));
+    }
+
+    HVX_Vector v_theta_final = hvx_vec_mul_f32_f32(v_theta, v_freq_scale);
+    HVX_Vector vcos;
+    HVX_Vector vsin;
+    hvx_vec_sincos_f32(v_theta_final, &vcos, &vsin);
+    vcos = hvx_vec_mul_f32_f32(vcos, v_mscale);
+    vsin = hvx_vec_mul_f32_f32(vsin, v_mscale);
+
+    if (((uintptr_t) (cache + i0)) % 128 == 0) {
+        hvx_vmem(cache + i0 + 0)  = vcos;
+        hvx_vmem(cache + i0 + 32) = vsin;
+    } else {
+        hvx_vec_store_u(cache + i0 + 0,  32 * sizeof(float), vcos);
+        hvx_vec_store_u(cache + i0 + 32, 32 * sizeof(float), vsin);
+    }
 }

 static __attribute__((noinline)) void rope_cache_init(const float    theta_base,
                            const float    freq_scale,
                            const float *  freq_factors,
                            float *        corr_dims,
-                            const uint32_t ne0,
+                            const uint32_t n_cache,
                            const float    ext_factor,
                            const float    mscale,
                            float *        cache,
-                            const float    theta_scale) {
+                            const float    theta_scale,
+                            const float *  theta_powers,
+                            const float    theta_scale_32) {
    // ref: https://github.com/jquesnelle/yarn/blob/master/scaled_rope/LlamaYaRNScaledRotaryEmbedding.py
-#if __HVX_ARCH__ >= 79
-    const bool is_v79_or_newer = true;
-#else
-    const bool is_v79_or_newer = false;
-#endif
-
-    if (is_v79_or_newer && ext_factor == 0.0f) {
+    if (ext_factor == 0.0f) {
        // Fast path: fully vectorized
        // We process 32 pairs (64 elements) per iteration.
-        const uint32_t n_blocks = ne0 / 64;
-
-        // Initialize theta scale powers: [1.0f, theta_scale, theta_scale^2, ..., theta_scale^31]
-        float __attribute__((aligned(128))) theta_powers[32];
-        theta_powers[0] = 1.0f;
-        for (int j = 1; j < 32; j++) {
-            theta_powers[j] = theta_powers[j - 1] * theta_scale;
-        }
-        HVX_Vector v_theta_powers = hvx_vmem(theta_powers);
+        const uint32_t n_blocks = n_cache / 64;

+        HVX_Vector v_theta_powers = hvx_vmemu(theta_powers);
        HVX_Vector v_freq_scale = hvx_vec_splat_f32(freq_scale);
        HVX_Vector v_mscale = hvx_vec_splat_f32(mscale);

-        // Base theta starts at theta_base
        float theta_block = theta_base;
-        // The scale factor for the next block is theta_scale^32
-        float theta_scale_32 = 1.0f;
-        for (int j = 0; j < 32; j++) {
-            theta_scale_32 *= theta_scale;
-        }

        for (uint32_t b = 0; b < n_blocks; b++) {
            uint32_t i0 = b * 64;
            HVX_Vector v_theta_base = hvx_vec_splat_f32(theta_block);
            HVX_Vector v_theta = hvx_vec_mul_f32_f32(v_theta_base, v_theta_powers);
-
-            if (freq_factors) {
-                // Load 32 elements of freq_factors
-                HVX_Vector v_ff = hvx_vmemu(freq_factors + i0 / 2);
-                HVX_Vector v_inv_ff = hvx_vec_inverse_f32(v_ff);
-                v_theta = hvx_vec_mul_f32_f32(v_theta, v_inv_ff);
-            }
-
-            HVX_Vector v_theta_final = hvx_vec_mul_f32_f32(v_theta, v_freq_scale);
-
-            HVX_Vector vcos = hvx_vec_cos_f32(v_theta_final);
-            HVX_Vector vsin = hvx_vec_sin_f32(v_theta_final);
-
-            vcos = hvx_vec_mul_f32_f32(vcos, v_mscale);
-            vsin = hvx_vec_mul_f32_f32(vsin, v_mscale);
-
-            HVX_VectorPair vstore = Q6_W_vshuff_VVR(vsin, vcos, -4);
-
-            if (((uintptr_t)cache) % 128 == 0) {
-                hvx_vmem(cache + i0 + 0)  = Q6_V_lo_W(vstore);
-                hvx_vmem(cache + i0 + 32) = Q6_V_hi_W(vstore);
-            } else {
-                hvx_vec_store_u(cache + i0 + 0,  32 * sizeof(float), Q6_V_lo_W(vstore));
-                hvx_vec_store_u(cache + i0 + 32, 32 * sizeof(float), Q6_V_hi_W(vstore));
-            }
-
+            rope_cache_hvx_32(cache, i0, v_theta, freq_factors, v_freq_scale, v_mscale);
            theta_block *= theta_scale_32;
        }

        // Leftovers
        float theta = theta_block;
-        for (uint32_t i0 = n_blocks * 64; i0 < ne0; i0 += 2) {
+        for (uint32_t i0 = n_blocks * 64; i0 < n_cache; i0 += 2) {
            const float ff = freq_factors ? freq_factors[i0 / 2] : 1.0f;
            rope_yarn_one(theta / ff, freq_scale, corr_dims, i0, ext_factor, mscale, cache);
            theta *= theta_scale;
        }
    } else {
-        // Fallback to original scalar loop
        float theta = theta_base;
-        for (uint32_t i0 = 0; i0 < ne0; i0 += 2) {
+        for (uint32_t i0 = 0; i0 < n_cache; i0 += 2) {
            const float ff = freq_factors ? freq_factors[i0 / 2] : 1.0f;
            rope_yarn_one(theta / ff, freq_scale, corr_dims, i0, ext_factor, mscale, cache);
            theta *= theta_scale;
@@ -207,6 +191,72 @@ static __attribute__((noinline)) void rope_cache_init(const float    theta_base,
    }
 }

+static inline float mrope_pick_theta(float theta_t, float theta_h, float theta_w, float theta_e,
+                                     int sector, const int32_t sections[4], int sec_w, int sec_e,
+                                     bool is_imrope) {
+    if (is_imrope) {
+        if      (sector % 3 == 0 && sector < 3 * sections[0]) { return theta_t; }
+        else if (sector % 3 == 1 && sector < 3 * sections[1]) { return theta_h; }
+        else if (sector % 3 == 2 && sector < 3 * sections[2]) { return theta_w; }
+        else                                                   { return theta_e; }
+    }
+    if      (sector < sections[0]) { return theta_t; }
+    else if (sector < sec_w)       { return theta_h; }
+    else if (sector < sec_e)       { return theta_w; }
+    else                           { return theta_e; }
+}
+
+// lane j is 1 when (j % 3) == rem
+static const float __attribute__((aligned(128))) mrope_mod3_eq0[32] = {
+    1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0
+};
+static const float __attribute__((aligned(128))) mrope_mod3_eq1[32] = {
+    0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1
+};
+static const float __attribute__((aligned(128))) mrope_mod3_eq2[32] = {
+    0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0,1,0,0
+};
+
+static const float __attribute__((aligned(128))) mrope_k_ramp[32] = {
+    0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,
+    16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31
+};
+
+static inline HVX_VectorPred mrope_mask_eq1(const float * m) {
+    return Q6_Q_vcmp_gt_VsfVsf(hvx_vmemu(m), Q6_V_vzero());
+}
+
+// IMROPE without wrap: theta[k] = pos[k % 3] * scale^k
+static inline HVX_Vector mrope_thetas_imrope_mod3(float pos_t, float pos_h, float pos_w,
+                                                  uint32_t k0, HVX_Vector v_powers, float scale_block) {
+    const int r = (int) (k0 % 3);
+    const float * mt = (r == 0) ? mrope_mod3_eq0 : (r == 1) ? mrope_mod3_eq2 : mrope_mod3_eq1;
+    const float * mh = (r == 0) ? mrope_mod3_eq1 : (r == 1) ? mrope_mod3_eq0 : mrope_mod3_eq2;
+
+    HVX_Vector v = hvx_vec_splat_f32(pos_w);
+    v = Q6_V_vmux_QVV(mrope_mask_eq1(mh), hvx_vec_splat_f32(pos_h), v);
+    v = Q6_V_vmux_QVV(mrope_mask_eq1(mt), hvx_vec_splat_f32(pos_t), v);
+    v = hvx_vec_mul_f32_f32(v, v_powers);
+    return hvx_vec_mul_f32_f32(v, hvx_vec_splat_f32(scale_block));
+}
+
+// Contiguous MROPE without wrap: theta[k] = pos[section(k)] * scale^k
+static inline HVX_Vector mrope_thetas_contig(float pos_t, float pos_h, float pos_w, float pos_e,
+                                             uint32_t k0, int s0, int sec_w, int sec_e,
+                                             HVX_Vector v_powers, float scale_block) {
+    HVX_Vector v_k = hvx_vec_add_f32_f32(hvx_vec_splat_f32((float) k0), hvx_vmemu(mrope_k_ramp));
+    HVX_VectorPred lt_s0 = Q6_Q_vcmp_gt_VsfVsf(hvx_vec_splat_f32((float) s0), v_k);
+    HVX_VectorPred lt_sw = Q6_Q_vcmp_gt_VsfVsf(hvx_vec_splat_f32((float) sec_w), v_k);
+    HVX_VectorPred lt_se = Q6_Q_vcmp_gt_VsfVsf(hvx_vec_splat_f32((float) sec_e), v_k);
+
+    HVX_Vector v = hvx_vec_splat_f32(pos_e);
+    v = Q6_V_vmux_QVV(lt_se, hvx_vec_splat_f32(pos_w), v);
+    v = Q6_V_vmux_QVV(lt_sw, hvx_vec_splat_f32(pos_h), v);
+    v = Q6_V_vmux_QVV(lt_s0, hvx_vec_splat_f32(pos_t), v);
+    v = hvx_vec_mul_f32_f32(v, v_powers);
+    return hvx_vec_mul_f32_f32(v, hvx_vec_splat_f32(scale_block));
+}
+
 // pos_t/h/w/e: the four position ids for this sequence step (t=time, h=height, w=width, e=extra).
 // sections[4]: number of head dims assigned to each position component.
 static __attribute__((noinline)) void mrope_cache_init(const float    pos_t,
@@ -219,23 +269,71 @@ static __attribute__((noinline)) void mrope_cache_init(const float    pos_t,
                             const float    freq_scale,
                             const float *  freq_factors,
                             float *        corr_dims,
-                             const uint32_t ne0,
+                             const uint32_t n_cache,
                             const float    ext_factor,
                             const float    mscale,
                             float *        cache,
-                             const float    theta_scale) {
+                             const float    theta_scale,
+                             const float *  theta_powers,
+                             const float    theta_scale_32) {
    const int sect_dims = sections[0] + sections[1] + sections[2] + sections[3];
    const int sec_w     = sections[0] + sections[1];
    const int sec_e     = sec_w + sections[2];
+    const uint32_t n_pairs = n_cache / 2;
+
+    const bool no_wrap = (sect_dims > 0) && (n_pairs <= (uint32_t) sect_dims);
+    const bool imrope_mod3 = is_imrope && !indep_sects && no_wrap
+        && sections[0] > 0 && sections[1] > 0 && sections[2] > 0
+        && n_pairs <= (uint32_t) (3 * sections[0])
+        && n_pairs <= (uint32_t) (3 * sections[1])
+        && n_pairs <= (uint32_t) (3 * sections[2]);
+    const bool contig = !is_imrope && !indep_sects && no_wrap;
+
+    if (ext_factor == 0.0f && (imrope_mod3 || contig)) {
+        HVX_Vector v_powers     = hvx_vmemu(theta_powers);
+        HVX_Vector v_freq_scale = hvx_vec_splat_f32(freq_scale);
+        HVX_Vector v_mscale     = hvx_vec_splat_f32(mscale);
+        float scale_block = 1.0f;
+        const uint32_t n_blocks = n_cache / 64;
+
+        for (uint32_t b = 0; b < n_blocks; b++) {
+            const uint32_t i0 = b * 64;
+            const uint32_t k0 = b * 32;
+            HVX_Vector v_theta = imrope_mod3
+                ? mrope_thetas_imrope_mod3(pos_t, pos_h, pos_w, k0, v_powers, scale_block)
+                : mrope_thetas_contig(pos_t, pos_h, pos_w, pos_e, k0, sections[0], sec_w, sec_e,
+                                      v_powers, scale_block);
+            rope_cache_hvx_32(cache, i0, v_theta, freq_factors, v_freq_scale, v_mscale);
+            scale_block *= theta_scale_32;
+        }
+
+        float theta_k = scale_block;
+        for (uint32_t k = n_blocks * 32; k < n_pairs; k++) {
+            const uint32_t i0 = 2 * k;
+            const float pos = mrope_pick_theta(pos_t, pos_h, pos_w, pos_e,
+                                               (int) k, sections, sec_w, sec_e, is_imrope);
+            const float ff = freq_factors ? freq_factors[k] : 1.0f;
+            rope_yarn_one(pos * theta_k / ff, freq_scale, corr_dims, i0, ext_factor, mscale, cache);
+            theta_k *= theta_scale;
+        }
+        return;
+    }

    float theta_t = pos_t;
    float theta_h = pos_h;
    float theta_w = pos_w;
    float theta_e = pos_e;

-    for (uint32_t i0 = 0; i0 < ne0; i0 += 2) {
-        const float ff     = freq_factors ? freq_factors[i0 / 2] : 1.0f;
-        const int   sector = (i0 / 2) % sect_dims;
+    const bool use_hvx = (ext_factor == 0.0f);
+    float __attribute__((aligned(128))) thetas[32];
+    uint32_t n_thetas = 0;
+    uint32_t block_i0 = 0;
+
+    HVX_Vector v_freq_scale = hvx_vec_splat_f32(freq_scale);
+    HVX_Vector v_mscale     = hvx_vec_splat_f32(mscale);
+
+    for (uint32_t i0 = 0; i0 < n_cache; i0 += 2) {
+        const int sector = (i0 / 2) % sect_dims;

        if (indep_sects) {
            // Reset theta when crossing into a new section.
@@ -245,28 +343,34 @@ static __attribute__((noinline)) void mrope_cache_init(const float    pos_t,
            else if (sector == sec_e)       { theta_e = pos_e; }
        }

-        float theta;
-        if (is_imrope) {
-            // Interleaved: sector mod 3 selects component
-            if      (sector % 3 == 0 && sector < 3 * sections[0]) { theta = theta_t; }
-            else if (sector % 3 == 1 && sector < 3 * sections[1]) { theta = theta_h; }
-            else if (sector % 3 == 2 && sector < 3 * sections[2]) { theta = theta_w; }
-            else                                                   { theta = theta_e; }
-        } else {
-            // Contiguous sections
-            if      (sector < sections[0]) { theta = theta_t; }
-            else if (sector < sec_w)       { theta = theta_h; }
-            else if (sector < sec_e)       { theta = theta_w; }
-            else                           { theta = theta_e; }
-        }
+        const float theta = mrope_pick_theta(theta_t, theta_h, theta_w, theta_e,
+                                             sector, sections, sec_w, sec_e, is_imrope);

-        rope_yarn_one(theta / ff, freq_scale, corr_dims, i0, ext_factor, mscale, cache);
+        if (use_hvx) {
+            if (n_thetas == 0) {
+                block_i0 = i0;
+            }
+            thetas[n_thetas++] = theta;
+            if (n_thetas == 32) {
+                rope_cache_hvx_32(cache, block_i0, hvx_vmemu(thetas), freq_factors, v_freq_scale, v_mscale);
+                n_thetas = 0;
+            }
+        } else {
+            const float ff = freq_factors ? freq_factors[i0 / 2] : 1.0f;
+            rope_yarn_one(theta / ff, freq_scale, corr_dims, i0, ext_factor, mscale, cache);
+        }

        theta_t *= theta_scale;
        theta_h *= theta_scale;
        theta_w *= theta_scale;
        theta_e *= theta_scale;
    }
+
+    for (uint32_t k = 0; k < n_thetas; k++) {
+        const uint32_t i0 = block_i0 + 2 * k;
+        const float ff = freq_factors ? freq_factors[i0 / 2] : 1.0f;
+        rope_yarn_one(thetas[k] / ff, freq_scale, corr_dims, i0, ext_factor, mscale, cache);
+    }
 }

 #define M_PI 3.1415926535897932384626433
@@ -283,52 +387,54 @@ static void rope_corr_dims(int     n_dims,
    dims[1]     = MIN(n_dims - 1, end);
 }

+static inline void hvx_rope_neox_mul(HVX_Vector v0, HVX_Vector v1, HVX_Vector vcos, HVX_Vector vsin,
+                                     HVX_Vector * o0, HVX_Vector * o1) {
+    HVX_Vector vx0_c = Q6_Vqf32_vmpy_VsfVsf(v0, vcos);
+    HVX_Vector vx0_s = Q6_Vqf32_vmpy_VsfVsf(v0, vsin);
+    HVX_Vector vx1_c = Q6_Vqf32_vmpy_VsfVsf(v1, vcos);
+    HVX_Vector vx1_s = Q6_Vqf32_vmpy_VsfVsf(v1, vsin);
+    *o0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(vx0_c, vx1_s));
+    *o1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(vx0_s, vx1_c));
+}
+
+// theta_cache full 32-pair blocks are deinterleaved [cos | sin].
 static inline void hvx_rope_neox_f32_aa(float * restrict dst, const float * restrict src0, uint32_t ne, const float * restrict theta_cache) {
    const uint32_t he = ne / 2;
    const uint32_t nvec = he / 32;
    const uint32_t nloe = he % 32;

+    if (nloe == 0) {
+        const HVX_Vector * vs = (const HVX_Vector *) src0;
+        const HVX_Vector * vt = (const HVX_Vector *) theta_cache;
+        HVX_Vector * vd = (HVX_Vector *) dst;
+        for (uint32_t i = 0; i < nvec; i++) {
+            HVX_Vector o0, o1;
+            hvx_rope_neox_mul(vs[i], vs[nvec + i], vt[i * 2 + 0], vt[i * 2 + 1], &o0, &o1);
+            vd[i] = o0;
+            vd[nvec + i] = o1;
+        }
+        return;
+    }
+
    for (uint32_t i = 0; i < nvec; i++) {
-        HVX_Vector v0 = ((const HVX_Vector *) src0)[i];
-        HVX_Vector v1 = hvx_vmemu(src0 + he + i * 32);
-
-        HVX_Vector v2 = ((const HVX_Vector *) theta_cache)[i * 2 + 0];
-        HVX_Vector v3 = ((const HVX_Vector *) theta_cache)[i * 2 + 1];
-
-        HVX_VectorPair vcos_sin = Q6_W_vdeal_VVR(v3, v2, -4);
-
-        HVX_Vector vx0_c = Q6_Vqf32_vmpy_VsfVsf(v0, Q6_V_lo_W(vcos_sin));
-        HVX_Vector vx0_s = Q6_Vqf32_vmpy_VsfVsf(v0, Q6_V_hi_W(vcos_sin));
-        HVX_Vector vx1_c = Q6_Vqf32_vmpy_VsfVsf(v1, Q6_V_lo_W(vcos_sin));
-        HVX_Vector vx1_s = Q6_Vqf32_vmpy_VsfVsf(v1, Q6_V_hi_W(vcos_sin));
-
-        HVX_Vector v4 = Q6_Vqf32_vsub_Vqf32Vqf32(vx0_c, vx1_s);
-        HVX_Vector v5 = Q6_Vqf32_vadd_Vqf32Vqf32(vx0_s, vx1_c);
-
-        ((HVX_Vector *) dst)[i] = Q6_Vsf_equals_Vqf32(v4);
-        hvx_vmemu(dst + he + i * 32) = Q6_Vsf_equals_Vqf32(v5);
+        HVX_Vector o0, o1;
+        hvx_rope_neox_mul(((const HVX_Vector *) src0)[i],
+                          hvx_vmemu(src0 + he + i * 32),
+                          ((const HVX_Vector *) theta_cache)[i * 2 + 0],
+                          ((const HVX_Vector *) theta_cache)[i * 2 + 1],
+                          &o0, &o1);
+        ((HVX_Vector *) dst)[i] = o0;
+        hvx_vmemu(dst + he + i * 32) = o1;
    }

-    if (nloe > 0) {
-        HVX_Vector v0 = hvx_vmemu(src0 + nvec * 32);
-        HVX_Vector v1 = hvx_vmemu(src0 + he + nvec * 32);
-
-        HVX_Vector v2 = ((const HVX_Vector *) theta_cache)[nvec * 2 + 0];
-        HVX_Vector v3 = ((const HVX_Vector *) theta_cache)[nvec * 2 + 1];
-
-        HVX_VectorPair vcos_sin = Q6_W_vdeal_VVR(v3, v2, -4);
-
-        HVX_Vector vx0_c = Q6_Vqf32_vmpy_VsfVsf(v0, Q6_V_lo_W(vcos_sin));
-        HVX_Vector vx0_s = Q6_Vqf32_vmpy_VsfVsf(v0, Q6_V_hi_W(vcos_sin));
-        HVX_Vector vx1_c = Q6_Vqf32_vmpy_VsfVsf(v1, Q6_V_lo_W(vcos_sin));
-        HVX_Vector vx1_s = Q6_Vqf32_vmpy_VsfVsf(v1, Q6_V_hi_W(vcos_sin));
-
-        HVX_Vector v4 = Q6_Vqf32_vsub_Vqf32Vqf32(vx0_c, vx1_s);
-        HVX_Vector v5 = Q6_Vqf32_vadd_Vqf32Vqf32(vx0_s, vx1_c);
-
-        hvx_vec_store_u(dst + nvec * 32, nloe * sizeof(float), Q6_Vsf_equals_Vqf32(v4));
-        hvx_vec_store_u(dst + he + nvec * 32, nloe * sizeof(float), Q6_Vsf_equals_Vqf32(v5));
-    }
+    HVX_Vector v0 = hvx_vmemu(src0 + nvec * 32);
+    HVX_Vector v1 = hvx_vmemu(src0 + he + nvec * 32);
+    HVX_Vector vcos = hvx_vmemu(theta_cache + nvec * 64);
+    HVX_Vector vsin = hvx_vmemu(theta_cache + nvec * 64 + 32);
+    HVX_Vector o0, o1;
+    hvx_rope_neox_mul(v0, v1, vcos, vsin, &o0, &o1);
+    hvx_vec_store_u(dst + nvec * 32, nloe * sizeof(float), o0);
+    hvx_vec_store_u(dst + he + nvec * 32, nloe * sizeof(float), o1);
 }

 static inline void hvx_rope_f32_aa(float * restrict dst, const float * restrict src0, uint32_t ne, const float * restrict theta_cache) {
@@ -339,16 +445,15 @@ static inline void hvx_rope_f32_aa(float * restrict dst, const float * restrict
        HVX_Vector v0 = ((const HVX_Vector *) src0)[i * 2 + 0];
        HVX_Vector v1 = ((const HVX_Vector *) src0)[i * 2 + 1];

-        HVX_Vector v2 = ((const HVX_Vector *) theta_cache)[i * 2 + 0];
-        HVX_Vector v3 = ((const HVX_Vector *) theta_cache)[i * 2 + 1];
+        HVX_Vector vcos = ((const HVX_Vector *) theta_cache)[i * 2 + 0];
+        HVX_Vector vsin = ((const HVX_Vector *) theta_cache)[i * 2 + 1];

-        HVX_VectorPair vx0_x1   = Q6_W_vdeal_VVR(v1, v0, -4);
-        HVX_VectorPair vcos_sin = Q6_W_vdeal_VVR(v3, v2, -4);
+        HVX_VectorPair vx0_x1 = Q6_W_vdeal_VVR(v1, v0, -4);

-        HVX_Vector vx0_c = Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(vx0_x1), Q6_V_lo_W(vcos_sin));
-        HVX_Vector vx0_s = Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(vx0_x1), Q6_V_hi_W(vcos_sin));
-        HVX_Vector vx1_c = Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(vx0_x1), Q6_V_lo_W(vcos_sin));
-        HVX_Vector vx1_s = Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(vx0_x1), Q6_V_hi_W(vcos_sin));
+        HVX_Vector vx0_c = Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(vx0_x1), vcos);
+        HVX_Vector vx0_s = Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(vx0_x1), vsin);
+        HVX_Vector vx1_c = Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(vx0_x1), vcos);
+        HVX_Vector vx1_s = Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(vx0_x1), vsin);

        HVX_Vector v4 = Q6_Vqf32_vsub_Vqf32Vqf32(vx0_c, vx1_s);
        HVX_Vector v5 = Q6_Vqf32_vadd_Vqf32Vqf32(vx0_s, vx1_c);
@@ -362,15 +467,15 @@ static inline void hvx_rope_f32_aa(float * restrict dst, const float * restrict
    if (nloe > 0) {
        if (nloe <= 32) {
            HVX_Vector v0 = hvx_vmemu(src0 + nvec * 64);
-            HVX_Vector v2 = hvx_vmemu(theta_cache + nvec * 64);
+            HVX_Vector vcos = hvx_vmemu(theta_cache + nvec * 64);
+            HVX_Vector vsin = hvx_vmemu(theta_cache + nvec * 64 + 32);

-            HVX_VectorPair vx0_x1   = Q6_W_vdeal_VVR(Q6_V_vzero(), v0, -4);
-            HVX_VectorPair vcos_sin = Q6_W_vdeal_VVR(Q6_V_vzero(), v2, -4);
+            HVX_VectorPair vx0_x1 = Q6_W_vdeal_VVR(Q6_V_vzero(), v0, -4);

-            HVX_Vector vx0_c = Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(vx0_x1), Q6_V_lo_W(vcos_sin));
-            HVX_Vector vx0_s = Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(vx0_x1), Q6_V_hi_W(vcos_sin));
-            HVX_Vector vx1_c = Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(vx0_x1), Q6_V_lo_W(vcos_sin));
-            HVX_Vector vx1_s = Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(vx0_x1), Q6_V_hi_W(vcos_sin));
+            HVX_Vector vx0_c = Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(vx0_x1), vcos);
+            HVX_Vector vx0_s = Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(vx0_x1), vsin);
+            HVX_Vector vx1_c = Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(vx0_x1), vcos);
+            HVX_Vector vx1_s = Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(vx0_x1), vsin);

            HVX_Vector v4 = Q6_Vqf32_vsub_Vqf32Vqf32(vx0_c, vx1_s);
            HVX_Vector v5 = Q6_Vqf32_vadd_Vqf32Vqf32(vx0_s, vx1_c);
@@ -382,16 +487,15 @@ static inline void hvx_rope_f32_aa(float * restrict dst, const float * restrict
            HVX_Vector v0 = hvx_vmemu(src0 + nvec * 64);
            HVX_Vector v1 = hvx_vmemu(src0 + nvec * 64 + 32);

-            HVX_Vector v2 = hvx_vmemu(theta_cache + nvec * 64);
-            HVX_Vector v3 = hvx_vmemu(theta_cache + nvec * 64 + 32);
+            HVX_Vector vcos = hvx_vmemu(theta_cache + nvec * 64);
+            HVX_Vector vsin = hvx_vmemu(theta_cache + nvec * 64 + 32);

-            HVX_VectorPair vx0_x1   = Q6_W_vdeal_VVR(v1, v0, -4);
-            HVX_VectorPair vcos_sin = Q6_W_vdeal_VVR(v3, v2, -4);
+            HVX_VectorPair vx0_x1 = Q6_W_vdeal_VVR(v1, v0, -4);

-            HVX_Vector vx0_c = Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(vx0_x1), Q6_V_lo_W(vcos_sin));
-            HVX_Vector vx0_s = Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(vx0_x1), Q6_V_hi_W(vcos_sin));
-            HVX_Vector vx1_c = Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(vx0_x1), Q6_V_lo_W(vcos_sin));
-            HVX_Vector vx1_s = Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(vx0_x1), Q6_V_hi_W(vcos_sin));
+            HVX_Vector vx0_c = Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(vx0_x1), vcos);
+            HVX_Vector vx0_s = Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(vx0_x1), vsin);
+            HVX_Vector vx1_c = Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(vx0_x1), vcos);
+            HVX_Vector vx1_s = Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(vx0_x1), vsin);

            HVX_Vector v4 = Q6_Vqf32_vsub_Vqf32Vqf32(vx0_c, vx1_s);
            HVX_Vector v5 = Q6_Vqf32_vadd_Vqf32Vqf32(vx0_s, vx1_c);
@@ -404,54 +508,23 @@ static inline void hvx_rope_f32_aa(float * restrict dst, const float * restrict
    }
 }

-static void inline rope_basic_f32(struct htp_rope_context * rctx, uint8_t * restrict dst, uint8_t * restrict src,
-                   uint32_t nr, uint32_t ne0, const float * restrict theta_cache) {
-    const uint32_t n_offs = rctx->n_offs; // VLEN-aligned (enforced by supports_op)
+static void inline rope_basic_f32_inplace(struct htp_rope_context * rctx, uint8_t * src,
+                   uint32_t nr, const float * restrict theta_cache) {
+    const uint32_t n_offs = rctx->n_offs;
    #pragma unroll(4)
    for (uint32_t i = 0; i < nr; i++) {
-        float * d = (float *) (dst + i * rctx->dst_row_size_aligned);
        float * s = (float *) (src + i * rctx->src0_row_size_aligned);
-
-        hvx_rope_f32_aa(d + n_offs, s + n_offs, rctx->n_dims, theta_cache);
-
-        // fill the remain channels with data from src tensor
-        if (n_offs > 0) {
-            hvx_copy_f32_uu((uint8_t *) d, (uint8_t *) s, n_offs);
-        }
-        if (n_offs + rctx->n_dims < ne0) {
-            hvx_copy_f32_uu((uint8_t *)(d + n_offs + rctx->n_dims), (uint8_t *)(s + n_offs + rctx->n_dims), ne0 - n_offs - rctx->n_dims);
-        }
+        hvx_rope_f32_aa(s + n_offs, s + n_offs, rctx->n_dims, theta_cache);
    }
 }

-static void inline rope_neox_f32(struct htp_rope_context * rctx, uint8_t * restrict dst, uint8_t * restrict src,
-                   uint32_t nr, uint32_t ne0, const float * restrict theta_cache) {
-    const uint32_t n_offs = rctx->n_offs; // VLEN-aligned (enforced by supports_op)
+static void inline rope_neox_f32_inplace(struct htp_rope_context * rctx, uint8_t * src,
+                   uint32_t nr, uint32_t ne, const float * restrict theta_cache) {
+    const uint32_t n_offs = rctx->n_offs;
    #pragma unroll(4)
    for (uint32_t i = 0; i < nr; i++) {
-        float * d = (float *) (dst + i * rctx->dst_row_size_aligned);
        float * s = (float *) (src + i * rctx->src0_row_size_aligned);
-
-        hvx_rope_neox_f32_aa(d + n_offs, s + n_offs, rctx->n_dims, theta_cache);
-
-        // fill the remain channels with data from src tensor
-        if (n_offs > 0) {
-            hvx_copy_f32_uu((uint8_t *) d, (uint8_t *) s, n_offs);
-        }
-        if (n_offs + rctx->n_dims < ne0) {
-            hvx_copy_f32_uu((uint8_t *)(d + n_offs + rctx->n_dims), (uint8_t *)(s + n_offs + rctx->n_dims), ne0 - n_offs - rctx->n_dims);
-        }
-    }
-}
-
-static void inline rope_vision_f32(struct htp_rope_context * rctx, uint8_t * restrict dst, uint8_t * restrict src,
-                   uint32_t nr, uint32_t ne0, const float * restrict theta_cache) {
-    #pragma unroll(4)
-    for (uint32_t i = 0; i < nr; i++) {
-        float * d = (float *) (dst + i * rctx->dst_row_size_aligned);
-        float * s = (float *) (src + i * rctx->src0_row_size_aligned);
-
-        hvx_rope_neox_f32_aa(d, s, ne0, theta_cache);
+        hvx_rope_neox_f32_aa(s + n_offs, s + n_offs, ne, theta_cache);
    }
 }

@@ -477,20 +550,18 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {
        return;
    }

-    uint64_t tt = HAP_perf_get_qtimer_count();
-
    const int32_t mode    = rctx->mode;
    // MROPE, IMROPE and VISION use NEOX-style pairing for the rotation
    const bool    is_neox = (mode & HTP_ROPE_TYPE_NEOX) || (mode & HTP_ROPE_TYPE_MROPE);
    const bool    is_vision = (mode == HTP_ROPE_TYPE_VISION);

    // VTCM setup
-    uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
+    uint8_t * src0_spad_base = rctx->vtcm_base + (ith * rctx->spad_per_thread);
    float *   theta_cache    = (float *) (src0_spad_base);
              src0_spad_base = src0_spad_base + rctx->theta_cache_offset;
-    uint8_t * dst_spad_base  = octx->dst_spad.data + (ith * octx->dst_spad.size_per_thread);

    dma_queue * dma_queue = octx->ctx->dma[ith];
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
    const int32_t * pos = (const int32_t *) src1->data;
    const float * freq_factors = src2 ? (const float *) src2->data : NULL;

@@ -501,6 +572,7 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {

    uint32_t ir = src0_start_row;
    uint32_t prev_i2 = (uint32_t) -1;
+    uint32_t cur_slot = 0;

    for (uint32_t i3 = i3_start; i3 < ne3; i3++) { // batch
        const uint32_t i2_init = (i3 == i3_start) ? i2_start : 0;
@@ -513,35 +585,30 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {
                const uint32_t nrows = MIN(src0_end_row - ir, ne1 - i1);

                // Depth before prefetch
-                uint32_t dma_depth = dma_queue_depth(dma_queue);
+                const uint32_t dma_depth = dma_queue_depth(dma_queue);

-                // FARF(HIGH, "rope-block %u: ir %u n-rows %u dma-depth %u : usec %u", ith, ir, nrows, dma_depth,
-                //             (unsigned) HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - rctx->t_start));
+                // Prefetch up to 2 blocks
+                const uint32_t p_nrows = MIN(nrows, 2 * HTP_ROPE_SPAD_BLOCK);
+                for (uint32_t pr = 0; pr < p_nrows; pr += HTP_ROPE_SPAD_BLOCK) {
+                    const uint32_t pnr = MIN(nrows - pr, HTP_ROPE_SPAD_BLOCK);
+                    const uint32_t slot = (cur_slot + pr / HTP_ROPE_SPAD_BLOCK) % HTP_ROPE_SPAD_NSLOTS;
+                    uint8_t * spad_slot = rope_spad_slot(src0_spad_base, slot, rctx->src0_row_size_aligned);
+                    const uint8_t * src_addr = (const uint8_t *) src0->data + i3 * nb03 + i2 * nb02 + (i1 + pr) * nb01;

-                // Prefetch loop
-                for (uint32_t pnr = 0, pr = 0; pr < nrows && pr < HTP_ROPE_SPAD_NROWS; pr += pnr) {
-                    pnr = MIN(nrows - pr, HTP_ROPE_SPAD_BLOCK);
+                    // Dummy DMA transaction for sequencing (interleaving wr, rd, wr, rd, ...)
+                    dma_queue_push(dma_queue, dma_make_ptr((void *) dst->data, spad_slot), 0, 0, 0, 0);

-                    uint32_t pi1 = i1 + pr;
-                    uint32_t pir = ir + pr;
-
-                    // Dummy DMA transaction for sequencing (interleaving dst,src,dst,...)
-                    dma_queue_push_vtcm_to_ddr(dma_queue, dma_make_ptr((void *) dst->data, dst_spad_base + pr * rctx->dst_row_size_aligned), 0, 0, 0);
-
-                    const uint8_t * src_addr = (const uint8_t *) src0->data + i3 * nb03 + i2 * nb02 + pi1 * nb01;
-                          uint8_t * src_spad = src0_spad_base + pr * rctx->src0_row_size_aligned;
-
-                    // Copy only the row payload while striding the DDR source
-                    dma_queue_push(dma_queue, dma_make_ptr(src_spad, src_addr),
+                    dma_queue_push(dma_queue, dma_make_ptr(spad_slot, src_addr),
                        rctx->src0_row_size_aligned, rctx->src0_row_stride, rctx->src0_row_size, pnr);
-
-                    // FARF(HIGH, "rope-prefetch %u: pr %u i1 %u i2 %u i3 %u src-spad %p src-addr %p pnr %u", ith, pir, pi1, i2, i3, src_spad, src_addr, pnr);
                }

                // Update theta cache
                if (i2 != prev_i2) {
                    prev_i2 = i2;

+                    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, i2);
+                    // VISION rotates the full row; other modes only rotate n_dims.
+                    const uint32_t n_cache = is_vision ? ne0 : (uint32_t) rctx->n_dims;
                    const bool is_mrope = (rctx->mode & HTP_ROPE_TYPE_MROPE) != 0;
                    if (is_mrope) {
                        // src1 holds four position arrays stacked along ne0:
@@ -554,66 +621,71 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {
                            (float) pos[i2 + ne2 * 3],
                            rctx->sections, is_imrope, is_vision,
                            rctx->freq_scale, freq_factors, rctx->corr_dims,
-                            ne0, rctx->ext_factor, rctx->attn_factor,
-                            theta_cache, rctx->theta_scale);
+                            n_cache, rctx->ext_factor, rctx->attn_factor,
+                            theta_cache, rctx->theta_scale, rctx->theta_powers, rctx->theta_scale_32);
                    } else {
                       rope_cache_init(pos[i2], rctx->freq_scale, freq_factors, rctx->corr_dims,
-                                        ne0, rctx->ext_factor, rctx->attn_factor,
-                                        theta_cache, rctx->theta_scale);
+                                        n_cache, rctx->ext_factor, rctx->attn_factor,
+                                        theta_cache, rctx->theta_scale, rctx->theta_powers, rctx->theta_scale_32);
                    }
+                    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, i2);
                }

                // Skip output DMA transactions from prev block (if any)
-                // No need to wait for those here since we're explicitly waiting for the latest prefecthes below.
-                for (uint32_t d=0; d < dma_depth; d++) { dma_queue_pop_nowait(dma_queue); }
+                for (uint32_t d = 0; d < dma_depth; d++) { dma_queue_pop_nowait(dma_queue); }

                // Compute loop
-                for (uint32_t cnr = 0, cr = 0; cr < nrows; cr += cnr, ir += cnr, i1 += cnr) {
-                    // Number of rows to compute
+                const uint32_t ne = is_vision ? ne0 : rctx->n_dims;
+                const uint32_t base_i1 = i1;
+                const uint32_t base_ir = ir;
+
+                for (uint32_t cnr = 0, cr = 0; cr < nrows; cr += cnr) {
                    cnr = MIN(nrows - cr, HTP_ROPE_SPAD_BLOCK);
+                    const uint32_t slot = (cur_slot + cr / HTP_ROPE_SPAD_BLOCK) % HTP_ROPE_SPAD_NSLOTS;
+                    const uint32_t cur_ir   = base_ir + cr;
+                    const uint32_t cur_i1   = base_i1 + cr;

-                    uint8_t * dst_spad = (uint8_t *) dma_queue_pop(dma_queue).src;
-                    uint8_t * src_spad = (uint8_t *) dma_queue_pop(dma_queue).dst;
+                    dma_queue_pop(dma_queue);
+                    uint8_t * cur_spad = (uint8_t *) dma_queue_pop(dma_queue).dst;

-                    // FARF(HIGH, "rope-compute %u: ir %u i1 %u i2 %u i3 %u src-spad %p cnr %u : usec %u", ith, ir, i1, i2, i3, src_spad, cnr,
-                    //         (unsigned) HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - rctx->t_start));
-
-                    if (is_vision) {
-                        rope_vision_f32(rctx, dst_spad, src_spad, cnr, ne0, theta_cache);
-                    } else if (is_neox) {
-                        rope_neox_f32(rctx, dst_spad, src_spad, cnr, ne0, theta_cache);
+                    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, cur_ir);
+                    if (is_neox || is_vision) {
+                        rope_neox_f32_inplace(rctx, cur_spad, cnr, ne, theta_cache);
                    } else {
-                        rope_basic_f32(rctx, dst_spad, src_spad, cnr, ne0, theta_cache);
+                        rope_basic_f32_inplace(rctx, cur_spad, cnr, theta_cache);
                    }
+                    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, cur_ir);

-                    uint8_t * dst_addr = (uint8_t *) dst->data + i3 * nb3 + i2 * nb2 + i1 * nb1;
+                    uint8_t * dst_addr = (uint8_t *) dst->data + i3 * nb3 + i2 * nb2 + cur_i1 * nb1;
+                    dma_queue_push(dma_queue, dma_make_ptr(dst_addr, cur_spad),
+                        rctx->dst_row_stride, rctx->src0_row_size_aligned, rctx->dst_row_size, cnr);

-                    // Write only the row payload while striding the DDR dst
-                    dma_queue_push(dma_queue, dma_make_ptr(dst_addr, dst_spad),
-                        rctx->dst_row_stride, rctx->dst_row_size_aligned, rctx->dst_row_size, cnr);
+                    // Prefetch 2 blocks ahead into the slot just freed
+                    if ((cr + 2 * HTP_ROPE_SPAD_BLOCK) < nrows) {
+                        const uint32_t p_cr   = cr + 2 * HTP_ROPE_SPAD_BLOCK;
+                        const uint32_t pnr    = MIN(nrows - p_cr, HTP_ROPE_SPAD_BLOCK);
+                        const uint32_t p_slot = (cur_slot + p_cr / HTP_ROPE_SPAD_BLOCK) % HTP_ROPE_SPAD_NSLOTS;
+                        uint8_t * p_spad      = rope_spad_slot(src0_spad_base, p_slot, rctx->src0_row_size_aligned);
+                        const uint8_t * src_addr = (const uint8_t *) src0->data + i3 * nb03 + i2 * nb02 + (base_i1 + p_cr) * nb01;

-                    // Prefetch more rows (if any)
-                    if ((cr + HTP_ROPE_SPAD_NROWS) < nrows) {
-                        uint32_t pnr = MIN(nrows - (cr + HTP_ROPE_SPAD_NROWS), HTP_ROPE_SPAD_BLOCK);
-                        uint32_t pi1 = i1 + HTP_ROPE_SPAD_NROWS;
-                        uint32_t pir = ir + HTP_ROPE_SPAD_NROWS;
-
-                        const uint8_t * src_addr = (const uint8_t *) src0->data + i3 * nb03 + i2 * nb02 + pi1 * nb01;
-                        dma_queue_push(dma_queue, dma_make_ptr(src_spad, src_addr),
+                        dma_queue_push(dma_queue, dma_make_ptr(p_spad, src_addr),
                            rctx->src0_row_size_aligned, rctx->src0_row_stride, rctx->src0_row_size, pnr);
-
-                        // FARF(HIGH, "rope-prefetch %u: pr %u i1 %u i2 %u i3 %u src-spad %p src-addr %p pnr %u", ith, pir, pi1, i2, i3, src_spad, src_addr, pnr);
                    }
                }
+
+                const uint32_t n_chunks = (nrows + HTP_ROPE_SPAD_BLOCK - 1) / HTP_ROPE_SPAD_BLOCK;
+                cur_slot = (cur_slot + n_chunks) % HTP_ROPE_SPAD_NSLOTS;
+
+                ir += nrows;
+                i1 += nrows;
            }
        }
    }

 done:
    dma_queue_flush(dma_queue);
-    tt = HAP_perf_get_qtimer_count() - tt;

-    FARF(HIGH, "rope-f32: %d/%d: (%u:%u) usec %u\n", ith, nth, src0_start_row, src0_end_row, (unsigned) HAP_perf_qtimer_count_to_us(tt));
+    FARF(HIGH, "rope-f32: %d/%d: (%u:%u)\n", ith, nth, src0_start_row, src0_end_row);
 }

 static int execute_op_rope_f32(struct htp_ops_context * octx) {
@@ -624,8 +696,6 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
    const struct htp_tensor * src2 = octx->src[2];
    const struct htp_tensor * dst  = octx->dst;

-    const char * op_type = "rope-f32";
-
    switch (octx->op) {
        case HTP_OP_ROPE:
            break;
@@ -635,48 +705,23 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
            return HTP_STATUS_NO_SUPPORT;
    }

-    const uint32_t ne0 = dst->ne[0];
-    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
-    const uint32_t n_threads = MIN(octx->n_threads, src0_nrows);
+    const struct htp_rope_kernel_params * kparams = (const struct htp_rope_kernel_params *) octx->kernel_params;
+    assert(kparams->n_threads > 0);
+    assert(octx->ctx->vtcm_size >= kparams->vtcm_size);

+    const uint32_t ne0 = dst->ne[0];
    const size_t src0_row_size   = src0->ne[0] * sizeof(float);
    const size_t src0_row_stride = src0->nb[1];
    const size_t dst_row_size    = dst->ne[0] * sizeof(float);
    const size_t dst_row_stride  = dst->nb[1];

-    // Aligned row sizes for VTCM
-    const size_t src0_row_size_aligned    = hex_round_up(src0_row_size, VLEN);
-    const size_t dst_row_size_aligned     = hex_round_up(dst_row_stride, VLEN);
-    const size_t theta_cache_size_aligned = hex_round_up(src0->ne[0] * sizeof(float), 256);
-
-    // Calculate spad sizes per thread
-    size_t src0_spad_per_thread = theta_cache_size_aligned + HTP_ROPE_SPAD_NROWS * src0_row_size_aligned;
-    size_t dst_spad_per_thread  = HTP_ROPE_SPAD_NROWS * dst_row_size_aligned;
-    size_t spad_per_thread = src0_spad_per_thread + dst_spad_per_thread;
-
-    // Check if we fit in VTCM
-    size_t total_vtcm_needed = spad_per_thread * n_threads;
-    if (octx->ctx->vtcm_size < total_vtcm_needed) {
-        FARF(ERROR, "%s : current VTCM reservation %zu is too small, needed %zu\n", op_type, octx->ctx->vtcm_size, total_vtcm_needed);
-        return HTP_STATUS_VTCM_TOO_SMALL;
-    }
-
-    octx->src0_spad.size_per_thread = src0_spad_per_thread;
-    octx->dst_spad.size_per_thread  = dst_spad_per_thread;
-    octx->src0_spad.size = n_threads * src0_spad_per_thread;
-    octx->dst_spad.size  = n_threads * dst_spad_per_thread;
-    octx->src1_spad.size = 0;
-
-    octx->src0_spad.data = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
-    octx->src1_spad.data = NULL;                                        octx->src1_spad.src = NULL;
-    octx->dst_spad.data  = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src  = NULL;
-
    struct htp_rope_context rctx;
    memset(&rctx, 0, sizeof(struct htp_rope_context));

-    rctx.t_start = HAP_perf_get_qtimer_count();
-
-    rctx.octx = octx;
+    rctx.octx                  = octx;
+    rctx.vtcm_base             = (uint8_t *) octx->ctx->vtcm_base;
+    rctx.spad_per_thread       = kparams->spad_per_thread;
+    rctx.theta_cache_offset    = kparams->theta_cache_offset;

    const int32_t * op_params = &octx->op_params[0];
    rctx.n_dims     = ((const int32_t *) op_params)[1];
@@ -693,31 +738,29 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
    memcpy(&rctx.sections,    (int32_t *) op_params + 11, sizeof(int) * 4);

    rctx.theta_scale = powf(rctx.freq_base, -2.0f / rctx.n_dims);
+    rctx.theta_powers[0] = 1.0f;
+    for (int j = 1; j < 32; j++) {
+        rctx.theta_powers[j] = rctx.theta_powers[j - 1] * rctx.theta_scale;
+    }
+    rctx.theta_scale_32 = rctx.theta_powers[31] * rctx.theta_scale;

    rope_corr_dims(rctx.n_dims, rctx.n_ctx_orig, rctx.freq_base, rctx.beta_fast, rctx.beta_slow, rctx.corr_dims);

-    rctx.src0_row_size   = src0_row_size;
-    rctx.src0_row_stride = src0_row_stride;
-    rctx.dst_row_size    = dst_row_size;
-    rctx.dst_row_stride  = dst_row_stride;
-    rctx.src0_row_size_aligned = src0_row_size_aligned;
-    rctx.dst_row_size_aligned  = dst_row_size_aligned;
-    rctx.theta_cache_offset    = theta_cache_size_aligned;
+    rctx.src0_row_size         = src0_row_size;
+    rctx.src0_row_stride       = src0_row_stride;
+    rctx.dst_row_size          = dst_row_size;
+    rctx.dst_row_stride        = dst_row_stride;
+    rctx.src0_row_size_aligned = kparams->src0_row_size_aligned;

-    rctx.src0_nrows = src0_nrows;
-    rctx.src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads;
-
-    if (src0_nrows > 0) {
-        rctx.div_ne2_ne1 = init_fastdiv_values(dst->ne[2] * dst->ne[1]);
-        rctx.div_ne1     = init_fastdiv_values(dst->ne[1]);
-    }
+    rctx.src0_nrows            = kparams->src0_nrows;
+    rctx.src0_nrows_per_thread = kparams->src0_nrows_per_thread;
+    rctx.div_ne2_ne1           = kparams->div_ne2_ne1;
+    rctx.div_ne1               = kparams->div_ne1;

    FARF(HIGH, "rope-f32 n-rows %u n-dims %d ne0 %u ext-factor %.6f theta-scale %.6f attn-factor %.6f\n", rctx.src0_nrows, rctx.n_dims, ne0,
         rctx.ext_factor, rctx.theta_scale, rctx.attn_factor);

-    if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
-        worker_pool_run_func(octx->ctx->worker_pool, rope_job_f32, &rctx, n_threads);
-    }
+    work_queue_run(octx->ctx->work_queue, rope_job_f32, &rctx, kparams->n_threads);

    return err;
 }
@@ -0,0 +1,56 @@
+#ifndef HTP_ROPE_OPS_H
+#define HTP_ROPE_OPS_H
+
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+
+#define HTP_ROPE_SPAD_BLOCK  8
+#define HTP_ROPE_SPAD_NSLOTS 4
+#define HTP_ROPE_SPAD_NROWS  (HTP_ROPE_SPAD_BLOCK * HTP_ROPE_SPAD_NSLOTS)
+
+struct htp_rope_kernel_params {
+    uint32_t n_threads;
+    uint32_t src0_nrows;
+    uint32_t src0_nrows_per_thread;
+    uint32_t vtcm_size;
+    uint32_t spad_per_thread;
+    uint32_t theta_cache_offset;
+    uint32_t src0_row_size_aligned;
+
+    struct fastdiv_values div_ne2_ne1;
+    struct fastdiv_values div_ne1;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_rope_kernel_params) <= 128, "htp_rope_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_rope_kernel_params) <= 128, "htp_rope_kernel_params is too large for kernel_params blob");
+#endif
+
+struct htp_rope_vtcm_layout {
+    size_t total_bytes;
+    size_t bytes_per_thread;
+    size_t theta_cache_size_aligned;
+    size_t src0_row_size_aligned;
+};
+
+static inline void htp_rope_vtcm_layout_build(
+    struct htp_rope_vtcm_layout * layout,
+    uint32_t ne00,
+    uint32_t n_threads
+) {
+    const size_t src0_row_size            = ne00 * sizeof(float);
+    const size_t src0_row_size_aligned    = hex_round_up((uint32_t) src0_row_size, 128);
+    const size_t theta_cache_size_aligned = hex_round_up((uint32_t) src0_row_size, 256);
+
+    layout->src0_row_size_aligned    = src0_row_size_aligned;
+    layout->theta_cache_size_aligned = theta_cache_size_aligned;
+    layout->bytes_per_thread         = theta_cache_size_aligned + HTP_ROPE_SPAD_NROWS * src0_row_size_aligned;
+    layout->total_bytes              = layout->bytes_per_thread * n_threads;
+}
+
+static inline uint8_t * rope_spad_slot(uint8_t * base, uint32_t slot, size_t row_size_aligned) {
+    return base + (slot * HTP_ROPE_SPAD_BLOCK) * row_size_aligned;
+}
+
+#endif // HTP_ROPE_OPS_H
@@ -70,17 +70,8 @@ list(APPEND GGML_SOURCES_ROCM ${SRCS})
 file(GLOB   SRCS "../ggml-cuda/template-instances/mmf*.cu")
 list(APPEND GGML_SOURCES_ROCM ${SRCS})

-if (GGML_CUDA_FA_ALL_QUANTS)
-    file(GLOB   SRCS "../ggml-cuda/template-instances/fattn-vec*.cu")
-    list(APPEND GGML_SOURCES_ROCM ${SRCS})
-    add_compile_definitions(GGML_CUDA_FA_ALL_QUANTS)
-else()
-    list(APPEND GGML_SOURCES_ROCM
-        ../ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu
-        ../ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu
-        ../ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu
-        ../ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu)
-endif()
+ggml_cuda_fattn_vec_instances(${CMAKE_CURRENT_SOURCE_DIR}/../ggml-cuda SRCS)
+list(APPEND GGML_SOURCES_ROCM ${SRCS})

 ggml_add_backend_library(ggml-hip
                         ${GGML_HEADERS_ROCM}
@@ -10,6 +10,7 @@ ggml_add_backend_library(ggml-metal
                         ggml-metal-device.cpp
                         ggml-metal-common.cpp
                         ggml-metal-context.m
+                         ggml-metal-fusion.cpp
                         ggml-metal-ops.cpp
                         ggml-metal-tuning.cpp
                        )
@@ -1,4 +1,5 @@
 #include "ggml-metal-common.h"
+#include "ggml-metal-fusion.h"

 #include "ggml.h"
 #include "ggml-impl.h"
@@ -390,59 +391,31 @@ static std::vector<int> ggml_metal_graph_optimize_reorder(const std::vector<node
 }

 void ggml_graph_optimize(ggml_cgraph * gf) {
-    constexpr int MAX_FUSE = 16;
-
    const int n = gf->n_nodes;

-    enum ggml_op ops[MAX_FUSE];
-
    std::vector<node_info> nodes;
    nodes.reserve(gf->n_nodes);

    // fuse nodes:
    // we don't want to make reorders that break fusing, so we first pack all fusable tensors
    //   and perform the reorder over the fused nodes. after the reorder is done, we unfuse
+    //
+    // the fusable sequences are declared in the fusion table (ggml-metal-fuse.cpp), so the
+    // packing here is driven by the same patterns that the op encoders will later use
    for (int i = 0; i < n; i++) {
        node_info node = {
            /*.node =*/ gf->nodes[i],
            /*.fused =*/ {},
        };

-        // fuse only ops that start with these operations
-        // can be expanded when needed
-        if (node.op() == GGML_OP_ADD ||
-            node.op() == GGML_OP_NORM ||
-            node.op() == GGML_OP_RMS_NORM) {
-            ops[0] = node.op();
+        const int f = ggml_metal_fusion_max(gf, i);

-            int f = i + 1;
-            while (f < n && f < i + MAX_FUSE) {
-                // conservatively allow fusing only these ops
-                // can be expanded when needed
-                if (gf->nodes[f]->op != GGML_OP_ADD &&
-                    gf->nodes[f]->op != GGML_OP_MUL &&
-                    gf->nodes[f]->op != GGML_OP_NORM &&
-                    gf->nodes[f]->op != GGML_OP_RMS_NORM) {
-                    break;
-                }
-                ops[f - i] = gf->nodes[f]->op;
-                f++;
-            }
+        // add the fused tensors into the node info so we can unfuse them later
+        for (int k = 1; k < f; k++) {
+            ++i;

-            f -= i;
-            for (; f > 1; f--) {
-                if (ggml_can_fuse(gf, i, ops, f)) {
-                    break;
-                }
-            }
-
-            // add the fused tensors into the node info so we can unfuse them later
-            for (int k = 1; k < f; k++) {
-                ++i;
-
-                // the .dst() becomes the last fused tensor
-                node.add_fused(gf->nodes[i]);
-            }
+            // the .dst() becomes the last fused tensor
+            node.add_fused(gf->nodes[i]);
        }

        nodes.push_back(std::move(node));
@@ -33,6 +33,7 @@ ggml_metal_event_t ggml_metal_get_ev_cpy(ggml_metal_t ctx);

 void ggml_metal_set_n_cb            (ggml_metal_t ctx, int n_cb);
 void ggml_metal_set_abort_callback  (ggml_metal_t ctx, ggml_abort_callback abort_callback, void * user_data);
+
 bool ggml_metal_supports_family     (ggml_metal_t ctx, int family);
 void ggml_metal_capture_next_compute(ggml_metal_t ctx);

@@ -6,6 +6,7 @@
 #import "ggml-metal-impl.h"
 #import "ggml-metal-common.h"
 #import "ggml-metal-ops.h"
+#import "ggml-metal-fusion.h"

 #import <Foundation/Foundation.h>

@@ -36,15 +37,12 @@ struct ggml_metal {
    // additional, inference-time compiled pipelines
    ggml_metal_pipelines_t pipelines_ext;

-    bool use_fusion;
    bool use_concurrency;
    bool use_graph_optimize;

    int debug_graph;
-    int debug_fusion;

-    // how many times a given op was fused
-    uint64_t fuse_cnt[GGML_OP_COUNT];
+    struct ggml_metal_fusion_info * finfo;

    // capture state
    int capture_compute;
@@ -139,7 +137,6 @@ ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) {

        res->d_queue = dispatch_queue_create("ggml-metal", DISPATCH_QUEUE_CONCURRENT);

-        res->use_fusion      = getenv("GGML_METAL_FUSION_DISABLE") == nil;
        res->use_concurrency = getenv("GGML_METAL_CONCURRENCY_DISABLE") == nil;

        {
@@ -147,20 +144,19 @@ ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) {
            res->debug_graph = val ? atoi(val) : 0;
        }

-        {
-            const char * val = getenv("GGML_METAL_FUSION_DEBUG");
-            res->debug_fusion = val ? atoi(val) : 0;
-        }
-
        res->use_graph_optimize = true;

        if (getenv("GGML_METAL_GRAPH_OPTIMIZE_DISABLE") != NULL) {
            res->use_graph_optimize = false;
        }

-        memset(res->fuse_cnt, 0, sizeof(res->fuse_cnt));
+        res->finfo = ggml_metal_device_get_fusion_info(dev);
+        if (ggml_metal_fusion_info_stats(res->finfo)) {
+            ggml_metal_fusion_info_labels_init(res->finfo);
+            res->n_cb = 0;
+        }

-        GGML_LOG_INFO("%s: use fusion         = %s\n", __func__, res->use_fusion         ? "true" : "false");
+        GGML_LOG_INFO("%s: use fusion         = %s\n", __func__, ggml_metal_fusion_info_enabled(res->finfo) ? "true" : "false");
        GGML_LOG_INFO("%s: use concurrency    = %s\n", __func__, res->use_concurrency    ? "true" : "false");
        GGML_LOG_INFO("%s: use graph optimize = %s\n", __func__, res->use_graph_optimize ? "true" : "false");

@@ -222,15 +218,18 @@ void ggml_metal_free(ggml_metal_t ctx) {
        ctx->pipelines_ext = nil;
    }

-    if (ctx->debug_fusion > 0) {
+    if (ggml_metal_fusion_info_debug(ctx->finfo) > 0) {
        GGML_LOG_DEBUG("%s: fusion stats:\n", __func__);
-        for (int i = 0; i < GGML_OP_COUNT; i++) {
-            if (ctx->fuse_cnt[i] == 0) {
+
+        const int n_fusions = ggml_metal_fusion_info_n_fusions(ctx->finfo);
+        for (int i = 0; i < n_fusions; i++) {
+            const uint64_t count = ggml_metal_fusion_info_count(ctx->finfo, i);
+            if (count == 0) {
                continue;
            }

            // note: cannot use ggml_log here
-            GGML_LOG_DEBUG("%s: - %s: %" PRIu64 "\n", __func__, ggml_op_name((enum ggml_op) i), ctx->fuse_cnt[i]);
+            GGML_LOG_DEBUG("%s: - %s: %" PRIu64 "\n", __func__, ggml_metal_fusion_info_label(ctx->finfo, i), count);
        }
    }

@@ -481,10 +480,17 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
    @autoreleasepool {
        ctx->gf = gf;

-        ctx->n_nodes_0 = MIN(n_main, gf->n_nodes);
-        ctx->n_nodes_1 = gf->n_nodes - ctx->n_nodes_0;
+        if (ctx->n_cb == 0) {
+            // single-threaded encoding: the whole graph is encoded by one command buffer
+            ctx->n_nodes_0      = gf->n_nodes;
+            ctx->n_nodes_1      = 0;
+            ctx->n_nodes_per_cb = 0;
+        } else {
+            ctx->n_nodes_0      = MIN(n_main, gf->n_nodes);
+            ctx->n_nodes_1      = gf->n_nodes - ctx->n_nodes_0;

-        ctx->n_nodes_per_cb = (ctx->n_nodes_1 + ctx->n_cb - 1) / ctx->n_cb;
+            ctx->n_nodes_per_cb = (ctx->n_nodes_1 + ctx->n_cb - 1) / ctx->n_cb;
+        }

        if (ctx->capture_compute >= 0) {
            ctx->capture_compute--;
@@ -682,6 +688,12 @@ ggml_metal_event_t ggml_metal_get_ev_cpy(ggml_metal_t ctx) {
 }

 void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) {
+    // when fusion stats are collected the graph must be encoded by a single thread so the
+    // counters are race-free; override whatever the caller requested
+    if (ggml_metal_fusion_info_stats(ctx->finfo)) {
+        n_cb = 0;
+    }
+
    if (ctx->n_cb != n_cb) {
        ctx->n_cb = MIN(n_cb, GGML_METAL_MAX_COMMAND_BUFFERS);

@@ -717,13 +729,12 @@ void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) {
            ctx->dev,
            cmd_buf,
            ctx->gf,
+            ctx->finfo,
            idx_start,
            idx_end,
-            ctx->use_fusion,
            ctx->use_concurrency,
            ctx->capture_compute,
-            ctx->debug_graph,
-            ctx->debug_fusion);
+            ctx->debug_graph);

        for (int idx = 0; idx < ggml_metal_op_n_nodes(ctx_op); ++idx) {
            const int res = ggml_metal_op_encode(ctx_op, idx);
@@ -932,12 +932,24 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
                nsg = N_SG_IQ2_XXS;
                nr0 = N_R0_IQ2_XXS;
                smem = 256*8+128;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ2_XXS_SPLIT;
+                    split = true;
+                }
            } break;
        case GGML_TYPE_IQ2_XS:
            {
                nsg = N_SG_IQ2_XS;
                nr0 = N_R0_IQ2_XS;
                smem = 512*8+128;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ2_XS_SPLIT;
+                    split = true;
+                }
            } break;
        case GGML_TYPE_IQ3_XXS:
            {
@@ -957,21 +969,45 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
                nsg = N_SG_IQ3_S;
                nr0 = N_R0_IQ3_S;
                smem = 512*4;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ3_S_SPLIT;
+                    split = true;
+                }
            } break;
        case GGML_TYPE_IQ2_S:
            {
                nsg = N_SG_IQ2_S;
                nr0 = N_R0_IQ2_S;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ2_S_SPLIT;
+                    split = true;
+                }
            } break;
        case GGML_TYPE_IQ1_S:
            {
                nsg = N_SG_IQ1_S;
                nr0 = N_R0_IQ1_S;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ1_S_SPLIT;
+                    split = true;
+                }
            } break;
        case GGML_TYPE_IQ1_M:
            {
                nsg = N_SG_IQ1_M;
                nr0 = N_R0_IQ1_M;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ1_M_SPLIT;
+                    split = true;
+                }
            } break;
        case GGML_TYPE_IQ4_NL:
            {
@@ -1177,12 +1213,24 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m
                nsg = N_SG_IQ2_XXS;
                nr0 = N_R0_IQ2_XXS;
                smem = 256*8+128;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ2_XXS_SPLIT;
+                    split = true;
+                }
            } break;
        case GGML_TYPE_IQ2_XS:
            {
                nsg = N_SG_IQ2_XS;
                nr0 = N_R0_IQ2_XS;
                smem = 512*8+128;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ2_XS_SPLIT;
+                    split = true;
+                }
            } break;
        case GGML_TYPE_IQ3_XXS:
            {
@@ -1202,21 +1250,45 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m
                nsg = N_SG_IQ3_S;
                nr0 = N_R0_IQ3_S;
                smem = 512*4;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ3_S_SPLIT;
+                    split = true;
+                }
            } break;
        case GGML_TYPE_IQ2_S:
            {
                nsg = N_SG_IQ2_S;
                nr0 = N_R0_IQ2_S;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ2_S_SPLIT;
+                    split = true;
+                }
            } break;
        case GGML_TYPE_IQ1_S:
            {
                nsg = N_SG_IQ1_S;
                nr0 = N_R0_IQ1_S;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ1_S_SPLIT;
+                    split = true;
+                }
            } break;
        case GGML_TYPE_IQ1_M:
            {
                nsg = N_SG_IQ1_M;
                nr0 = N_R0_IQ1_M;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ1_M_SPLIT;
+                    split = true;
+                }
            } break;
        case GGML_TYPE_IQ4_NL:
            {
@@ -325,6 +325,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te

 const struct ggml_metal_device_props * ggml_metal_device_get_props(ggml_metal_device_t dev);

+struct ggml_metal_fusion_info;
+
+// the device-owned fusion debugging context (NULL unless fusion debugging is enabled)
+struct ggml_metal_fusion_info * ggml_metal_device_get_fusion_info(ggml_metal_device_t dev);
+
 //
 // device buffers
 //
@@ -1,4 +1,5 @@
 #import "ggml-metal-device.h"
+#import "ggml-metal-fusion.h"

 #import "ggml-impl.h"
 #import "ggml-backend-impl.h"
@@ -896,6 +897,9 @@ struct ggml_metal_device {

    struct ggml_metal_device_props props;

+    // shared fusion debugging context
+    struct ggml_metal_fusion_info * finfo;
+
    // virtual address for GPU memory allocations
    atomic_uintptr_t addr_virt;
 };
@@ -1274,6 +1278,13 @@ ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) {
                    dev->props.max_working_set_size   = dev->mtl_device.maxBufferLength;
                }

+                {
+                    const char * val = getenv("GGML_METAL_FUSION_DEBUG");
+                    dev->finfo = ggml_metal_fusion_info_init(
+                            getenv("GGML_METAL_FUSION_DISABLE") == nil,
+                            val ? atoi(val) : 0);
+                }
+
                snprintf(dev->props.name, sizeof(dev->props.name), "%s%d", "MTL", device);
                const char * gpu_name = [[dev->mtl_device name] UTF8String];
                if (n_devices > 1) {
@@ -1348,6 +1359,8 @@ void ggml_metal_device_free(ggml_metal_device_t dev) {
    assert(dev != NULL);

    @autoreleasepool {
+        ggml_metal_fusion_info_free(dev->finfo);
+
        ggml_metal_rsets_free(dev->rsets);

        ggml_metal_library_free(dev->library);
@@ -1935,6 +1948,10 @@ static void ggml_metal_device_disable_tensor(ggml_metal_device_t dev) {
    dev->props.has_tensor = false;
 }

+struct ggml_metal_fusion_info * ggml_metal_device_get_fusion_info(ggml_metal_device_t dev) {
+    return dev->finfo;
+}
+
 //
 // device buffers
 //
@@ -0,0 +1,502 @@
+#include "ggml-metal-fusion.h"
+
+#include "ggml-backend-impl.h"
+#include "ggml-metal-device.h"
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+// ---- helpers -------------------------------------------------------------
+
+// true if two tensors live in the same Metal buffer
+static bool ggml_metal_fusion_same_buffer(const ggml_tensor * a, const ggml_tensor * b) {
+    if (!a || !b) {
+        return false;
+    }
+
+    ggml_backend_buffer_t ba = a->view_src ? a->view_src->buffer : a->buffer;
+    ggml_backend_buffer_t bb = b->view_src ? b->view_src->buffer : b->buffer;
+
+    ggml_metal_buffer_t ca = (ggml_metal_buffer_t) ba->context;
+    ggml_metal_buffer_t cb = (ggml_metal_buffer_t) bb->context;
+
+    return ggml_metal_buffer_get_id(ca, a).metal == ggml_metal_buffer_get_id(cb, b).metal;
+}
+
+// ---- pattern checks ------------------------------------------------------
+
+// NORM/RMS_NORM + MUL + ADD: the weight/bias of each fused step must match the norm input
+// width, be contiguous rows, and the fused outputs must stay F32
+static bool ggml_metal_fusion_check_norm(
+        const ggml_metal_fusion      * fusion,
+        const ggml_tensor * const    * nodes,
+              ggml_metal_fusion_mode   mode) {
+    GGML_UNUSED(mode);
+
+    GGML_ASSERT(fusion->n_ops >= 2);
+
+    for (int j = 1; j < fusion->n_ops; j++) {
+        // the fused MUL/ADD must read the previous node as src0
+        if (nodes[j]->src[0] != nodes[j - 1]) {
+            return false;
+        }
+
+        // the weight/bias must have the same row width as the norm input
+        if (nodes[j]->src[1]->ne[0] != nodes[0]->ne[0]) {
+            return false;
+        }
+
+        if (!ggml_is_contiguous_rows(nodes[j]->src[1])) {
+            return false;
+        }
+
+        if (nodes[j]->type != GGML_TYPE_F32) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+// ADD x N: each ADD reads the previous ADD as src0, and all addends must share layout
+// (and, in FULL mode, live in the same Metal buffer)
+static bool ggml_metal_fusion_check_add_chain(
+        const ggml_metal_fusion      * fusion,
+        const ggml_tensor * const    * nodes,
+              ggml_metal_fusion_mode   mode) {
+    GGML_ASSERT(fusion->n_ops >= 2);
+
+    for (int j = 1; j < fusion->n_ops; j++) {
+        if (nodes[j]->src[0] != nodes[j - 1]) {
+            return false;
+        }
+
+        if (!ggml_are_same_layout(nodes[j]->src[1], nodes[j - 1]->src[1])) {
+            return false;
+        }
+
+        if (mode == GGML_METAL_FUSION_FULL) {
+            if (!ggml_metal_fusion_same_buffer(nodes[j]->src[1], nodes[0]->src[1])) {
+                return false;
+            }
+        }
+    }
+
+    return true;
+}
+
+// GATED_DELTA_NET + CPY: the trailing cpy scatters the gdn state snapshots into the recurrent
+// cache, so the gdn kernel writes them straight to the cache and the cpy is elided.
+// mirrors ggml_metal_op_can_fuse_gdn_cache (PR #25788). the gdn output has other consumers (the
+// attn scores view), so unlike the other patterns this is not an elision chain: the structural
+// checks live entirely in this callback (unsafe = true).
+static bool ggml_metal_fusion_check_gdn_cache(
+        const ggml_metal_fusion      * fusion,
+        const ggml_tensor * const    * nodes,
+              ggml_metal_fusion_mode    mode) {
+    GGML_UNUSED(fusion);
+
+    const ggml_tensor * gdn = nodes[0];
+    const ggml_tensor * cpy = nodes[1];
+
+    // the kernel skips the snapshot tail, so the gdn output must not be a graph output
+    if (gdn->type != GGML_TYPE_F32 || (gdn->flags & GGML_TENSOR_FLAG_OUTPUT)) {
+        return false;
+    }
+
+    if (cpy->op != GGML_OP_CPY || (cpy->flags & GGML_TENSOR_FLAG_OUTPUT)) {
+        return false;
+    }
+
+    const int64_t S_v      = gdn->src[2]->ne[0];
+    const int64_t H        = gdn->src[2]->ne[1];
+    const int64_t n_tokens = gdn->src[2]->ne[2];
+    const int64_t n_seqs   = gdn->src[2]->ne[3];
+    const int64_t K        = ggml_get_op_params_i32(gdn, 0);
+    const size_t  tail_off = ggml_row_size(GGML_TYPE_F32, S_v * H * n_tokens * n_seqs);
+
+    const int64_t D         = S_v * S_v * H;
+    const int64_t n_written = std::min<int64_t>(n_tokens, K);
+
+    const ggml_tensor * src = cpy->src[0]; // gdn snapshot tail view
+    const ggml_tensor * dst = cpy->src[1]; // cache view
+
+    // src must be this gdn's snapshot tail (contiguous, at the tail offset)
+    if (src->op != GGML_OP_VIEW || src->view_src != gdn ||
+        src->view_offs != tail_off || !ggml_is_contiguous(src)) {
+        return false;
+    }
+
+    const int64_t expected_ne[GGML_MAX_DIMS] = { D, n_seqs, n_written, 1 };
+    if (dst->type != GGML_TYPE_F32 ||
+        !std::equal(expected_ne, expected_ne + GGML_MAX_DIMS, dst->ne) ||
+        dst->nb[0] != ggml_type_size(GGML_TYPE_F32) ||
+        dst->nb[1] != ggml_row_size(GGML_TYPE_F32, D)) {
+        return false;
+    }
+
+    if (mode == GGML_METAL_FUSION_FULL) {
+        // the cache must be allocated so the kernel can write straight to its buffer
+        if (dst->data == nullptr) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+// MUL + SIN + SQR + MUL + ADD (snake activation)
+static bool ggml_metal_fusion_check_snake(
+        const ggml_metal_fusion      * fusion,
+        const ggml_tensor * const    * nodes,
+              ggml_metal_fusion_mode   mode) {
+    GGML_UNUSED(fusion);
+    GGML_UNUSED(mode);
+
+    const ggml_tensor * mul0     = nodes[0];
+    const ggml_tensor * sin_node = nodes[1];
+    const ggml_tensor * sqr      = nodes[2];
+    const ggml_tensor * mul1     = nodes[3];
+    const ggml_tensor * add      = nodes[4];
+
+    // x carries the full activation shape, a is the broadcast operand
+    const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1];
+    const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0];
+
+    // mul1 reads sqr and inv_b in either operand order
+    const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0];
+
+    // closure check: the trailing add reads the same x as the leading mul
+    const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0];
+
+    // x is in the supported whitelist and every chain intermediate shares x's type.
+    // a and inv_b bind as device const float * in the kernel, so they stay F32.
+    const bool types_ok =
+        (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) &&
+        (a->type    == GGML_TYPE_F32) && (inv_b->type    == GGML_TYPE_F32) &&
+        (mul0->type == x->type)       && (sin_node->type == x->type) &&
+        (sqr->type  == x->type)       && (mul1->type     == x->type) &&
+        (add->type  == x->type);
+
+    // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D
+    const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1];
+    const bool dim_ok =
+        (x->ne[2]     == 1) && (x->ne[3]     == 1) &&
+        (add->ne[2]   == 1) && (add->ne[3]   == 1) &&
+        (a->ne[2]     == 1) && (a->ne[3]     == 1) &&
+        (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1);
+
+    // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous
+    const bool contig_ok =
+        ggml_is_contiguous(x) && ggml_is_contiguous(add) &&
+        ggml_is_contiguous(a) && ggml_is_contiguous(inv_b);
+
+    return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x;
+}
+
+// ---- patterns ------------------------------------------------------------
+
+static const ggml_op ops_norm_mul[]         = { GGML_OP_NORM, GGML_OP_MUL };
+static const ggml_op ops_norm_mul_add[]     = { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD };
+static const ggml_op ops_rms_norm_mul[]     = { GGML_OP_RMS_NORM, GGML_OP_MUL };
+static const ggml_op ops_rms_norm_mul_add[] = { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD };
+
+static const ggml_op ops_add_2[] = { GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_add_3[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_add_4[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_add_5[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_add_6[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_add_7[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_snake[] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD };
+
+static const ggml_op ops_gdn_cache[] = { GGML_OP_GATED_DELTA_NET, GGML_OP_CPY };
+
+static const ggml_metal_fusion ggml_metal_fusions[] = {
+    { GGML_METAL_FUSION_NORM_MUL,     ops_norm_mul,         2, false, ggml_metal_fusion_check_norm },
+    { GGML_METAL_FUSION_NORM_MUL_ADD, ops_norm_mul_add,     3, false, ggml_metal_fusion_check_norm },
+    { GGML_METAL_FUSION_NORM_MUL,     ops_rms_norm_mul,     2, false, ggml_metal_fusion_check_norm },
+    { GGML_METAL_FUSION_NORM_MUL_ADD, ops_rms_norm_mul_add, 3, false, ggml_metal_fusion_check_norm },
+    { GGML_METAL_FUSION_ADD_CHAIN,    ops_add_2,            2, false, ggml_metal_fusion_check_add_chain },
+    { GGML_METAL_FUSION_ADD_CHAIN,    ops_add_3,            3, false, ggml_metal_fusion_check_add_chain },
+    { GGML_METAL_FUSION_ADD_CHAIN,    ops_add_4,            4, false, ggml_metal_fusion_check_add_chain },
+    { GGML_METAL_FUSION_ADD_CHAIN,    ops_add_5,            5, false, ggml_metal_fusion_check_add_chain },
+    { GGML_METAL_FUSION_ADD_CHAIN,    ops_add_6,            6, false, ggml_metal_fusion_check_add_chain },
+    { GGML_METAL_FUSION_ADD_CHAIN,    ops_add_7,            7, false, ggml_metal_fusion_check_add_chain },
+    { GGML_METAL_FUSION_SNAKE,        ops_snake,            5, false, ggml_metal_fusion_check_snake },
+    { GGML_METAL_FUSION_GDN_CACHE,    ops_gdn_cache,        2, true,  ggml_metal_fusion_check_gdn_cache },
+};
+
+const ggml_metal_fusion * ggml_metal_fusion_all(int * n) {
+    *n = (int) sizeof(ggml_metal_fusions) / sizeof(ggml_metal_fusions[0]);
+
+    return ggml_metal_fusions;
+}
+
+// ---- shared fusion info ---------------------------------------------------
+
+static std::string ggml_metal_fusion_label(const ggml_metal_fusion * fusion) {
+    GGML_ASSERT(fusion != nullptr);
+
+    std::string label;
+    for (int j = 0; j < fusion->n_ops; j++) {
+        if (j > 0) {
+            label += '+';
+        }
+        label += ggml_op_name(fusion->ops[j]);
+    }
+    return label;
+}
+
+struct ggml_metal_fusion_info {
+    std::vector<std::string> labels;
+    std::vector<uint64_t>    counts;
+    bool enabled;
+    bool stats;
+    bool labels_set;
+    int  debug;
+};
+
+struct ggml_metal_fusion_info * ggml_metal_fusion_info_init(bool enabled, int debug) {
+    ggml_metal_fusion_info * finfo = new ggml_metal_fusion_info;
+    finfo->enabled    = enabled;
+    finfo->stats      = debug > 0;
+    finfo->labels_set = false;
+    finfo->debug      = debug;
+
+    if (finfo->stats) {
+        ggml_metal_fusion_info_labels_init(finfo);
+    }
+
+    return finfo;
+}
+
+void ggml_metal_fusion_info_free(struct ggml_metal_fusion_info * finfo) {
+    delete finfo;
+}
+
+bool ggml_metal_fusion_info_enabled(const struct ggml_metal_fusion_info * finfo) {
+    return finfo->enabled;
+}
+
+bool ggml_metal_fusion_info_stats(const struct ggml_metal_fusion_info * finfo) {
+    return finfo->stats;
+}
+
+int ggml_metal_fusion_info_debug(const struct ggml_metal_fusion_info * finfo) {
+    return finfo->debug;
+}
+
+int ggml_metal_fusion_info_n_fusions(const struct ggml_metal_fusion_info * finfo) {
+    return (int) finfo->labels.size();
+}
+
+const char * ggml_metal_fusion_info_label(const struct ggml_metal_fusion_info * finfo, int idx) {
+    GGML_ASSERT(idx >= 0 && idx < (int) finfo->labels.size());
+    return finfo->labels[idx].c_str();
+}
+
+uint64_t ggml_metal_fusion_info_count(const struct ggml_metal_fusion_info * finfo, int idx) {
+    GGML_ASSERT(idx >= 0 && idx < (int) finfo->counts.size());
+    return finfo->counts[idx];
+}
+
+void ggml_metal_fusion_info_count_fusion(struct ggml_metal_fusion_info * finfo, const struct ggml_metal_fusion * fusion) {
+    if (!finfo->stats || fusion == nullptr) {
+        return;
+    }
+
+    int n = 0;
+    const ggml_metal_fusion * all = ggml_metal_fusion_all(&n);
+
+    int idx = -1;
+    for (int i = 0; i < n; i++) {
+        if (&all[i] == fusion) {
+            idx = i;
+            break;
+        }
+    }
+
+    if (idx >= 0 && idx < (int) finfo->counts.size()) {
+        finfo->counts[idx]++;
+    }
+}
+
+void ggml_metal_fusion_info_set_enabled(struct ggml_metal_fusion_info * finfo, bool enabled) {
+    finfo->enabled = enabled;
+}
+
+void ggml_metal_fusion_info_labels_init(struct ggml_metal_fusion_info * finfo) {
+    if (finfo->labels_set) {
+        return;
+    }
+
+    int n = 0;
+    const ggml_metal_fusion * all = ggml_metal_fusion_all(&n);
+
+    finfo->labels.clear();
+    finfo->counts.assign(n, 0);
+    finfo->labels.reserve(n);
+
+    for (int i = 0; i < n; i++) {
+        finfo->labels.emplace_back(ggml_metal_fusion_label(&all[i]));
+    }
+
+    finfo->labels_set = true;
+}
+
+void ggml_metal_fusion_info_stats_init(struct ggml_metal_fusion_info * finfo) {
+    finfo->stats = true;
+    ggml_metal_fusion_info_labels_init(finfo);
+}
+
+void ggml_metal_fusion_info_stats_reset(struct ggml_metal_fusion_info * finfo) {
+    std::fill(finfo->counts.begin(), finfo->counts.end(), 0);
+}
+
+int ggml_metal_fusion_info_stats_get(const struct ggml_metal_fusion_info * finfo, const char ** labels, uint64_t * counts, int n) {
+    const int n_fusions = (int) finfo->labels.size();
+
+    if (labels == nullptr) {
+        return n_fusions;
+    }
+
+    const int n_fill = std::min(n, n_fusions);
+    for (int i = 0; i < n_fill; i++) {
+        labels[i] = finfo->labels[i].c_str();
+        if (counts != nullptr) {
+            counts[i] = finfo->counts[i];
+        }
+    }
+
+    return n_fill;
+}
+
+// ---- queries -------------------------------------------------------------
+
+// find the longest pattern matching the node sequence starting at idx
+// (idx is a position in node_idxs, which maps to graph node indices)
+const ggml_metal_fusion * ggml_metal_fusion_next(
+        const ggml_cgraph * gf,
+        const int * node_idxs,
+        int n_idxs,
+        int idx,
+        ggml_metal_fusion_mode mode,
+        int * n_out) {
+    int n = 0;
+    const ggml_metal_fusion * all = ggml_metal_fusion_all(&n);
+
+    const ggml_metal_fusion * res = nullptr;
+    int best = 1;
+
+    for (int i = 0; i < n; i++) {
+        const ggml_metal_fusion * fusion = &all[i];
+
+        // only look for a longer match than the current best
+        if (fusion->n_ops <= best) {
+            continue;
+        }
+        if (idx + fusion->n_ops > n_idxs) {
+            continue;
+        }
+
+        const ggml_tensor * nodes[GGML_METAL_FUSION_MAX];
+
+        // the op sequence must match exactly
+        bool ok = true;
+        for (int j = 0; j < fusion->n_ops; j++) {
+            nodes[j] = gf->nodes[node_idxs[idx + j]];
+            if (nodes[j]->op != fusion->ops[j]) {
+                ok = false;
+                break;
+            }
+        }
+        if (!ok) {
+            continue;
+        }
+
+        if (!fusion->unsafe) {
+            // common element-wise chain constraints: each node reads the previous one,
+            // and all nodes have the same shape
+            for (int j = 1; j < fusion->n_ops && ok; j++) {
+                if (nodes[j]->src[0] != nodes[j - 1] && nodes[j]->src[1] != nodes[j - 1]) {
+                    ok = false;
+                    break;
+                }
+                if (!ggml_are_same_shape(nodes[j], nodes[j - 1])) {
+                    ok = false;
+                    break;
+                }
+            }
+            if (!ok) {
+                continue;
+            }
+
+            // all current fusions are single-output elision chains, so the last node is the only output
+            // TODO: multi-output fusions: store pattern-relative offsets in the table and translate them here
+            int outputs_buf[1];
+            outputs_buf[0] = node_idxs[idx + fusion->n_ops - 1];
+
+            // structural subgraph checks (op sequence, elidable uses, view containment)
+            if (!ggml_can_fuse_subgraph_ext(gf, node_idxs + idx, fusion->n_ops, fusion->ops, outputs_buf, 1)) {
+                continue;
+            }
+        }
+
+        // pattern-specific checks (the sole validator for unsafe patterns)
+        if (fusion->check && !fusion->check(fusion, nodes, mode)) {
+            continue;
+        }
+
+        best = fusion->n_ops;
+        res = fusion;
+    }
+
+    *n_out = best;
+
+    return res;
+}
+
+// optimize phase: maximum number of nodes starting at idx (a raw sequential graph index) that
+// could be fused, chaining patterns back-to-back. matching runs on the same filtered (view
+// transparent) node sequence that the compute phase uses, so the returned count is the raw index
+// span from idx to the last matched node (intermediate views are packed along).
+int ggml_metal_fusion_max(const ggml_cgraph * gf, int idx) {
+    // an empty/view node cannot start a pattern - pack it alone
+    if (ggml_op_is_empty(gf->nodes[idx]->op) || ggml_is_empty(gf->nodes[idx])) {
+        return 1;
+    }
+
+    // collect the non-empty node indices starting at idx
+    int idxs[GGML_METAL_FUSION_MAX];
+    int n_idxs = 0;
+    for (int i = idx; i < gf->n_nodes && n_idxs < GGML_METAL_FUSION_MAX; i++) {
+        if (!ggml_op_is_empty(gf->nodes[i]->op) && !ggml_is_empty(gf->nodes[i])) {
+            idxs[n_idxs++] = i;
+        }
+    }
+    if (n_idxs == 0) {
+        return 1;
+    }
+
+    int total = 0;
+    int i_f = 0;
+
+    while (i_f < n_idxs && total < GGML_METAL_FUSION_MAX) {
+        int len = 1;
+        const ggml_metal_fusion * fusion = ggml_metal_fusion_next(gf, idxs, n_idxs, i_f, GGML_METAL_FUSION_STRUCTURAL, &len);
+        if (!fusion || total + len > GGML_METAL_FUSION_MAX) {
+            break;
+        }
+
+        total += len;
+        i_f += len;
+    }
+
+    if (i_f == 0) {
+        return 1;
+    }
+
+    // map the matched non-empty nodes back to the raw index span (views are included)
+    return std::min(GGML_METAL_FUSION_MAX, idxs[i_f - 1] - idx + 1);
+}
@@ -0,0 +1,104 @@
+// single source of truth for the fusions supported by the Metal backend
+//
+// every fusable subgraph is declared exactly once as a ggml_metal_fusion entry in
+// the table in ggml-metal-fusion.cpp. both the graph optimizer (ggml_metal_fusion_max)
+// and the op encoders (ggml_metal_fusion_next) consult this same table, so the two
+// phases can never disagree about what can be fused.
+
+#pragma once
+
+#include "ggml-impl.h"
+
+#include <stdbool.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// the maximum number of nodes that can be fused in a single kernel
+// (also the maximum length of a packed fusion group during graph optimization)
+#define GGML_METAL_FUSION_MAX 16
+
+typedef enum ggml_metal_fusion_mode {
+    // structural checks only; used by the graph optimizer, at which point the graph
+    // tensors are not allocated yet, so buffer placement cannot be verified
+    GGML_METAL_FUSION_STRUCTURAL = 0,
+    // full checks, including buffer placement; used by the op encoders
+    GGML_METAL_FUSION_FULL,
+} ggml_metal_fusion_mode;
+
+// identifier of each fusion pattern so the op encoders know which kernel to use
+typedef enum ggml_metal_fusion_id {
+    GGML_METAL_FUSION_NONE = 0,
+    GGML_METAL_FUSION_NORM_MUL,     // NORM/RMS_NORM + MUL
+    GGML_METAL_FUSION_NORM_MUL_ADD, // NORM/RMS_NORM + MUL + ADD
+    GGML_METAL_FUSION_ADD_CHAIN,    // ADD x N (N in [2, 7])
+    GGML_METAL_FUSION_SNAKE,        // MUL + SIN + SQR + MUL + ADD
+    GGML_METAL_FUSION_GDN_CACHE,    // GATED_DELTA_NET + CPY (write snapshots into the recurrent cache)
+} ggml_metal_fusion_id;
+
+struct ggml_metal_fusion {
+    ggml_metal_fusion_id id;
+
+    const enum ggml_op * ops;        // op sequence (fixed length)
+    int                  n_ops;      // number of ops
+
+    // if unsafe: the generic chain/shape + ggml_can_fuse_subgraph checks are skipped and the
+    // check callback below is the sole validator (used for patterns that are not elision chains,
+    // e.g. the gdn + cache-cpy write-through fusion)
+    bool unsafe;
+
+    // extra backend constraints on top of ggml_can_fuse_subgraph
+    // nodes[j] is the j-th node of the pattern
+    bool (*check)(const struct ggml_metal_fusion   * fusion,
+                  const struct ggml_tensor * const * nodes,
+                        ggml_metal_fusion_mode       mode);
+};
+
+typedef struct ggml_metal_fusion ggml_metal_fusion;
+
+// the single table of all fusions supported by the Metal backend
+const ggml_metal_fusion * ggml_metal_fusion_all(int * n);
+
+// ---- shared fusion info ---------------------------------------------------
+
+// shared fusion debugging context, owned by the device; newly created backend contexts for that
+// device register with it so the fusion counters are race-free and accumulate across contexts.
+struct ggml_metal_fusion_info; // defined in ggml-metal-fusion.cpp
+
+struct ggml_metal_fusion_info * ggml_metal_fusion_info_init(bool enabled, int debug);
+void ggml_metal_fusion_info_free(struct ggml_metal_fusion_info * finfo);
+
+bool ggml_metal_fusion_info_enabled(const struct ggml_metal_fusion_info * finfo);
+bool ggml_metal_fusion_info_stats  (const struct ggml_metal_fusion_info * finfo);
+int  ggml_metal_fusion_info_debug  (const struct ggml_metal_fusion_info * finfo);
+
+int          ggml_metal_fusion_info_n_fusions(const struct ggml_metal_fusion_info * finfo);
+const char * ggml_metal_fusion_info_label    (const struct ggml_metal_fusion_info * finfo, int idx);
+uint64_t     ggml_metal_fusion_info_count    (const struct ggml_metal_fusion_info * finfo, int idx);
+
+void ggml_metal_fusion_info_count_fusion(struct ggml_metal_fusion_info * finfo, const struct ggml_metal_fusion * fusion);
+void ggml_metal_fusion_info_set_enabled (struct ggml_metal_fusion_info * finfo, bool enabled);
+
+void ggml_metal_fusion_info_stats_init (      struct ggml_metal_fusion_info * finfo);
+void ggml_metal_fusion_info_stats_reset(      struct ggml_metal_fusion_info * finfo);
+int  ggml_metal_fusion_info_stats_get  (const struct ggml_metal_fusion_info * finfo, const char ** labels, uint64_t * counts, int n);
+void ggml_metal_fusion_info_labels_init(      struct ggml_metal_fusion_info * finfo);
+
+// compute phase: longest fusion starting at idx (a position in node_idxs) that matches in `mode`.
+// returns the matching pattern (nullptr if no fusion) and sets *n_out to the number of nodes consumed.
+const ggml_metal_fusion * ggml_metal_fusion_next(
+        const struct ggml_cgraph * gf,
+        const int * node_idxs,
+        int n_idxs,
+        int idx,
+        ggml_metal_fusion_mode mode,
+        int * n_out);
+
+// optimize phase: maximum number of nodes starting at idx (a raw sequential graph index) that
+// could be fused, chaining patterns back-to-back. returns at least 1.
+int ggml_metal_fusion_max(const struct ggml_cgraph * gf, int idx);
+
+#ifdef __cplusplus
+}
+#endif
@@ -62,18 +62,23 @@

 #define N_R0_IQ1_S 4
 #define N_SG_IQ1_S 2
+#define N_R0_IQ1_S_SPLIT 8

 #define N_R0_IQ1_M 4
 #define N_SG_IQ1_M 2
+#define N_R0_IQ1_M_SPLIT 8

 #define N_R0_IQ2_XXS 4
 #define N_SG_IQ2_XXS 2
+#define N_R0_IQ2_XXS_SPLIT 8

 #define N_R0_IQ2_XS 4
 #define N_SG_IQ2_XS 2
+#define N_R0_IQ2_XS_SPLIT 8

 #define N_R0_IQ2_S 4
 #define N_SG_IQ2_S 2
+#define N_R0_IQ2_S_SPLIT 8

 #define N_R0_IQ3_XXS 4
 #define N_SG_IQ3_XXS 2
@@ -81,6 +86,7 @@

 #define N_R0_IQ3_S 4
 #define N_SG_IQ3_S 2
+#define N_R0_IQ3_S_SPLIT 8

 #define N_R0_IQ4_NL 2
 #define N_SG_IQ4_NL 2
@@ -979,6 +985,7 @@ typedef struct {
    uint64_t nb1;
    uint64_t nb2;
    uint64_t nb3;
+    uint64_t nb_out; // 0 => snapshots are appended after the attn scores (unfused)
 } ggml_metal_kargs_gated_delta_net;

 typedef struct {
@@ -7,6 +7,7 @@
 #include "ggml-metal-impl.h"
 #include "ggml-metal-common.h"
 #include "ggml-metal-device.h"
+#include "ggml-metal-fusion.h"
 #include "ggml-metal-tuning.h"

 #include <cassert>
@@ -31,24 +32,22 @@ struct ggml_metal_op {
        ggml_metal_device_t dev,
        ggml_metal_cmd_buf_t cmd_buf,
        ggml_cgraph * gf,
+        ggml_metal_fusion_info * finfo,
        int  idx_start,
        int  idx_end,
-        bool use_fusion,
        bool use_concurrency,
        bool use_capture,
-        int  debug_graph,
-        int  debug_fusion) {
+        int  debug_graph) {
        this->dev             = dev;
        this->lib             = ggml_metal_device_get_library(dev);
        this->enc             = ggml_metal_encoder_init(cmd_buf, use_concurrency);
        this->mem_ranges      = ggml_mem_ranges_init(debug_graph);
+        this->finfo           = finfo;
        this->idx_start       = idx_start;
        this->idx_end         = idx_end;
-        this->use_fusion      = use_fusion;
        this->use_concurrency = use_concurrency;
        this->use_capture     = use_capture;
        this->debug_graph     = debug_graph;
-        this->debug_fusion    = debug_fusion;
        this->gf              = gf;

        idxs.reserve(gf->n_nodes);
@@ -78,15 +77,24 @@ struct ggml_metal_op {
        return ggml_graph_node(gf, idxs[i]);
    }

-    bool can_fuse(int i0, const ggml_op * ops, int n_ops) const {
-        assert(use_fusion);
+    // consult the fusion table for the longest pattern starting at i0
+    // returns the matching pattern (nullptr if no fusion) and sets *n_out to the number of nodes
+    const ggml_metal_fusion * can_fuse(int i0, enum ggml_metal_fusion_mode mode, int * n_out) const {
+        assert(use_fusion());
        assert(i0 >= 0 && i0 < n_nodes());

-        if (i0 + n_ops > n_nodes()) {
-            return false;
-        }
+        return ggml_metal_fusion_next(gf, idxs.data(), (int) idxs.size(), i0, mode, n_out);
+    }

-        return ggml_can_fuse_ext(gf, idxs.data() + i0, ops, n_ops);
+    // whether to attempt fusion; the toggle lives in the shared fusion debugging context owned
+    // by the device (initialized from GGML_METAL_FUSION_DISABLE, overridable by the test)
+    bool use_fusion() const {
+        return ggml_metal_fusion_info_enabled(finfo);
+    }
+
+    // record that a fusion fired, indexed by the matching table entry
+    void count_fusions(const ggml_metal_fusion * fusion) const {
+        ggml_metal_fusion_info_count_fusion(finfo, fusion);
    }

    ggml_metal_device_t  dev;
@@ -94,12 +102,13 @@ struct ggml_metal_op {
    ggml_metal_encoder_t enc;
    ggml_mem_ranges_t    mem_ranges;

-    bool use_fusion;
+    // shared fusion debugging context
+    ggml_metal_fusion_info * finfo;
+
    bool use_concurrency;
    bool use_capture;

    int debug_graph;
-    int debug_fusion;

 private:
    ggml_cgraph * gf;
@@ -115,24 +124,22 @@ ggml_metal_op_t ggml_metal_op_init(
        ggml_metal_device_t dev,
        ggml_metal_cmd_buf_t cmd_buf,
        ggml_cgraph * gf,
+        ggml_metal_fusion_info * finfo,
        int idx_start,
        int idx_end,
-        bool use_fusion,
        bool use_concurrency,
        bool use_capture,
-        int debug_graph,
-        int debug_fusion) {
+        int debug_graph) {
    ggml_metal_op_t res = new ggml_metal_op(
        dev,
        cmd_buf,
        gf,
+        finfo,
        idx_start,
        idx_end,
-        use_fusion,
        use_concurrency,
        use_capture,
-        debug_graph,
-        debug_fusion);
+        debug_graph);

    return res;
 }
@@ -1868,6 +1875,8 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) {
    ggml_metal_library_t lib = ctx->lib;
    ggml_metal_encoder_t enc = ctx->enc;

+    const bool use_fusion = ctx->use_fusion();
+    const int  debug_fusion = ggml_metal_fusion_info_debug(ctx->finfo);

    GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
    GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
@@ -1880,6 +1889,31 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) {

    auto pipeline = ggml_metal_library_get_pipeline_gated_delta_net(lib, op);

+    // when fused with the trailing cache cpy, the snapshots are written straight into the
+    // recurrent cache and the cpy is skipped (see GGML_METAL_FUSION_GDN_CACHE)
+    ggml_metal_buffer_id bid_out = ggml_metal_get_buffer_id(op);
+    uint64_t nb_out = 0;
+    int n_fuse = 1;
+
+    if (use_fusion) {
+        int n = 1;
+        const ggml_metal_fusion * fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n);
+
+        if (fusion && fusion->id == GGML_METAL_FUSION_GDN_CACHE) {
+            const ggml_tensor * dst_cache = ctx->node(idx + 1)->src[1]; // cache view
+
+            bid_out = ggml_metal_get_buffer_id(dst_cache);
+            nb_out  = dst_cache->nb[2]/sizeof(float);
+            n_fuse = 2;
+
+            ctx->count_fusions(fusion);
+
+            if (debug_fusion > 1) {
+                GGML_LOG_DEBUG("%s: fuse: GATED_DELTA_NET + CPY\n", __func__);
+            }
+        }
+    }
+
    int ida = 0;

    ggml_metal_kargs_gated_delta_net args = {
@@ -1918,23 +1952,25 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) {
        /*.nb1  =*/ nb1,
        /*.nb2  =*/ nb2,
        /*.nb3  =*/ nb3,
+        /*.nb_out =*/ nb_out,
    };

    ggml_metal_encoder_set_pipeline(enc, pipeline);
-    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args),                  ida++);
+    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args),                  ida++); // args
    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), ida++); // q
    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[1]), ida++); // k
    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[2]), ida++); // v
    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[3]), ida++); // gate
    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[4]), ida++); // beta
    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[5]), ida++); // state
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         ida++); // dst
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         ida++); // dst (attn)
+    ggml_metal_encoder_set_buffer  (enc, bid_out,                              ida++); // state_out

    const int nsg = pipeline.nsg;

    ggml_metal_encoder_dispatch_threadgroups(enc, op->src[2]->ne[0]/nsg, op->src[2]->ne[1], op->src[2]->ne[3], 32, nsg, 1);

-    return 1;
+    return n_fuse;
 }

 int ggml_metal_op_solve_tri(ggml_metal_op_t ctx, int idx) {
@@ -3718,56 +3754,20 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
    return 1;
 }

-// Snake activation autofuse: mul -> sin -> sqr -> mul -> add
-static bool ggml_metal_op_can_fuse_snake(ggml_metal_op_t ctx, int idx) {
-    static constexpr ggml_op snake_ops[5] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD };
-
-    if (ctx->node(idx)->op != GGML_OP_MUL || !ctx->can_fuse(idx, snake_ops, 5)) {
-        return false;
-    }
-
-    const ggml_tensor * mul0     = ctx->node(idx + 0);
-    const ggml_tensor * sin_node = ctx->node(idx + 1);
-    const ggml_tensor * sqr      = ctx->node(idx + 2);
-    const ggml_tensor * mul1     = ctx->node(idx + 3);
-    const ggml_tensor * add      = ctx->node(idx + 4);
-
-    // x carries the full activation shape, a is the broadcast operand
-    const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1];
-    const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0];
-
-    // mul1 reads sqr and inv_b in either operand order
-    const ggml_tensor * inv_b    = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0];
-
-    // closure check: the trailing add reads the same x as the leading mul
-    const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0];
-
-    // x is in the supported whitelist and every chain intermediate shares x's type.
-    // a and inv_b bind as device const float * in the kernel, so they stay F32.
-    const bool types_ok =
-        (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) &&
-        (a->type    == GGML_TYPE_F32) && (inv_b->type    == GGML_TYPE_F32) &&
-        (mul0->type == x->type)       && (sin_node->type == x->type) &&
-        (sqr->type  == x->type)       && (mul1->type     == x->type) &&
-        (add->type  == x->type);
-    // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D
-    const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1];
-    const bool dim_ok =
-        (x->ne[2]     == 1) && (x->ne[3]     == 1) &&
-        (add->ne[2]   == 1) && (add->ne[3]   == 1) &&
-        (a->ne[2]     == 1) && (a->ne[3]     == 1) &&
-        (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1);
-    // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous
-    const bool contig_ok =
-        ggml_is_contiguous(x) && ggml_is_contiguous(add) &&
-        ggml_is_contiguous(a) && ggml_is_contiguous(inv_b);
-
-    return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x;
-}
-
 int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
-    if (ctx->use_fusion && ggml_metal_op_can_fuse_snake(ctx, idx)) {
-        return ggml_metal_op_snake_fused(ctx, idx);
+    int n_fuse = 1;
+    const ggml_metal_fusion * fusion = nullptr;
+
+    if (ctx->use_fusion()) {
+        int n = 1;
+        fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n);
+        n_fuse = n;
+
+        // snake activation autofuse: mul -> sin -> sqr -> mul -> add
+        if (fusion && fusion->id == GGML_METAL_FUSION_SNAKE) {
+            ctx->count_fusions(fusion);
+            return ggml_metal_op_snake_fused(ctx, idx);
+        }
    }

    ggml_tensor * op = ctx->node(idx);
@@ -3775,9 +3775,9 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
    ggml_metal_library_t lib = ctx->lib;
    ggml_metal_encoder_t enc = ctx->enc;

-    const bool use_fusion = ctx->use_fusion;
+    const bool use_fusion = ctx->use_fusion();

-    const int debug_fusion = ctx->debug_fusion;
+    const int debug_fusion = ggml_metal_fusion_info_debug(ctx->finfo);

    GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
    GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
@@ -3822,57 +3822,19 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
        /*.o1   =*/ { bid_src1.offs },
    };

-    ggml_op fops[8];
-
-    int n_fuse = 1;
-
    // c[0] = add(a,    b[0])
    // c[1] = add(c[0], b[1])
    // c[2] = add(c[1], b[2])
    // ...
-    if (use_fusion) {
-        fops[0] = GGML_OP_ADD;
-        fops[1] = GGML_OP_ADD;
-        fops[2] = GGML_OP_ADD;
-        fops[3] = GGML_OP_ADD;
-        fops[4] = GGML_OP_ADD;
-        fops[5] = GGML_OP_ADD;
-        fops[6] = GGML_OP_ADD;
-        fops[7] = GGML_OP_ADD;
-
-        // note: in metal, we sometimes encode the graph in parallel so we have to avoid fusing ops
-        //       across splits. idx_end indicates the last node in the current split
-        for (n_fuse = 0; n_fuse <= 6; ++n_fuse) {
-            if (!ctx->can_fuse(idx + n_fuse, fops + n_fuse, 2)) {
-                break;
-            }
-
-            ggml_tensor * f0 = ctx->node(idx + n_fuse);
-            ggml_tensor * f1 = ctx->node(idx + n_fuse + 1);
-
-            if (f0 != f1->src[0]) {
-                break;
-            }
-
-            // b[0] === b[1] === ...
-            if (!ggml_are_same_layout(f0->src[1], f1->src[1])) {
-                break;
-            }
-
-            // only fuse ops if src1 is in the same Metal buffer
-            ggml_metal_buffer_id bid_fuse = ggml_metal_get_buffer_id(f1->src[1]);
-            if (bid_fuse.metal != bid_src1.metal) {
-                break;
-            }
-
-            //ctx->fuse_cnt[ops[n_fuse + 1]->op]++;
-
-            args.o1[n_fuse + 1] = bid_fuse.offs;
+    if (use_fusion && fusion && fusion->id == GGML_METAL_FUSION_ADD_CHAIN) {
+        // the offsets of the fused addends are relative to the start of the src1 buffer
+        for (int i = 1; i < n_fuse; i++) {
+            args.o1[i] = ggml_metal_get_buffer_id(ctx->node(idx + i)->src[1]).offs;
        }

-        ++n_fuse;
+        ctx->count_fusions(fusion);

-        if (debug_fusion > 1 && n_fuse > 1) {
+        if (debug_fusion > 1) {
            GGML_LOG_DEBUG("%s: fuse: ADD x %d\n", __func__, n_fuse);
        }
    }
@@ -4080,9 +4042,9 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) {
    ggml_metal_library_t lib = ctx->lib;
    ggml_metal_encoder_t enc = ctx->enc;

-    const bool use_fusion = ctx->use_fusion;
+    const bool use_fusion = ctx->use_fusion();

-    const int debug_fusion = ctx->debug_fusion;
+    const int debug_fusion = ggml_metal_fusion_info_debug(ctx->finfo);

    GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
    GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
@@ -4110,8 +4072,6 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) {
        /*.nbf3   =*/ { nb03 },
    };

-    ggml_op fops[8];
-
    int n_fuse = 1;

    ggml_metal_buffer_id bid_fuse[2] = { bid_src0, bid_src0 };
@@ -4120,55 +4080,35 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) {
    // d[1] = mul(d[0], b)
    // d[2] = add(d[1], c)
    if (use_fusion) {
-        fops[0] = op->op;
-        fops[1] = GGML_OP_MUL;
-        fops[2] = GGML_OP_ADD;
+        int n = 1;
+        const ggml_metal_fusion * fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n);

-        for (n_fuse = 0; n_fuse <= 1; ++n_fuse) {
-            if (!ctx->can_fuse(idx + n_fuse, fops + n_fuse, 2)) {
-                break;
+        if (fusion && (fusion->id == GGML_METAL_FUSION_NORM_MUL || fusion->id == GGML_METAL_FUSION_NORM_MUL_ADD)) {
+            n_fuse = n;
+
+            ctx->count_fusions(fusion);
+
+            for (int i = 1; i < n_fuse; i++) {
+                const ggml_tensor * fn = ctx->node(idx + i);
+
+                bid_fuse[i - 1] = ggml_metal_get_buffer_id(fn->src[1]);
+
+                args.nef1[i] = fn->src[1]->ne[1];
+                args.nef2[i] = fn->src[1]->ne[2];
+                args.nef3[i] = fn->src[1]->ne[3];
+
+                args.nbf1[i] = fn->src[1]->nb[1];
+                args.nbf2[i] = fn->src[1]->nb[2];
+                args.nbf3[i] = fn->src[1]->nb[3];
            }

-            ggml_tensor * f0 = ctx->node(idx + n_fuse);
-            ggml_tensor * f1 = ctx->node(idx + n_fuse + 1);
-
-            if (f0 != f1->src[0]) {
-                break;
-            }
-
-            if (f1->src[1]->ne[0] != op->ne[0]) {
-                break;
-            }
-
-            if (!ggml_is_contiguous_rows(f1->src[1])) {
-                break;
-            }
-
-            if (f1->type != GGML_TYPE_F32) {
-                break;
-            }
-
-            //ctx->fuse_cnt[f1->op]++;
-
-            bid_fuse[n_fuse] = ggml_metal_get_buffer_id(f1->src[1]);
-
-            args.nef1[n_fuse + 1] = f1->src[1]->ne[1];
-            args.nef2[n_fuse + 1] = f1->src[1]->ne[2];
-            args.nef3[n_fuse + 1] = f1->src[1]->ne[3];
-
-            args.nbf1[n_fuse + 1] = f1->src[1]->nb[1];
-            args.nbf2[n_fuse + 1] = f1->src[1]->nb[2];
-            args.nbf3[n_fuse + 1] = f1->src[1]->nb[3];
-        }
-
-        ++n_fuse;
-
-        if (debug_fusion > 1 && n_fuse > 1) {
-            if (n_fuse == 2) {
-                GGML_LOG_DEBUG("%s: fuse: %s + MUL\n", __func__, ggml_op_name(op->op));
-            }
-            if (n_fuse == 3) {
-                GGML_LOG_DEBUG("%s: fuse: %s + MUL + ADD\n", __func__, ggml_op_name(op->op));
+            if (debug_fusion > 1) {
+                if (n_fuse == 2) {
+                    GGML_LOG_DEBUG("%s: fuse: %s + MUL\n", __func__, ggml_op_name(op->op));
+                }
+                if (n_fuse == 3) {
+                    GGML_LOG_DEBUG("%s: fuse: %s + MUL + ADD\n", __func__, ggml_op_name(op->op));
+                }
            }
        }
    }
@@ -8,17 +8,18 @@ extern "C" {

 typedef struct ggml_metal_op * ggml_metal_op_t;

+struct ggml_metal_fusion; // forward decl (ggml-metal-device.h)
+
 ggml_metal_op_t ggml_metal_op_init(
        ggml_metal_device_t dev,
        ggml_metal_cmd_buf_t cmd_buf,
        struct ggml_cgraph * gf,
+        struct ggml_metal_fusion_info * finfo,
        int  idx_start,
        int  idx_end,
-        bool use_fusion,
        bool use_concurrency,
        bool use_capture,
-        int  debug_graph,
-        int  debug_fusion);
+        int  debug_graph);

 void ggml_metal_op_free(ggml_metal_op_t ctx);

@@ -4,6 +4,7 @@
 #include "ggml-backend-impl.h"

 #include "ggml-metal-device.h"
+#include "ggml-metal-fusion.h"
 #include "ggml-metal-context.h"
 #include "ggml-metal-ops.h"
 #include "ggml-metal-tuning.h"
@@ -906,6 +907,30 @@ static const char * ggml_backend_metal_tuning_device_token(ggml_backend_dev_t de
    return ggml_metal_device_id_token(ggml_metal_device_get_props(ctx_dev)->device_id);
 }

+// generic fusion debugging API (ad-hoc proc-address mechanism): the test resolves the device
+// fusion context once and passes that opaque handle to the rest of the functions
+typedef void * ggml_backend_fusion_t;
+
+static ggml_backend_fusion_t ggml_backend_metal_fusion_get(ggml_backend_dev_t dev) {
+    return ggml_metal_device_get_fusion_info((ggml_metal_device_t)dev->context);
+}
+
+static void ggml_backend_metal_fusion_stats_init(ggml_backend_fusion_t finfo) {
+    ggml_metal_fusion_info_stats_init((struct ggml_metal_fusion_info *) finfo);
+}
+
+static void ggml_backend_metal_fusion_stats_reset(ggml_backend_fusion_t finfo) {
+    ggml_metal_fusion_info_stats_reset((struct ggml_metal_fusion_info *) finfo);
+}
+
+static int ggml_backend_metal_fusion_stats_get(ggml_backend_fusion_t finfo, const char ** labels, uint64_t * counts, int n) {
+    return ggml_metal_fusion_info_stats_get((struct ggml_metal_fusion_info *) finfo, labels, counts, n);
+}
+
+static void ggml_backend_metal_fusion_set_enabled(ggml_backend_fusion_t finfo, bool enabled) {
+    ggml_metal_fusion_info_set_enabled((struct ggml_metal_fusion_info *) finfo, enabled);
+}
+
 static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const char * name) {
    if (strcmp(name, "ggml_backend_get_features") == 0) {
        return (void *)ggml_backend_metal_get_features;
@@ -928,6 +953,23 @@ static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const
    if (strcmp(name, "ggml_backend_metal_tuning_device_token") == 0) {
        return (void *)ggml_backend_metal_tuning_device_token;
    }
+    // generic fusion debugging API (ad-hoc proc-address mechanism, not part of the official
+    // ggml backend interface yet; a backend that adopts it exports these exact names)
+    if (strcmp(name, "ggml_backend_fusion_get") == 0) {
+        return (void *)ggml_backend_metal_fusion_get;
+    }
+    if (strcmp(name, "ggml_backend_fusion_stats_init") == 0) {
+        return (void *)ggml_backend_metal_fusion_stats_init;
+    }
+    if (strcmp(name, "ggml_backend_fusion_stats_reset") == 0) {
+        return (void *)ggml_backend_metal_fusion_stats_reset;
+    }
+    if (strcmp(name, "ggml_backend_fusion_stats_get") == 0) {
+        return (void *)ggml_backend_metal_fusion_stats_get;
+    }
+    if (strcmp(name, "ggml_backend_fusion_set_enabled") == 0) {
+        return (void *)ggml_backend_metal_fusion_set_enabled;
+    }

    return NULL;

@@ -15,6 +15,7 @@ kernel void kernel_gated_delta_net_impl(
        device const char * b,
        device const char * s,
        device       char * dst,
+        device       char * dst_fuse,
        uint3 tgpig[[threadgroup_position_in_grid]],
        uint3 tpitg[[thread_position_in_threadgroup]],
        uint3   ntg[[threads_per_threadgroup]])  {
@@ -65,6 +66,12 @@ kernel void kernel_gated_delta_net_impl(
    // per-(seq,head) offset within a slot
    const uint state_out_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;

+    // when fused with the cache cpy, write the snapshots straight into the cache buffer using
+    // the slot stride; otherwise append them after the attn scores (nb_out == 0)
+    const bool fused = args.nb_out > 0;
+    const device float * state_out = fused ? (device float *)dst_fuse : (device float *)dst + attn_size;
+    const uint slot_stride = fused ? (uint)args.nb_out : state_size_per_snap;
+
    for (short t = 0; t < args.ne22; t++) {
        float s_k = 0.0f;

@@ -116,7 +123,7 @@ kernel void kernel_gated_delta_net_impl(
        if (K > 1) {
            const int target_slot = (int)args.ne22 - 1 - (int)t;
            if (target_slot >= 0 && target_slot < (int)K) {
-                device float * dst_state = (device float *) (dst) + attn_size + (uint)target_slot * state_size_per_snap + state_out_base;
+                device float * dst_state = (device float *)state_out + (uint)target_slot * slot_stride + state_out_base;
                FOR_UNROLL (short j = 0; j < NSG; j++) {
                    const short is = tx*NSG + j;
                    dst_state[is] = ls[j];
@@ -126,7 +133,7 @@ kernel void kernel_gated_delta_net_impl(
    }

    if (K == 1) {
-        device float * dst_state = (device float *) (dst) + attn_size + state_out_base;
+        device float * dst_state = (device float *)state_out + state_out_base;
        FOR_UNROLL (short j = 0; j < NSG; j++) {
            const short is = tx*NSG + j;
            dst_state[is] = ls[j];
@@ -158,6 +165,7 @@ kernel void kernel_gated_delta_net_impl(
        device const char * b,
        device const char * s,
        device       char * dst,
+        device       char * dst_fuse,
        uint3 tgpig[[threadgroup_position_in_grid]],
        uint3 tpitg[[thread_position_in_threadgroup]],
        uint3   ntg[[threads_per_threadgroup]])  {
@@ -230,7 +238,13 @@ kernel void kernel_gated_delta_net_impl(
        dst_attn += args.ne21*S_v;
    }

-    device float * dst_state  = (device float *) (dst) + args.ne23*args.ne22*args.ne21*S_v + (i23*args.ne21 + i21)*S_v*S_v + i20;
+    // when fused with the cache cpy, write the snapshots straight into the cache buffer using
+    // the slot stride; otherwise append them after the attn scores (nb_out == 0)
+    const bool fused = args.nb_out > 0;
+    const device float * state_out = fused ? (device float *)dst_fuse : (device float *)dst + args.ne23*args.ne22*args.ne21*S_v;
+    const uint slot_stride = fused ? (uint)args.nb_out : S_v*S_v;
+
+    device float * dst_state  = (device float *)state_out + (i23*args.ne21 + i21)*slot_stride + i20;
    device T     * dstt_state = (device T     *) (dst_state);

    FOR_UNROLL (short j = 0; j < NSG; j++) {
@@ -496,6 +496,13 @@ kernel void kernel_mul_mm_id(
        + args.nb11*i11
        + args.nb10*iy);

+    // skip the upper half of the token tile when the expert did not fill it
+    constexpr short NR1H = NR1/2;
+
+    const bool has_hi = nr1 > NR1H;
+
+    const short lb1 = (short) tiitg/NL1; // 0 .. NR1-1, this thread's row of the B tile
+
 #ifndef GGML_METAL_HAS_TENSOR
    S0_8x8 ma[4];
    S1_8x8 mb[2];
@@ -505,15 +512,22 @@ kernel void kernel_mul_mm_id(
    for (short i = 0; i < 8; i++){
        mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
    }
+
+    // simdgroups 2,3 own rows NR1H..NR1-1
+    const bool sg_active = has_hi || sgitg < 2;
 #else
-    auto tA = tensor<threadgroup S0, dextents<int32_t, 2>, tensor_inline>(sa, dextents<int32_t, 2>(NK,  NR0));
-    auto tB = tensor<threadgroup S1, dextents<int32_t, 2>, tensor_inline>(sb, dextents<int32_t, 2>(NR1, NK ));
+    auto tA  = tensor<threadgroup S0, dextents<int32_t, 2>, tensor_inline>(sa, dextents<int32_t, 2>(NK, NR0));
+
+    // sb is [NR1][NK] row-major
+    auto tB0 = tensor<threadgroup S1, dextents<int32_t, 2>, tensor_inline>(sb,             dextents<int32_t, 2>(NK, NR1H));
+    auto tB1 = tensor<threadgroup S1, dextents<int32_t, 2>, tensor_inline>(sb + NR1H*NK,   dextents<int32_t, 2>(NK, NR1H));

    mpp::tensor_ops::matmul2d<
-        mpp::tensor_ops::matmul2d_descriptor(NR1, NR0, NK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
+        mpp::tensor_ops::matmul2d_descriptor(NR1H, NR0, NK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
        execution_simdgroups<4>> mm;

-    auto cT = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB), float>();
+    auto cT0 = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB0), float>();
+    auto cT1 = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB1), float>();
 #endif

    for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
@@ -656,37 +670,45 @@ kernel void kernel_mul_mm_id(
        threadgroup_barrier(mem_flags::mem_threadgroup);

 #ifndef GGML_METAL_HAS_TENSOR
-        // load matrices from threadgroup memory and conduct outer products
-        threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
-        threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
+        if (sg_active) {
+            // load matrices from threadgroup memory and conduct outer products
+            threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
+            threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));

-        FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
-            simdgroup_barrier(mem_flags::mem_none);
+            FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
+                simdgroup_barrier(mem_flags::mem_none);

-            FOR_UNROLL (short i = 0; i < 4; i++) {
-                simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
+                FOR_UNROLL (short i = 0; i < 4; i++) {
+                    simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
+                }
+
+                simdgroup_barrier(mem_flags::mem_none);
+
+                FOR_UNROLL (short i = 0; i < 2; i++) {
+                    simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
+                }
+
+                simdgroup_barrier(mem_flags::mem_none);
+
+                FOR_UNROLL (short i = 0; i < 8; i++){
+                    simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
+                }
+
+                lsma += 8*64;
+                lsmb += 4*64;
            }
-
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 2; i++) {
-                simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
-            }
-
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 8; i++){
-                simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
-            }
-
-            lsma += 8*64;
-            lsmb += 4*64;
        }
 #else
-        auto sA = tA.slice(0, 0);
-        auto sB = tB.slice(0, 0);
+        auto sA  = tA.slice(0, 0);
+        auto sB0 = tB0.slice(0, 0);

-        mm.run(sB, sA, cT);
+        mm.run(sB0, sA, cT0);
+
+        if (has_hi) {
+            auto sB1 = tB1.slice(0, 0);
+
+            mm.run(sB1, sA, cT1);
+        }
 #endif
    }

@@ -694,13 +716,20 @@ kernel void kernel_mul_mm_id(
    threadgroup_barrier(mem_flags::mem_threadgroup);

 #ifdef GGML_METAL_HAS_TENSOR
-    auto tC = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc, dextents<int32_t, 2>(NR0, NR1));
-    cT.store(tC);
-#else
-    threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
+    auto tC0 = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc,             dextents<int32_t, 2>(NR0, NR1H));
+    cT0.store(tC0);

-    for (short i = 0; i < 8; i++) {
-        simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
+    if (has_hi) {
+        auto tC1 = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc + NR1H*NR0, dextents<int32_t, 2>(NR0, NR1H));
+        cT1.store(tC1);
+    }
+#else
+    if (sg_active) {
+        threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
+
+        for (short i = 0; i < 8; i++) {
+            simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
+        }
    }
 #endif

@@ -1889,8 +1889,19 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
    const uint i12 = im%FC_mul_mv_ne12;
    const uint i13 = im/FC_mul_mv_ne12;

-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;

    device const block_iq2_xxs * x = (device const block_iq2_xxs *) (src0 + offset0);
    device const float         * y = (device const float         *) (src1 + offset1);
@@ -1898,8 +1909,6 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
    float yl[32];
    float sumf[nr0]={0.f};

-    const int nb32 = nb * (QK_K / 32);
-
    threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem);
    threadgroup uint8_t  * ssigns  = (threadgroup uint8_t  *)(svalues + 256);
    {
@@ -1912,11 +1921,9 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

-    const int ix = tiisg;
-
    device const float * y4 = y + 32 * ix;

-    for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
        for (short i = 0; i < 32; ++i) {
            yl[i] = y4[i];
        }
@@ -1928,7 +1935,7 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
        device const uint16_t * q2 = xr->qs + 4 * ib;
        device const half * dh = &xr->d;

-        for (short row = 0; row < nr0; row++) {
+        for (short row = row0; row < row1; row++) {
            const float db = dh[0];
            device const uint8_t * aux8 = (device const uint8_t *)q2;
            const uint32_t aux32 = q2[2] | (q2[3] << 16);
@@ -1948,7 +1955,7 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
            q2 += args.nb01/2;
        }

-        y4 += 32 * 32;
+        y4 += 32 * ntx;
    }

    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -1961,6 +1968,23 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
    }
 }

+template<typename args_t>
+void kernel_mul_mv_iq2_xxs_f32_disp(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    if (FC_mul_mv_split) {
+        kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    }
+}
+
 [[host_name("kernel_mul_mv_iq2_xxs_f32")]]
 kernel void kernel_mul_mv_iq2_xxs_f32(
        constant ggml_metal_kargs_mul_mv & args,
@@ -1971,7 +1995,7 @@ kernel void kernel_mul_mv_iq2_xxs_f32(
        uint3  tgpig[[threadgroup_position_in_grid]],
        ushort tiisg[[thread_index_in_simdgroup]],
        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
-    kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    kernel_mul_mv_iq2_xxs_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
 }

 template<int nr0, typename args_t>
@@ -1997,8 +2021,19 @@ void kernel_mul_mv_iq2_xs_f32_impl(
    const uint i12 = im%FC_mul_mv_ne12;
    const uint i13 = im/FC_mul_mv_ne12;

-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;

    device const block_iq2_xs * x = (device const block_iq2_xs *) (src0 + offset0);
    device const float        * y = (device const float        *) (src1 + offset1);
@@ -2006,8 +2041,6 @@ void kernel_mul_mv_iq2_xs_f32_impl(
    float yl[32];
    float sumf[nr0]={0.f};

-    const int nb32 = nb * (QK_K / 32);
-
    threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem);
    threadgroup uint8_t  * ssigns  = (threadgroup uint8_t  *)(svalues + 512);
    {
@@ -2020,11 +2053,9 @@ void kernel_mul_mv_iq2_xs_f32_impl(
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

-    const int ix = tiisg;
-
    device const float * y4 = y + 32 * ix;

-    for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
        for (short i = 0; i < 32; ++i) {
            yl[i] = y4[i];
        }
@@ -2037,7 +2068,7 @@ void kernel_mul_mv_iq2_xs_f32_impl(
        device const uint8_t  * sc = xr->scales + ib;
        device const half * dh = &xr->d;

-        for (short row = 0; row < nr0; row++) {
+        for (short row = row0; row < row1; row++) {
            const float db = dh[0];
            const uint8_t ls1 = sc[0] & 0xf;
            const uint8_t ls2 = sc[0] >>  4;
@@ -2066,7 +2097,7 @@ void kernel_mul_mv_iq2_xs_f32_impl(
            sc += args.nb01;
        }

-        y4 += 32 * 32;
+        y4 += 32 * ntx;
    }

    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2079,6 +2110,23 @@ void kernel_mul_mv_iq2_xs_f32_impl(
    }
 }

+template<typename args_t>
+void kernel_mul_mv_iq2_xs_f32_disp(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    if (FC_mul_mv_split) {
+        kernel_mul_mv_iq2_xs_f32_impl<N_R0_IQ2_XS_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_iq2_xs_f32_impl<N_R0_IQ2_XS, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    }
+}
+
 [[host_name("kernel_mul_mv_iq2_xs_f32")]]
 kernel void kernel_mul_mv_iq2_xs_f32(
        constant ggml_metal_kargs_mul_mv & args,
@@ -2090,7 +2138,7 @@ kernel void kernel_mul_mv_iq2_xs_f32(
        ushort tiisg[[thread_index_in_simdgroup]],
        ushort sgitg[[simdgroup_index_in_threadgroup]]) {

-    kernel_mul_mv_iq2_xs_f32_impl<N_R0_IQ2_XS, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    kernel_mul_mv_iq2_xs_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
 }

 // FC_mul_mv_split: for nb32 < 32 (nb32 divides 32), 32/nb32 threads share each chunk and each takes a slice of the rows
@@ -2117,8 +2165,19 @@ void kernel_mul_mv_iq3_xxs_f32_impl(
    const uint i12 = im%FC_mul_mv_ne12;
    const uint i13 = im/FC_mul_mv_ne12;

-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;

    device const block_iq3_xxs * x = (device const block_iq3_xxs *) (src0 + offset0);
    device const float         * y = (device const float         *) (src1 + offset1);
@@ -2126,8 +2185,6 @@ void kernel_mul_mv_iq3_xxs_f32_impl(
    float yl[32];
    float sumf[nr0]={0.f};

-    const int nb32 = nb * (QK_K / 32);
-
    threadgroup uint32_t * svalues = (threadgroup uint32_t *)(shmem);
    threadgroup uint8_t  * ssigns  = (threadgroup uint8_t  *)(svalues + 256);
    {
@@ -2140,15 +2197,6 @@ void kernel_mul_mv_iq3_xxs_f32_impl(
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

-    const short ntx  = FC_mul_mv_split ? nb32 : 32;
-    const short nrep = 32 / ntx;
-
-    const short ix   = tiisg % ntx;
-    const short irep = tiisg / ntx;
-
-    const short row0 = (nr0 * irep      ) / nrep;
-    const short row1 = (nr0 * (irep + 1)) / nrep;
-
    device const float * y4 = y + 32 * ix;

    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
@@ -2160,9 +2208,9 @@ void kernel_mul_mv_iq3_xxs_f32_impl(
        const int ib  = ib32 % (QK_K / 32);

        device const block_iq3_xxs * xr = x + ibl;
-        device const uint8_t  * q3 = xr->qs + 8 * ib + (uint64_t) row0*args.nb01;
-        device const uint16_t * gas = (device const uint16_t *)(xr->qs + QK_K/4) + 2 * ib + (uint64_t) row0*args.nb01/2;
-        device const half * dh = &xr->d + (uint64_t) row0*args.nb01/2;
+        device const uint8_t  * q3 = xr->qs + 8 * ib;
+        device const uint16_t * gas = (device const uint16_t *)(xr->qs + QK_K/4) + 2 * ib;
+        device const half * dh = &xr->d;

        for (short row = row0; row < row1; row++) {
            const float db = dh[0];
@@ -2253,8 +2301,19 @@ void kernel_mul_mv_iq3_s_f32_impl(
    const uint i12 = im%FC_mul_mv_ne12;
    const uint i13 = im/FC_mul_mv_ne12;

-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;

    device const block_iq3_s * x = (device const block_iq3_s *) (src0 + offset0);
    device const float       * y = (device const float       *) (src1 + offset1);
@@ -2262,8 +2321,6 @@ void kernel_mul_mv_iq3_s_f32_impl(
    float yl[32];
    float sumf[nr0]={0.f};

-    const int nb32 = nb * (QK_K / 32);
-
    threadgroup uint32_t * svalues = (threadgroup uint32_t *) shmem;
    {
        int nval = 8;
@@ -2272,11 +2329,9 @@ void kernel_mul_mv_iq3_s_f32_impl(
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

-    const int ix = tiisg;
-
    device const float * y4 = y + 32 * ix;

-    for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
        for (short i = 0; i < 32; ++i) {
            yl[i] = y4[i];
        }
@@ -2291,7 +2346,7 @@ void kernel_mul_mv_iq3_s_f32_impl(
        device const uint8_t * signs = xr->signs + 4 * ib;
        device const half * dh = &xr->d;

-        for (short row = 0; row < nr0; row++) {
+        for (short row = row0; row < row1; row++) {
            const float db = dh[0];
            const float d = db * (1 + 2*((sc[0] >> 4*(ib%2)) & 0xf));

@@ -2315,7 +2370,7 @@ void kernel_mul_mv_iq3_s_f32_impl(
            signs += args.nb01;
        }

-        y4 += 32 * 32;
+        y4 += 32 * ntx;
    }

    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2328,6 +2383,23 @@ void kernel_mul_mv_iq3_s_f32_impl(
    }
 }

+template<typename args_t>
+void kernel_mul_mv_iq3_s_f32_disp(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    if (FC_mul_mv_split) {
+        kernel_mul_mv_iq3_s_f32_impl<N_R0_IQ3_S_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_iq3_s_f32_impl<N_R0_IQ3_S, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    }
+}
+
 [[host_name("kernel_mul_mv_iq3_s_f32")]]
 kernel void kernel_mul_mv_iq3_s_f32(
        constant ggml_metal_kargs_mul_mv & args,
@@ -2339,7 +2411,7 @@ kernel void kernel_mul_mv_iq3_s_f32(
        ushort tiisg[[thread_index_in_simdgroup]],
        ushort sgitg[[simdgroup_index_in_threadgroup]]) {

-    kernel_mul_mv_iq3_s_f32_impl<N_R0_IQ3_S, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    kernel_mul_mv_iq3_s_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
 }

 template<int nr0, typename args_t>
@@ -2365,8 +2437,19 @@ void kernel_mul_mv_iq2_s_f32_impl(
    const uint i12 = im%FC_mul_mv_ne12;
    const uint i13 = im/FC_mul_mv_ne12;

-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;

    device const block_iq2_s * x = (device const block_iq2_s *) (src0 + offset0);
    device const float       * y = (device const float       *) (src1 + offset1);
@@ -2374,8 +2457,6 @@ void kernel_mul_mv_iq2_s_f32_impl(
    float yl[32];
    float sumf[nr0]={0.f};

-    const int nb32 = nb * (QK_K / 32);
-
    //threadgroup uint64_t * svalues = (threadgroup uint64_t *) shmem;
    //{
    //    int nval = 32;
@@ -2384,11 +2465,9 @@ void kernel_mul_mv_iq2_s_f32_impl(
    //    threadgroup_barrier(mem_flags::mem_threadgroup);
    //}

-    const short ix = tiisg;
-
    device const float * y4 = y + 32 * ix;

-    for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
        for (short i = 0; i < 32; ++i) {
            yl[i] = y4[i];
        }
@@ -2403,7 +2482,7 @@ void kernel_mul_mv_iq2_s_f32_impl(
        device const uint8_t * signs = qs + QK_K/8;
        device const half * dh = &xr->d;

-        for (short row = 0; row < nr0; row++) {
+        for (short row = row0; row < row1; row++) {
            const float db = dh[0];
            const float d1 = db * (0.5f + (sc[0] & 0xf));
            const float d2 = db * (0.5f + (sc[0] >>  4));
@@ -2428,7 +2507,7 @@ void kernel_mul_mv_iq2_s_f32_impl(
            signs += args.nb01;
        }

-        y4 += 32 * 32;
+        y4 += 32 * ntx;
    }

    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2441,6 +2520,23 @@ void kernel_mul_mv_iq2_s_f32_impl(
    }
 }

+template<typename args_t>
+void kernel_mul_mv_iq2_s_f32_disp(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    if (FC_mul_mv_split) {
+        kernel_mul_mv_iq2_s_f32_impl<N_R0_IQ2_S_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_iq2_s_f32_impl<N_R0_IQ2_S, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    }
+}
+
 [[host_name("kernel_mul_mv_iq2_s_f32")]]
 kernel void kernel_mul_mv_iq2_s_f32(
        constant ggml_metal_kargs_mul_mv & args,
@@ -2452,7 +2548,7 @@ kernel void kernel_mul_mv_iq2_s_f32(
        ushort tiisg[[thread_index_in_simdgroup]],
        ushort sgitg[[simdgroup_index_in_threadgroup]]) {

-    kernel_mul_mv_iq2_s_f32_impl<N_R0_IQ2_S, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    kernel_mul_mv_iq2_s_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
 }

 template<int nr0, typename args_t>
@@ -2478,8 +2574,19 @@ void kernel_mul_mv_iq1_s_f32_impl(
    const uint i12 = im%FC_mul_mv_ne12;
    const uint i13 = im/FC_mul_mv_ne12;

-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;

    device const block_iq1_s * x = (device const block_iq1_s *) (src0 + offset0);
    device const float       * y = (device const float       *) (src1 + offset1);
@@ -2487,13 +2594,9 @@ void kernel_mul_mv_iq1_s_f32_impl(
    float yl[32];
    float sumf[nr0]={0.f};

-    const int nb32 = nb * (QK_K / 32);
-
-    const short ix = tiisg;
-
    device const float * y4 = y + 32 * ix;

-    for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
        float sumy = 0;
        for (short i = 0; i < 32; ++i) {
            yl[i] = y4[i];
@@ -2508,7 +2611,7 @@ void kernel_mul_mv_iq1_s_f32_impl(
        device const uint16_t * qh = xr->qh + ib;
        device const half     * dh = &xr->d;

-        for (short row = 0; row < nr0; row++) {
+        for (short row = row0; row < row1; row++) {
            constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700)));
            constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((qh[0] << 5) & 0x700)));
            constant uint8_t * grid3 = (constant uint8_t *)(iq1s_grid_gpu + (qs[2] | ((qh[0] << 2) & 0x700)));
@@ -2528,7 +2631,7 @@ void kernel_mul_mv_iq1_s_f32_impl(
            qh += args.nb01/2;
        }

-        y4 += 32 * 32;
+        y4 += 32 * ntx;
    }

    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2541,6 +2644,23 @@ void kernel_mul_mv_iq1_s_f32_impl(
    }
 }

+template<typename args_t>
+void kernel_mul_mv_iq1_s_f32_disp(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    if (FC_mul_mv_split) {
+        kernel_mul_mv_iq1_s_f32_impl<N_R0_IQ1_S_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_iq1_s_f32_impl<N_R0_IQ1_S, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    }
+}
+
 [[host_name("kernel_mul_mv_iq1_s_f32")]]
 kernel void kernel_mul_mv_iq1_s_f32(
        constant ggml_metal_kargs_mul_mv & args,
@@ -2551,7 +2671,7 @@ kernel void kernel_mul_mv_iq1_s_f32(
        ushort tiisg[[thread_index_in_simdgroup]],
        ushort sgitg[[simdgroup_index_in_threadgroup]]) {

-    kernel_mul_mv_iq1_s_f32_impl<N_R0_IQ1_S, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
+    kernel_mul_mv_iq1_s_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
 }

 template<int nr0, typename args_t>
@@ -2577,8 +2697,19 @@ void kernel_mul_mv_iq1_m_f32_impl(
    const uint i12 = im%FC_mul_mv_ne12;
    const uint i13 = im/FC_mul_mv_ne12;

-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;

    device const block_iq1_m * x = (device const block_iq1_m *) (src0 + offset0);
    device const float       * y = (device const float       *) (src1 + offset1);
@@ -2586,15 +2717,11 @@ void kernel_mul_mv_iq1_m_f32_impl(
    float yl[32];
    float sumf[nr0]={0.f};

-    const int nb32 = nb * (QK_K / 32);
-
-    const short ix = tiisg;
-
    device const float * y4 = y + 32 * ix;

    iq1m_scale_t scale;

-    for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
        float4 sumy = {0.f};
        for (short i = 0; i < 8; ++i) {
            yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0];
@@ -2611,7 +2738,7 @@ void kernel_mul_mv_iq1_m_f32_impl(
        device const uint8_t  * qh = xr->qh + 2 * ib;
        device const uint16_t * sc = (device const uint16_t *)xr->scales;

-        for (short row = 0; row < nr0; row++) {
+        for (short row = row0; row < row1; row++) {
            scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);

            constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700)));
@@ -2637,7 +2764,7 @@ void kernel_mul_mv_iq1_m_f32_impl(
            qh += args.nb01;
        }

-        y4 += 32 * 32;
+        y4 += 32 * ntx;
    }

    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2650,6 +2777,23 @@ void kernel_mul_mv_iq1_m_f32_impl(
    }
 }

+template<typename args_t>
+void kernel_mul_mv_iq1_m_f32_disp(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    if (FC_mul_mv_split) {
+        kernel_mul_mv_iq1_m_f32_impl<N_R0_IQ1_M_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_iq1_m_f32_impl<N_R0_IQ1_M, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    }
+}
+
 [[host_name("kernel_mul_mv_iq1_m_f32")]]
 kernel void kernel_mul_mv_iq1_m_f32(
        constant ggml_metal_kargs_mul_mv & args,
@@ -2660,7 +2804,7 @@ kernel void kernel_mul_mv_iq1_m_f32(
        ushort tiisg[[thread_index_in_simdgroup]],
        ushort sgitg[[simdgroup_index_in_threadgroup]]) {

-    kernel_mul_mv_iq1_m_f32_impl<N_R0_IQ1_M, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
+    kernel_mul_mv_iq1_m_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
 }

 template<int NR0, typename args_t>
@@ -3239,13 +3383,13 @@ template [[host_name("kernel_mul_mv_id_q3_K_f32")]]    kernel kernel_mul_mv_id_t
 template [[host_name("kernel_mul_mv_id_q4_K_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q4_K_f32_impl   <N_R0_Q4_K>>>;
 template [[host_name("kernel_mul_mv_id_q5_K_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q5_K_f32_impl   <N_R0_Q5_K>>>;
 template [[host_name("kernel_mul_mv_id_q6_K_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q6_K_f32_impl   <N_R0_Q6_K>>>;
-template [[host_name("kernel_mul_mv_id_iq1_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_s_f32_impl  <N_R0_IQ1_S>>>;
-template [[host_name("kernel_mul_mv_id_iq1_m_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_m_f32_impl  <N_R0_IQ1_M>>>;
-template [[host_name("kernel_mul_mv_id_iq2_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS>>>;
-template [[host_name("kernel_mul_mv_id_iq2_xs_f32")]]  kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xs_f32_impl <N_R0_IQ2_XS>>>;
+template [[host_name("kernel_mul_mv_id_iq1_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_s_f32_disp<ggml_metal_kargs_mul_mv>>>;
+template [[host_name("kernel_mul_mv_id_iq1_m_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_m_f32_disp<ggml_metal_kargs_mul_mv>>>;
+template [[host_name("kernel_mul_mv_id_iq2_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xxs_f32_disp<ggml_metal_kargs_mul_mv>>>;
+template [[host_name("kernel_mul_mv_id_iq2_xs_f32")]]  kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xs_f32_disp<ggml_metal_kargs_mul_mv>>>;
 template [[host_name("kernel_mul_mv_id_iq3_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq3_xxs_f32_disp<ggml_metal_kargs_mul_mv>>>;
-template [[host_name("kernel_mul_mv_id_iq3_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq3_s_f32_impl  <N_R0_IQ3_S>>>;
-template [[host_name("kernel_mul_mv_id_iq2_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_s_f32_impl  <N_R0_IQ2_S>>>;
+template [[host_name("kernel_mul_mv_id_iq3_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq3_s_f32_disp<ggml_metal_kargs_mul_mv>>>;
+template [[host_name("kernel_mul_mv_id_iq2_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_s_f32_disp<ggml_metal_kargs_mul_mv>>>;
 template [[host_name("kernel_mul_mv_id_iq4_nl_f32")]]  kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq4_nl_f32_impl <N_R0_IQ4_NL>>>;
 template [[host_name("kernel_mul_mv_id_iq4_xs_f32")]]  kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq4_xs_f32_impl <N_R0_IQ4_XS>>>;
 template [[host_name("kernel_mul_mv_id_tq2_0_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_tq2_0_f32_impl  <N_R0_TQ2_0>>>;
@@ -43,17 +43,8 @@ if (MUSAToolkit_FOUND)
        add_compile_definitions(GGML_MUSA_MUDNN_COPY)
    endif()

-    if (GGML_CUDA_FA_ALL_QUANTS)
-        file(GLOB   SRCS "../ggml-cuda/template-instances/fattn-vec*.cu")
-        list(APPEND GGML_SOURCES_MUSA ${SRCS})
-        add_compile_definitions(GGML_CUDA_FA_ALL_QUANTS)
-    else()
-        list(APPEND GGML_SOURCES_MUSA
-            ../ggml-cuda/template-instances/fattn-vec-instance-f16-f16.cu
-            ../ggml-cuda/template-instances/fattn-vec-instance-q4_0-q4_0.cu
-            ../ggml-cuda/template-instances/fattn-vec-instance-q8_0-q8_0.cu
-            ../ggml-cuda/template-instances/fattn-vec-instance-bf16-bf16.cu)
-    endif()
+    ggml_cuda_fattn_vec_instances(${CMAKE_CURRENT_SOURCE_DIR}/../ggml-cuda SRCS)
+    list(APPEND GGML_SOURCES_MUSA ${SRCS})

    set_source_files_properties(${GGML_SOURCES_MUSA} PROPERTIES LANGUAGE CXX)
    foreach(SOURCE ${GGML_SOURCES_MUSA})
@@ -170,6 +170,7 @@ set(GGML_OPENCL_KERNELS
    gemv_noshuffle_q4_0_f32
    gemv_noshuffle_q4_0_f32_spec
    gemm_noshuffle_q4_0_f32
+    gemv_noshuffle_q4_0_f32_32b_trans
    gemv_noshuffle_q4_1_f32
    gemm_noshuffle_q4_1_f32
    gemv_noshuffle_q5_0_f32
@@ -1160,6 +1160,8 @@ struct ggml_backend_opencl_context {
    cl_kernel kernel_gemm_noshuffle_q4_0_f32;
    cl_kernel kernel_gemv_noshuffle_q4_0_f32;
    cl_kernel kernel_gemv_noshuffle_q4_0_f32_mc3;  // multi-column (N=3) verify GEMV (spec/MTP)
+    cl_kernel kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin;
+    cl_kernel kernel_gemv_noshuffle_q4_0_f32_32b_trans;
    cl_kernel kernel_gemv_noshuffle_q4_0_f32_4096_1_11008;
    cl_kernel kernel_gemv_noshuffle_q4_0_f32_4096_1_4096;
    cl_kernel kernel_gemv_noshuffle_q4_0_f32_11008_1_4096;
@@ -3787,6 +3789,43 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
        GGML_LOG_CONT(".");
    }

+    backend_ctx->kernel_gemv_noshuffle_q4_0_f32_32b_trans = nullptr;
+    backend_ctx->kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin = nullptr;
+    if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E) {
+        {
+            std::string opts = std::string("-cl-std=") + opencl_c_std +
+                                           " -cl-mad-enable "
+                                           " -DSIMDGROUP_WIDTH=" +
+                                           std::to_string(backend_ctx->adreno_wave_size);
+#ifdef GGML_OPENCL_EMBED_KERNELS
+            const std::string kernel_src {
+                #include "gemv_noshuffle_q4_0_f32_32b_trans.cl.h"
+            };
+#else
+            const std::string kernel_src = read_file("gemv_noshuffle_q4_0_f32_32b_trans.cl");
+#endif
+            cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), opts);
+            CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q4_0_f32_32b_trans =
+                clCreateKernel(prog, "kernel_gemv_noshuffle_q4_0_f32_32b_trans", &err), err));
+            CL_CHECK(clReleaseProgram(prog));
+            GGML_LOG_CONT(".");
+        }
+
+        if (use_adreno_bin_kernels(backend_ctx)) {
+            size_t bin_size = 0;
+            const char * kernel_bin = (const char *)backend_ctx->get_adreno_bin_kernel("gemm_noshuffle_q4_0_f32_32b_trans_ila_a8", &bin_size);
+            if (kernel_bin && bin_size > 0) {
+                cl_program bin_prog =
+                    build_program_from_binary(backend_ctx->context, backend_ctx->device, kernel_bin, "", bin_size);
+
+                CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin =
+                    clCreateKernel(bin_prog, "kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8", &err), err));
+                CL_CHECK(clReleaseProgram(bin_prog));
+                GGML_LOG_CONT(".");
+            }
+        }
+    }
+
    // gemm_noshuffle_q4_1_f32
    {
 #ifdef GGML_OPENCL_EMBED_KERNELS
@@ -6725,11 +6764,10 @@ struct ggml_tensor_extra_cl_q4_0 {
            CL_CHECK(clReleaseMemObject(q_img));
            q_img = nullptr;
        }
-        // Currently, q_img and d_img are only initialized when SMALL_ALLOC is
-        // enabled. They point to the images in ggml_backend_opencl_buffer_context.
-        // So, there is no need to release them here.
-        // TODO: initialize them for non SMALL_PATH path, or remove them.
-        d_img = nullptr;
+        if (d_img != nullptr) {
+            CL_CHECK(clReleaseMemObject(d_img));
+            d_img = nullptr;
+        }
        size_q = 0;
        size_d = 0;
    }
@@ -8311,6 +8349,20 @@ inline bool enable_adreno_trans_weight_q5_K(const ggml_backend_opencl_context *b
           qh_img_width <= backend_ctx->image_max_buffer_size;
 }

+inline bool use_q4_0_ila_kernels(const ggml_backend_opencl_context *backend_ctx, const ggml_tensor *tensor) {
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+    if (!backend_ctx->kernel_gemv_noshuffle_q4_0_f32_32b_trans ||
+        !backend_ctx->kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin) {
+        return false;
+    }
+    return (tensor->ne[0] % 32 == 0) && (tensor->ne[1] % 64 == 0);
+#else
+    GGML_UNUSED(backend_ctx);
+    GGML_UNUSED(tensor);
+    return false;
+#endif
+}
+
 // The flat-GEMV large-m escape is OPT-IN (GGML_OPENCL_FLAT_LARGE_M=1) because it
 // is SLOWER than the route it replaces, not because it is unsafe. It was first
 // parked on the theory that it out-of-bounds-writes at vocab-scale shapes; that
@@ -9573,10 +9625,34 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,

            GGML_ASSERT(K % 32 == 0);

-            // Transpose q as ushort
-            transpose_2d_as_16b(backend_ctx, extra->q, extra->q, size_q, K/4, M);
-            // Transpose d as ushort
-            transpose_2d_as_16b(backend_ctx, extra->d, extra->d, size_d, K/32, M);
+            if (use_q4_0_ila_kernels(backend_ctx, tensor)) {
+                cl_int err;
+                cl_image_format wimg_fmt;
+                cl_image_desc   wimg_desc;
+
+                // transpose quants as 32-bit words (M-first)
+                GGML_ASSERT(M % 64 == 0);
+                transpose_2d_as_32b(backend_ctx, extra->q, extra->q, size_q, K / 8, M);
+                transpose_2d_as_16b(backend_ctx, extra->d, extra->d, size_d, K / 32, M);
+
+                wimg_fmt = { CL_R, CL_UNSIGNED_INT32 };
+                memset(&wimg_desc, 0, sizeof(wimg_desc));
+                wimg_desc.image_type  = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+                wimg_desc.image_width = (size_t)M * K / 8;
+                wimg_desc.buffer      = extra->q;
+                CL_CHECK((extra->q_img = clCreateImage(context, CL_MEM_READ_ONLY, &wimg_fmt, &wimg_desc, NULL, &err), err));
+
+                wimg_fmt = { CL_R, CL_HALF_FLOAT };
+                memset(&wimg_desc, 0, sizeof(wimg_desc));
+                wimg_desc.image_type  = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+                wimg_desc.image_width = (size_t)M * K / 32;
+                wimg_desc.buffer      = extra->d;
+                CL_CHECK((extra->d_img = clCreateImage(context, CL_MEM_READ_ONLY, &wimg_fmt, &wimg_desc, NULL, &err), err));
+            } else {
+                // Transpose q and d as ushort
+                transpose_2d_as_16b(backend_ctx, extra->q, extra->q, size_q, K/4, M);
+                transpose_2d_as_16b(backend_ctx, extra->d, extra->d, size_d, K/32, M);
+            }
        }
 #endif // GGML_OPENCL_USE_ADRENO_KERNELS
        return;
@@ -11104,7 +11180,11 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
            buf_trans_d.allocate(backend_ctx->context, size_d);
            buf_unpacked.allocate(backend_ctx->context, ggml_nbytes(tensor));

-            transpose_2d_as_16b(backend_ctx, extra->q, buf_trans_q.buffer, size_q, M, K/4);
+            if (use_q4_0_ila_kernels(backend_ctx, tensor)) {
+                transpose_2d_as_32b(backend_ctx, extra->q, buf_trans_q.buffer, size_q, M, K / 8);
+            } else {
+                transpose_2d_as_16b(backend_ctx, extra->q, buf_trans_q.buffer, size_q, M, K / 4);
+            }
            transpose_2d_as_16b(backend_ctx, extra->d, buf_trans_d.buffer, size_d, M, K/32);

            cl_uchar mask_0F = 0x0F;
@@ -18347,6 +18427,166 @@ static void ggml_cl_mul_mat_q1_0_f32_adreno(ggml_backend_t backend, const ggml_t
 #endif
 }

+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+static void ggml_cl_mul_mat_q4_0_f32_adreno_ila(ggml_backend_t backend, const ggml_tensor * src0,
+                                                const ggml_tensor * src1, ggml_tensor * dst) {
+    GGML_ASSERT(src0);
+    GGML_ASSERT(src0->extra);
+    GGML_ASSERT(src1);
+    GGML_ASSERT(src1->extra);
+    GGML_ASSERT(dst);
+    GGML_ASSERT(dst->extra);
+
+    ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context;
+
+    ggml_tensor_extra_cl * extra1 = (ggml_tensor_extra_cl *)src1->extra;
+    ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra;
+    ggml_tensor_extra_cl_q4_0 * extra0_q4_0 = (ggml_tensor_extra_cl_q4_0 *)src0->extra;
+
+    cl_ulong offset1 = extra1->offset + src1->view_offs;
+    cl_ulong offsetd = extrad->offset + dst->view_offs;
+
+    const int ne00 = src0->ne[0];
+    const int ne01 = src0->ne[1];
+
+    const int ne1 = dst->ne[1];
+
+    GGML_ASSERT(ne00 % ggml_blck_size(src0->type) == 0);
+
+    cl_context context = backend_ctx->context;
+    cl_kernel kernel;
+
+    cl_int              err;
+    cl_image_format     img_fmt;
+    cl_image_desc       img_desc;
+    cl_buffer_region    region;
+
+    int M = ne01;
+    int N = ne1;
+    int K = ne00;
+
+    if (ne1 == 1) {
+        cl_mem b_sub_buf = nullptr;
+        cl_mem b_img     = nullptr;
+
+        region.origin = offset1;
+        region.size   = (size_t)K * N * sizeof(float);
+        CL_CHECK((b_sub_buf = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, &region, &err), err));
+
+        img_fmt = { CL_RGBA, CL_FLOAT };
+        memset(&img_desc, 0, sizeof(img_desc));
+        img_desc.image_type  = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+        img_desc.image_width = (size_t)K * N / 4;
+        img_desc.buffer      = b_sub_buf;
+        CL_CHECK((b_img = clCreateImage(context, CL_MEM_READ_ONLY, &img_fmt, &img_desc, NULL, &err), err));
+
+        kernel = backend_ctx->kernel_gemv_noshuffle_q4_0_f32_32b_trans;
+        CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem),   &extra0_q4_0->q_img));
+        CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem),   &extra0_q4_0->d));
+        CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem),   &b_img));
+        CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_mem),   &extrad->data_device));
+        CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_ulong), &offsetd));
+        CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_int),   &K));
+        CL_CHECK(clSetKernelArg(kernel, 6, sizeof(cl_int),   &M));
+
+        size_t wavesize = backend_ctx->adreno_wave_size;
+        size_t local_work_size[3]  = { wavesize, 4, 1 };
+        size_t global_work_size[3] = { (size_t)CEIL_DIV(M, 64) * 64, 4, 1 };
+        backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+
+        CL_CHECK(clReleaseMemObject(b_sub_buf));
+        CL_CHECK(clReleaseMemObject(b_img));
+    } else {
+        const int gemm_tile_n = 64;
+        int N_pad = (N + gemm_tile_n - 1) & ~(gemm_tile_n - 1);
+
+        cl_mem a_img = extra0_q4_0->q_img;
+        cl_mem s_img = extra0_q4_0->d_img;
+        GGML_ASSERT(a_img && s_img && "ILA Q4_0 weight images missing; set_tensor should have built them");
+
+        // Pad B through a zero-filled scratch buffer when N needs
+        // padding, since the GEMM kernel always reads a full N-tile.
+        const bool need_pad = N_pad > N;
+        cl_mem b_sub_buf = nullptr;
+        cl_mem b_padded  = nullptr;
+        if (need_pad) {
+            CL_CHECK((b_padded = clCreateBuffer(context, CL_MEM_READ_WRITE,
+                (size_t)K * N_pad * sizeof(float), NULL, &err), err));
+            const float zero = 0.0f;
+            CL_CHECK(clEnqueueFillBuffer(backend_ctx->queue, b_padded, &zero, sizeof(zero),
+                0, (size_t)K * N_pad * sizeof(float), 0, NULL, NULL));
+            CL_CHECK(clEnqueueCopyBuffer(backend_ctx->queue, extra1->data_device, b_padded,
+                offset1, 0, (size_t)K * N * sizeof(float), 0, NULL, NULL));
+        } else {
+            region.origin = offset1;
+            region.size   = (size_t)K * N * sizeof(float);
+            CL_CHECK((b_sub_buf = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, &region, &err), err));
+        }
+
+        img_fmt = { CL_R, CL_FLOAT };
+        memset(&img_desc, 0, sizeof(img_desc));
+        img_desc.image_type  = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+        img_desc.image_width = need_pad ? (size_t)K * N_pad : (size_t)K * N;
+        img_desc.buffer      = need_pad ? b_padded : b_sub_buf;
+        cl_mem b_img;
+        CL_CHECK((b_img = clCreateImage(context, CL_MEM_READ_ONLY, &img_fmt, &img_desc, NULL, &err), err));
+
+        region.origin = offsetd;
+        region.size   = (size_t)M * N * sizeof(float);
+        cl_mem d_sub_buf;
+        CL_CHECK((d_sub_buf = clCreateSubBuffer(extrad->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, &region, &err), err));
+
+        img_fmt = { CL_R, CL_FLOAT };
+        memset(&img_desc, 0, sizeof(img_desc));
+        img_desc.image_type  = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+        img_desc.image_width = (size_t)M * N;
+        img_desc.buffer      = d_sub_buf;
+        cl_mem d_img;
+        CL_CHECK((d_img = clCreateImage(context, CL_MEM_WRITE_ONLY, &img_fmt, &img_desc, NULL, &err), err));
+
+        int line_stride_matrix_A_in_bytes = M * 4;
+        int line_stride_matrix_S_in_bytes = M * 2;
+        int line_stride_matrix_B_in_bytes = K * 4;
+        int line_stride_matrix_C_in_bytes = M * 4;
+
+        int c_offset_for_kernel = 0;
+        int b_offset_for_kernel = 0;
+
+        kernel = backend_ctx->kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin;
+
+        cl_uint k_arg = 0;
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &a_img));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &s_img));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &b_img));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int),    &b_offset_for_kernel));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &d_img));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &c_offset_for_kernel));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &K));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &line_stride_matrix_A_in_bytes));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &line_stride_matrix_S_in_bytes));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &line_stride_matrix_B_in_bytes));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &line_stride_matrix_C_in_bytes));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &M));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &N));
+
+        size_t local_work_size[3]  = { 64, 2, 2 };
+        size_t m_tiles = (size_t)CEIL_DIV(M, 64);
+        size_t global_work_size[3] = { 64, m_tiles, (size_t)CEIL_DIV(N_pad, gemm_tile_n) };
+        backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+
+        CL_CHECK(clReleaseMemObject(b_img));
+        if (b_sub_buf) {
+            CL_CHECK(clReleaseMemObject(b_sub_buf));
+        }
+        if (b_padded) {
+            CL_CHECK(clReleaseMemObject(b_padded));
+        }
+        CL_CHECK(clReleaseMemObject(d_img));
+        CL_CHECK(clReleaseMemObject(d_sub_buf));
+    }
+}
+#endif // GGML_OPENCL_USE_ADRENO_KERNELS
+
 static void ggml_cl_mul_mat_q4_0_f32_adreno(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
    GGML_ASSERT(src0);
@@ -18399,6 +18639,20 @@ static void ggml_cl_mul_mat_q4_0_f32_adreno(ggml_backend_t backend, const ggml_t
    static const bool q40_mc3 = (getenv("GGML_OPENCL_Q40_MC3") != nullptr);
    const bool use_q40_mc3 = q40_mc3 && (ne1 >= 2 && ne1 <= 4) && (ne01 < 32768);

+    const bool use_ila = use_q4_0_ila_kernels(backend_ctx, src0);
+
+    if (use_ila) {
+        if (use_q40_mc3) {
+            static bool warned = false;
+            if (!warned) {
+                GGML_LOG_WARN("ggml_opencl: GGML_OPENCL_Q40_MC3 is bypassed by Q4_0 binary kernels\n");
+                warned = true;
+            }
+        }
+        ggml_cl_mul_mat_q4_0_f32_adreno_ila(backend, src0, src1, dst);
+        return;
+    }
+
    if (ne1 == 1 || use_q40_mc3) {
        cl_mem q_img = nullptr;
        cl_mem b_sub_buf = nullptr;
@@ -0,0 +1,137 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+
+#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 QK4_0 32
+#define N_SIMDGROUP 4
+
+#define dequantizeBlockAccum_ila_1row_hi(total_sum, bits4, scale, y) \
+    float shared_y; \
+    shared_y = sub_group_broadcast(y.s0, 0); \
+    total_sum += ((bits4.s0 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s1, 0); \
+    total_sum += (((bits4.s0 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s2, 0); \
+    total_sum += (((bits4.s0 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s3, 0); \
+    total_sum += (((bits4.s0 & 0xF000) >> 12) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s4, 0); \
+    total_sum += ((bits4.s1 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s5, 0); \
+    total_sum += (((bits4.s1 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s6, 0); \
+    total_sum += (((bits4.s1 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s7, 0); \
+    total_sum += (((bits4.s1 & 0xF000) >> 12) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s0, 1); \
+    total_sum += ((bits4.s2 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s1, 1); \
+    total_sum += (((bits4.s2 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s2, 1); \
+    total_sum += (((bits4.s2 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s3, 1); \
+    total_sum += (((bits4.s2 & 0xF000) >> 12) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s4, 1); \
+    total_sum += ((bits4.s3 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s5, 1); \
+    total_sum += (((bits4.s3 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s6, 1); \
+    total_sum += (((bits4.s3 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s7, 1); \
+    total_sum += (((bits4.s3 & 0xF000) >> 12) - 8) * scale * shared_y;
+
+#define dequantizeBlockAccum_ila_1row_lo(total_sum, bits4, scale, y) \
+    shared_y = sub_group_broadcast(y.s0, 2); \
+    total_sum += ((bits4.s4 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s1, 2); \
+    total_sum += (((bits4.s4 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s2, 2); \
+    total_sum += (((bits4.s4 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s3, 2); \
+    total_sum += (((bits4.s4 & 0xF000) >> 12) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s4, 2); \
+    total_sum += ((bits4.s5 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s5, 2); \
+    total_sum += (((bits4.s5 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s6, 2); \
+    total_sum += (((bits4.s5 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s7, 2); \
+    total_sum += (((bits4.s5 & 0xF000) >> 12) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s0, 3); \
+    total_sum += ((bits4.s6 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s1, 3); \
+    total_sum += (((bits4.s6 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s2, 3); \
+    total_sum += (((bits4.s6 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s3, 3); \
+    total_sum += (((bits4.s6 & 0xF000) >> 12) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s4, 3); \
+    total_sum += ((bits4.s7 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s5, 3); \
+    total_sum += (((bits4.s7 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s6, 3); \
+    total_sum += (((bits4.s7 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s7, 3); \
+    total_sum += (((bits4.s7 & 0xF000) >> 12) - 8) * scale * shared_y;
+
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+__kernel void kernel_gemv_noshuffle_q4_0_f32_32b_trans(
+        __read_only  image1d_buffer_t src0_q,
+        global half  * src0_d,
+        __read_only  image1d_buffer_t src1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01)
+{
+    uint groupId = get_local_id(1);
+    uint gid     = get_global_id(0);
+    ushort slid  = get_sub_group_local_id();
+
+    uint K = ne00;
+    uint M = ne01;
+
+    __private uint4  regA;
+    __private half   regS;
+    __private float8 regB;
+    __private float  totalSum = 0.0f;
+
+    for (uint k = groupId; k < (K / QK4_0); k += N_SIMDGROUP) {
+        regS = src0_d[k * M + gid];
+        if (slid < 4) {
+            regB.s0123 = read_imagef(src1, (slid * 2 + k * 8));
+            regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8));
+        }
+        regA.s0 = read_imageui(src0_q, ((k * 4 + 0) * M + gid)).x;
+        regA.s1 = read_imageui(src0_q, ((k * 4 + 1) * M + gid)).x;
+        regA.s2 = read_imageui(src0_q, ((k * 4 + 2) * M + gid)).x;
+        regA.s3 = read_imageui(src0_q, ((k * 4 + 3) * M + gid)).x;
+
+        dequantizeBlockAccum_ila_1row_hi(totalSum, as_ushort8(regA), regS, regB);
+        dequantizeBlockAccum_ila_1row_lo(totalSum, as_ushort8(regA), regS, regB);
+    }
+
+    __local float reduceLM[SIMDGROUP_WIDTH * 3];
+    if (groupId == 1) reduceLM[SIMDGROUP_WIDTH * 0 + slid] = totalSum;
+    if (groupId == 2) reduceLM[SIMDGROUP_WIDTH * 1 + slid] = totalSum;
+    if (groupId == 3) reduceLM[SIMDGROUP_WIDTH * 2 + slid] = totalSum;
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (groupId == 0) totalSum += reduceLM[SIMDGROUP_WIDTH * 0 + slid];
+    if (groupId == 0) totalSum += reduceLM[SIMDGROUP_WIDTH * 1 + slid];
+    if (groupId == 0) totalSum += reduceLM[SIMDGROUP_WIDTH * 2 + slid];
+
+    if (groupId == 0) {
+        dst = (global float*)((global char*)dst + offsetd);
+        if (gid < M) {
+            dst[gid] = totalSum;
+        }
+    }
+}
@@ -33,7 +33,11 @@ void argsort(bool needs_bounds_check, const uint row) {
    const uint row_offset = row * p.ncols;

    // initialize indices
-    dst_row[col] = ivec2(col, floatBitsToInt(data_a[row_offset + col]));
+    ivec2 value = ivec2(col, 0);
+    if (!needs_bounds_check || col < p.ncols) {
+        value.y = floatBitsToInt(data_a[row_offset + col]);
+    }
+    dst_row[col] = value;
    barrier();

    uint num_outer_loop_iters = NCOLS_PADDED_LOG2;
@@ -42,18 +46,20 @@ void argsort(bool needs_bounds_check, const uint row) {
        [[unroll]] for (uint j = k / 2, inner_idx = 0; inner_idx < num_inner_loop_iters; j /= 2, inner_idx++) {
            const int ixj = int(col ^ j);

-            int idx_0 = (col & k) == 0 ? col : ixj;
-            int idx_1 = (col & k) == 0 ? ixj : col;
+            if (ixj > col) {
+                int idx_0 = (col & k) == 0 ? col : ixj;
+                int idx_1 = (col & k) == 0 ? ixj : col;

-            ivec2 sh_idx_0 = dst_row[idx_0];
-            ivec2 sh_idx_1 = dst_row[idx_1];
-            bool idx_0_oob = needs_bounds_check ? sh_idx_0.x >= p.ncols : false;
-            bool idx_1_oob = needs_bounds_check ? sh_idx_1.x >= p.ncols : false;
+                ivec2 sh_idx_0 = dst_row[idx_0];
+                ivec2 sh_idx_1 = dst_row[idx_1];
+                bool idx_0_oob = needs_bounds_check ? sh_idx_0.x >= p.ncols : false;
+                bool idx_1_oob = needs_bounds_check ? sh_idx_1.x >= p.ncols : false;

-            if ((idx_0_oob ||
-                (!idx_1_oob && intBitsToFloat(sh_idx_0.y) > intBitsToFloat(sh_idx_1.y))) && (ixj > col)) {
-                dst_row[idx_0] = sh_idx_1;
-                dst_row[idx_1] = sh_idx_0;
+                if (idx_0_oob ||
+                    (!idx_1_oob && intBitsToFloat(sh_idx_0.y) > intBitsToFloat(sh_idx_1.y))) {
+                    dst_row[idx_0] = sh_idx_1;
+                    dst_row[idx_1] = sh_idx_0;
+                }
            }

            barrier();
@@ -42,7 +42,10 @@ void argsort(bool needs_bounds_check, const uint row) {
        [[unroll]] for (int u = 0; u < WG_UNROLL_FACTOR; ++u) {
            uint c = u*BLOCK_SIZE + col;
            if (c < p.ncols_padded) {
-                ivec2 v = ivec2(c, floatBitsToInt(data_a[row_offset + c]));
+                ivec2 v = ivec2(c, 0);
+                if (!needs_bounds_check || c < p.ncols) {
+                    v.y = floatBitsToInt(data_a[row_offset + c]);
+                }
                tmp_idx[idx_offset + c] = v;
            }
        }
@@ -260,6 +260,20 @@ float16_t dequantFuncTQ1_0(const in decodeBufTQ1_0 bl, const in uint blockCoords
    return bl.block.d * (float16_t(int(xi)) - float16_t(1.0));
 }

+f16vec4 dequantFuncTQ1_0_v(const in decodeBufTQ1_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
+{
+    const uint e = coordInBlock[1];
+    f16vec4 v;
+    [[unroll]] for (uint k = 0u; k < 4u; ++k) {
+        const uint ee = e + k;
+        const uint bidx = tq1_0_byte_of(ee);
+        const uint qbyte = uint(bidx < 48u ? bl.block.qs[bidx] : bl.block.qh[bidx - 48u]);
+        const uint xi = tq1_0_trit(qbyte, tq1_0_digit_of(ee));
+        v[k] = bl.block.d * (float16_t(int(xi)) - float16_t(1.0));
+    }
+    return v;
+}
+
 layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufTQ2_0 {
   block_tq2_0 block;
 };
@@ -1054,7 +1068,7 @@ float16_t dequantFuncIQ2_S(const in decodeBufIQ2_S bl, const in uint blockCoords
    const uint scale = (bl.block.scales[ib32] >> ((idx & 0x10) >> 2)) & 0xf;
    const uint qs = bl.block.qs[ib8];
    const uint qh = bl.block.qh[ib32];
-    const uint sign = bl.block.qs[QUANT_K / 8 + ib8] >> (idx & 0x6);
+    const uint sign = bl.block.qs[QUANT_K_IQ2_S / 8 + ib8] >> (idx & 0x6);

    const float d = float(bl.block.d);
    const float db = d * 0.25 * (0.5 + scale);
@@ -1076,7 +1090,7 @@ f16vec4 dequantFuncIQ2_S_v(const in decodeBufIQ2_S bl, const in uint blockCoords
    const uint scale = (bl.block.scales[ib32] >> ((idx & 0x10) >> 2)) & 0xf;
    const uint qs    = bl.block.qs[ib8];
    const uint qh    = bl.block.qh[ib32];
-    const uint sb    = uint(bl.block.qs[QUANT_K / 8 + ib8]) >> (idx & 0x6u);
+    const uint sb    = uint(bl.block.qs[QUANT_K_IQ2_S / 8 + ib8]) >> (idx & 0x6u);

    const float d  = float(bl.block.d);
    const float db = d * 0.25 * (0.5 + scale);
@@ -1107,7 +1121,7 @@ float16_t dequantFuncIQ3_XXS(const in decodeBufIQ3_XXS bl, const in uint blockCo
    uint idx = coordInBlock[1];

    const uint iqs = (idx & 0xFC) >> 2;             // 0..63
-    const uint is = QUANT_K / 4 + ((idx & 0xE0) >> 3);// 8 values
+    const uint is = QUANT_K_IQ3_XXS / 4 + ((idx & 0xE0) >> 3);// 8 values

    const float d = float(bl.block.d);
    const uint qs = bl.block.qs[iqs];
@@ -1130,7 +1144,7 @@ f16vec4 dequantFuncIQ3_XXS_v(const in decodeBufIQ3_XXS bl, const in uint blockCo
    const uint idx = coordInBlock[1];

    const uint iqs = idx >> 2;
-    const uint is  = QUANT_K / 4 + ((idx & 0xE0) >> 3);
+    const uint is  = QUANT_K_IQ3_XXS / 4 + ((idx & 0xE0) >> 3);

    const float d     = float(bl.block.d);
    const uint  qs    = bl.block.qs[iqs];
@@ -1,32 +1,22 @@
 #if !defined(GGML_FA_TYPES_COMP)
 #define GGML_FA_TYPES_COMP

-// FaTypeK / FaTypeV spec constant values. These mirror enum ggml_type so the
-// host can pass the type directly. Keep in sync with ggml.h.
-#define FA_TYPE_F32   0u
-#define FA_TYPE_F16   1u
-#define FA_TYPE_Q4_0  2u
-#define FA_TYPE_Q4_1  3u
-#define FA_TYPE_Q5_0  6u
-#define FA_TYPE_Q5_1  7u
-#define FA_TYPE_Q8_0  8u
-#define FA_TYPE_IQ4_NL 20u
-#define FA_TYPE_BF16 30u
+#include "ggml_type_ids.glsl"

 // Number of matrix elements per buffer block, derived from the K/V type spec
 // constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1
 // and bypasses the dequant path entirely. Quants follow their ggml block sizes.
 uint fa_block_elems(uint ty) {
    switch (ty) {
-        case FA_TYPE_F32:  return 4u;
-        case FA_TYPE_F16:  return 1u;
-        case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0);
-        case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1);
-        case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0);
-        case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1);
-        case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0);
-        case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL);
-        case FA_TYPE_BF16: return 1u;
+        case GGML_TYPE_F32:  return 4u;
+        case GGML_TYPE_F16:  return 1u;
+        case GGML_TYPE_Q4_0: return uint(QUANT_K_Q4_0);
+        case GGML_TYPE_Q4_1: return uint(QUANT_K_Q4_1);
+        case GGML_TYPE_Q5_0: return uint(QUANT_K_Q5_0);
+        case GGML_TYPE_Q5_1: return uint(QUANT_K_Q5_1);
+        case GGML_TYPE_Q8_0: return uint(QUANT_K_Q8_0);
+        case GGML_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL);
+        case GGML_TYPE_BF16: return 1u;
        default:           return 1u;
    }
 }
@@ -36,18 +26,18 @@ uint fa_block_elems(uint ty) {
 // of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R.
 uint fa_quant_r_mmq(uint ty) {
    switch (ty) {
-        case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0);
-        case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1);
-        case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0);
-        case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1);
-        case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0);
+        case GGML_TYPE_Q4_0: return uint(QUANT_R_Q4_0);
+        case GGML_TYPE_Q4_1: return uint(QUANT_R_Q4_1);
+        case GGML_TYPE_Q5_0: return uint(QUANT_R_Q5_0);
+        case GGML_TYPE_Q5_1: return uint(QUANT_R_Q5_1);
+        case GGML_TYPE_Q8_0: return uint(QUANT_R_Q8_0);
        default:           return 1u;
    }
 }

 bool fa_type_needs_shmem(uint ty) {
    switch (ty) {
-        case FA_TYPE_IQ4_NL: return true;
+        case GGML_TYPE_IQ4_NL: return true;
        default:             return false;
    }
 }
@@ -8,7 +8,9 @@ layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
 layout (binding = 0) writeonly buffer D {D_TYPE data_d[];};

 void main() {
-    const uint i = gl_GlobalInvocationID.x;
+    // 2D grid flattening: each x workgroup covers gl_WorkGroupSize.x elements,
+    // each y workgroup covers gl_NumWorkGroups.x * gl_WorkGroupSize.x elements.
+    const uint i = (gl_GlobalInvocationID.y * gl_NumWorkGroups.x * gl_WorkGroupSize.x) + gl_GlobalInvocationID.x;

    if (i >= p.KX) {
        return;
@@ -134,7 +134,7 @@ void main() {

        // Q8_0 K only needs (qd, _); the asymmetric Q4_*/Q5_* family also stores
        // the row-sum scaled by qd, used in k_dot_correction.
-        if (FaTypeK == FA_TYPE_Q8_0) {
+        if (FaTypeK == GGML_TYPE_Q8_0) {
            if (buf_iqs == 0) {
                Qf[buf_ib].ds = FLOAT_TYPEV2(qd, 0.0f);
            }
@@ -367,7 +367,7 @@ void main() {
                // Q4_*/Q5_* take the block-8 fast path when one step covers a full
                // block; Q8_0 always goes through the per-int get_k_qs* helpers
                // (its qs is byte-packed, not nibble-packed).
-                const bool block8_fast = (d_per_step == 8) && (FaTypeK != FA_TYPE_Q8_0);
+                const bool block8_fast = (d_per_step == 8) && (FaTypeK != GGML_TYPE_Q8_0);

                if (SHMEM_STAGING != 0) {
                    const uint k_block_idx = (d_tid * (HSK_per_thread / 4) + d_block) / 8;
@@ -375,7 +375,7 @@ void main() {
                    k_dm = ACC_TYPEV2(kblocksh[buf_ib].dm);

                    if (block8_fast) {
-                        const bool has_qh = (FaTypeK == FA_TYPE_Q5_0) || (FaTypeK == FA_TYPE_Q5_1);
+                        const bool has_qh = (FaTypeK == GGML_TYPE_Q5_0) || (FaTypeK == GGML_TYPE_Q5_1);
                        [[unroll]] for (uint32_t d = 0; d < 4; d++) {
                            uint vui = kblocksh[buf_ib].qs[d];
                            k_quants[d    ] = int32_t( vui       & 0x0F0F0F0F);
@@ -105,8 +105,8 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
 #define BLOCK_SIZE_V fa_block_elems(FaTypeV)
 // F16 reads f16 elements directly from the binding; everything else routes
 // through dequantize4 / the MMQ helpers to unpack from the packed block layout.
-#define USE_DECODE_K (FaTypeK != FA_TYPE_F16)
-#define USE_DECODE_V (FaTypeV != FA_TYPE_F16)
+#define USE_DECODE_K (FaTypeK != GGML_TYPE_F16)
+#define USE_DECODE_V (FaTypeV != GGML_TYPE_F16)

 #define CEIL_DIV(a, b) (((a) + (b) - 1) / (b))

@@ -40,26 +40,28 @@ layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_
 #if !defined(BFLOAT16)
 float16_t faDecodeK(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
    switch (FaTypeK) {
-        case FA_TYPE_F32:  return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_F32:  return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
        default: return float16_t(0);
    }
 }

 float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
    switch (FaTypeV) {
-        case FA_TYPE_F32:  return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_F32:  return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
        default: return float16_t(0);
    }
 }
@@ -67,26 +69,26 @@ float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const
 // V=4 vector decode for K/V; dispatches to per-format _v decoders.
 f16vec4 faDecodeKVector(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
    switch (FaTypeK) {
-        case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
-        case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
+        case GGML_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
        default: return f16vec4(0);
    }
 }

 f16vec4 faDecodeVVector(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
    switch (FaTypeV) {
-        case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
-        case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
-        case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
+        case GGML_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
+        case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
        default: return f16vec4(0);
    }
 }
@@ -121,25 +121,25 @@ layout (binding = 1) readonly buffer K_PACKED_Q5_1_P32 { block_q5_1_packed32 dat
 FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) {
    if (binding_idx == BINDING_IDX_K) {
        switch (FaTypeK) {
-            case FA_TYPE_F32:  FA_DEQUANT4_F32 (k_packed_f32)
-            case FA_TYPE_Q4_0: FA_DEQUANT4_Q4_0(k_packed_q4_0)
-            case FA_TYPE_Q4_1: FA_DEQUANT4_Q4_1(k_packed_q4_1)
-            case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(k_packed_q5_0)
-            case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(k_packed_q5_1)
-            case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(k_packed_q8_0)
-            case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl)
-            case FA_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16)
+            case GGML_TYPE_F32:  FA_DEQUANT4_F32 (k_packed_f32)
+            case GGML_TYPE_Q4_0: FA_DEQUANT4_Q4_0(k_packed_q4_0)
+            case GGML_TYPE_Q4_1: FA_DEQUANT4_Q4_1(k_packed_q4_1)
+            case GGML_TYPE_Q5_0: FA_DEQUANT4_Q5_0(k_packed_q5_0)
+            case GGML_TYPE_Q5_1: FA_DEQUANT4_Q5_1(k_packed_q5_1)
+            case GGML_TYPE_Q8_0: FA_DEQUANT4_Q8_0(k_packed_q8_0)
+            case GGML_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl)
+            case GGML_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16)
        }
    } else {
        switch (FaTypeV) {
-            case FA_TYPE_F32:  FA_DEQUANT4_F32 (v_packed_f32)
-            case FA_TYPE_Q4_0: FA_DEQUANT4_Q4_0(v_packed_q4_0)
-            case FA_TYPE_Q4_1: FA_DEQUANT4_Q4_1(v_packed_q4_1)
-            case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(v_packed_q5_0)
-            case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(v_packed_q5_1)
-            case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(v_packed_q8_0)
-            case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(v_packed_iq4_nl)
-            case FA_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16)
+            case GGML_TYPE_F32:  FA_DEQUANT4_F32 (v_packed_f32)
+            case GGML_TYPE_Q4_0: FA_DEQUANT4_Q4_0(v_packed_q4_0)
+            case GGML_TYPE_Q4_1: FA_DEQUANT4_Q4_1(v_packed_q4_1)
+            case GGML_TYPE_Q5_0: FA_DEQUANT4_Q5_0(v_packed_q5_0)
+            case GGML_TYPE_Q5_1: FA_DEQUANT4_Q5_1(v_packed_q5_1)
+            case GGML_TYPE_Q8_0: FA_DEQUANT4_Q8_0(v_packed_q8_0)
+            case GGML_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(v_packed_iq4_nl)
+            case GGML_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16)
        }
    }
    return FLOAT_TYPEV4(0);
@@ -4,20 +4,20 @@

 int32_t get_k_qs(uint ib, uint iqs, uint a_offset) {
    switch (FaTypeK) {
-        case FA_TYPE_Q4_0: {
+        case GGML_TYPE_Q4_0: {
            uint vui = pack32(u16vec2(k_packed_q4_0.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 0],
                                      k_packed_q4_0.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 1]));
            uint shift = (iqs & 0x10) >> 2;
            vui >>= shift;
            return int32_t(vui & 0x0F0F0F0F);
        }
-        case FA_TYPE_Q4_1: { // uses packed32 alias
+        case GGML_TYPE_Q4_1: { // uses packed32 alias
            uint vui = k_packed_q4_1_p32.data[a_offset + ib].qs[(iqs & 0xF) / 4];
            uint shift = (iqs & 0x10) >> 2;
            vui >>= shift;
            return int32_t(vui & 0x0F0F0F0F);
        }
-        case FA_TYPE_Q5_0: {
+        case GGML_TYPE_Q5_0: {
            uint vui = pack32(u16vec2(k_packed_q5_0.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 0],
                                      k_packed_q5_0.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 1]));
            uint qh = pack32(u16vec2(k_packed_q5_0.data[a_offset + ib].qh[0],
@@ -27,7 +27,7 @@ int32_t get_k_qs(uint ib, uint iqs, uint a_offset) {
            uint qh_bits = (qh >> iqs) & 0xF;
            return int32_t(vui & 0x0F0F0F0F) | int32_t((qh_bits * 0x02040810u) & 0x10101010u);
        }
-        case FA_TYPE_Q5_1: { // qs via packed32, qh via packed16
+        case GGML_TYPE_Q5_1: { // qs via packed32, qh via packed16
            uint vui = k_packed_q5_1_p32.data[a_offset + ib].qs[(iqs & 0xF) / 4];
            uint qh  = k_packed_q5_1.data[a_offset + ib].qh;
            uint shift = (iqs & 0x10) >> 2;
@@ -35,7 +35,7 @@ int32_t get_k_qs(uint ib, uint iqs, uint a_offset) {
            uint qh_bits = (qh >> iqs) & 0xF;
            return int32_t(vui & 0x0F0F0F0F) | int32_t((qh_bits * 0x02040810u) & 0x10101010u);
        }
-        case FA_TYPE_Q8_0: {
+        case GGML_TYPE_Q8_0: {
            return pack32(i16vec2(k_packed_q8_0.data[a_offset + ib].qs[iqs / 2],
                                  k_packed_q8_0.data[a_offset + ib].qs[iqs / 2 + 1]));
        }
@@ -47,11 +47,11 @@ int32_t get_k_qs(uint ib, uint iqs, uint a_offset) {
 // return (d, 0) so call sites always see the same shape.
 FLOAT_TYPEV2 get_k_scale(uint ib, uint a_offset) {
    switch (FaTypeK) {
-        case FA_TYPE_Q4_0: return FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q4_0.data[a_offset + ib].d), 0.0);
-        case FA_TYPE_Q4_1: return FLOAT_TYPEV2(k_packed_q4_1_p32.data[a_offset + ib].dm);
-        case FA_TYPE_Q5_0: return FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q5_0.data[a_offset + ib].d), 0.0);
-        case FA_TYPE_Q5_1: return FLOAT_TYPEV2(k_packed_q5_1_p32.data[a_offset + ib].dm);
-        case FA_TYPE_Q8_0: return FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q8_0.data[a_offset + ib].d), 0.0);
+        case GGML_TYPE_Q4_0: return FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q4_0.data[a_offset + ib].d), 0.0);
+        case GGML_TYPE_Q4_1: return FLOAT_TYPEV2(k_packed_q4_1_p32.data[a_offset + ib].dm);
+        case GGML_TYPE_Q5_0: return FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q5_0.data[a_offset + ib].d), 0.0);
+        case GGML_TYPE_Q5_1: return FLOAT_TYPEV2(k_packed_q5_1_p32.data[a_offset + ib].dm);
+        case GGML_TYPE_Q8_0: return FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q8_0.data[a_offset + ib].d), 0.0);
        default: return FLOAT_TYPEV2(0);
    }
 }
@@ -61,16 +61,16 @@ void k_block_to_shmem(const uint buf_ib, const uint global_ib, const uint iqs, c
    // explicit casts. The bit pattern is what we care about here -- the actual
    // signed/unsigned interpretation happens downstream in the dot product.
    switch (FaTypeK) {
-        case FA_TYPE_Q4_0: {
+        case GGML_TYPE_Q4_0: {
            kblocksh[buf_ib].qs[iqs] = int32_t(pack32(u16vec2(k_packed_q4_0.data[a_offset + global_ib].qs[iqs * 2],
                                                              k_packed_q4_0.data[a_offset + global_ib].qs[iqs * 2 + 1])));
            break;
        }
-        case FA_TYPE_Q4_1: {
+        case GGML_TYPE_Q4_1: {
            kblocksh[buf_ib].qs[iqs] = int32_t(k_packed_q4_1_p32.data[a_offset + global_ib].qs[iqs]);
            break;
        }
-        case FA_TYPE_Q5_0: {
+        case GGML_TYPE_Q5_0: {
            kblocksh[buf_ib].qs[iqs] = int32_t(pack32(u16vec2(k_packed_q5_0.data[a_offset + global_ib].qs[iqs * 2],
                                                              k_packed_q5_0.data[a_offset + global_ib].qs[iqs * 2 + 1])));
            if (iqs == 0) {
@@ -79,14 +79,14 @@ void k_block_to_shmem(const uint buf_ib, const uint global_ib, const uint iqs, c
            }
            break;
        }
-        case FA_TYPE_Q5_1: {
+        case GGML_TYPE_Q5_1: {
            kblocksh[buf_ib].qs[iqs] = int32_t(k_packed_q5_1_p32.data[a_offset + global_ib].qs[iqs]);
            if (iqs == 0) {
                kblocksh[buf_ib].qh = k_packed_q5_1.data[a_offset + global_ib].qh;
            }
            break;
        }
-        case FA_TYPE_Q8_0: {
+        case GGML_TYPE_Q8_0: {
            kblocksh[buf_ib].qs[iqs] = pack32(i16vec2(k_packed_q8_0.data[a_offset + global_ib].qs[iqs * 2],
                                                      k_packed_q8_0.data[a_offset + global_ib].qs[iqs * 2 + 1]));
            break;
@@ -96,11 +96,11 @@ void k_block_to_shmem(const uint buf_ib, const uint global_ib, const uint iqs, c
    if (iqs == 0) {
        // Q4_0/Q5_0/Q8_0 store dm.x = d; Q4_1/Q5_1 store dm = (d, m) pair.
        switch (FaTypeK) {
-            case FA_TYPE_Q4_0: kblocksh[buf_ib].dm = FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q4_0.data[a_offset + global_ib].d), 0.0); break;
-            case FA_TYPE_Q4_1: kblocksh[buf_ib].dm = FLOAT_TYPEV2(k_packed_q4_1_p32.data[a_offset + global_ib].dm); break;
-            case FA_TYPE_Q5_0: kblocksh[buf_ib].dm = FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q5_0.data[a_offset + global_ib].d), 0.0); break;
-            case FA_TYPE_Q5_1: kblocksh[buf_ib].dm = FLOAT_TYPEV2(k_packed_q5_1_p32.data[a_offset + global_ib].dm); break;
-            case FA_TYPE_Q8_0: kblocksh[buf_ib].dm = FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q8_0.data[a_offset + global_ib].d), 0.0); break;
+            case GGML_TYPE_Q4_0: kblocksh[buf_ib].dm = FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q4_0.data[a_offset + global_ib].d), 0.0); break;
+            case GGML_TYPE_Q4_1: kblocksh[buf_ib].dm = FLOAT_TYPEV2(k_packed_q4_1_p32.data[a_offset + global_ib].dm); break;
+            case GGML_TYPE_Q5_0: kblocksh[buf_ib].dm = FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q5_0.data[a_offset + global_ib].d), 0.0); break;
+            case GGML_TYPE_Q5_1: kblocksh[buf_ib].dm = FLOAT_TYPEV2(k_packed_q5_1_p32.data[a_offset + global_ib].dm); break;
+            case GGML_TYPE_Q8_0: kblocksh[buf_ib].dm = FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q8_0.data[a_offset + global_ib].d), 0.0); break;
        }
    }
 }
@@ -121,31 +121,31 @@ struct fa_k_qs_block8 {
 fa_k_qs_block8 get_k_qs_block8(uint ib, uint a_offset) {
    fa_k_qs_block8 r;
    uint qh = 0;
-    if (FaTypeK == FA_TYPE_Q5_0) {
+    if (FaTypeK == GGML_TYPE_Q5_0) {
        qh = pack32(u16vec2(k_packed_q5_0.data[a_offset + ib].qh[0],
                            k_packed_q5_0.data[a_offset + ib].qh[1]));
-    } else if (FaTypeK == FA_TYPE_Q5_1) {
+    } else if (FaTypeK == GGML_TYPE_Q5_1) {
        qh = k_packed_q5_1.data[a_offset + ib].qh;
    }
-    const bool has_qh = (FaTypeK == FA_TYPE_Q5_0) || (FaTypeK == FA_TYPE_Q5_1);
+    const bool has_qh = (FaTypeK == GGML_TYPE_Q5_0) || (FaTypeK == GGML_TYPE_Q5_1);
    [[unroll]] for (uint32_t d = 0; d < 4; d++) {
        uint vui = 0;
        switch (FaTypeK) {
-            case FA_TYPE_Q4_0: { // packed16
+            case GGML_TYPE_Q4_0: { // packed16
                vui = pack32(u16vec2(k_packed_q4_0.data[a_offset + ib].qs[d * 2 + 0],
                                     k_packed_q4_0.data[a_offset + ib].qs[d * 2 + 1]));
                break;
            }
-            case FA_TYPE_Q4_1: { // packed32 alias
+            case GGML_TYPE_Q4_1: { // packed32 alias
                vui = k_packed_q4_1_p32.data[a_offset + ib].qs[d];
                break;
            }
-            case FA_TYPE_Q5_0: { // packed16
+            case GGML_TYPE_Q5_0: { // packed16
                vui = pack32(u16vec2(k_packed_q5_0.data[a_offset + ib].qs[d * 2 + 0],
                                     k_packed_q5_0.data[a_offset + ib].qs[d * 2 + 1]));
                break;
            }
-            case FA_TYPE_Q5_1: { // packed32 alias
+            case GGML_TYPE_Q5_1: { // packed32 alias
                vui = k_packed_q5_1_p32.data[a_offset + ib].qs[d];
                break;
            }
@@ -164,21 +164,21 @@ fa_k_qs_block8 get_k_qs_block8(uint ib, uint a_offset) {

 int32_t get_k_qs_shmem(const uint buf_ib, const uint pos) {
    switch (FaTypeK) {
-        case FA_TYPE_Q4_0:
-        case FA_TYPE_Q4_1: {
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q4_1: {
            uint sub = pos % 4;
            uint shift = ((pos % 8) >= 4) ? 4u : 0u;
            return int32_t((uint(kblocksh[buf_ib].qs[sub]) >> shift) & 0x0F0F0F0Fu);
        }
-        case FA_TYPE_Q5_0:
-        case FA_TYPE_Q5_1: {
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q5_1: {
            uint sub = pos % 4;
            uint shift = ((pos % 8) >= 4) ? 4u : 0u;
            int32_t result = int32_t((uint(kblocksh[buf_ib].qs[sub]) >> shift) & 0x0F0F0F0Fu);
            uint qh_bits = (kblocksh[buf_ib].qh >> (pos * 4u)) & 0xFu;
            return result | int32_t((qh_bits * 0x02040810u) & 0x10101010u);
        }
-        case FA_TYPE_Q8_0: {
+        case GGML_TYPE_Q8_0: {
            return kblocksh[buf_ib].qs[pos];
        }
        default: return 0;
@@ -187,10 +187,10 @@ int32_t get_k_qs_shmem(const uint buf_ib, const uint pos) {

 ACC_TYPE k_dot_correction(const uint qib, const ACC_TYPEV2 k_dm) {
    switch (FaTypeK) {
-        case FA_TYPE_Q4_0: return -ACC_TYPE(8.0)  * ACC_TYPE(Qf[qib].ds.y) * k_dm.x;
-        case FA_TYPE_Q5_0: return -ACC_TYPE(16.0) * ACC_TYPE(Qf[qib].ds.y) * k_dm.x;
-        case FA_TYPE_Q4_1:
-        case FA_TYPE_Q5_1: return ACC_TYPE(Qf[qib].ds.y) * k_dm.y;
+        case GGML_TYPE_Q4_0: return -ACC_TYPE(8.0)  * ACC_TYPE(Qf[qib].ds.y) * k_dm.x;
+        case GGML_TYPE_Q5_0: return -ACC_TYPE(16.0) * ACC_TYPE(Qf[qib].ds.y) * k_dm.x;
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q5_1: return ACC_TYPE(Qf[qib].ds.y) * k_dm.y;
        default: return ACC_TYPE(0.0);
    }
 }
@@ -0,0 +1,34 @@
+#if !defined(GGML_TYPE_IDS_COMP)
+#define GGML_TYPE_IDS_COMP
+
+// ggml_type enum values — must match ggml.h
+#define GGML_TYPE_F32      0u
+#define GGML_TYPE_F16      1u
+#define GGML_TYPE_Q4_0     2u
+#define GGML_TYPE_Q4_1     3u
+#define GGML_TYPE_Q5_0     6u
+#define GGML_TYPE_Q5_1     7u
+#define GGML_TYPE_Q8_0     8u
+#define GGML_TYPE_Q2_K    10u
+#define GGML_TYPE_Q3_K    11u
+#define GGML_TYPE_Q4_K    12u
+#define GGML_TYPE_Q5_K    13u
+#define GGML_TYPE_Q6_K    14u
+#define GGML_TYPE_IQ2_XXS 16u
+#define GGML_TYPE_IQ2_XS  17u
+#define GGML_TYPE_IQ3_XXS 18u
+#define GGML_TYPE_IQ1_S   19u
+#define GGML_TYPE_IQ4_NL  20u
+#define GGML_TYPE_IQ3_S   21u
+#define GGML_TYPE_IQ2_S   22u
+#define GGML_TYPE_IQ4_XS  23u
+#define GGML_TYPE_IQ1_M   29u
+#define GGML_TYPE_BF16    30u
+#define GGML_TYPE_TQ1_0   34u
+#define GGML_TYPE_TQ2_0   35u
+#define GGML_TYPE_MXFP4   39u
+#define GGML_TYPE_NVFP4   40u
+#define GGML_TYPE_Q1_0    41u
+#define GGML_TYPE_Q2_0    42u
+
+#endif // !defined(GGML_TYPE_IDS_COMP)
@@ -0,0 +1,2 @@
+void init_iq_shmem(uvec3 wgsize) {
+}
@@ -12,9 +12,9 @@

 #include "types.glsl"
 #include "fa_types.glsl"
-#define FaTypeV FA_TYPE_F32
+#define FaTypeV GGML_TYPE_F32

-layout(constant_id = 0) const uint FaTypeK = FA_TYPE_F32;
+layout(constant_id = 0) const uint FaTypeK = GGML_TYPE_F32;
 layout(constant_id = 1) const uint FaBlockBytesK = 4;
 layout(constant_id = 2) const uint SUBGROUP_SIZE = 32;

@@ -84,11 +84,11 @@ void main() {
    const uint k_block_elems = fa_block_elems(FaTypeK);
    const uint k_elem_bytes = FaBlockBytesK / k_block_elems;

-    if (FaTypeK == FA_TYPE_F16) {
+    if (FaTypeK == GGML_TYPE_F16) {
        k_row[tid] = float(k_f16[k_offset / k_elem_bytes + tid]);
-    } else if (FaTypeK == FA_TYPE_F32) {
+    } else if (FaTypeK == GGML_TYPE_F32) {
        k_row[tid] = k_f32[k_offset / k_elem_bytes + tid];
-    } else if (FaTypeK == FA_TYPE_BF16) {
+    } else if (FaTypeK == GGML_TYPE_BF16) {
        k_row[tid] = bf16_to_fp32(uint(k_bf16[k_offset / k_elem_bytes + tid]));
    } else if (4 * tid < HEAD_SIZE) {
        const uint coord = 4 * tid;
@@ -9,6 +9,9 @@
 #if defined(DATA_A_IQ1_M)
 #extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
 #endif
+#if !defined(DATA_A_F32) && !defined(DATA_A_F16) && !defined(DATA_A_BF16)
+#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
+#endif

 #if defined(DATA_A_BF16) && defined(COOPMAT)
 #extension GL_EXT_bfloat16 : enable
@@ -28,24 +31,54 @@
 #extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
 #endif

+#ifdef MULMAT_QUANT
+#include "ggml_type_ids.glsl"
+layout (constant_id = 12) const uint MmTypeA = 0;
+#endif
+
 #include "types.glsl"
 #include "dot_product_funcs.glsl"

+#ifndef MULMAT_QUANT
 #ifndef LOAD_VEC_A
 #define LOAD_VEC_A 1
 #endif
+#endif
 #ifndef LOAD_VEC_B
 #define LOAD_VEC_B 1
 #endif

 layout (constant_id = 11) const uint ALIGNED = 0;

+#ifdef MULMAT_QUANT
+
+uint mm_load_vec_a() {
+    switch (MmTypeA) {
+    case GGML_TYPE_Q1_0:
+    case GGML_TYPE_Q4_0:
+    case GGML_TYPE_Q4_1:
+    case GGML_TYPE_Q5_1:
+        return 8u;
+    case GGML_TYPE_Q2_0:
+    case GGML_TYPE_Q5_0:
+    case GGML_TYPE_Q8_0:
+    case GGML_TYPE_Q2_K:
+    case GGML_TYPE_Q4_K:
+    case GGML_TYPE_Q5_K:
+        return 4u;
+    default:
+        return 2u;
+    }
+}
+#endif
+
 #if !defined(TO_FLOAT_TYPE)
 #define TO_FLOAT_TYPE FLOAT_TYPE
 #endif

 layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;

+#ifndef MULMAT_QUANT
 layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
 #if defined(DATA_A_F32)
 layout (binding = 0) readonly buffer A_SCALAR {float data_a_scalar[];};
@@ -60,6 +93,30 @@ layout (binding = 0) readonly buffer A_PACKED16 {A_TYPE_PACKED16 data_a_packed16
 #if defined(A_TYPE_PACKED32)
 layout (binding = 0) readonly buffer A_PACKED32 {A_TYPE_PACKED32 data_a_packed32[];};
 #endif
+#else
+// Unpacked struct aliases
+layout (binding = 0) readonly buffer BUF_Q1_0     { block_q1_0 data[];     } a_q1_0;
+layout (binding = 0) readonly buffer BUF_Q2_0     { block_q2_0 data[];     } a_q2_0;
+layout (binding = 0) readonly buffer BUF_Q2_K     { block_q2_K data[];     } a_q2_k;
+layout (binding = 0) readonly buffer BUF_Q3_K     { block_q3_K data[];     } a_q3_k;
+layout (binding = 0) readonly buffer BUF_Q4_K     { block_q4_K data[];     } a_q4_k;
+layout (binding = 0) readonly buffer BUF_Q5_K     { block_q5_K data[];     } a_q5_k;
+layout (binding = 0) readonly buffer BUF_Q6_K     { block_q6_K data[];     } a_q6_k;
+layout (binding = 0) readonly buffer BUF_TQ1_0    { block_tq1_0 data[];    } a_tq1_0;
+layout (binding = 0) readonly buffer BUF_TQ2_0    { block_tq2_0 data[];    } a_tq2_0;
+// Packed16 aliases
+layout (binding = 0) readonly buffer BUF_Q4_0_P16   { block_q4_0_packed16 data[];   } a_q4_0_p16;
+layout (binding = 0) readonly buffer BUF_Q5_0_P16   { block_q5_0_packed16 data[];   } a_q5_0_p16;
+layout (binding = 0) readonly buffer BUF_Q8_0_P16   { block_q8_0_packed16 data[];   } a_q8_0_p16;
+layout (binding = 0) readonly buffer BUF_Q3_K_P16   { block_q3_K_packed16 data[];   } a_q3_k_p16;
+layout (binding = 0) readonly buffer BUF_Q6_K_P16   { block_q6_K_packed16 data[];   } a_q6_k_p16;
+// Packed32 aliases
+layout (binding = 0) readonly buffer BUF_Q4_1_P32   { block_q4_1_packed32 data[];   } a_q4_1_p32;
+layout (binding = 0) readonly buffer BUF_Q5_1_P32   { block_q5_1_packed32 data[];   } a_q5_1_p32;
+layout (binding = 0) readonly buffer BUF_Q2_K_P32   { block_q2_K_packed32 data[];   } a_q2_k_p32;
+layout (binding = 0) readonly buffer BUF_Q4_K_P32   { block_q4_K_packed32 data[];   } a_q4_k_p32;
+layout (binding = 0) readonly buffer BUF_Q5_K_P32   { block_q5_K_packed32 data[];   } a_q5_k_p32;
+#endif

 layout (binding = 1) readonly buffer B {B_TYPE data_b[];};
 layout (binding = 1) readonly buffer B_SCALAR {B_TYPE_SCALAR data_b_scalar[];};
@@ -121,8 +178,13 @@ layout (constant_id = 3) const uint BK = 16;  // Assumed to be 32 if working wit
 #endif

 #ifdef COOPMAT
+#ifdef MULMAT_QUANT
+layout(constant_id = 13) const uint SHMEM_STRIDE_PAD = 4;
+layout(constant_id = 14) const bool APPLY_SLM_A_RESHAPE = false;
+#else
 layout(constant_id = 12) const uint SHMEM_STRIDE_PAD = 4;
 layout(constant_id = 13) const bool APPLY_SLM_A_RESHAPE = false;
+#endif
 #else
 const uint SHMEM_STRIDE_PAD = 1;
 const bool APPLY_SLM_A_RESHAPE = false;
@@ -141,6 +203,10 @@ shared ACC_TYPE coopmat_stage[TM * TN * NUM_WARPS];
 #include "mul_mm_id_funcs.glsl"
 #include "mul_mm_funcs.glsl"

+#ifdef MULMAT_QUANT
+#include "iq_shmem_init.glsl"
+#endif
+
 void main() {
    const uint ic = gl_WorkGroupID.y;

@@ -150,7 +216,7 @@ void main() {
        return;
    }
 #endif
-#ifdef NEEDS_INIT_IQ_SHMEM
+#if defined(NEEDS_INIT_IQ_SHMEM) || defined(MULMAT_QUANT)
    init_iq_shmem(gl_WorkGroupSize);
 #endif

@@ -200,9 +266,12 @@ void main() {
 #if defined(DATA_A_F32) || defined(DATA_A_F16) || defined(DATA_A_BF16)
    const uint LOAD_VEC_A_EFF = (ALIGNED != 0) ? LOAD_VEC_A : 1;
    const uint LOAD_VEC_BATCH_A = (ALIGNED != 0) ? 1 : 2;
-#else
+#elif !defined(MULMAT_QUANT)
    const uint LOAD_VEC_A_EFF = LOAD_VEC_A;
    const uint LOAD_VEC_BATCH_A = 1;
+#else
+    const uint LOAD_VEC_A_EFF = mm_load_vec_a();
+    const uint LOAD_VEC_BATCH_A = 1;
 #endif
    const uint LOAD_VEC_B_EFF = (ALIGNED != 0) ? LOAD_VEC_B : 1;
    const uint LOAD_VEC_BATCH_B = (ALIGNED != 0) ? 1 : 2;
@@ -21,6 +21,13 @@
 #extension GL_EXT_bfloat16 : enable
 #endif

+#include "ggml_type_ids.glsl"
+
+#ifdef MULMAT_QUANT
+layout (constant_id = 7) const uint MmTypeA = 0;
+layout (constant_id = 8) const uint MmABlockBytes = 2;
+#endif
+
 #include "types.glsl"
 #include "utils.glsl"

@@ -37,6 +44,24 @@ layout (constant_id = 4) const bool enable_smaller_matrices = false;
 const uint BNover2 = enable_smaller_matrices ? (BN / 2) : BN;
 const uint BNover4 = enable_smaller_matrices ? (BN / 4) : BN;
 layout (constant_id = 5) const uint ALIGNED = 0;
+layout (constant_id = 6) const uint subgroup_size = 32;
+
+#ifdef MULMAT_QUANT
+
+uint mm_quant_k() {
+    switch (MmTypeA) {
+    case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: case GGML_TYPE_Q5_1:
+    case GGML_TYPE_Q8_0:
+        return 32u;
+    case GGML_TYPE_Q1_0:
+        return 128u;
+    case GGML_TYPE_Q2_0:
+        return 64u;
+    default:
+        return 256u;
+    }
+}
+#endif

 layout (push_constant) uniform parameter
 {
@@ -72,21 +97,73 @@ layout (push_constant) uniform parameter
 } p;


+#ifndef MULMAT_QUANT
 layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
+#else
+layout (binding = 0) readonly buffer A {uint8_t data_a[];};
+#endif
 layout (binding = 1) readonly buffer B {B_TYPE data_b[];};
 layout (binding = 2) writeonly buffer D {D_TYPE data_d[];};
 #if defined(MUL_MAT_ID) && defined(GGML_VULKAN_COOPMAT2_DECODE_VECTOR)
 layout (binding = 1) readonly buffer B4 {B_TYPEV4 data_b_v4[];};
 #endif

-#if QUANT_K > 1
+#if defined(MULMAT_QUANT) || QUANT_K > 1
 #include "dequant_funcs_cm2.glsl"
+#ifndef MULMAT_QUANT
+// Per-type path: use the alias set by dequant_funcs_cm2.glsl
 #if defined(dequantFuncA_v) && defined(GGML_VULKAN_COOPMAT2_DECODE_VECTOR)
 #define DECODEFUNCA , dequantFuncA, dequantFuncA_v
 #else
 #define DECODEFUNCA , dequantFuncA
 #endif
 #else
+layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufA {
+    uint8_t raw[MmABlockBytes];
+};
+
+float16_t mmDecodeA(const in decodeBufA bl_in, const in uint blockCoords[2], const in uint coordInBlock[2]) {
+    switch (MmTypeA) {
+    case GGML_TYPE_Q1_0:  return dequantFuncQ1_0 (decodeBufQ1_0 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q2_0:  return dequantFuncQ2_0 (decodeBufQ2_0 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q4_0:  return dequantFuncQ4_0 (decodeBufQ4_0 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q4_1:  return dequantFuncQ4_1 (decodeBufQ4_1 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q5_0:  return dequantFuncQ5_0 (decodeBufQ5_0 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q5_1:  return dequantFuncQ5_1 (decodeBufQ5_1 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q8_0:  return dequantFuncQ8_0 (decodeBufQ8_0 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q2_K:  return dequantFuncQ2_K (decodeBufQ2_K (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q3_K:  return dequantFuncQ3_K (decodeBufQ3_K (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q6_K:  return dequantFuncQ6_K (decodeBufQ6_K (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_TQ1_0: return dequantFuncTQ1_0(decodeBufTQ1_0(bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_TQ2_0: return dequantFuncTQ2_0(decodeBufTQ2_0(bl_in), blockCoords, coordInBlock);
+    default:              return float16_t(0);
+    }
+}
+
+#ifdef GGML_VULKAN_COOPMAT2_DECODE_VECTOR
+f16vec4 mmDecodeA_v(const in decodeBufA bl_in, const in uint blockCoords[2], const in uint coordInBlock[2]) {
+    switch (MmTypeA) {
+    case GGML_TYPE_Q1_0:  return dequantFuncQ1_0_v (decodeBufQ1_0 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q2_0:  return dequantFuncQ2_0_v (decodeBufQ2_0 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q4_0:  return dequantFuncQ4_0_v (decodeBufQ4_0 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q4_1:  return dequantFuncQ4_1_v (decodeBufQ4_1 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q5_0:  return dequantFuncQ5_0_v (decodeBufQ5_0 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q5_1:  return dequantFuncQ5_1_v (decodeBufQ5_1 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q8_0:  return dequantFuncQ8_0_v (decodeBufQ8_0 (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q2_K:  return dequantFuncQ2_K_v (decodeBufQ2_K (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q3_K:  return dequantFuncQ3_K_v (decodeBufQ3_K (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_Q6_K:  return dequantFuncQ6_K_v (decodeBufQ6_K (bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_TQ1_0: return dequantFuncTQ1_0_v(decodeBufTQ1_0(bl_in), blockCoords, coordInBlock);
+    case GGML_TYPE_TQ2_0: return dequantFuncTQ2_0_v(decodeBufTQ2_0(bl_in), blockCoords, coordInBlock);
+    default:              return f16vec4(0);
+    }
+}
+#define DECODEFUNCA , mmDecodeA, mmDecodeA_v
+#else
+#define DECODEFUNCA , mmDecodeA
+#endif
+#endif
+#else
 #define DECODEFUNCA
 #endif

@@ -114,7 +191,6 @@ layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufB {
 };

 uint _ne1;
-layout (constant_id = 6) const uint subgroup_size = 32;
 shared uvec4 ballots_sh[BLOCK_SIZE / subgroup_size];

 B_TYPE decodeFuncB(const in decodeBufB bl, const in uint blockCoords[2], const in uint coordInBlock[2])
@@ -246,6 +322,10 @@ void load_row_ids_hoisted(uint expert_idx, uint ic) {
 }
 #endif

+#ifdef MULMAT_QUANT
+#include "iq_shmem_init.glsl"
+#endif
+
 void main() {
    const uint tid = gl_LocalInvocationIndex;
    const uint ic = gl_WorkGroupID.y;
@@ -264,7 +344,7 @@ void main() {
 #endif
 #endif

-#ifdef NEEDS_INIT_IQ_SHMEM
+#if defined(NEEDS_INIT_IQ_SHMEM) || defined(MULMAT_QUANT)
    init_iq_shmem(gl_WorkGroupSize);
 #endif

@@ -305,22 +385,33 @@ void main() {
    const uint end_k = min(p.K, (ik + 1) * p.k_split);
 #endif

+#ifdef MULMAT_QUANT
+    const uint qk = mm_quant_k();
+#else
+    const uint qk = QUANT_K;
+#endif
+
 #ifdef MUL_MAT_ID
-    uint pos_a = expert_idx * (p.batch_stride_a / QUANT_K);
+    uint pos_a = expert_idx * (p.batch_stride_a / qk);
    uint pos_b = 0;
 #else
-    uint pos_a = batch_idx_a * (p.batch_stride_a / QUANT_K);
+    uint pos_a = batch_idx_a * (p.batch_stride_a / qk);
    uint pos_b = batch_idx * p.batch_stride_b;
    uint pos_d = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * p.num_batches;
 #endif

-    uint stride_a = p.stride_a / QUANT_K;
+#ifdef MULMAT_QUANT
+    // pos_a is a byte offset into the raw buffer; strides stay in block units
+    pos_a *= MmABlockBytes;
+#endif
+
+    uint stride_a = p.stride_a / qk;
    uint stride_b = p.stride_b;

    // Hint to the compiler that values are aligned (want 16B alignment).
    // Quants are always block-aligned, no alignment needed.
    if (ALIGNED != 0) {
-#if QUANT_K == 1
+#if !defined(MULMAT_QUANT) && QUANT_K == 1
        stride_a &= ~7;
 #endif
        stride_b &= ~7;
@@ -335,10 +426,8 @@ void main() {
 #endif
    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);

-#if QUANT_K > 1
-    tensorLayoutA = setTensorLayoutBlockSizeNV(tensorLayoutA, 1, QUANT_K);
-    tensorLayoutAClamp = setTensorLayoutBlockSizeNV(tensorLayoutAClamp, 1, QUANT_K);
-#endif
+    tensorLayoutA = setTensorLayoutBlockSizeNV(tensorLayoutA, 1, qk);
+    tensorLayoutAClamp = setTensorLayoutBlockSizeNV(tensorLayoutAClamp, 1, qk);
 #if defined(MUL_MAT_ID) && defined(GGML_VULKAN_COOPMAT2_DECODE_VECTOR)
    tensorLayoutB = setTensorLayoutBlockSizeNV(tensorLayoutB, 1, BK);
 #endif
@@ -368,19 +457,19 @@ void main() {
    const uint START_ALIGN_K = 256;
    // For Qi_K (block size 256), unroll whole 256 element tiles.
    // For legacy quants (block size 32), unroll 8x.
-    const uint UNROLL_K = (QUANT_K == 256) ? 256 : (BK * 8);
+    const uint UNROLL_K = (qk == 256) ? 256 : (BK * 8);
    const uint unroll_count = UNROLL_K / BK;

    // Detect a fast path where all loads are entirely in bounds and no clamping is required
    if ((ir + 1) * BM <= p.M && (ic + 1) * BN <= p.padded_N && (start_k % START_ALIGN_K) == 0 && (end_k % BK) == 0 &&
-#if QUANT_K == 1
+#if !defined(MULMAT_QUANT) && QUANT_K == 1
        (stride_a % 8) == 0 &&
 #endif
        (stride_b % 8) == 0) {
        // Hint to the compiler that values are aligned (want 16B alignment)
        start_k &= ~(START_ALIGN_K-1);
        stride_b &= ~7;
-#if QUANT_K == 1
+#if !defined(MULMAT_QUANT) && QUANT_K == 1
        stride_a &= ~7;
 #endif

@@ -551,10 +640,10 @@ void main() {
            [[dont_unroll]]
            for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) {

-                if ((block_k % QUANT_K) == 0) {
+                if ((block_k % qk) == 0) {
                    store_scales(tid);
                }
-                if (block_k + BK < end_k && ((block_k + BK) % QUANT_K) == 0) {
+                if (block_k + BK < end_k && ((block_k + BK) % qk) == 0) {
                    fetch_scales(ir * BM, pos_a, stride_a, block_k + BK, tid, false);
                }

@@ -595,10 +684,10 @@ void main() {
            [[dont_unroll]]
            for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) {

-                if ((block_k % QUANT_K) == 0) {
+                if ((block_k % qk) == 0) {
                    store_scales(tid);
                }
-                if (block_k + BK < end_k && ((block_k + BK) % QUANT_K) == 0) {
+                if (block_k + BK < end_k && ((block_k + BK) % qk) == 0) {
                    fetch_scales(ir * BM, pos_a, stride_a, block_k + BK, tid, false);
                }

@@ -639,10 +728,10 @@ void main() {
        [[dont_unroll]]
        for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) {

-            if ((block_k % QUANT_K) == 0) {
+            if ((block_k % qk) == 0) {
                store_scales(tid);
            }
-            if (block_k + BK < end_k && ((block_k + BK) % QUANT_K) == 0) {
+            if (block_k + BK < end_k && ((block_k + BK) % qk) == 0) {
                fetch_scales(ir * BM, pos_a, stride_a, block_k + BK, tid, false);
            }

@@ -955,6 +955,7 @@ shared uint16_t iq1s_grid[2048];
 shared uint32_t iq1s_grid_gpu[2048];
 #endif

+#if defined(DATA_A_IQ1_S) || defined(DATA_A_IQ1_M)
 #define NEEDS_INIT_IQ_SHMEM
 void init_iq_shmem(uvec3 wgsize)
 {
@@ -978,6 +979,17 @@ void init_iq_shmem(uvec3 wgsize)
    barrier();
 }
 #endif
+#endif
+
+#if defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S)
+#if defined(DATA_A_IQ2_S)
+shared uvec2 iq2s_grid[1024];
+#elif defined(DATA_A_IQ2_XS)
+shared uvec2 iq2xs_grid[512];
+#else
+shared uvec2 iq2xxs_grid[256];
+#endif
+#endif

 #define QUANT_K_IQ2_XXS 256
 #define QUANT_R_IQ2_XXS 1
@@ -1063,8 +1075,7 @@ const uvec2[256] iq2xxs_grid_const = {
    uvec2(0x08080808, 0x2b2b082b), uvec2(0x08192b08, 0x2b2b1908), uvec2(0x19190808, 0x2b2b2b08), uvec2(0x08081908, 0x2b2b2b19)
 };

-shared uvec2 iq2xxs_grid[256];
-
+#if defined(DATA_A_IQ2_XXS)
 #define NEEDS_INIT_IQ_SHMEM
 void init_iq_shmem(uvec3 wgsize)
 {
@@ -1076,12 +1087,15 @@ void init_iq_shmem(uvec3 wgsize)
    }
    barrier();
 }
+#endif

+#if defined(DATA_A_IQ2_XXS)
 #define QUANT_K QUANT_K_IQ2_XXS
 #define QUANT_R QUANT_R_IQ2_XXS
 #define A_TYPE block_iq2_xxs
 #define A_TYPE_PACKED16 block_iq2_xxs_packed16
 #endif
+#endif

 #define QUANT_K_IQ2_XS 256
 #define QUANT_R_IQ2_XS 1
@@ -1233,8 +1247,7 @@ const uvec2 iq2xs_grid_const[512] = {
    uvec2(0x082b2b08, 0x2b2b2b2b), uvec2(0x082b2b2b, 0x2b2b2b2b), uvec2(0x2b190819, 0x2b2b2b2b), uvec2(0x2b2b2b2b, 0x2b2b2b2b),
 };

-shared uvec2 iq2xs_grid[512];
-
+#if defined(DATA_A_IQ2_XS)
 #define NEEDS_INIT_IQ_SHMEM
 void init_iq_shmem(uvec3 wgsize)
 {
@@ -1246,12 +1259,15 @@ void init_iq_shmem(uvec3 wgsize)
    }
    barrier();
 }
+#endif

+#if defined(DATA_A_IQ2_XS)
 #define QUANT_K QUANT_K_IQ2_XS
 #define QUANT_R QUANT_R_IQ2_XS
 #define A_TYPE block_iq2_xs
 #define A_TYPE_PACKED16 block_iq2_xs_packed16
 #endif
+#endif

 #define QUANT_K_IQ2_S 256
 #define QUANT_R_IQ2_S 1
@@ -1533,8 +1549,7 @@ const uvec2 iq2s_grid_const[1024] = {
    uvec2(0x082b082b, 0x2b2b2b2b), uvec2(0x082b2b08, 0x2b2b2b2b), uvec2(0x2b082b08, 0x2b2b2b2b), uvec2(0x2b2b2b2b, 0x2b2b2b2b)
 };

-shared uvec2 iq2s_grid[1024];
-
+#if defined(DATA_A_IQ2_S)
 #define NEEDS_INIT_IQ_SHMEM
 void init_iq_shmem(uvec3 wgsize)
 {
@@ -1546,12 +1561,23 @@ void init_iq_shmem(uvec3 wgsize)
    }
    barrier();
 }
+#endif

+#if defined(DATA_A_IQ2_S)
 #define QUANT_K QUANT_K_IQ2_S
 #define QUANT_R QUANT_R_IQ2_S
 #define A_TYPE block_iq2_s
 #define A_TYPE_PACKED16 block_iq2_s_packed16
 #endif
+#endif
+
+#if defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S)
+#if defined(DATA_A_IQ3_S)
+shared uint32_t iq3s_grid[512];
+#else
+shared uint32_t iq3xxs_grid[256];
+#endif
+#endif

 #define QUANT_K_IQ3_XXS 256
 #define QUANT_R_IQ3_XXS 1
@@ -1605,8 +1631,7 @@ const uint32_t iq3xxs_grid_const[256] = {
    0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c, 0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04,
 };

-shared uint32_t iq3xxs_grid[256];
-
+#if defined(DATA_A_IQ3_XXS)
 #define NEEDS_INIT_IQ_SHMEM
 void init_iq_shmem(uvec3 wgsize)
 {
@@ -1618,12 +1643,15 @@ void init_iq_shmem(uvec3 wgsize)
    }
    barrier();
 }
+#endif

+#if defined(DATA_A_IQ3_XXS)
 #define QUANT_K QUANT_K_IQ3_XXS
 #define QUANT_R QUANT_R_IQ3_XXS
 #define A_TYPE block_iq3_xxs
 #define A_TYPE_PACKED16 block_iq3_xxs_packed16
 #endif
+#endif

 #define QUANT_K_IQ3_S 256
 #define QUANT_R_IQ3_S 1
@@ -1715,8 +1743,7 @@ const uint32_t iq3s_grid_const[512] = {
    0x0f090307, 0x0f090501, 0x0f090b01, 0x0f0b0505, 0x0f0b0905, 0x0f0d0105, 0x0f0d0703, 0x0f0f0101,
 };

-shared uint32_t iq3s_grid[512];
-
+#if defined(DATA_A_IQ3_S)
 #define NEEDS_INIT_IQ_SHMEM
 void init_iq_shmem(uvec3 wgsize)
 {
@@ -1728,12 +1755,15 @@ void init_iq_shmem(uvec3 wgsize)
    }
    barrier();
 }
+#endif

+#if defined(DATA_A_IQ3_S)
 #define QUANT_K QUANT_K_IQ3_S
 #define QUANT_R QUANT_R_IQ3_S
 #define A_TYPE block_iq3_s
 #define A_TYPE_PACKED16 block_iq3_s_packed16
 #endif
+#endif

 #define QUANT_K_IQ4_XS 256
 #define QUANT_R_IQ4_XS 1
@@ -1847,6 +1877,7 @@ const int8_t kvalues_iq4nl_const[16] = {

 shared FLOAT_TYPE kvalues_iq4nl[16];

+#if defined(DATA_A_IQ4_NL) || defined(DATA_A_IQ4_XS)
 #define NEEDS_INIT_IQ_SHMEM
 void init_iq_shmem(uvec3 wgsize)
 {
@@ -1857,6 +1888,7 @@ void init_iq_shmem(uvec3 wgsize)
    barrier();
 }
 #endif
+#endif

 #if defined(DATA_A_MXFP4) || defined(DATA_A_NVFP4)
 #if !defined(USE_OCP_FP4)
@@ -1886,7 +1918,7 @@ float ue4m3_to_fp32_build(uint u) {
 }
 #endif

-#if !defined(USE_OCP_FP4)
+#if (defined(DATA_A_MXFP4) || defined(DATA_A_NVFP4)) && !defined(USE_OCP_FP4)
 #define NEEDS_INIT_IQ_SHMEM
 void init_iq_shmem(uvec3 wgsize)
 {
@@ -246,6 +246,17 @@ bool is_iq_quant(const std::string& type_name) {
    return string_starts_with(type_name, "iq");
 }

+bool is_lut_quant(const std::string& type_name) {
+    return is_iq_quant(type_name) || type_name == "mxfp4" || type_name == "nvfp4";
+}
+
+std::string lut_load_vec_a(const std::string& type_name) {
+    if (type_name == "iq1_s" || type_name == "iq1_m" || type_name == "iq2_xxs" || type_name == "iq2_xs" || type_name == "iq2_s") {
+        return "8";
+    }
+    return "4";
+}
+
 static const char path_separator = '/';

 std::string join_paths(const std::string& path1, const std::string& path2) {
@@ -583,20 +594,28 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c
    }

    for (const auto& tname : type_names) {
-        std::string load_vec_quant = "2";
-        if ((tname == "q1_0") || (tname == "q4_0") || (tname == "q4_1") || (tname == "q5_1") || (tname == "iq1_s") || (tname == "iq1_m") || (tname == "iq2_xxs") || (tname == "iq2_xs") || (tname == "iq2_s"))
-            load_vec_quant = "8";
-        else if ((tname == "q2_0") || (tname == "q5_0") || (tname == "q8_0") || (tname == "q2_k") || (tname == "q4_k") || (tname == "q5_k") || (tname == "iq3_xxs") || (tname == "iq3_s") || (tname == "iq4_xs") || (tname == "iq4_nl") || (tname == "mxfp4") || (tname == "nvfp4"))
-            load_vec_quant = "4";
-
        if (tname == "bf16") {
            continue;
        }

-        std::string data_a_key = "DATA_A_" + to_uppercase(tname);
-        // For aligned matmul loads
-        std::string load_vec_a = (coopmat2 || tname == "f32" || tname == "f16" || tname == "bf16") ? load_vec : load_vec_quant;
+        // Float types keep per-type compilation (different accumulation loop structure)
+        if (tname == "f32" || tname == "f16") {
+            std::string data_a_key = "DATA_A_" + to_uppercase(tname);

+            const std::map<std::string, std::string> float_type_dict = {
+                {"FLOAT_TYPE",   FLOAT_TYPE(1, tname)},
+                {"FLOAT_TYPEV2", FLOAT_TYPE(2, tname)},
+                {"FLOAT_TYPEV4", FLOAT_TYPE(4, tname)},
+                {"FLOAT_TYPEV8", FLOAT_TYPE(8, tname)},
+            };
+
+            if (!coopmat2) {
+                string_to_spv(shader_name + "_" + tname + "_f32" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
+            }
+            continue;
+        }
+
+        std::string data_a_key = "DATA_A_" + to_uppercase(tname);
        const std::map<std::string, std::string> float_type_dict = {
            {"FLOAT_TYPE",   FLOAT_TYPE(1, tname)},
            {"FLOAT_TYPEV2", FLOAT_TYPE(2, tname)},
@@ -604,30 +623,52 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c
            {"FLOAT_TYPEV8", FLOAT_TYPE(8, tname)},
        };

-        // don't generate f32 variants for coopmat2
-        if (!coopmat2) {
-            string_to_spv(shader_name + "_" + tname + "_f32" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
-        }
-
-        if (tname != "f16" && tname != "f32") {
-            string_to_spv(shader_name + "_" + tname + "_f16" + dot2_sfx, source_name,  merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
-        }
-
-#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
-        if ((coopmat || coopmat2) && (tname == "mxfp4" || tname == "nvfp4")) {
-            if (!coopmat2) {
-                string_to_spv(shader_name + "_" + tname + "_f32_ocp" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
-            }
-            string_to_spv(shader_name + "_" + tname + "_f16_ocp" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
-        }
-#endif
-
 #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
-        // Integer dot mmq performs better with f32 accumulators (different shader, skip for dot2)
        if (!f16acc && !coopmat && !coopmat2 && !dot2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4")) {
            string_to_spv(shader_name + "_" + tname + "_q8_1", "mul_mmq.comp", merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"D_TYPE", "float"},}), fp16, coopmat, coopmat2, f16acc);
        }
 #endif
+
+        if (is_lut_quant(tname)) {
+            std::string lva = lut_load_vec_a(tname);
+
+            string_to_spv(shader_name + "_" + tname + "_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", lva}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
+
+            if (!coopmat2) {
+                string_to_spv(shader_name + "_" + tname + "_f32" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", lva}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
+            }
+
+#if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT)
+            if ((tname == "mxfp4" || tname == "nvfp4") && (coopmat || coopmat2)) {
+                string_to_spv(shader_name + "_" + tname + "_f16_ocp" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"LOAD_VEC_A", lva}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
+                if (!coopmat2) {
+                    string_to_spv(shader_name + "_" + tname + "_f32_ocp" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"LOAD_VEC_A", lva}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
+                }
+            }
+#endif
+            continue;
+        }
+
+        // dedicated shader needed due to regression on Ampere
+        if (coopmat2 && (tname == "q4_k" || tname == "q5_k")) {
+            string_to_spv(shader_name + "_" + tname + "_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
+        }
+    }
+
+    // Quant shader: one SPIR-V for all quant types, selected via MmTypeA spec constant
+    {
+        const std::map<std::string, std::string> quant_float_type_dict = {
+            {"FLOAT_TYPE",   FLOAT_TYPE(1, "q4_0")},
+            {"FLOAT_TYPEV2", FLOAT_TYPE(2, "q4_0")},
+            {"FLOAT_TYPEV4", FLOAT_TYPE(4, "q4_0")},
+            {"FLOAT_TYPEV8", FLOAT_TYPE(8, "q4_0")},
+        };
+
+        string_to_spv(shader_name + "_quant_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, quant_float_type_dict), {{"MULMAT_QUANT", "1"}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
+
+        if (!coopmat2) {
+            string_to_spv(shader_name + "_quant_f32" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, quant_float_type_dict), {{"MULMAT_QUANT", "1"}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
+        }
    }
 }

@@ -7,7 +7,7 @@ dynamic = ["classifiers"]
 readme = "README.md"
 authors = [{name = "GGML", email = "ggml@ggml.ai"}]
 requires-python = '>=3.10'
-dependencies = ['numpy (>=1.17)', 'tqdm (>=4.27)', 'pyyaml (>=5.1)', 'requests (>=2.25)']
+dependencies = ['numpy (>=2.2.6)', 'tqdm (>=4.27)', 'pyyaml (>=5.1)', 'requests (>=2.25)']
 classifiers = [
    "Programming Language :: Python :: 3",
    "License :: OSI Approved :: MIT License",
@@ -1357,7 +1357,7 @@ extern "C" {
    LLAMA_API struct llama_sampler * llama_sampler_chain_get(      struct llama_sampler * chain, int32_t i);

    // the total number of samplers in the chain
-    LLAMA_API int                    llama_sampler_chain_n  (const struct llama_sampler * chain);
+    LLAMA_API int32_t                llama_sampler_chain_n  (const struct llama_sampler * chain);

    // after removing a sampler, the chain will no longer own it, and it will not be freed when the chain is freed
    LLAMA_API struct llama_sampler * llama_sampler_chain_remove(   struct llama_sampler * chain, int32_t i);
@@ -8,7 +8,7 @@ readme = "README.md"
 authors = [{name = "GGML", email = "ggml@ggml.ai"}]
 requires-python = '>=3.10,<3.15'
 dependencies = [
-    'numpy (>=1.26.4,<3.0.0)',
+    'numpy (>=2.2.6,<3.0.0)',
    'sentencepiece (>=0.1.98,<0.3.0)',
    'transformers (==4.57.6)',
    'protobuf (>=4.21.0,<5.0.0)',
@@ -1,4 +1,4 @@
-numpy~=1.26.4
+numpy~=2.2.6
 sentencepiece>=0.1.98,<0.3.0

 transformers==4.57.6
@@ -1,3 +1,3 @@
-numpy~=1.26.4
+numpy~=2.2.6
 PySide6~=6.9.0
 gguf>=0.17.0
@@ -1,5 +1,5 @@
 datasets~=4.8.0
 matplotlib~=3.10.0
-numpy~=1.26.4
+numpy~=2.2.6
 requests~=2.32.3
 tqdm~=4.67.1
@@ -1,7 +1,7 @@
 aiohttp~=3.9.3
 pytest~=8.3.3
 matplotlib~=3.10.0
-numpy~=1.26.4
+numpy~=2.2.6
 openai~=2.14.0
 pandas~=2.2.3
 prometheus-client~=0.20.0
@@ -0,0 +1,136 @@
+# Compile-time profiling using clang -ftime-trace + ClangBuildAnalyzer.
+#
+# Usage:
+#   .\scripts\build-profile.ps1 [-Full] [-Jobs N]
+#
+# -Full : include Server, Tools, and Tests (default: minimal build)
+# -Jobs :  number of parallel jobs (default: all cores)
+#
+# Requires ClangBuildAnalyzer:
+#   https://github.com/aras-p/ClangBuildAnalyzer
+
+param(
+    [switch]$Full,
+    [int]$Jobs = [Environment]::ProcessorCount
+)
+
+$ErrorActionPreference = "Stop"
+
+$ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path
+$RootDir = Split-Path -Parent $ScriptDir
+
+if ($Full) {
+    $BuildDir = Join-Path $RootDir "build-profile-full"
+    $Report = Join-Path $BuildDir "profile-report-full.txt"
+} else {
+    $BuildDir = Join-Path $RootDir "build-profile-baseline"
+    $Report = Join-Path $BuildDir "profile-report.txt"
+}
+
+$OutputBin = Join-Path $BuildDir "clang_analysis.bin"
+
+if (-not (Get-Command clang++ -ErrorAction SilentlyContinue)) {
+    Write-Error "clang++ not found"
+    exit 1
+}
+
+if (-not (Get-Command ninja -ErrorAction SilentlyContinue)) {
+    Write-Error "ninja not found (required so cmake does not fall back to the Visual Studio/MSVC generator)"
+    exit 1
+}
+
+if (-not (Get-Command ClangBuildAnalyzer -ErrorAction SilentlyContinue)) {
+    Write-Error "ClangBuildAnalyzer not found`n  https://github.com/aras-p/ClangBuildAnalyzer/releases"
+    exit 1
+}
+
+$ClangVer = (clang++ --version | Select-Object -First 1)
+Write-Host "compiler : $ClangVer"
+Write-Host "build dir: $BuildDir"
+Write-Host "output   : $OutputBin"
+Write-Host "jobs     : $Jobs"
+Write-Host ""
+
+if (Get-Command ccache -ErrorAction SilentlyContinue) {
+    Write-Host "clearing ccache..."
+    ccache -C -z
+}
+
+$env:CCACHE_DISABLE = "1"
+
+$TestsFlag = if ($Full) { "ON" } else { "OFF" }
+$ToolsFlag = if ($Full) { "ON" } else { "OFF" }
+$ServerFlag = if ($Full) { "ON" } else { "OFF" }
+
+cmake --fresh `
+    -S $RootDir `
+    -B $BuildDir `
+    -G "Ninja" `
+    -DCMAKE_BUILD_TYPE=Release `
+    -DCMAKE_C_COMPILER=clang `
+    -DCMAKE_CXX_COMPILER=clang++ `
+    -DCMAKE_C_FLAGS="-ftime-trace" `
+    -DCMAKE_CXX_FLAGS="-ftime-trace" `
+    -DGGML_CCACHE=OFF `
+    -DGGML_OPENMP=ON `
+    -DGGML_NATIVE=OFF `
+    "-DLLAMA_BUILD_TESTS=$TestsFlag" `
+    -DLLAMA_BUILD_EXAMPLES=OFF `
+    "-DLLAMA_BUILD_TOOLS=$ToolsFlag" `
+    "-DLLAMA_BUILD_SERVER=$ServerFlag" `
+    -DLLAMA_BUILD_APP=OFF
+
+if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE }
+
+$StrayTrace = Join-Path $RootDir "-.json"
+if (Test-Path $StrayTrace) {
+    Remove-Item $StrayTrace -Force
+}
+
+Write-Host ""
+Write-Host "Initializing ClangBuildAnalyzer..."
+ClangBuildAnalyzer --start $BuildDir
+Write-Host ""
+
+Write-Host "building..."
+Write-Host ""
+
+$StartTime = Get-Date
+
+cmake --build $BuildDir --clean-first -j $Jobs
+
+if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE }
+
+$Elapsed = (Get-Date) - $StartTime
+
+Write-Host ""
+Write-Host ("build time: {0}s ({1}m {2}s)" -f [int]$Elapsed.TotalSeconds, [int]$Elapsed.TotalMinutes, $Elapsed.Seconds)
+Write-Host ""
+
+Write-Host "Aggregating profile metrics..."
+ClangBuildAnalyzer --stop $BuildDir $OutputBin | Out-Null
+
+Write-Host ""
+Write-Host ("=" * 80)
+
+$TUs = "?"
+if (Test-Path $Report) {
+    $Match = Select-String -Path $Report -Pattern "Compilation \((\d+)" | Select-Object -First 1
+    if ($Match) { $TUs = $Match.Matches[0].Groups[1].Value }
+}
+
+ClangBuildAnalyzer --analyze $OutputBin | Tee-Object -FilePath $Report
+
+Write-Host ""
+Write-Host "translation units: $TUs"
+Write-Host ""
+Write-Host "largest trace files (top 20 by size):"
+
+Get-ChildItem -Path $BuildDir -Recurse -Filter "*.json" |
+    Where-Object { $_.Name -ne "compile_commands.json" } |
+    Sort-Object Length -Descending |
+    Select-Object -First 20 |
+    ForEach-Object { "{0,8:F1} KB  {1}" -f ($_.Length / 1024), $_.FullName }
+
+Write-Host ""
+Write-Host "ClangBuildAnalyzer report was generated: $Report"
@@ -0,0 +1,122 @@
+#!/usr/bin/env bash
+# Compile-time profiling using clang -ftime-trace + ClangBuildAnalyzer.
+#
+# Usage:
+#   ./scripts/build-profile.sh [--full] [-jN]
+#
+# --full: include Server, Tools, and Tests (default: minimal build)
+# -jN   :  number of parallel jobs (default: all cores)
+#
+# Requires ClangBuildAnalyzer:
+#   macOS: brew install clang-build-analyzer
+#   Linux: https://github.com/aras-p/ClangBuildAnalyzer.git
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+ROOT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
+
+FULL=0
+JOBS="-j$(nproc 2>/dev/null || sysctl -n hw.ncpu)"
+
+for arg in "$@"; do
+    case "${arg}" in
+        --full) FULL=1 ;;
+        -j*)    JOBS="${arg}" ;;
+        *) echo "error: unknown argument: ${arg}" >&2; exit 1 ;;
+    esac
+done
+
+if [ "${FULL}" -eq 1 ]; then
+    BUILD_DIR="${ROOT_DIR}/build-profile-full"
+    REPORT="${BUILD_DIR}/profile-report-full.txt"
+else
+    BUILD_DIR="${ROOT_DIR}/build-profile-baseline"
+    REPORT="${BUILD_DIR}/profile-report.txt"
+fi
+
+OUTPUT_BIN="${BUILD_DIR}/clang_analysis.bin"
+
+if ! command -v clang++ &>/dev/null; then
+    echo "error: clang++ not found" >&2
+    exit 1
+fi
+
+if ! command -v ClangBuildAnalyzer &>/dev/null; then
+    echo "error: ClangBuildAnalyzer not found" >&2
+    echo "  brew install clangbuildanalyzer  (macOS)" >&2
+    echo "  or: https://github.com/aras-p/ClangBuildAnalyzer/releases" >&2
+    exit 1
+fi
+
+CLANG_VER=$(clang++ --version | head -1)
+echo "compiler : ${CLANG_VER}"
+echo "build dir: ${BUILD_DIR}"
+echo "output   : ${OUTPUT_BIN}"
+echo "jobs     : ${JOBS}"
+echo
+
+if command -v ccache &>/dev/null; then
+    echo "clearing ccache..."
+    ccache -C -z
+fi
+
+export CCACHE_DISABLE=1
+
+cmake --fresh \
+    -S "${ROOT_DIR}" \
+    -B "${BUILD_DIR}" \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCMAKE_C_COMPILER=clang \
+    -DCMAKE_CXX_COMPILER=clang++ \
+    -DCMAKE_C_FLAGS="-ftime-trace" \
+    -DCMAKE_CXX_FLAGS="-ftime-trace" \
+    -DGGML_CCACHE=OFF \
+    -DGGML_OPENMP=ON \
+    -DGGML_NATIVE=OFF \
+    -DLLAMA_BUILD_TESTS=$([ "${FULL}" -eq 1 ] && echo ON || echo OFF) \
+    -DLLAMA_BUILD_EXAMPLES=OFF \
+    -DLLAMA_BUILD_TOOLS=$([ "${FULL}" -eq 1 ] && echo ON || echo OFF) \
+    -DLLAMA_BUILD_SERVER=$([ "${FULL}" -eq 1 ] && echo ON || echo OFF) \
+    -DLLAMA_BUILD_APP=OFF
+
+echo
+
+echo "Initializing ClangBuildAnalyzer..."
+ClangBuildAnalyzer --start "${BUILD_DIR}"
+echo
+
+echo "building..."
+echo
+
+START=$(date +%s)
+
+cmake --build "${BUILD_DIR}" --clean-first "${JOBS}"
+
+END=$(date +%s)
+ELAPSED=$((END - START))
+
+echo
+printf "build time: %ds (%dm %ds)\n" "${ELAPSED}" "$((ELAPSED / 60))" "$((ELAPSED % 60))"
+echo
+
+echo "Aggregating profile metrics..."
+ClangBuildAnalyzer --stop "${BUILD_DIR}" "${OUTPUT_BIN}" > /dev/null
+
+echo
+echo "================================================================================"
+TUS=$(grep -oP "Compilation \(\K[0-9]+" "${REPORT}" 2>/dev/null || echo "?")
+ClangBuildAnalyzer --analyze "${OUTPUT_BIN}" | tee "${REPORT}"
+
+echo
+echo "translation units: ${TUS}"
+echo
+echo "largest trace files (top 20 by size):"
+find "${BUILD_DIR}" -name "*.json" ! -name "compile_commands.json" \
+    | xargs ls -l 2>/dev/null \
+    | awk 'NF>5 {print $5, $NF}' \
+    | sort -rn \
+    | awk 'NR<=20 {printf "%8.1f KB  %s\n", $1/1024, $2}'
+
+echo
+echo "ClangBuildAnalyzer report was generated: ${REPORT}"
@@ -22,9 +22,9 @@ for arg in "$@"; do
    esac
 done

-MAJOR=$(grep "set(LLAMA_VERSION_MAJOR" "$REPO_ROOT/CMakeLists.txt" | grep -oP '\d+')
-MINOR=$(grep "set(LLAMA_VERSION_MINOR" "$REPO_ROOT/CMakeLists.txt" | grep -oP '\d+')
-PATCH=$(grep "set(LLAMA_VERSION_PATCH" "$REPO_ROOT/CMakeLists.txt" | grep -oP '\d+')
+MAJOR=$(grep "set(LLAMA_VERSION_MAJOR" "$REPO_ROOT/CMakeLists.txt" | sed 's/.*MAJOR \([0-9]*\).*/\1/')
+MINOR=$(grep "set(LLAMA_VERSION_MINOR" "$REPO_ROOT/CMakeLists.txt" | sed 's/.*MINOR \([0-9]*\).*/\1/')
+PATCH=$(grep "set(LLAMA_VERSION_PATCH" "$REPO_ROOT/CMakeLists.txt" | sed 's/.*PATCH \([0-9]*\).*/\1/')
 VERSION="v${MAJOR}.${MINOR}.${PATCH}"
 echo "Determined version: ${VERSION}"
 if [[ -n "${GITHUB_OUTPUT:-}" ]]; then
@@ -91,9 +91,9 @@ else
    fi
 fi

-MAJOR=$(grep "set(GGML_VERSION_MAJOR" "$REPO_ROOT/ggml/CMakeLists.txt" | grep -oP '\d+')
-MINOR=$(grep "set(GGML_VERSION_MINOR" "$REPO_ROOT/ggml/CMakeLists.txt" | grep -oP '\d+')
-PATCH=$(grep "set(GGML_VERSION_PATCH" "$REPO_ROOT/ggml/CMakeLists.txt" | grep -oP '\d+')
+MAJOR=$(grep "set(GGML_VERSION_MAJOR" "$REPO_ROOT/ggml/CMakeLists.txt" | sed 's/.*MAJOR \([0-9]*\).*/\1/')
+MINOR=$(grep "set(GGML_VERSION_MINOR" "$REPO_ROOT/ggml/CMakeLists.txt" | sed 's/.*MINOR \([0-9]*\).*/\1/')
+PATCH=$(grep "set(GGML_VERSION_PATCH" "$REPO_ROOT/ggml/CMakeLists.txt" | sed 's/.*PATCH \([0-9]*\).*/\1/')
 GGML_VERSION="v${MAJOR}.${MINOR}.${PATCH}"
 echo "Local ggml version: ${GGML_VERSION}"

@@ -8,40 +8,44 @@ llama_add_compile_flags()

 file(GLOB LLAMA_MODELS_SOURCES "models/*.cpp")

+set(LLAMA_CORE_SOURCES
+    llama.cpp
+    llama-adapter.cpp
+    llama-arch.cpp
+    llama-batch.cpp
+    llama-chat.cpp
+    llama-context.cpp
+    llama-cparams.cpp
+    llama-grammar.cpp
+    llama-graph.cpp
+    llama-hparams.cpp
+    llama-impl.cpp
+    llama-io.cpp
+    llama-kv-cache.cpp
+    llama-kv-cache-iswa.cpp
+    llama-kv-cache-dsa.cpp
+    llama-kv-cache-dsa-iswa.cpp
+    llama-kv-cache-msa.cpp
+    llama-kv-cache-dsv4.cpp
+    llama-memory.cpp
+    llama-memory-hybrid.cpp
+    llama-memory-hybrid-iswa.cpp
+    llama-memory-hybrid-idx.cpp
+    llama-memory-recurrent.cpp
+    llama-mmap.cpp
+    llama-model-loader.cpp
+    llama-model-saver.cpp
+    llama-model.cpp
+    llama-quant.cpp
+    llama-sampler.cpp
+    llama-vocab.cpp
+    unicode-data.cpp
+    unicode.cpp
+)
+
 add_library(llama
            ../include/llama.h
-            llama.cpp
-            llama-adapter.cpp
-            llama-arch.cpp
-            llama-batch.cpp
-            llama-chat.cpp
-            llama-context.cpp
-            llama-cparams.cpp
-            llama-grammar.cpp
-            llama-graph.cpp
-            llama-hparams.cpp
-            llama-impl.cpp
-            llama-io.cpp
-            llama-kv-cache.cpp
-            llama-kv-cache-iswa.cpp
-            llama-kv-cache-dsa.cpp
-            llama-kv-cache-dsa-iswa.cpp
-            llama-kv-cache-msa.cpp
-            llama-kv-cache-dsv4.cpp
-            llama-memory.cpp
-            llama-memory-hybrid.cpp
-            llama-memory-hybrid-iswa.cpp
-            llama-memory-hybrid-idx.cpp
-            llama-memory-recurrent.cpp
-            llama-mmap.cpp
-            llama-model-loader.cpp
-            llama-model-saver.cpp
-            llama-model.cpp
-            llama-quant.cpp
-            llama-sampler.cpp
-            llama-vocab.cpp
-            unicode-data.cpp
-            unicode.cpp
+            ${LLAMA_CORE_SOURCES}
            unicode.h
            ${LLAMA_MODELS_SOURCES}
            )
@@ -50,13 +54,20 @@ set_target_properties(llama PROPERTIES
    VERSION ${LLAMA_VERSION_BASE}
    SOVERSION ${LLAMA_VERSION_MAJOR}
    MACHO_CURRENT_VERSION 0 # keep macOS linker from seeing oversized version number
+    UNITY_BUILD ON
+    UNITY_BUILD_BATCH_SIZE 16
 )

+# exclude non-model sources from unity build
+set_source_files_properties(${LLAMA_CORE_SOURCES} ../include/llama.h unicode.h
+    PROPERTIES SKIP_UNITY_BUILD_INCLUSION ON)
+
 configure_file(llama-version.h.in ${CMAKE_CURRENT_BINARY_DIR}/llama-version.h @ONLY)

 target_include_directories(llama PRIVATE . ${CMAKE_CURRENT_BINARY_DIR})
 target_include_directories(llama PUBLIC ../include)
 target_compile_features   (llama PRIVATE cxx_std_17) # don't bump
+target_precompile_headers (llama PRIVATE models/models.h)

 target_link_libraries(llama PUBLIC ggml)

@@ -666,7 +666,9 @@ void llama_context::sched_reserve() {
        //       need to implement a more robust mechanism that tries a few different inputs and analyzes the results
        ggml_cgraph * gf = nullptr;
        switch (model.arch) {
+            case LLM_ARCH_KIMI_LINEAR:
            case LLM_ARCH_MINIMAX_01:
+                // [TAG_RESERVE_DIAG_DECAY]
                // the `inp_diag_decay` tensor size scales with `n_seq_tokens^2` which
                // makes `n_seqs == 1` use more memory for the compute graph compared to `n_seqs > 1`
                gf = graph_reserve(n_tokens, 1,      n_outputs_pp, mctx.get(), model.hparams.no_alloc);
@@ -55,6 +55,10 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx(
        // K-shift must not rotate them while the stream copies in the same update still apply
        hparams_idx.rope_type = LLAMA_ROPE_TYPE_NONE;

+        // fool llama_kv_cache into thinking this is a MLA cache, so it won't cache V tensors
+        hparams_idx.n_embd_head_k_mla_impl = model.hparams.indexer_head_size;
+        hparams_idx.n_embd_head_v_mla_impl = model.hparams.indexer_head_size;
+
        LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);

        return new llama_kv_cache(
@@ -936,6 +936,7 @@ const char * llm_type_name(llm_type type) {
        case LLM_TYPE_17B_128E:      return "17Bx128E (Maverick)";
        case LLM_TYPE_A13B:          return "A13B";
        case LLM_TYPE_1B_A400M:      return "1B.A400M";
+        case LLM_TYPE_3B_A800M:      return "3B.A800M";
        case LLM_TYPE_7B_A1B:        return "7B.A1B";
        case LLM_TYPE_8B_A1B:        return "8B.A1B";
        case LLM_TYPE_7_9B_A1_3B:    return "7.9B.A1.3B";
@@ -946,6 +947,7 @@ const char * llm_type_name(llm_type type) {
        case LLM_TYPE_26B_A4B:       return "26B.A4B";
        case LLM_TYPE_30B_A3B:       return "30B.A3B";
        case LLM_TYPE_31B_A3_5B:     return "31B.A3.5B";
+        case LLM_TYPE_32B_A9B:       return "32B.A9B";
        case LLM_TYPE_35B_A3B:       return "35B.A3B";
        case LLM_TYPE_48B_A3B:       return "48B.A3B";
        case LLM_TYPE_75B_A9B:       return "75B.A9B";
@@ -2642,9 +2644,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                        filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
                    }

-                    if ((arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_HY_V3 || arch == LLM_ARCH_GLM_DSA ||
-                            arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_DEEPSEEK32) &&
-                            hparams.n_layer_nextn > 0) {
+                    // don't filter when n_layer_nextn is repurposed for a router layer the trunk attends
+                    // or when a model is entirely n_layer_nextn layers and has no trunk
+                    if (hparams.n_layer_nextn > 0 && hparams.n_layer() > 0 && hparams.router_layer < 0) {
                        if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
                            filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
                        } else {
@@ -117,6 +117,7 @@ enum llm_type {
    LLM_TYPE_17B_128E, // llama4 Maverick
    LLM_TYPE_A13B,
    LLM_TYPE_1B_A400M, // Granite3 MoE
+    LLM_TYPE_3B_A800M, // Granite3 MoE
    LLM_TYPE_7B_A1B,
    LLM_TYPE_8B_A1B, // lfm2moe
    LLM_TYPE_7_9B_A1_3B, // Ling-3.0-tiny
@@ -127,6 +128,7 @@ enum llm_type {
    LLM_TYPE_26B_A4B, // Gemma4
    LLM_TYPE_30B_A3B,
    LLM_TYPE_31B_A3_5B,
+    LLM_TYPE_32B_A9B, // Granite4 Hybrid
    LLM_TYPE_35B_A3B, // Qwen3.5
    LLM_TYPE_48B_A3B, // Kimi Linear
    LLM_TYPE_75B_A9B, // Nemotron 3 Puzzle
@@ -1006,7 +1006,7 @@ struct llama_sampler * llama_sampler_chain_remove(struct llama_sampler * chain,
    return result;
 }

-int llama_sampler_chain_n(const struct llama_sampler * chain) {
+int32_t llama_sampler_chain_n(const struct llama_sampler * chain) {
    const auto * p = (const llama_sampler_chain *) chain->ctx;

    return p->samplers.size();
@@ -29,15 +29,13 @@ void llama_model_bert::load_arch_tensors(llama_model_loader &) {
    tok_embd     = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD,  "weight"), {n_embd, n_vocab}, 0);
    type_embd    = create_tensor(tn(LLM_TENSOR_TOKEN_TYPES, "weight"), {n_embd, n_token_types}, TENSOR_NOT_REQUIRED);

-    if (arch == LLM_ARCH_BERT) {
-        pos_embd = create_tensor(tn(LLM_TENSOR_POS_EMBD,    "weight"), {n_embd, n_ctx_train}, 0);
+    pos_embd = create_tensor(tn(LLM_TENSOR_POS_EMBD,    "weight"), {n_embd, n_ctx_train}, 0);

-        cls   = create_tensor(tn(LLM_TENSOR_CLS, "weight"), {n_embd, n_embd}, TENSOR_NOT_REQUIRED);
-        cls_b = create_tensor(tn(LLM_TENSOR_CLS, "bias"),   {n_embd},         TENSOR_NOT_REQUIRED);
+    cls   = create_tensor(tn(LLM_TENSOR_CLS, "weight"), {n_embd, n_embd}, TENSOR_NOT_REQUIRED);
+    cls_b = create_tensor(tn(LLM_TENSOR_CLS, "bias"),   {n_embd},         TENSOR_NOT_REQUIRED);

-        cls_out   = create_tensor(tn(LLM_TENSOR_CLS_OUT, "weight"), {n_embd, hparams.n_cls_out}, TENSOR_NOT_REQUIRED);
-        cls_out_b = create_tensor(tn(LLM_TENSOR_CLS_OUT, "bias"),   {hparams.n_cls_out},         TENSOR_NOT_REQUIRED);
-    }
+    cls_out   = create_tensor(tn(LLM_TENSOR_CLS_OUT, "weight"), {n_embd, hparams.n_cls_out}, TENSOR_NOT_REQUIRED);
+    cls_out_b = create_tensor(tn(LLM_TENSOR_CLS_OUT, "bias"),   {hparams.n_cls_out},         TENSOR_NOT_REQUIRED);

    tok_norm   = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD_NORM, "weight", 0), {n_embd}, 0);
    tok_norm_b = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD_NORM, "bias",   0), {n_embd}, 0);
@@ -53,20 +51,10 @@ void llama_model_bert::load_arch_tensors(llama_model_loader &) {
        layer.attn_out_norm   = create_tensor(tn(LLM_TENSOR_ATTN_OUT_NORM, "weight", i), {n_embd}, 0);
        layer.attn_out_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_NORM, "bias", i),   {n_embd}, 0);

-        if (hparams.moe_every_n_layers > 0 && i % hparams.moe_every_n_layers == 1) {
-            layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {  n_embd, n_ff,   n_expert}, 0);
-            layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {  n_ff,   n_embd, n_expert}, 0);
-            layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,   "weight", i), {n_embd, n_expert}, 0);
-        } else {
-            layer.ffn_up     = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff}, 0);
-            layer.ffn_up_b   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "bias", i),   {n_ff}, TENSOR_NOT_REQUIRED);
-            layer.ffn_down   = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
-            layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i),   {n_embd}, TENSOR_NOT_REQUIRED);
-
-            if (arch == LLM_ARCH_NOMIC_BERT) {
-                layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
-            }
-        }
+        layer.ffn_up     = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff}, 0);
+        layer.ffn_up_b   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "bias", i),   {n_ff}, TENSOR_NOT_REQUIRED);
+        layer.ffn_down   = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
+        layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i),   {n_embd}, TENSOR_NOT_REQUIRED);

        layer.layer_out_norm   = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", i), {n_embd}, 0);
        layer.layer_out_norm_b = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "bias", i),   {n_embd}, 0);
@@ -82,7 +82,7 @@ std::unique_ptr<llm_graph_context> llama_model_gemma3n::build_arch_graph(const l
 }

 // get 2D slice view from a 3D tensor, the idx corresponds to the 3rd dim
-static ggml_tensor * ggml_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, int idx) {
+static ggml_tensor * gemma3n_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, int idx) {
    GGML_ASSERT(idx < (int) x->ne[2]);
    return ggml_view_2d(ctx0, x, x->ne[0], x->ne[1], ggml_row_size(x->type, x->ne[0]),
                        idx * x->ne[0] * x->ne[1] * ggml_element_size(x));
@@ -139,7 +139,7 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par
        ggml_tensor * predictions = altup_predict(cur, il);  // [n_embd, n_tokens, n_altup]

        // predicted value will go through self-attention and laurel
-        ggml_tensor * active_prediction = ggml_view_2d_slice(ctx0, predictions, i_altup_act);  // [n_embd, n_tokens]
+        ggml_tensor * active_prediction = gemma3n_view_2d_slice(ctx0, predictions, i_altup_act);  // [n_embd, n_tokens]
        cur = active_prediction;
        cb(cur, "active_prediction", il);

@@ -236,13 +236,13 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par

        ggml_tensor * first_prediction;                                                   // [n_embd, n_tokens]
        {
-            first_prediction = ggml_view_2d_slice(ctx0, corrected, i_altup_act);          // [n_embd, n_tokens]
+            first_prediction = gemma3n_view_2d_slice(ctx0, corrected, i_altup_act);       // [n_embd, n_tokens]
            first_prediction = ggml_mul(ctx0, first_prediction, model.layers[il].altup_correct_scale);
            first_prediction = build_lora_mm(model.layers[il].per_layer_inp_gate, first_prediction);
            first_prediction = ggml_gelu(ctx0, first_prediction);                 // [n_embd_altup, n_tokens]
            cb(first_prediction, "first_prediction_gated", il);

-            ggml_tensor * inp_this_layer = ggml_view_2d_slice(ctx0, inp_per_layer, il);   // [n_embd_altup, n_tokens]
+            ggml_tensor * inp_this_layer = gemma3n_view_2d_slice(ctx0, inp_per_layer, il);   // [n_embd_altup, n_tokens]
            first_prediction = ggml_mul(ctx0, first_prediction, inp_this_layer);  // [n_embd_altup, n_tokens]
            cb(first_prediction, "first_prediction_scaled", il);

@@ -253,7 +253,7 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par
        }
        // equivalent to python code: corrected_predictions[1:] += first_prediction
        {
-            ggml_tensor * slice_first = ggml_view_2d_slice(ctx0, corrected, 0);
+            ggml_tensor * slice_first = gemma3n_view_2d_slice(ctx0, corrected, 0);
            ggml_tensor * slice_rest  = ggml_view_3d(
                ctx0, corrected, n_embd, n_tokens, n_altup - 1, ggml_row_size(corrected->type, n_embd),
                ggml_row_size(corrected->type, n_embd * n_tokens), n_embd * n_tokens * ggml_element_size(corrected));
@@ -271,7 +271,7 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par

    // cur now has multiple altup(s), we want to merge them back to 1 altup
    {
-        ggml_tensor * target_magnitude = calc_magnitude(ggml_view_2d_slice(ctx0, cur, i_altup_act));  // [n_embd, n_tokens]
+        ggml_tensor * target_magnitude = calc_magnitude(gemma3n_view_2d_slice(ctx0, cur, i_altup_act));  // [n_embd, n_tokens]
        // do a view to skip the first slice (active altup)
        ggml_tensor * alt_slice =
            ggml_view_3d(ctx0, cur, n_embd, n_tokens, n_altup - 1, ggml_row_size(cur->type, n_embd),
@@ -283,9 +283,9 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par
        cb(altup_unembd, "altup_unembd", -1);

        // equivalent to torch.mean(hidden_states, dim=0)
-        cur = ggml_view_2d_slice(ctx0, cur, 0);  // [n_embd, n_tokens]
+        cur = gemma3n_view_2d_slice(ctx0, cur, 0);  // [n_embd, n_tokens]
        for (int i = 0; i < n_altup - 1; ++i) {
-            cur = ggml_add(ctx0, cur, ggml_view_2d_slice(ctx0, altup_unembd, i));
+            cur = ggml_add(ctx0, cur, gemma3n_view_2d_slice(ctx0, altup_unembd, i));
        }
        cur = ggml_scale(ctx0, cur, 1.0f / float(n_altup));  // [n_embd, n_tokens]
        cb(cur, "unembd_merged", -1);
@@ -419,7 +419,7 @@ ggml_tensor * llama_model_gemma3n::graph::altup_compute_router_modalities(ggml_t
 // input cur shape: [n_embd, n_tokens, n_altup]
 // output    shape: [n_embd, n_tokens, n_altup]
 ggml_tensor * llama_model_gemma3n::graph::altup_predict(ggml_tensor * cur, int il) {
-    ggml_tensor * activated  = ggml_view_2d_slice(ctx0, cur, i_altup_act);      // [n_embd, n_tokens]
+    ggml_tensor * activated  = gemma3n_view_2d_slice(ctx0, cur, i_altup_act);   // [n_embd, n_tokens]
    ggml_tensor * modalities = altup_compute_router_modalities(activated, il);  // [n_altup, n_tokens]
    cb(modalities, "modalities", il);

@@ -447,7 +447,7 @@ ggml_tensor * llama_model_gemma3n::graph::altup_correct(ggml_tensor * prediction
    ggml_tensor * modalities = altup_compute_router_modalities(activated, il);  // [n_altup, n_tokens]
    cb(modalities, "modalities", il);

-    ggml_tensor * active_prediction = ggml_view_2d_slice(ctx0, predictions, i_altup_act);
+    ggml_tensor * active_prediction = gemma3n_view_2d_slice(ctx0, predictions, i_altup_act);
    ggml_tensor * innovation        = ggml_sub(ctx0, activated, active_prediction);  // [n_embd, n_tokens]
    cb(innovation, "innovation", il);

@@ -145,7 +145,7 @@ std::unique_ptr<llm_graph_context> llama_model_gemma4::build_arch_graph(const ll
 }

 // get 2D slice view from a 3D tensor, the idx corresponds to the 3rd dim
-static ggml_tensor * ggml_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, int idx) {
+static ggml_tensor * gemma4_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, int idx) {
    GGML_ASSERT(idx < (int) x->ne[2]);
    return ggml_view_2d(ctx0, x, x->ne[0], x->ne[1], ggml_row_size(x->type, x->ne[0]),
                        idx * x->ne[0] * x->ne[1] * ggml_element_size(x));
@@ -372,7 +372,7 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
            cur = build_lora_mm(model.layers[il].per_layer_inp_gate, cur); // [n_embd_per_layer, n_tokens]
            cur = ggml_gelu(ctx0, cur);

-            ggml_tensor * inp_this_layer = ggml_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_per_layer, n_tokens]
+            ggml_tensor * inp_this_layer = gemma4_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_per_layer, n_tokens]

            // TODO @ngxson : improve this
            if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
@@ -30,7 +30,7 @@ void llama_model_granite_hybrid::load_arch_hparams(llama_model_loader & ml) {
        case 768: type = LLM_TYPE_350M; break;
        case 1536: type = (hparams.n_ff() == 512 ? LLM_TYPE_7B_A1B : LLM_TYPE_1B); break;
        case 2048: case 2560: type = LLM_TYPE_3B; break;
-        case 4096: type = LLM_TYPE_32B; break;
+        case 4096: type = LLM_TYPE_32B_A9B; break;
        default: type = LLM_TYPE_UNKNOWN;
    }

@@ -9,8 +9,7 @@ void llama_model_granite_moe::load_arch_hparams(llama_model_loader & ml) {

    switch (hparams.n_layer()) {
        case 24: type = LLM_TYPE_1B_A400M; break;
-        case 32: type = LLM_TYPE_3B; break;
-        case 40: type = LLM_TYPE_3B; break;
+        case 32: type = LLM_TYPE_3B_A800M; break;
        // Add additional layer/vocab/etc checks here for other model sizes
        default: type = LLM_TYPE_UNKNOWN;
    }
@@ -38,7 +38,16 @@ void llama_model_granite::load_arch_hparams(llama_model_loader & ml) {

    switch (hparams.n_layer()) {
        case 32: type = LLM_TYPE_3B; break;
-        case 40: type = LLM_TYPE_3B; break;
+        case 40: {
+            switch (hparams.n_embd) {
+                case 2048: type = LLM_TYPE_2B; break;
+                case 2560: type = LLM_TYPE_3B; break;
+                case 4096: type = LLM_TYPE_8B; break;
+                default: type = LLM_TYPE_UNKNOWN;
+            }
+            break;
+        }
+        case 64: type = LLM_TYPE_30B; break;
        // Add additional layer/vocab/etc checks here for other model sizes
        default: type = LLM_TYPE_UNKNOWN;
    }
@@ -19,16 +19,6 @@ void llama_model_jina_bert_v3::load_arch_tensors(llama_model_loader &) {
    tok_embd     = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD,  "weight"), {n_embd, n_vocab}, 0);
    type_embd    = create_tensor(tn(LLM_TENSOR_TOKEN_TYPES, "weight"), {n_embd, n_token_types}, TENSOR_NOT_REQUIRED);

-    if (arch == LLM_ARCH_BERT) {
-        pos_embd = create_tensor(tn(LLM_TENSOR_POS_EMBD,    "weight"), {n_embd, n_ctx_train}, 0);
-
-        cls   = create_tensor(tn(LLM_TENSOR_CLS, "weight"), {n_embd, n_embd}, TENSOR_NOT_REQUIRED);
-        cls_b = create_tensor(tn(LLM_TENSOR_CLS, "bias"),   {n_embd},         TENSOR_NOT_REQUIRED);
-
-        cls_out   = create_tensor(tn(LLM_TENSOR_CLS_OUT, "weight"), {n_embd, hparams.n_cls_out}, TENSOR_NOT_REQUIRED);
-        cls_out_b = create_tensor(tn(LLM_TENSOR_CLS_OUT, "bias"),   {hparams.n_cls_out},         TENSOR_NOT_REQUIRED);
-    }
-
    tok_norm   = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD_NORM, "weight", 0), {n_embd}, 0);
    tok_norm_b = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD_NORM, "bias",   0), {n_embd}, 0);

@@ -43,20 +33,10 @@ void llama_model_jina_bert_v3::load_arch_tensors(llama_model_loader &) {
        layer.attn_out_norm   = create_tensor(tn(LLM_TENSOR_ATTN_OUT_NORM, "weight", i), {n_embd}, 0);
        layer.attn_out_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_NORM, "bias", i),   {n_embd}, 0);

-        if (hparams.moe_every_n_layers > 0 && i % hparams.moe_every_n_layers == 1) {
-            layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {  n_embd, n_ff,   n_expert}, 0);
-            layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {  n_ff,   n_embd, n_expert}, 0);
-            layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,   "weight", i), {n_embd, n_expert}, 0);
-        } else {
-            layer.ffn_up     = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff}, 0);
-            layer.ffn_up_b   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "bias", i),   {n_ff}, TENSOR_NOT_REQUIRED);
-            layer.ffn_down   = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
-            layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i),   {n_embd}, TENSOR_NOT_REQUIRED);
-
-            if (arch == LLM_ARCH_NOMIC_BERT) {
-                layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
-            }
-        }
+        layer.ffn_up     = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff}, 0);
+        layer.ffn_up_b   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "bias", i),   {n_ff}, TENSOR_NOT_REQUIRED);
+        layer.ffn_down   = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
+        layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i),   {n_embd}, TENSOR_NOT_REQUIRED);

        layer.layer_out_norm   = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", i), {n_embd}, 0);
        layer.layer_out_norm_b = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "bias", i),   {n_embd}, 0);
@@ -229,6 +229,7 @@ llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_
    ggml_set_input(inp->inp_k_decay);
    cb(inp->inp_k_decay, "k_decay_exp", -1);

+    // [TAG_RESERVE_DIAG_DECAY]
    inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs);
    ggml_set_input(inp->inp_diag_decay);
    cb(inp->inp_diag_decay, "diag_decay_exp", -1);
--- a/Show More
+++ b/Show More