vulkan: fix missing view-alias dependencies in ggml_vk_graph_optimize (#27812 )

* vulkan: fix missing view-alias dependencies in ggml_vk_graph_optimize is_src_of doesn't treat two views of one tensor as dependent, so the optimizer reorders nodes across aliased reads and writes. Result: silently wrong tokens under greedy decoding, different output on every server start, and invalid speculative-decoding acceptance, with nothing logged. Hits Qwen3.8's recurrent state (and any model with view-aliased state) on AMD and NVIDIA Vulkan. CUDA is clean. Compare view_src bases on both sides. Fixes #27805 * vulkan: don't treat view/no-op nodes as aliasing dependencies Nodes whose op is NONE, RESHAPE, TRANSPOSE, VIEW or PERMUTE execute nothing, so aliasing through them is not a real dependency. The previous base comparison matched them anyway, which only costs the optimizer reordering freedom. Co-authored-by: Jeff Bolz <jbolz@nvidia.com> * vulkan: make the lambda parameter const and capture is_empty in is_src_of Code will not compile without these changes. is_src_of has an empty capture list, so is_empty was not visible inside it, and is_empty took a non-const pointer, while is_src_of receives const ones. Other call sites pass non-const pointers, which still convert as usual. --------- Co-authored-by: Jeff Bolz <jbolz@nvidia.com>
ggml : fix conv_transpose_2d for multiple batches (#26132 )
2026-08-28 19:17:47 +02:00 · 2026-08-28 19:12:33 +02:00 · 2026-08-28 20:09:08 +03:00 · 2026-08-28 16:52:49 +02:00 · 2026-08-28 16:34:26 +03:00 · 2026-08-28 15:37:37 +03:00
255 changed files with 15459 additions and 6273 deletions
@@ -1,12 +1,12 @@
-ARG OPENVINO_VERSION_MAJOR=2026.3
-ARG OPENVINO_VERSION_FULL=2026.3.0.22451.bd8d6542e3c
+ARG OPENVINO_VERSION_MAJOR=2026.3.1
+ARG OPENVINO_VERSION_FULL=2026.3.1.22476.56d9685302d
 ARG UBUNTU_VERSION=24.04

 # Intel GPU driver versions. https://github.com/intel/compute-runtime/releases
-ARG IGC_VERSION=v2.38.2
-ARG IGC_VERSION_FULL=2_2.38.2+22051
-ARG COMPUTE_RUNTIME_VERSION=26.27.39122.11
-ARG COMPUTE_RUNTIME_VERSION_FULL=26.27.39122.11-0
+ARG IGC_VERSION=v2.40.13
+ARG IGC_VERSION_FULL=2_2.40.13+22418
+ARG COMPUTE_RUNTIME_VERSION=26.31.39395.13
+ARG COMPUTE_RUNTIME_VERSION_FULL=26.31.39395.13-0
 ARG IGDGMM_VERSION=22.10.0

 # Intel NPU driver versions. https://github.com/intel/linux-npu-driver/releases
@@ -41,8 +41,8 @@ jobs:

    env:
      # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3"
-      OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c"
+      OPENVINO_VERSION_MAJOR: "2026.3.1"
+      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"

    steps:
      - name: Clone
@@ -69,8 +69,8 @@ jobs:

    env:
      # Sync versions in build.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3"
-      OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c"
+      OPENVINO_VERSION_MAJOR: "2026.3.1"
+      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"

    steps:
      - name: Clone
@@ -32,6 +32,8 @@ env:
  LLAMA_ARG_LOG_COLORS: 1
  LLAMA_ARG_LOG_PREFIX: 1
  LLAMA_ARG_LOG_TIMESTAMPS: 1
+  # TODO: fix and re-enable the `test-llama-archs` and `test-recurrent-state-rollback`
+  CTEST_EXCLUDE: "test-llama-archs|^test-recurrent-state-rollback"

 jobs:
  ubuntu-24-openvino:
@@ -39,8 +41,8 @@ jobs:

    env:
      # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3"
-      OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c"
+      OPENVINO_VERSION_MAJOR: "2026.3.1"
+      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"

    steps:
      - name: Clone
@@ -78,26 +80,24 @@ jobs:

      - name: Test (CPU)
        id: cmake_test_cpu
-        # TODO: fix and re-enable the `test-llama-archs` test below
        run: |
          cd ${{ github.workspace }}
-          ctest --test-dir build/ReleaseOV -L main -E "test-llama-archs|test-recurrent-state-rollback-nemotron-h" --verbose --timeout 2000
+          ctest --test-dir build/ReleaseOV -L main -E "${{ env.CTEST_EXCLUDE }}" --verbose --timeout 3000

      - name: Test (GPU)
        id: cmake_test_gpu
-        # TODO: fix and re-enable the `test-llama-archs` test below
        run: |
          cd ${{ github.workspace }}
          export GGML_OPENVINO_DEVICE=GPU
-          ctest --test-dir build/ReleaseOV -L main -E "test-llama-archs|test-recurrent-state-rollback-nemotron-h" --verbose --timeout 3000
+          ctest --test-dir build/ReleaseOV -L main -E "${{ env.CTEST_EXCLUDE }}" --verbose --timeout 3000

  openvino-windows-2022:
    runs-on: windows-2022

    env:
      # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3"
-      OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c"
+      OPENVINO_VERSION_MAJOR: "2026.3.1"
+      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"

    steps:
      - name: Clone
@@ -159,14 +159,13 @@ jobs:
      - name: Test (CPU)
        id: cmake_test_cpu
        shell: cmd
-        # TODO: fix and re-enable the `test-llama-archs` test below
        run: |
          REM Find extracted OpenVINO folder dynamically
          for /d %%i in (openvino_toolkit\*) do set OPENVINO_ROOT=%%i
          call "%OPENVINO_ROOT%\setupvars.bat"

          cd build
-          ctest --test-dir ReleaseOV -L main -E "test-llama-archs|test-recurrent-state-rollback-nemotron-h" -C Release --verbose --timeout 3000
+          ctest --test-dir ReleaseOV -L main -E "${{ env.CTEST_EXCLUDE }}" -C Release --verbose --timeout 3000

      - name: ccache-clear
        uses: ./.github/actions/ccache-clear
@@ -288,8 +288,8 @@ jobs:

    env:
      # Sync versions in build.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3"
-      OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c"
+      OPENVINO_VERSION_MAJOR: "2026.3.1"
+      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"

    steps:
      - name: Clone
@@ -64,7 +64,7 @@ jobs:
    needs: create_tag
    uses: ./.github/workflows/ui-build.yml
    with:
-      hf_ui_version: ${{ needs.create_tag.outputs.source_tag }}
+      ui_version: ${{ needs.create_tag.outputs.source_tag }}

  prepare_matrices:
    name: Prepare Docker matrices
@@ -162,7 +162,7 @@ jobs:
        if: ${{ matrix.config.prebuilt_ui == true }}
        uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist

      - name: Set up QEMU
@@ -61,31 +61,8 @@ jobs:
            echo "should_release=false" >> $GITHUB_OUTPUT
          fi

-  get-version:
-    runs-on: ubuntu-slim
-    outputs:
-      ui_version: ${{ steps.version.outputs.ui_version }}
-    steps:
-      - uses: actions/checkout@v6
-        with:
-          fetch-depth: 0
-      - id: version
-        run: |
-          # Resolve UI version: BUILD_NUMBER from cmake/build-info.cmake > git hash + epoch > fallback
-          version=""
-          if grep -q "BUILD_NUMBER" cmake/build-info.cmake; then
-            build_number=$(grep "set(BUILD_NUMBER" cmake/build-info.cmake | grep -oP '\d+')
-            if [ -n "$build_number" ] && [ "$build_number" -gt 0 ]; then
-              version="b${build_number}"
-            fi
-          fi
-          if [ -z "$version" ]; then
-            version=$(git rev-parse --short HEAD)-$(date +%s)
-          fi
-          echo "ui_version=${version}" >> $GITHUB_OUTPUT
-
  macos-cpu:
-    needs: [check-release, get-version]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
    strategy:
      matrix:
@@ -119,12 +96,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
@@ -141,7 +117,6 @@ jobs:
            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
            -DLLAMA_FATAL_WARNINGS=ON \
            -DLLAMA_BUILD_BORINGSSL=ON \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
            ${{ env.CMAKE_ARGS }}
          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu)

@@ -167,7 +142,7 @@ jobs:
          key: release-${{ matrix.os }}-${{ matrix.arch }}

  ubuntu-cpu:
-    needs: [check-release, get-version]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
    strategy:
      matrix:
@@ -191,12 +166,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: Dependencies
        id: depends
@@ -227,7 +201,6 @@ jobs:
            -DGGML_NATIVE=OFF \
            -DGGML_CPU_ALL_VARIANTS=ON \
            -DLLAMA_FATAL_WARNINGS=ON \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
            ${{ env.CMAKE_ARGS }}
          cmake --build build --config Release -j $(nproc)

@@ -254,7 +227,7 @@ jobs:
          key: release-${{ matrix.os }}-cpu

  ubuntu-vulkan:
-    needs: [check-release, get-version]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    strategy:
@@ -277,12 +250,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: Dependencies
        id: depends
@@ -314,7 +286,6 @@ jobs:
            -DGGML_NATIVE=OFF \
            -DGGML_CPU_ALL_VARIANTS=ON \
            -DGGML_VULKAN=ON \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
            ${{ env.CMAKE_ARGS }}
          cmake --build build --config Release -j $(nproc)

@@ -340,7 +311,7 @@ jobs:
          key: release-${{ matrix.os }}-vulkan

  android-arm64:
-    needs: [check-release, get-version]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: ubuntu-latest
@@ -358,12 +329,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: Set up JDK
        uses: actions/setup-java@v5
@@ -407,7 +377,6 @@ jobs:
            -DLLAMA_FATAL_WARNINGS=ON \
            -DGGML_OPENMP=OFF \
            -DLLAMA_BUILD_BORINGSSL=ON \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
            ${{ env.CMAKE_ARGS }}
          cmake --build build --config Release -j $(nproc)

@@ -433,7 +402,7 @@ jobs:
          name: llama-bin-android-arm64.tar.gz

  ubuntu-24-openvino:
-    needs: [check-release, get-version]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: ubuntu-24.04
@@ -446,8 +415,8 @@ jobs:

    env:
      # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3"
-      OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c"
+      OPENVINO_VERSION_MAJOR: "2026.3.1"
+      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"

    steps:
      - name: Set OpenVINO version output
@@ -460,12 +429,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
@@ -508,7 +476,6 @@ jobs:
            -DGGML_OPENVINO=ON \
            -DCMAKE_INSTALL_RPATH='$ORIGIN' \
            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
            ${{ env.CMAKE_ARGS }}
          cmake --build build/ReleaseOV --config Release --parallel

@@ -552,7 +519,7 @@ jobs:
          key: release-ubuntu-24.04-openvino-release-no-preset-v1

  windows-openvino:
-    needs: [check-release]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: windows-2022
@@ -562,8 +529,8 @@ jobs:

    env:
      # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile
-      OPENVINO_VERSION_MAJOR: "2026.3"
-      OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c"
+      OPENVINO_VERSION_MAJOR: "2026.3.1"
+      OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d"

    steps:
      - name: Set OpenVINO version output
@@ -577,12 +544,11 @@ jobs:
        with:
            fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
@@ -682,7 +648,7 @@ jobs:

  windows-cpu:
    name: windows-cpu / ${{ matrix.arch }}
-    needs: [check-release]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: windows-2025-vs2026
@@ -702,12 +668,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: Install Ninja
        run: |
@@ -749,6 +714,8 @@ jobs:
        with:
          key: release-windows-2025-vs2026-${{ matrix.arch }}-cpu

+  # note: builds only the ggml-hip backend - llama-server is injected from the
+  #       windows-cpu zip during the release "Merge artifacts" step
  windows-rocm:
    needs: [check-release]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
@@ -769,6 +736,10 @@ jobs:
        with:
          fetch-depth: 0

+      - name: Install Ninja
+        run: |
+          choco install ninja
+
      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
        with:
@@ -822,33 +793,28 @@ jobs:

      - name: Build
        run: |
-          mkdir build
-          cd build
-          cmake .. `
-            -G "Unix Makefiles" `
+          cmake -S . -B build `
+            -G "Ninja Multi-Config" `
            -DCMAKE_PREFIX_PATH="${env:HIP_PATH}" `
-            -DCMAKE_BUILD_TYPE=Release `
            -DGGML_BACKEND_DL=ON `
            -DGGML_NATIVE=OFF `
-            -DGGML_CPU=ON `
-            -DGGML_CPU_ALL_VARIANTS=ON `
+            -DGGML_CPU=OFF `
            -DGGML_HIP=ON `
            -DCMAKE_C_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" `
            -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang++.exe" `
            -DCMAKE_C_FLAGS="-Wno-error=incompatible-pointer-types" `
            -DCMAKE_HIP_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" `
            -DHIP_PATH="${env:HIP_PATH}" `
-            -DGGML_HIP_ROCWMMA_FATTN=ON `
            -DAMDGPU_TARGETS="${{ matrix.gpu_targets }}"
-          cmake --build . --config Release --parallel ${env:NUMBER_OF_PROCESSORS}
+          cmake --build build --config Release --parallel ${env:NUMBER_OF_PROCESSORS} --target ggml-hip

      - name: Verify HIP backend was built
        run: |
-          $hipDll = Get-ChildItem -Path build\bin -Filter "ggml-hip*.dll" -ErrorAction SilentlyContinue
+          $hipDll = Get-ChildItem -Path build\bin\Release -Filter "ggml-hip*.dll" -ErrorAction SilentlyContinue
          if (-not $hipDll) {
            Write-Host "##[error]ggml-hip*.dll was NOT produced. The HIP backend silently failed to build."
-            Write-Host "Contents of build\bin:"
-            Get-ChildItem build\bin | Format-Table -AutoSize
+            Write-Host "Contents of build\bin\Release:"
+            Get-ChildItem build\bin\Release | Format-Table -AutoSize
            exit 1
          }
          Write-Host "HIP backend artifact found:"
@@ -863,10 +829,40 @@ jobs:
          $rocmVersionShort = ('${{ matrix.ROCM_VERSION }}'.Split('.')[0..1] -join '.')
          echo "ROCM_VERSION_SHORT=$rocmVersionShort" >> $env:GITHUB_ENV

+      - name: Bundle HIP runtime DLLs (amdhip64_7.dll, rocm_kpack.dll, amd_comgr.dll)
+        run: |
+          $ErrorActionPreference = "Stop"
+          # See issue https://github.com/ggml-org/llama.cpp/issues/26929.
+          # ggml-hip.dll loads amdhip64_7.dll at run time. The Adrenalin driver
+          # ships an amdhip64_7.dll in System32, which the loader searches before PATH,
+          # so a matching DLL from PATH cannot win. Copy amdhip64 next to the
+          # binaries (exe directory is searched before System32) so the correct
+          # runtime is used. rocm_kpack.dll is amdhip64_7's direct dependency, so
+          # copy the matching version too. amd_comgr is copied as well to keep it
+          # in sync with the bundled amdhip64, avoiding a version mismatch with a
+          # amd_comgr from System32.
+          # rocblas/hipblaslt kernels resolve fine via PATH and are not copied.
+          $binPath = (rocm-sdk path --bin).Trim()
+          if (-not $binPath) { throw "rocm-sdk path --bin returned empty" }
+          write-host "ROCm bin path: $binPath"
+
+          $patterns = @("amdhip64_7.dll", "rocm_kpack.dll", "amd_comgr.dll")
+          foreach ($pattern in $patterns) {
+            $files = Get-ChildItem -Path $binPath -Filter $pattern -ErrorAction SilentlyContinue
+            if (-not $files) { throw "no match for $pattern in $binPath" }
+            foreach ($f in $files) {
+              Copy-Item $f.FullName -Destination build\bin\Release -Force
+              write-host "  copied $($f.Name)"
+            }
+          }
+
      - name: Pack artifacts
        run: |
-          cp "LICENSE" "build\bin\"
-          7z a -snl llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip .\build\bin\*
+          7z a -snl llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip `
+            .\build\bin\Release\ggml-hip.dll `
+            .\build\bin\Release\amdhip64_7.dll `
+            .\build\bin\Release\rocm_kpack.dll `
+            .\build\bin\Release\amd_comgr.dll

      - name: Upload artifacts
        uses: actions/upload-artifact@v6
@@ -879,6 +875,8 @@ jobs:
        with:
          key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}

+  # note: builds only the backend library - llama-server (with the embedded UI)
+  #       is injected from the windows-cpu zip during the release "Merge artifacts" step
  windows:
    needs: [check-release]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
@@ -909,13 +907,6 @@ jobs:
        id: checkout
        uses: actions/checkout@v6

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
-        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
-
      - name: Install Vulkan SDK
        id: get_vulkan
        if: ${{ matrix.backend == 'vulkan' }}
@@ -978,6 +969,8 @@ jobs:
          path: llama-bin-win-${{ matrix.backend }}-${{ matrix.arch }}.zip
          name: llama-bin-win-${{ matrix.backend }}-${{ matrix.arch }}.zip

+  # note: builds only the ggml-cuda backend - llama-server is injected from the
+  #       windows-cpu zip during the release "Merge artifacts" step
  windows-cuda:
    name: windows-cuda (${{ matrix.cuda }}, ${{ matrix.arch }})
    needs: [check-release]
@@ -1006,13 +999,6 @@ jobs:
        id: checkout
        uses: actions/checkout@v6

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
-        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
-
      - name: Install Cuda Toolkit
        uses: ./.github/actions/windows-setup-cuda
        with:
@@ -1084,6 +1070,8 @@ jobs:
        with:
          key: release-windows-2022-${{ matrix.arch }}-cuda-${{ matrix.cuda }}

+  # note: builds only the ggml-sycl backend - llama-server is injected from the
+  #       windows-cpu zip during the release "Merge artifacts" step
  windows-sycl:
    needs: [check-release]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
@@ -1118,13 +1106,6 @@ jobs:
          Expand-Archive -Path "level-zero-win-sdk.zip" -DestinationPath "C:/level-zero-sdk" -Force
          "LEVEL_ZERO_V1_SDK_PATH=C:/level-zero-sdk" | Out-File -FilePath $env:GITHUB_ENV -Append

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
-        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
-
      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
        with:
@@ -1195,7 +1176,7 @@ jobs:
          key: release-windows-2022-x64-sycl

  ubuntu-24-sycl:
-    needs: [check-release]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    strategy:
@@ -1237,12 +1218,11 @@ jobs:
          wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/level-zero-devel_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O level-zero-devel.deb
          sudo apt-get install -y ./level-zero.deb ./level-zero-devel.deb

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
@@ -1288,7 +1268,7 @@ jobs:
          key: release-ubuntu-24.04-sycl-${{ matrix.build }}

  ubuntu-24-rocm:
-    needs: [check-release, get-version]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: ubuntu-24.04
@@ -1310,12 +1290,11 @@ jobs:
        with:
          fetch-depth: 0

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
+      - name: Download UI build
+        uses: actions/download-artifact@v7
        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
+          name: llama-ui.zip
+          path: tools/ui/dist

      - name: Free up disk space
        uses: ggml-org/free-disk-space@v1.3.1
@@ -1388,7 +1367,6 @@ jobs:
            -DGPU_TARGETS="${{ matrix.gpu_targets }}" \
            -DGGML_HIP=ON \
            -DHIP_PLATFORM=amd \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
            ${{ env.CMAKE_ARGS }}
          cmake --build build --config Release -j $(nproc)

@@ -1417,7 +1395,7 @@ jobs:
          key: release-ubuntu-24.04-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}

  ios-xcode:
-    needs: [check-release, get-version]
+    needs: [check-release]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
    runs-on: macos-26

@@ -1445,8 +1423,7 @@ jobs:
            -DLLAMA_BUILD_SERVER=OFF \
            -DCMAKE_SYSTEM_NAME=iOS \
            -DCMAKE_OSX_DEPLOYMENT_TARGET=16.0 \
-            -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }}
+            -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml
          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu) -- CODE_SIGNING_ALLOWED=NO

      - name: xcodebuild for swift package
@@ -1569,11 +1546,9 @@ jobs:
 #          name: llama-bin-${{ matrix.chip_type }}-openEuler-${{ matrix.arch }}${{ matrix.use_acl_graph == 'on' && '-aclgraph' || '' }}.tar.gz

  ui-build:
-    needs: [check-release, get-version]
+    needs: [check-release]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}
    uses: ./.github/workflows/ui-build.yml
-    with:
-      hf_ui_version: ${{ needs.get-version.outputs.ui_version }}

  release:
    if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
@@ -1588,7 +1563,6 @@ jobs:
    runs-on: ubuntu-slim

    needs:
-      - get-version
      - windows
      - windows-cpu
      - windows-cuda
@@ -1628,24 +1602,27 @@ jobs:
          path: ./artifact
          merge-multiple: true

-      - name: Move artifacts
+      - name: Merge artifacts
        id: move_artifacts
        run: |
          mkdir -p release

-          echo "Adding CPU backend files to existing zips..."
+          # the windows-cpu zip contains the full toolset (llama-server with the embedded
+          # UI, ggml-cpu) - inject it into the other windows zips so that every archive
+          # ships the same binaries, only with a different backend library on top
+          echo "Injecting windows-cpu binaries (llama-server + CPU backend) into the backend zips..."
          for arch in x64 arm64; do
            cpu_zip="artifact/llama-bin-win-cpu-${arch}.zip"
            temp_dir=$(mktemp -d)
-            echo "Extracting CPU backend for $arch..."
+            echo "Extracting windows-cpu-${arch} package..."
            unzip "$cpu_zip" -d "$temp_dir"

-            echo "Adding CPU files to $arch zips..."
+            echo "Merging into $arch zips..."
            for target_zip in artifact/llama-bin-win-*-${arch}.zip; do
              if [[ "$target_zip" == "$cpu_zip" ]]; then
                continue
              fi
-              echo "Adding CPU backend to $(basename "$target_zip")"
+              echo "Injecting into $(basename "$target_zip")"
              realpath_target_zip=$(realpath "$target_zip")
              (cd "$temp_dir" && zip -r "$realpath_target_zip" .)
            done
@@ -1669,7 +1646,7 @@ jobs:
        id: download_ui
        uses: actions/download-artifact@v7
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: ./ui-dist

      - name: Package UI
@@ -73,13 +73,6 @@ jobs:
          fetch-depth: 0
          ref: ${{ github.event.inputs.sha || github.event.pull_request.head.sha || github.sha || github.head_ref || github.ref_name }}

-      - name: Setup Node.js
-        uses: actions/setup-node@v6
-        with:
-          node-version: "24"
-          cache: "npm"
-          cache-dependency-path: "tools/ui/package-lock.json"
-
      - name: Build
        id: cmake_build
        run: |
@@ -31,6 +31,6 @@ jobs:
      - name: Upload built UI
        uses: actions/upload-artifact@v6
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/
          retention-days: 1
@@ -3,8 +3,8 @@ name: UI Build
 on:
  workflow_call:
    inputs:
-      hf_ui_version:
-        description: 'Version string for version.json (e.g. 12345)'
+      ui_version:
+        description: 'Version string embedded in build.json (e.g. b1234); defaults to b<commit-count>'
        required: false
        type: string

@@ -17,6 +17,17 @@ jobs:
    steps:
      - name: Checkout code
        uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+
+      - name: Resolve UI version
+        id: version
+        run: |
+          version="${{ inputs.ui_version }}"
+          if [ -z "$version" ]; then
+            version="b$(git rev-list --count HEAD)"
+          fi
+          echo "ui_version=${version}" >> $GITHUB_OUTPUT

      - name: Setup Node.js
        uses: actions/setup-node@v6
@@ -31,8 +42,7 @@ jobs:

      - name: Build application
        env:
-          HF_UI_VERSION: ${{ inputs.hf_ui_version || '' }}
-          LLAMA_BUILD_NUMBER: ${{ inputs.hf_ui_version || 'b0000' }}
+          LLAMA_BUILD_NUMBER: ${{ steps.version.outputs.ui_version }}
        run: npm run build
        working-directory: tools/ui

@@ -43,6 +53,6 @@ jobs:
      - name: Upload built UI
        uses: actions/upload-artifact@v6
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/
          retention-days: 1
@@ -37,7 +37,7 @@ jobs:
      - name: Download UI build artifact
        uses: actions/download-artifact@v7
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/

      - name: Create distribution archive
@@ -64,7 +64,7 @@ jobs:
      - name: Download built UI artifacts
        uses: actions/download-artifact@v6
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/

      - name: Run type checking
@@ -106,7 +106,7 @@ jobs:
      - name: Download built UI artifacts
        uses: actions/download-artifact@v6
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/

      - name: Build Storybook
@@ -63,7 +63,7 @@ jobs:
      - name: Download built UI artifacts
        uses: actions/download-artifact@v6
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/

      - name: Install dependencies
@@ -126,7 +126,7 @@ jobs:
      - name: Download built UI artifacts (reuses ui-build)
        uses: actions/download-artifact@v6
        with:
-          name: ui-build
+          name: llama-ui.zip
          path: tools/ui/dist/

      - name: Install Playwright browsers
@@ -134,8 +134,8 @@ option(LLAMA_BUILD_TOOLS     "llama: build tools"
 option(LLAMA_BUILD_EXAMPLES  "llama: build examples"                                                             ${LLAMA_STANDALONE})
 option(LLAMA_BUILD_SERVER    "llama: build server example"                                                       ${LLAMA_STANDALONE})
 option(LLAMA_BUILD_APP       "llama: build the unified binary"                                                   ${LLAMA_STANDALONE})
-option(LLAMA_BUILD_UI        "llama: build the embedded Web UI for server"                                       ON)
-option(LLAMA_USE_PREBUILT_UI "llama: use prebuilt UI from HF Bucket when available (requires LLAMA_BUILD_UI=ON)" ON)
+option(LLAMA_BUILD_UI        "llama: build the embedded Web UI for server"                                       OFF)
+option(LLAMA_USE_PREBUILT_UI "llama: use prebuilt UI from HF Bucket when available" ON)

 option(LLAMA_TOOLS_INSTALL "llama: install tools" ${LLAMA_TOOLS_INSTALL_DEFAULT})
 option(LLAMA_TESTS_INSTALL "llama: install tests" ON)
@@ -189,8 +189,8 @@ if [ ! -z ${GG_BUILD_OPENVINO} ]; then
    fi
    CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_OPENVINO=ON"

-    # TODO: fix and re-enable the `test-llama-archs` test below
-    CTEST_EXTRA="-E test-llama-archs|test-recurrent-state-rollback-nemotron-h"
+    # TODO: fix and re-enable the `test-llama-archs` and `test-recurrent-state-rollback*`
+    CTEST_EXTRA="-E test-llama-archs|^test-recurrent-state-rollback"
 fi

 ## helpers
@@ -1643,6 +1643,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            }
        }
    ).set_env("LLAMA_ARG_CTX_SIZE"));
+    add_opt(common_arg(
+        { "--kv-unified-per-slot" }, "N",
+        "context limit per parallel slot (default: unset, behavior unchanged).\n"
+        "when set without -c/--ctx-size, the shared KV pool is sized to n_parallel*N",
+        [](common_params & params, int value) {
+            params.kv_unified_per_slot = value;
+        }
+    ).set_env("LLAMA_ARG_KV_UNIFIED_PER_SLOT").set_examples({ LLAMA_EXAMPLE_SERVER }));
    add_opt(common_arg(
        {"-n", "--predict", "--n-predict"}, "N",
        string_format(
@@ -2644,6 +2652,27 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            params.mtmd_batch_max_tokens = value;
        }
    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MTMD_BATCH_MAX_TOKENS"));
+    add_opt(common_arg(
+        {"--video-fps"}, "N",
+        string_format("target video frame rate (default: %.1f)", params.video_fps),
+        [](common_params & params, const std::string & value) {
+            params.video_fps = std::stof(value);
+        }
+    ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_FPS"));
+    add_opt(common_arg(
+        {"--video-timestamp-interval"}, "N",
+        string_format("interval in milliseconds between text timestamps (default: %" PRId64 ")", params.video_timestamp_interval_ms),
+        [](common_params & params, int value) {
+            params.video_timestamp_interval_ms = value;
+        }
+    ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_TIMESTAMP_INTERVAL"));
+    add_opt(common_arg(
+        {"--video-ffmpeg-dir"}, "DIR",
+        "path to the directory containing ffmpeg and ffprobe (default: search in PATH)",
+        [](common_params & params, const std::string & value) {
+            params.video_ffmpeg_bin_dir = value;
+        }
+    ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_FFMPEG_DIR"));
    if (params.is_gen_docs || llama_supports_rpc()) {
        add_opt(common_arg(
            {"--rpc"}, "SERVERS",
@@ -2699,6 +2728,19 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            else { throw std::invalid_argument("invalid value"); }
        }
    ).set_env("LLAMA_ARG_LOAD_MODE"));
+    add_opt(common_arg(
+        {"--tensor-read-lazy"}, "MODE",
+        "on-demand reading of certain tensors, for example per-layer embeddings (default: auto)\n"
+        "- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)\n"
+        "- auto: on, but only for tensors larger than 4 GiB\n"
+        "- off: always keep them resident",
+        [](common_params & params, const std::string & value) {
+            /**/ if (value == "on")   { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_ON;   }
+            else if (value == "auto") { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_AUTO; }
+            else if (value == "off")  { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_OFF;  }
+            else { throw std::invalid_argument("invalid value"); }
+        }
+    ).set_env("LLAMA_ARG_TENSOR_READ_LAZY"));
    add_opt(common_arg(
        {"--numa"}, "TYPE",
        "attempt optimizations that help on some NUMA systems\n"
@@ -2750,14 +2792,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            if (value < 0) {
                throw std::invalid_argument("invalid value");
            }
-            for (int i = 0; i < value; ++i) {
-                // keep strings alive and avoid leaking memory by storing them in a static vector
-                static std::list<std::string> buft_overrides;
-                buft_overrides.push_back(llm_ffn_exps_block_regex(i));
-                params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()});
-            }
+            llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.tensor_buft_overrides);
        }
    ).set_env("LLAMA_ARG_N_CPU_MOE"));
+    add_opt(common_arg(
+        {"-ncffn", "--n-cpu-ffn"}, "N",
+        "keep the dense FFN weights of the first N layers in the CPU\n"
+        "(dense models; for MoE expert weights use --n-cpu-moe)",
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("invalid value");
+            }
+            llm_add_n_cpu_ffn_overrides(value, LLM_FFN_DENSE_REGEX, params.tensor_buft_overrides);
+        }
+    ).set_env("LLAMA_ARG_N_CPU_FFN"));
    GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0
    add_opt(common_arg(
        {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N",
@@ -4084,11 +4132,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            if (value < 0) {
                throw std::invalid_argument("invalid value");
            }
-            for (int i = 0; i < value; ++i) {
-                static std::list<std::string> buft_overrides_draft;
-                buft_overrides_draft.push_back(llm_ffn_exps_block_regex(i));
-                params.speculative.draft.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type()});
-            }
+            llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.speculative.draft.tensor_buft_overrides);
        }
    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE"));

@@ -4109,6 +4153,38 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            params.speculative.draft.n_min = value;
        }
    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MIN"));
+    add_opt(common_arg(
+        {"--spec-synth-len"}, "L",
+        "target mean synthetic acceptance length, including the target token (benchmarking only)",
+        [](common_params & params, const std::string & value) {
+            const std::string text = string_strip(value);
+            size_t pos = 0;
+            const double length = std::stod(text, &pos);
+            if (pos != text.size() || length == -1.0) {
+                throw std::invalid_argument("invalid value");
+            }
+            params.speculative.synth_len = length;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_SYNTH_LEN"));
+    add_opt(common_arg(
+        {"--spec-synth-rates"}, "P0,P1,...",
+        "comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)",
+        [](common_params & params, const std::string & value) {
+            const auto values = string_split<std::string>(value, ',');
+            std::vector<double> rates;
+            rates.reserve(values.size());
+            for (const auto & raw : values) {
+                const std::string text = string_strip(raw);
+                size_t pos = 0;
+                const double rate = std::stod(text, &pos);
+                if (pos != text.size()) {
+                    throw std::invalid_argument("invalid value");
+                }
+                rates.push_back(rate);
+            }
+            params.speculative.synth_rates = std::move(rates);
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_SYNTH_RATES"));

    add_opt(common_arg(
        {"--spec-draft-p-split", "--draft-p-split"}, "P",
@@ -1688,6 +1688,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
    mparams.main_gpu        = params.main_gpu;
    mparams.split_mode      = params.split_mode;
    mparams.load_mode       = params.load_mode;
+    mparams.tensor_read_lazy = params.tensor_read_lazy;
    mparams.tensor_split    = params.tensor_split;
    mparams.check_tensors   = params.check_tensors;
    mparams.use_extra_bufts = !params.no_extra_bufts;
@@ -8,6 +8,7 @@
 #include "ggml.h"
 #include "llama.h"

+#include <list>
 #include <set>
 #include <sstream>
 #include <string>
@@ -369,6 +370,9 @@ struct common_params_speculative_ngram_cache {
 struct common_params_speculative {
    std::vector<enum common_speculative_type> types = { COMMON_SPECULATIVE_TYPE_NONE };

+    double synth_len = -1.0;
+    std::vector<double> synth_rates;
+
    // used by Simple, MTP, Eagle3, etc. - all methods that require some kind of draft model
    common_params_speculative_draft draft;

@@ -383,6 +387,10 @@ struct common_params_speculative {
        return !draft.mparams.empty();
    }

+    bool has_synth() const {
+        return synth_len != -1.0 || !synth_rates.empty();
+    }
+
    uint32_t need_n_rs_seq() const {
        bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) {
            return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;
@@ -475,6 +483,8 @@ struct common_params {
    enum llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER; // how to split the model across GPUs
    enum llama_load_mode  load_mode  = LLAMA_LOAD_MODE_AUTO; // how to load the model

+    enum llama_tensor_read_lazy tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_AUTO; // on-demand reading of tensors marked by the arch
+
    common_cpu_params cpuparams;
    common_cpu_params cpuparams_batch;

@@ -589,6 +599,11 @@ struct common_params {
    int image_max_tokens = -1;
    int mtmd_batch_max_tokens = 1024;

+    // for video input
+    float       video_fps                   = 4.0f;
+    int64_t     video_timestamp_interval_ms = 5000;
+    std::string video_ffmpeg_bin_dir        = "";
+
    // finetune
    struct lr_opt lr;
    enum ggml_opt_optimizer_type optimizer = GGML_OPT_OPTIMIZER_TYPE_ADAMW;
@@ -612,6 +627,7 @@ struct common_params {
    bool    cache_prompt        = true;  // whether to enable prompt caching
    bool    cache_idle_slots    = true;  // save and clear idle slots upon starting a new task
    int32_t n_ctx_checkpoints   = 32;    // max number of context checkpoints per slot
+    int32_t kv_unified_per_slot = 0;     // max context per parallel slot; 0 = unset
    int32_t checkpoint_min_step = 8192;  // minimum spacing between context checkpoints
    int32_t cache_ram_mib       = 8192;  // -1 = no limit, 0 - disable, 1 = 1 MiB, etc.

@@ -1108,19 +1124,30 @@ const char * const LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count";
 }

 //
-// MoE utils
+// FFN offload utils
 //

 const char * const LLM_FFN_EXPS_REGEX = "\\.ffn_(up|down|gate|gate_up)_(ch|)exps";

-inline std::string llm_ffn_exps_block_regex(int idx) {
-    return string_format("blk\\.%d%s", idx, LLM_FFN_EXPS_REGEX);
+const char * const LLM_FFN_DENSE_REGEX = "\\.ffn_(up|down|gate)\\.";
+
+inline std::string llm_ffn_block_regex(int idx, const char * ffn_regex) {
+    return string_format("blk\\.%d%s", idx, ffn_regex);
 }

 inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() {
    return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() };
 }

+inline void llm_add_n_cpu_ffn_overrides(int n, const char * ffn_regex, std::vector<llama_model_tensor_buft_override> & overrides) {
+    // keep strings alive and avoid leaking memory by storing them in a static list
+    static std::list<std::string> buft_override_strings;
+    for (int i = 0; i < n; ++i) {
+        buft_override_strings.push_back(llm_ffn_block_regex(i, ffn_regex));
+        overrides.push_back({buft_override_strings.back().c_str(), ggml_backend_cpu_buffer_type()});
+    }
+}
+
 //
 // training utils
 //
@@ -14,6 +14,7 @@

 #include <algorithm>
 #include <cassert>
+#include <cmath>
 #include <cstring>
 #include <iomanip>
 #include <map>
@@ -138,6 +139,7 @@ struct common_speculative_impl {
    const common_speculative_type type;

    uint32_t n_seq;
+    int32_t n_max; // maximum draft length after implementation-specific limits

    size_t n_call_begin  = 0; // number of times this implementation was called for refresh.
    size_t n_call_draft  = 0; // number of times this implementation was called for generation.
@@ -157,7 +159,7 @@ struct common_speculative_impl {
    int64_t t_draft_us  = 0; // total time spent in generating drafts in this implementation in microseconds.
    int64_t t_accept_us = 0; // total time spent in accumulation of this implementation in microseconds.

-    common_speculative_impl(common_speculative_type type, uint32_t n_seq) : type(type), n_seq(n_seq) {}
+    common_speculative_impl(common_speculative_type type, uint32_t n_seq, int32_t n_max) : type(type), n_seq(n_seq), n_max(n_max) {}

    virtual ~common_speculative_impl() = default;

@@ -182,7 +184,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
    std::vector<common_sampler_ptr> smpls;

    common_speculative_impl_draft_simple(const common_params_speculative & params, uint32_t n_seq)
-        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq)
+        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq, params.draft.n_max)
        , params(params.draft)
    {
        auto * ctx_dft = this->params.ctx_dft;
@@ -452,7 +454,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
    std::vector<float> g_embd_buf;

    common_speculative_impl_draft_eagle3(const common_params_speculative & params, uint32_t n_seq)
-        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq)
+        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq, params.draft.n_max)
        , params(params.draft)
    {
        SPC_TRC("%s", "adding speculative implementation 'draft-eagle3'\n");
@@ -923,12 +925,19 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
    int32_t     block_size    = 0;
    llama_token mask_token_id = 0;

+    bool    is_dflash2     = false;
+    bool    is_mrope       = false;
+    int32_t selector_top_k = 0;
+
    // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout
    const bool is_dspark;

    // dspark speculators
    bool sample_from_anchor = true;

+    // block-internal attention
+    bool causal_attn = false;
+
    const int32_t * target_layer_ids   = nullptr; // model_dft's extract layer indices
    uint32_t        target_layer_ids_n = 0;

@@ -937,7 +946,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

    common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq,
            common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)
-        : common_speculative_impl(type, n_seq)
+        : common_speculative_impl(type, n_seq, params.draft.n_max)
        , params(params.draft)
        , is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)
    {
@@ -966,9 +975,25 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
            if (llama_model_meta_val_str(model_dft, "dflash.sample_from_anchor", buf, sizeof(buf)) >= 0) {
                sample_from_anchor = std::strcmp(buf, "true") == 0;
            }
+            if (llama_model_meta_val_str(model_dft, "dflash.attention.causal", buf, sizeof(buf)) >= 0) {
+                causal_attn = std::strcmp(buf, "true") == 0;
+            }
        }
+
+        selector_top_k = llama_model_dflash_selector_top_k(model_dft);
+        is_dflash2     = selector_top_k > 0;
        mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft));

+        if (is_dspark && this->params.p_min > 0.0f) {
+            char buf[16] = {};
+            const bool has_conf =
+                llama_model_meta_val_str(model_dft, "dflash.has_confidence_head", buf, sizeof(buf)) < 0 ||
+                std::strcmp(buf, "true") == 0;
+            if (!has_conf) {
+                throw std::runtime_error("DSpark draft has no confidence head: please set --spec-draft-p-min 0");
+            }
+        }
+
        LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str());
        LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min);
        LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__,
@@ -983,10 +1008,18 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
            this->params.n_max = std::min(this->params.n_max, n_draft_max);
            this->params.n_min = std::min(this->params.n_min, n_draft_max);
        }
+        this->n_max = this->params.n_max;

        batch        = llama_batch_init(llama_n_batch(ctx_dft), 0,          n_seq);
        batch_inject = llama_batch_init(llama_n_batch(ctx_dft), n_embd_dec, n_seq);

+        // embd batches on an M-RoPE draft need 4 position rows per token
+        is_mrope = llama_model_rope_type(model_dft) == LLAMA_ROPE_TYPE_MROPE;
+        if (is_mrope) {
+            free(batch_inject.pos);
+            batch_inject.pos = (llama_pos *) malloc(sizeof(llama_pos) * 4 * llama_n_batch(ctx_dft));
+        }
+
        smpls.resize(n_seq);
        for (auto & s : smpls) {
            common_params_sampling sparams;
@@ -998,7 +1031,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

        // offload draft sampling to the backend
        backend_chains.assign(n_seq, nullptr);
-        if (this->params.backend_sampling) {
+        if (this->params.backend_sampling && !is_dflash2) {
            for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
                llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params());
                llama_sampler_chain_add(chain, llama_sampler_init_top_k(10));
@@ -1017,8 +1050,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
            llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
        }

-        llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
-        llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention
+        // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits.
+        llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2);
+        llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise
    }

    ~common_speculative_impl_draft_dflash() override {
@@ -1118,11 +1152,24 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
                }

                // fuse extracted features through DFlash encoder
+                // M-RoPE drafts read 4 position rows per token from embd batches, so pass them explicitly
+                std::vector<llama_pos> enc_pos;
+                if (is_mrope) {
+                    enc_pos.resize((size_t) 4 * n_chunk);
+                    for (int32_t i = 0; i < n_chunk; ++i) {
+                        const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i];
+                        enc_pos[0 * n_chunk + i] = p;
+                        enc_pos[1 * n_chunk + i] = p;
+                        enc_pos[2 * n_chunk + i] = p;
+                        enc_pos[3 * n_chunk + i] = 0;
+                    }
+                }
+
                llama_batch enc_batch = {
                    /*.n_tokens =*/ n_chunk,
                    /*.token    =*/ nullptr,
                    /*.embd     =*/ features_buf.data(),
-                    /*.pos      =*/ nullptr,
+                    /*.pos      =*/ is_mrope ? enc_pos.data() : nullptr,
                    /*.n_seq_id =*/ nullptr,
                    /*.seq_id   =*/ nullptr,
                    /*.logits   =*/ nullptr,
@@ -1143,7 +1190,13 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
                std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float));

                for (int32_t i = 0; i < n_chunk; ++i) {
-                    batch_inject.pos[i]       = batch_in.pos[i_batch_beg[seq_id] + offset + i];
+                    const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i];
+                    batch_inject.pos[i] = p;
+                    if (is_mrope) {
+                        batch_inject.pos[1 * n_chunk + i] = p;
+                        batch_inject.pos[2 * n_chunk + i] = p;
+                        batch_inject.pos[3 * n_chunk + i] = 0;
+                    }
                    batch_inject.n_seq_id[i]  = 1;
                    batch_inject.seq_id[i][0] = seq_id;
                    batch_inject.logits[i]    = false;
@@ -1186,7 +1239,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
            i_block_beg[seq_id] = batch.n_tokens;
            n_block    [seq_id] = n_block_tokens;
            for (int32_t i = 0; i < n_block_tokens; ++i) {
-                common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true);
+                common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, !is_dflash2);
            }
        }

@@ -1214,6 +1267,36 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

            auto & result = *dp.result;

+            if (is_dflash2) {
+                const float * lattice = llama_get_embeddings_nextn(ctx_dft);
+                GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");
+
+                int32_t predecessor = 0;
+                for (int32_t i = 1; i < n_block_tokens; ++i) {
+                    const float * row = lattice + (size_t) (beg + i) * n_embd_dec;
+                    const float * scores = row + selector_top_k + (size_t) predecessor * selector_top_k;
+
+                    predecessor = (int32_t) std::distance(scores,
+                            std::max_element(scores, scores + selector_top_k));
+                    if (params.p_min > 0.0f) {
+                        // softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max))
+                        float sum = 0.0f;
+                        for (int32_t k = 0; k < selector_top_k; ++k) {
+                            sum += std::exp(scores[k] - scores[predecessor]);
+                        }
+                        if (1.0f / sum < params.p_min) {
+                            break;
+                        }
+                    }
+                    result.push_back((llama_token) row[predecessor]);
+                }
+
+                if (result.size() < (size_t) params.n_min) {
+                    result.clear();
+                }
+                continue;
+            }
+
            if (is_dspark) {
                // DSpark: read from the first draft slot, truncate below the confidence threshold
                const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr;
@@ -1315,7 +1398,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
    std::vector<std::vector<float>> chain_h;

    common_speculative_impl_draft_mtp(const common_params_speculative & params, uint32_t n_seq)
-        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq)
+        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq, params.draft.n_max)
        , params(params.draft)
    {
        auto * ctx_tgt = this->params.ctx_tgt;
@@ -1382,6 +1465,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
                c.reserve((size_t) (this->params.n_max + 1) * n_embd);
            }
        }
+        this->n_max = this->params.n_max;

        pending_h.assign(n_seq, std::vector<float>(n_embd, 0.0f));

@@ -1726,7 +1810,7 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl {
    common_speculative_impl_ngram_simple(
            const common_params_speculative & params, uint32_t n_seq,
            common_ngram_simple_config config)
-        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, n_seq)
+        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, n_seq, params.ngram_simple.size_m)
        , params(params.ngram_simple)
        , config(config)
    {
@@ -1770,7 +1854,7 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl {
            const common_ngram_map & config,
            uint32_t n_seq)
        : common_speculative_impl(config.key_only ? COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K
-            : COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, n_seq)
+            : COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, n_seq, config.size_value)
    {
        for (uint32_t i = 0; i < n_seq; i++) {
            this->config.push_back(config);
@@ -1841,7 +1925,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl {
    common_speculative_impl_ngram_mod(
            const common_params_speculative & params,
            uint32_t n_seq)
-        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, n_seq)
+        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, n_seq, params.ngram_mod.n_max)
        , params(params.ngram_mod)
        , mod(params.ngram_mod.n_match, 4*1024*1024)
        , verbose(std::getenv("LLAMA_TRACE") != nullptr) {
@@ -2017,7 +2101,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
            const std::string & path_dynamic,
            bool save_dynamic,
            bool save_static)
-        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq)
+        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq, n_draft)
        , params(params.ngram_cache)
        , n_draft(n_draft)
        , save_dynamic(save_dynamic)
@@ -2138,6 +2222,8 @@ struct common_speculative {

    // which implementaion was used for a given seq_id
    std::vector<common_speculative_impl *> impl_last;
+
+    std::vector<double> synth_probs;
 };

 static common_ngram_map get_common_ngram_map(
@@ -2316,6 +2402,101 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) {
    return n_max;
 }

+int32_t common_speculative_n_max(const common_speculative * spec) {
+    int32_t n_max = 0;
+
+    if (spec == nullptr) {
+        return n_max;
+    }
+
+    for (const auto & impl : spec->impls) {
+        n_max = std::max(n_max, std::max(0, impl->n_max));
+    }
+
+    return n_max;
+}
+
+std::vector<double> common_speculative_synth_rates_resolve(const common_params_speculative * spec, int32_t n_max) {
+    const bool has_length = spec->synth_len != -1.0;
+    const bool has_rates  = !spec->synth_rates.empty();
+
+    if (!has_length && !has_rates) {
+        return {};
+    }
+    if (has_length && has_rates) {
+        throw std::invalid_argument("synthetic acceptance length and rates are mutually exclusive");
+    }
+
+    if (n_max <= 0) {
+        throw std::invalid_argument("synthetic acceptance requires at least one speculative token");
+    }
+
+    if (has_rates) {
+        const auto & rates = spec->synth_rates;
+        if (rates.size() != (size_t) n_max) {
+            throw std::invalid_argument(string_format(
+                    "synthetic acceptance rates must contain %d values, got %zu", n_max, rates.size()));
+        }
+
+        for (size_t i = 0; i < rates.size(); ++i) {
+            if (!std::isfinite(rates[i]) || rates[i] < 0.0 || rates[i] > 1.0) {
+                throw std::invalid_argument("synthetic acceptance rates must be finite and within [0, 1]");
+            }
+            if (i > 0 && rates[i] > rates[i - 1]) {
+                throw std::invalid_argument("synthetic acceptance rates must be monotonically non-increasing");
+            }
+        }
+
+        return rates;
+    }
+
+    const double length = spec->synth_len;
+    const double length_max = (double) n_max + 1.0;
+    if (!std::isfinite(length) || length < 1.0 || length > length_max) {
+        throw std::invalid_argument(string_format(
+                "synthetic acceptance length must be finite and within [1, %.0f]", length_max));
+    }
+
+    double p = 0.0;
+    if (length == length_max) {
+        p = 1.0;
+    } else if (length > 1.0) {
+        double p_min = 0.0;
+        double p_max = 1.0;
+        for (int i = 0; i < 32; ++i) {
+            const double p_mid = 0.5 * (p_min + p_max);
+            double sum = 0.0;
+            double term = p_mid;
+            for (int32_t j = 0; j < n_max; ++j) {
+                sum += term;
+                term *= p_mid;
+            }
+
+            if (sum < length - 1.0) {
+                p_min = p_mid;
+            } else {
+                p_max = p_mid;
+            }
+        }
+        p = 0.5 * (p_min + p_max);
+    }
+
+    std::vector<double> rates;
+    rates.reserve(n_max);
+    double rate = p;
+    for (int32_t i = 0; i < n_max; ++i) {
+        rates.push_back(rate);
+        rate *= p;
+    }
+
+    return rates;
+}
+
+const std::vector<double> & common_speculative_get_synth_probs(const common_speculative * spec) {
+    GGML_ASSERT(spec);
+    return spec->synth_probs;
+}
+
 common_params common_base_params_to_speculative(const common_params & params) {
    const bool has_draft = params.speculative.has_dft();

@@ -2568,13 +2749,39 @@ common_speculative * common_speculative_init(common_params_speculative & params,
        return nullptr;
    }

-    auto * result = new common_speculative {
-        /* .dparams   = */ common_speculative_draft_params_vec(n_seq),
-        /* .impls     = */ std::move(impls),
-        /* .impl_last = */ std::vector<common_speculative_impl *>(n_seq, nullptr)
-    };
+    common_speculative_ptr result(new common_speculative {
+        /* .dparams     = */ common_speculative_draft_params_vec(n_seq),
+        /* .impls       = */ std::move(impls),
+        /* .impl_last   = */ std::vector<common_speculative_impl *>(n_seq, nullptr),
+        /* .synth_probs = */ {},
+    });

-    return result;
+    const int32_t n_max_configured = common_speculative_n_max(&params);
+    const int32_t n_max_effective  = common_speculative_n_max(result.get());
+    const auto rates = common_speculative_synth_rates_resolve(&params, n_max_effective);
+
+    std::vector<std::string> rates_str;
+    rates_str.reserve(rates.size());
+    result->synth_probs.reserve(rates.size());
+    double rate_prev = 1.0;
+    double acceptance_length = 1.0;
+    for (const double rate : rates) {
+        result->synth_probs.push_back(rate_prev > 0.0 ? rate / rate_prev : 0.0);
+        rates_str.push_back(string_format("%.6g", rate));
+        rate_prev = rate;
+        acceptance_length += rate;
+    }
+    if (!result->synth_probs.empty()) {
+        SPC_WRN("%s", "synthetic speculative acceptance is enabled for benchmarking; generated output is not valid\n");
+        if (n_max_effective != n_max_configured) {
+            SPC_WRN("synthetic acceptance draft limit was reduced from %d to %d by the initialized speculative implementations\n",
+                    n_max_configured, n_max_effective);
+        }
+        SPC_INF("synthetic acceptance: n_max = %zu, mean length = %.6f, rates = [%s]\n",
+                rates.size(), acceptance_length, string_join(rates_str, ", ").c_str());
+    }
+
+    return result.release();
 }

 void common_speculative_free(common_speculative * spec) {
@@ -26,6 +26,15 @@ std::string common_speculative_type_to_str(enum common_speculative_type type);
 // return the max number of draft tokens based on the speculative parameters
 int32_t common_speculative_n_max(const common_params_speculative * spec);

+// return the max number of draft tokens from the initialized implementations
+int32_t common_speculative_n_max(const common_speculative * spec);
+
+// validate and resolve the unconditional synthetic acceptance rates
+std::vector<double> common_speculative_synth_rates_resolve(const common_params_speculative * spec, int32_t n_max);
+
+// return the conditional synthetic acceptance probabilities
+const std::vector<double> & common_speculative_get_synth_probs(const common_speculative * spec);
+
 common_params common_base_params_to_speculative(const common_params & params);

 struct common_speculative_output_limits {
@@ -54,6 +54,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "DeepseekV3ForCausalLM": "deepseek",
    "DeepseekV32ForCausalLM": "deepseek",
    "DFlashDraftModel": "qwen",
+    "DFlash2DraftModel": "qwen",
    "Qwen3DSparkModel": "qwen",
    "DSparkDraftModel": "qwen",
    "DSparkSpeculator": "qwen",
@@ -235,6 +236,8 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "Qwen3_5ForConditionalGeneration": "qwen",
    "Qwen3_5MoeForCausalLM": "qwen",
    "Qwen3_5MoeForConditionalGeneration": "qwen",
+    "Qwen4ExpForCausalLM": "qwen4exp",
+    "Qwen4ExpForConditionalGeneration": "qwen4exp",
    "RND1": "qwen",
    "RWForCausalLM": "falcon",
    "RWKV6Qwen2ForCausalLM": "rwkv",
@@ -332,6 +335,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "Qwen3VLMoeForConditionalGeneration": "qwen3vl",
    "Qwen3_5ForConditionalGeneration": "qwen3vl",
    "Qwen3_5MoeForConditionalGeneration": "qwen3vl",
+    "Qwen4ExpForConditionalGeneration": "qwen4exp",
    "RADIOModel": "nemotron",
    "Sarashina2VisionForCausalLM": "sarashina2",
    "SmolVLMForConditionalGeneration": "smolvlm",
@@ -1006,12 +1006,16 @@ class ModelBase:
                    else:
                        raise ValueError(f"Unknown file type: {self.ftype.name}")

+                # a chunked tensor quantizes as one chunk at a time, while it is written
+                quantize = data.quantize if isinstance(data, gguf.LazyChunkedTensor) else (
+                    lambda qtype, d=data: gguf.quants.quantize(d, qtype))
+
                try:
-                    data = gguf.quants.quantize(data, data_qtype)
+                    data = quantize(data_qtype)
                except gguf.QuantError as e:
                    logger.warning("%s, %s", e, "falling back to F16")
                    data_qtype = gguf.GGMLQuantizationType.F16
-                    data = gguf.quants.quantize(data, data_qtype)
+                    data = quantize(data_qtype)

                shape = gguf.quant_shape_from_byte_shape(data.shape, data_qtype) if data.dtype == np.uint8 else data.shape

@@ -302,6 +302,10 @@ class NemotronHModel(GraniteHybridModel):
            )
            if not keep:
                return None
+        # PEFT names adapter tensors using model.layers.*, while Nemotron-H checkpoints
+        # and the GGUF tensor map use backbone.layers.*
+        if name.startswith("model.layers.") and ".mixer." in name:
+            name = name.replace("model.layers.", "backbone.layers.", 1)
        return super().filter_tensors((name, gen))

    def prepare_metadata(self, vocab_only: bool):
@@ -639,7 +639,7 @@ class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
    model_arch = gguf.MODEL_ARCH.QWEN35MOE


-@ModelBase.register("DFlashDraftModel")
+@ModelBase.register("DFlashDraftModel", "DFlash2DraftModel")
@ModelBase.example("z-lab/Qwen3.5-9B-DFlash")
 class DFlashModel(Qwen3Model):
    model_arch = gguf.MODEL_ARCH.DFLASH
@@ -678,34 +678,98 @@ class DFlashModel(Qwen3Model):
    def set_gguf_parameters(self):
        super().set_gguf_parameters()

-        block_size = self.hparams.get("block_size", 16)
-        self.gguf_writer.add_block_size(block_size)
        dflash_config = self.hparams.get("dflash_config", {})
+        block_size = dflash_config.get("block_size", self.hparams.get("block_size", 16))
+        self.gguf_writer.add_block_size(block_size)
+
+        if "conv_kernel_size" in dflash_config:
+            self.gguf_writer.add_conv_kernel_size(int(dflash_config["conv_kernel_size"]))
+            self.gguf_writer.add_conv_group_size(int(dflash_config["conv_group_size"]))
+            self.gguf_writer.add_selector_rank(int(dflash_config["selector_rank"]))
+            self.gguf_writer.add_selector_top_k(int(dflash_config["selector_top_k"]))
+
+        output_multiplier = dflash_config.get(
+            "output_multiplier", self.hparams.get("output_multiplier")
+        )
+        if output_multiplier is not None:
+            self.gguf_writer.add_logit_scale(float(output_multiplier))
+        softcap = dflash_config.get(
+            "final_logit_softcapping", self.hparams.get("final_logit_softcapping")
+        )
+        if softcap is not None and float(softcap) > 0:
+            self.gguf_writer.add_final_logit_softcapping(float(softcap))
+        embedding_scale = dflash_config.get(
+            "input_embedding_scale", self.hparams.get("input_embedding_scale")
+        )
+        if embedding_scale is not None:
+            self.gguf_writer.add_embedding_scale(float(embedding_scale))

        target_layer_ids = dflash_config.get("target_layer_ids", [])
        if target_layer_ids:
            extract_layer_ids = [i + 1 for i in target_layer_ids]
            self.gguf_writer.add_target_layers(extract_layer_ids)

-        use_sliding_window = self.hparams.get("use_sliding_window", False)
-        sliding_window = self.hparams.get("sliding_window")
+        use_sliding_window = self.hparams.get("use_sliding_window", False) or dflash_config.get("use_swa", False)
+        sliding_window = dflash_config.get("swa_window_size") or self.hparams.get("sliding_window")
        layer_types = self.hparams.get("layer_types")
        if use_sliding_window and sliding_window and layer_types:
            is_swa = [lt == "sliding_attention" for lt in layer_types]
            self.gguf_writer.add_sliding_window(sliding_window)
            self.gguf_writer.add_sliding_window_pattern(is_swa)

+        causal = self.hparams.get("is_causal")
+        if causal is None:
+            causal = dflash_config.get("causal")
+        if causal is not None:
+            self.gguf_writer.add_causal_attention(bool(causal))
+
+        # M-RoPE target: the draft ropes on the temporal dim only, so write
+        # degenerate sections [n_rot/2, 0, 0, 0]
+        if self._target_uses_mrope():
+            head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
+            self.gguf_writer.add_rope_dimension_sections([head_dim // 2, 0, 0, 0])
+
+    def _target_uses_mrope(self) -> bool:
+        if self.target_model_dir is None:
+            return False
+        with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f:
+            cfg = json.load(f)
+        cfg = cfg.get("text_config", cfg)
+        rope = cfg.get("rope_parameters") or cfg.get("rope_scaling") or {}
+        return "mrope_section" in rope
+
    @classmethod
    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
        name, gen = item
        if not name.startswith("model."):
            name = "model." + name
+        if "sink" in name and not name.endswith(".weight"):
+            name += ".weight"
        return super().filter_tensors((name, gen))

+    _ROPE_PERMUTE_SUFFIXES = (
+        "self_attn.q_proj.weight",
+        "self_attn.k_proj.weight",
+        "self_attn.q_norm.weight",
+        "self_attn.k_norm.weight",
+    )
+
    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
        if name == "model.embed_tokens.weight" and not self.hparams.get("has_embed_tokens", True):
            return

+        # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd
+        if not self.hparams.get("rope_is_neox_style", True) and name.endswith(self._ROPE_PERMUTE_SUFFIXES):
+            head_dim = self.hparams["head_dim"]
+            shape = data_torch.shape
+            data_torch = data_torch.reshape(-1, head_dim // 2, 2, *shape[1:]).transpose(1, 2).reshape(shape)
+
+        if name in (
+            "model.candidate_selector.predecessor_codebook",
+            "model.candidate_selector.successor_codebook",
+        ):
+            name += ".weight"
+
        yield from super().modify_tensors(data_torch, name, bid)


@@ -759,6 +823,10 @@ class DSparkModel(DFlashModel):
        super().set_gguf_parameters()
        self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor)

+        # confidence head is optional: vanilla-markov exports ship without it
+        has_conf = any("confidence_head.proj" in name for name in self.model_tensors)
+        self.gguf_writer.add_has_confidence_head(has_conf)
+
    @classmethod
    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
        if item[0] == "t2d":  # not used at runtime
@@ -777,7 +845,7 @@ class DSparkModel(DFlashModel):
            self._d2t = data_torch
            return

-        if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith(("embed_tokens.weight", "lm_head.weight")):
+        if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith("lm_head.weight"):
            return

        # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd
@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+from typing import Iterable, cast
+
+import torch
+from torch import Tensor
+
+import gguf
+import numpy as np
+
+from .base import ModelBase
+from .qwen import _LinearAttentionVReorderBase, _Qwen35MRopeMixin
+from .qwen3vl import Qwen3VLVisionModel
+
+
+@ModelBase.register("Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM")
+@ModelBase.example("Qwen/Qwen3.8-Flash-Next")
+class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
+    """Qwen3.8-Flash-Next.
+
+    Shares the Qwen3.5 gated delta net and interleaved mrope, and adds three things:
+    hyper-connections in place of every layer norm, QSA sparse attention on the full
+    attention layers, and PLE n-gram hash embeddings on a single layer.
+    """
+
+    model_arch = gguf.MODEL_ARCH.QWEN4EXP
+
+    # the MTP block is a separate draft head; vLLM drops it too
+    supports_mtp_export = False
+    no_mtp = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # only the shard names, so the table itself is never held
+        self._ple_shards: dict[int, str] = {}
+        self._ple_row_dim: int | None = None
+
+    def _read_hash_constants(self, suffix: str) -> list[int]:
+        """Read an int64 PLE constant straight from the checkpoint.
+
+        prepare_tensors() casts every non-float dtype to float32 before
+        modify_tensors() sees it (base.py), which would silently round these
+        45-bit multipliers. Reading the lazy tensor here bypasses that.
+        """
+        for name, gen in self.model_tensors.items():
+            if name.endswith(suffix):
+                t = gen()
+                if t.dtype != torch.int64:
+                    t = t.to(torch.int64)
+                return [int(x) for x in t.tolist()]
+        raise ValueError(f"PLE constant {suffix!r} missing from the checkpoint")
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        hp = self.hparams
+
+        self.gguf_writer.add_hyper_connection_count(hp["hc_count"])
+        self.gguf_writer.add_hyper_connection_low_rank(hp["hc_lowrank"])
+
+        n_layer = hp["num_hidden_layers"]
+        self.gguf_writer.add_indexer_head_count(hp["indexer_n_heads"])
+        self.gguf_writer.add_indexer_key_length(hp["indexer_head_dim"])
+        self.gguf_writer.add_indexer_top_k(hp["indexer_budget"])
+        ratio = hp["indexer_compress_ratio"]
+        layer_types = hp["layer_types"]
+        self.gguf_writer.add_attention_compress_ratios(
+            [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)]
+        )
+
+        # ple_layer_ids is 1-based in the HF config; empty means no n-gram table,
+        # so emit no PLE keys rather than optional ones
+        ple_layers = [i - 1 for i in hp["ple_layer_ids"]]
+        if not ple_layers:
+            return
+        self.gguf_writer.add_ple_layers(ple_layers)
+        self.gguf_writer.add_ple_ngram_size(hp["ngram_size"])
+        self.gguf_writer.add_ple_heads_per_ngram(hp["heads_per_ngram"])
+        self.gguf_writer.add_ple_conv_kernel(hp["ple_conv_kernel_size"])
+        self.gguf_writer.add_ple_eos_token_id(self._eos_token_id())
+        # an image is decoded as an embeddings-only batch, so the graph has no placeholder
+        # ids to hash; carry the id and let it stand in for those positions
+        _img = self._image_token_id()
+        if _img is not None:
+            self.gguf_writer.add_ple_image_token_id(int(_img))
+        if self._ple_row_dim is not None:
+            self.gguf_writer.add_embedding_length_per_layer_input(self._ple_row_dim)
+
+        self.gguf_writer.add_ple_layer_multipliers(
+            self._read_hash_constants("ple_embedding.layer_multipliers"))
+        self.gguf_writer.add_ple_head_offsets(
+            self._read_hash_constants("ple_embedding.ngram_heads_offsets"))
+        self.gguf_writer.add_ple_head_vocab_sizes(
+            self._read_hash_constants("ple_embedding.ngram_heads_vocab_sizes"))
+
+    def _image_token_id(self) -> int | None:
+        img = self.hparams.get("image_token_id")
+        return None if img is None else int(img)
+
+    def _eos_token_id(self) -> int:
+        eos = self.hparams.get("eos_token_id")
+        if isinstance(eos, list):
+            # the PLE hash resets n-grams on the primary EOS
+            return int(eos[-1])
+        if eos is None:
+            raise ValueError("eos_token_id is required: the PLE hash resets its n-grams on it")
+        return int(eos)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # int64 hash constants must stay exact; 1-D tensors force F32, so use KV
+        if name.endswith("ple_embedding.layer_multipliers"):
+            self._ple_multipliers = [int(x) for x in data_torch.tolist()]
+            return []
+        if name.endswith("ple_embedding.ngram_heads_offsets"):
+            self._ple_head_offsets = [int(x) for x in data_torch.tolist()]
+            return []
+        if name.endswith("ple_embedding.ngram_heads_vocab_sizes"):
+            self._ple_head_vocab_sizes = [int(x) for x in data_torch.tolist()]
+            return []
+
+        if ".ngram_embedding.shard_" in name:
+            return self._place_ple_shard(data_torch, name)
+
+        # one projection feeds indexer q and k; split it, as minimax-m3 does
+        if ".indexer.index_qk_proj.weight" in name:
+            n_q = self.hparams["indexer_n_heads"] * self.hparams["indexer_head_dim"]
+            q = data_torch[:n_q]
+            k = data_torch[n_q:]
+            return [
+                (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid, ".weight"), q),
+                (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid, ".weight"), k),
+            ]
+
+        # Gemma zero-centred gammas the inherited norm.weight rule misses
+        if name.endswith((".ple.norm_key.weight", ".ple.norm_query.weight", ".ple.norm_conv.weight",
+                          ".indexer.q_layernorm.weight", ".indexer.k_layernorm.weight")):
+            return [(self.map_tensor_name(name), data_torch + 1)]
+
+        if name.endswith(".ple.conv1d.weight"):
+            return [(self.map_tensor_name(name), data_torch.squeeze())]
+
+        return super().modify_tensors(data_torch, name, bid)
+
+    # the shards concatenate into a tensor of well over 100 GB
+    # use LazyChunkedTensor here, a single shard resident at a time
+    def _place_ple_shard(self, data_torch: Tensor, name: str) -> Iterable[tuple[str, Tensor]]:
+
+        idx = int(name.rpartition(".shard_")[2].partition(".")[0])
+        n_parts = self.hparams["split_ngram_parts"]
+
+        self._ple_shards[idx] = name
+        self._ple_row_dim = int(data_torch.shape[-1])
+
+        if len(self._ple_shards) < n_parts:
+            return []
+
+        # the checkpoint may yield the shards in any order, the row order is by index
+        shards = [self._ple_shards[i] for i in sorted(self._ple_shards)]
+        rows = 0
+        for shard in shards:
+            shape = self.model_tensors[shard]().shape
+            if int(shape[-1]) != self._ple_row_dim:
+                raise ValueError(
+                    f"PLE shard {shard} has row dim {int(shape[-1])}, expected {self._ple_row_dim}")
+            rows += int(shape[0])
+
+        table = gguf.LazyChunkedTensor(
+            [self._load_ple_shard(shard) for shard in shards],
+            shape=(rows, self._ple_row_dim),
+            dtype=np.float32,
+        )
+        gguf_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.PER_LAYER_TOKEN_EMBD]
+        return [(gguf_name + ".weight", cast(Tensor, table))]
+
+    def _load_ple_shard(self, name: str):
+        def load() -> np.ndarray:
+            from .base import LazyTorchTensor
+
+            # a fresh lazy tensor every call, or to_eager() memoizes every shard
+            eager = LazyTorchTensor.to_eager(self.model_tensors[name]())
+            return eager.to(torch.float32).contiguous().numpy()
+        return load
+
+    def prepare_tensors(self):
+        super().prepare_tensors()
+        n_parts = self.hparams.get("split_ngram_parts", 0)
+        if self._ple_shards and len(self._ple_shards) != n_parts:
+            raise ValueError(
+                f"got {len(self._ple_shards)} PLE embedding shards, expected {n_parts}"
+            )
+
+
+@ModelBase.register("Qwen4ExpForConditionalGeneration")
+@ModelBase.example("Qwen/Qwen3.8-Flash-Next")
+class Qwen4ExpVisionModel(Qwen3VLVisionModel):
+    """The vision tower is an unmodified Qwen3-VL ViT."""
@@ -22,8 +22,8 @@ The OpenVINO backend is implemented in `ggml/src/ggml-openvino` and provides a t
  - [0. Prerequisites](#0-prerequisites)
  - [1. Install OpenVINO Runtime](#1-install-openvino-runtime)
  - [2. Build llama.cpp with OpenVINO Backend](#2-build-llamacpp-with-openvino-backend)
-    - [Automated Ubuntu Build Script](#automated-ubuntu-build-script)
-    - [Automated Windows Build Script](#automated-windows-build-script)
+    - [Ubuntu Build Script](#ubuntu-build-script)
+    - [Windows Build Script](#windows-build-script)
  - [3. Download Sample Model](#3-download-sample-model)
  - [4. Run Inference with OpenVINO Backend](#4-run-inference-with-openvino-backend)
  - [5. Docker Build](#5-docker-build)
@@ -96,7 +96,7 @@ Although, the validated models below were tested with `llama-cli` using the `Q4_
  - **SL** = Stateless (`GGML_OPENVINO_STATEFUL_EXECUTION=0`)
  - **SF** = Stateful (`GGML_OPENVINO_STATEFUL_EXECUTION=1`)
  - Note: The NPU operates in stateless mode only.
- **Validation system:** Intel® Core™ Ultra 5 238V (Lunar Lake) | 32 GB RAM | Ubuntu 24.04 | Intel OpenCL GPU Driver 26.18.38308.1 | Intel NPU Driver 1.33.0.
+- **Validation system:** Intel® Core™ Ultra 5 238V (Lunar Lake) | 32 GB RAM | Ubuntu 24.04 | Intel OpenCL GPU Driver 26.31.39395.13-0 | Intel NPU Driver 1.35.0.
 - See [Known Limitations](#known-limitations) for context on observed failures.

 | Model | CPU (SL / SF) | GPU (SL / SF) | NPU (SL) |
@@ -105,27 +105,32 @@ Although, the validated models below were tested with `llama-cli` using the `Q4_
 | [bartowski/Llama-3.2-3B-Instruct-Q4_K_M](https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 | [bartowski/Meta-Llama-3.1-8B-Instruct-Q4_K_M](https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 |  |  |  |  |
-| [Qwen/qwen2.5-1.5b-instruct-q4_k_m](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ |
-| [Qwen/qwen2.5-coder-7b-instruct-q4_k_m](https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ |
-| [bartowski/Qwen_Qwen3-0.6B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3-0.6B-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ |
-| [bartowski/Qwen_Qwen3-1.7B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3-1.7B-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ |
-| [Qwen/Qwen3-4B-Q4_K_M](https://huggingface.co/Qwen/Qwen3-4B-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ |
-| [lm-kit/Qwen3-8B-Q4_K_M](https://huggingface.co/lm-kit/qwen-3-8b-instruct-gguf) | ✓ / ✓ | ✓ / ✗ | ✓ |
+| [Qwen/qwen2.5-1.5b-instruct-q4_k_m](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [Qwen/qwen2.5-coder-7b-instruct-q4_k_m](https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [bartowski/Qwen_Qwen3-0.6B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3-0.6B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [bartowski/Qwen_Qwen3-1.7B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3-1.7B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [Qwen/Qwen3-4B-Q4_K_M](https://huggingface.co/Qwen/Qwen3-4B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [lm-kit/Qwen3-8B-Q4_K_M](https://huggingface.co/lm-kit/qwen-3-8b-instruct-gguf) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [bartowski/Qwen_Qwen3.5-0.8B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3.5-0.8B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ |
+| [bartowski/Qwen_Qwen3.5-2B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3.5-2B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ |
+| [bartowski/Qwen_Qwen3.5-4B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3.5-4B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ |
+| [lmstudio-community/Qwen3.5-9B-Q4_K_M](https://huggingface.co/lmstudio-community/Qwen3.5-9B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ |
 |  |  |  |  |
-| [unsloth/gemma-3-4b-it-Q4_K_M](https://huggingface.co/unsloth/gemma-3-4b-it-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ |
-| [bartowski/google_gemma-4-E2B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E2B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✓ |
+| [unsloth/gemma-3-4b-it-Q4_K_M](https://huggingface.co/unsloth/gemma-3-4b-it-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [bartowski/google_gemma-4-E2B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E2B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ |
 | [bartowski/google_gemma-4-E4B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E4B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✓ |
-| [bartowski/gemma-4-12B-it-Q4_K_M](https://huggingface.co/bartowski/gemma-4-12B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ |
+| [bartowski/gemma-4-12B-it-Q4_K_M](https://huggingface.co/bartowski/gemma-4-12B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✓ |
 |  |  |  |  |
-| [bartowski/Phi-3-mini-4k-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3-mini-4k-instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ |
-| [bartowski/Phi-3.5-mini-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3.5-mini-instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ |
+| [bartowski/Phi-3-mini-4k-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3-mini-4k-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [bartowski/Phi-3.5-mini-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3.5-mini-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [bartowski/microsoft_Phi-4-mini-instruct-Q4_K_M](https://huggingface.co/bartowski/microsoft_Phi-4-mini-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 |  |  |  |  |
 | [bartowski/Mistral-7B-Instruct-v0.3-Q4_K_M](https://huggingface.co/bartowski/Mistral-7B-Instruct-v0.3-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 | [QuantFactory/Ministral-3b-instruct.Q4_K_M](https://huggingface.co/QuantFactory/Ministral-3b-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 | [bartowski/Ministral-8B-Instruct-2410-Q4_K_M](https://huggingface.co/bartowski/Ministral-8B-Instruct-2410-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 |  |  |  |  |
 | [bartowski/DeepSeek-R1-Distill-Llama-8B-Q4_K_M](https://huggingface.co/bartowski/DeepSeek-R1-Distill-Llama-8B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
-| [bartowski/DeepSeek-R1-Distill-Qwen-7B-Q4_K_M](https://huggingface.co/bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ |
+| [bartowski/DeepSeek-R1-Distill-Qwen-7B-Q4_K_M](https://huggingface.co/bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 |  |  |  |  |
 | [ibm-granite/granite-4.0-350m-Q4_K_M](https://huggingface.co/ibm-granite/granite-4.0-350m-GGUF) | ✓ / ✓ | ✗ / ✗ | ✓ |
 | [ibm-granite/granite-4.0-micro-Q4_K_M](https://huggingface.co/ibm-granite/granite-4.0-micro-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
@@ -133,10 +138,10 @@ Although, the validated models below were tested with `llama-cli` using the `Q4_
 | [ibm-research/granite-3.2-8b-instruct-Q4_K_M](https://huggingface.co/ibm-research/granite-3.2-8b-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 |  |  |  |  |
 | [HuggingFaceTB/smollm2-1.7b-instruct-q4_k_m](https://huggingface.co/HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
-| [openbmb/MiniCPM-V-2_6-Q4_K_M](https://huggingface.co/openbmb/MiniCPM-V-2_6-gguf) | ✓ / ✓ | ✓ / ✗ | ✓ |
-| [bartowski/tencent_Hunyuan-7B-Instruct-Q4_K_M](https://huggingface.co/bartowski/tencent_Hunyuan-7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ |
-| [LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-Q4_K_M](https://huggingface.co/LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ |
-| [bartowski/prism-ml_Bonsai-8B-unpacked-Q4_K_M](https://huggingface.co/bartowski/prism-ml_Bonsai-8B-unpacked-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ |
+| [openbmb/MiniCPM-V-2_6-Q4_K_M](https://huggingface.co/openbmb/MiniCPM-V-2_6-gguf) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [bartowski/tencent_Hunyuan-7B-Instruct-Q4_K_M](https://huggingface.co/bartowski/tencent_Hunyuan-7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-Q4_K_M](https://huggingface.co/LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
+| [bartowski/prism-ml_Bonsai-8B-unpacked-Q4_K_M](https://huggingface.co/bartowski/prism-ml_Bonsai-8B-unpacked-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ |
 |  |  |  |  |
 | [gpustack/bge-m3-Q4_K_M.gguf](https://huggingface.co/gpustack/bge-m3-GGUF) | ✓ | ✗ | ✗ |

@@ -217,18 +222,18 @@ cmake --build build\ReleaseOV --parallel
 > [!NOTE]
 > The Windows install path is `C:\Intel\openvino` (no spaces) to avoid quoting problems some CMake/Ninja toolchains have with `C:\Program Files (x86)\...`. Adjust to wherever you installed OpenVINO Runtime. From `cmd`, run `C:\Intel\openvino\setupvars.bat`; from PowerShell, run `& "C:\Intel\openvino\setupvars.ps1"` instead. Once the build is finished you can launch the binaries from any `cmd` or `PowerShell` window after sourcing the matching `setupvars` script for that shell.

-#### Automated Ubuntu Build Script
+#### Ubuntu Build Script

 For Ubuntu24 users, the following shell script automates the prerequisite installs (build tools, OpenCL ICD), the OpenVINO Runtime download/extract/setup, and the Ninja-based llama.cpp build.
-Save the following as `ubuntu-llamacpp-ov-install.sh` next to where you want the `llama.cpp` folder to land, then run it:
+Save the following as `build-llamacpp-ov.sh` next to where you want the `llama.cpp` folder to land, then run it:

 ```bash
-chmod +x ubuntu-llamacpp-ov-install.sh
-./ubuntu-llamacpp-ov-install.sh
+chmod +x build-llamacpp-ov.sh
+./build-llamacpp-ov.sh
 ```

 <details>
-<summary>Click to expand <code>ubuntu-llamacpp-ov-install.sh</code></summary>
+<summary>Click to expand <code>build-llamacpp-ov.sh</code></summary>

 ```bash
 #!/usr/bin/env bash
@@ -237,8 +242,8 @@ chmod +x ubuntu-llamacpp-ov-install.sh
 # ============================================
 set -euo pipefail

-OPENVINO_VERSION_MAJOR="2026.3"
-OPENVINO_VERSION_FULL="2026.3.0.22451.bd8d6542e3c"
+OPENVINO_VERSION_MAJOR="2026.3.1"
+OPENVINO_VERSION_FULL="2026.3.1.22476.56d9685302d"

 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 OPENVINO_INSTALL_DIR="/opt/intel/openvino_${OPENVINO_VERSION_MAJOR}"
@@ -313,8 +318,9 @@ fi
 echo "============================================"
 echo "Configuring with CMake..."
 echo "============================================"
-# shellcheck disable=SC1091
+set +u
 source "${OPENVINO_ROOT}/setupvars.sh"
+set -u

 cmake -B build/ReleaseOV -G Ninja \
    -DCMAKE_BUILD_TYPE=Release \
@@ -334,27 +340,27 @@ echo "  ./build/ReleaseOV/bin/llama-cli -m model.gguf"
 ```

 > [!NOTE]
-> The script pins OpenVINO `2026.3` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release.
+> The script pins OpenVINO `2026.3.1` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release.

 </details>

-#### Automated Windows Build Script
+#### Windows Build Script

 For Windows users, the following `.bat` script automates the prerequisite installs (Git, Ninja, CMake, Visual Studio 2022 Build Tools, vcpkg + OpenCL), the OpenVINO Runtime download/extract, and the Ninja-based llama.cpp build.
-Save the following as `windows-llamacpp-ov-install.bat` next to where you want the `llama.cpp` to land, then run it from either **Command Prompt** or **PowerShell**:
+Save the following as `build-llamacpp-ov.bat` next to where you want the `llama.cpp` to land, then run it from either **Command Prompt** or **PowerShell**:

 ```cmd
 :: Command Prompt
-windows-llamacpp-ov-install.bat
+build-llamacpp-ov.bat
 ```

 ```powershell
 # PowerShell
-.\windows-llamacpp-ov-install.bat
+.\build-llamacpp-ov.bat
 ```

 <details>
-<summary>Click to expand <code>windows-llamacpp-ov-install.bat</code></summary>
+<summary>Click to expand <code>build-llamacpp-ov.bat</code></summary>

 ```bat
@echo off
@@ -364,8 +370,8 @@ REM ============================================
 REM llama.cpp OpenVINO Build Script (Ninja)
 REM ============================================

-set "OPENVINO_VERSION_MAJOR=2026.3"
-set "OPENVINO_VERSION_FULL=2026.3.0.22451.bd8d6542e3c"
+set "OPENVINO_VERSION_MAJOR=2026.3.1"
+set "OPENVINO_VERSION_FULL=2026.3.1.22476.56d9685302d"

 set "SCRIPT_DIR=%~dp0"
 set "VCPKG_DIR=C:\vcpkg"
@@ -453,9 +459,6 @@ if exist "%OPENVINO_INSTALL_DIR%\setupvars.bat" (
    )

    REM Move the single top-level folder contents into the versioned install dir.
-    REM NOTE: delayed expansion (!VAR!) is required because the surrounding else( ... )
-    REM block is parsed once up-front, so %OPENVINO_EXTRACTED% would expand to "" here
-    REM and xcopy would then treat "\*" as C:\* and fail with "Cannot perform a cyclic copy".
    set "OPENVINO_EXTRACTED="
    for /d %%i in ("%OPENVINO_EXTRACT_TMP%\*") do set "OPENVINO_EXTRACTED=%%i"
    if not defined OPENVINO_EXTRACTED (
@@ -547,7 +550,7 @@ endlocal
 ```

 > [!NOTE]
-> The script pins OpenVINO `2026.3` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release. From any new shell, source the matching `setupvars` script via the junction — `call "C:\Intel\openvino\setupvars.bat"` from `cmd`, or `& "C:\Intel\openvino\setupvars.ps1"` from PowerShell. If `winget` cannot register Visual Studio Build Tools on first run, install them once manually and re-run the script from an elevated **Developer Command Prompt for VS 2022**.
+> The script pins OpenVINO `2026.3.1` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release. From any new shell, source the matching `setupvars` script via the junction — `call "C:\Intel\openvino\setupvars.bat"` from `cmd`, or `& "C:\Intel\openvino\setupvars.ps1"` from PowerShell. If `winget` cannot register Visual Studio Build Tools on first run, install them once manually and re-run the script from an elevated **Developer Command Prompt for VS 2022**.

 </details>

@@ -712,6 +715,7 @@ Boolean flags follow a uniform convention: set to a **positive integer** (e.g. `
 | `GGML_OPENVINO_CACHE_DIR`         | String    | `not set`  | Directory for OpenVINO model caching (recommended: `/tmp/ov_cache`). Enables model caching when set. **Not supported on NPU devices.** |
 | `GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR` | String | `not set` | Directory for the frontend compiled-model cache. When set, OpenVINO compiled models are exported as blobs and imported on later runs to skip weight requantization, graph conversion, and compilation for matching single-graph models. |
 | `GGML_OPENVINO_PREFILL_CHUNK_SIZE`| Integer   | `256`      | Token chunk size for **NPU** prefill (NPU-only; ignored on CPU/GPU). Must be a positive integer; otherwise the default is used. |
+| `GGML_OPENVINO_NPU_COMPILE_CONFIG` | String | `not set` | NPU-only compiler mode parameters forwarded to OpenVINO as `NPU_COMPILATION_MODE_PARAMS`, for example `optimization-level=3`. |
 | `GGML_OPENVINO_STATEFUL_EXECUTION`| Boolean   | `0`        | Enable stateful KV cache for better performance. Recommended on CPU, GPU.                                   |
 | `GGML_OPENVINO_DISABLE_CACHE`     | Boolean   | `0`        | Disable the in-process compiled-model / decoder cache (cache is on by default). Set to `1` to disable.      |
 | `GGML_OPENVINO_DISABLE_KV_SLICE`  | Boolean   | `0`        | Disable the KV-cache input-tensor slicing optimization (slicing is on by default on CPU/GPU). Set to `1` to disable. |
@@ -725,9 +729,11 @@ Boolean flags follow a uniform convention: set to a **positive integer** (e.g. `
 | `GGML_OPENVINO_DEBUG_INPUT`       | Boolean   | `0`        | Enable input debugging and print input tensor info.                                                         |
 | `GGML_OPENVINO_DEBUG_OUTPUT`      | Boolean   | `0`        | Enable output debugging and print output tensor info.                                                       |
 | `GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS` | Boolean | `0` | Print tensor address map once.                                                                           |
+| `GGML_OPENVINO_LOG_UNSUPPORTED_OPS`| Boolean   | `0`        | Log warning messages with tensor details and rejection reasons for any ops not supported by the OpenVINO backend. Emits at `WARN` level (requires `--log-verbosity >= 2`, enabled by default). |

 > [!NOTE]
->`GGML_OPENVINO_STATEFUL_EXECUTION` is an **Experimental** feature to allow stateful execution for managing the KV cache internally inside the OpenVINO model, improving performance on CPUs and GPUs. Stateful execution is not effective on NPUs, and not all models currently support this feature. This feature is experimental and has been validated only with the llama-simple, llama-cli, llama-bench, and llama-run applications and is recommended to enable for the best performance. Other applications, such as llama-server and llama-perplexity, are not yet supported.
+> - `GGML_OPENVINO_STATEFUL_EXECUTION` is an **Experimental** feature to allow stateful execution for managing the KV cache internally inside the OpenVINO model, improving performance on CPUs and GPUs. Stateful execution is not effective on NPUs, and not all models currently support this feature. This feature is experimental and has been validated only with the llama-simple, llama-cli, llama-bench, and llama-run applications and is recommended to enable for the best performance. Other applications, such as llama-server and llama-perplexity, are not yet supported.
+> - `GGML_OPENVINO_LOG_UNSUPPORTED_OPS` emits logs at `WARN` level (`GGML_LOG_WARN`), which requires application log verbosity `--log-verbosity >= 2` (or `-lv 2`).

 ### Example Usage

@@ -8,7 +8,7 @@
        "toolset":      { "value": "host=x86_64", "strategy": "external" },
        "cacheVariables": {
            "ANDROID_ABI":      "arm64-v8a",
-            "ANDROID_PLATFORM": "android-31",
+            "ANDROID_PLATFORM": "android-34",
            "CMAKE_TOOLCHAIN_FILE": "$env{ANDROID_NDK_ROOT}/build/cmake/android.toolchain.cmake",
            "CMAKE_C_FLAGS":   "-march=armv8.7a+fp16+dotprod+i8mm -fvectorize -ffp-model=fast -fno-finite-math-only -flto -D_GNU_SOURCE",
            "CMAKE_CXX_FLAGS": "-march=armv8.7a+fp16+dotprod+i8mm -fvectorize -ffp-model=fast -fno-finite-math-only -flto -D_GNU_SOURCE",
@@ -2,39 +2,47 @@

 ## Setup

-### Android
+The cross-compilation toolchain images are provided by the
+[Qualcomm Snapdragon Toolchain registry](https://github.com/snapdragon-toolchain).
+These Docker images include the Android NDK, OpenCL SDK, Hexagon SDK, CMake, and the necessary cross-compilers:

-The easiest way to build llama.cpp for a Snapdragon-based Android device is using the toolchain Docker image (see github.com/snapdragon-toolchain).
-This image includes Android NDK, OpenCL SDK, Hexagon SDK, CMake, etc.
+* **Android toolchain**: `ghcr.io/snapdragon-toolchain/arm64-android:v0.7`
+* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7`

-This method works on Linux, macOS, and Windows. macOS and Windows users should install Docker Desktop.
-
-```
-~/src/llama.cpp$ docker run -it -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-android:v0.7
-[d]/> cd /workspace
-```
-
-Note: The rest of the **Android** build process assumes that you're running inside the toolchain container.
-
-### Windows On Snapdragon
-
-Native Windows 11 arm64 builds has the following tools dependencies:
- MS Visual Studio 2026 (Community Edition or Pro)
-  - MSVC arm64 standard and runtime libraries
-  - UCRT and Driver Kit
- LLVM core libraries and Clang compiler (winget)
- CMake, Git, Python (winget)
- Hexagon SDK Community Edition 6.6 or later (see windows.md)
- OpenCL SDK 2.3 or later (see windows.md)
-
-Note: The rest of the **Windows** build process assumes that you're running natively in Powershell.
-Adapt below build commands accordingly.
+The unified build utility (`scripts/snapdragon/build.py`) automatically pulls
+and orchestrates these containers to perform target compilation.
+You only need to ensure that Docker (or Docker Desktop on macOS/Windows) is running on your host machine.
+Specific setup, build, and installation details for Linux and Windows on Snapdragon platforms are documented in:
+* [Linux on Snapdragon guide](linux.md)
+* [Windows on Snapdragon guide](windows.md)

 ## How to Build

-Let's build llama.cpp with CPU, OpenCL, and Hexagon backends via CMake presets:
+### Using build.py script (Recommended)

+The easiest way to build llama.cpp is by using the `scripts/snapdragon/build.py` script. It automatically copies the CMake presets,
+launches the correct compilation Docker container, builds the libraries and tools,
+installs them, and optionally pushes them to your ADB device.
+
+Build and deploy for Android target (accepts `android` or `adb` alias):
 ```
+$ ./scripts/snapdragon/build.py --target adb --push
+```
+
+Build and deploy for Linux target (accepts `linux` or `lnx` alias):
+```
+$ ./scripts/snapdragon/build.py --target linux:user@host --push
+```
+
+### Manual CMake Build
+
+Alternatively, you can build llama.cpp manually by entering the cross-compilation Docker container and running the CMake commands:
+
+```bash
+# Start the cross-compilation container manually:
+~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-android:v0.7
+
+# Inside the container, build the project using presets:
 [d]/workspace> cp docs/backend/snapdragon/CMakeUserPresets.json .

 [d]/workspace> cmake --preset arm64-android-snapdragon-release -B build-snapdragon
@@ -68,19 +76,19 @@ Preset CMake variables:
 To generate an installable "package" simply use cmake --install:

 ```
-[d]/workspace> cmake --install build-snapdragon --prefix pkg-snapdragon/llama.cpp
+[d]/workspace> cmake --install build-snapdragon --prefix pkg-android/llama.cpp
 -- Install configuration: "Release"
-- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-cpu.so
-- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-opencl.so
-- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-hexagon.so
-- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v73.so
-- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v75.so
-- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v79.so
-- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v81.so
-- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-cpu.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-opencl.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-hexagon.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v73.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v75.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v79.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v81.so
+-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml.so
 ...
-- Installing: /workspace/pkg-snapdragon/llama.cpp/bin/llama-bench
-- Installing: /workspace/pkg-snapdragon/llama.cpp/bin/llama-cli
+-- Installing: /workspace/pkg-android/llama.cpp/bin/llama-bench
+-- Installing: /workspace/pkg-android/llama.cpp/bin/llama-cli
 ...
 ```

@@ -91,14 +99,14 @@ To generate an installable "package" simply use cmake --install:
 For this step, your device needs to be configured for on-device development.
 Please see https://developer.android.com/studio/debug/dev-options for details.

-Once ADB is enabled, use `adb push` to install `pkg-snapdragon` on the device.
+Once ADB is enabled, use `adb push` to install `pkg-android` on the device.
 **Note that the toolchain Docker image doesn't have ADB and doesn't set up the ADB bridge. Please use native ADB on the host.**

 ```
-~/src/llama.cpp$ adb push pkg-snapdragon/llama.cpp /data/local/tmp/
-pkg-snapdragon/llama.cpp/bin/: 67 files pushed, 0 skipped. 190.2 MB/s (919095042 bytes in 4.607s)
-pkg-snapdragon/llama.cpp/include/: 19 files pushed, 0 skipped. 20.5 MB/s (255173 bytes in 0.012s)
-pkg-snapdragon/llama.cpp/lib/: 16 files pushed, 0 skipped. 144.4 MB/s (43801382 bytes in 0.289s)
+~/src/llama.cpp$ adb push pkg-android/llama.cpp /data/local/tmp/
+pkg-android/llama.cpp/bin/: 67 files pushed, 0 skipped. 190.2 MB/s (919095042 bytes in 4.607s)
+pkg-android/llama.cpp/include/: 19 files pushed, 0 skipped. 20.5 MB/s (255173 bytes in 0.012s)
+pkg-android/llama.cpp/lib/: 16 files pushed, 0 skipped. 144.4 MB/s (43801382 bytes in 0.289s)
 102 files pushed, 0 skipped. 186.9 MB/s (963151597 bytes in 4.914s)
 ```

@@ -115,24 +123,44 @@ Llama-3.2-1B-Instruct-Q4_0.gguf: 1 file pushed, 0 skipped. 38.3 MB/s (773025920

 ### Windows

-All artifacts are already installed in the `pkg-snapdragon` folder.
-To run, adapt below instructions to use Powershell scripts in `scripts/snapdragon/windows`.
+All artifacts are already installed in the `pkg-wos` folder.
+To run, you can use the `scripts/snapdragon/run.py` runner script (see details below).

 ## How to Run

-The easiest way to run llama.cpp cli tools is using provided wrapper scripts that properly set up all required environment variables.
+The easiest way to run llama.cpp cli tools is using the provided `scripts/snapdragon/run.py` wrapper script. This script automatically
+maps CLI options to environment variables, resolves executable paths, and runs the command locally, via ADB, or remotely via SSH on the
+target device.

-llama.cpp supports three backends on Snapdragon-based devices: CPU, Adreno GPU (GPUOpenCL), and Hexagon NPU (HTP0-4).
-You can select which backend to run the model on using the `D=` variable, which maps to the `--device` option.
+llama.cpp supports three backends on Snapdragon-based devices: CPU, Adreno GPU (GPUOpenCL), and Hexagon NPU.
+You can select which backend(s) to run the model on using the `--device` option of the tool (or `--devices` option in `run.py`).

 Hexagon NPU behaves as a "GPU" device when it comes to `-ngl` and other offload-related options.

-Here are some examples of running various llama.cpp tools via ADB.
+Here are some examples of running various llama.cpp tools.

-Simple question for Llama-3.2-1B
+Generating a completion with Gemma on Android (relying on default `HTP0:0` device and default thread count `-t 6`):

 ```
-~/src/llama.cpp$ M=Llama-3.2-1B-Instruct-Q4_0.gguf D=HTP0 ./scripts/snapdragon/adb/run-completion.sh -p "what is the most popular cookie in the world?"
+~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb -- llama-completion -m models/gemma-2-2b-it-Q4_0.gguf -f prompts/sample_prompt_1024.txt --jinja -st
+...
+ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1
+ggml-hex: Hexagon Arch version v79
+ggml-hex: allocating new session: HTP0:0
+...
+load_tensors: offloading output layer to GPU
+load_tensors: offloaded 27/27 layers to GPU
+load_tensors:          CPU model buffer size =   300.00 MiB
+load_tensors:      HTP0:0 model buffer size  =  1400.26 MiB
+...
+llama_perf_context_print: prompt eval time =     320.00 ms /  1024 tokens (    0.31 ms per token,  3200.00 tokens per second)
+llama_perf_context_print:        eval time =     2100.00 ms /   100 runs   (   21.00 ms per token,    47.62 tokens per second)
+```
+
+Simple question for Llama-3.2-1B:
+
+```
+~/src/llama.cpp$ ./scripts/snapdragon/run.py --target android --devices HTP0 -- llama-cli -m Llama-3.2-1B-Instruct-Q4_0.gguf -p "what is the most popular cookie in the world?"
 ...
 ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1
 ggml-hex: Hexagon Arch version v79
@@ -142,8 +170,7 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v
 load_tensors: offloading output layer to GPU
 load_tensors: offloaded 17/17 layers to GPU
 load_tensors:          CPU model buffer size =   225.49 MiB
-load_tensors:         HTP0 model buffer size =     0.26 MiB
-load_tensors:  HTP0-REPACK model buffer size =   504.00 MiB
+load_tensors:         HTP0 model buffer size =   504.26 MiB
 ...
 I hope this helps you understand the world's most popular cookies! [end of text]
 ...
@@ -156,60 +183,25 @@ llama_perf_context_print:    graphs reused =        473
 llama_memory_breakdown_print: | memory breakdown [MiB] | total   free    self   model   context   compute    unaccounted |
 llama_memory_breakdown_print: |   - HTP0 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
 llama_memory_breakdown_print: |   - Host               |                  439 =   225 +     136 +      77                |
-llama_memory_breakdown_print: |   - HTP0-REPACK        |                  504 =   504 +       0 +       0                |
 ```

-Summary request for OLMoE-1B-7B. This is a large model that requires two HTP sessions/devices
+Op test for MUL_MAT:

 ```
-~/src/llama.cpp$ M=OLMoE-1B-7B-0125-Instruct-Q4_0.gguf NDEV=2 D=HTP0,HTP1 ./scripts/snapdragon/adb/run-completion.sh -f surfing.txt
+~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --hex-hostbuf 0 --devices HTP0:0 -- test-backend-ops -b HTP0:0 -o MUL_MAT
 ...
-ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1
-ggml-hex: Hexagon Arch version v81
-ggml-hex: allocating new session: HTP0
-ggml-hex: allocating new session: HTP1
-...
-load_tensors: offloading output layer to GPU
-load_tensors: offloaded 17/17 layers to GPU
-load_tensors:          CPU model buffer size =   143.86 MiB
-load_tensors:         HTP1 model buffer size =     0.23 MiB
-load_tensors:  HTP1-REPACK model buffer size =  1575.00 MiB
-load_tensors:         HTP0 model buffer size =     0.28 MiB
-load_tensors:  HTP0-REPACK model buffer size =  2025.00 MiB
-...
-llama_context:        CPU  output buffer size =     0.19 MiB
-llama_kv_cache:       HTP1 KV buffer size =   238.00 MiB
-llama_kv_cache:       HTP0 KV buffer size =   306.00 MiB
-llama_kv_cache: size =  544.00 MiB (  8192 cells,  16 layers,  1/1 seqs), K (q8_0):  272.00 MiB, V (q8_0):  272.00 MiB
-llama_context:       HTP0 compute buffer size =    15.00 MiB
-llama_context:       HTP1 compute buffer size =    15.00 MiB
-llama_context:        CPU compute buffer size =    24.56 MiB
-...
-llama_perf_context_print: prompt eval time =    1730.57 ms /   212 tokens (    8.16 ms per token,   122.50 tokens per second)
-llama_perf_context_print:        eval time =    5624.75 ms /   257 runs   (   21.89 ms per token,    45.69 tokens per second)
-llama_perf_context_print:       total time =    7377.33 ms /   469 tokens
-llama_perf_context_print:    graphs reused =        255
-llama_memory_breakdown_print: | memory breakdown [MiB] | total   free    self   model   context   compute    unaccounted |
-llama_memory_breakdown_print: |   - HTP0 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP1 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - Host               |                  742 =   144 +     544 +      54                |
-llama_memory_breakdown_print: |   - HTP1-REPACK        |                 1575 =  1575 +       0 +       0                |
-llama_memory_breakdown_print: |   - HTP0-REPACK        |                 2025 =  2025 +       0 +       0                |
-```
-
-Op test for MUL_MAT
-
-```
-~/src/llama.cpp$ HB=0 ./scripts/snapdragon/adb/run-tool.sh test-backend-ops -b HTP0 -o MUL_MAT
-...
-Backend 2/3: HTP0
+Backend 2/3: HTP0:0
 Device description: Hexagon
 Device memory: 2048 MB (2048 MB free)
 MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
 MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=2,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
 MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=3,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
+```

-~/src/llama.cpp-hexagon$ M=Llama-3.2-1B-Instruct-Q4_0.gguf ./scripts/snapdragon/adb/run-bench.sh -p 128 -n 64
+Llama benchmark:
+
+```
+~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --devices HTP0 -- llama-bench -p 128 -n 64 -m Llama-3.2-1B-Instruct-Q4_0.gguf
 ...
 ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1
 ggml-hex: Hexagon Arch version v79
@@ -219,15 +211,20 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v
 | ---------------| ---------: | -----: | ---------- | --: | ------: | ------: | ---: | ----: | ------------: |
 | llama 1B Q4_0  | 729.75 MiB | 1.24 B | HTP        |  99 |       4 |     128 |    0 | pp128 | 169.42 ± 1.75 |
 | llama 1B Q4_0  | 729.75 MiB | 1.24 B | HTP        |  99 |       4 |     128 |    0 |  tg64 |  51.54 ± 1.13 |
-
-build: 6a8cf8914 (6733)
 ```

 ## Environment variables

- `GGML_HEXAGON_NDEV=1`
-  Controls the number of devices/sessions to allocate. The default is 1.
-  Most quantized models under 4B fit into a single session; an 8B model needs two, and a 20B model needs four.
+- `GGML_HEXAGON_DEVICES` (default: not set, defaults to HTP0 session)
+  Controls which NPU devices and sessions to allocate. Can be configured as:
+  - A single integer `N`: Allocates `N` sessions named `HTP0`, `HTP1`, ..., `HTP<N-1>` (behaves identically to `GGML_HEXAGON_NDEV=N`).
+  - A comma-separated list of device names in `HTP<physical_idx>:<virtual_idx>` format (or legacy `HTP<idx>` format). For example, `HTP0:0,HTP0:1` creates two virtual
+    sessions on the first physical NPU (useful for memory limits). `HTP0:0,HTP1:0` allocates one session on each of the two physical NPUs
+    on a dual-NPU device.
+
+- `GGML_HEXAGON_NDEV` (deprecated)
+  Replaced by `GGML_HEXAGON_DEVICES`. Controls the number of virtual sessions to allocate on physical NPU `0`.
+  Allocates sessions named `HTP0`, `HTP1`, etc.

 - `GGML_HEXAGON_NHVX=0`
  Controls the number of HVX hardware threads to use. The default is all (actual number varies depending on the hardware version).
@@ -255,26 +252,17 @@ build: 6a8cf8914 (6733)
  - `2` Extended profile with per-op `usecs`, `cycles` and default PMU counter data
  - `0x1,...,0x8` Extended profile with per-op `usecs`, `cycles` and custom PMU counter data

-  The logging output can be either saved into a file for post-processing or it can be piped directly into the post-processing tool to generate the report.
+  The logging output can be either saved into a file for post-processing or it can be piped directly into the post-processing tool
+  to generate the report.
  Examples:

-      `GGML_HEXAGON_PROFILE=1 llama-completion ... |& ./scripts/snapdragon/ggml-hexagon-profile.py -`
-
- `GGML_HEXAGON_OPSTAGE=0x0`
-  Allows enabling specific stages of the Op processing pipeline:
-
-  - `0x1` Enable Op Queue (i.e., queuing Ops into NPU)
-  - `0x2` Enable Op Compute (MUL_MAT, etc.)
-
-  Examples:
-
-      `GGML_HEXAGON_OPSTAGE=0x1 llama-completion ...` - Ops are enqueued to the NPU but dma & compute are disabled
-      `GGML_HEXAGON_OPSTAGE=0x3 llama-completion ...` - Full queuing and processing of Ops (default)
+      `GGML_HEXAGON_PROFILE=1 ./scripts/snapdragon/run.py --target adb -- llama-cli ... |& ./scripts/snapdragon/ggml-hexagon-profile.py -`

 - `GGML_HEXAGON_OPFILTER=regex`
  Allows filtering (disabling) Ops that match the regex pattern:

  Examples:

-      `GGML_HEXAGON_OPFILTER="FLASH_ATTN_EXT" llama-completion ...` - Disable Flash Attention on Hexagon (falls back to CPU or GPU)
-      `GGML_HEXAGON_OPFILTER="ADD\|SUB" llama-completion ...` - Disable ADD and SUB on Hexagon (fall back to CPU or GPU)
+      `GGML_HEXAGON_OPFILTER="FLASH_ATTN_EXT" ./scripts/snapdragon/run.py --target adb -- llama-cli ...` - Disable Flash Attention on Hexagon (falls back to CPU or GPU)
+      `GGML_HEXAGON_OPFILTER="ADD\|SUB" ./scripts/snapdragon/run.py --target adb -- llama-cli ...` - Disable ADD and SUB on Hexagon (fall back to CPU or GPU)
+
@@ -39,22 +39,21 @@ the repacking.

 ## Large model handling

-Hexagon NPU session (aka Process Domain (PD) in the Hexagon docs) is limited to a memory mapping of around 3.5GB.
-In llama.cpp/GGML the Hexagon session is mapped to a single GGML backend device (HTP0, HTP1, etc).
+Hexagon NPU sessions (aka Process Domains (PD) in the Hexagon SDK) are limited to a maximum memory mapping window of around 3.5GB.
+In llama.cpp/GGML, each Hexagon session is mapped to a single GGML backend device (e.g., `HTP0:0`, `HTP0:1`, etc. when using
+`GGML_HEXAGON_DEVICES`, or `HTP0`, `HTP1` in legacy mode).

-In order to map models larger than 3.5GB we need to allocate multiple devices and split the model.
-For this we're taking advantage of the llama.cpp/GGML multi-GPU layer-splitting support.
-Each Hexagon device behaves like a GPU from the offload and model splitting perspective.
+To support running models larger than 3.5GB on a single device, the Hexagon backend dynamically maps and unmaps execution buffers
+during the graph execution cycle to stay within the Process Domain window. This enables large models to run successfully on a single
+NPU device.

-Here is an example of running GPT-OSS-20B model on a newer Snapdragon device with 16GB of DDR.
+Alternatively, users can choose to use standard llama.cpp/GGML layer-splitting mode to partition and split the model across
+multiple Hexagon devices or virtual sessions (which behave like multiple GPUs from the offload and splitting perspective).
+
+Here is an example of running GPT-OSS-20B model on a Snapdragon device using 4 virtual sessions on a single NPU (physical index 0).

 ```
-M=gpt-oss-20b-Q4_0.gguf NDEV=4 D=HTP0,HTP1,HTP2,HTP3 P=surfing.txt scripts/snapdragon/adb/run-completion.sh -f surfing.txt -n 32
-...
-LD_LIBRARY_PATH=/data/local/tmp/llama.cpp/lib
-ADSP_LIBRARY_PATH=/data/local/tmp/llama.cpp/lib
-GGML_HEXAGON_NDEV=4 ./bin/llama-cli --load-mode none -m /data/local/tmp/llama.cpp/../gguf/gpt-oss-20b-Q4_0.gguf
-      -t 4 --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 --device HTP0,HTP1,HTP2,HTP3 -no-cnv -f surfing.txt
+~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --devices HTP0:0,HTP0:1,HTP0:2,HTP0:3 -- llama-cli --load-mode none -m /data/local/tmp/gguf/gpt-oss-20b-Q4_0.gguf -t 4 --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 -no-cnv -f surfing.txt
 ...
 llama_model_loader: - type  f32:  289 tensors
 llama_model_loader: - type q4_0:   96 tensors
@@ -63,33 +62,29 @@ llama_model_loader: - type mxfp4:  72 tensors
 ...
 load_tensors: offloaded 25/25 layers to GPU
 load_tensors:          CPU model buffer size =  1182.09 MiB
-load_tensors:         HTP1 model buffer size =     6.64 MiB
-load_tensors:  HTP1-REPACK model buffer size =  2505.94 MiB
-load_tensors:         HTP3 model buffer size =     5.55 MiB
-load_tensors:  HTP3-REPACK model buffer size =  2088.28 MiB
-load_tensors:         HTP0 model buffer size =     7.75 MiB
-load_tensors:  HTP0-REPACK model buffer size =  2923.59 MiB
-load_tensors:         HTP2 model buffer size =     6.64 MiB
-load_tensors:  HTP2-REPACK model buffer size =  2505.94 MiB
+load_tensors:       HTP0:1 model buffer size =  2512.58 MiB
+load_tensors:       HTP0:3 model buffer size =  2093.83 MiB
+load_tensors:       HTP0:0 model buffer size =  2931.34 MiB
+load_tensors:       HTP0:2 model buffer size =  2512.58 MiB
 ...
 llama_context: n_ctx_per_seq (8192) < n_ctx_train (131072) -- the full capacity of the model will not be utilized
 llama_context:        CPU  output buffer size =     0.77 MiB
 llama_kv_cache_iswa: creating non-SWA KV cache, size = 8192 cells
-llama_kv_cache:       HTP1 KV buffer size =    25.50 MiB
-llama_kv_cache:       HTP3 KV buffer size =    25.50 MiB
-llama_kv_cache:       HTP0 KV buffer size =    25.50 MiB
-llama_kv_cache:       HTP2 KV buffer size =    25.50 MiB
+llama_kv_cache:     HTP0:1 KV buffer size =    25.50 MiB
+llama_kv_cache:     HTP0:3 KV buffer size =    25.50 MiB
+llama_kv_cache:     HTP0:0 KV buffer size =    25.50 MiB
+llama_kv_cache:     HTP0:2 KV buffer size =    25.50 MiB
 llama_kv_cache: size =  102.00 MiB (  8192 cells,  12 layers,  1/1 seqs), K (q8_0):   51.00 MiB, V (q8_0):   51.00 MiB
 llama_kv_cache_iswa: creating     SWA KV cache, size = 256 cells
-llama_kv_cache:       HTP1 KV buffer size =     0.80 MiB
-llama_kv_cache:       HTP3 KV buffer size =     0.53 MiB
-llama_kv_cache:       HTP0 KV buffer size =     1.06 MiB
-llama_kv_cache:       HTP2 KV buffer size =     0.80 MiB
+llama_kv_cache:     HTP0:1 KV buffer size =     0.80 MiB
+llama_kv_cache:     HTP0:3 KV buffer size =     0.53 MiB
+llama_kv_cache:     HTP0:0 KV buffer size =     1.06 MiB
+llama_kv_cache:     HTP0:2 KV buffer size =     0.80 MiB
 llama_kv_cache: size =    3.19 MiB (   256 cells,  12 layers,  1/1 seqs), K (q8_0):    1.59 MiB, V (q8_0):    1.59 MiB
-llama_context:       HTP0 compute buffer size =    16.06 MiB
-llama_context:       HTP1 compute buffer size =    16.06 MiB
-llama_context:       HTP2 compute buffer size =    16.06 MiB
-llama_context:       HTP3 compute buffer size =    16.06 MiB
+llama_context:     HTP0:0 compute buffer size =    16.06 MiB
+llama_context:     HTP0:1 compute buffer size =    16.06 MiB
+llama_context:     HTP0:2 compute buffer size =    16.06 MiB
+llama_context:     HTP0:3 compute buffer size =    16.06 MiB
 llama_context:        CPU compute buffer size =    98.19 MiB
 ...
 llama_perf_context_print: prompt eval time =    3843.67 ms /   197 tokens ( 19.51 ms per token, 51.25 tokens per second)
@@ -97,13 +92,9 @@ llama_perf_context_print:        eval time =    1686.13 ms /    31 runs   ( 54.3
 llama_perf_context_print:       total time =    6266.30 ms /   228 tokens
 llama_perf_context_print:    graphs reused =         30
 llama_memory_breakdown_print: | memory breakdown [MiB] | total   free    self   model   context   compute    unaccounted |
-llama_memory_breakdown_print: |   - HTP0 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP1 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP2 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP3 (Hexagon)     |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:0 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:1 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:2 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:3 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
 llama_memory_breakdown_print: |   - Host               |                 1476 =  1208 +     105 +     162                |
-llama_memory_breakdown_print: |   - HTP1-REPACK        |                 2505 =  2505 +       0 +       0                |
-llama_memory_breakdown_print: |   - HTP3-REPACK        |                 2088 =  2088 +       0 +       0                |
-llama_memory_breakdown_print: |   - HTP0-REPACK        |                 2923 =  2923 +       0 +       0                |
-llama_memory_breakdown_print: |   - HTP2-REPACK        |                 2505 =  2505 +       0 +       0                |
 ```
@@ -1,25 +1,37 @@
 # Snapdragon-based Linux devices

-## Docker Setup
+The cross-compilation is performed using the Snapdragon Linux Docker toolchain image (see
+[github.com/snapdragon-toolchain](https://github.com/snapdragon-toolchain)):

-The easiest way to build llama.cpp for a Snapdragon-based Linux device is using the toolchain Docker image (see [github.com/snapdragon-toolchain](https://github.com/snapdragon-toolchain)).
-This image includes OpenCL SDK, Hexagon SDK, CMake, and the ARM64 Linux cross-compilation toolchain.
+* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7`

-Cross-compilation is supported on **Linux X86** hosts. The resulting binaries are deployed to and run on the target **Qualcomm Snapdragon ARM64 Linux** device.
-
-```
-~/src/llama.cpp$ docker run -it -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-linux:v0.1
-[d]/> cd /workspace
-```
-
-Note: The rest of the **Linux** build process assumes that you're running inside the toolchain container.
+The unified build utility (`scripts/snapdragon/build.py`) automatically pulls
+and orchestrates this container to perform target compilation. You only need to
+ensure that Docker is running on your host machine.


 ## How to Build

-Let's build llama.cpp with CPU, OpenCL, and Hexagon backends via CMake presets:
+### Using build.py script (Recommended)

+The easiest way to build llama.cpp is by using the `scripts/snapdragon/build.py` script. It automatically copies the CMake presets,
+launches the correct compilation Docker container, builds the libraries and tools,
+installs them, and optionally pushes them to your target device.
+
+Build and deploy for a Linux target (using SSH deployment alias `lnx` or `linux`):
 ```
+$ ./scripts/snapdragon/build.py --target lnx:user@host --push
+```
+
+### Manual CMake Build
+
+Alternatively, you can build llama.cpp manually by entering the cross-compilation Docker container and running the CMake commands:
+
+```bash
+# Start the cross-compilation container manually:
+~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-linux:v0.7
+
+# Inside the container, build the project using presets:
 [d]/workspace> cp docs/backend/snapdragon/CMakeUserPresets.json .

 [d]/workspace> cmake --preset arm64-linux-snapdragon-release -B build-snapdragon
@@ -30,17 +42,19 @@ Let's build llama.cpp with CPU, OpenCL, and Hexagon backends via CMake presets:
 To generate an installable "package" simply use cmake --install, then zip it:

 ```
-[d]/workspace> cmake --install build-snapdragon --prefix pkg-snapdragon
-[d]/workspace> zip -r pkg-snapdragon.zip pkg-snapdragon
+[d]/workspace> cmake --install build-snapdragon --prefix pkg-linux
+[d]/workspace> zip -r pkg-linux.zip pkg-linux
 ```

 ## How to Install

-For this step, you will deploy the built binaries and libraries to the target Linux device. Transfer `pkg-snapdragon.zip` to the target device, then unzip it and set up the environment variables:
+For this step, you will deploy the built binaries and libraries to the target
+Linux device. Transfer `pkg-linux.zip` to the target device, then unzip it
+and set up the environment variables:

 ```
-$ unzip pkg-snapdragon.zip
-$ cd pkg-snapdragon
+$ unzip pkg-linux.zip
+$ cd pkg-linux
 $ export LD_LIBRARY_PATH=./lib
 $ export ADSP_LIBRARY_PATH=./lib
 ```
@@ -52,7 +66,28 @@ $ wget https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/
 ```

 ## How to Run
-Next, since we have setup the environment variables, we can run the llama-cli with the Hexagon backends:
+You can run locally on the Snapdragon Linux device:
+```
+$ ./scripts/snapdragon/run.py --devices HTP0 -- llama-cli -m Llama-3.2-3B-Instruct-Q4_0.gguf -ngl 99 -p "what is the most popular cookie in the world?"
+```
+
+Or run remotely from your host development machine using the SSH target option:
+```
+$ ./scripts/snapdragon/run.py --target lnx:user@host --devices HTP0 -- llama-cli -m Llama-3.2-3B-Instruct-Q4_0.gguf -ngl 99 -p "what is the most popular cookie in the world?"
+```
+
+For multi-NPU systems, you can run a tensor split completion command targeting a remote Linux system:
+```
+$ ./scripts/snapdragon/run.py --target ubuntu:maxk@192.168.1.87 --device HTP0:0,HTP1:0 -- llama-completion -m models/gemma-2b-it-Q4_0.gguf -f prompts/sample_prompt_1024.txt --jinja -st --split-mode tensor --ctx-size 8192
+```
+
+This translates to the following command being executed remotely via SSH:
+```
+ ssh maxk@192.168.1.87 "cd ~/llama.cpp && ulimit -c unlimited && LD_LIBRARY_PATH=./lib ADSP_LIBRARY_PATH=./lib GGML_HEXAGON_DEVICES=HTP0:0,HTP1:0 GGML_HEXAGON_OPPOLL=1 ./bin/llama-completion -m models/gemma-2b-it-Q4_0.gguf -f prompts/sample_prompt_1024.txt --jinja -st --split-mode tensor --ctx-size 8192 -v -n 16 --device HTP0:0,HTP1:0 -ngl 99 --ubatch-size 1024 -fa on -t 6"
+```
+
+Alternatively, you can run the binary directly on the device:
 ```
 $ ./bin/llama-cli -m Llama-3.2-3B-Instruct-Q4_0.gguf --device HTP0 -ngl 99 -p "what is the most popular cookie in the world?"
 ```
+
@@ -1,3 +1,18 @@
+# Snapdragon-based Windows devices
+
+## Tool Dependencies
+
+Native Windows 11 arm64 builds have the following tool dependencies:
+- MS Visual Studio 2026 (Community Edition or Pro)
+  - MSVC arm64 standard and runtime libraries
+  - UCRT and Driver Kit
+- LLVM core libraries and Clang compiler (winget)
+- CMake, Git, Python (winget)
+- Hexagon SDK Community Edition 6.6 or later (see below)
+- OpenCL SDK 2.3 or later (see below)
+
+Note: The rest of the **Windows** build process assumes that you're running natively in Powershell.
+
 ## Overview

 The document covers procedures for installing the latest GPU and NPU drivers, and OpenCL and Hexagon SDKs.
@@ -53,7 +68,8 @@ Download the driver from

    https://softwarecenter.qualcomm.com/catalog/item/Qualcomm_HND

-After the automated installation and reboot please make sure that the Hexagon NPU device shows up in the `Device Manager` (under `Neural Processors`).
+After the automated installation and reboot please make sure that the Hexagon NPU device shows up in the `Device Manager`
+(under `Neural Processors`).

 If the device is not available you can try installing all components (`qcnspmcdm8380`, `qcnspmcdm8380_ext`) manually.
 The components are extracted into
@@ -130,12 +146,12 @@ However, additional settings are required for generating and signing HTP Ops lib

 > cmake --preset arm64-windows-snapdragon-release -B build-wos
 ...
-> cmake --install build-wos --prefix pkg-snapdragon
+> cmake --install build-wos --prefix pkg-wos
 ```

 Once the build is complete HTP ops libraries will be installed like this
 ```
-> dir pkg-snapdragon/lib
+> dir pkg-wos/lib
 ...
 -a----         1/22/2026   6:01 PM         187656 libggml-htp-v73.so
 -a----         1/22/2026   6:01 PM         191752 libggml-htp-v75.so
@@ -147,8 +163,8 @@ Once the build is complete HTP ops libraries will be installed like this
 The .cat file, the signature and proper certificate installation can be verified with

 ```
-> signtool.exe verify /v /pa .\pkg-snapdragon\lib\libggml-htp.cat
-Verifying: .\pkg-snapdragon\lib\libggml-htp.cat
+> signtool.exe verify /v /pa .\pkg-wos\lib\libggml-htp.cat
+Verifying: .\pkg-wos\lib\libggml-htp.cat

 Signature Index: 0 (Primary Signature)
 Hash of file (sha256): 9820C664DA59D5EAE31DBB664127FCDAEF59CDC31502496BC567544EC2F401CF
@@ -156,6 +172,6 @@ Hash of file (sha256): 9820C664DA59D5EAE31DBB664127FCDAEF59CDC31502496BC567544EC
 Signing Certificate Chain:
        Issued to: GGML.HTP.v1
 ...
-Successfully verified: .\pkg-snapdragon\lib\libggml-htp.cat
+Successfully verified: .\pkg-wos\lib\libggml-htp.cat
 ...
 ```
@@ -35,8 +35,8 @@ Legend:
 |                              COS | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                      COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              CPY | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
-|               CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|               CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -19292,10 +19292,10 @@
 "Vulkan0","FLASH_ATTN_EXT","hsk=128,hsv=128,nh=8,nr23=[4,1],kv=4096,nb=512,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=f16,type_V=f16,permute=[0,1,2,3]","support","1","yes","Vulkan"
 "Vulkan0","FLASH_ATTN_EXT","hsk=256,hsv=256,nh=4,nr23=[6,1],kv=16384,nb=512,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=f16,type_V=f16,permute=[0,1,2,3]","support","1","yes","Vulkan"
 "Vulkan0","FLASH_ATTN_EXT","hsk=128,hsv=128,nh=8,nr23=[4,1],kv=16384,nb=512,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=f16,type_V=f16,permute=[0,1,2,3]","support","1","yes","Vulkan"
-"Vulkan0","CROSS_ENTROPY_LOSS","type=f32,ne=[10,5,4,3]","support","0","no","Vulkan"
-"Vulkan0","CROSS_ENTROPY_LOSS","type=f32,ne=[30000,1,1,1]","support","0","no","Vulkan"
-"Vulkan0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[10,5,4,3]","support","0","no","Vulkan"
-"Vulkan0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[30000,1,1,1]","support","0","no","Vulkan"
+"Vulkan0","CROSS_ENTROPY_LOSS","type=f32,ne=[10,5,4,3]","support","1","yes","Vulkan"
+"Vulkan0","CROSS_ENTROPY_LOSS","type=f32,ne=[30000,1,1,1]","support","1","yes","Vulkan"
+"Vulkan0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[10,5,4,3]","support","1","yes","Vulkan"
+"Vulkan0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[30000,1,1,1]","support","1","yes","Vulkan"
 "Vulkan0","OPT_STEP_ADAMW","type=f32,ne=[10,5,4,3]","support","1","yes","Vulkan"
 "Vulkan0","OPT_STEP_SGD","type=f32,ne=[10,5,4,3]","support","1","yes","Vulkan"
 "Vulkan0","GATED_DELTA_NET","type=f32,head_count=32,head_size=128,n_seq_tokens=1,n_seqs=1,v_repeat=1,permuted=0,kda=0,K=1","support","1","yes","Vulkan"
@@ -212,6 +212,15 @@ Use `--backend-sampling` to run supported target-model samplers on the model bac

 Unsupported samplers and device layouts fall back to CPU sampling. Tensor split mode does not support backend sampling. A fixed seed produces repeatable random draws, but stochastic CPU and backend sampling can still select different tokens because floating-point operations can differ between implementations and devices. Use greedy sampling when exact output matching is required.

+### Synthetic Acceptance
+
+`llama-server` and `llama-cli` can replace normal speculative verification with synthetic decisions for benchmarking. The generated output is not valid model output because accepted draft tokens do not have to match the target model.
+
+Use exactly one of these options:
+
+- `--spec-synth-rates P0,P1,...` sets unconditional per-position acceptance probabilities. Entry `i` is the probability that the first `i+1` draft tokens are all accepted. The number of entries must match the effective maximum draft length. Values must be finite, within `[0, 1]`, and monotonically non-increasing.
+- `--spec-synth-len L` sets the target mean acceptance length, including the target token. For `K` maximum draft tokens, `L` must be within `[1, K+1]`. The server finds a constant conditional probability `p` such that `p + p^2 + ... + p^K = L - 1`, then uses unconditional rates `[p, p^2, ..., p^K]`.
+
 ### General Speculative Parameters

 ```
@@ -6,8 +6,8 @@
 extern "C" {
 #endif

-#define RPC_PROTO_MAJOR_VERSION    5
-#define RPC_PROTO_MINOR_VERSION    1
+#define RPC_PROTO_MAJOR_VERSION    6
+#define RPC_PROTO_MINOR_VERSION    0
 #define RPC_PROTO_PATCH_VERSION    0

 #ifdef  __cplusplus
@@ -2936,12 +2936,13 @@ struct ggml_cplan ggml_graph_plan(
                        const int64_t ne10 = node->src[1]->ne[0]; // W
                        const int64_t ne11 = node->src[1]->ne[1]; // H
                        const int64_t ne12 = node->src[1]->ne[2]; // Channels In
+                        const int64_t ne13 = node->src[1]->ne[3]; // Batch

                        GGML_ASSERT(node->src[0]->type == GGML_TYPE_F16 || node->src[0]->type == GGML_TYPE_F32);
                        GGML_ASSERT(node->src[1]->type == GGML_TYPE_F32);

                        cur += ggml_type_size(node->src[0]->type) * ne00 * ne01 * ne02 * ne03;
-                        cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12;
+                        cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12 * ne13;

                    } break;
                case GGML_OP_TOP_K:
@@ -7267,18 +7267,21 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
            }
        }

-        // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh)
+        // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh), for all batches
        {
            kernel_t * const wdata = (kernel_t *) params->wdata + nk;
-            for (int i12 = 0; i12 < ne12; i12++) {
-                for (int i11 = 0; i11 < ne11; i11++) {
-                    const float * const src = (float *)((char *) src1->data + i12*nb12 + i11*nb11);
-                    kernel_t * dst_data = wdata + i11*ne10*ne12;
-                    for (int i10 = 0; i10 < ne10; i10++) {
-                        if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
-                            dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]);
-                        } else {
-                            dst_data[i10*ne12 + i12] = src[i10];
+            for (int i13 = 0; i13 < ne13; i13++) {
+                kernel_t * const wdata_b = wdata + i13*ne10*ne11*ne12;
+                for (int i12 = 0; i12 < ne12; i12++) {
+                    for (int i11 = 0; i11 < ne11; i11++) {
+                        const float * const src = (float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11);
+                        kernel_t * dst_data = wdata_b + i11*ne10*ne12;
+                        for (int i10 = 0; i10 < ne10; i10++) {
+                            if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
+                                dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]);
+                            } else {
+                                dst_data[i10*ne12 + i12] = src[i10];
+                            }
                        }
                    }
                }
@@ -7305,24 +7308,27 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
    kernel_t * const wdata_src = wdata + nk;

    for (int i2 = ip0; i2 < ip1; i2++) { // Cout
-        float * dst_data = (float *)((char *) dst->data + i2*nb2);
        kernel_t * wdata_kernel = wdata + i2*ne01*ne00*ne03;
-        for (int i11 = 0; i11 < ne11; i11++) {
-            for (int i10 = 0; i10 < ne10; i10++) {
-                const int i1n = i11*ne10*ne12 + i10*ne12;
-                for (int i01 = 0; i01 < ne01; i01++) {
-                    for (int i00 = 0; i00 < ne00; i00++) {
-                        float v = 0;
-                        if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
-                            ggml_vec_dot_f16(ne03, &v, 0,
-                                    wdata_src + i1n, 0,
-                                    wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
-                        } else {
-                            ggml_vec_dot_f32(ne03, &v, 0,
-                                    wdata_src + i1n, 0,
-                                    wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
+        for (int i3 = 0; i3 < ne3; i3++) { // batch
+            float * dst_data = (float *)((char *) dst->data + i3*nb3 + i2*nb2);
+            kernel_t * wdata_src_b = wdata_src + i3*ne10*ne11*ne12;
+            for (int i11 = 0; i11 < ne11; i11++) {
+                for (int i10 = 0; i10 < ne10; i10++) {
+                    const int i1n = i11*ne10*ne12 + i10*ne12;
+                    for (int i01 = 0; i01 < ne01; i01++) {
+                        for (int i00 = 0; i00 < ne00; i00++) {
+                            float v = 0;
+                            if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
+                                ggml_vec_dot_f16(ne03, &v, 0,
+                                        wdata_src_b + i1n, 0,
+                                        wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
+                            } else {
+                                ggml_vec_dot_f32(ne03, &v, 0,
+                                        wdata_src_b + i1n, 0,
+                                        wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
+                            }
+                            dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v;
                        }
-                        dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v;
                    }
                }
            }
@@ -8,60 +8,107 @@
 #include <algorithm>
 #include <string>
 #include <vector>
+#include <memory>
 #include <stdio.h>
 #include "htp-ops.h"
 #include "htp/matmul-ops.h"
 #include "htp/flash-attn-ops.h"
 #include "htp/unary-ops.h"
+#include "htp/allreduce-ops.h"

 struct htp_opnode {
-    ggml_tensor * node = nullptr;
+    ggml_tensor * node   { nullptr };
+    htp_op_code   opcode { HTP_OP_INVALID };
+    int32_t       kernel_params[HTP_OP_MAX_KERN_PARAMS] {0};

-    std::vector<ggml_tensor *> fused;
+    std::vector<ggml_tensor *>                fused;
+    std::vector<std::shared_ptr<ggml_tensor>> dummy;

-    htp_op_code opcode = HTP_OP_INVALID;
+    std::vector<const ggml_tensor *> inputs;
+    std::vector<const ggml_tensor *> outputs;
+    std::string                      name;

-    std::vector<ggml_tensor *> extra_dsts;
-
-    int32_t kernel_params[HTP_OP_MAX_KERN_PARAMS] = {0};
-
-    htp_opnode(ggml_tensor * node = nullptr, std::vector<ggml_tensor *> fused = {}, htp_op_code opcode = HTP_OP_INVALID, std::vector<ggml_tensor *> extra_dsts = {})
-        : node(node), fused(std::move(fused)), opcode(opcode), extra_dsts(std::move(extra_dsts)) {}
-
-    ggml_op op() const {
-        return node->op;
+    int n_active_src(const ggml_tensor * t) const {
+        if (!t) return 0;
+        for (int i = GGML_MAX_SRC - 1; i >= 0; i--) {
+            if (t->src[i]) {
+                return i + 1;
+            }
+        }
+        return 0;
    }

-    const ggml_tensor * dst() const {
-        return fused.empty() ? node : fused.back();
+    void init(ggml_tensor * node) {
+        this->node = node;
+        if (this->node) {
+            this->name = ggml_op_desc(this->node);
+
+            // Build inputs (preserving optional nullptrs)
+            int n_inputs = n_active_src(this->node);
+            this->inputs.resize(n_inputs, nullptr);
+            for (int i = 0; i < n_inputs; i++) {
+                this->inputs[i] = this->node->src[i];
+            }
+
+            // Build outputs
+            this->outputs.push_back(this->dst());
+        }
+    }
+
+    htp_opnode(htp_op_code opcode = HTP_OP_INVALID, ggml_tensor * node = nullptr) : opcode(opcode) {
+        init(node);
+    }
+
+    ggml_op             op()   const { return node->op; }
+    const ggml_tensor * src0() const { return node->src[0]; }
+    const ggml_tensor * src1() const { return node->src[1]; }
+    const ggml_tensor * dst()  const { return outputs.empty() ? node : outputs.back(); }
+
+    ggml_tensor * add_dummy(const ggml_tensor & t) {
+        dummy.push_back(std::make_shared<ggml_tensor>(t));
+        return dummy.back().get();
    }

    void add_fused(ggml_tensor * t, bool extra_dst = false) {
        fused.push_back(t);
-        if (extra_dst) {
-            extra_dsts.push_back(t);
-        }
-    }

-    std::vector<const ggml_tensor *> get_outputs() const {
-        std::vector<const ggml_tensor *> res;
-        if (extra_dsts.empty()) {
-            res.push_back(dst());
+        name += "+";
+        name += ggml_op_desc(t);
+
+        if (extra_dst) {
+            outputs.push_back(t);
        } else {
-            res.push_back(node);
-            for (const auto * x : extra_dsts) {
-                res.push_back(x);
+            outputs.clear();
+            outputs.push_back(t);
+        }
+
+        // Remove the newly fused intermediate output tensor t from inputs (if it was there)
+        inputs.erase(std::remove(inputs.begin(), inputs.end(), t), inputs.end());
+
+        // Append new inputs from t, preserving middle nullptrs
+        int n_inputs = n_active_src(t);
+        for (int i = 0; i < n_inputs; i++) {
+            const auto * src = t->src[i];
+            if (!src) {
+                inputs.push_back(nullptr);
+            } else if (src != node &&
+                       std::find(fused.begin(), fused.end(), src) == fused.end() &&
+                       std::find(inputs.begin(), inputs.end(), src) == inputs.end()) {
+                inputs.push_back(src);
            }
        }
-        return res;
    }

-    const ggml_tensor * src0() const {
-        return node->src[0];
+    const std::vector<const ggml_tensor *> & get_inputs() const {
+        return inputs;
    }

-    const ggml_tensor * src1() const {
-        return node->src[1];
+    const std::vector<const ggml_tensor *> & get_outputs() const {
+        return outputs;
+    }
+
+    std::string op_name() const {
+        return name;
    }

    bool is_empty() const {
@@ -81,75 +128,6 @@ struct htp_opnode {
    bool same_input(const htp_opnode& n) const {
        return n.src1() == this->src1();
    }
-
-    std::vector<const ggml_tensor *> get_inputs() const {
-        if (fused.empty()) {
-            int last_non_null = -1;
-            for (int i = 0; i < GGML_MAX_SRC; i++) {
-                if (node->src[i]) {
-                    last_non_null = i;
-                }
-            }
-            std::vector<const ggml_tensor *> inputs(last_non_null + 1, nullptr);
-            for (int i = 0; i <= last_non_null; i++) {
-                inputs[i] = node->src[i];
-            }
-            return inputs;
-        }
-
-        std::vector<const ggml_tensor *> inputs(GGML_MAX_SRC, nullptr);
-        std::vector<const ggml_tensor *> outputs;
-        outputs.push_back(node);
-        for (const auto * f : fused) {
-            outputs.push_back(f);
-        }
-
-        auto contains = [&](const std::vector<const ggml_tensor *> & vec, const ggml_tensor * t) {
-            for (const auto * x : vec) {
-                if (x == t) return true;
-            }
-            return false;
-        };
-
-        int count = 0;
-        auto add_input = [&](const ggml_tensor * t) {
-            if (t && !contains(outputs, t) && !contains(inputs, t)) {
-                if (count < (int)inputs.size()) {
-                    inputs[count++] = t;
-                } else {
-                    inputs.push_back(t);
-                }
-            }
-        };
-
-        for (int i = 0; i < GGML_MAX_SRC; i++) {
-            if (node->src[i]) {
-                add_input(node->src[i]);
-            }
-        }
-        for (const auto * f : fused) {
-            for (int i = 0; i < GGML_MAX_SRC; i++) {
-                if (f->src[i]) {
-                    add_input(f->src[i]);
-                }
-            }
-        }
-
-        inputs.resize(count);
-        return inputs;
-    }
-
-    std::string op_name() const {
-        if (fused.empty()) {
-            return ggml_op_desc(node);
-        }
-        std::string name = ggml_op_desc(node);
-        for (const auto * f : fused) {
-            name += "+";
-            name += ggml_op_desc(f);
-        }
-        return name;
-    }
 };

 struct htp_opformat {
@@ -337,8 +315,7 @@ struct htp_opformat {
    }
    void format_kernel_params(char * str, size_t max_size, const htp_opnode & node) {
        if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID ||
-            node.opcode == HTP_OP_MUL_MAT_QKV || node.opcode == HTP_OP_MUL_MAT_FFN ||
-            node.opcode == HTP_OP_MUL_MAT_ADD) {
+            node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ADD) {
            const auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params;
            const char * path = "unknown";
            int32_t type = kparams->kernel_type;
@@ -43,6 +43,7 @@ add_library(${HTP_LIB} SHARED
    pad-ops.c
    argsort-ops.c
    im2col-ops.c
+    allreduce-ops.c
 )

 target_compile_definitions(${HTP_LIB} PRIVATE
@@ -183,6 +183,53 @@ static void swiglu_oai_f32(const float * restrict src0,
 static const float GELU_COEF_A     = 0.044715f;
 static const float SQRT_2_OVER_PI  = 0.79788456080286535587989211986876f;

+static inline HVX_Vector hvx_vec_fast_sigmoid_f32_2it(HVX_Vector v) {
+    v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
+    v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3));
+
+    HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
+    HVX_Vector x      = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
+    HVX_Vector xx     = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x);
+
+    HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2));
+    v1            = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
+
+    HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1));
+    v2            = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx);
+    v2            = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x);
+
+    HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
+    v3            = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
+
+    HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
+    HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
+
+    // Newton-Raphson with 2 iterations
+    HVX_Vector two_sf = hvx_vec_splat_f32(2.0f);
+    HVX_Vector i_sf   = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(0x7EEEEBB3), v5);
+    HVX_Vector r_qf   = Q6_Vqf32_vmpy_VsfVsf(
+        i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
+    r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
+        r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
+    HVX_Vector res = Q6_Vsf_equals_Vqf32(r_qf);
+
+    res = Q6_Vqf32_vmpy_VsfVsf(v3, res);
+
+    return Q6_Vsf_equals_Vqf32(res);
+}
+
+static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard_2it(HVX_Vector v,
+                                                            HVX_Vector one,
+                                                            HVX_Vector max_exp,
+                                                            HVX_Vector min_exp) {
+    const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v);
+    const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp);
+
+    HVX_Vector out = hvx_vec_fast_sigmoid_f32_2it(v);
+    out            = Q6_V_vmux_QVV(pred_max, out, one);
+    return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero());
+}
+
 static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
    assert((unsigned long) dst  % 128 == 0);
    assert((unsigned long) src0 % 128 == 0);
@@ -200,20 +247,13 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest

    const HVX_Vector v_coef_a_times_sqrt = hvx_vec_splat_f32(GELU_COEF_A_TIMES_SQRT);
    const HVX_Vector v_sqrt_2_pi         = hvx_vec_splat_f32(SQRT_2_OVER_PI);
-    const HVX_Vector v_half              = hvx_vec_splat_f32(0.5f);
    const HVX_Vector v_one               = hvx_vec_splat_f32(1.0f);
-    const HVX_Vector v_two               = hvx_vec_splat_f32(2.0f);
-
-    // Hoisted fast sigmoid / inverse constants to avoid loop-internal overhead
-    const HVX_Vector v_log2f             = Q6_V_vsplat_R(FAST_SIGMOID_LOG2F);
-    const HVX_Vector v_c1                = Q6_V_vsplat_R(FAST_SIGMOID_C1);
-    const HVX_Vector v_c2                = Q6_V_vsplat_R(FAST_SIGMOID_C2);
-    const HVX_Vector v_inv_aprox         = Q6_V_vsplat_R(0x7EEEEBB3);
    const HVX_Vector v_max_exp           = hvx_vec_splat_f32(87.0f);
    const HVX_Vector v_min_exp           = hvx_vec_splat_f32(-87.0f);

    uint32_t i = 0;

+    _Pragma("unroll(4)")
    for (; i < nvec; i++) {
        HVX_Vector x = vsrc0[i];
        HVX_Vector g = vsrc1[i];
@@ -223,56 +263,13 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest
        coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi);
        HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef);

-        // y2 = 2 * inner
-        HVX_Vector y2 = hvx_vec_mul_f32_f32(inner, v_two);
+        // y2 = 2 * inner = inner + inner
+        HVX_Vector y2 = hvx_vec_add_f32_f32(inner, inner);

-        // Sigmoid guard check predicates
-        HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(v_max_exp, y2);
-        HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(y2, v_min_exp);
-
-        // Fast sigmoid approximation
-        HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(y2, v_log2f);
-        v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), v_half);
-
-        HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
-        HVX_Vector x_sig  = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
-        HVX_Vector xx_sig = Q6_Vqf32_vmpy_Vqf32Vqf32(x_sig, x_sig);
-
-        HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx_sig), v_c2);
-        v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, v_log2f);
-
-        HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x_sig), v_c1);
-        v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx_sig);
-        v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x_sig);
-
-        HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
-        v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
-
-        HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
-        HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
-
-        // Fast division (Newton-Raphson with 2 iterations)
-        HVX_Vector i_sf = Q6_Vw_vsub_VwVw(v_inv_aprox, v5);
-        HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf(
-            i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
-        r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
-            r_qf, Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
-        HVX_Vector res_inv = Q6_Vsf_equals_Vqf32(r_qf);
-
-        HVX_Vector sig2y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, res_inv));
-
-        // Sigmoid guards
-        sig2y = Q6_V_vmux_QVV(pred_max, sig2y, v_one);
-        sig2y = Q6_V_vmux_QVV(pred_min, sig2y, Q6_V_vzero());
-
-        // tanh(inner) = 2 * sigmoid(2 * inner) - 1
-        HVX_Vector tanh_val = hvx_vec_mul_f32_f32(sig2y, v_two);
-        tanh_val = hvx_vec_sub_f32_f32(tanh_val, v_one);
-
-        HVX_Vector tanh_plus_one = hvx_vec_add_f32_f32(tanh_val, v_one);
-        HVX_Vector half_x = hvx_vec_mul_f32_f32(x, v_half);
-        HVX_Vector gelu_x = hvx_vec_mul_f32_f32(half_x, tanh_plus_one);
+        // Fast sigmoid approximation (2 iterations)
+        HVX_Vector sig2y = hvx_vec_fast_sigmoid_f32_guard_2it(y2, v_one, v_max_exp, v_min_exp);

+        HVX_Vector gelu_x = hvx_vec_mul_f32_f32(x, sig2y);
        vdst[i] = hvx_vec_mul_f32_f32(gelu_x, g);
    }

@@ -285,50 +282,11 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest
        coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi);
        HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef);

-        HVX_Vector y2 = hvx_vec_mul_f32_f32(inner, v_two);
+        HVX_Vector y2 = hvx_vec_add_f32_f32(inner, inner);

-        HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(v_max_exp, y2);
-        HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(y2, v_min_exp);
-
-        HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(y2, v_log2f);
-        v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), v_half);
-
-        HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
-        HVX_Vector x_sig  = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
-        HVX_Vector xx_sig = Q6_Vqf32_vmpy_Vqf32Vqf32(x_sig, x_sig);
-
-        HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx_sig), v_c2);
-        v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, v_log2f);
-
-        HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x_sig), v_c1);
-        v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx_sig);
-        v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x_sig);
-
-        HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
-        v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
-
-        HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
-        HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
-
-        HVX_Vector i_sf = Q6_Vw_vsub_VwVw(v_inv_aprox, v5);
-        HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf(
-            i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
-        r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
-            r_qf, Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
-        HVX_Vector res_inv = Q6_Vsf_equals_Vqf32(r_qf);
-
-        HVX_Vector sig2y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, res_inv));
-
-        sig2y = Q6_V_vmux_QVV(pred_max, sig2y, v_one);
-        sig2y = Q6_V_vmux_QVV(pred_min, sig2y, Q6_V_vzero());
-
-        HVX_Vector tanh_val = hvx_vec_mul_f32_f32(sig2y, v_two);
-        tanh_val = hvx_vec_sub_f32_f32(tanh_val, v_one);
-
-        HVX_Vector tanh_plus_one = hvx_vec_add_f32_f32(tanh_val, v_one);
-        HVX_Vector half_x = hvx_vec_mul_f32_f32(x, v_half);
-        HVX_Vector gelu_x = hvx_vec_mul_f32_f32(half_x, tanh_plus_one);
+        HVX_Vector sig2y = hvx_vec_fast_sigmoid_f32_guard_2it(y2, v_one, v_max_exp, v_min_exp);

+        HVX_Vector gelu_x = hvx_vec_mul_f32_f32(x, sig2y);
        HVX_Vector res = hvx_vec_mul_f32_f32(gelu_x, g);
        hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), res);
    }
@@ -0,0 +1,398 @@
+#pragma clang diagnostic ignored "-Wunused-variable"
+#pragma clang diagnostic ignored "-Wunused-function"
+#pragma clang diagnostic ignored "-Wunused-but-set-variable"
+
+#include <HAP_farf.h>
+#include <HAP_perf.h>
+#include <stdatomic.h>
+#include <math.h>
+#include <string.h>
+
+#define GGML_COMMON_DECL_C
+#include "ggml-common.h"
+#include "htp-ctx.h"
+#include "htp-ops.h"
+#include "hvx-utils.h"
+#include "htp-tensor.h"
+#include "hex-dma.h"
+#include "hex-profile.h"
+#include "allreduce-ops.h"
+
+struct htp_allreduce_context {
+    struct htp_ops_context * octx;
+    uint32_t n_ranks;
+    uint32_t n_dsts;
+    uint32_t nelem;
+    uint32_t ne0;
+    uint32_t ne1;
+    uint32_t row_size_aligned;
+    uint32_t rank_elem_start;
+    uint32_t rank_nelem;
+    uint32_t elems_per_thread;
+    uint32_t block_elems;
+    uint32_t vtcm_size_per_thread;
+    bool     is_row_bcast;
+    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];
+    uint8_t * dst_spad_base;
+    uint8_t * res_spad_base;
+};
+
+#define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD)                                             \
+static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                       \
+    struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data;                                      \
+    struct htp_ops_context * octx = actx->octx;                                                                       \
+                                                                                                                      \
+    const uint32_t n_ranks     = actx->n_ranks;                                                                       \
+    const uint32_t n_dsts      = actx->n_dsts;                                                                        \
+    const uint32_t block_elems = actx->block_elems;                                                                   \
+                                                                                                                      \
+    const uint32_t dr  = actx->elems_per_thread;                                                                      \
+    const uint32_t ir0 = actx->rank_elem_start + dr * ith;                                                            \
+    const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem);                                     \
+    if (ir0 >= ir1) return;                                                                                           \
+                                                                                                                      \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                            \
+    dma_queue * q = octx->ctx->dma[ith];                                                                              \
+                                                                                                                      \
+    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];                                                                 \
+    for (uint32_t s = 0; s < n_ranks; s++) {                                                                          \
+        src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread);                               \
+    }                                                                                                                 \
+    uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread);                               \
+    uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + (ith * actx->vtcm_size_per_thread)) : NULL;            \
+                                                                                                                      \
+    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                          \
+    uint32_t ir_prefetch = ir0;                                                                                       \
+    int spad_idx = 0;                                                                                                 \
+                                                                                                                      \
+    for (int k = 0; k < 2 && ir_prefetch < ir1; k++) {                                                                \
+        uint32_t cur_elems = MIN(block_elems, ir1 - ir_prefetch);                                                     \
+        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                                \
+        uint8_t * d_spad = dst_spad_base + spad_idx * spad_half;                                                      \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                       \
+            uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir_prefetch * sizeof(TYPE);                           \
+            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0);                       \
+        }                                                                                                             \
+        for (uint32_t s = 0; s < n_ranks; s++) {                                                                      \
+            uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half;                                               \
+            const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE);                \
+            dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1);                       \
+        }                                                                                                             \
+        if (HAS_ADD) {                                                                                                \
+            uint8_t * r_spad = res_spad_base + spad_idx * spad_half;                                                  \
+            const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE);      \
+            dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), cur_bytes, cur_bytes, cur_bytes, 1);                       \
+        }                                                                                                             \
+        ir_prefetch += cur_elems;                                                                                     \
+        spad_idx ^= 1;                                                                                                \
+    }                                                                                                                 \
+                                                                                                                      \
+    for (uint32_t ir = ir0; ir < ir1; ) {                                                                             \
+        uint32_t cur_elems = MIN(block_elems, ir1 - ir);                                                              \
+        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                                \
+        uint8_t * d_spad = NULL;                                                                                      \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                       \
+            d_spad = (uint8_t *) dma_queue_pop(q).src;                                                                \
+        }                                                                                                             \
+        uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS];                                                                    \
+        for (uint32_t s = 0; s < n_ranks; s++) {                                                                      \
+            s_spad[s] = (uint8_t *) dma_queue_pop(q).dst;                                                             \
+        }                                                                                                             \
+        uint8_t * r_spad = HAS_ADD ? (uint8_t *) dma_queue_pop(q).dst : NULL;                                         \
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                             \
+        HVX_ADD_FN(d_spad, s_spad[0], s_spad[1], cur_elems);                                                          \
+        for (uint32_t s = 2; s < n_ranks; s++) {                                                                      \
+            HVX_ADD_FN(d_spad, d_spad, s_spad[s], cur_elems);                                                         \
+        }                                                                                                             \
+        if (HAS_ADD) {                                                                                                \
+            HVX_ADD_FN(d_spad, d_spad, r_spad, cur_elems);                                                            \
+        }                                                                                                             \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                              \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                       \
+            uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir * sizeof(TYPE);                                    \
+            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 1);                       \
+        }                                                                                                             \
+        if (ir_prefetch < ir1) {                                                                                      \
+            uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch);                                                \
+            size_t   next_bytes = next_elems * sizeof(TYPE);                                                          \
+            for (uint32_t s = 0; s < n_ranks; s++) {                                                                  \
+                const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE);           \
+                dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), next_bytes, next_bytes, next_bytes, 1);            \
+            }                                                                                                         \
+            if (HAS_ADD) {                                                                                            \
+                const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \
+                dma_queue_push(q, dma_make_ptr(r_spad, r_next), next_bytes, next_bytes, next_bytes, 1);               \
+            }                                                                                                         \
+            ir_prefetch += next_elems;                                                                                \
+        }                                                                                                             \
+        ir += cur_elems;                                                                                              \
+    }                                                                                                                 \
+    dma_queue_flush(q);                                                                                               \
+}
+
+DEFINE_ALLREDUCE_THREAD_DMA_1D(f16,     __fp16, hvx_add_f16_aaa, 0)
+DEFINE_ALLREDUCE_THREAD_DMA_1D(f32,     float,  hvx_add_f32_aaa, 0)
+DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f16, __fp16, hvx_add_f16_aaa, 1)
+DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f32, float,  hvx_add_f32_aaa, 1)
+
+#define DEFINE_ALLREDUCE_THREAD_DMA_2D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD, IS_ROW_BCAST)                                                           \
+static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                                                   \
+    struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data;                                                                  \
+    struct htp_ops_context * octx = actx->octx;                                                                                                   \
+                                                                                                                                                  \
+    const uint32_t n_ranks          = actx->n_ranks;                                                                                              \
+    const uint32_t n_dsts           = actx->n_dsts;                                                                                               \
+    const uint32_t ne0              = actx->ne0;                                                                                                  \
+    const uint32_t block_rows       = actx->block_elems;                                                                                          \
+    const uint32_t row_size_aligned = actx->row_size_aligned;                                                                                     \
+    const uint32_t row_bytes        = ne0 * sizeof(TYPE);                                                                                         \
+                                                                                                                                                  \
+    const uint32_t dr  = actx->elems_per_thread;                                                                                                  \
+    const uint32_t r0  = actx->rank_elem_start + dr * ith;                                                                                        \
+    const uint32_t r1  = MIN(r0 + dr, actx->rank_elem_start + actx->rank_nelem);                                                                  \
+    if (r0 >= r1) return;                                                                                                                         \
+                                                                                                                                                  \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                                        \
+    dma_queue * q = octx->ctx->dma[ith];                                                                                                          \
+                                                                                                                                                  \
+    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];                                                                                             \
+    for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                      \
+        src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread);                                                           \
+    }                                                                                                                                             \
+    uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread);                                                           \
+    uint8_t * res_spad_base = HAS_ADD ? (IS_ROW_BCAST ? actx->res_spad_base : (actx->res_spad_base + (ith * actx->vtcm_size_per_thread))) : NULL; \
+                                                                                                                                                  \
+    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                                                      \
+    uint32_t r_prefetch = r0;                                                                                                                     \
+    int spad_idx = 0;                                                                                                                             \
+                                                                                                                                                  \
+    for (int k = 0; k < 2 && r_prefetch < r1; k++) {                                                                                              \
+        uint32_t cur_rows = MIN(block_rows, r1 - r_prefetch);                                                                                     \
+        uint8_t * d_spad = dst_spad_base + spad_idx * spad_half;                                                                                  \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
+            uint8_t * d_ddr  = (uint8_t *) octx->dsts[d]->data + r_prefetch * octx->dsts[d]->nb[1];                                               \
+            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, 0);                                 \
+        }                                                                                                                                         \
+        for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                  \
+            uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half;                                                                           \
+            const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                      \
+            dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), row_size_aligned, octx->src[s]->nb[1], row_bytes, cur_rows);                           \
+        }                                                                                                                                         \
+        if (HAS_ADD && !IS_ROW_BCAST) {                                                                                                           \
+            uint8_t * r_spad = res_spad_base + spad_idx * spad_half;                                                                              \
+            const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1];                  \
+            dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, cur_rows);                 \
+        }                                                                                                                                         \
+        r_prefetch += cur_rows;                                                                                                                   \
+        spad_idx ^= 1;                                                                                                                            \
+    }                                                                                                                                             \
+                                                                                                                                                  \
+    for (uint32_t r = r0; r < r1; ) {                                                                                                             \
+        uint32_t cur_rows = MIN(block_rows, r1 - r);                                                                                              \
+        uint8_t * d_spad = NULL;                                                                                                                  \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
+            d_spad = (uint8_t *) dma_queue_pop(q).src;                                                                                            \
+        }                                                                                                                                         \
+        uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS];                                                                                                \
+        for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                  \
+            s_spad[s] = (uint8_t *) dma_queue_pop(q).dst;                                                                                         \
+        }                                                                                                                                         \
+        uint8_t * r_spad = (HAS_ADD && !IS_ROW_BCAST) ? (uint8_t *) dma_queue_pop(q).dst : NULL;                                                  \
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);                                                                          \
+        for (uint32_t row = 0; row < cur_rows; row++) {                                                                                           \
+            uint8_t * d_row = d_spad + row * row_size_aligned;                                                                                    \
+            const uint8_t * s0_row = s_spad[0] + row * row_size_aligned;                                                                          \
+            const uint8_t * s1_row = s_spad[1] + row * row_size_aligned;                                                                          \
+            HVX_ADD_FN(d_row, s0_row, s1_row, ne0);                                                                                               \
+            for (uint32_t s = 2; s < n_ranks; s++) {                                                                                              \
+                const uint8_t * ss_row = s_spad[s] + row * row_size_aligned;                                                                      \
+                HVX_ADD_FN(d_row, d_row, ss_row, ne0);                                                                                            \
+            }                                                                                                                                     \
+            if (HAS_ADD) {                                                                                                                        \
+                const uint8_t * res_row = IS_ROW_BCAST ? res_spad_base : (r_spad + row * row_size_aligned);                                       \
+                HVX_ADD_FN(d_row, d_row, res_row, ne0);                                                                                           \
+            }                                                                                                                                     \
+        }                                                                                                                                         \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);                                                                           \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
+            uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + r * octx->dsts[d]->nb[1];                                                         \
+            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, cur_rows);                          \
+        }                                                                                                                                         \
+        if (r_prefetch < r1) {                                                                                                                    \
+            uint32_t next_rows = MIN(block_rows, r1 - r_prefetch);                                                                                \
+            for (uint32_t s = 0; s < n_ranks; s++) {                                                                                              \
+                const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                 \
+                dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows);                  \
+            }                                                                                                                                     \
+            if (HAS_ADD && !IS_ROW_BCAST) {                                                                                                       \
+                const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1];             \
+                dma_queue_push(q, dma_make_ptr(r_spad, r_next), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, next_rows);           \
+            }                                                                                                                                     \
+            r_prefetch += next_rows;                                                                                                              \
+        }                                                                                                                                         \
+        r += cur_rows;                                                                                                                            \
+    }                                                                                                                                             \
+    dma_queue_flush(q);                                                                                                                           \
+}
+
+DEFINE_ALLREDUCE_THREAD_DMA_2D(f16,           __fp16, hvx_add_f16_aaa, 0, 0)
+DEFINE_ALLREDUCE_THREAD_DMA_2D(f32,           float,  hvx_add_f32_aaa, 0, 0)
+DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f16,       __fp16, hvx_add_f16_aaa, 1, 0)
+DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f32,       float,  hvx_add_f32_aaa, 1, 0)
+DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f16, __fp16, hvx_add_f16_aaa, 1, 1)
+DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f32, float,  hvx_add_f32_aaa, 1, 1)
+
+int op_allreduce(struct htp_ops_context * octx) {
+    const struct htp_allreduce_kernel_params * kparams = (const struct htp_allreduce_kernel_params *) octx->kernel_params;
+    const struct htp_tensor * dst = octx->dst;
+
+    const uint32_t rank    = (uint32_t) kparams->rank;
+    const uint32_t n_ranks = (uint32_t) kparams->n_ranks;
+
+    if (n_ranks < 2 || n_ranks > HTP_ALLREDUCE_MAX_RANKS || rank >= n_ranks) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    if (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    const uint32_t nelem = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
+    const uint32_t fence_seq_entry = (uint32_t) octx->op_params[0];
+    const uint32_t fence_seq_exit  = (uint32_t) octx->op_params[1];
+
+    // 1. Entry Barrier: Synchronize all ranks before reading
+    struct htp_thread_trace * tr0 = &octx->ctx->trace[0];
+    htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
+
+    const struct htp_tensor * my_sync = octx->src[n_ranks + rank];
+    atomic_uint * my_fence = (atomic_uint *) my_sync->data;
+
+    atomic_store(&my_fence[0], fence_seq_entry);
+    asm volatile ("syncht" : : : "memory");
+    Q6_dccleaninva_A((void *) my_fence);
+
+    for (uint32_t j = 0; j < n_ranks; j++) {
+        if (j == rank) continue;
+        const struct htp_tensor * peer_sync = octx->src[n_ranks + j];
+        atomic_uint * peer_fence = (atomic_uint *) peer_sync->data;
+        uint64_t spins = 0;
+        while (1) {
+            Q6_dccleaninva_A((void *) peer_fence);
+            uint32_t val = atomic_load(&peer_fence[0]);
+            if (val == fence_seq_entry || val == fence_seq_exit) {
+                break;
+            }
+            if (++spins > HTP_FENCE_TIMEOUT) {
+                FARF(ERROR, "ggml-hex: allreduce entry fence-wait TIMEOUT: rank %u waiting on %u (fence %p seq %u)\n", rank, j, peer_fence, fence_seq_entry);
+                return HTP_STATUS_INTERNAL_ERR;
+            }
+            hex_pause();
+        }
+    }
+    asm volatile ("syncht" : : : "memory");
+
+    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
+
+    // 2. Multi-threaded Reduction across assigned rank chunk
+    if (nelem > 0) {
+        const uint32_t n_threads            = (uint32_t) kparams->n_threads;
+        const uint32_t block_elems          = (uint32_t) kparams->block_elems;
+        const uint32_t elems_per_thread     = (uint32_t) kparams->elems_per_thread;
+        const uint32_t vtcm_size_per_thread = (uint32_t) kparams->vtcm_size_per_thread;
+
+        const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
+
+        struct htp_allreduce_context actx;
+        actx.octx                 = octx;
+        actx.n_ranks              = n_ranks;
+        actx.n_dsts               = (uint32_t) kparams->n_dsts ? (uint32_t) kparams->n_dsts : n_ranks;
+        actx.nelem                = nelem;
+        actx.ne0                  = (uint32_t) kparams->ne0;
+        actx.ne1                  = (uint32_t) kparams->ne1;
+        actx.row_size_aligned     = (uint32_t) kparams->row_size_aligned;
+        actx.rank_elem_start      = (uint32_t) kparams->rank_elem_start;
+        actx.rank_nelem           = (uint32_t) kparams->rank_nelem;
+        actx.elems_per_thread     = elems_per_thread;
+        actx.block_elems          = block_elems;
+        actx.vtcm_size_per_thread = vtcm_size_per_thread;
+        actx.is_row_bcast         = (kparams->is_row_bcast != 0);
+
+        work_queue_func_t reduce_fun = NULL;
+        switch (kparams->kernel_type) {
+            case HTP_ALLREDUCE_KERNEL_DMA_1D:
+                if (has_add) {
+                    reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_1d_add_f16 : allreduce_thread_dma_1d_add_f32;
+                } else {
+                    reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_1d_f16 : allreduce_thread_dma_1d_f32;
+                }
+                break;
+            case HTP_ALLREDUCE_KERNEL_DMA_2D:
+                if (has_add) {
+                    if (kparams->is_row_bcast) {
+                        reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_add_bcast_f16 : allreduce_thread_dma_2d_add_bcast_f32;
+                    } else {
+                        reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_add_f16 : allreduce_thread_dma_2d_add_f32;
+                    }
+                } else {
+                    reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_f16 : allreduce_thread_dma_2d_f32;
+                }
+                break;
+            default:
+                return HTP_STATUS_NO_SUPPORT;
+        }
+
+        uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
+        for (uint32_t s = 0; s < n_ranks; s++) {
+            actx.src_spad_base[s] = vtcm_ptr;
+            vtcm_ptr += n_threads * vtcm_size_per_thread;
+        }
+        actx.dst_spad_base = vtcm_ptr;
+        vtcm_ptr += n_threads * vtcm_size_per_thread;
+        if (has_add) {
+            actx.res_spad_base = vtcm_ptr;
+            vtcm_ptr += (actx.is_row_bcast ? 1 : n_threads) * vtcm_size_per_thread;
+        }
+
+        if (has_add && actx.is_row_bcast) {
+            const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data;
+            const uint32_t row_bytes = actx.ne0 * (dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float));
+            dma_queue * q = octx->ctx->dma[0];
+            dma_queue_push(q, dma_make_ptr(actx.res_spad_base, r_ddr), actx.row_size_aligned, 0, row_bytes, 1);
+            dma_queue_pop(q);
+        }
+
+        work_queue_run(octx->ctx->work_queue, reduce_fun, &actx, n_threads);
+    }
+
+    // 4. Exit Barrier: Synchronize all ranks after writing
+    htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
+
+    atomic_store(&my_fence[0], fence_seq_exit);
+    asm volatile ("syncht" : : : "memory");
+    Q6_dccleaninva_A((void *) my_fence);
+
+    for (uint32_t j = 0; j < n_ranks; j++) {
+        if (j == rank) continue;
+        const struct htp_tensor * peer_sync = octx->src[n_ranks + j];
+        atomic_uint * peer_fence = (atomic_uint *) peer_sync->data;
+        uint64_t spins = 0;
+        while (1) {
+            Q6_dccleaninva_A((void *) peer_fence);
+            uint32_t val = atomic_load(&peer_fence[0]);
+            if (val == fence_seq_exit) {
+                break;
+            }
+            if (++spins > HTP_FENCE_TIMEOUT) {
+                FARF(ERROR, "ggml-hex: allreduce exit fence-wait TIMEOUT: rank %u waiting on %u (fence %p seq %u)\n", rank, j, peer_fence, fence_seq_exit);
+                return HTP_STATUS_INTERNAL_ERR;
+            }
+            hex_pause();
+        }
+    }
+    asm volatile ("syncht" : : : "memory");
+
+    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
+
+    return HTP_STATUS_OK;
+}
@@ -0,0 +1,40 @@
+#ifndef ALLREDUCE_OPS_H
+#define ALLREDUCE_OPS_H
+
+#include <stdint.h>
+
+#define HTP_ALLREDUCE_MAX_RANKS 4
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+enum htp_allreduce_kernel_type {
+    HTP_ALLREDUCE_KERNEL_UNSUPPORTED = 0,
+    HTP_ALLREDUCE_KERNEL_DMA_1D,
+    HTP_ALLREDUCE_KERNEL_DMA_2D,
+};
+
+struct htp_allreduce_kernel_params {
+    int32_t rank;
+    int32_t n_ranks;
+    int32_t n_threads;
+    int32_t block_elems;          // 1D: block_elems, 2D: block_rows
+    int32_t elems_per_thread;     // 1D: nelem_per_thread, 2D: nrows_per_thread
+    int32_t vtcm_size_per_thread;
+    int32_t vtcm_size;
+    int32_t kernel_type;
+    int32_t ne0;
+    int32_t ne1;
+    int32_t row_size_aligned;
+    int32_t rank_elem_start;
+    int32_t rank_nelem;
+    int32_t n_dsts;
+    int32_t is_row_bcast;
+};
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* ALLREDUCE_OPS_H */
@@ -4,6 +4,7 @@

 #include <HAP_farf.h>
 #include <HAP_perf.h>
+#include <qurt_memory.h>

 #include <math.h>
 #include <string.h>
@@ -14,6 +15,7 @@
 #include "htp-ops.h"
 #include "htp-ops.h"
 #include "hvx-utils.h"
+#include "htp-tensor.h"

 struct htp_copy_context {
    struct htp_ops_context * octx;
@@ -78,7 +80,7 @@ static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, vo
    }                                                                                                          \
 }

-DEFINE_CPY_SAMESHAPE(f32, float, 4)
+DEFINE_CPY_SAMESHAPE(f32,  float, 4)
 DEFINE_CPY_SAMESHAPE(f16, __fp16, 2)

 #define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                                         \
@@ -179,7 +181,7 @@ static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void
    }                                                                                                          \
 }

-DEFINE_CPY_RESHAPE(f32, float, 4)
+DEFINE_CPY_RESHAPE(f32,  float, 4)
 DEFINE_CPY_RESHAPE(f16, __fp16, 2)

 static void cpy_thread_f16_f32_sameshape(unsigned int nth, unsigned int ith, void * data) {
@@ -232,6 +234,41 @@ static void cpy_thread_f32_f16_sameshape(unsigned int nth, unsigned int ith, voi
    }
 }

+static inline void cpy_dma_sametype_sameshape(
+    struct htp_ops_context * octx,
+    const struct htp_tensor * dst,
+    const struct htp_tensor * src0,
+    uint32_t elem_size,
+    uint32_t ne00, uint32_t ne01, uint32_t ne02, uint32_t ne03,
+    uint32_t nb01, uint32_t nb02, uint32_t nb03,
+    uint32_t  nb1, uint32_t  nb2, uint32_t nb3
+) {
+    const bool contiguous_outer =
+        (ne02 == 1 || (nb02 == ne01 * nb01 && nb2 == ne01 * nb1)) &&
+        (ne03 == 1 || (nb03 == ne02 * nb02 && nb3 == ne02 * nb2));
+
+    dma_queue * q = octx->ctx->dma[0];
+
+    if (contiguous_outer) {
+        dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03);
+        dma_queue_pop(q);
+        return;
+    }
+
+    for (uint32_t i03 = 0; i03 < ne03; i03++) {
+        for (uint32_t i02 = 0; i02 < ne02; i02++) {
+            uint8_t* dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
+            uint8_t* src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
+            if (!dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01)) {
+                dma_queue_flush(q);
+                dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01);
+            }
+        }
+    }
+
+    dma_queue_flush(q);
+}
+
 int op_cpy(struct htp_ops_context * octx) {
    cpy_preamble;

@@ -264,14 +301,11 @@ int op_cpy(struct htp_ops_context * octx) {

    ct.src0_nrows_per_thread = (nr + n_threads - 1) / n_threads;

-    worker_callback_t copy_fun;
+    worker_callback_t copy_fun = NULL;
+    bool use_dma = false;

    if (sametype && sameshape) {
-        if (src0->type == HTP_TYPE_F32) {
-            copy_fun = cpy_thread_f32_sameshape;
-        } else {
-            copy_fun = cpy_thread_f16_sameshape;
-        }
+        use_dma = true;
    } else if (sameshape) {
        /**/ if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32)
            copy_fun = cpy_thread_f16_f32_sameshape;
@@ -289,7 +323,28 @@ int op_cpy(struct htp_ops_context * octx) {
        return HTP_STATUS_NO_SUPPORT;
    }

-    worker_pool_run_func(octx->ctx->worker_pool, copy_fun, &ct, n_threads);
+    if (use_dma) {
+        cpy_dma_sametype_sameshape(octx, dst, src0, ct.src0_type_size, ne00, ne01, ne02, ne03, nb01, nb02, nb03, nb1, nb2, nb3);
+    } else {
+        worker_pool_run_func(octx->ctx->worker_pool, copy_fun, &ct, n_threads);
+    }
+
+    const struct htp_tensor *sync = octx->src[1];
+    if (sync) {
+        if (!use_dma) {
+            // htp_tensor_flush_all(octx->ctx, octx->dsts, 1);
+            qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
+        }
+
+        atomic_uint * sync_fence = (atomic_uint *) sync->data;
+        const uint32_t seq = (uint32_t) octx->op_params[0];
+
+        atomic_store(&sync_fence[0], seq);
+        asm volatile ("syncht" : : : "memory");
+        Q6_dccleaninva_A((void *) sync_fence);
+
+        FARF(HIGH, "ggml-hex: sync-release : fence %p seq %u\n", sync_fence, seq);
+    }

    return HTP_STATUS_OK;
 }
@@ -244,17 +244,18 @@ static inline dma_ptr dma_queue_pop(dma_queue * q) {
        return dptr;
    }

-    dma_descriptor_2d * desc = &r->desc[r->pop_idx];
+    dptr = r->dptr[r->pop_idx];
+
+    volatile dma_descriptor_2d * desc = &r->desc[r->pop_idx];

    // Wait for desc to complete
    if (!desc->done) {
+        // FARF(ALWAYS, "dma-poll: idx %u dst %p src %p", r->pop_idx, dptr.dst, dptr.src);
        while (!desc->done) {
            dmpoll();
        }
    }

-    dptr = r->dptr[r->pop_idx];
-
    htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);

    r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
@@ -30,6 +30,8 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
+#include "hvx-quant.h"

 #include "flash-attn-ops.h"
 #include "hvx-fa-kernels.h"
@@ -85,12 +87,17 @@ struct htp_fa_context {
    uint8_t * spad_m;
    uint8_t * spad_a;

+    const struct htp_tensor * k;
+    const struct htp_tensor * v;
+
    uint64_t t_start;
 };

 struct hmx_fa_context {
    const struct htp_ops_context * octx;
    const struct htp_tensor *      sinks;  // attention sinks (src[4]), NULL if absent
+    const struct htp_tensor *      k;
+    const struct htp_tensor *      v;
    bool         pipeline;  // true when n_kv_blocks >= FA_MIN_KV_BLOCKS && n_threads >= 2
    uint32_t     n_threads;

@@ -214,8 +221,8 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
    const uint32_t DV = nev0;

    const size_t size_q_row = DK * ((q->type == HTP_TYPE_F32) ? 4 : 2);
-    const size_t size_k_row = DK * sizeof(__fp16);
-    const size_t size_v_row = DV * sizeof(__fp16);
+    const size_t size_k_row = htp_tensor_get_row_size(k->type, DK);
+    const size_t size_v_row = htp_tensor_get_row_size(v->type, DV);

    // Scratchpad buffers for Q, K, V, Mask, and VKQ32 accumulator
    uint8_t * spad_q = factx->spad_q + factx->size_q_block * ith;
@@ -364,6 +371,23 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
            uint8_t * v_base = dma_queue_pop(dma).dst; // V
            __fp16  * m_base = mask ? dma_queue_pop(dma).dst : NULL; // M

+            if (factx->k->type == HTP_TYPE_Q8_0) {
+                htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir);
+                for (uint32_t r = 0; r < current_block_size; ++r) {
+                    __fp16 * row_k = (__fp16 *)(k_base + r * factx->size_k_row_padded);
+                    hvx_dequantize_row_q8_0_f16(row_k, row_k, DK);
+                }
+                htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir);
+            }
+            if (factx->v->type == HTP_TYPE_Q8_0) {
+                htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir);
+                for (uint32_t r = 0; r < current_block_size; ++r) {
+                    __fp16 * row_v = (__fp16 *)(v_base + r * factx->size_v_row_padded);
+                    hvx_dequantize_row_q8_0_f16(row_v, row_v, DV);
+                }
+                htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir);
+            }
+
            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_QK, ir);

            // Inner loop processing the block from VTCM
@@ -625,6 +649,12 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)

    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
+    if (factx->k->type == HTP_TYPE_Q8_0) {
+        for (uint32_t r = start; r < end; ++r) {
+            __fp16 * row_k = (__fp16 *)((char *)args->curr_k + r * args->src_stride * sizeof(__fp16));
+            hvx_dequantize_row_q8_0_f16(row_k, row_k, factx->DK);
+        }
+    }
    hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK,
                             args->src_stride, start, end);
    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
@@ -673,6 +703,12 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data)

    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
+    if (factx->v->type == HTP_TYPE_Q8_0) {
+        for (uint32_t r = start; r < end; ++r) {
+            __fp16 * row_v = (__fp16 *)((char *)args->v_src + r * args->src_stride * sizeof(__fp16));
+            hvx_dequantize_row_q8_0_f16(row_v, row_v, factx->DV);
+        }
+    }
    hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV,
                             args->src_stride, (uint32_t) args->n_col_tiles, start, end);
    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
@@ -1809,6 +1845,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
    memset(&factx, 0, sizeof(factx));
    factx.octx           = octx;
    factx.sinks          = octx->src[4];  // NULL if this op has no attention sinks
+    factx.k              = k;
+    factx.v              = v;
    factx.n_threads      = kparams->n_threads;
    factx.DK             = DK;
    factx.DV             = DV;
@@ -1853,10 +1891,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
    // ======== VTCM allocation (GQA-aware) ========
    // K/V row sizes drive the DMA descriptors (not the VTCM layout) and are used
    // throughout the KV loop below.
-    const size_t size_k_row        = DK * sizeof(__fp16);
-    const size_t size_v_row        = DV * sizeof(__fp16);
-    const size_t size_k_row_padded = hex_round_up(size_k_row, 128);
-    const size_t size_v_row_padded = hex_round_up(size_v_row, 128);
+    const size_t size_k_row        = htp_tensor_get_row_size(k->type, DK);
+    const size_t size_v_row        = htp_tensor_get_row_size(v->type, DV);
+    const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128);
+    const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128);

    // Build the VTCM layout once (shared with the host estimator) and place every
    // scratch buffer at its computed offset.
@@ -2348,7 +2386,9 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
    const struct htp_tensor * dst  = octx->dst;

    // Check support
-    if ((q->type != HTP_TYPE_F16 && q->type != HTP_TYPE_F32) || k->type != HTP_TYPE_F16 || v->type != HTP_TYPE_F16) {
+    if ((q->type != HTP_TYPE_F16 && q->type != HTP_TYPE_F32) ||
+        (k->type != HTP_TYPE_F16 && k->type != HTP_TYPE_Q8_0) ||
+        (v->type != HTP_TYPE_F16 && v->type != HTP_TYPE_Q8_0)) {
        return HTP_STATUS_NO_SUPPORT;
    }

@@ -2364,6 +2404,8 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {

    struct htp_fa_context factx;
    factx.octx = octx;
+    factx.k = k;
+    factx.v = v;

    factx.t_start = HAP_perf_get_qtimer_count();

@@ -12,18 +12,17 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
+#include "htp-tensor.h"
 #include "hvx-utils.h"
+#include "hvx-quant.h"
+#include "get-rows-ops.h"
+#include "work-queue.h"

 struct get_rows_context {
    struct htp_ops_context * octx;
-    uint32_t tasks_per_thread;
-    uint32_t total_tasks;
-    uint32_t chunks_per_row;
-    uint32_t chunk_size;
-    struct fastdiv_values get_rows_div_ne10;
-    struct fastdiv_values get_rows_div_ne10_ne11;
-    struct fastdiv_values get_rows_div_chunks_per_row;
+    const struct htp_get_rows_kernel_params * kparams;
+    struct htp_get_rows_vtcm_layout vtcm_layout;
+    uint8_t * vtcm_base;
 };

 #define get_rows_preamble \
@@ -56,102 +55,161 @@ struct get_rows_context {
                                               \
    const uint32_t nr = ne10 * ne11 * ne12;

-static void get_rows_thread_f32_f32_dma(unsigned int nth, unsigned int ith, void *data) {
-    struct get_rows_context * grctx = (struct get_rows_context *)data;
-    struct htp_ops_context * octx = grctx->octx;
-    get_rows_preamble;
-
-    uint64_t qt = HAP_perf_get_qtimer_count();
-
-    const uint32_t dr  = grctx->tasks_per_thread;
-    const uint32_t ir0 = dr * ith;
-    if (ir0 >= grctx->total_tasks) {
-        return;
-    }
-    const uint32_t ir1 = MIN(ir0 + dr, grctx->total_tasks);
-
-    const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);
-
-    dma_queue * dma_queue = octx->ctx->dma[ith];
-    for (uint32_t i = ir0; i < ir1; ++i) {
-        const uint32_t i12 = fastdiv(i, &grctx->get_rows_div_ne10_ne11);
-        const uint32_t rem = i - i12 * ne11 * ne10;
-        const uint32_t i11 = fastdiv(rem, &grctx->get_rows_div_ne10);
-        const uint32_t i10 = rem - i11 * ne10;
-
-        const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12;
-        uint32_t i01 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr;
-
-        if (i01 >= ne01) {
-            continue;
-        }
-
-        const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i11*nb02 + i12*nb03;
-        const uintptr_t dst_ptr  = octx->dst->data    + i10*nb1  + i11*nb2  + i12*nb3;
-
-        while (!dma_queue_push(dma_queue, dma_make_ptr((void *)dst_ptr, (const void *)src0_ptr), nb1, nb01, ne00 * sizeof(float), 1)) {
-            dma_queue_pop(dma_queue);
-        }
-    }
-    dma_queue_flush(dma_queue);
-
-    qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt);
-    FARF(HIGH, "get-rows-f32-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth,
-         ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt);
+#define GET_ROWS_THREAD_ST_FN(IDX_TYPE)                                                                                \
+static void get_rows_thread_st_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) {                            \
+    struct get_rows_context * grctx = (struct get_rows_context *)data;                                                 \
+    struct htp_ops_context * octx = grctx->octx;                                                                       \
+    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;                                                \
+    get_rows_preamble;                                                                                                 \
+    const uint32_t dr  = kparams->tasks_per_thread;                                                                    \
+    const uint32_t ir0 = dr * ith;                                                                                     \
+    if (ir0 >= kparams->total_tasks) {                                                                                 \
+        return;                                                                                                        \
+    }                                                                                                                  \
+    const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks);                                                          \
+    const uint32_t row_size_bytes = htp_tensor_get_row_size(octx->src[0]->type, ne00);                                 \
+    dma_queue * dma_queue = octx->ctx->dma[ith];                                                                       \
+    for (uint32_t i = ir0; i < ir1; ++i) {                                                                             \
+        const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11);                                                      \
+        const uint32_t rem = i - i12 * ne11 * ne10;                                                                    \
+        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);                                                         \
+        const uint32_t i10 = rem - i11 * ne10;                                                                         \
+        const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12);           \
+        const uint32_t i01 = (uint32_t)*src1_ptr;                                                                      \
+        assert(i01 < ne01);                                                                                            \
+        const uint32_t q02 = fastdiv(i11, &kparams->div_ne02);                                                         \
+        const uint32_t i02 = i11 - q02 * ne02;                                                                         \
+        const uint32_t q03 = fastdiv(i12, &kparams->div_ne03);                                                         \
+        const uint32_t i03 = i12 - q03 * ne03;                                                                         \
+        const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03;                                \
+        const uintptr_t dst_ptr  = octx->dst->data    + i10*nb1  + i11*nb2  + i12*nb3;                                 \
+        while (!dma_queue_push(dma_queue, dma_make_ptr((void *)dst_ptr, (const void *)src0_ptr), nb1, nb01,            \
+                               row_size_bytes, 1)) {                                                                   \
+            dma_queue_pop(dma_queue);                                                                                  \
+        }                                                                                                              \
+    }                                                                                                                  \
+    dma_queue_flush(dma_queue);                                                                                        \
 }

-static void get_rows_thread_f32_f32_hvx(unsigned int nth, unsigned int ith, void *data) {
-    struct get_rows_context * grctx = (struct get_rows_context *)data;
-    struct htp_ops_context * octx = grctx->octx;
-    get_rows_preamble;
+GET_ROWS_THREAD_ST_FN(int32_t)
+GET_ROWS_THREAD_ST_FN(int64_t)

-    uint64_t qt = HAP_perf_get_qtimer_count();
-
-    const uint32_t dr  = grctx->tasks_per_thread;
-    const uint32_t ir0 = dr * ith;
-    if (ir0 >= grctx->total_tasks) {
-        return;
-    }
-    const uint32_t ir1 = MIN(ir0 + dr, grctx->total_tasks);
-
-    const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);
-
-    const uint32_t chunks_per_row = grctx->chunks_per_row;
-    const uint32_t chunk_size     = grctx->chunk_size;
-    for (uint32_t i = ir0; i < ir1; ++i) {
-        const uint32_t row_idx   = fastdiv(i, &grctx->get_rows_div_chunks_per_row);
-        const uint32_t chunk_idx = i - row_idx * chunks_per_row;
-
-        const uint32_t i12 = fastdiv(row_idx, &grctx->get_rows_div_ne10_ne11);
-        const uint32_t rem = row_idx - i12 * ne11 * ne10;
-        const uint32_t i11 = fastdiv(rem, &grctx->get_rows_div_ne10);
-        const uint32_t i10 = rem - i11 * ne10;
-
-        const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12;
-        uint32_t i01 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr;
-
-        if (i01 >= ne01) {
-            continue;
-        }
-
-        const uint32_t offset = chunk_idx * chunk_size;
-        if (offset < ne00) {
-            const uint32_t copy_size = MIN(chunk_size, ne00 - offset);
-            const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i11*nb02 + i12*nb03 + offset * sizeof(float);
-            const uintptr_t dst_ptr  = octx->dst->data    + i10*nb1  + i11*nb2  + i12*nb3  + offset * sizeof(float);
-            hvx_copy_f32_uu((uint8_t *)dst_ptr, (const uint8_t *)src0_ptr, copy_size);
-        }
-    }
-
-    qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt);
-    FARF(HIGH, "get-rows-f32-f32-hvx %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth,
-         ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt);
+#define GET_ROWS_THREAD_DT_FN(TYPE_NAME, SRC0_SIZE_EXPR, IDX_TYPE, COMPUTE_EXPR)                                       \
+static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) {                 \
+    struct get_rows_context * grctx = (struct get_rows_context *)data;                                                 \
+    struct htp_ops_context * octx = grctx->octx;                                                                       \
+    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;                                                \
+    get_rows_preamble;                                                                                                 \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                             \
+    const uint32_t dr  = kparams->tasks_per_thread;                                                                    \
+    const uint32_t ir0 = dr * ith;                                                                                     \
+    if (ir0 >= kparams->total_tasks) {                                                                                 \
+        return;                                                                                                        \
+    }                                                                                                                  \
+    const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks);                                                          \
+    const uint32_t chunks_per_row = kparams->chunks_per_row;                                                           \
+    const uint32_t chunk_size     = kparams->chunk_size;                                                               \
+    dma_queue * dma_queue = octx->ctx->dma[ith];                                                                       \
+    const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout;                                         \
+    uint8_t * vtcm_src0 = grctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread;         \
+    uint8_t * vtcm_dst  = grctx->vtcm_base + vtcm_layout->off_dst  + ith * vtcm_layout->dst_bytes_per_thread;          \
+    for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, spad_idx++) {                      \
+        const uint32_t i = ir0 + step;                                                                                 \
+        const uint32_t row_idx   = fastdiv(i, &kparams->div_chunks_per_row);                                           \
+        const uint32_t chunk_idx = i - row_idx * chunks_per_row;                                                       \
+        const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11);                                                \
+        const uint32_t rem = row_idx - i12 * ne11 * ne10;                                                              \
+        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);                                                         \
+        const uint32_t i10 = rem - i11 * ne10;                                                                         \
+        const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12);           \
+        const uint32_t i01 = (uint32_t)*src1_ptr;                                                                      \
+        assert(i01 < ne01);                                                                                            \
+        const uint32_t q02 = fastdiv(i11, &kparams->div_ne02);                                                         \
+        const uint32_t i02 = i11 - q02 * ne02;                                                                         \
+        const uint32_t q03 = fastdiv(i12, &kparams->div_ne03);                                                         \
+        const uint32_t i03 = i12 - q03 * ne03;                                                                         \
+        const uint32_t offset = chunk_idx * chunk_size;                                                                \
+        const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0;                               \
+        const uint32_t cur_src0_bytes = SRC0_SIZE_EXPR(cur_elems);                                                     \
+        const uint32_t cur_dst_bytes  = cur_elems * sizeof(float);                                                     \
+        const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03 + SRC0_SIZE_EXPR(offset);       \
+        dma_queue_push(dma_queue,                                                                                      \
+                       dma_make_ptr((void *)(uintptr_t)octx->dst->data,                                                \
+                                    vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size),                            \
+                       cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 0);                              \
+        dma_queue_push(dma_queue,                                                                                      \
+                       dma_make_ptr((void *)(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size),                 \
+                                    (const void *)src0_ptr),                                                           \
+                       vtcm_layout->src0_spad_half_size, cur_src0_bytes, cur_src0_bytes, 1);                           \
+    }                                                                                                                  \
+    for (uint32_t step = 0; step < ir1 - ir0; ++step) {                                                                \
+        const uint32_t i = ir0 + step;                                                                                 \
+        void * dst_spad = (void *) dma_queue_pop(dma_queue).src;                                                       \
+        void * src_spad = (void *) dma_queue_pop(dma_queue).dst;                                                       \
+        const uint32_t row_idx   = fastdiv(i, &kparams->div_chunks_per_row);                                           \
+        const uint32_t chunk_idx = i - row_idx * chunks_per_row;                                                       \
+        const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11);                                                \
+        const uint32_t rem = row_idx - i12 * ne11 * ne10;                                                              \
+        const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);                                                         \
+        const uint32_t i10 = rem - i11 * ne10;                                                                         \
+        const uint32_t offset = chunk_idx * chunk_size;                                                                \
+        const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0;                               \
+        const uint32_t cur_dst_bytes  = cur_elems * sizeof(float);                                                     \
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, i);                                                          \
+        COMPUTE_EXPR;                                                                                                  \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, i);                                                           \
+        const uintptr_t dst_ptr  = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3 + offset * sizeof(float);             \
+        dma_queue_push(dma_queue,                                                                                      \
+                       dma_make_ptr((void *)dst_ptr, (const void *)dst_spad),                                          \
+                       cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 1);                              \
+        const uint32_t next_step = step + 2;                                                                           \
+        if (next_step < ir1 - ir0) {                                                                                   \
+            const uint32_t pi = ir0 + next_step;                                                                       \
+            const uint32_t prow_idx   = fastdiv(pi, &kparams->div_chunks_per_row);                                     \
+            const uint32_t pchunk_idx = pi - prow_idx * chunks_per_row;                                                \
+            const uint32_t pi12 = fastdiv(prow_idx, &kparams->div_ne10_ne11);                                          \
+            const uint32_t prem = prow_idx - pi12 * ne11 * ne10;                                                       \
+            const uint32_t pi11 = fastdiv(prem, &kparams->div_ne10);                                                   \
+            const uint32_t pi10 = prem - pi11 * ne10;                                                                  \
+            const IDX_TYPE * psrc1_ptr = (const IDX_TYPE *)(octx->src[1]->data + pi10*nb10 + pi11*nb11 + pi12*nb12);   \
+            const uint32_t pi01 = (uint32_t)*psrc1_ptr;                                                                \
+            assert(pi01 < ne01);                                                                                       \
+            const uint32_t pq02 = fastdiv(pi11, &kparams->div_ne02);                                                   \
+            const uint32_t pi02 = pi11 - pq02 * ne02;                                                                  \
+            const uint32_t pq03 = fastdiv(pi12, &kparams->div_ne03);                                                   \
+            const uint32_t pi03 = pi12 - pq03 * ne03;                                                                  \
+            const uint32_t poffset = pchunk_idx * chunk_size;                                                          \
+            const uint32_t pcur_elems = (poffset < ne00) ? MIN(chunk_size, ne00 - poffset) : 0;                        \
+            const uint32_t pcur_src0_bytes = SRC0_SIZE_EXPR(pcur_elems);                                               \
+            const uintptr_t psrc0_ptr =                                                                                \
+                octx->src[0]->data + pi01*nb01 + pi02*nb02 + pi03*nb03 + SRC0_SIZE_EXPR(poffset);                      \
+            dma_queue_push(dma_queue,                                                                                  \
+                           dma_make_ptr((void *)src_spad, (const void *)psrc0_ptr),                                    \
+                           vtcm_layout->src0_spad_half_size, pcur_src0_bytes, pcur_src0_bytes, 1);                     \
+        }                                                                                                              \
+    }                                                                                                                  \
+    dma_queue_flush(dma_queue);                                                                                        \
 }

+#define F32_BYTES(n)  ((n) * sizeof(float))
+#define F16_BYTES(n)  ((n) * sizeof(__fp16))
+#define Q8_0_BYTES(n) (((n) / 32) * sizeof(block_q8_0))
+
+GET_ROWS_THREAD_DT_FN(f32,  F32_BYTES,  int32_t, { if (cur_elems > 0) hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, cur_elems); })
+GET_ROWS_THREAD_DT_FN(f32,  F32_BYTES,  int64_t, { if (cur_elems > 0) hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, cur_elems); })
+
+GET_ROWS_THREAD_DT_FN(f16,  F16_BYTES,  int32_t, { hvx_dequantize_row_f16_f32((float *)dst_spad, src_spad, ne00); })
+GET_ROWS_THREAD_DT_FN(f16,  F16_BYTES,  int64_t, { hvx_dequantize_row_f16_f32((float *)dst_spad, src_spad, ne00); })
+
+GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int32_t, { hvx_dequantize_row_q8_0_f32((float *)dst_spad, src_spad, ne00); })
+GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int64_t, { hvx_dequantize_row_q8_0_f32((float *)dst_spad, src_spad, ne00); })
+
 int op_get_rows(struct htp_ops_context * octx) {
-    get_rows_preamble;
+    const struct htp_get_rows_kernel_params * kparams = (const struct htp_get_rows_kernel_params *) octx->kernel_params;

-    if (octx->src[0]->type != HTP_TYPE_F32) {
+    if (octx->src[0]->type != HTP_TYPE_F32 &&
+        octx->src[0]->type != HTP_TYPE_F16 &&
+        octx->src[0]->type != HTP_TYPE_Q8_0) {
        return HTP_STATUS_NO_SUPPORT;
    }

@@ -167,52 +225,28 @@ int op_get_rows(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

-    const uint32_t nb00 = octx->src[0]->nb[0];
-    const uint32_t nb0  = octx->dst->nb[0];
-
-    const bool can_use_dma = (nb00 == sizeof(float)) && (nb0 == sizeof(float));
-    const bool use_dma = can_use_dma && (ne00 >= 2048);
-
    struct get_rows_context grctx;
    grctx.octx = octx;
-    grctx.get_rows_div_ne10      = init_fastdiv_values(octx->src[1]->ne[0]);
-    grctx.get_rows_div_ne10_ne11 = init_fastdiv_values(octx->src[1]->ne[0] * octx->src[1]->ne[1]);
+    grctx.kparams = kparams;
+    grctx.vtcm_base = (uint8_t *)octx->ctx->vtcm_base;

-    if (use_dma) {
-        grctx.chunks_per_row = 1;
-        grctx.chunk_size = ne00;
-        grctx.total_tasks = nr;
-        grctx.get_rows_div_chunks_per_row = init_fastdiv_values(1);
+    const uint32_t ne00 = octx->src[0]->ne[0];
+    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, octx->src[0]->type, ne00, kparams->n_threads);

-        const uint32_t n_threads = MIN(nr, octx->n_threads);
-        grctx.tasks_per_thread = (nr + n_threads - 1) / n_threads;
+    const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);

-        worker_pool_run_func(octx->ctx->worker_pool, get_rows_thread_f32_f32_dma, &grctx, n_threads);
+    work_queue_func_t q_func = NULL;
+    if (kparams->use_dma) {
+        q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_st_int32_t : get_rows_thread_st_int64_t);
    } else {
-        uint32_t chunks_per_row = 1;
-        uint32_t chunk_size = ne00;
-        uint32_t total_tasks = nr;
-
-        if (nr < octx->n_threads) {
-            const uint32_t min_chunk_size = 1024;
-            uint32_t max_chunks = ne00 / min_chunk_size;
-            if (max_chunks == 0) {
-                max_chunks = 1;
-            }
-            chunks_per_row = MIN((octx->n_threads + nr - 1) / nr, max_chunks);
-            chunk_size = (ne00 + chunks_per_row - 1) / chunks_per_row;
-            total_tasks = nr * chunks_per_row;
+        switch (octx->src[0]->type) {
+            case HTP_TYPE_F32:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f32_int32_t  : get_rows_thread_f32_int64_t);  break;
+            case HTP_TYPE_F16:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f16_int32_t  : get_rows_thread_f16_int64_t);  break;
+            case HTP_TYPE_Q8_0: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_q8_0_int32_t : get_rows_thread_q8_0_int64_t); break;
+            default:            return HTP_STATUS_NO_SUPPORT;
        }
-
-        grctx.chunks_per_row = chunks_per_row;
-        grctx.chunk_size = chunk_size;
-        grctx.total_tasks = total_tasks;
-        grctx.get_rows_div_chunks_per_row = init_fastdiv_values(chunks_per_row);
-
-        const uint32_t n_threads = MIN(total_tasks, octx->n_threads);
-        grctx.tasks_per_thread = (total_tasks + n_threads - 1) / n_threads;
-
-        worker_pool_run_func(octx->ctx->worker_pool, get_rows_thread_f32_f32_hvx, &grctx, n_threads);
    }
+
+    work_queue_run(octx->ctx->work_queue, q_func, &grctx, kparams->n_threads);
    return HTP_STATUS_OK;
 }
@@ -0,0 +1,77 @@
+#ifndef HTP_GET_ROWS_OPS_H
+#define HTP_GET_ROWS_OPS_H
+
+#include "hex-fastdiv.h"
+
+struct htp_get_rows_kernel_params {
+    int32_t  n_threads;
+    int32_t  use_dma;
+    int32_t  chunks_per_row;
+    int32_t  chunk_size;
+    int32_t  total_tasks;
+    int32_t  tasks_per_thread;
+    int32_t  vtcm_size;
+
+    // Fastdiv helpers
+    struct fastdiv_values div_ne10;
+    struct fastdiv_values div_ne10_ne11;
+    struct fastdiv_values div_chunks_per_row;
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_ne03;
+};
+
+struct htp_get_rows_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_dst;
+
+    size_t src0_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+
+    size_t src0_spad_half_size;
+    size_t dst_spad_half_size;
+};
+
+static inline void htp_get_rows_vtcm_layout_build(
+    struct htp_get_rows_vtcm_layout * vtcm_layout,
+    int type,
+    uint32_t ne00,
+    uint32_t n_threads) {
+
+    uint32_t src0_row_size = 0;
+    switch (type) {
+        case 0: // HTP_TYPE_F32
+            src0_row_size = ne00 * 4;
+            break;
+        case 1: // HTP_TYPE_F16
+            src0_row_size = ne00 * 2;
+            break;
+        case 8: // HTP_TYPE_Q8_0
+            src0_row_size = (ne00 / 32) * 34;
+            break;
+        default:
+            src0_row_size = 0;
+            break;
+    }
+
+    size_t src0_row_size_aligned = (src0_row_size + 255) & ~255;
+    size_t dst_row_size_aligned  = (ne00 * sizeof(float) + 255) & ~255;
+
+    vtcm_layout->src0_spad_half_size = src0_row_size_aligned;
+    vtcm_layout->dst_spad_half_size  = dst_row_size_aligned;
+
+    vtcm_layout->src0_bytes_per_thread = src0_row_size_aligned * 2;
+    vtcm_layout->dst_bytes_per_thread  = dst_row_size_aligned * 2;
+
+    vtcm_layout->off_src0 = 0;
+    vtcm_layout->off_dst  = vtcm_layout->off_src0 + vtcm_layout->src0_bytes_per_thread * n_threads;
+    vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads;
+}
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_get_rows_kernel_params) <= 128, "htp_get_rows_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_get_rows_kernel_params) <= 128, "htp_get_rows_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_GET_ROWS_OPS_H
@@ -39,17 +39,22 @@ static inline void hex_l2fetch_block(const void * addr, size_t size) {

 #define HEX_L2_LINE_SIZE           128
 #define HEX_L2_BLOCK_SIZE          (HEX_L2_LINE_SIZE * 4) // flush granularity (lines per loop iteration)
+#define HEX_L2_FLUSH_IL_THRESHOLD  1024                   // inline flush threshold
 #define HEX_L2_FLUSH_WQ_THRESHOLD  (4 * 1024)
 #define HEX_L2_FLUSH_ALL_THRESHOLD (4 * 1024 * 1024)

 static inline void hex_l2flush(void * addr, size_t size) {
    const uint32_t s = ((uint32_t) addr) & ~(HEX_L2_LINE_SIZE - 1);
    const uint32_t e = (((uint32_t) addr) + size + HEX_L2_LINE_SIZE - 1) & ~(HEX_L2_LINE_SIZE - 1);
-    for (uint32_t i = s; i < e; i += HEX_L2_BLOCK_SIZE) {
-        Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 0);
-        Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 1);
-        Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 2);
-        Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 3);
+    const uint32_t eb = s + ((e - s) & ~(HEX_L2_BLOCK_SIZE - 1));
+    for (uint32_t i = s; i < eb; i += HEX_L2_BLOCK_SIZE) {
+        Q6_dccleaninva_A((void *) (i + HEX_L2_LINE_SIZE * 0));
+        Q6_dccleaninva_A((void *) (i + HEX_L2_LINE_SIZE * 1));
+        Q6_dccleaninva_A((void *) (i + HEX_L2_LINE_SIZE * 2));
+        Q6_dccleaninva_A((void *) (i + HEX_L2_LINE_SIZE * 3));
+    }
+    for (uint32_t i = eb; i < e; i += HEX_L2_LINE_SIZE) {
+        Q6_dccleaninva_A((void *) i);
    }
 }

@@ -117,8 +117,7 @@ struct htp_context {

 int op_matmul(struct htp_ops_context * octx);
 int op_matmul_id(struct htp_ops_context * octx);
-int op_matmul_qkv(struct htp_ops_context * octx);
-int op_matmul_ffn(struct htp_ops_context * octx);
+int op_matmul_nx(struct htp_ops_context * octx);
 int op_binary(struct htp_ops_context * octx);
 int op_unary(struct htp_ops_context * octx);
 int op_sum_rows(struct htp_ops_context * octx);
@@ -141,5 +140,6 @@ int op_solve_tri(struct htp_ops_context * octx);
 int op_gated_delta_net(struct htp_ops_context * octx);
 int op_pad(struct htp_ops_context * octx);
 int op_im2col(struct htp_ops_context * octx);
+int op_allreduce(struct htp_ops_context * octx);

 #endif /* HTP_CTX_H */
@@ -43,13 +43,6 @@ enum htp_data_type {



-// Mask to enable various stages of the Ops.
-// Used for debugging and profiling.
-enum htp_op_stage {
-    HTP_OPSTAGE_QUEUE    = (1 << 0),  // Enable Queueing (ie calls into NPU)
-    HTP_OPSTAGE_COMPUTE  = (1 << 1),  // Enable Compute
-};
-
 // Do not reorder first 4 (used as an index)
 enum htp_op_code {
    HTP_OP_MUL = 0,
@@ -58,8 +51,7 @@ enum htp_op_code {
    HTP_OP_DIV = 3,
    HTP_OP_MUL_MAT,
    HTP_OP_MUL_MAT_ID,
-    HTP_OP_MUL_MAT_QKV,
-    HTP_OP_MUL_MAT_FFN,
+    HTP_OP_MUL_MAT_NX,
    HTP_OP_MUL_MAT_ADD,
    HTP_OP_RMS_NORM,
    HTP_OP_RMS_NORM_MUL,
@@ -70,6 +62,8 @@ enum htp_op_code {
    HTP_OP_UNARY_NEG,
    HTP_OP_UNARY_SOFTPLUS,
    HTP_OP_UNARY_TANH,
+    HTP_OP_UNARY_ABS,
+    HTP_OP_UNARY_LOG,
    HTP_OP_GLU_SWIGLU,
    HTP_OP_GLU_SWIGLU_OAI,
    HTP_OP_GLU_GEGLU,
@@ -99,12 +93,15 @@ enum htp_op_code {
    HTP_OP_CONCAT,
    HTP_OP_CLAMP,
    HTP_OP_IM2COL,
+    HTP_OP_FENCE,
+    HTP_OP_ALLREDUCE,
+    HTP_OP_ALLREDUCE_ADD,

    HTP_OP_INVALID
 };

 #define HTP_OP_MAX_DIMS    4    // aka GGML_MAX_DIMS
-#define HTP_OP_MAX_INPUTS  6    // aka GGML_MAX_SRCS
+#define HTP_OP_MAX_INPUTS  10   // aka GGML_MAX_SRCS
 #define HTP_OP_MAX_OUTPUTS 4
 #define HTP_OP_MAX_PARAMS  16   // aka GGML_MAX_OP_PARAMS
 #define HTP_OP_MAX_KERN_PARAMS 32
@@ -112,13 +109,16 @@ enum htp_op_code {
 #define HTP_OP_MAX_BUFS    16
 #define HTP_OP_MAX_TENSORS 8192 // must stay under 64K (uint16)

+#define HTP_FENCE_TIMEOUT  (1000000000ULL)
+
 #define HTP_OP_MAX_VMEM_DEFAULT (3355443200u)

 #define HTP_MMAP_MAX_VMEM  (2147483648u)

 enum htp_tensor_flags {
-    HTP_TENSOR_COMPUTE = (1U << 0), // Tensor buffer temporal compute data (not weights)
-    HTP_TENSOR_DIRTY   = (1U << 1)  // Tensor buffer is dirty and needs to be flushed
+    HTP_TENSOR_WEIGHT  = (1U << 0), // Tensor buffer model weight data (not compute)
+    HTP_TENSOR_REPACK  = (1U << 1), // Tensor is in repacked tiled format
+    HTP_TENSOR_FENCE   = (1U << 2)  // Tensor is synchronization fence (explicitly managed)
 };

 // Tensor descriptor
@@ -175,6 +175,7 @@ enum htp_trace_event_id {
    HTP_TRACE_EVT_L2FLUSH             = 1,
    HTP_TRACE_EVT_INIT                = 2,
    HTP_TRACE_EVT_BUFF                = 3,
+    HTP_TRACE_EVT_FENCE               = 4,

    HTP_TRACE_EVT_HVX_COMP            = 20,
    HTP_TRACE_EVT_HVX_A_QUANT         = 21,
@@ -215,6 +216,7 @@ struct htp_opbatch_req {
    uint32_t n_ops;       // Number of ops
    uint32_t n_traces;    // Number of trace descriptors per thread
    uint32_t pad;         // unused
+    uint64_t seq;         // Sequence number
    // struct htp_buf_desc  bufs[];    -- dspqueue buf 0
    // struct htp_tensor    tensors[]; -- dspqueue buf 0
    // struct htp_op_desc   ops[];     -- dspqueue buf 0
@@ -231,6 +233,7 @@ struct htp_opbatch_rsp {
    uint32_t pad;            // align to 8 bytes
    uint64_t cycles_start;   // Start cycle counter
    uint64_t cycles_stop;    // Stop cycle counter
+    uint64_t seq;            // Sequence number
    // struct htp_prof_desc profs[];  -- dspqueue buf 0
 };

@@ -79,7 +79,14 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co

    for (uint32_t i = 0; i < n; i++) {
        const struct htp_tensor * t = tensors[i];
-        if (!t) continue;
+        if (!t || (t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE))) {
+            continue;
+        }
+
+        if (t->size <= HEX_L2_FLUSH_IL_THRESHOLD) {
+            hex_l2flush((void *) (uintptr_t) t->data, t->size);
+            continue;
+        }

        uint32_t t_start = t->data;
        uint32_t t_end   = t_start + t->size;
@@ -242,7 +249,7 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co

    for (uint32_t i = 0; i < n; i++) {
        const struct htp_tensor * t = tensors[i];
-        if (t && (t->flags & HTP_TENSOR_COMPUTE) && is_tensor_dirty(ctx, t)) {
+        if (t && !(t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE)) && is_tensor_dirty(ctx, t)) {
            dirty_tensors[n_dirty++] = t;
            total_dirty += t->size;
        }
@@ -13,6 +13,15 @@ static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) {
    return (uint32_t *) &t->flags;
 }

+static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) {
+    switch (type) {
+        case HTP_TYPE_F32:  return ne00 * 4;
+        case HTP_TYPE_F16:  return ne00 * 2;
+        case HTP_TYPE_Q8_0: return (ne00 / 32) * 34;
+        default:            return 0;
+    }
+}
+
 struct htp_context;
 void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
 void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
@@ -17,9 +17,9 @@

 #define hvx_arith_loop_body(dst_type, src0_type, src1_type, elem_size, vec_store, vec_op) \
    do {                                                                       \
-        dst_type * restrict vdst  = (dst_type *) dst;                          \
-        src0_type * restrict vsrc0 = (src0_type *) src0;                       \
-        src1_type * restrict vsrc1 = (src1_type *) src1;                       \
+        dst_type * vdst  = (dst_type *) dst;                                   \
+        src0_type * vsrc0 = (src0_type *) src0;                                \
+        src1_type * vsrc1 = (src1_type *) src1;                                \
                                                                               \
        const uint32_t epv  = 128 / (elem_size);                               \
        const uint32_t nvec = n / epv;                                         \
@@ -57,40 +57,40 @@

 // Generic macro to define alignment permutations for an op
 #define DEFINE_HVX_BINARY_OP_VARIANTS(OP_NAME, OP_MACRO, ELEM_TYPE) \
-static inline void OP_NAME##_aaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
+static inline void OP_NAME##_aaa(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
    assert((uintptr_t) dst % 128 == 0); \
    assert((uintptr_t) src0 % 128 == 0); \
    assert((uintptr_t) src1 % 128 == 0); \
    hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \
 } \
-static inline void OP_NAME##_aau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
+static inline void OP_NAME##_aau(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
    assert((uintptr_t) dst % 128 == 0); \
    assert((uintptr_t) src0 % 128 == 0); \
    hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \
 } \
-static inline void OP_NAME##_aua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
+static inline void OP_NAME##_aua(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
    assert((uintptr_t) dst % 128 == 0); \
    assert((uintptr_t) src1 % 128 == 0); \
    hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \
 } \
-static inline void OP_NAME##_auu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
+static inline void OP_NAME##_auu(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
    assert((uintptr_t) dst % 128 == 0); \
    hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \
 } \
-static inline void OP_NAME##_uaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
+static inline void OP_NAME##_uaa(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
    assert((uintptr_t) src0 % 128 == 0); \
    assert((uintptr_t) src1 % 128 == 0); \
    hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
 } \
-static inline void OP_NAME##_uau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
+static inline void OP_NAME##_uau(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
    assert((uintptr_t) src0 % 128 == 0); \
    hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
 } \
-static inline void OP_NAME##_uua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
+static inline void OP_NAME##_uua(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
    assert((uintptr_t) src1 % 128 == 0); \
    hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
 } \
-static inline void OP_NAME##_uuu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \
+static inline void OP_NAME##_uuu(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \
    hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
 } \

@@ -358,6 +358,34 @@ static inline void hvx_clamp_scalar_f32(uint8_t * restrict dst, const uint8_t *
    }
 }

+//
+// Abs
+//
+
+static inline void hvx_abs_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    assert((unsigned long) src % 128 == 0);
+
+    HVX_Vector * restrict vdst = (HVX_Vector *) dst;
+    HVX_Vector * restrict vsrc = (HVX_Vector *) src;
+
+    const uint32_t elem_size = sizeof(float);
+    const uint32_t epv       = 128 / elem_size;
+    const uint32_t nvec      = n / epv;
+    const uint32_t nloe      = n % epv;
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        vdst[i] = hvx_vec_abs_f32(vsrc[i]);
+    }
+    if (nloe) {
+        HVX_Vector v = hvx_vec_abs_f32(vsrc[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v);
+    }
+}
+
 //
 // Square
 //
@@ -62,4 +62,28 @@ static inline HVX_Vector hvx_vec_log_f32(HVX_Vector x) {
    return hvx_vec_add_f32_f32(term_e, res);
 }

+static inline void hvx_log_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    assert((unsigned long) src % 128 == 0);
+
+    HVX_Vector * restrict vdst = (HVX_Vector *) dst;
+    HVX_Vector * restrict vsrc = (HVX_Vector *) src;
+
+    const uint32_t elem_size = sizeof(float);
+    const uint32_t epv       = 128 / elem_size;
+    const uint32_t nvec      = n / epv;
+    const uint32_t nloe      = n % epv;
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        vdst[i] = hvx_vec_log_f32(vsrc[i]);
+    }
+    if (nloe) {
+        HVX_Vector v = hvx_vec_log_f32(vsrc[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v);
+    }
+}
+
 #endif /* HVX_LOG_H */
@@ -0,0 +1,165 @@
+#ifndef HVX_QUANT_H
+#define HVX_QUANT_H
+
+#include <math.h>
+#include <stdint.h>
+#include <string.h>
+
+#include "hvx-arith.h"
+#include "hvx-base.h"
+#include "hvx-reduce.h"
+#include "hvx-repl.h"
+#include "hvx-utils.h"
+
+#ifndef GGML_COMMON_DECL_C
+#define GGML_COMMON_DECL_C
+#endif
+#include "ggml-common.h"
+#include "ggml-impl.h"
+
+static inline void hvx_quantize_row_q8_0_f32(void * restrict dst_ptr, const float * restrict src_ptr, int n) {
+    const int nb = n / QK8_0;
+    block_q8_0 * dst = (block_q8_0 *) dst_ptr;
+    HVX_Vector zero = Q6_V_vzero();
+
+    int i = 0;
+    for (; i + 3 < nb; i += 4) {
+        HVX_Vector * vx = (HVX_Vector *) (src_ptr + i * QK8_0);
+
+        HVX_Vector vmax0_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[0]));
+        HVX_Vector vmax1_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[1]));
+        HVX_Vector vmax2_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[2]));
+        HVX_Vector vmax3_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[3]));
+
+        HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero);
+        HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero);
+        HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero);
+        HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero);
+
+        HVX_Vector vmax0_qf = Q6_Vqf32_vsub_VsfVsf(vmax0_sf, zero);
+        HVX_Vector vmax1_qf = Q6_Vqf32_vsub_VsfVsf(vmax1_sf, zero);
+        HVX_Vector vmax2_qf = Q6_Vqf32_vsub_VsfVsf(vmax2_sf, zero);
+        HVX_Vector vmax3_qf = Q6_Vqf32_vsub_VsfVsf(vmax3_sf, zero);
+
+        HVX_Vector vmax01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax1_qf, vmax0_qf)));
+        HVX_Vector vmax23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax3_qf, vmax2_qf)));
+
+        HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf)));
+        HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf)));
+
+        HVX_Vector vd01_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax01_hf, Q6_Vh_vsplat_R(0x2008));  // 1.0 / 127.0
+        HVX_Vector vd23_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax23_hf, Q6_Vh_vsplat_R(0x2008));  // 1.0 / 127.0
+        HVX_Vector vd01_hf   = Q6_Vhf_equals_Vqf16(vd01_qf16);
+        HVX_Vector vd23_hf   = Q6_Vhf_equals_Vqf16(vd23_qf16);
+
+        HVX_Vector vd01_inv_hf = hvx_vec_inverse_f16(vd01_hf);
+        HVX_Vector vd23_inv_hf = hvx_vec_inverse_f16(vd23_hf);
+        vx01_hf              = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd01_inv_hf));
+        vx23_hf              = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd23_inv_hf));
+
+        HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf);
+        HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf);
+        HVX_Vector vx_i8    = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16);
+
+        hvx_vec_store_u(&dst[i + 0].d, 2, vd01_hf);
+        hvx_vec_store_u(dst[i + 0].qs, 32, vx_i8);
+
+        hvx_vec_store_u(&dst[i + 1].d, 2, Q6_V_vror_VR(vd01_hf, 64));
+        hvx_vec_store_u(dst[i + 1].qs, 32, Q6_V_vror_VR(vx_i8, 32));
+
+        hvx_vec_store_u(&dst[i + 2].d, 2, vd23_hf);
+        hvx_vec_store_u(dst[i + 2].qs, 32, Q6_V_vror_VR(vx_i8, 64));
+
+        hvx_vec_store_u(&dst[i + 3].d, 2, Q6_V_vror_VR(vd23_hf, 64));
+        hvx_vec_store_u(dst[i + 3].qs, 32, Q6_V_vror_VR(vx_i8, 96));
+    }
+
+    for (; i < nb; i++) {
+        const float * block_src = src_ptr + i * QK8_0;
+        HVX_Vector vx = *(const HVX_UVector *) block_src;
+        HVX_Vector v_abs = hvx_vec_abs_f32(vx);
+        HVX_Vector v_max = hvx_vec_reduce_max_f32(v_abs);
+        float amax = hvx_vec_get_f32(v_max);
+
+        const float d = amax / 127.0f;
+        const float id = d ? (1.0f / d) : 0.0f;
+        dst[i].d = GGML_FP32_TO_FP16(d);
+
+        HVX_Vector vid = hvx_vec_splat_f32(id);
+        HVX_Vector v_scaled = hvx_vec_mul_f32_f32(vx, vid);
+        HVX_Vector v_scaled_qf = Q6_Vqf32_vsub_VsfVsf(v_scaled, zero);
+        HVX_Vector v_scaled_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(zero, v_scaled_qf)));
+        HVX_Vector v_i16 = hvx_vec_i16_from_hf_rnd_sat(v_scaled_hf);
+        HVX_Vector v_i8  = Q6_Vb_vpack_VhVh_sat(zero, v_i16);
+
+        hvx_vec_store_u(dst[i].qs, 32, v_i8);
+    }
+}
+
+static inline void hvx_dequantize_row_q8_0_f32(float * restrict dst_ptr, const void * restrict src_ptr, int n) {
+    const int nb = n / QK8_0;
+    const block_q8_0 * src = (const block_q8_0 *) src_ptr;
+
+    for (int i = 0; i < nb; i++) {
+        HVX_Vector vd_f16     = Q6_Vh_vsplat_R(*(const int16_t *) &src[i].d);
+        HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(vd_f16);
+        HVX_Vector vd         = Q6_V_lo_W(vp_f32);
+
+        HVX_Vector vq_i8 = *(const HVX_UVector *) src[i].qs;
+
+        HVX_VectorPair p16   = Q6_Wh_vunpack_Vb(vq_i8);
+        HVX_Vector     v_i16 = Q6_V_lo_W(p16);
+        HVX_VectorPair p32   = Q6_Ww_vunpack_Vh(v_i16);
+        HVX_Vector     v_i32 = Q6_V_lo_W(p32);
+
+        HVX_Vector v_f32 = Q6_Vsf_equals_Vw(v_i32);
+        HVX_Vector res   = hvx_vec_mul_f32_f32(v_f32, vd);
+
+        float * block_dst = dst_ptr + i * QK8_0;
+        hvx_vmem(block_dst) = res;
+    }
+}
+
+static inline void hvx_dequantize_row_q8_0_f16(__fp16 * restrict dst_ptr, const void * restrict src_ptr, int n) {
+    const int nb = n / QK8_0;
+    const block_q8_0 * src = (const block_q8_0 *) src_ptr;
+
+    for (int i = nb - 1; i >= 0; i--) {
+        HVX_Vector vd_f16     = Q6_Vh_vsplat_R(*(const int16_t *) &src[i].d);
+        HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(vd_f16);
+        HVX_Vector vd         = Q6_V_lo_W(vp_f32);
+
+        HVX_Vector vq_i8 = *(const HVX_UVector *) src[i].qs;
+
+        HVX_VectorPair p16   = Q6_Wh_vunpack_Vb(vq_i8);
+        HVX_Vector     v_i16 = Q6_V_lo_W(p16);
+        HVX_VectorPair p32   = Q6_Ww_vunpack_Vh(v_i16);
+        HVX_Vector     v_i32 = Q6_V_lo_W(p32);
+
+        HVX_Vector v_f32 = Q6_Vsf_equals_Vw(v_i32);
+        HVX_Vector res_f32 = hvx_vec_mul_f32_f32(v_f32, vd);
+
+        HVX_Vector res_f16 = hvx_vec_f32_to_f16(res_f32, Q6_V_vzero());
+
+        __fp16 * block_dst = dst_ptr + i * QK8_0;
+        hvx_vec_store_u(block_dst, QK8_0 * sizeof(__fp16), res_f16);
+    }
+}
+
+static inline void hvx_dequantize_row_f16_f32(float * restrict dst_ptr, const void * restrict src_ptr, int n) {
+    const int nb = n / 32;
+    const _Float16 * src = (const _Float16 *) src_ptr;
+
+    for (int i = 0; i < nb; i++) {
+        HVX_Vector v_f16 = *(const HVX_UVector *) (src + i * 32);
+        HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(v_f16);
+        HVX_Vector res = Q6_V_lo_W(vp_f32);
+
+        float * block_dst = dst_ptr + i * 32;
+        hvx_vmem(block_dst) = res;
+    }
+}
+
+
+
+#endif // HVX_QUANT_H
@@ -18,6 +18,7 @@
 #include <qurt_memory.h>
 #include <remote.h>
 #include <string.h>
+#include <stdatomic.h>

 #include "hex-utils.h"
 #include "hex-dma.h"
@@ -32,6 +33,7 @@
 #include "htp_iface.h"
 #include "work-queue.h"
 #include "hex-profile.h"
+#include "allreduce-ops.h"

 #define HMX_QUEUE_CAPACITY     16
 #define HMX_QUEUE_STACK_SIZE   16384
@@ -46,6 +48,36 @@ struct htp_handle {
    struct htp_context * ctx;
 };

+static inline void * htp_mmap(uint32_t fd, uint32_t size) {
+    void * va = (void *)-1;
+    for (int retry = 0; retry < 2; retry++) {
+#if __HVX_ARCH__ > 73
+        va = HAP_mmap2(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0);
+#else
+        if (size > HTP_MMAP_MAX_VMEM) {
+            FARF(ERROR, "mmap failed : size %u exceeds 2GB limit for HAP_mmap", (uint32_t) size);
+            abort();
+        }
+        va = HAP_mmap(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0);
+#endif
+        if (va != (void *)-1 && va != NULL) {
+            return va;
+        }
+        if (retry == 0) {
+            FARF(HIGH, "mmap failed first try (va %p fd %u size %u), retrying...", va, fd, size);
+        }
+    }
+    return NULL;
+}
+
+static inline void htp_munmap(void * va, uint32_t size) {
+#if __HVX_ARCH__ > 73
+    HAP_munmap2(va, size);
+#else
+    HAP_munmap(va, size);
+#endif
+}
+
 AEEResult htp_iface_open(const char * uri, remote_handle64 * handle) {
    (void) uri;
    struct htp_handle * h = calloc(1, sizeof(*h));
@@ -127,11 +159,7 @@ AEEResult htp_iface_close(remote_handle64 handle) {
        // release the mmaps (if any)
        for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
            if (ctx->mmap[i].size) {
-#if __HVX_ARCH__ > 73
-                HAP_munmap2((void *) ctx->mmap[i].base, ctx->mmap[i].size);
-#else
-                HAP_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size);
-#endif
+                htp_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size);
                ctx->mmap[i].size = 0;
                ctx->mmap[i].base = NULL;
                ctx->mmap[i].fd   = -1;
@@ -175,18 +203,9 @@ AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint32_t size) {
        struct htp_mmap *m = &ctx->mmap[i];
        if (!m->size) {
            FARF(HIGH, "mmap : fd %u size %u", fd, size);
-#if __HVX_ARCH__ > 73
-            void *va = HAP_mmap2(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0);
-#else
-            if (size > HTP_MMAP_MAX_VMEM) { // HAP_mmap has a size limit of 2GB
-                FARF(ERROR, "mmap failed : size %u exceeds 2GB limit for HAP_mmap", (uint32_t) size);
-                abort(); // can't do much else at this point
-            }
-
-            void *va = HAP_mmap(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0);
-#endif
-            if (va == (void*)-1) {
-                FARF(ERROR, "mmap failed : va %p fd %u size %u", va, fd, (uint32_t) size);
+            void *va = htp_mmap(fd, size);
+            if (va == NULL) {
+                FARF(ERROR, "mmap failed : fd %u size %u", fd, (uint32_t) size);
                return AEE_EFAILED;
            }

@@ -212,11 +231,7 @@ AEEResult htp_iface_munmap(remote_handle64 handle, uint32 fd) {
        struct htp_mmap *m = &ctx->mmap[i];
        if (fd < 0 || m->fd == fd) {
            FARF(HIGH, "unmmap : base %p fd %u size %u", (void*) m->base, m->fd, (uint32_t) m->size);
-#if __HVX_ARCH__ > 73
-            HAP_munmap2((void *) m->base, m->size);
-#else
-            HAP_munmap((void *) m->base, m->size);
-#endif
+            htp_munmap((void *) m->base, m->size);
            m->size   = 0;
            m->base   = NULL;
            m->fd     = -1;
@@ -228,7 +243,7 @@ AEEResult htp_iface_munmap(remote_handle64 handle, uint32 fd) {

 static void vtcm_acquire(struct htp_context * ctx) {
    if (!ctx->vtcm_valid) {
-        int err = HAP_compute_res_acquire_cached(ctx->vtcm_rctx, 1000000u);
+        int err = HAP_compute_res_acquire_cached(ctx->vtcm_rctx, 10000000u);
        if (err != 0) {
            FARF(ERROR, "ggml-hex: failed to acquire VTCM: 0x%08x", (unsigned)err);
            abort();
@@ -692,8 +707,45 @@ static inline void profile_stop(uint32_t mode, struct profile_data * d) {
    }
 }

+static int op_fence(struct htp_ops_context * octx) {
+    struct htp_context *ctx = octx->ctx;
+    struct htp_thread_trace * tr = &ctx->trace[0];
+    const uint32_t seq = (uint32_t) octx->op_params[0];
+
+    htp_trace_event_start(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);
+
+    const struct htp_tensor * sync = octx->src[0];
+    atomic_uint * sync_fence = (atomic_uint *) sync->data;
+    uint64_t spins = 0;
+    while (1) {
+        Q6_dccleaninva_A((void *) sync_fence);
+        asm volatile ("syncht" : : : "memory");
+        uint32_t val = atomic_load(&sync_fence[0]);
+        if ((int32_t)(val - seq) >= 0) {
+            break;
+        }
+        if (++spins > HTP_FENCE_TIMEOUT) {
+            FARF(ERROR, "ggml-hex: sync-wait TIMEOUT : fence %p spins %llu seq %u\n", sync_fence, spins, seq);
+            break;
+        }
+        hex_pause();
+    }
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);
+
+    FARF(HIGH, "ggml-hex: sync-done : fence %p spins %llu seq %u\n", sync_fence, spins, seq);
+    return HTP_STATUS_OK;
+}
+
 static int execute_op(struct htp_ops_context * octx) {
    switch (octx->op) {
+        case HTP_OP_FENCE:
+            return op_fence(octx);
+
+        case HTP_OP_ALLREDUCE:
+        case HTP_OP_ALLREDUCE_ADD:
+            return op_allreduce(octx);
+
        case HTP_OP_MUL_MAT:
        case HTP_OP_MUL_MAT_ADD:
            return op_matmul(octx);
@@ -701,11 +753,8 @@ static int execute_op(struct htp_ops_context * octx) {
        case HTP_OP_MUL_MAT_ID:
            return op_matmul_id(octx);

-        case HTP_OP_MUL_MAT_QKV:
-            return op_matmul_qkv(octx);
-
-        case HTP_OP_MUL_MAT_FFN:
-            return op_matmul_ffn(octx);
+        case HTP_OP_MUL_MAT_NX:
+            return op_matmul_nx(octx);

        case HTP_OP_MUL:
        case HTP_OP_ADD:
@@ -728,6 +777,8 @@ static int execute_op(struct htp_ops_context * octx) {
        case HTP_OP_UNARY_NEG:
        case HTP_OP_UNARY_EXP:
        case HTP_OP_UNARY_TANH:
+        case HTP_OP_UNARY_ABS:
+        case HTP_OP_UNARY_LOG:
        case HTP_OP_L2_NORM:
            return op_unary(octx);

@@ -818,12 +869,8 @@ static inline bool reuse_buf(struct htp_context *ctx, uint32_t *m_reuse, struct

 static inline void drop_mmap(struct htp_context *ctx, struct htp_mmap *m) {
    if (m->size) {
-        FARF(HIGH, "unmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size);
-#if __HVX_ARCH__ > 73
-        HAP_munmap2((void *) m->base, m->size);
-#else
-        HAP_munmap((void *) m->base, m->size);
-#endif
+        FARF(ALWAYS, "unmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size);
+        htp_munmap((void *) m->base, m->size);
        m->size = 0;
        m->base = 0;
        m->fd   = -1;
@@ -837,18 +884,9 @@ static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) {
    for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) {
        struct htp_mmap *m = &ctx->mmap[i];
        if (!m->size) {
-#if __HVX_ARCH__ > 73
-            void *va = HAP_mmap2(NULL, b->size, HAP_PROT_READ | HAP_PROT_WRITE, 0, b->fd, 0);
-#else
-            if (b->size > HTP_MMAP_MAX_VMEM) { // HAP_mmap has a size limit of 2GB
-                FARF(ERROR, "mmap failed : size %u exceeds 2GB limit for HAP_mmap", (uint32_t) b->size);
-                abort(); // can't do much else at this point
-            }
-
-            void *va = HAP_mmap(NULL, b->size, HAP_PROT_READ | HAP_PROT_WRITE, 0, b->fd, 0);
-#endif
-            if (va == (void*)-1) {
-                FARF(ERROR, "mmap failed : va %p fd %u size %u", va, b->fd, (uint32_t) b->size);
+            void *va = htp_mmap(b->fd, b->size);
+            if (va == NULL) {
+                FARF(ERROR, "mmap failed : fd %u size %u", b->fd, (uint32_t) b->size);
                abort(); // can't do much else at this point
            }

@@ -856,10 +894,13 @@ static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) {
            m->fd     = b->fd;
            m->size   = b->size;

-            FARF(HIGH, "mmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size);
+            FARF(ALWAYS, "mmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size);
            return;
        }
    }
+
+    FARF(ERROR, "mmap failed : exceeded mapping capacity limit of %u", HTP_MAX_MMAPS);
+    abort();
 }

 static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uint32_t n_bufs) {
@@ -1081,6 +1122,7 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
    rsp.usecs        = batch_prof.usecs;
    rsp.cycles_start = batch_prof.cycles_start;
    rsp.cycles_stop  = batch_prof.cycles_stop;
+    rsp.seq          = req->seq;

    if (ctx->profiler == HTP_PROF_TRACE) {
        for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
@@ -88,6 +88,7 @@ struct htp_mm_kernel_params {
    int32_t  vtcm_src2_size;     // src2 scratchpad size in VTCM (fused only)
    int32_t  vtcm_src3_size;     // src3 scratchpad size in VTCM (fused only)
    int32_t  vtcm_dst_size;      // dst scratchpad size in VTCM
+    int32_t  n_weights;          // Number of weights for fused NX

    // Precomputed division values
    struct fastdiv_values div_ne12_ne1;
@@ -463,8 +464,7 @@ static inline void htp_mm_hvx_vtcm_layout_build(
    size_t src2_row_size,
    uint32_t n_prefetch,
    bool is_matmul_id,
-    bool is_fused_qkv,
-    bool is_fused_ffn
+    bool is_fused_nx
 ) {
    size_t src0_sz = 0;
    size_t src1_sz = 0;
@@ -476,44 +476,33 @@ static inline void htp_mm_hvx_vtcm_layout_build(
                            wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL ||
                            wtype == HTP_TYPE_MXFP4);

-    if (is_fused_qkv || is_fused_ffn) {
+    if (is_fused_nx) {
        const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);
        const size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads;

-        size_t src0_sz_per_thread = 0;
-        size_t src2_sz_per_thread = 0;
-        size_t src3_sz_per_thread = 0;
+        size_t weight_sz_per_thread = 0;

        if (is_repack) {
            uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
            uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32;
            uint32_t tile_row_size = n_k_tiles * aligned_tile_size;

-            src0_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128);
-            src2_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128);
-            if (is_fused_qkv) {
-                src3_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128);
-            }
+            weight_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128);
        } else {
-            src0_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128);
-            src2_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128);
-            if (is_fused_qkv) {
-                src3_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128);
-            }
+            weight_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128);
        }

-        size_t flat_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
-        size_t tiled_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
+        size_t flat_act_row_size  = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10)  : htp_mm_q8_0_flat_row_size(ne10);
+        size_t tiled_act_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);

-        if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) {
-            src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128);
-        } else {
-            src1_sz = hex_round_up(tiled_src1_row_size * src1_nrows, 128);
-        }
+        size_t act_sz = (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT)
+            ? hex_round_up(flat_act_row_size  * src1_nrows, 128)
+            : hex_round_up(tiled_act_row_size * src1_nrows, 128);

-        src0_sz = src0_sz_per_thread * n_threads;
-        src2_sz = src2_sz_per_thread * n_threads;
-        src3_sz = src3_sz_per_thread * n_threads;
+        src0_sz = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer
+        src1_sz = act_sz;                           // quantized activation buffer
+        src2_sz = 0;
+        src3_sz = 0;
        dst_sz  = quant_scratch_size;
    } else if (is_matmul_id) {
        const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
@@ -616,8 +605,8 @@ static inline void htp_mm_hvx_vtcm_layout_build(
    }

    size_t off = 0;
-    VTCM_LAYOUT_ALLOC(off, off_src1, src1_sz);
    VTCM_LAYOUT_ALLOC(off, off_src0, src0_sz);
+    VTCM_LAYOUT_ALLOC(off, off_src1, src1_sz);
    VTCM_LAYOUT_ALLOC(off, off_src2, src2_sz);
    VTCM_LAYOUT_ALLOC(off, off_src3, src3_sz);
    VTCM_LAYOUT_ALLOC(off, off_dst,  dst_sz);
@@ -8,14 +8,20 @@
 #include <math.h>
 #include <string.h>

-#include "hex-dma.h"
+#include "dma-queue.h"
+#include "work-queue.h"
 #include "hvx-utils.h"
+#include "hex-utils.h"
+#include "hvx-copy.h"
+#include "hvx-quant.h"

 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
+#include "htp-tensor.h"
+#include "htp/set-rows-ops.h"

 #define set_rows_preamble                      \
    const uint32_t ne00 = octx->src[0]->ne[0]; \
@@ -47,116 +53,142 @@
                                               \
    const uint32_t nr  = ne01;

-struct htp_set_rows_context {
+struct set_rows_context {
    struct htp_ops_context * octx;
-    struct fastdiv_values div_ne12;
-    struct fastdiv_values div_ne11;
-    uint32_t src0_nrows_per_thread;
+    const struct htp_set_rows_kernel_params * kparams;
+    struct htp_set_rows_vtcm_layout vtcm_layout;
+    uint8_t * vtcm_base;
 };

-static void set_rows_thread_f32_f32(unsigned int nth, unsigned int ith, void *data) {
-    struct htp_set_rows_context * srctx = (struct htp_set_rows_context *)data;
-    struct htp_ops_context * octx = srctx->octx;
-
-    set_rows_preamble;
-
-    uint64_t qt = HAP_perf_get_qtimer_count();
-
-    // parallelize by rows of src0
-    const uint32_t dr  = srctx->src0_nrows_per_thread;
-    const uint32_t ir0 = dr * ith;
-    if (ir0 >= nr) {
-        return;
-    }
-    const uint32_t ir1 = (ir0 + dr < nr) ? (ir0 + dr) : nr;
-
-    const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);
-
-    for (uint32_t i03 = 0; i03 < ne03; ++i03) {
-        for (uint32_t i02 = 0; i02 < ne02; ++i02) {
-            for (uint32_t i = ir0; i < ir1; ++i) {
-                const uint32_t i12 = fastmodulo(i03, ne12, &srctx->div_ne12);
-                const uint32_t i11 = fastmodulo(i02, ne11, &srctx->div_ne11);
-                const uint32_t i10 = i;
-
-                const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12;
-
-                uint32_t i1 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr;
-                if (i1 >= ne1) {
-                    // ignore invalid indices
-                    continue;
-                }
-
-                const uintptr_t src0_ptr = octx->src[0]->data + i*nb01 + i02*nb02 + i03*nb03;
-                const uintptr_t dst_ptr  = octx->dst->data  + i1*nb1 + i02*nb2  + i03*nb3;
-
-                // copy row
-                hvx_copy_f32_uu((uint8_t *)dst_ptr, (const uint8_t *)src0_ptr, ne00);
-            }
-        }
-    }
-
-    qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt);
-    FARF(HIGH, "set-rows-f32-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth,
-         ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt);
+#define SET_ROWS_THREAD_DMA_FN(TYPE_NAME, IDX_TYPE, COMPUTE_EXPR)                                                \
+static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) {       \
+    struct set_rows_context * srctx = (struct set_rows_context *)data;                                           \
+    struct htp_ops_context * octx = srctx->octx;                                                                 \
+    const struct htp_set_rows_kernel_params * kparams = srctx->kparams;                                          \
+    set_rows_preamble;                                                                                           \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                       \
+    const uint32_t dr  = kparams->tasks_per_thread;                                                              \
+    const uint32_t ir0 = dr * ith;                                                                               \
+    if (ir0 >= kparams->total_tasks) {                                                                           \
+        return;                                                                                                  \
+    }                                                                                                            \
+    const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks);                                                    \
+    dma_queue * dma_queue = octx->ctx->dma[ith];                                                                 \
+    const struct htp_set_rows_vtcm_layout * vtcm_layout = &srctx->vtcm_layout;                                   \
+    uint8_t * vtcm_src0 = srctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread;   \
+    uint8_t * vtcm_dst  = srctx->vtcm_base + vtcm_layout->off_dst  + ith * vtcm_layout->dst_bytes_per_thread;    \
+    const uint32_t src0_row_size = ne00 * sizeof(float);                                                         \
+    const uint32_t dst_row_size  = htp_tensor_get_row_size(octx->dst->type, ne00);                               \
+    const uint32_t nrows_per_thread = ir1 - ir0;                                                                 \
+    const uint32_t total_steps = ne03 * ne02 * nrows_per_thread;                                                 \
+    uint32_t pi_step = 0;                                                                                        \
+    uint32_t pi02 = 0;                                                                                           \
+    uint32_t pi03 = 0;                                                                                           \
+    for (uint32_t step = 0, spad_idx = 0; step < total_steps && spad_idx < 2; ++step, spad_idx++) {              \
+        uint32_t i = ir0 + pi_step;                                                                              \
+        const uintptr_t src0_ptr = octx->src[0]->data + i*nb01 + pi02*nb02 + pi03*nb03;                          \
+        dma_queue_push(dma_queue,                                                                                \
+                       dma_make_ptr((void *)octx->dst->data,                                                     \
+                                    vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size),                      \
+                       dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 0);                          \
+        dma_queue_push(dma_queue,                                                                                \
+                       dma_make_ptr((void *)(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size),           \
+                                    (const void *)src0_ptr),                                                     \
+                       vtcm_layout->src0_spad_half_size, src0_row_size, src0_row_size, 1);                       \
+        pi_step++;                                                                                               \
+        if (pi_step == nrows_per_thread) {                                                                       \
+            pi_step = 0;                                                                                         \
+            pi02++;                                                                                              \
+            if (pi02 == ne02) {                                                                                  \
+                pi02 = 0;                                                                                        \
+                pi03++;                                                                                          \
+            }                                                                                                    \
+        }                                                                                                        \
+    }                                                                                                            \
+    uint32_t ci_step = 0;                                                                                        \
+    uint32_t ci02 = 0;                                                                                           \
+    uint32_t ci03 = 0;                                                                                           \
+    uint32_t ci11_base = 0;                                                                                      \
+    uint32_t ci12_base = 0;                                                                                      \
+    for (uint32_t step = 0; step < total_steps; ++step) {                                                        \
+        void * dst_spad = (void *) dma_queue_pop(dma_queue).src;                                                 \
+        void * src_spad = (void *) dma_queue_pop(dma_queue).dst;                                                 \
+        uint32_t i = ir0 + ci_step;                                                                              \
+        const uintptr_t src1_addr = octx->src[1]->data + i*nb10 + ci11_base*nb11 + ci12_base*nb12;               \
+        const IDX_TYPE i1 = *(const IDX_TYPE *)src1_addr;                                                        \
+        const bool valid_i1 = ((uint64_t)i1 < (uint64_t)ne1);                                                    \
+        const uint32_t target_i1 = (uint32_t)i1;                                                                 \
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, step);                                                 \
+        if (valid_i1) {                                                                                          \
+            COMPUTE_EXPR;                                                                                        \
+        }                                                                                                        \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, step);                                                  \
+        if (valid_i1) {                                                                                          \
+            const uintptr_t dst_ptr = octx->dst->data + target_i1*nb1 + ci02*nb2 + ci03*nb3;                     \
+            dma_queue_push(dma_queue,                                                                            \
+                           dma_make_ptr((void *)dst_ptr, (const void *)dst_spad),                                \
+                           dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 1);                      \
+        } else {                                                                                                 \
+            dma_queue_push(dma_queue,                                                                            \
+                           dma_make_ptr((void *)octx->dst->data, (const void *)dst_spad),                        \
+                           dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 0);                      \
+        }                                                                                                        \
+        const uint32_t next_step = step + 2;                                                                     \
+        if (next_step < total_steps) {                                                                           \
+            uint32_t ni = ir0 + pi_step;                                                                         \
+            const uintptr_t psrc0_ptr = octx->src[0]->data + ni*nb01 + pi02*nb02 + pi03*nb03;                    \
+            dma_queue_push(dma_queue,                                                                            \
+                           dma_make_ptr((void *)src_spad, (const void *)psrc0_ptr),                              \
+                           vtcm_layout->src0_spad_half_size, src0_row_size, src0_row_size, 1);                   \
+            pi_step++;                                                                                           \
+            if (pi_step == nrows_per_thread) {                                                                   \
+                pi_step = 0;                                                                                     \
+                pi02++;                                                                                          \
+                if (pi02 == ne02) {                                                                              \
+                    pi02 = 0;                                                                                    \
+                    pi03++;                                                                                      \
+                }                                                                                                \
+            }                                                                                                    \
+        }                                                                                                        \
+        ci_step++;                                                                                               \
+        if (ci_step == nrows_per_thread) {                                                                       \
+            ci_step = 0;                                                                                         \
+            ci02++;                                                                                              \
+            ci11_base++;                                                                                         \
+            if (ci11_base == ne11) {                                                                             \
+                ci11_base = 0;                                                                                   \
+            }                                                                                                    \
+            if (ci02 == ne02) {                                                                                  \
+                ci02 = 0;                                                                                        \
+                ci03++;                                                                                          \
+                ci12_base++;                                                                                     \
+                if (ci12_base == ne12) {                                                                         \
+                    ci12_base = 0;                                                                               \
+                }                                                                                                \
+            }                                                                                                    \
+        }                                                                                                        \
+    }                                                                                                            \
+    dma_queue_flush(dma_queue);                                                                                  \
 }

-static void set_rows_thread_f16_f32(unsigned int nth, unsigned int ith, void *data) {
-    struct htp_set_rows_context * srctx = (struct htp_set_rows_context *)data;
-    struct htp_ops_context * octx = srctx->octx;
+SET_ROWS_THREAD_DMA_FN(f32,  int32_t, { hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); })
+SET_ROWS_THREAD_DMA_FN(f32,  int64_t, { hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); })

-    set_rows_preamble;
+SET_ROWS_THREAD_DMA_FN(f16,  int32_t, { hvx_copy_f16_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); })
+SET_ROWS_THREAD_DMA_FN(f16,  int64_t, { hvx_copy_f16_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); })

-    uint64_t qt = HAP_perf_get_qtimer_count();
-
-    // parallelize by rows of src0
-    const uint32_t dr  = srctx->src0_nrows_per_thread;
-    const uint32_t ir0 = dr * ith;
-    if (ir0 >= nr) {
-        return;
-    }
-    const uint32_t ir1 = (ir0 + dr < nr) ? (ir0 + dr) : nr;
-
-    const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);
-
-    for (uint32_t i03 = 0; i03 < ne03; ++i03) {
-        for (uint32_t i02 = 0; i02 < ne02; ++i02) {
-            for (uint32_t i = ir0; i < ir1; ++i) {
-                const uint32_t i12 = fastmodulo(i03, ne12, &srctx->div_ne12);
-                const uint32_t i11 = fastmodulo(i02, ne11, &srctx->div_ne11);
-                const uint32_t i10 = i;
-
-                const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12;
-
-                uint32_t i1 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr;
-                if (i1 >= ne1) {
-                    // ignore invalid indices
-                    continue;
-                }
-
-                const uint8_t* src0_ptr = (const uint8_t *) octx->src[0]->data + i*nb01 + i02*nb02 + i03*nb03;
-                uint8_t*       dst_ptr  = (uint8_t *)       octx->dst->data  + i1*nb1 + i02*nb2  + i03*nb3;
-
-                hvx_copy_f16_f32_uu(dst_ptr, src0_ptr, ne00);
-            }
-        }
-    }
-
-    qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt);
-    FARF(HIGH, "set-rows-f16-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth,
-         ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt);
-}
+SET_ROWS_THREAD_DMA_FN(q8_0, int32_t, { hvx_quantize_row_q8_0_f32(dst_spad, (const float *)src_spad, ne00); })
+SET_ROWS_THREAD_DMA_FN(q8_0, int64_t, { hvx_quantize_row_q8_0_f32(dst_spad, (const float *)src_spad, ne00); })

 int op_set_rows(struct htp_ops_context * octx) {
+    const struct htp_set_rows_kernel_params * kparams = (const struct htp_set_rows_kernel_params *)octx->kernel_params;
    set_rows_preamble;

-    const uint32_t n_threads = MIN(nr, octx->n_threads);
-
    if (octx->src[0]->type != HTP_TYPE_F32) {
        return HTP_STATUS_NO_SUPPORT;
    }

-    if (octx->dst->type != HTP_TYPE_F32 && octx->dst->type != HTP_TYPE_F16) {
+    if (octx->dst->type != HTP_TYPE_F32 && octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_Q8_0) {
        return HTP_STATUS_NO_SUPPORT;
    }

@@ -164,27 +196,27 @@ int op_set_rows(struct htp_ops_context * octx) {
        return HTP_STATUS_NO_SUPPORT;
    }

-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
-        return HTP_STATUS_OK;
-    }
+    // l2fetch the src1 (indices) tensor in the main thread
+    hex_l2fetch_block((const void *)octx->src[1]->data, octx->src[1]->ne[3] * octx->src[1]->nb[3]);

-    struct htp_set_rows_context srctx;
+    struct set_rows_context srctx;
    srctx.octx = octx;
-    srctx.div_ne12 = init_fastdiv_values(ne12);
-    srctx.div_ne11 = init_fastdiv_values(ne11);
+    srctx.kparams = kparams;

-    srctx.src0_nrows_per_thread = (nr + n_threads - 1) / n_threads;
+    htp_set_rows_vtcm_layout_build(&srctx.vtcm_layout, octx->dst->type, ne00, kparams->n_threads);
+    srctx.vtcm_base = (uint8_t *)octx->ctx->vtcm_base;

-    switch(octx->dst->type) {
-    case HTP_TYPE_F32:
-        worker_pool_run_func(octx->ctx->worker_pool, set_rows_thread_f32_f32, &srctx, n_threads);
-        break;
-    case HTP_TYPE_F16:
-        worker_pool_run_func(octx->ctx->worker_pool, set_rows_thread_f16_f32, &srctx, n_threads);
-        break;
-    default:
-        return HTP_STATUS_NO_SUPPORT;
+    work_queue_func_t q_func = NULL;
+    const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);
+
+    switch (octx->dst->type) {
+        case HTP_TYPE_F32:  q_func = is_i32 ? set_rows_thread_dma_f32_int32_t  : set_rows_thread_dma_f32_int64_t;  break;
+        case HTP_TYPE_F16:  q_func = is_i32 ? set_rows_thread_dma_f16_int32_t  : set_rows_thread_dma_f16_int64_t;  break;
+        case HTP_TYPE_Q8_0: q_func = is_i32 ? set_rows_thread_dma_q8_0_int32_t : set_rows_thread_dma_q8_0_int64_t; break;
+        default:            return HTP_STATUS_NO_SUPPORT;
    }

+    work_queue_run(octx->ctx->work_queue, q_func, &srctx, kparams->n_threads);
+
    return HTP_STATUS_OK;
 }
@@ -0,0 +1,74 @@
+#ifndef HTP_SET_ROWS_OPS_H
+#define HTP_SET_ROWS_OPS_H
+
+#include "hex-fastdiv.h"
+
+struct htp_set_rows_kernel_params {
+    int32_t  n_threads;
+    int32_t  total_tasks;
+    int32_t  tasks_per_thread;
+    int32_t  vtcm_size;
+
+    // Fastdiv helpers
+    struct fastdiv_values div_ne11;
+    struct fastdiv_values div_ne12;
+    struct fastdiv_values div_tasks_per_thread;
+    struct fastdiv_values div_ne02;
+};
+
+struct htp_set_rows_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_dst;
+
+    size_t src0_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+
+    size_t src0_spad_half_size;
+    size_t dst_spad_half_size;
+};
+
+static inline void htp_set_rows_vtcm_layout_build(
+    struct htp_set_rows_vtcm_layout * vtcm_layout,
+    int dst_type,
+    uint32_t ne00,
+    uint32_t n_threads) {
+
+    size_t src0_row_size = ne00 * 4;
+    size_t dst_row_size = 0;
+    switch (dst_type) {
+        case 0: // HTP_TYPE_F32
+            dst_row_size = ne00 * 4;
+            break;
+        case 1: // HTP_TYPE_F16
+            dst_row_size = ne00 * 2;
+            break;
+        case 8: // HTP_TYPE_Q8_0
+            dst_row_size = (ne00 / 32) * 34;
+            break;
+        default:
+            dst_row_size = 0;
+            break;
+    }
+
+    size_t src0_row_size_aligned = (src0_row_size + 255) & ~255;
+    size_t dst_row_size_aligned  = (dst_row_size + 255) & ~255;
+
+    vtcm_layout->src0_spad_half_size = src0_row_size_aligned;
+    vtcm_layout->dst_spad_half_size  = dst_row_size_aligned;
+
+    vtcm_layout->src0_bytes_per_thread = src0_row_size_aligned * 2;
+    vtcm_layout->dst_bytes_per_thread  = dst_row_size_aligned * 2;
+
+    vtcm_layout->off_src0 = 0;
+    vtcm_layout->off_dst  = vtcm_layout->off_src0 + vtcm_layout->src0_bytes_per_thread * n_threads;
+    vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads;
+}
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_set_rows_kernel_params) <= 128, "htp_set_rows_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_set_rows_kernel_params) <= 128, "htp_set_rows_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_SET_ROWS_OPS_H
@@ -443,6 +443,34 @@ static void tanh_f32(const float * restrict src,
    }
 }

+static void abs_f32(const float * restrict src,
+                    float * restrict dst,
+                    const uint32_t num_rows,
+                    const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_abs_f32_aa(dst_local, src_local, ne0);
+    }
+}
+
+static void log_f32(const float * restrict src,
+                    float * restrict dst,
+                    const uint32_t num_rows,
+                    const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_log_f32_aa(dst_local, src_local, ne0);
+    }
+}
+
 #define DEFINE_UNARY_TASK(NAME, IS_RMS_NORM_MUL, IS_TRI, CORE_EXPR)                                                 \
 static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * data) {                                \
    const struct htp_unary_context * uctx = (const struct htp_unary_context *) data;                                \
@@ -478,6 +506,9 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
    const uint32_t nb11 = src1 ? src1->nb[1] : 0;                                                                   \
    const uint32_t nb12 = src1 ? src1->nb[2] : 0;                                                                   \
    const uint32_t nb13 = src1 ? src1->nb[3] : 0;                                                                   \
+    const uint32_t nb11_bc = (src1 && src1->ne[1] > 1) ? nb11 : 0;                                                  \
+    const uint32_t nb12_bc = (src1 && src1->ne[2] > 1) ? nb12 : 0;                                                  \
+    const uint32_t nb13_bc = (src1 && src1->ne[3] > 1) ? nb13 : 0;                                                  \
    const bool src1_contig = src1 ? ((nb12 == (size_t)ne01 * nb11) && (nb13 == (size_t)ne02 * nb12)) : false;       \
                                                                                                                    \
    uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread);                           \
@@ -497,8 +528,12 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
    const struct fastdiv_values * div_ne02  = &uctx->kparams->div_ne02;                                             \
    const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012;                                            \
                                                                                                                    \
-    const uint32_t src0_max_block = src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01);                   \
-    const uint32_t dst_max_block  = dst_contig  ? uctx->block : MIN((uint32_t)uctx->block, ne1);                    \
+    const bool src1_needs_row_clip = (IS_RMS_NORM_MUL) && !uctx->broadcast_weight && !src1_contig;                  \
+    const bool block_src0_contig = src0_contig && !src1_needs_row_clip;                                             \
+    const bool block_dst_contig  = dst_contig  && !src1_needs_row_clip;                                             \
+                                                                                                                    \
+    const uint32_t src0_max_block = block_src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01);             \
+    const uint32_t dst_max_block  = block_dst_contig  ? uctx->block : MIN((uint32_t)uctx->block, ne1);              \
    const uint32_t BLOCK = MIN(src0_max_block, dst_max_block);                                                      \
    if (BLOCK == 0) {                                                                                               \
        FARF(ERROR, "unary-f32 : current VTCM reservation %zu is too small, needed at least %zu\n",                 \
@@ -515,8 +550,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
    }                                                                                                               \
                                                                                                                    \
    for (uint32_t ir = src0_start_row, vtcm_idx = 0; ir < src0_end_row && vtcm_idx < 2; vtcm_idx++) {               \
-        const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01,        \
-                                                     div_ne01);                                                     \
+        const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig,  \
+                                                     ne01, div_ne01);                                               \
                                                                                                                    \
        dma_queue_push(dma_queue,                                                                                   \
            dma_make_ptr(data_dst, dst_vtcm_data + (vtcm_idx * dst_vtcm_half_size)),                                \
@@ -530,7 +565,7 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
                                                                                                                    \
        if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) {                                                         \
            const size_t src1_off = src1_contig ? (ir * nb11) :                                                     \
-                unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13);                  \
+                unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, nb13_bc);         \
            dma_queue_push(dma_queue,                                                                               \
                dma_make_ptr(src1_vtcm_data + (vtcm_idx * src1_vtcm_half_size), data_src1 + src1_off),              \
                uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, block_size);                           \
@@ -540,8 +575,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
    }                                                                                                               \
                                                                                                                    \
    for (uint32_t ir = src0_start_row; ir < src0_end_row; ) {                                                       \
-        const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01,        \
-                                                     div_ne01);                                                     \
+        const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig,  \
+                                                     ne01, div_ne01);                                               \
                                                                                                                    \
        float * dst_vtcm  = (float *) dma_queue_pop(dma_queue).src;                                                 \
        float * src0_vtcm = (float *) dma_queue_pop(dma_queue).dst;                                                 \
@@ -562,12 +597,12 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
                                                                                                                    \
        const uint32_t next_ir = ir + block_size;                                                                   \
        if (next_ir < src0_end_row) {                                                                               \
-            const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, src0_contig, dst_contig,\
-                                                              ne01, div_ne01);                                      \
+            const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, block_src0_contig,      \
+                                                              block_dst_contig, ne01, div_ne01);                    \
            const uint32_t pref_ir = next_ir + next_block_size;                                                     \
            if (pref_ir < src0_end_row) {                                                                           \
-                const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, src0_contig,        \
-                                                                  dst_contig, ne01, div_ne01);                      \
+                const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, block_src0_contig,  \
+                                                                  block_dst_contig, ne01, div_ne01);                \
                const size_t src0_pref_off = src0_contig ? (pref_ir * nb01) :                                       \
                    unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03);         \
                dma_queue_push(dma_queue,                                                                           \
@@ -576,7 +611,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
                                                                                                                    \
                if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) {                                                 \
                    const size_t src1_pref_off = src1_contig ? (pref_ir * nb11) :                                   \
-                        unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13);     \
+                        unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc,      \
+                                          nb13_bc);                                                                 \
                    dma_queue_push(dma_queue,                                                                       \
                        dma_make_ptr(src1_vtcm, data_src1 + src1_pref_off),                                         \
                        uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, pref_block_size);              \
@@ -603,6 +639,8 @@ DEFINE_UNARY_TASK(unary_silu,     false, false, silu_f32(src0_vtcm, dst_vtcm, bl
 DEFINE_UNARY_TASK(unary_gelu,     false, false, gelu_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(unary_softplus, false, false, softplus_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(unary_tanh,     false, false, tanh_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_abs,      false, false, abs_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_log,      false, false, log_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(l2_norm,        false, false, l2_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(tri,            false, true,  tri_f32(src0_vtcm, dst_vtcm, block_size, ir, uctx))

@@ -850,6 +888,8 @@ DEFINE_UNARY_TILED_TASK(unary_silu,     false, tile_silu_f32(dst_vtcm, src_vtcm,
 DEFINE_UNARY_TILED_TASK(unary_gelu,     false, tile_gelu_f32(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(unary_softplus, false, tile_unary_softplus_f32(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(unary_tanh,     false, hvx_tanh_f32_aa(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(unary_abs,      false, hvx_abs_f32_aa(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(unary_log,      false, hvx_log_f32_aa(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(tri,            true,  tri_apply_tile_f32(src_vtcm, dst_vtcm, tw, col, i01, ne0, tri_ttype))

 static int execute_op_unary_f32(struct htp_ops_context * octx) {
@@ -875,6 +915,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
        case HTP_OP_UNARY_GELU:      op_type = "gelu-f32";         break;
        case HTP_OP_UNARY_SOFTPLUS:  op_type = "softplus-f32";     break;
        case HTP_OP_UNARY_TANH:      op_type = "tanh-f32";         break;
+        case HTP_OP_UNARY_ABS:       op_type = "abs-f32";          break;
+        case HTP_OP_UNARY_LOG:       op_type = "log-f32";          break;
        case HTP_OP_L2_NORM:         op_type = "l2norm-f32";       break;
        case HTP_OP_TRI:             op_type = "tri-f32";          break;

@@ -973,6 +1015,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
                case HTP_OP_UNARY_GELU:      task_func = unary_task_f32_tiled_unary_gelu;     break;
                case HTP_OP_UNARY_SOFTPLUS:  task_func = unary_task_f32_tiled_unary_softplus; break;
                case HTP_OP_UNARY_TANH:      task_func = unary_task_f32_tiled_unary_tanh;     break;
+                case HTP_OP_UNARY_ABS:       task_func = unary_task_f32_tiled_unary_abs;      break;
+                case HTP_OP_UNARY_LOG:       task_func = unary_task_f32_tiled_unary_log;      break;
                case HTP_OP_TRI:             task_func = unary_task_f32_tiled_tri;            break;
                default:                     break;
            }
@@ -992,6 +1036,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
                case HTP_OP_UNARY_GELU:      task_func = unary_task_f32_unary_gelu;           break;
                case HTP_OP_UNARY_SOFTPLUS:  task_func = unary_task_f32_unary_softplus;       break;
                case HTP_OP_UNARY_TANH:      task_func = unary_task_f32_unary_tanh;           break;
+                case HTP_OP_UNARY_ABS:       task_func = unary_task_f32_unary_abs;            break;
+                case HTP_OP_UNARY_LOG:       task_func = unary_task_f32_unary_log;            break;
                case HTP_OP_L2_NORM:         task_func = unary_task_f32_l2_norm;              break;
                case HTP_OP_TRI:             task_func = unary_task_f32_tri;                  break;
                default:                     break;
@@ -55,6 +55,8 @@ static inline bool htp_op_is_unary(uint32_t opcode) {
        case HTP_OP_UNARY_GELU:
        case HTP_OP_UNARY_SOFTPLUS:
        case HTP_OP_UNARY_TANH:
+        case HTP_OP_UNARY_ABS:
+        case HTP_OP_UNARY_LOG:
        case HTP_OP_L2_NORM:
        case HTP_OP_TRI:
            return true;
@@ -84,106 +84,108 @@ struct ggml_metal {
 ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) {
    GGML_LOG_INFO("%s: allocating\n", __func__);

+    @autoreleasepool {
 #if TARGET_OS_OSX && !GGML_METAL_NDEBUG
-    // Show all the Metal device instances in the system
-    NSArray * devices = MTLCopyAllDevices();
-    for (id<MTLDevice> device in devices) {
-        GGML_LOG_INFO("%s: found device: %s\n", __func__, [[device name] UTF8String]);
-    }
-    [devices release]; // since it was created by a *Copy* C method
+        // Show all the Metal device instances in the system
+        NSArray * devices = MTLCopyAllDevices();
+        for (id<MTLDevice> device in devices) {
+            GGML_LOG_INFO("%s: found device: %s\n", __func__, [[device name] UTF8String]);
+        }
+        [devices release]; // since it was created by a *Copy* C method
 #endif

-    // init context
-    ggml_metal_t res = calloc(1, sizeof(struct ggml_metal));
+        // init context
+        ggml_metal_t res = calloc(1, sizeof(struct ggml_metal));

-    id<MTLDevice> device = ggml_metal_device_get_obj(dev);
+        id<MTLDevice> device = ggml_metal_device_get_obj(dev);

-    GGML_LOG_INFO("%s: picking default device: %s\n", __func__, [[device name] UTF8String]);
-
-    // TODO: would it be better to have one queue for the backend and one queue for the device?
-    //       the graph encoders and async ops would use the backend queue while the sync ops would use the device queue?
-    //res->queue = [device newCommandQueue]; [TAG_QUEUE_PER_BACKEND]
-    id<MTLCommandQueue> queue = ggml_metal_device_get_queue(dev);
-    if (queue == nil) {
-        GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__);
-        return NULL;
-    }
-
-    res->dev = dev;
-    res->lib = ggml_metal_device_get_library(dev);
-    if (res->lib == NULL) {
-        GGML_LOG_WARN("%s: the device does not have a precompiled Metal library - this is unexpected\n", __func__);
-        GGML_LOG_WARN("%s: will try to compile it on the fly\n", __func__);
-
-        res->lib = ggml_metal_library_init(dev);
-        if (res->lib == NULL) {
-            GGML_LOG_ERROR("%s: error: failed to initialize the Metal library\n", __func__);
-
-            free(res);
+        GGML_LOG_INFO("%s: picking default device: %s\n", __func__, [[device name] UTF8String]);

+        // TODO: would it be better to have one queue for the backend and one queue for the device?
+        //       the graph encoders and async ops would use the backend queue while the sync ops would use the device queue?
+        //res->queue = [device newCommandQueue]; [TAG_QUEUE_PER_BACKEND]
+        id<MTLCommandQueue> queue = ggml_metal_device_get_queue(dev);
+        if (queue == nil) {
+            GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__);
            return NULL;
        }
-    }

-    res->ev_cpy = ggml_metal_device_event_init(dev);
+        res->dev = dev;
+        res->lib = ggml_metal_device_get_library(dev);
+        if (res->lib == NULL) {
+            GGML_LOG_WARN("%s: the device does not have a precompiled Metal library - this is unexpected\n", __func__);
+            GGML_LOG_WARN("%s: will try to compile it on the fly\n", __func__);

-    const struct ggml_metal_device_props * props_dev = ggml_metal_device_get_props(dev);
+            res->lib = ggml_metal_library_init(dev);
+            if (res->lib == NULL) {
+                GGML_LOG_ERROR("%s: error: failed to initialize the Metal library\n", __func__);

-    snprintf(res->name, sizeof(res->name), "%s", props_dev->name);
+                free(res);

-    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;
-
-    {
-        const char * val = getenv("GGML_METAL_GRAPH_DEBUG");
-        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));
-
-    GGML_LOG_INFO("%s: use fusion         = %s\n", __func__, res->use_fusion         ? "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");
-
-    res->capture_compute = 0;
-    res->capture_started = false;
-    res->capture_scope = nil;
-
-    {
-        const char * val = getenv("GGML_METAL_CAPTURE_COMPUTE");
-        if (val) {
-            res->capture_compute = atoi(val);
+                return NULL;
+            }
        }
+
+        res->ev_cpy = ggml_metal_device_event_init(dev);
+
+        const struct ggml_metal_device_props * props_dev = ggml_metal_device_get_props(dev);
+
+        snprintf(res->name, sizeof(res->name), "%s", props_dev->name);
+
+        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;
+
+        {
+            const char * val = getenv("GGML_METAL_GRAPH_DEBUG");
+            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));
+
+        GGML_LOG_INFO("%s: use fusion         = %s\n", __func__, res->use_fusion         ? "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");
+
+        res->capture_compute = 0;
+        res->capture_started = false;
+        res->capture_scope = nil;
+
+        {
+            const char * val = getenv("GGML_METAL_CAPTURE_COMPUTE");
+            if (val) {
+                res->capture_compute = atoi(val);
+            }
+        }
+
+        res->has_error = false;
+
+        res->gf = nil;
+        res->encode_async = nil;
+        for (int i = 0; i < GGML_METAL_MAX_COMMAND_BUFFERS; ++i) {
+            res->cmd_bufs[i].obj = nil;
+        }
+
+        res->cmd_bufs_ext = [[NSMutableArray alloc] init];
+
+        res->cmd_buf_last = nil;
+
+        res->pipelines_ext = ggml_metal_pipelines_init();
+
+        return res;
    }
-
-    res->has_error = false;
-
-    res->gf = nil;
-    res->encode_async = nil;
-    for (int i = 0; i < GGML_METAL_MAX_COMMAND_BUFFERS; ++i) {
-        res->cmd_bufs[i].obj = nil;
-    }
-
-    res->cmd_bufs_ext = [[NSMutableArray alloc] init];
-
-    res->cmd_buf_last = nil;
-
-    res->pipelines_ext = ggml_metal_pipelines_init();
-
-    return res;
 }

 void ggml_metal_free(ggml_metal_t ctx) {
@@ -778,7 +778,9 @@ void ggml_metal_encoder_free(ggml_metal_encoder_t encoder) {
 }

 void ggml_metal_encoder_debug_group_push(ggml_metal_encoder_t encoder, const char * name) {
-    [encoder->obj pushDebugGroup:[NSString stringWithCString:name encoding:NSUTF8StringEncoding]];
+    @autoreleasepool {
+        [encoder->obj pushDebugGroup:[NSString stringWithCString:name encoding:NSUTF8StringEncoding]];
+    }
 }

 void ggml_metal_encoder_debug_group_pop (ggml_metal_encoder_t encoder) {
@@ -1023,249 +1025,251 @@ ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) {

    assert(dev != NULL);

-    if (dev->mtl_device == nil) {
-        dev->mtl_device = MTLCreateSystemDefaultDevice();
+    @autoreleasepool {
+        if (dev->mtl_device == nil) {
+            dev->mtl_device = MTLCreateSystemDefaultDevice();

-        if (dev->mtl_device) {
-            dev->mtl_queue = [dev->mtl_device newCommandQueue];
-            if (dev->mtl_queue == nil) {
-                GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__);
-            }
+            if (dev->mtl_device) {
+                dev->mtl_queue = [dev->mtl_device newCommandQueue];
+                if (dev->mtl_queue == nil) {
+                    GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__);
+                }

-            dev->addr_virt = 0x000000400ULL;
+                dev->addr_virt = 0x000000400ULL;

-            dev->props.device = device;
+                dev->props.device = device;

-            // the Metal backend uses the system default device as the single physical device;
-            // additional (virtual) devices are emulated on top of it via GGML_METAL_DEVICES
-            dev->props.device_phys = 0;
-            dev->props.device_virt = device;
+                // the Metal backend uses the system default device as the single physical device;
+                // additional (virtual) devices are emulated on top of it via GGML_METAL_DEVICES
+                dev->props.device_phys = 0;
+                dev->props.device_virt = device;

-            dev->props.has_simdgroup_reduction  = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7];
-            dev->props.has_simdgroup_reduction |= [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML];
+                dev->props.has_simdgroup_reduction  = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7];
+                dev->props.has_simdgroup_reduction |= [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML];

-            dev->props.has_simdgroup_mm = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7];
-            dev->props.has_unified_memory = dev->mtl_device.hasUnifiedMemory;
+                dev->props.has_simdgroup_mm = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7];
+                dev->props.has_unified_memory = dev->mtl_device.hasUnifiedMemory;

-            dev->props.has_bfloat  = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML];
-            dev->props.has_bfloat |= [dev->mtl_device supportsFamily:MTLGPUFamilyApple6];
-            if (getenv("GGML_METAL_BF16_DISABLE") != NULL) {
-                dev->props.has_bfloat = false;
-            }
+                dev->props.has_bfloat  = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML];
+                dev->props.has_bfloat |= [dev->mtl_device supportsFamily:MTLGPUFamilyApple6];
+                if (getenv("GGML_METAL_BF16_DISABLE") != NULL) {
+                    dev->props.has_bfloat = false;
+                }

-            dev->props.has_tensor = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal4_GGML];
-            if (getenv("GGML_METAL_TENSOR_DISABLE") != NULL) {
-                dev->props.has_tensor = false;
-            }
-
-            // note: disable the tensor API by default for old chips because with the current implementation it is not useful
-            // - M2 Ultra:   ~5% slower
-            // - M4, M4 Max: no significant difference
-            //
-            // TODO: try to update the tensor API kernels to at least match the simdgroup performance
-            if (getenv("GGML_METAL_TENSOR_ENABLE") == NULL &&
-                ![[dev->mtl_device name] containsString:@"M5"] &&
-                ![[dev->mtl_device name] containsString:@"M6"] &&
-                ![[dev->mtl_device name] containsString:@"A19"] &&
-                ![[dev->mtl_device name] containsString:@"A20"]) {
-                GGML_LOG_INFO("%s: tensor API disabled for pre-M5 and pre-A19 devices\n", __func__);
-                dev->props.has_tensor = false;
-            }
-
-            // double-check that the tensor API compiles
-            if (dev->props.has_tensor) {
-                const char * src_tensor_f16 = "\n"
-                    "#include <metal_stdlib> \n"
-                    "#include <metal_tensor> \n"
-                    "#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h> \n"
-                    " \n"
-                    "using namespace metal; \n"
-                    "using namespace mpp::tensor_ops; \n"
-                    " \n"
-                    "kernel void dummy_kernel( \n"
-                    "    tensor<device  half, dextents<int32_t, 2>> A [[buffer(0)]], \n"
-                    "    tensor<device  half, dextents<int32_t, 2>> B [[buffer(1)]], \n"
-                    "    device float * C [[buffer(2)]], \n"
-                    "    uint2 tgid [[threadgroup_position_in_grid]]) \n"
-                    "{ \n"
-                    "    auto tA = A.slice(0, (int)tgid.y); \n"
-                    "    auto tB = B.slice((int)tgid.x, 0); \n"
-                    " \n"
-                    "    matmul2d< \n"
-                    "        matmul2d_descriptor(16, 16, dynamic_extent), \n"
-                    "        execution_simdgroups<4>> mm; \n"
-                    " \n"
-                    "    auto cT = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB), float>(); \n"
-                    " \n"
-                    "    auto sA = tA.slice(0, 0); \n"
-                    "    auto sB = tB.slice(0, 0); \n"
-                    "    mm.run(sB, sA, cT); \n"
-                    " \n"
-                    "    auto tC = tensor<device float, dextents<int32_t, 2>, tensor_inline>(C, dextents<int32_t, 2>(16, 16)); \n"
-                    " \n"
-                    "    cT.store(tC); \n"
-                    "}";
-
-                GGML_LOG_INFO("%s: testing tensor API for f16 support\n", __func__);
-                ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_f16, false);
-                if (lib == NULL) {
-                    GGML_LOG_WARN("%s: - the tensor API is not supported in this environment - disabling\n", __func__);
+                dev->props.has_tensor = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal4_GGML];
+                if (getenv("GGML_METAL_TENSOR_DISABLE") != NULL) {
                    dev->props.has_tensor = false;
-                } else {
-                    struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil);
-                    if (!ppl.pipeline) {
+                }
+
+                // note: disable the tensor API by default for old chips because with the current implementation it is not useful
+                // - M2 Ultra:   ~5% slower
+                // - M4, M4 Max: no significant difference
+                //
+                // TODO: try to update the tensor API kernels to at least match the simdgroup performance
+                if (getenv("GGML_METAL_TENSOR_ENABLE") == NULL &&
+                    ![[dev->mtl_device name] containsString:@"M5"] &&
+                    ![[dev->mtl_device name] containsString:@"M6"] &&
+                    ![[dev->mtl_device name] containsString:@"A19"] &&
+                    ![[dev->mtl_device name] containsString:@"A20"]) {
+                    GGML_LOG_INFO("%s: tensor API disabled for pre-M5 and pre-A19 devices\n", __func__);
+                    dev->props.has_tensor = false;
+                }
+
+                // double-check that the tensor API compiles
+                if (dev->props.has_tensor) {
+                    const char * src_tensor_f16 = "\n"
+                        "#include <metal_stdlib> \n"
+                        "#include <metal_tensor> \n"
+                        "#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h> \n"
+                        " \n"
+                        "using namespace metal; \n"
+                        "using namespace mpp::tensor_ops; \n"
+                        " \n"
+                        "kernel void dummy_kernel( \n"
+                        "    tensor<device  half, dextents<int32_t, 2>> A [[buffer(0)]], \n"
+                        "    tensor<device  half, dextents<int32_t, 2>> B [[buffer(1)]], \n"
+                        "    device float * C [[buffer(2)]], \n"
+                        "    uint2 tgid [[threadgroup_position_in_grid]]) \n"
+                        "{ \n"
+                        "    auto tA = A.slice(0, (int)tgid.y); \n"
+                        "    auto tB = B.slice((int)tgid.x, 0); \n"
+                        " \n"
+                        "    matmul2d< \n"
+                        "        matmul2d_descriptor(16, 16, dynamic_extent), \n"
+                        "        execution_simdgroups<4>> mm; \n"
+                        " \n"
+                        "    auto cT = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB), float>(); \n"
+                        " \n"
+                        "    auto sA = tA.slice(0, 0); \n"
+                        "    auto sB = tB.slice(0, 0); \n"
+                        "    mm.run(sB, sA, cT); \n"
+                        " \n"
+                        "    auto tC = tensor<device float, dextents<int32_t, 2>, tensor_inline>(C, dextents<int32_t, 2>(16, 16)); \n"
+                        " \n"
+                        "    cT.store(tC); \n"
+                        "}";
+
+                    GGML_LOG_INFO("%s: testing tensor API for f16 support\n", __func__);
+                    ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_f16, false);
+                    if (lib == NULL) {
                        GGML_LOG_WARN("%s: - the tensor API is not supported in this environment - disabling\n", __func__);
                        dev->props.has_tensor = false;
+                    } else {
+                        struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil);
+                        if (!ppl.pipeline) {
+                            GGML_LOG_WARN("%s: - the tensor API is not supported in this environment - disabling\n", __func__);
+                            dev->props.has_tensor = false;
+                        }
+
+                        ggml_metal_library_free(lib);
                    }
-
-                    ggml_metal_library_free(lib);
                }
-            }

-            // try to compile a dummy kernel to determine if the tensor API is supported for bfloat
-            if (dev->props.has_tensor && dev->props.has_bfloat) {
-                const char * src_tensor_bf16 = "\n"
-                    "#include <metal_stdlib> \n"
-                    "#include <metal_tensor> \n"
-                    "#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h> \n"
-                    " \n"
-                    "using namespace metal; \n"
-                    "using namespace mpp::tensor_ops; \n"
-                    " \n"
-                    "kernel void dummy_kernel( \n"
-                    "    tensor<device bfloat, dextents<int32_t, 2>> A [[buffer(0)]], \n"
-                    "    tensor<device bfloat, dextents<int32_t, 2>> B [[buffer(1)]], \n"
-                    "    device float * C [[buffer(2)]], \n"
-                    "    uint2 tgid [[threadgroup_position_in_grid]]) \n"
-                    "{ \n"
-                    "    auto tA = A.slice(0, (int)tgid.y); \n"
-                    "    auto tB = B.slice((int)tgid.x, 0); \n"
-                    " \n"
-                    "    matmul2d< \n"
-                    "        matmul2d_descriptor(16, 16, dynamic_extent), \n"
-                    "        execution_simdgroups<4>> mm; \n"
-                    " \n"
-                    "    auto cT = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB), float>(); \n"
-                    " \n"
-                    "    auto sA = tA.slice(0, 0); \n"
-                    "    auto sB = tB.slice(0, 0); \n"
-                    "    mm.run(sB, sA, cT); \n"
-                    " \n"
-                    "    auto tC = tensor<device float, dextents<int32_t, 2>, tensor_inline>(C, dextents<int32_t, 2>(16, 16)); \n"
-                    " \n"
-                    "    cT.store(tC); \n"
-                    "}";
+                // try to compile a dummy kernel to determine if the tensor API is supported for bfloat
+                if (dev->props.has_tensor && dev->props.has_bfloat) {
+                    const char * src_tensor_bf16 = "\n"
+                        "#include <metal_stdlib> \n"
+                        "#include <metal_tensor> \n"
+                        "#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h> \n"
+                        " \n"
+                        "using namespace metal; \n"
+                        "using namespace mpp::tensor_ops; \n"
+                        " \n"
+                        "kernel void dummy_kernel( \n"
+                        "    tensor<device bfloat, dextents<int32_t, 2>> A [[buffer(0)]], \n"
+                        "    tensor<device bfloat, dextents<int32_t, 2>> B [[buffer(1)]], \n"
+                        "    device float * C [[buffer(2)]], \n"
+                        "    uint2 tgid [[threadgroup_position_in_grid]]) \n"
+                        "{ \n"
+                        "    auto tA = A.slice(0, (int)tgid.y); \n"
+                        "    auto tB = B.slice((int)tgid.x, 0); \n"
+                        " \n"
+                        "    matmul2d< \n"
+                        "        matmul2d_descriptor(16, 16, dynamic_extent), \n"
+                        "        execution_simdgroups<4>> mm; \n"
+                        " \n"
+                        "    auto cT = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB), float>(); \n"
+                        " \n"
+                        "    auto sA = tA.slice(0, 0); \n"
+                        "    auto sB = tB.slice(0, 0); \n"
+                        "    mm.run(sB, sA, cT); \n"
+                        " \n"
+                        "    auto tC = tensor<device float, dextents<int32_t, 2>, tensor_inline>(C, dextents<int32_t, 2>(16, 16)); \n"
+                        " \n"
+                        "    cT.store(tC); \n"
+                        "}";

-                GGML_LOG_INFO("%s: testing tensor API for bfloat support\n", __func__);
-                ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_bf16, false);
-                if (lib == NULL) {
-                    GGML_LOG_WARN("%s: - the tensor API does not support bfloat - disabling bfloat support\n", __func__);
-                    dev->props.has_bfloat = false;
-                } else {
-                    struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil);
-                    if (!ppl.pipeline) {
+                    GGML_LOG_INFO("%s: testing tensor API for bfloat support\n", __func__);
+                    ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_bf16, false);
+                    if (lib == NULL) {
                        GGML_LOG_WARN("%s: - the tensor API does not support bfloat - disabling bfloat support\n", __func__);
                        dev->props.has_bfloat = false;
+                    } else {
+                        struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil);
+                        if (!ppl.pipeline) {
+                            GGML_LOG_WARN("%s: - the tensor API does not support bfloat - disabling bfloat support\n", __func__);
+                            dev->props.has_bfloat = false;
+                        }
+
+                        ggml_metal_library_free(lib);
                    }
-
-                    ggml_metal_library_free(lib);
                }
-            }

-            dev->props.use_residency_sets = true;
+                dev->props.use_residency_sets = true;
 #if defined(GGML_METAL_HAS_RESIDENCY_SETS)
-            dev->props.use_residency_sets = getenv("GGML_METAL_NO_RESIDENCY") == nil;
+                dev->props.use_residency_sets = getenv("GGML_METAL_NO_RESIDENCY") == nil;
 #endif

-            dev->props.use_shared_buffers = dev->props.has_unified_memory;
+                dev->props.use_shared_buffers = dev->props.has_unified_memory;
 #if TARGET_OS_OSX
-            // In case of eGPU, shared memory may be preferable.
-            dev->props.use_shared_buffers |= [dev->mtl_device location] == MTLDeviceLocationExternal;
+                // In case of eGPU, shared memory may be preferable.
+                dev->props.use_shared_buffers |= [dev->mtl_device location] == MTLDeviceLocationExternal;
 #endif
-            if (getenv("GGML_METAL_SHARED_BUFFERS_DISABLE") != NULL) {
-                dev->props.use_shared_buffers = false;
-            }
-            if (getenv("GGML_METAL_SHARED_BUFFERS_ENABLE") != NULL) {
-                dev->props.use_shared_buffers = true;
-            }
+                if (getenv("GGML_METAL_SHARED_BUFFERS_DISABLE") != NULL) {
+                    dev->props.use_shared_buffers = false;
+                }
+                if (getenv("GGML_METAL_SHARED_BUFFERS_ENABLE") != NULL) {
+                    dev->props.use_shared_buffers = true;
+                }

-            dev->props.supports_gpu_family_apple7 = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7];
+                dev->props.supports_gpu_family_apple7 = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7];

-            dev->props.device_id = ggml_metal_device_id_parse([[dev->mtl_device name] UTF8String]);
+                dev->props.device_id = ggml_metal_device_id_parse([[dev->mtl_device name] UTF8String]);

-            dev->props.op_offload_min_batch_size  = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32;
+                dev->props.op_offload_min_batch_size  = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32;

-            dev->props.max_buffer_size            = dev->mtl_device.maxBufferLength;
-            dev->props.max_theadgroup_memory_size = dev->mtl_device.maxThreadgroupMemoryLength;
-            if (@available(macOS 10.12, iOS 16.0, *)) {
-                dev->props.max_working_set_size   = dev->mtl_device.recommendedMaxWorkingSetSize;
-            } else {
-                dev->props.max_working_set_size   = dev->mtl_device.maxBufferLength;
-            }
+                dev->props.max_buffer_size            = dev->mtl_device.maxBufferLength;
+                dev->props.max_theadgroup_memory_size = dev->mtl_device.maxThreadgroupMemoryLength;
+                if (@available(macOS 10.12, iOS 16.0, *)) {
+                    dev->props.max_working_set_size   = dev->mtl_device.recommendedMaxWorkingSetSize;
+                } else {
+                    dev->props.max_working_set_size   = dev->mtl_device.maxBufferLength;
+                }

-            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) {
-                snprintf(dev->props.desc, sizeof(dev->props.desc), "%s (dev p%d/v%d)",
-                         gpu_name, dev->props.device_phys, dev->props.device_virt);
-            } else {
-                snprintf(dev->props.desc, sizeof(dev->props.desc), "%s", gpu_name);
-            }
+                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) {
+                    snprintf(dev->props.desc, sizeof(dev->props.desc), "%s (dev p%d/v%d)",
+                             gpu_name, dev->props.device_phys, dev->props.device_virt);
+                } else {
+                    snprintf(dev->props.desc, sizeof(dev->props.desc), "%s", gpu_name);
+                }

-            dev->library = ggml_metal_library_init(dev);
-            if (!dev->library) {
-                GGML_LOG_ERROR("%s: error: failed to create library\n", __func__);
-            }
+                dev->library = ggml_metal_library_init(dev);
+                if (!dev->library) {
+                    GGML_LOG_ERROR("%s: error: failed to create library\n", __func__);
+                }

-            if (dev->props.use_residency_sets) {
-                dev->rsets = ggml_metal_rsets_init(dev);
-            } else {
-                dev->rsets = nil;
-            }
+                if (dev->props.use_residency_sets) {
+                    dev->rsets = ggml_metal_rsets_init(dev);
+                } else {
+                    dev->rsets = nil;
+                }

-            // print MTL GPU family:
-            GGML_LOG_INFO("%s: GPU name:   %s (%s)\n", __func__, dev->props.name, dev->props.desc);
+                // print MTL GPU family:
+                GGML_LOG_INFO("%s: GPU name:   %s (%s)\n", __func__, dev->props.name, dev->props.desc);

-            // determine max supported GPU family
-            // https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
-            // https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf
-            {
-                for (int i = MTLGPUFamilyApple1 + 20; i >= MTLGPUFamilyApple1; --i) {
-                    if ([dev->mtl_device supportsFamily:i]) {
-                        dev->props.gpu_family = i - (int) MTLGPUFamilyApple1 + 1;
-                        GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyApple%d  (%d)\n", __func__, dev->props.gpu_family, i);
-                        break;
+                // determine max supported GPU family
+                // https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
+                // https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf
+                {
+                    for (int i = MTLGPUFamilyApple1 + 20; i >= MTLGPUFamilyApple1; --i) {
+                        if ([dev->mtl_device supportsFamily:i]) {
+                            dev->props.gpu_family = i - (int) MTLGPUFamilyApple1 + 1;
+                            GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyApple%d  (%d)\n", __func__, dev->props.gpu_family, i);
+                            break;
+                        }
+                    }
+
+                    for (int i = MTLGPUFamilyCommon1 + 5; i >= MTLGPUFamilyCommon1; --i) {
+                        if ([dev->mtl_device supportsFamily:i]) {
+                            GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyCommon%d (%d)\n", __func__, i - (int) MTLGPUFamilyCommon1 + 1, i);
+                            break;
+                        }
+                    }
+
+                    for (int i = MTLGPUFamilyMetal3_GGML + 5; i >= MTLGPUFamilyMetal3_GGML; --i) {
+                        if ([dev->mtl_device supportsFamily:i]) {
+                            GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyMetal%d  (%d)\n", __func__, i - (int) MTLGPUFamilyMetal3_GGML + 3, i);
+                            break;
+                        }
                    }
                }

-                for (int i = MTLGPUFamilyCommon1 + 5; i >= MTLGPUFamilyCommon1; --i) {
-                    if ([dev->mtl_device supportsFamily:i]) {
-                        GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyCommon%d (%d)\n", __func__, i - (int) MTLGPUFamilyCommon1 + 1, i);
-                        break;
-                    }
-                }
-
-                for (int i = MTLGPUFamilyMetal3_GGML + 5; i >= MTLGPUFamilyMetal3_GGML; --i) {
-                    if ([dev->mtl_device supportsFamily:i]) {
-                        GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyMetal%d  (%d)\n", __func__, i - (int) MTLGPUFamilyMetal3_GGML + 3, i);
-                        break;
-                    }
-                }
-            }
-
-            GGML_LOG_INFO("%s: simdgroup reduction   = %s\n", __func__, dev->props.has_simdgroup_reduction ? "true" : "false");
-            GGML_LOG_INFO("%s: simdgroup matrix mul. = %s\n", __func__, dev->props.has_simdgroup_mm        ? "true" : "false");
-            GGML_LOG_INFO("%s: has unified memory    = %s\n", __func__, dev->props.has_unified_memory      ? "true" : "false");
-            GGML_LOG_INFO("%s: has bfloat            = %s\n", __func__, dev->props.has_bfloat              ? "true" : "false");
-            GGML_LOG_INFO("%s: has tensor            = %s\n", __func__, dev->props.has_tensor              ? "true" : "false");
-            GGML_LOG_INFO("%s: use residency sets    = %s\n", __func__, dev->props.use_residency_sets      ? "true" : "false");
-            GGML_LOG_INFO("%s: use shared buffers    = %s\n", __func__, dev->props.use_shared_buffers      ? "true" : "false");
+                GGML_LOG_INFO("%s: simdgroup reduction   = %s\n", __func__, dev->props.has_simdgroup_reduction ? "true" : "false");
+                GGML_LOG_INFO("%s: simdgroup matrix mul. = %s\n", __func__, dev->props.has_simdgroup_mm        ? "true" : "false");
+                GGML_LOG_INFO("%s: has unified memory    = %s\n", __func__, dev->props.has_unified_memory      ? "true" : "false");
+                GGML_LOG_INFO("%s: has bfloat            = %s\n", __func__, dev->props.has_bfloat              ? "true" : "false");
+                GGML_LOG_INFO("%s: has tensor            = %s\n", __func__, dev->props.has_tensor              ? "true" : "false");
+                GGML_LOG_INFO("%s: use residency sets    = %s\n", __func__, dev->props.use_residency_sets      ? "true" : "false");
+                GGML_LOG_INFO("%s: use shared buffers    = %s\n", __func__, dev->props.use_shared_buffers      ? "true" : "false");

 #if TARGET_OS_OSX || (TARGET_OS_IOS && __clang_major__ >= 15)
-            if (@available(macOS 10.12, iOS 16.0, *)) {
-                GGML_LOG_INFO("%s: recommendedMaxWorkingSetSize  = %8.2f MB\n", __func__, dev->props.max_working_set_size / 1e6);
-            }
+                if (@available(macOS 10.12, iOS 16.0, *)) {
+                    GGML_LOG_INFO("%s: recommendedMaxWorkingSetSize  = %8.2f MB\n", __func__, dev->props.max_working_set_size / 1e6);
+                }
 #endif
+            }
        }
    }

@@ -660,6 +660,7 @@ typedef struct {
    uint64_t nb0;
    uint64_t nb1;
    uint64_t nb2;
+    uint64_t nb3;
 } ggml_metal_kargs_conv_transpose_2d;

 typedef struct {
@@ -4645,6 +4645,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
    const int32_t OW = op->ne[0];
    const int32_t OH = op->ne[1];
    const int32_t OC = op->ne[2];
+    const int32_t N  = op->src[1]->ne[3];

    ggml_metal_kargs_conv_transpose_2d args = {
        /*.IC  =*/ IC,
@@ -4657,6 +4658,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
        /*.nb0 =*/ nb0,
        /*.nb1 =*/ nb1,
        /*.nb2 =*/ nb2,
+        /*.nb3 =*/ nb3,
    };

    auto pipeline = ggml_metal_library_get_pipeline_conv_transpose_2d(lib, op);
@@ -4671,7 +4673,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
    const size_t smem = GGML_PAD(KW * KH * sizeof(float), 16);
    ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);

-    ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC, KW, KH, 1);
+    ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC * N, KW, KH, 1);

    return 1;
 }
@@ -66,6 +66,7 @@ fa_vec_cfg_t fa_vec_baseline_cfg(int dk, int dv) {
 // One row per kept bucket, plus per-(dtype,dk,dv) ne11-collapsed domain defaults
 // (ne11_b = FA_VEC_NE11_DEFAULT, ne01_b = domain). To retune or add a device, re-run the
 // sweep and paste its output. See ggml-metal-tuning.h for the row/lookup semantics.
+// ref: https://github.com/ggml-org/llama.cpp/pull/27824
 constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
    { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } },
    { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
@@ -449,6 +450,385 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
    { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
    { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },

+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 3, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 3, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 3, 3 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 3, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 128, 128, 2, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 192, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 256, 256, 2, 3 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 320, 256, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 2, 0 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 3, 0 }, { 4, 1 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 1, 3 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 2, 1 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 2, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 3, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 3, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 2, 0 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 2, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 3, 1 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 1, 4 }, { 1, 2 } },
+
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, 2, 2 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, 2, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 192, 3, 0 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 256, 256, 3, 0 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 512, 512, 3, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 512, 512, 3, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 512, 512, 3, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 512, 512, 3, 3 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 1, 3 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 1, 3 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 1, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 256, 256, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 256, 256, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 320, 256, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 320, 256, 2, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 320, 256, 3, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 3 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 96, 96, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 96, 96, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 96, 96, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 256, 256, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 256, 256, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, 2, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, 3, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 3 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 576, 512, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 576, 512, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
+
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 192, 128, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 256, 256, 3, 0 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 512, 512, 3, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 3 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 4 }, { 1, 2 } },
+
    { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } },
    { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } },
    { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } },
@@ -640,7 +1020,215 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
    { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
    { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },

-        { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 3, 0 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 192, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 3, 2 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 3, 3 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, 1, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 2, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 2, 3 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 2 }, { 4, 1 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 4 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 2, 0 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 2, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 320, 256, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 512, 512, 3, 0 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 512, 512, 2, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 512, 512, 3, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
+
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 1, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 3, 4 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 128, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 128, 3, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 256, 256, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 320, 256, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 320, 256, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 512, 512, 2, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 512, 512, 2, 3 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 512, 512, 3, 3 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, -1, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 3 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 2, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, 1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, 2, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, 2, 3 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
+
+    { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, -1, 1 }, { 2, 4 } },
    { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, 1, 2 }, { 1, 4 } },
    { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, 1, 4 }, { 1, 4 } },
    { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, 2, 2 }, { 4, 4 } },
@@ -366,7 +366,8 @@ kernel void kernel_conv_transpose_2d(

    const int64_t out_x = tgpig[0];
    const int64_t out_y = tgpig[1];
-    const int64_t out_c = tgpig[2];
+    const int64_t batch = tgpig[2] / args.OC;
+    const int64_t out_c = tgpig[2] % args.OC;

    const int64_t kw = tpitg[0];
    const int64_t kh = tpitg[1];
@@ -390,7 +391,7 @@ kernel void kernel_conv_transpose_2d(

        if (in_x >= args.IW) continue;

-        const int64_t input_idx = (args.IW * args.IH) * in_c + (args.IW) * in_y + in_x;
+        const int64_t input_idx = (args.IW * args.IH) * (args.IC * batch + in_c) + (args.IW) * in_y + in_x;
        const int64_t kernel_idx = (args.KH * args.KW * args.OC) * in_c + (args.KH * args.KW) * out_c + (args.KW) * kh + kw;

        v += (float)src0[kernel_idx] * src1[input_idx];
@@ -408,7 +409,7 @@ kernel void kernel_conv_transpose_2d(
            total += shared_sum[i];
        }

-        device float * dst_ptr = (device float *) (dst + out_x*args.nb0 + out_y * args.nb1 + out_c*args.nb2);
+        device float * dst_ptr = (device float *) (dst + batch*args.nb3 + out_c*args.nb2 + out_y * args.nb1 + out_x*args.nb0);
        dst_ptr[0] = total;
    }
 }
@@ -903,6 +903,8 @@ struct ggml_backend_opencl_context {
    cl_kernel kernel_gemv_moe_mxfp4_f32_ns_wimg = nullptr;      // weight-as-texture MoE decode GEMV
    cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a = nullptr;   // dp4a (int8) mxfp4 MoE prefill GEMM
    cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a = nullptr;    // dp4a (int8) q4_0 MoE prefill GEMM
+    cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a_bin = nullptr;   // binary dp4a (int8) mxfp4 MoE prefill GEMM
+    cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a_bin = nullptr;    // binary dp4a (int8) q4_0 MoE prefill GEMM
    cl_kernel kernel_moe_reorder_b;
    cl_kernel kernel_moe_histogram, kernel_moe_scan, kernel_moe_fill, kernel_moe_scatter;
    cl_kernel kernel_moe_scatter_stable = nullptr;   // deterministic slot assignment
@@ -4248,6 +4250,24 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
        GGML_LOG_CONT(".");
    }

+    // gemm_moe_mxfp4_q8_1_dp4a_bin (dp4a prefill GEMM)
+    if (backend_ctx->has_integer_dot) {
+        size_t bin_size = 0;
+        backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin = nullptr;
+
+        if (use_adreno_bin_kernels(backend_ctx)) {
+            const char * kernel_bin = (const char *)backend_ctx->get_adreno_bin_kernel("gemm_moe_mxfp4_q8_1_dp4a_ila", &bin_size);
+            if (kernel_bin && bin_size > 0) {
+                cl_program prog =
+                    build_program_from_binary(backend_ctx->context, backend_ctx->device, kernel_bin, CL_moe_compile_opts, bin_size);
+
+                CL_CHECK((backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin = clCreateKernel(prog, "kernel_gemm_moe_mxfp4_q8_1_dp4a_ila", &err), err));
+                CL_CHECK(clReleaseProgram(prog));
+                GGML_LOG_CONT(".");
+            }
+        }
+    }
+
    // gemm_moe_q4_0_q8_1_dp4a (dp4a prefill GEMM)
    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
@@ -4265,6 +4285,24 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
        GGML_LOG_CONT(".");
    }

+    // gemm_moe_q4_0_q8_1_dp4a_bin (dp4a prefill GEMM)
+    if (backend_ctx->has_integer_dot) {
+        size_t bin_size = 0;
+        backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin = nullptr;
+
+        if (use_adreno_bin_kernels(backend_ctx)) {
+            const char * kernel_bin = (const char *)backend_ctx->get_adreno_bin_kernel("gemm_moe_q4_0_q8_1_dp4a_ila", &bin_size);
+            if (kernel_bin && bin_size > 0) {
+                cl_program prog =
+                    build_program_from_binary(backend_ctx->context, backend_ctx->device, kernel_bin, CL_moe_compile_opts, bin_size);
+
+                CL_CHECK((backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin = clCreateKernel(prog, "kernel_gemm_moe_q4_0_q8_1_dp4a_ila", &err), err));
+                CL_CHECK(clReleaseProgram(prog));
+                GGML_LOG_CONT(".");
+            }
+        }
+    }
+
    // gemm_moe_q8_1_dp4a (generic dp4a MoE GEMM; MOE_QT=80 -> q8_0 expert variant)
    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
@@ -21519,7 +21557,9 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
                    // dot prod has to be available
                    use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
                    // bin kernel takes precedence
-                    use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin == nullptr;
+                    if (backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin == nullptr) {
+                        use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin == nullptr;
+                    }

                    cl_buffer_region region;
                    region.origin = 0;
@@ -21625,6 +21665,10 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,

                        // dp4a GEMM
                        cl_kernel dk = backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a;
+                        if (backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin) {
+                            dk = backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin;
+                        }
+
                        int aidx = 0;
                        CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q4_0->q_img));
                        CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q4_0->d));
@@ -23463,8 +23507,10 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
                        : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
                    // dot prod has to be available
                    use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
-                    // bin kernel takes precedence
-                    use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_mxfp4_f32_ns_bin == nullptr;
+                    // bin kernel takes precedence, dp4a bin kernel has higher priority than normal bin kernel
+                    if (backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin == nullptr) {
+                        use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_mxfp4_f32_ns_bin == nullptr;
+                    }

                    cl_buffer_region region;
                    region.origin = 0;
@@ -23573,6 +23619,10 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,

                        // dp4a GEMM
                        cl_kernel dk = backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a;
+                        if (backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin) {
+                            dk = backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin;
+                        }
+
                        int aidx = 0;
                        CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_mxfp4->q_img));
                        CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_mxfp4->e));
@@ -1,6 +1,8 @@
 find_package(OpenVINO REQUIRED COMPONENTS Runtime Threading)
 find_package(OpenCL REQUIRED)

+message(STATUS "Found OpenVINO: ${OpenVINO_DIR} (found version \"${OpenVINO_VERSION}\")")
+
 file(GLOB_RECURSE GGML_HEADERS_OPENVINO "*.h" "*.hpp")
 file(GLOB_RECURSE GGML_SOURCES_OPENVINO "*.cpp")

@@ -357,6 +357,18 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {
        break;
    }
    case GGML_OP_VIEW: {
+        if (m_is_static && node->src[0] != nullptr &&
+            (node->src[0]->op == GGML_OP_GATED_DELTA_NET || node->src[0]->op == GGML_OP_CONCAT)) {
+            // VIEW slicing a GATED_DELTA_NET combined [attn|state] output, or the conv_input
+            // CONCAT. The consuming CPY/RMS_NORM op recovers the true window at runtime via
+            // ssm_state_size / the fixed conv kernel width, so this VIEW must stay an identity
+            // pass-through of the full source here too (it already is on the dynamic path);
+            // otherwise the generic static-mode Slice below would bake in the *captured*
+            // cgraph's token count, which is wrong once the compiled static model runs with a
+            // different token count (prefill chunk size or 1).
+            op_case = 1;
+            break;
+        }
        if (node->src[0]->op == GGML_OP_VIEW) {
            auto * src = node->src[0];
            if (ggml_nelements(node) != ggml_nelements(src)) {
@@ -408,6 +420,23 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {
        }
        break;
    }
+    case GGML_OP_POOL_2D: {
+        const ggml_op_pool pool_mode = static_cast<ggml_op_pool>(node->op_params[0]);
+        switch (pool_mode) {
+        case GGML_OP_POOL_MAX: {
+            op_case = 1;
+            break;
+        }
+        case GGML_OP_POOL_AVG: {
+            op_case = 2;
+            break;
+        }
+        default:
+            op_case = 0;
+            break;
+        }
+        break;
+    }
    case GGML_OP_CPY: {
        if (node->src[0]->op == GGML_OP_VIEW) {
            if (node->src[0]->src[0]->op == GGML_OP_GATED_DELTA_NET) {
@@ -425,6 +454,31 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {
                   is_kvcache(node->src[1]->view_src, nullptr)) {
            // s_copy defrag remainder writeback: gathered extra state rows copied back into the cache
            op_case = 3;
+        } else if (node->src[1] != nullptr && node->src[1]->op == GGML_OP_VIEW && node->src[1]->view_src != nullptr) {
+            // op_case 5: KV write for decoder self-attention (dynamic write offset)
+            // op_case 6: KV write for encoder self-attn or cross-attn (static offset)
+            const ggml_tensor * kv_buf = node->src[1]->view_src;
+            if (kv_buf->ne[1] == 1 && kv_buf->ne[2] == 1 && kv_buf->ne[3] == 1) {
+                op_case = 6;
+                // Forward-scan the graph for a FLASH_ATTN_EXT that reads from
+                // the same buffer. Having a mask (src[3] != nullptr) implies
+                // decoder self-attention and the write offset is dynamic.
+                for (int i = 0; i < m_cgraph->n_nodes; i++) {
+                    const ggml_tensor * n = m_cgraph->nodes[i];
+                    if (n->op != GGML_OP_FLASH_ATTN_EXT) {
+                        continue;
+                    }
+                    // K (src[1]) and V (src[2]) are 3-D views whose view_src is
+                    // the flat KV buffer we are writing to.
+                    if ((n->src[1] != nullptr && n->src[1]->view_src == kv_buf) ||
+                        (n->src[2] != nullptr && n->src[2]->view_src == kv_buf)) {
+                        if (n->src[3] != nullptr) {
+                            op_case = 5;  // decoder self-attention: mask present
+                        }
+                        break;
+                    }
+                }
+            }
        }
        break;
    }
@@ -448,6 +502,15 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const {
        }
        break;
    }
+    case GGML_OP_FLASH_ATTN_EXT: {
+        if (node->src[1] != nullptr && node->src[1]->op == GGML_OP_VIEW && node->src[1]->view_src != nullptr) {
+            const ggml_tensor * kv_buf = node->src[1]->view_src;
+            if (kv_buf->ne[1] == 1 && kv_buf->ne[2] == 1 && kv_buf->ne[3] == 1) {
+                op_case = (node->src[3] != nullptr) ? 1 : 2;
+            }
+        }
+        break;
+    }
    default:
        break;
    }
@@ -479,23 +542,35 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr

        switch (node->op) {
        case GGML_OP_FLASH_ATTN_EXT:
-            if (node->src[0] == nullptr || node->src[1] == nullptr || node->src[3] == nullptr) {
+            if (node->src[0] == nullptr || node->src[1] == nullptr) {
                return -1;
            }
            switch (node->src[1]->op) {
            case GGML_OP_PERMUTE:
-                // case 0: node op is FLASH_ATTN_EXT, src 1 not null & op is PERMUTE & the permuted tensor src is the view of cache k
-                if (node->src[1]->src[0] != nullptr && node->src[1]->src[0]->op == GGML_OP_VIEW) {
+                // case 0: src[1] is PERMUTE of a cache VIEW, mask required
+                if (node->src[3] != nullptr && node->src[1]->src[0] != nullptr &&
+                    node->src[1]->src[0]->op == GGML_OP_VIEW) {
                    return 0;
                }
                break;
            case GGML_OP_CPY:
-                // case 1: node op is FLASH_ATTN_EXT, src 1 not null & op is CPY & the copied tensor src is PERMUTE & the permuted tensor src is the view of cache k
-                if (node->src[1]->src[0] != nullptr && node->src[1]->src[0]->op == GGML_OP_PERMUTE &&
-                    node->src[1]->src[0]->src[0] != nullptr && node->src[1]->src[0]->src[0]->op == GGML_OP_VIEW) {
+                // case 1: src[1] is CPY of a PERMUTE(VIEW), mask required
+                if (node->src[3] != nullptr && node->src[1]->src[0] != nullptr &&
+                    node->src[1]->src[0]->op == GGML_OP_PERMUTE && node->src[1]->src[0]->src[0] != nullptr &&
+                    node->src[1]->src[0]->src[0]->op == GGML_OP_VIEW) {
                    return 1;
                }
                break;
+            case GGML_OP_VIEW:
+                // cases 4/5/6: whisper - K is a direct non-contiguous VIEW_3D of a KV cache
+                if (node->src[1]->view_src != nullptr) {
+                    if (node->src[3] != nullptr) {
+                        return 4;  // decoder self-attention
+                    } else {
+                        return 5;  // cross-attention or encoder self-attention
+                    };
+                }
+                break;
            default:
                break;
            }
@@ -548,6 +623,18 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr
                cache_k_permute = node->src[0]->src[0]->src[0];
                mask = node->src[1];
                break;
+            case 4:
+            case 5: {
+                // whisper: K is a direct VIEW_3D of the KV buffer, no PERMUTE node
+                auto * cache_k_view = node->src[1];  // VIEW_3D of kv_self.k or kv_cross.k`
+                compute_params.token_len_per_seq = node->src[0]->ne[1];
+                if (attention_pattern_case == 4) {
+                    compute_params.attention_size = cache_k_view->ne[1];
+                } else {
+                    compute_params.attention_size_static = cache_k_view->ne[1];
+                }
+                continue;
+            }
            default:
                break;
            }
@@ -654,10 +741,8 @@ std::pair<ModelParams, ComputeParams> GgmlOvDecoder::compute_llm_params(ggml_cgr
                ComputeParams::RsWriteback writeback;
                writeback.slot_begin = (int) (dest_view->view_offs / row_bytes);
                if (is_conv) {
-                    // conv_input column the copied window starts at
                    writeback.src_begin = (int) (node->src[0]->view_offs / node->src[0]->view_src->nb[0]);
                } else if (is_gdn) {
-                    // first row of the state part of the gated-delta-net output
                    writeback.src_begin = (int) (node->src[0]->view_offs / node->src[0]->view_src->nb[1]);
                }
                compute_params.rs_writebacks[get_tensor_ov_name(cgraph, node)] = writeback;
@@ -718,11 +803,15 @@ ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * op,
    } else if (is_kvcache(input, op)) {
        // kvcache
        input_shape = ov::PartialShape{get_shape(input)};
-        if (!m_is_static) {
+        // Whisper.cpp uses a fixed size 1D KV buffer [N, 1, 1, 1] (GGML) or [1, 1, 1, N] (OV).
+        // the token fill level is handled by token_len_per_seq + dynamic mask input.
+        // skip dynamic dim and stateful reshape for this layout.
+        const bool is_flat_kv = (input->ne[1] == 1 && input->ne[2] == 1 && input->ne[3] == 1);
+        if (!m_is_static && !is_flat_kv) {
            // do not fix ctx size to make llama-bench work across test params
            input_shape[2] = -1;
        }
-        if (is_stateful()) {
+        if (is_stateful() && !is_flat_kv) {
            // Convert stateless KV cache layout [1, 1, seq, n_heads_kv * head_size]
            // to stateful layout [1, seq, n_heads_kv, head_size].
            assert(input_shape.size() == 4 && input_shape[0] == 1 && input_shape[1] == 1 &&
@@ -738,7 +827,9 @@ ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * op,
        input_shape = ov::PartialShape{1, 1, 1, len};

    } else if (is_inp_s_copy(input, op) || is_s_copy_leaf(input)) {
-        input_shape = ov::PartialShape{1, 1, 1, -1};
+        // On NPU the total slot count (n_seq_max) is fixed at translation time, so the s_copy
+        // index list has a static length; on CPU/GPU it may change across compiles (defrag).
+        input_shape = m_is_static ? ov::PartialShape{get_shape(input)} : ov::PartialShape{1, 1, 1, -1};

    } else {
        input_shape = ov::PartialShape{get_shape(input)};
@@ -790,13 +881,16 @@ void GgmlOvDecoder::add_extra_inputs() {
    //     see llama_kv_cache_unified::get_n_kv and llama_kv_cache_unified::get_padding.
    // 2. `n_seq_active` and `seq_active_start`, used in FLASH_ATTN_EXT to indicate the active sequences in the batch

-    auto create_1d_input = [this](const std::string & name, int64_t value) {
-        m_model_extra_inputs[name] = {ov::element::i64, ov::Shape{1}, value, !m_is_static};
+    auto create_1d_input = [this](const std::string & name, int64_t value, bool force_parameter = false) {
+        m_model_extra_inputs[name] = {ov::element::i64, ov::Shape{1}, value, force_parameter || !m_is_static};
    };

    if (m_compute_params.attention_size != -1) {
        create_1d_input("attention_size", m_compute_params.attention_size);
    }
+    if (m_compute_params.attention_size_static != -1) {
+        create_1d_input("attention_size_static", m_compute_params.attention_size_static);
+    }
    if (m_compute_params.attention_size_swa != -1) {
        create_1d_input("attention_size_swa", m_compute_params.attention_size_swa);
    }
@@ -809,17 +903,32 @@ void GgmlOvDecoder::add_extra_inputs() {
    // create_1d_input("token_len", m_compute_params.token_len_per_seq * m_compute_params.n_seq_active);

    if (m_compute_params.cache_rs_reset_idx != -1) {
-        create_1d_input("cache_rs_reset_idx", m_compute_params.cache_rs_reset_idx);
-        create_1d_input("cache_rs_reset_len", m_compute_params.cache_rs_reset_len);
+        // Whether/which cache slot to reset varies per compute call (e.g. a new sequence starting
+        // vs. continued decoding). can_reuse_statically() does not invalidate the cached static
+        // model on ComputeParams changes, so these must stay runtime Parameters even when static
+        // (scale.cpp op_case 1 only uses them in value comparisons, never as Slice bounds, so this
+        // does not reintroduce dynamic shapes).
+        create_1d_input("cache_rs_reset_idx", m_compute_params.cache_rs_reset_idx, /*force_parameter=*/true);
+        create_1d_input("cache_rs_reset_len", m_compute_params.cache_rs_reset_len, /*force_parameter=*/true);
    }

    if (m_compute_params.s_copy_active_slot_len != -1) {
        create_1d_input("s_copy_active_slot_len", m_compute_params.s_copy_active_slot_len);
+        if (m_is_static) {
+            // Number of real tokens in the current prefill chunk. The last chunk is padded with
+            // fabricated token ids; attention masks them out, but the recurrent (GDN/conv) path
+            // would otherwise fold them into cache_r/cache_s permanently. Varies per chunk, so it
+            // must stay a runtime Parameter; it is only compared against a Range or used as Gather
+            // indices, so it does not make any shape dynamic.
+            create_1d_input("chunk_valid_len", get_static_n_tokens(), /*force_parameter=*/true);
+        }
    }

    for (const auto & [node_name, writeback] : m_compute_params.rs_writebacks) {
        create_1d_input("rs_slot_begin_" + node_name, writeback.slot_begin);
-        create_1d_input("rs_src_begin_" + node_name, writeback.src_begin);
+        if (!m_is_static) {
+            create_1d_input("rs_src_begin_" + node_name, writeback.src_begin);
+        }
    }
 }

@@ -1785,13 +1894,23 @@ void GgmlOvDecoder::compute_node_dynamic_dims() {
                    auto dynamic_dim_stride = src_logical_nb[dynamic_dim_idx] / ggml_type_size(node->src[0]->type) *
                                              ggml_type_size(node->type);
                    int matched_dim_count = 0;
+                    int first_matched_dim = -1;
                    for (int i = 0; i < GGML_MAX_DIMS; i++) {
                        if (node->nb[i] == dynamic_dim_stride && node->ne[i] == node->src[0]->ne[dynamic_dim_idx]) {
+                            if (first_matched_dim == -1) {
+                                first_matched_dim = i;
+                            }
                            m_node_dynamic_dims[node] = i;
                            matched_dim_count++;
                        }
                    }
-                    if (matched_dim_count != 1) {
+                    if (matched_dim_count > 1 && node->src[0]->ne[dynamic_dim_idx] == 1) {
+                        // Single-token capture: every trailing dim is size 1 with the same stride, so
+                        // the match is ambiguous. The lowest index is the real axis; the rest are
+                        // ggml's size-1 padding. Bailing out here would bake the captured token count
+                        // into the static prefill model, which then runs with a different one.
+                        m_node_dynamic_dims[node] = first_matched_dim;
+                    } else if (matched_dim_count != 1) {
                        m_node_dynamic_dims[node] = -1;
                        GGML_LOG_WARN("ggml-openvino: cannot determine dynamic dim for CONT node '%s', src[0]: '%s'\n",
                                      node->name, node->src[0]->name);
@@ -47,6 +47,7 @@ struct ComputeParams {
    int seq_active_start = 0;
    int attention_size = -1;
    int attention_size_swa = -1;
+    int attention_size_static = -1;  // encoder/cross-attn KV fill level (whisper)
    int input_len = -1;
    int token_len_per_seq = -1;
    int past_kv_len = -1;
@@ -84,14 +85,15 @@ struct ComputeParams {

    struct RsWriteback {
        int slot_begin = 0;  // first cache slot written by the CPY
-        int src_begin = 0;   // where the copied data starts in the source tensor (in rows of it)
+        int src_begin = 0;   // first source row or column copied by the CPY
    };

    std::map<std::string, RsWriteback> rs_writebacks;
-    // Offsets of the state cache writeback CPY nodes, keyed by node name. They change with the
-    // batch (kv head, active sequence count, token count) and, with rollback enabled
-    // (cparams.n_rs_seq > 0), the conv state is written back once per snapshot slot, each snapshot
-    // taking a different conv_input window. Passed to the cached model as runtime inputs.
+    // Destination slot offset of each state cache writeback CPY node, keyed by node name. It
+    // changes with the batch (kv head, active sequence count) and, with rollback enabled
+    // (cparams.n_rs_seq > 0), the conv state is written back once per snapshot slot. Passed to the
+    // cached model as a runtime input. Dynamic models also receive the source-side offset; static
+    // models use a fixed end-anchored offset in the translator.
 };

 class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder {
@@ -32,6 +32,8 @@ void ggml_openvino_device_config::init() {
        "GGML_OPENVINO_DEVICE",
        "GGML_OPENVINO_CACHE_DIR",
        "GGML_OPENVINO_DEBUG_NODE",
+        "GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR",
+        "GGML_OPENVINO_NPU_COMPILE_CONFIG",
        // Integer values (use ggml_openvino_getenv_int)
        "GGML_OPENVINO_PREFILL_CHUNK_SIZE",
        // Boolean toggles (treated as int flags via ggml_openvino_getenv_int)
@@ -41,6 +43,9 @@ void ggml_openvino_device_config::init() {
        "GGML_OPENVINO_DUMP_IR",
        "GGML_OPENVINO_DEBUG_INPUT",
        "GGML_OPENVINO_DEBUG_OUTPUT",
+        // Force the static (NPU-shape) compute path on any device, e.g. GGML_OPENVINO_DEVICE=CPU,
+        // to test the static-shape translation without NPUW/real NPU hardware in the loop.
+        "GGML_OPENVINO_FORCE_STATIC",
        "GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS",
        "GGML_OPENVINO_ENABLE_CACHE",
        "GGML_OPENVINO_DISABLE_CACHE",
@@ -50,7 +55,7 @@ void ggml_openvino_device_config::init() {
        "GGML_OPENVINO_MEMORY_OPTIMIZE",
        "GGML_OPENVINO_RELEASE_WEIGHTS",
        "GGML_OPENVINO_REDUCE_COMPILE_MEM",
-        "GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR",
+        "GGML_OPENVINO_LOG_UNSUPPORTED_OPS",
    };

    for (const char * const & env_var : env_var_names) {
@@ -85,6 +90,11 @@ void ggml_openvino_device_config::init() {
            compile_config["NPUW_CACHE_DIR"] = cache_dir;
            compile_config.insert(ov::cache_mode(ov::CacheMode::OPTIMIZE_SIZE));
        }
+        const char * compilation_mode_params =
+            ggml_openvino_getenv_str("GGML_OPENVINO_NPU_COMPILE_CONFIG");
+        if (compilation_mode_params && strlen(compilation_mode_params) > 0) {
+            compile_config["NPU_COMPILATION_MODE_PARAMS"] = compilation_mode_params;
+        }
    } else if (cache_dir && strlen(cache_dir) > 0) {
        compile_config.insert(ov::cache_dir(cache_dir));
        compile_config.insert(ov::cache_mode(ov::CacheMode::OPTIMIZE_SIZE));
@@ -908,11 +908,27 @@ static bool has_non_contiguous_view_input(const ggml_tensor * op) {
 }

 static bool is_supported_flash_attn_pattern(const ggml_tensor * op) {
-    // pattern of q,k,v should be q->op==PERMUTE, q->src[0]->op==VIEW, q->src[0]->src[0]->view_src==nullptr
+    // Each Q/K/V input must follow one of:
+    //   PERMUTE -> VIEW  -> base (view_src==nullptr)   (llama KV-cache path)
+    //   PERMUTE -> RESHAPE -> base (view_src==nullptr)  (whisper Q)
+    //   VIEW -> base (view_src==nullptr)                (whisper K/V from kv_pad)
    for (int i = 0; i < 3; i++) {
        const ggml_tensor * src = op->src[i];
-        if (src->op != GGML_OP_PERMUTE || src->src[0] == nullptr || src->src[0]->op != GGML_OP_VIEW ||
-            src->src[0]->src[0] == nullptr || src->src[0]->src[0]->view_src != nullptr) {
+        if (src->op == GGML_OP_PERMUTE) {
+            if (src->src[0] == nullptr) {
+                return false;
+            }
+            if (src->src[0]->op != GGML_OP_VIEW && src->src[0]->op != GGML_OP_RESHAPE) {
+                return false;
+            }
+            if (src->src[0]->src[0] == nullptr || src->src[0]->src[0]->view_src != nullptr) {
+                return false;
+            }
+        } else if (src->op == GGML_OP_VIEW) {
+            if (src->src[0] == nullptr || src->src[0]->view_src != nullptr) {
+                return false;
+            }
+        } else {
            return false;
        }
    }
@@ -1030,18 +1046,29 @@ static bool is_msa_block_mask_expansion(const ggml_tensor * op) {
    return tensor_name_starts_with(src, "msa_block_mask");
 }

-static bool is_op_unsupported_case(const ggml_tensor * op) {
+namespace {
+struct ggml_openvino_op_support {
+    bool is_supported = true;
+    std::string reason;
+
+    operator bool() const {
+        return is_supported;
+    }
+};
+} // namespace
+
+static ggml_openvino_op_support is_op_supported_case(const ggml_tensor * op) {
    if (is_msa_block_mask_expansion(op)) {
-        return true;
+        return {false, "MSA block mask expansion is not supported"};
    }

    switch (op->op) {
    case GGML_OP_CONCAT: {
        if (op->type == GGML_TYPE_I64) {
-            return true;
+            return {false, "CONCAT with I64 type is not supported"};
        }
        if (ggml_openvino_get_device_name() == "GPU" && op->type == GGML_TYPE_BF16 && has_view_op_input(op)) {
-            return true;
+            return {false, "CONCAT with BF16 type and VIEW input is not supported on GPU"};
        }
        break;
    }
@@ -1052,24 +1079,21 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {

        // OpenVINO SET translation currently supports dst layouts that match src0 strides.
        if (op->src[0] == nullptr || nb1 != op->src[0]->nb[1] || nb2 != op->src[0]->nb[2] || nb3 != op->src[0]->nb[3]) {
-            // std::cout << "Unsupported SET op with dst nb1=" << nb1 << ", nb2=" << nb2 << ", nb3=" << nb3
-            //           << " that does not match src0 strides nb[1]="
-            //           << (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[1]) : "null")
-            //           << ", nb[2]=" << (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[2]) : "null")
-            //           << ", nb[3]=" << (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[3]) : "null")
-            //           << std::endl;
-            return true;
+            return {false, "SET op with dst nb1=" + std::to_string(nb1) + ", nb2=" + std::to_string(nb2) + ", nb3=" + std::to_string(nb3) +
+                           " that does not match src0 strides nb[1]=" + (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[1]) : "null") +
+                           ", nb[2]=" + (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[2]) : "null") +
+                           ", nb[3]=" + (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[3]) : "null")};
        }
        break;
    }
    case GGML_OP_GET_ROWS:
    case GGML_OP_SET_ROWS: {
        if (op->ne[3] != 1) {
-            return true;
+            return {false, "GET_ROWS/SET_ROWS with ne[3] != 1 (ne[3]=" + std::to_string(op->ne[3]) + ") is not supported"};
        }
        if (op->op == GGML_OP_GET_ROWS && ggml_openvino_get_device_name() == "GPU" &&
            op->src[0]->type == GGML_TYPE_BF16) {
-            return true;
+            return {false, "GET_ROWS with BF16 src0 is not supported on GPU"};
        }
        if (op->ne[0] == 256 && (op->src[0]->type == GGML_TYPE_Q4_K || op->src[0]->type == GGML_TYPE_Q5_K ||
                                 op->src[0]->type == GGML_TYPE_Q4_1 || op->src[0]->type == GGML_TYPE_Q5_1)) {
@@ -1078,14 +1102,14 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
            // make_int8_weights/make_int4_weights: dequant is done in f16, not f32, to keep the
            // Convert/Subtract/Multiply chain fusable into GatherMatmulCompressed/FullyConnectedCompressed
            // for the shared non-test code paths).
-            return true;
+            return {false, "GET_ROWS/SET_ROWS with ne[0] == 256 and type " + std::string(ggml_type_name(op->src[0]->type)) +
+                           " rejected due to f16-arithmetic dequant rounding errors that intermittently exceed 1e-7 NMSE threshold"};
        }
-
        break;
    }
    case GGML_OP_RESHAPE: {
        if (strncmp(op->name, "ffn_norm_exps", sizeof("ffn_norm_exps") - 1) == 0) {
-            return true;
+            return {false, "RESHAPE for ffn_norm_exps is not supported"};
        }
        break;
    }
@@ -1093,11 +1117,13 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
    case GGML_OP_MUL:
    case GGML_OP_SUB: {
        if (op->src[1]->op == GGML_OP_PERMUTE) {
-            return true;
+            return {false, "ADD/MUL/SUB with PERMUTE src1 is not supported"};
        }
        for (int i = 0; i < 4; i++) {
            if (op->src[0]->ne[i] != op->src[1]->ne[i] && (op->src[0]->ne[i] != 1 && op->src[1]->ne[i] != 1)) {
-                return true;
+                return {false, "ADD/MUL/SUB with incompatible broadcast shapes: src0->ne[" + std::to_string(i) + "]=" +
+                               std::to_string(op->src[0]->ne[i]) + ", src1->ne[" + std::to_string(i) + "]=" +
+                               std::to_string(op->src[1]->ne[i])};
            }
        }
        break;
@@ -1106,7 +1132,7 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
        // Keep support aligned with the CPU backend implementation, which only handles f32 inputs/output and i32 ids.
        if (op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32 ||
            op->src[2]->type != GGML_TYPE_I32) {
-            return true;
+            return {false, "ADD_ID only supports F32 inputs/output and I32 ids"};
        }
        break;
    }
@@ -1116,14 +1142,27 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
        // until the fused GPU kernel is reliable. (falied case llama-arch-test mpt)
        if (ggml_openvino_get_device_name() == "GPU" && op->src[1]->ne[0] == op->ne[0] &&
            op->src[1]->ne[1] == 1 && op->src[1]->ne[2] == 1 && op->src[1]->ne[3] == 1) {
-            return true;
+            return {false, "DIV per-channel scale broadcast is not supported on GPU"};
+        }
+        break;
+    }
+    case GGML_OP_POOL_2D: {
+        const auto& name = ggml_openvino_get_device_name();
+        if (name == "GPU") {
+            const int32_t * params = op->op_params;
+            const int k0 = params[1];
+            const int k1 = params[2];
+            const int p0 = params[5];
+            const int p1 = params[6];
+            if ((p0 > 0 || p1 > 0) && (k0 < 3 || k1 < 3)) {
+                return {false, "POOL_2D with padding and kernel size < 3 is not supported on " + name};
+            }
        }
        break;
    }
    case GGML_OP_SUM_ROWS: {
-        // if the input is PERMUTE skip
        if (op->src[0]->op == GGML_OP_PERMUTE) {
-            return true;
+            return {false, "SUM_ROWS with PERMUTE input is not supported"};
        }
        break;
    }
@@ -1140,54 +1179,51 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
        // accuracy drift in the OpenVINO path. Restrict by scale=1.0 to avoid
        // affecting non-gemma3n models such as Llama-3.2.
        if (fabsf(scale - 1.0f) < 1e-6f && is_gemma3n_flash_attn_pattern(op)) {
-            return true;
+            return {false, "FLASH_ATTN_EXT gemma3n pattern on GPU is not supported"};
        }

        if (op->src[4] != nullptr) {
-            // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with sinks\n");
-            return true;
+            return {false, "FLASH_ATTN_EXT with sinks is not supported"};
        }
        if (!is_supported_flash_attn_pattern(op)) {
-            return true;
+            return {false, "FLASH_ATTN_EXT unsupported attention pattern"};
        }
        if (max_bias > 0) {
-            // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with max_bias > 0\n");
-            return true;
+            return {false, "FLASH_ATTN_EXT with max_bias > 0 (max_bias=" + std::to_string(max_bias) + ") is not supported"};
        }
        if (logit_softcap != 0) {
-            // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with logit_softcap != 0\n");
-            return true;
+            return {false, "FLASH_ATTN_EXT with logit_softcap != 0 (logit_softcap=" + std::to_string(logit_softcap) + ") is not supported"};
        }
        break;
    }
    case GGML_OP_PERMUTE: {
-        if (op->type == GGML_TYPE_BF16) {
-            // err msg: [GPU] Could not find a suitable kernel for transpose
-            // GGML_LOG_WARN("OpenVINO backend does not support PERMUTE with BF16 type\n");
-            return true;
+        if (op->type == GGML_TYPE_BF16 && ggml_openvino_get_device_name() == "GPU") {
+            return {false, "PERMUTE with BF16 type is not supported on GPU"};
        }
        break;
    }
    case GGML_OP_CPY: {
        if (op->src[0]->type == GGML_TYPE_BF16 || op->src[1]->type == GGML_TYPE_BF16) {
-            // GGML_LOG_WARN("OpenVINO backend does not support CPY with non-contiguous data or bf16 types\n");
-            return true;
+            return {false, "CPY with BF16 src type is not supported"};
        }
        // CPY to a quantized destination (e.g. f32 -> q4_0) is numerically unstable with OpenVINO backend.
        if (ggml_is_quantized(op->type)) {
-            return true;
+            return {false, "CPY to quantized destination (e.g. f32 -> q4_0) is numerically unstable"};
        }
        if (ggml_nelements(op->src[0]) != ggml_nelements(op->src[1])) {
-            return true;
+            return {false, "CPY with mismatched element counts is not supported: src0=" + std::to_string(ggml_nelements(op->src[0])) +
+                           " != src1=" + std::to_string(ggml_nelements(op->src[1]))};
        }
        // op test case with non-contiguous src or dst
        if ((op->ne[0] == 3 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2) ||
            (op->ne[0] == 1 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2) ||
            (op->ne[0] == 2 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2)) {
-            return true;
+            return {false, "CPY with non-contiguous shape [" + std::to_string(op->ne[0]) + ", " +
+                           std::to_string(op->ne[1]) + ", " + std::to_string(op->ne[2]) + ", " +
+                           std::to_string(op->ne[3]) + "] is not supported"};
        }
        if (!cpy_output_view_is_supported(op)) {
-            return true;
+            return {false, "CPY with non-contiguous output view is not supported"};
        }
        break;
    }
@@ -1196,13 +1232,14 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
            ggml_is_quantized(op->src[0]->type) && strcmp(op->src[0]->name, "a") == 0 &&
            strcmp(op->src[1]->name, "b") == 0 && op->src[0]->ne[1] == 1 && op->src[1]->ne[1] == 64 &&
            op->src[0]->ne[0] == 256 && op->src[1]->ne[0] == 256) {
-            return true;
+            return {false, "MUL_MAT quantized benchmark test case on GPU is not supported"};
        }
        if (op->src[0]->ne[3] != op->src[1]->ne[3] && op->src[0]->ne[3] != 1 && op->src[1]->ne[3] != 1) {
-            return true;
+            return {false, "MUL_MAT with incompatible broadcast on ne[3]: src0->ne[3]=" + std::to_string(op->src[0]->ne[3]) +
+                           ", src1->ne[3]=" + std::to_string(op->src[1]->ne[3])};
        }
        if (op->src[0]->op == GGML_OP_VIEW && op->src[1]->op == GGML_OP_VIEW) {
-            return true;
+            return {false, "MUL_MAT with both inputs as VIEW is not supported"};
        }
        break;
    }
@@ -1210,16 +1247,17 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
        // Single-expert (or empty) MUL_MAT_ID is a degenerate shape that stresses GatherMatmul edge
        // cases and never occurs in real MoE; let it fall back to CPU.
        if (op->src[0] != nullptr && op->src[0]->ne[2] <= 1) {
-            return true;
+            return {false, "MUL_MAT_ID with single-expert or empty ne[2] <= 1 (ne[2]=" +
+                           std::to_string(op->src[0]->ne[2]) + ") is not supported"};
        }
        if (ggml_openvino_get_device_name() == "GPU" && op->src[0] != nullptr && op->src[0]->type == GGML_TYPE_BF16) {
-            return true;
+            return {false, "MUL_MAT_ID with BF16 weights on GPU is not supported"};
        }
        // GPU MUL_MAT_ID uses a Gather+MatMul fallback because the GPU plugin rejects internal
        // GatherMatmul for these test shapes. Skip cases that would materialize a large selected
        // expert-weight temporary.
        if (ggml_openvino_get_device_name() == "GPU" && mul_mat_id_requires_large_tmp(op)) {
-            return true;
+            return {false, "MUL_MAT_ID requires large temporary on GPU"};
        }
        break;
    }
@@ -1229,51 +1267,46 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
        const int mode = op_params[2];
        if (op_params[15] != 0) {
            // FIXME: support ggml_rope_set_offset
-            return true;
+            return {false, "ggml_rope_set_offset is not supported"};
        }
        if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX && mode != GGML_ROPE_TYPE_IMROPE) {
-            // GGML_LOG_WARN("OpenVINO backend does not support ROPE with mode %d\n", mode);
-            return true;
+            return {false, "ROPE with mode " + std::to_string(mode) + " is not supported"};
        }
        const int64_t head_dim = op->src[0]->ne[0];
        const int64_t rope_dims = n_dims == 0 ? head_dim : n_dims;
        if (rope_dims <= 0 || rope_dims > head_dim || (rope_dims % 2) != 0) {
-            // GGML_LOG_WARN("OpenVINO backend does not support ROPE with n_dims %d and src[0]->ne[0] %ld\n", n_dims,
-            //               op->src[0]->ne[0]);
-            return true;
+            return {false, "ROPE with n_dims=" + std::to_string(n_dims) + ", head_dim=" + std::to_string(head_dim) + " is not supported"};
        }
        if (op->type != GGML_TYPE_F32 && op->type != GGML_TYPE_F16) {
-            // GGML_LOG_WARN("OpenVINO backend does not support ROPE with type %s\n", ggml_type_name(op->type));
-            return true;
+            return {false, "ROPE with type " + std::string(ggml_type_name(op->type)) + " is not supported"};
        }
        if (op->src[0]->op == GGML_OP_VIEW) {
-            if (op->src[0]->view_src->ne[1] != op->src[0]->ne[2]) {
-                // GGML_LOG_WARN(
-                //     "OpenVINO backend does not support ROPE with src[0]->view_src->ne[1] %ld != src[0]->ne[2] "
-                //     "%ld\n",
-                //     op->src[0]->view_src->ne[1], op->src[0]->ne[2]);
-                return true;
+            const struct ggml_tensor * view = op->src[0];
+            const struct ggml_tensor * view_src = view->view_src;
+            if (view_src->ne[1] != view->ne[1] || view_src->ne[2] != view->ne[2] || view_src->ne[3] != view->ne[3]) {
+                return {false, "ROPE with view_src->ne [" + std::to_string(view_src->ne[1]) + ", " +
+                               std::to_string(view_src->ne[2]) + ", " + std::to_string(view_src->ne[3]) +
+                               "] != view->ne [" + std::to_string(view->ne[1]) + ", " +
+                               std::to_string(view->ne[2]) + ", " + std::to_string(view->ne[3]) +
+                               "] is not supported"};
            }
        }
        if (mode == GGML_ROPE_TYPE_IMROPE &&
            (op->src[2] != 0 || ((const float *) op_params)[6] != 1 || ((const float *) op_params)[7] != 0 ||
             ((const float *) op_params)[8] != 1)) {
-            // GGML_LOG_WARN("OpenVINO backend does not support IMROPE with freq_factors, freq_scale, ext_factor, and attn_factor\n");
-            return true;
+            return {false, "IMROPE with freq_factors, freq_scale, ext_factor, and attn_factor is not supported"};
        }
        break;
    }
    case GGML_OP_TRANSPOSE: {
-        // if the type is bf16, will return true
        if (op->type == GGML_TYPE_BF16) {
-            // GGML_LOG_WARN("OpenVINO backend does not support CONT with BF16 type\n");
-            return true;
+            return {false, "TRANSPOSE with BF16 type is not supported"};
        }
        break;
    }
    case GGML_OP_REPEAT: {
        if (ggml_openvino_get_device_name() == "GPU" && op->type == GGML_TYPE_BF16) {
-            return true;
+            return {false, "REPEAT with BF16 type is not supported on GPU"};
        }
        break;
    }
@@ -1285,15 +1318,15 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
        //     return true;
        // }
        if (op->src[2]->op == GGML_OP_PERMUTE) {
-            return true;
+            return {false, "GATED_DELTA_NET with PERMUTE src2 is not supported"};
        }
        // kda (per-key-dimension gating) not supported by fused GatedDeltaNet op
        if (op->src[3]->ne[0] != 1) {
-            return true;
+            return {false, "GATED_DELTA_NET with kda (per-key-dimension gating) is not supported"};
        }
        // K > 1 (multiple state snapshots) not supported by fused op
        if (((const int32_t *) op->op_params)[0] > 1) {
-            return true;
+            return {false, "GATED_DELTA_NET with K > 1 (multiple state snapshots) is not supported"};
        }
        break;
    }
@@ -1307,17 +1340,17 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
        // Skip TOPK_MOE fused tests until it is fully supported.
        // The argsort_top_k VIEW wrapping ARGSORT is named "selected_experts" in test_topk_moe.
        if (strcmp(op->name, "selected_experts") == 0) {
-            return true;
+            return {false, "VIEW for selected_experts (argsort_top_k) is not supported"};
        }
        break;
    }
    default:
        break;
    }
-    return false;
+    return {true, ""};
 }

-static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) {
+static ggml_openvino_op_support ggml_backend_openvino_device_supports_op_impl(ggml_backend_dev_t dev, const ggml_tensor * op) {
    GGML_ASSERT(dev->reg != nullptr);

    static std::unordered_set<ggml_type> supported_types{
@@ -1367,48 +1400,41 @@ static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, con
    case GGML_OP_UNARY: {
        auto supported = supported_unary_ops.find(ggml_get_unary_op(op)) != supported_unary_ops.end();
        if (!supported) {
-            // GGML_LOG_WARN("OpenVINO backend does not support unary op %s\n", ggml_unary_op_name(ggml_get_unary_op(op)));
-            return false;
+            return {false, "unary op " + std::string(ggml_unary_op_name(ggml_get_unary_op(op))) + " has no op translator"};
        }
        if (ggml_get_unary_op(op) == GGML_UNARY_OP_EXP && op->type == GGML_TYPE_F32) {
-            return false;
+            return {false, "UNARY_EXP with F32 type is not supported"};
        }
        break;
    }
    case GGML_OP_GLU: {
        auto supported = supported_glu_ops.find(ggml_get_glu_op(op)) != supported_glu_ops.end();
        if (!supported) {
-            // GGML_LOG_WARN("OpenVINO backend does not support GLU op %s\n", ggml_glu_op_name(ggml_get_glu_op(op)));
-            return false;
+            return {false, "GLU op " + std::string(ggml_glu_op_name(ggml_get_glu_op(op))) + " has no op translator"};
        }
        // if (has_view_op_input(op)) {
-        //     // GGML_LOG_WARN("OpenVINO backend does not support unary op %s with view input\n",
-        //     //               ggml_glu_op_name(ggml_get_glu_op(op)));
-        //     return false;
+        //     return {false, "GLU op " + std::string(ggml_glu_op_name(ggml_get_glu_op(op))) + " with view input is not supported"};
        // }
        if (op->src[1] == nullptr && op->src[0]->ne[0] % 2 != 0) {
            // triggers bug in ov gpu
-            return false;
+            return {false, "GLU op with odd src0 ne[0] and null src1 is not supported"};
        }
        break;
    }
    default: {
        auto supported = supported_ops.find(op->op) != supported_ops.end();
        if (!supported) {
-            // GGML_LOG_WARN("OpenVINO backend does not support op %s\n", ggml_op_name(op->op));
-            return false;
+            return {false, "op " + std::string(ggml_op_name(op->op)) + " has no op translator"};
        }
        static std::set<ggml_op> ops_not_support_view_input{};
        if (ops_not_support_view_input.find(op->op) != ops_not_support_view_input.end() && has_view_op_input(op)) {
-            // GGML_LOG_WARN("OpenVINO backend does not support op %s with view input\n", ggml_op_name(op->op));
-            return false;
+            return {false, "op " + std::string(ggml_op_name(op->op)) + " with VIEW input is not supported"};
        }
    }
    }

    if (supported_types.find(op->type) == supported_types.end()) {
-        // GGML_LOG_WARN("OpenVINO backend does not support tensor type %s\n", ggml_type_name(op->type));
-        return false;
+        return {false, "tensor type " + std::string(ggml_type_name(op->type)) + " is not supported"};
    }
    for (int i = 0; i < GGML_MAX_SRC; i++) {
        auto * src = op->src[i];
@@ -1416,21 +1442,32 @@ static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, con
            break;
        }
        if (supported_types.find(src->type) == supported_types.end()) {
-            // GGML_LOG_WARN("OpenVINO backend does not support tensor type %s\n", ggml_type_name(src->type));
-            return false;
+            return {false, "src[" + std::to_string(i) + "] type " + std::string(ggml_type_name(src->type)) + " is not supported"};
        }
        const bool is_supported_3d_moe_expert =
            op->op == GGML_OP_MUL_MAT_ID && i == 0 && (src->type == GGML_TYPE_MXFP4 || src->ne[3] == 1);
        if (ggml_is_quantized(src->type) && src->ne[2] != 1 && !is_supported_3d_moe_expert) {
-            // GGML_LOG_WARN("OpenVINO backend does not support 3D quantized tensors\n");
-            return false;
+            return {false, "3D quantized tensor for src[" + std::to_string(i) + "] is not supported"};
        }
    }

-    if (is_op_unsupported_case(op)) {
-        return false;
+    auto op_support_case = is_op_supported_case(op);
+    if (!op_support_case.is_supported) {
+        return op_support_case;
    }
-    return true;
+    return {true, ""};
+}
+
+static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) {
+    auto res = ggml_backend_openvino_device_supports_op_impl(dev, op);
+    if (!res.is_supported) {
+        static const bool log_unsupported = ggml_openvino_getenv_int("GGML_OPENVINO_LOG_UNSUPPORTED_OPS") != 0;
+        if (log_unsupported) {
+            GGML_LOG_WARN("OpenVINO op unsupported: op '%s' (%s), type %s: %s\n",
+                          op->name, ggml_op_name(op->op), ggml_type_name(op->type), res.reason.c_str());
+        }
+    }
+    return res.is_supported;
 }

 static bool ggml_backend_openvino_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) {
@@ -3,8 +3,11 @@
 #include "../utils.h"

 #include <climits>
+#include <cstdint>
+#include <cstdio>
 #include <memory>
-#include <vector>
+#include <numeric>
+#include <openvino/frontend/exception.hpp>
 #include <openvino/op/add.hpp>
 #include <openvino/op/concat.hpp>
 #include <openvino/op/constant.hpp>
@@ -12,9 +15,14 @@
 #include <openvino/op/gather.hpp>
 #include <openvino/op/multiply.hpp>
 #include <openvino/op/negative.hpp>
+#include <openvino/op/range.hpp>
 #include <openvino/op/reshape.hpp>
+#include <openvino/op/scatter_update.hpp>
 #include <openvino/op/shape_of.hpp>
 #include <openvino/op/slice.hpp>
+#include <openvino/op/squeeze.hpp>
+#include <openvino/op/subtract.hpp>
+#include <vector>

 namespace ov {
 namespace frontend {
@@ -61,10 +69,27 @@ OutputVector translate_cpy(const NodeContext & context) {
        return rename_outputs_with_suffix({res}, context.get_name());
    }

-    // Recurrent state cache writeback into a slot block of the cache. Where the block starts and
-    // where the copied data starts in the source are runtime inputs, so the cached model works for
-    // any kv head, active sequence count and token count. The result is the full updated cache.
+    // Recurrent state cache writeback into a slot block of the cache. Where the block starts is a
+    // runtime input, so the cached model works for any kv head and active sequence count. The
+    // result is the full updated cache.
    // op_case 1: gated-delta-net state, op_case 2: conv state, op_case 3: defrag remainder.
+    if (op_case == 3) {
+        // With -np 1 (and generally whenever there is no defrag remainder) this GET_ROWS gathers
+        // zero rows: nothing to write back, and the cache is unchanged. NPU rejects zero-size
+        // tensors, so short-circuit instead of building a degenerate Slice/Concat chain.
+        bool is_empty = false;
+        if (input_shape.rank().is_static()) {
+            for (const auto & d : input_shape) {
+                if (d.is_static() && d.get_length() == 0) {
+                    is_empty = true;
+                    break;
+                }
+            }
+        }
+        if (is_empty) {
+            return {context.get_input(1)};
+        }
+    }
    const std::string slot_begin_name = "rs_slot_begin_" + context.get_name();
    const bool slice_assign =
        context.has_input(slot_begin_name) && !context.is_stateful() && (op_case >= 1 && op_case <= 3);
@@ -81,19 +106,49 @@ OutputVector translate_cpy(const NodeContext & context) {
        ov::Output<ov::Node> begin = context.get_input(slot_begin_name);
        auto base = context.get_input(1);
        if (op_case == 1) {
-            // GDN packs [attn | state snapshots]; the state part runs from src_begin to the end.
-            auto src_begin = context.get_input("rs_src_begin_" + context.get_name());
-            auto state_part = std::make_shared<ov::op::v8::Slice>(context.get_input(0), src_begin, int_max, one, axis);
+            ov::Output<ov::Node> state_begin;
+            const std::string src_begin_name = "rs_src_begin_" + context.get_name();
+            if (context.has_input(src_begin_name)) {
+                state_begin = context.get_input(src_begin_name);
+            } else {
+                auto ssm_state_size = context.get_ssm_state_size();
+                if (context.has_input("s_copy_active_slot_len")) {
+                    auto len = context.get_input("s_copy_active_slot_len");
+                    auto state_rows = std::make_shared<ov::op::v1::Multiply>(
+                        ov::op::v0::Constant::create(ov::element::i64, {1}, {ssm_state_size}), len);
+                    state_begin = std::make_shared<ov::op::v0::Negative>(state_rows);
+                } else {
+                    state_begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {-ssm_state_size});
+                }
+            }
+            auto state_part =
+                std::make_shared<ov::op::v8::Slice>(context.get_input(0), state_begin, int_max, one, axis);
            src = std::make_shared<ov::op::v1::Reshape>(state_part, feature, false);
        } else if (op_case == 2) {
-            // conv_input is [previous conv state | new tokens]; copy the conv_kernel_size - 1 wide
-            // window starting at src_begin, which is the snapshot this writeback corresponds to.
+            // conv_input is [previous conv state | new tokens]; the snapshot is the conv_kernel_size - 1
+            // columns ending at the last *valid* token. Gather (rather than Slice) keeps the output
+            // shape static even though the window start is a runtime value.
            auto window_size = (int64_t) input_shape[3].get_length();
-            auto src_begin = context.get_input("rs_src_begin_" + context.get_name());
-            auto src_end = std::make_shared<ov::op::v1::Add>(
-                src_begin, ov::op::v0::Constant::create(ov::element::i64, {1}, {window_size}));
-            auto window = std::make_shared<ov::op::v8::Slice>(context.get_input(0), src_begin, src_end, one,
-                                                              ov::op::v0::Constant::create(ov::element::i64, {1}, {3}));
+            ov::Output<ov::Node> window;
+            auto col_axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {3});
+            const std::string src_begin_name = "rs_src_begin_" + context.get_name();
+            if (context.has_input(src_begin_name)) {
+                auto src_begin = context.get_input(src_begin_name);
+                auto src_end = std::make_shared<ov::op::v1::Add>(
+                    src_begin, ov::op::v0::Constant::create(ov::element::i64, {1}, {window_size}));
+                window = std::make_shared<ov::op::v8::Slice>(context.get_input(0), src_begin, src_end, one, col_axis);
+            } else if (context.has_input("chunk_valid_len")) {
+                std::vector<int64_t> offsets(window_size);
+                std::iota(offsets.begin(), offsets.end(), 0);
+                auto indices = std::make_shared<ov::op::v1::Add>(
+                    ov::op::v0::Constant::create(ov::element::i64, {(size_t) window_size}, offsets),
+                    context.get_input("chunk_valid_len"));
+                window = std::make_shared<ov::op::v8::Gather>(context.get_input(0), indices, col_axis);
+            } else {
+                auto window_begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {-window_size});
+                window =
+                    std::make_shared<ov::op::v8::Slice>(context.get_input(0), window_begin, int_max, one, col_axis);
+            }
            const auto base_shape = base.get_partial_shape();
            FRONT_END_OP_CONVERSION_CHECK(base_shape.rank().is_static() && base_shape.rank().get_length() == 4,
                                          "CPY conv state cache update requires rank-4 base cache");
@@ -157,6 +212,63 @@ OutputVector translate_cpy(const NodeContext & context) {

    auto input = process_view_input_new(context, 0);

+    if (op_case == 5 || op_case == 6) {
+        auto input_shape = context.get_input_shape(0);
+        auto output_shape = context.get_output_shape();
+        auto dst_ggml_shape = context.get_view_input_ggml_shape(1, 0);
+        auto dst_stride = context.get_view_input_stride(1, 0);
+        size_t offset_bytes = context.get_view_input_offset(1, 0);
+        auto n_state = (int64_t) context.get_input_shape(0)[3].get_length();
+        auto n_state_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {n_state});
+        auto kv_buf = context.get_input(1);  // shape {1,1,1,N}
+
+        Output<Node> token_len_per_seq;
+        Output<Node> n_write_dyn;
+        if (context.has_input("token_len_per_seq")) {
+            token_len_per_seq = context.get_input("token_len_per_seq");
+            n_write_dyn = std::make_shared<ov::op::v1::Multiply>(token_len_per_seq, n_state_c);
+        } else {
+            n_write_dyn = ov::op::v0::Constant::create(ov::element::i64, {1}, {(int64_t) dst_ggml_shape[3]});
+        }
+        size_t elem_size = dst_stride[3];
+        FRONT_END_OP_CONVERSION_CHECK(elem_size > 0, "CPY KV cache view update has invalid element size");
+        int64_t start_elem = (int64_t) (offset_bytes / elem_size);
+        // op_case 5: decoder self-attention – write offset advances each step.
+        // op_case 6: encoder self-attn or cross-attn – offset fixed at compile time.
+        const bool is_decoder_self_attn = (op_case == 5);
+        auto ones_c = ov::op::v0::Constant::create(ov::element::i64, {3}, std::vector<int64_t>{1, 1, 1});
+        auto new_shape = std::make_shared<ov::op::v0::Concat>(ov::OutputVector{ones_c, n_write_dyn}, 0);
+
+        auto reshaped = std::make_shared<ov::op::v1::Reshape>(input, new_shape, false);
+        auto data = std::make_shared<ov::op::v0::Convert>(reshaped, context.get_output_type());
+        // Indices [start_elem .. start_elem + n_write) on axis 3 of {1,1,1,N}
+        // For decoder self-attention the write offset advances each step, so compute it
+        // dynamically from the model inputs: start = (attention_size - token_len_per_seq) * n_state.
+        // For encoder self-attn and cross-attn the offset is fixed at graph-compile time.
+        ov::Output<ov::Node> start;
+        if (is_decoder_self_attn && context.has_input("attention_size") && context.has_input("token_len_per_seq")) {
+            auto attention_size_in = context.get_input("attention_size");
+            auto token_len_in = context.get_input("token_len_per_seq");
+            auto past_tokens = std::make_shared<ov::op::v1::Subtract>(attention_size_in, token_len_in);
+            auto new_start = std::make_shared<ov::op::v1::Multiply>(past_tokens, n_state_c);
+            start = std::make_shared<ov::op::v1::Add>(
+                new_start, ov::op::v0::Constant::create(ov::element::i64, {1}, {start_elem}));
+        } else {
+            start = ov::op::v0::Constant::create(ov::element::i64, {1}, {start_elem});
+        }
+        auto start_squeezed = std::make_shared<ov::op::v0::Squeeze>(start);
+        auto end = std::make_shared<ov::op::v1::Add>(start_squeezed, n_write_dyn);
+        auto end_squeezed = std::make_shared<ov::op::v0::Squeeze>(end);
+        auto step = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
+        auto step_squeezed = std::make_shared<ov::op::v0::Squeeze>(step);
+        auto indices =
+            std::make_shared<ov::op::v4::Range>(start_squeezed, end_squeezed, step_squeezed, ov::element::i64);
+        auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {3});
+
+        auto kv_updated = std::make_shared<ov::op::v3::ScatterUpdate>(kv_buf, indices, data, axis);
+        return rename_outputs_with_suffix({kv_updated}, context.get_name());
+    }
+
    if (input_shape != output_shape) {
        auto new_shape = ov::op::v0::Constant::create(
            ov::element::i64, {static_cast<size_t>(output_shape.rank().get_length())}, output_shape.to_shape());
@@ -3,8 +3,8 @@
 #include "../utils.h"
 #include "ggml-openvino/ggml-openvino-extra.h"

+#include <cstddef>
 #include <cstdint>
-#include <cstdlib>
 #include <memory>
 #include <openvino/op/add.hpp>
 #include <openvino/op/broadcast.hpp>
@@ -15,6 +15,7 @@
 #include <openvino/op/multiply.hpp>
 #include <openvino/op/reshape.hpp>
 #include <openvino/op/scaled_dot_product_attention.hpp>
+#include <openvino/op/slice.hpp>
 #include <openvino/op/softmax.hpp>
 #include <openvino/op/transpose.hpp>
 #include <openvino/op/unsqueeze.hpp>
@@ -24,13 +25,62 @@ namespace ov {
 namespace frontend {
 namespace ggml {
 namespace op {
+static ov::Output<ov::Node> reshape_flat_kv(const ov::Output<ov::Node> & kv_flat,
+                                            size_t view_offset_bytes,
+                                            size_t nb1_bytes,
+                                            int64_t n_head,
+                                            int64_t head_size,
+                                            const ov::Output<ov::Node> & attention_size) {
+    int64_t n_state = n_head * head_size;
+    int64_t layer_start_elem = (int64_t) (view_offset_bytes / (nb1_bytes / n_state));
+    // Dynamic slice: [layer_start_elem, layer_start_elem + n_kv * n_state)
+    auto start_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {layer_start_elem});
+    auto n_state_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {n_state});
+    // end = start + attention_size * n_state  (both static + dynamic)
+    auto kv_len_elems = std::make_shared<ov::op::v1::Multiply>(attention_size, n_state_c);
+    auto end_c = std::make_shared<ov::op::v1::Add>(start_c, kv_len_elems);
+    auto step_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
+    auto axis_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {3});
+    auto sliced = std::make_shared<ov::op::v8::Slice>(kv_flat, start_c, end_c, step_c, axis_c);
+
+    // KV cache is laid out as {n_kv, n_head, head_size} in memory
+    // Reshape to {1, n_kv, n_head, head_size}, then transpose to {1, n_head, n_kv, head_size}
+    // as required by SDPA.
+    auto one_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
+    auto n_head_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {n_head});
+    auto head_size_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {head_size});
+    // reshape: {n_kv*n_state} -> {1, n_kv, n_head, head_size}
+    auto new_shape =
+        std::make_shared<ov::op::v0::Concat>(ov::OutputVector{one_c, attention_size, n_head_c, head_size_c}, 0);
+    auto reshaped = std::make_shared<ov::op::v1::Reshape>(sliced, new_shape, false);
+    // transpose: {1, n_kv, n_head, head_size} -> {1, n_head, n_kv, head_size}
+    auto perm = ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3});
+    auto ret = std::make_shared<ov::op::v1::Transpose>(reshaped, perm);
+    return ret;
+}

 OutputVector translate_flash_attn_ext(const NodeContext & context) {
-    num_inputs_check(context, 4, 4);
+    num_inputs_check(context, 3, 4);
+    const bool has_mask = context.get_input_size() == 4;
    auto q_f32 = context.get_input(0);
    auto k = context.get_input(1);
    auto v = context.get_input(2);
-    auto mask = context.get_input(3);
+    const int op_case = context.get_op_case();
+
+    if (op_case == 1 || op_case == 2) {
+        int64_t n_state_head = (int64_t) context.get_view_input_ggml_shape(1, 0)[3];
+        int64_t n_head = (int64_t) context.get_view_input_ggml_shape(1, 0)[1];
+        size_t nb1 = context.get_view_input_stride(1, 0)[2];
+        size_t offset = context.get_view_input_offset(1, 0);
+        ov::Output<ov::Node> attention_size;
+        if (op_case == 1) {
+            attention_size = context.get_input("attention_size");
+        } else {
+            attention_size = context.get_input("attention_size_static");
+        }
+        k = reshape_flat_kv(k, offset, nb1, n_head, n_state_head, attention_size);
+        v = reshape_flat_kv(v, offset, nb1, n_head, n_state_head, attention_size);
+    }

    float * params = reinterpret_cast<float *>(context.get_output_op_params());
    float scale = params[0];
@@ -43,16 +93,19 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) {
    ov::Output<ov::Node> res;

    // For stateful
-    std::string mask_name = "KQ_mask_sliced";
-    if (context.get_input_names()[3].find("swa") != std::string::npos) {
-        mask_name = "KQ_mask_swa_sliced";
-    }
-    if (context.has_input(mask_name)) {
-        mask = context.get_input(mask_name);
-    }
-
-    if (mask.get_element_type() != ov::element::f16) {
-        mask = std::make_shared<ov::op::v0::Convert>(mask, ov::element::f16);
+    ov::Output<ov::Node> mask;
+    if (has_mask) {
+        mask = context.get_input(3);
+        std::string mask_name = "KQ_mask_sliced";
+        if (context.get_input_names()[3].find("swa") != std::string::npos) {
+            mask_name = "KQ_mask_swa_sliced";
+        }
+        if (context.has_input(mask_name)) {
+            mask = context.get_input(mask_name);
+        }
+        if (mask.get_element_type() != ov::element::f16) {
+            mask = std::make_shared<ov::op::v0::Convert>(mask, ov::element::f16);
+        }
    }

    //auto tile_kv = [&](int64_t num_heads, int64_t num_heads_kv, int64_t head_size, ov::Output<Node> kv) {
@@ -108,10 +161,14 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) {
        // get [B, 1, 1, S_q, S_k], which NUMPY-broadcasts cleanly against the
        // [B, num_heads_kv, factor, S_q, S_k] scores: B==B, then 1→num_heads_kv and
        // 1→factor on the head dims.
-        auto mask_unsq1 =
-            std::make_shared<ov::op::v0::Unsqueeze>(mask, ov::op::v0::Constant::create(ov::element::i64, {1}, {2}));
-        // mask_unsq1: [B, 1, 1, S_q, S_k] (rank 5)
-        ov::Output<ov::Node> qk_masked = std::make_shared<ov::op::v1::Add>(qk_scaled, mask_unsq1);
+        ov::Output<ov::Node> qk_masked;
+        if (has_mask) {
+            auto mask_unsq1 =
+                std::make_shared<ov::op::v0::Unsqueeze>(mask, ov::op::v0::Constant::create(ov::element::i64, {1}, {2}));
+            qk_masked = std::make_shared<ov::op::v1::Add>(qk_scaled, mask_unsq1);
+        } else {
+            qk_masked = qk_scaled;
+        }

        auto softmax = std::make_shared<ov::op::v8::Softmax>(qk_masked, /*axis=*/-1);

@@ -164,9 +221,16 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) {
    k = tile_kv(num_heads, num_heads_kv, head_size, k);
    v = tile_kv(num_heads, num_heads_kv, head_size, v);

-    auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(q, k, v, mask, scale_node, false);
-    res = std::make_shared<ov::op::v1::Transpose>(sdpa,
-                                                  ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3}));
+    constexpr auto causal = false;
+    if (has_mask) {
+        auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(q, k, v, mask, scale_node, causal);
+        res = std::make_shared<ov::op::v1::Transpose>(
+            sdpa, ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3}));
+    } else {
+        auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(q, k, v, scale_node, causal);
+        res = std::make_shared<ov::op::v1::Transpose>(
+            sdpa, ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3}));
+    }
    res = std::make_shared<ov::op::v0::Convert>(res, ov::element::f32);
    return rename_outputs_with_suffix({res}, context.get_name());
 }
@@ -7,12 +7,15 @@
 #include <cmath>
 #include <cstdint>
 #include <memory>
+#include <numeric>
 #include <openvino/op/add.hpp>
 #include <openvino/op/broadcast.hpp>
 #include <openvino/op/concat.hpp>
 #include <openvino/op/constant.hpp>
+#include <openvino/op/convert.hpp>
 #include <openvino/op/exp.hpp>
 #include <openvino/op/gather.hpp>
+#include <openvino/op/less.hpp>
 #include <openvino/op/loop.hpp>
 #include <openvino/op/matmul.hpp>
 #include <openvino/op/multiply.hpp>
@@ -80,6 +83,28 @@ OutputVector translate_gated_delta_net(const NodeContext & context) {
    g = std::make_shared<ov::op::v0::Squeeze>(g, ov::op::v0::Constant::create(ov::element::i64, {1}, {3}));
    beta = std::make_shared<ov::op::v0::Squeeze>(beta, ov::op::v0::Constant::create(ov::element::i64, {1}, {3}));

+    if (context.has_input("chunk_valid_len")) {
+        // The last prefill chunk is padded with fabricated tokens. The recurrence is
+        //   S_t = S_{t-1} * exp(g_t) + k_t (x) ((v_t - S_{t-1}^T k_t) * beta_t)
+        // so forcing g = 0 and beta = 0 makes a padded step an exact identity and keeps the final
+        // state equal to the state after the last real token. Attention output at those positions
+        // is garbage but never read.
+        const auto & g_ps = g.get_partial_shape();
+        FRONT_END_OP_CONVERSION_CHECK(g_ps.rank().is_static() && g_ps.rank().get_length() == 3 && g_ps[1].is_static(),
+                                      "GATED_DELTA_NET pad masking requires a static token dimension");
+        const int64_t n_tokens = g_ps[1].get_length();
+        std::vector<int64_t> positions(n_tokens);
+        std::iota(positions.begin(), positions.end(), 0);
+        auto valid = std::make_shared<ov::op::v1::Less>(
+            ov::op::v0::Constant::create(ov::element::i64, {(size_t) n_tokens}, positions),
+            context.get_input("chunk_valid_len"));
+        auto mask = std::make_shared<ov::op::v0::Unsqueeze>(
+            std::make_shared<ov::op::v0::Convert>(valid, g.get_element_type()),
+            ov::op::v0::Constant::create(ov::element::i64, {2}, std::vector<int64_t>{0, 2}));
+        g = std::make_shared<ov::op::v1::Multiply>(g, mask);
+        beta = std::make_shared<ov::op::v1::Multiply>(beta, mask);
+    }
+
    // std::cout << "GatedDeltaNet input shapes: q=" << q.get_partial_shape() << ", k=" << k.get_partial_shape()
    //           << ", v=" << v.get_partial_shape() << ", g=" << g.get_partial_shape()
    //           << ", beta=" << beta.get_partial_shape() << ", state=" << state.get_partial_shape() << std::endl;
@@ -0,0 +1,64 @@
+#include "../node_context.h"
+#include "../op_table.h"
+#include "../utils.h"
+
+#include <memory>
+#include <openvino/core/node_output.hpp>
+#include <openvino/op/constant.hpp>
+#include <openvino/op/multiply.hpp>
+#include <openvino/op/sigmoid.hpp>
+#include <openvino/op/slice.hpp>
+
+namespace ov {
+namespace frontend {
+namespace ggml {
+namespace op {
+
+OutputVector translate_glu_geglu_quick(const NodeContext & context) {
+    num_inputs_check(context, 1, 2);
+
+    ov::Output<ov::Node> src0;
+    ov::Output<ov::Node> src1;
+    if (context.get_input_size() == 2) {
+        src0 = process_view_input_new(context, 0);
+        src1 = process_view_input_new(context, 1);
+    } else {
+        // split along last axis, nc = ne[0] / 2
+        auto combined = process_view_input_new(context, 0);
+        auto combined_shape = combined.get_partial_shape();
+        int64_t last_dim_val = combined_shape[combined_shape.rank().get_length() - 1].get_length();
+        int64_t nc = last_dim_val / 2;
+
+        auto axis  = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1});
+        auto step  = ov::op::v0::Constant::create(ov::element::i64, {1}, {1});
+        auto start0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {0});
+        auto stop0  = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc});
+        auto start1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc});
+        auto stop1  = ov::op::v0::Constant::create(ov::element::i64, {1}, {2 * nc});
+
+        src0 = std::make_shared<ov::op::v8::Slice>(combined, start0, stop0, step, axis);
+        src1 = std::make_shared<ov::op::v8::Slice>(combined, start1, stop1, step, axis);
+    }
+
+    int32_t * params = context.get_output_op_params();
+    const int32_t swapped = params[1];
+    if (swapped) {
+        std::swap(src0, src1);
+    }
+
+    // GELU_QUICK(x) = x * sigmoid(1.702 * x)
+    // Create the constant in the same type as src0 to avoid f16/f32 mismatch.
+    auto input_type = src0.get_element_type();
+    auto coef = ov::op::v0::Constant::create(input_type, ov::Shape{}, {1.702f});
+    auto scaled  = std::make_shared<ov::op::v1::Multiply>(src0, coef);
+    auto sigmoid = std::make_shared<ov::op::v0::Sigmoid>(scaled);
+    auto gated   = std::make_shared<ov::op::v1::Multiply>(src0, sigmoid);
+    auto res     = std::make_shared<ov::op::v1::Multiply>(gated, src1);
+
+    return rename_outputs_with_suffix({res}, context.get_name());
+}
+
+}  // namespace op
+}  // namespace ggml
+}  // namespace frontend
+}  // namespace ov
@@ -0,0 +1,53 @@
+#include "../node_context.h"
+#include "../op_table.h"
+#include "../utils.h"
+
+#include <openvino/op/avg_pool.hpp>
+#include <openvino/op/max_pool.hpp>
+#include <openvino/op/convert.hpp>
+
+namespace ov {
+namespace frontend {
+namespace ggml {
+namespace op {
+
+OutputVector translate_pool_2d(const NodeContext & context) {
+    num_inputs_check(context, 1, 1);
+    const int32_t * params = context.get_output_op_params();
+
+    const int k0 = params[1];
+    const int k1 = params[2];
+    const int s0 = params[3];
+    const int s1 = params[4];
+    const int p0 = params[5];
+    const int p1 = params[6];
+
+    const int op_case = context.get_op_case();
+    ov::Output<Node> input = context.get_input(0);
+    ov::Strides strides{static_cast<size_t>(s1), static_cast<size_t>(s0)};
+    ov::Shape pads_begin{static_cast<size_t>(p1), static_cast<size_t>(p0)};
+    ov::Shape pads_end{static_cast<size_t>(p1), static_cast<size_t>(p0)};
+    ov::Shape kernel{static_cast<size_t>(k1), static_cast<size_t>(k0)};
+    ov::Output<Node> res;
+
+    switch (op_case) {
+    case 1:  // GGML_OP_POOL_MAX
+    {
+        res = std::make_shared<ov::op::v1::MaxPool>(input, strides, pads_begin, pads_end, kernel);
+        break;
+    }
+    case 2:  // GGML_OP_POOL_AVG
+    {
+        res = std::make_shared<ov::op::v1::AvgPool>(input, strides, pads_begin, pads_end, kernel, false);
+        break;
+    }
+    default:
+        break;
+    }
+    return rename_outputs_with_suffix({res}, context.get_name());
+}
+
+}  // namespace op
+}  // namespace ggml
+}  // namespace frontend
+}  // namespace ov
@@ -0,0 +1,36 @@
+#include "../node_context.h"
+#include "../op_table.h"
+#include "../utils.h"
+
+#include <openvino/op/roll.hpp>
+#include <openvino/op/constant.hpp>
+
+namespace ov {
+namespace frontend {
+namespace ggml {
+namespace op {
+
+OutputVector translate_roll(const NodeContext & context) {
+    num_inputs_check(context, 1, 1);
+    const int32_t * params = context.get_output_op_params();
+
+    int64_t s0 = params[0];
+    int64_t s1 = params[1];
+    int64_t s2 = params[2];
+    int64_t s3 = params[3];
+
+    auto input = context.get_input(0);
+
+    auto shift = ov::op::v0::Constant::create(
+        ov::element::i64, ov::Shape{4}, std::vector<int64_t>{s3, s2, s1, s0});
+    auto axes  = ov::op::v0::Constant::create(
+        ov::element::i64, ov::Shape{4}, std::vector<int64_t>{0, 1, 2, 3});
+
+    auto roll = std::make_shared<ov::op::v7::Roll>(input, shift, axes);
+    return rename_outputs_with_suffix({roll}, context.get_name());
+}
+
+}  // namespace op
+}  // namespace ggml
+}  // namespace frontend
+}  // namespace ov
@@ -17,6 +17,13 @@ namespace op {
 OutputVector translate_view(const NodeContext & context) {
    num_inputs_check(context, 1, 1);

+    if (context.get_op_case() == 1) {
+        // Static-mode identity pass-through for VIEWs over a GATED_DELTA_NET combined output or
+        // the conv_input CONCAT; the consuming op (CPY/RMS_NORM) does its own runtime-correct
+        // slicing on the full tensor (see ggml-decoder.cpp compute_op_case, GGML_OP_VIEW).
+        return {context.get_input(0)};
+    }
+
    if (!context.is_static()) {
        // On the stateless/non-static path VIEW is normally a no-op (consumers re-slice).
        // EXCEPTION: the MoE expert aggregation slices each expert plane out of
@@ -10,6 +10,7 @@
 #include <openvino/op/matmul.hpp>
 #include <openvino/op/multiply.hpp>
 #include <openvino/op/negative.hpp>
+#include <openvino/op/relu.hpp>
 #include <openvino/op/sigmoid.hpp>
 #include <openvino/op/subtract.hpp>
 #include <openvino/op/tanh.hpp>
@@ -55,10 +56,12 @@ std::unordered_map<std::string, CreatorFunction> get_supported_ops() {
        {"GGML_UNARY_OP_SIGMOID",   op::translate_1to1_match_1_input<v0::Sigmoid>  },
        {"GGML_UNARY_OP_EXP",       op::translate_1to1_match_1_input<v0::Exp>      },
        {"GGML_UNARY_OP_NEG",       op::translate_1to1_match_1_input<v0::Negative> },
+        {"GGML_UNARY_OP_RELU",      op::translate_1to1_match_1_input<v0::Relu>     },
        {"GGML_OP_VIEW",            op::translate_view                             },
        {"GGML_GLU_OP_SWIGLU",      op::translate_glu_swiglu                       },
        {"GGML_GLU_OP_SWIGLU_OAI",  op::translate_glu_swiglu_oai                   },
        {"GGML_GLU_OP_GEGLU",       op::translate_glu_geglu                        },
+        {"GGML_GLU_OP_GEGLU_QUICK", op::translate_glu_geglu_quick                  },
        {"GGML_OP_SET_ROWS",        op::translate_set_rows                         },
        {"GGML_OP_CPY",             op::translate_cpy                              },
        {"GGML_OP_FLASH_ATTN_EXT",  op::translate_flash_attn_ext                   },
@@ -72,6 +75,8 @@ std::unordered_map<std::string, CreatorFunction> get_supported_ops() {
        {"GGML_OP_DIAG",            op::translate_diag                             },
        {"GGML_OP_TRI",             op::translate_tri                              },
        {"GGML_OP_SET",             op::translate_set                              },
+        {"GGML_OP_POOL_2D",         op::translate_pool_2d                          },
+        {"GGML_OP_ROLL",            op::translate_roll                             },
        // solve_tri has accuracy issues on GPU
        // {"GGML_OP_SOLVE_TRI",       op::translate_solve_tri                        },
    };
@@ -38,6 +38,7 @@ GGML_OP_CONVERTER(translate_view);
 GGML_OP_CONVERTER(translate_glu_swiglu);
 GGML_OP_CONVERTER(translate_glu_swiglu_oai);
 GGML_OP_CONVERTER(translate_glu_geglu);
+GGML_OP_CONVERTER(translate_glu_geglu_quick);
 GGML_OP_CONVERTER(translate_set_rows);
 GGML_OP_CONVERTER(translate_cpy);
 GGML_OP_CONVERTER(translate_argsort);
@@ -53,6 +54,8 @@ GGML_OP_CONVERTER(translate_set);
 GGML_OP_CONVERTER(translate_diag);
 GGML_OP_CONVERTER(translate_tri);
 GGML_OP_CONVERTER(translate_solve_tri);
+GGML_OP_CONVERTER(translate_pool_2d);
+GGML_OP_CONVERTER(translate_roll);

 }  // namespace op

@@ -0,0 +1,212 @@
+#include "fuse_to_conv.h"
+
+#include <openvino/core/graph_util.hpp>
+#include <openvino/core/rt_info.hpp>
+#include <openvino/op/add.hpp>
+#include <openvino/op/constant.hpp>
+#include <openvino/op/convert.hpp>
+#include <openvino/op/convolution.hpp>
+#include <openvino/op/extractimagepatches.hpp>
+#include <openvino/op/matmul.hpp>
+#include <openvino/op/pad.hpp>
+#include <openvino/op/reshape.hpp>
+#include <openvino/op/transpose.hpp>
+#include <openvino/pass/pattern/op/label.hpp>
+#include <openvino/pass/pattern/op/pattern.hpp>
+#include <openvino/pass/pattern/op/wrap_type.hpp>
+
+namespace opp = ov::pass::pattern;
+
+namespace ov {
+namespace frontend {
+namespace ggml {
+namespace pass {
+
+// This pass fuses an IM2COL + MatMul convolution into OpenVINO's Convolution op for performance gains.
+// Reference the im2col.cpp translator for reference on the pattern being matched.
+
+FuseToConv::FuseToConv() {
+    const auto m_wei = opp::any_input();
+    const auto m_act = opp::any_input();
+    const auto m_matmul = opp::wrap_type<ov::op::v0::MatMul>({m_wei, m_act});
+
+    const auto callback = [=](ov::pass::pattern::Matcher & m) {
+        const auto & pm = m.get_pattern_value_map();
+
+        auto matmul_node = ov::as_type_ptr<ov::op::v0::MatMul>(pm.at(m_matmul).get_node_shared_ptr());
+        if (!matmul_node || matmul_node->get_transpose_a() || !matmul_node->get_transpose_b()) {
+            return false;
+        }
+
+        auto trace = matmul_node->input_value(1);
+
+        // Optional Convert
+        if (auto n = ov::as_type_ptr<ov::op::v0::Convert>(trace.get_node_shared_ptr())) {
+            trace = n->input_value(0);
+        }
+
+        for (int i = 0; i < 2; ++i) {
+            auto n = ov::as_type_ptr<ov::op::v1::Reshape>(trace.get_node_shared_ptr());
+            if (!n) {
+                return false;
+            }
+            trace = n->input_value(0);
+        }
+
+        if (auto n = ov::as_type_ptr<ov::op::v1::Transpose>(trace.get_node_shared_ptr())) {
+            trace = n->input_value(0);
+        } else {
+            return false;
+        }
+
+        if (auto n = ov::as_type_ptr<ov::op::v1::Reshape>(trace.get_node_shared_ptr())) {
+            trace = n->input_value(0);
+        } else {
+            return false;
+        }
+
+        if (auto n = ov::as_type_ptr<ov::op::v1::Transpose>(trace.get_node_shared_ptr())) {
+            trace = n->input_value(0);
+        } else {
+            return false;
+        }
+
+        auto eip = ov::as_type_ptr<ov::op::v3::ExtractImagePatches>(trace.get_node_shared_ptr());
+        if (!eip) {
+            return false;
+        }
+        const auto eip_strides = eip->get_strides();  // {stride_h, stride_w}
+        const auto eip_rates = eip->get_rates();      // {dil_h, dil_w}
+
+        auto pad = ov::as_type_ptr<ov::op::v1::Pad>(eip->input_value(0).get_node_shared_ptr());
+        if (!pad) {
+            return false;
+        }
+        auto pads_begin_const =
+            ov::as_type_ptr<ov::op::v0::Constant>(pad->input_value(1).get_node_shared_ptr());
+
+        const auto pads_begin_vals = pads_begin_const->cast_vector<int64_t>();  // {0, 0, pad_h, pad_w}
+        const std::ptrdiff_t pad_h = static_cast<std::ptrdiff_t>(pads_begin_vals[2]);
+        const std::ptrdiff_t pad_w = static_cast<std::ptrdiff_t>(pads_begin_vals[3]);
+
+        auto image_input = pad->input_value(0);  // [N, IC, 1, IW] NCHW
+
+        auto w_trace = matmul_node->input_value(0);
+        if (auto n = ov::as_type_ptr<ov::op::v0::Convert>(w_trace.get_node_shared_ptr())) {
+            w_trace = n->input_value(0);
+        }
+        for (int i = 0; i < 2; ++i) {
+            auto n = ov::as_type_ptr<ov::op::v1::Reshape>(w_trace.get_node_shared_ptr());
+            if (!n) {
+                break;
+            }
+            w_trace = n->input_value(0);
+        }
+
+        auto weight_const = ov::as_type_ptr<ov::op::v0::Constant>(w_trace.get_node_shared_ptr());
+        if (!weight_const) {
+            return false;
+        }
+
+        // Reshape weight to [OC, IC, 1, KW] (OIHW).
+        const auto w_shape = weight_const->get_shape();
+        ov::Shape conv_w_shape;
+        if (w_shape.size() == 3) {
+            conv_w_shape = {w_shape[0], w_shape[1], 1, w_shape[2]};
+        } else if (w_shape.size() == 4) {
+            conv_w_shape = {w_shape[1], w_shape[2], 1, w_shape[3]};
+        } else {
+            return false;
+        }
+
+        auto weight_reshaped = register_new_node<ov::op::v0::Constant>(weight_const->get_element_type(), conv_w_shape,
+                                                                       weight_const->get_data_ptr());
+
+        ov::Output<Node> weight_input = weight_reshaped;
+        if (weight_reshaped->get_element_type() != image_input.get_element_type()) {
+            weight_input = register_new_node<ov::op::v0::Convert>(weight_reshaped, image_input.get_element_type());
+        }
+
+        auto conv = register_new_node<ov::op::v1::Convolution>(
+            image_input, weight_input,
+            ov::Strides{static_cast<size_t>(eip_strides[0]), static_cast<size_t>(eip_strides[1])},
+            ov::CoordinateDiff{pad_h, pad_w}, ov::CoordinateDiff{pad_h, pad_w},
+            ov::Strides{static_cast<size_t>(eip_rates[0]), static_cast<size_t>(eip_rates[1])},
+            ov::op::PadType::EXPLICIT);
+
+        constexpr auto target_type = ov::element::f32;
+        ov::Output<Node> conv_out = conv;
+        if (conv_out.get_element_type() != target_type) {
+            conv_out = register_new_node<ov::op::v0::Convert>(conv_out, target_type);
+        }
+
+        std::shared_ptr<ov::op::v1::Add> add_node;
+        ov::Output<Node> bias_input;
+        for (const auto & consumer_in : matmul_node->output(0).get_target_inputs()) {
+            auto cast = ov::as_type_ptr<ov::op::v0::Convert>(consumer_in.get_node()->shared_from_this());
+            if (!cast) {
+                continue;
+            }
+            for (const auto & add_in : cast->output(0).get_target_inputs()) {
+                auto add = ov::as_type_ptr<ov::op::v1::Add>(add_in.get_node()->shared_from_this());
+                if (!add) {
+                    continue;
+                }
+                for (size_t i = 0; i < 2; ++i) {
+                    if (ov::as_type_ptr<ov::op::v0::Constant>(add->input_value(i).get_node_shared_ptr())) {
+                        bias_input = add->input_value(i);
+                        add_node = add;
+                        break;
+                    }
+                }
+                if (add_node) {
+                    break;
+                }
+            }
+            if (add_node) {
+                break;
+            }
+        }
+
+        ov::Output<Node> final_out;
+        std::shared_ptr<Node> target_node;
+
+        if (add_node) {
+            // Reshape bias [OC, 1] → [1, OC, 1, 1] for NCHW broadcasting.
+            ov::Output<Node> bias = bias_input;
+            if (bias.get_element_type() != target_type) {
+                bias = register_new_node<ov::op::v0::Convert>(bias, target_type);
+            }
+            const auto oc = static_cast<int64_t>(conv_w_shape[0]);
+            auto bias_shape = register_new_node<ov::op::v0::Constant>(ov::element::i64, ov::Shape{4},
+                                                                      std::vector<int64_t>{1, oc, 1, 1});
+            bias = register_new_node<ov::op::v1::Reshape>(bias, bias_shape, false);
+            final_out = register_new_node<ov::op::v1::Add>(conv_out, bias);
+            target_node = add_node;
+        } else {
+            final_out = conv_out;
+            target_node = matmul_node;
+        }
+
+        // Reshape final output back to the target node's original shape if needed.
+        auto orig_shape = target_node->get_output_partial_shape(0);
+        if (orig_shape.is_static() && final_out.get_partial_shape() != orig_shape) {
+            auto shape_const = register_new_node<ov::op::v0::Constant>(ov::element::i64, ov::Shape{orig_shape.size()},
+                                                                       orig_shape.to_shape());
+            final_out = register_new_node<ov::op::v1::Reshape>(final_out, shape_const, false);
+        }
+
+        final_out.get_node_shared_ptr()->set_friendly_name(target_node->get_friendly_name());
+        ov::copy_runtime_info(m.get_matched_nodes(), final_out.get_node_shared_ptr());
+        ov::replace_node(target_node, final_out.get_node_shared_ptr());
+
+        return true;
+    };
+
+    register_matcher(std::make_shared<opp::Matcher>(m_matmul, "ov::frontend::ggml::pass::FuseToConv"), callback);
+}
+
+}  // namespace pass
+}  // namespace ggml
+}  // namespace frontend
+}  // namespace ov
@@ -0,0 +1,17 @@
+#include "openvino/pass/matcher_pass.hpp"
+
+namespace ov {
+namespace frontend {
+namespace ggml {
+namespace pass {
+
+class FuseToConv : public ov::pass::MatcherPass {
+public:
+    OPENVINO_MATCHER_PASS_RTTI("ov::frontend::ggml::pass::FuseToConv")
+    FuseToConv();
+};
+
+}  // namespace pass
+}  // namespace ggml
+}  // namespace frontend
+}  // namespace ov
@@ -5,6 +5,7 @@
 #include "ggml-openvino/openvino/node_context.h"
 #include "ggml-openvino/openvino/utils.h"
 #include "input_model.h"
+#include "pass/fuse_to_conv.h"
 #include "pass/mark_decompression_convert_constant_folding.h"
 #include "pass/mark_dequantization_subgraph.h"
 #include "pass/squeeze_matmul.h"
@@ -109,7 +110,8 @@ ov::pass::MakeStateful::ParamResPairs get_kv_param_res_pairs(
 void add_sliced_mask_stateful(TensorMap & tensor_map) {
    auto create_sliced_mask = [&](const std::string & mask_name, const std::string & sliced_name) {
        if ((tensor_map.find(mask_name) != tensor_map.end()) &&
-            (tensor_map.find("token_len_per_seq") != tensor_map.end())) {
+            (tensor_map.find("token_len_per_seq") != tensor_map.end()) &&
+            (tensor_map.find("inp_pos") != tensor_map.end())) {
            auto token_len_per_seq = tensor_map.at("token_len_per_seq").get_node_shared_ptr();
            auto mask = tensor_map.at(mask_name).get_node_shared_ptr();
            std::shared_ptr<ov::Node> mask_sliced = mask;
@@ -137,6 +139,7 @@ void add_sliced_mask_stateful(TensorMap & tensor_map) {
    };

    create_sliced_mask("self_kq_mask", "KQ_mask_sliced");
+    create_sliced_mask("KQ_mask", "KQ_mask_sliced");
    create_sliced_mask("self_kq_mask_swa", "KQ_mask_swa_sliced");
 }

@@ -395,6 +398,7 @@ std::shared_ptr<Model> TranslateSession::apply_transformations(std::shared_ptr<M
        // is_decompression_multiply() recognizes GatherMatmul as a valid consumer.
        manager.register_pass<ov::pass::MarkDequantization>(
            std::vector<ov::element::Type>{ov::element::u8, ov::element::i8, ov::element::u4, ov::element::i4});
+        manager.register_pass<pass::FuseToConv>();

        if (ggml_model_decoder->is_stateful()) {
            const auto kv_param_res_names = ggml_model_decoder->get_kv_param_res_names();
@@ -72,6 +72,7 @@ OutputVector rename_outputs_with_suffix(const OutputVector & outputs, const std:
        name += "_";
        name += suffix;
        node->set_friendly_name(name);
+        // Uncomment to dump every node's inferred shape (used to hunt down dynamic dims on NPU).
        // std::cout << name << "  " << output.get_partial_shape() << std::endl;
    }
    return outputs;
@@ -16,6 +16,8 @@
 #include <cstdlib>
 #include <cstring>
 #include <fstream>
+#include <functional>
+#include <future>
 #include <iomanip>
 #include <iostream>
 #include <memory>
@@ -48,7 +50,7 @@ enum ggml_status ov_graph_compute(ggml_cgraph * cgraph, ggml_backend_t backend)
            GgmlOvDecoder::dump_cgraph(cgraph, filename);
        }

-        const auto is_static = ggml_openvino_is_npu();
+        const auto is_static = ggml_openvino_is_npu() || ggml_openvino_getenv_int("GGML_OPENVINO_FORCE_STATIC");

        GGML_ASSERT(ctx->runtime_context != nullptr);
        std::shared_ptr<ov_runtime_context> r_ctx = std::static_pointer_cast<ov_runtime_context>(ctx->runtime_context);
@@ -168,13 +170,24 @@ ov::Tensor create_ov_output_tensor(std::shared_ptr<GgmlOvDecoder> ggml_decoder,

    auto output_type = ggml_decoder->get_ov_type(ggml_tensor);
    ov::Shape output_shape;
+    void * output_data = ggml_tensor->data;
    if (ggml_decoder->is_static()) {
        output_shape = infer_request->get_output_tensor(output_index).get_shape();
    } else {
-        output_shape = ggml_decoder->get_shape(ggml_tensor);
+        // For a CPY into a padded view_src (e.g. a padded KV cache buffer), the
+        // OV ScatterUpdate node outputs the full view_src shape, not the CPY node's
+        // own (smaller) shape.  Using the CPY shape here causes set_output_tensor to
+        // fail with a shape-incompatibility error.  Use view_src's shape and data
+        // pointer instead so the OV tensor matches the model output exactly.
+        if (ggml_tensor->op == GGML_OP_CPY && ggml_tensor->view_src != nullptr &&
+            ggml_nbytes(ggml_tensor) != ggml_nbytes(ggml_tensor->view_src)) {
+            output_shape = ggml_decoder->get_shape(ggml_tensor->view_src);
+            output_data = ggml_tensor->view_src->data;
+        } else {
+            output_shape = ggml_decoder->get_shape(ggml_tensor);
+        }
    }
-
-    ov::Tensor output_tensor(output_type, output_shape, ggml_tensor->data);
+    ov::Tensor output_tensor(output_type, output_shape, output_data);
    return output_tensor;
 }

@@ -583,7 +596,9 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptr<o
        return chunk_size;
    };

-    static std::string device = "NPU";
+    // Normally NPU, but honors GGML_OPENVINO_DEVICE so GGML_OPENVINO_FORCE_STATIC can run the
+    // static-shape path on CPU/GPU to isolate translation bugs from NPUW/NPU-driver issues.
+    static std::string device = ggml_openvino_get_device_name();
    static auto is_static = true;
    static auto stateful = false;

@@ -603,7 +618,7 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptr<o
    std::tie(m_params, c_params) = GgmlOvDecoder::compute_llm_params(cgraph, is_static);

    const auto * inp_pos = get_inp_pos_tensor(cgraph);
-    const auto is_prefill = get_is_prefill(inp_pos);
+    const auto is_prefill = get_is_prefill(cgraph, inp_pos);
    graph_key key(cgraph);
    static const bool cache_enabled = !ggml_openvino_getenv_int("GGML_OPENVINO_DISABLE_CACHE");
    bool cache_hit = false;
@@ -687,38 +702,55 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptr<o
                                                                   stateful, false, false, prefill_chunk_size);
        decoder_end_time = ggml_time_us();

-        auto input_model_prefill = std::make_shared<ov::frontend::ggml::InputModel>(ggml_decoder_prefill);
-        auto input_model_decode = std::make_shared<ov::frontend::ggml::InputModel>(ggml_decoder_decode);
+        const bool dump_ir = ggml_openvino_getenv_int("GGML_OPENVINO_DUMP_IR");
+        const auto dump_ir_timestamp = static_cast<long long>(ggml_time_us());

-        auto model_prefill = ov::frontend::ggml::FrontEnd::convert(input_model_prefill);
-        ggml_decoder_prefill->clear_model_weights();
-        auto model_decode = ov::frontend::ggml::FrontEnd::convert(input_model_decode);
-        ggml_decoder_decode->clear_model_weights();
-        conversion_end_time = ggml_time_us();
+        auto build_static_model = [&core, &config, dump_ir, dump_ir_timestamp](
+                          std::shared_ptr<GgmlOvDecoder> decoder,
+                          const char * tag,
+                          std::shared_ptr<ov::Model> & model,
+                          ov::CompiledModel & compiled_model,
+                          std::shared_ptr<ov::InferRequest> & infer_request,
+                          int64_t & local_conversion_end_time,
+                          int64_t & local_compile_end_time) {
+            auto input_model = std::make_shared<ov::frontend::ggml::InputModel>(decoder);
+            model = ov::frontend::ggml::FrontEnd::convert(input_model);
+            decoder->clear_model_weights();
+            local_conversion_end_time = ggml_time_us();

-        if (ggml_openvino_getenv_int("GGML_OPENVINO_DUMP_IR")) {
-            char timestamped_filename[64];
-            auto timestamp = (long long) ggml_time_us();
-            snprintf(timestamped_filename, sizeof(timestamped_filename), "model_prefill_%lld.xml", timestamp);
-            ov::serialize(model_prefill, timestamped_filename);
-            snprintf(timestamped_filename, sizeof(timestamped_filename), "model_decode_%lld.xml", timestamp);
-            ov::serialize(model_decode, timestamped_filename);
-        }
+            if (dump_ir) {
+                char timestamped_filename[64];
+                snprintf(timestamped_filename, sizeof(timestamped_filename), "model_%s_%lld.xml", tag,
+                         dump_ir_timestamp);
+                ov::serialize(model, timestamped_filename);
+            }

+            compiled_model = core.compile_model(model, device, config);
+            infer_request = std::make_shared<ov::InferRequest>(compiled_model.create_infer_request());
+            local_compile_end_time = ggml_time_us();
+        };
+        std::shared_ptr<ov::Model> model_prefill;
+        std::shared_ptr<ov::Model> model_decode;
        ov::CompiledModel compiled_model_prefill;
        ov::CompiledModel compiled_model_decode;
-        auto remote_context = ggml_openvino_get_remote_context();
-        if (remote_context.has_value()) {
-            compiled_model_prefill = core.compile_model(model_prefill, remote_context.value(), config);
-            compiled_model_decode = core.compile_model(model_decode, remote_context.value(), config);
-        } else {
-            compiled_model_prefill = core.compile_model(model_prefill, device, config);
-            compiled_model_decode = core.compile_model(model_decode, device, config);
-        }
-
-        auto infer_request_prefill = std::make_shared<ov::InferRequest>(compiled_model_prefill.create_infer_request());
-        auto infer_request_decode = std::make_shared<ov::InferRequest>(compiled_model_decode.create_infer_request());
-        compile_end_time = ggml_time_us();
+        std::shared_ptr<ov::InferRequest> infer_request_prefill;
+        std::shared_ptr<ov::InferRequest> infer_request_decode;
+        int64_t prefill_conversion_end_time;
+        int64_t decode_conversion_end_time;
+        int64_t prefill_compile_end_time;
+        int64_t decode_compile_end_time;
+        auto prefill_future = std::async(std::launch::async, build_static_model, ggml_decoder_prefill, "prefill",
+                                         std::ref(model_prefill), std::ref(compiled_model_prefill),
+                                         std::ref(infer_request_prefill), std::ref(prefill_conversion_end_time),
+                                         std::ref(prefill_compile_end_time));
+        auto decode_future = std::async(std::launch::async, build_static_model, ggml_decoder_decode, "decode",
+                                        std::ref(model_decode), std::ref(compiled_model_decode),
+                                        std::ref(infer_request_decode), std::ref(decode_conversion_end_time),
+                                        std::ref(decode_compile_end_time));
+        prefill_future.get();
+        decode_future.get();
+        conversion_end_time = std::max(prefill_conversion_end_time, decode_conversion_end_time);
+        compile_end_time = std::max(prefill_compile_end_time, decode_compile_end_time);

        model = is_prefill ? model_prefill : model_decode;
        ggml_decoder = is_prefill ? ggml_decoder_prefill : ggml_decoder_decode;
@@ -742,7 +774,7 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptr<o
    }

    if (is_prefill) {
-        auto inp_len = inp_pos->ne[0];
+        auto inp_len = get_inp_pos_n_tokens(cgraph, inp_pos);
        for (int chunk_index = 0; chunk_index * prefill_chunk_size < inp_len; chunk_index++) {
            for (size_t i = 0; i < ov_input_names_local.size(); i++) {
                auto param_name = ov_input_names_local[i];
@@ -762,6 +794,11 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptr<o
                    continue;
                }
                auto * ggml_tensor = model_output_it->second;
+                if (ggml_nbytes(ggml_tensor) == 0) {
+                    // Zero-row in-place writeback (e.g. the empty s_copy defrag remainder). The OV
+                    // Result is the full cache, so binding it over this 0-byte buffer overflows it.
+                    continue;
+                }
                auto output_tensor = create_ov_output_tensor(ggml_decoder, infer_request, i, ggml_tensor);
                infer_request->set_output_tensor(i, output_tensor);
            }
@@ -798,6 +835,9 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptr<o
                continue;
            }
            auto * ggml_tensor = model_output_it->second;
+            if (ggml_nbytes(ggml_tensor) == 0) {
+                continue;
+            }
            auto output_tensor = create_ov_output_tensor(ggml_decoder, infer_request, i, ggml_tensor);
            infer_request->set_output_tensor(i, output_tensor);
        }
@@ -1074,6 +1114,9 @@ ov::Tensor get_ov_input_tensor(std::shared_ptr<GgmlOvDecoder> ggml_decoder, cons
 ov::Tensor get_ov_input_tensor_static_decode(std::shared_ptr<GgmlOvDecoder> ggml_decoder,
                                             const std::string & param_name) {
    // NPU decoding stage
+    if (ggml_decoder->get_model_extra_inputs().count(param_name)) {
+        return get_ov_input_tensor(ggml_decoder, param_name);
+    }
    const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name);
    const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor);

@@ -1123,14 +1166,30 @@ ov::Tensor get_ov_input_tensor_static_prefill(std::shared_ptr<GgmlOvDecoder> ggm
                                              const std::string & param_name,
                                              int chunk_index) {
    // NPU prompt processing stage
-    const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name);
-    const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor);
-
    const size_t input_len = ggml_decoder->get_input_len();
    const size_t chunk_size = ggml_decoder->m_prefill_chunk_size;
    const size_t chunk_valid_size = std::min(chunk_size, input_len - chunk_index * chunk_size);
    const size_t chunk_pad_size = chunk_size - chunk_valid_size;

+    if (param_name == "chunk_valid_len") {
+        ov::Tensor input_tensor(ov::element::i64, ov::Shape{1});
+        *input_tensor.data<int64_t>() = (int64_t) chunk_valid_size;
+        return input_tensor;
+    }
+    if (chunk_index > 0 && param_name == "cache_rs_reset_len") {
+        // The recurrent-state clear belongs to the start of the sequence. Re-applying it on every
+        // chunk would wipe the state accumulated by the preceding chunks, so disable it (a zero
+        // length makes scale.cpp's keep-mask select every slot) after the first chunk.
+        ov::Tensor input_tensor(ov::element::i64, ov::Shape{1});
+        *input_tensor.data<int64_t>() = 0;
+        return input_tensor;
+    }
+    if (ggml_decoder->get_model_extra_inputs().count(param_name)) {
+        return get_ov_input_tensor(ggml_decoder, param_name);
+    }
+    const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name);
+    const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor);
+
    if (GgmlOvDecoder::is_inp_pos(ggml_tensor, op) && GgmlOvDecoder::get_inp_pos_n_planes(op) > 1) {
        // IMROPE: inp_pos stacks n_planes (t/h/w/e) position planes, each of length
        // input_len; pad every plane independently so they stay aligned to chunk_size.
@@ -1306,7 +1365,7 @@ void print_input_tensor_info(const std::string & name, const ov::Tensor & tensor
              << std::endl;
    switch (tensor.get_element_type()) {
    case ov::element::f32: {
-        if (name.find("self_kq_mask") == std::string::npos) {
+        if (name.find("self_kq_mask") == std::string::npos && name.find("KQ_mask") == std::string::npos) {
            std::cout << *(tensor.data<float>()) << std::endl;
        } else {
            size_t rows = tensor.get_shape()[2];
@@ -1414,8 +1473,24 @@ const ggml_tensor * get_inp_pos_tensor(ggml_cgraph * cgraph) {
    throw std::runtime_error("get_inp_pos_tensor: inp_pos not found in cgraph");
 }

-bool get_is_prefill(const ggml_tensor * inp_pos) {
-    return inp_pos->ne[0] > 1;
+int64_t get_inp_pos_n_tokens(ggml_cgraph * cgraph, const ggml_tensor * inp_pos) {
+    // IMROPE stacks n_planes (t/h/w/e) position planes into inp_pos, so ne[0] is
+    // n_planes * n_tokens. Callers that need a token count must divide the planes out.
+    int n_planes = 1;
+    for (int i = 0; i < cgraph->n_nodes; ++i) {
+        auto * op = cgraph->nodes[i];
+        for (int j = 0; j < GGML_MAX_SRC; ++j) {
+            if (op->src[j] == inp_pos) {
+                n_planes = GgmlOvDecoder::get_inp_pos_n_planes(op);
+                break;
+            }
+        }
+    }
+    return inp_pos->ne[0] / n_planes;
+}
+
+bool get_is_prefill(ggml_cgraph * cgraph, const ggml_tensor * inp_pos) {
+    return get_inp_pos_n_tokens(cgraph, inp_pos) > 1;
 }

 #pragma GCC diagnostic pop
@@ -164,7 +164,9 @@ std::vector<T> pad_input(const ggml_tensor * tensor, size_t padded_rows, size_t

 const ggml_tensor * get_inp_pos_tensor(struct ggml_cgraph * cgraph);

-bool get_is_prefill(const ggml_tensor * inp_pos);
+int64_t get_inp_pos_n_tokens(struct ggml_cgraph * cgraph, const ggml_tensor * inp_pos);
+
+bool get_is_prefill(struct ggml_cgraph * cgraph, const ggml_tensor * inp_pos);

 ov::Tensor get_ov_input_tensor(std::shared_ptr<GgmlOvDecoder> ggml_decoder, const std::string & param_name);
 ov::Tensor get_ov_input_tensor_static_decode(std::shared_ptr<GgmlOvDecoder> ggml_decoder,
@@ -9,6 +9,9 @@
 #include <optional>
 #include <string>
 #include <vector>
+#include <queue>
+#include <condition_variable>
+#include <future>
 #include <memory>
 #include <mutex>
 #include <unordered_map>
@@ -17,6 +20,8 @@
 #include <fstream>
 #include <filesystem>
 #include <algorithm>
+#include <atomic>
+#include <thread>

 static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG");

@@ -72,6 +77,7 @@ enum rpc_cmd {
    RPC_CMD_DEVICE_COUNT,
    RPC_CMD_GRAPH_RECOMPUTE,
    RPC_CMD_MEMSET_TENSOR,
+    RPC_CMD_NONE,
    RPC_CMD_COUNT,
 };

@@ -223,24 +229,24 @@ struct ggml_backend_rpc_buffer_type_context {
    size_t      max_size;
 };

+class rpc_dispatcher;
 struct ggml_backend_rpc_context {
-    std::string endpoint;
-    uint32_t    device;
-    std::string name;
+    std::shared_ptr<rpc_dispatcher> dispatcher;
+    uint32_t                        device;
+    std::string                     name;
 };

 struct ggml_backend_rpc_buffer_context {
-    std::shared_ptr<socket_t> sock;
-    void * base_ptr;
-    uint64_t remote_ptr;
+    std::shared_ptr<rpc_dispatcher>   dispatcher;
+    void                            * base_ptr;
+    uint64_t                          remote_ptr;
 };

 // RPC helper functions

 // Computes FNV-1a hash of the data
-static uint64_t fnv_hash(const uint8_t * data, size_t len) {
+static uint64_t fnv_hash(const uint8_t * data, size_t len, uint64_t hash = 0xcbf29ce484222325ULL) {
    const uint64_t fnv_prime = 0x100000001b3ULL;
-    uint64_t hash = 0xcbf29ce484222325ULL;

    for (size_t i = 0; i < len; ++i) {
        hash ^= data[i];
@@ -357,44 +363,248 @@ static bool negotiate_hello(const std::shared_ptr<socket_t> & sock) {
    return true;
 }

-static std::shared_ptr<socket_t> get_socket(const std::string & endpoint) {
-    static std::mutex mutex;
-    std::lock_guard<std::mutex> lock(mutex);
-    static std::unordered_map<std::string, std::weak_ptr<socket_t>> sockets;
+template <typename T>
+class message_queue {
+public:
+    message_queue() {}

-    auto it = sockets.find(endpoint);
-    if (it != sockets.end()) {
-        if (auto sock = it->second.lock()) {
-            return sock;
+    bool push(const T &value) {
+        std::unique_lock<std::mutex> lock(mutex);
+        if (interrupted) {
+            return false;
        }
+        queue.push(value);
+        cvar.notify_all();
+        return true;
    }
+
+    bool pop(T* out) {
+        std::unique_lock<std::mutex> lock(mutex);
+        cvar.wait(lock, [this] { return !queue.empty() || interrupted; });
+        if (interrupted) {
+            return false;
+        }
+        *out = queue.front();
+        queue.pop();
+        return true;
+    }
+
+    void interrupt() {
+        std::unique_lock<std::mutex> lock(mutex);
+        interrupted = true;
+        lock.unlock();
+        cvar.notify_all();
+    }
+
+private:
+    bool interrupted = false;
+    std::queue<T> queue;
+    std::mutex mutex;
+    std::condition_variable cvar;
+};
+
+class rpc_dispatcher {
+public:
+    rpc_dispatcher() {
+    }
+
+    void send(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size);
+    void send(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size, void * output, size_t output_size);
+    void send_async(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size);
+    void send_async(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size, void * output, size_t output_size);
+
+    ggml_backend_event_t event_new(ggml_backend_dev_t dev);
+    void event_free(ggml_backend_event_t event);
+    void event_synchronize(ggml_backend_event_t event);
+    void event_record(ggml_backend_event_t event);
+    void synchronize();
+
+    void start(const std::string & endpoint);
+    void work();
+
+    ~rpc_dispatcher();
+
+private:
+    struct rpc_msg {
+        rpc_cmd                       cmd;
+        std::shared_ptr<const void>   input;
+        size_t                        input_size;
+        void                        * output;
+        size_t                        output_size;
+        std::promise<void>            completion;
+    };
+    using rpc_msg_ptr   = std::shared_ptr<rpc_msg>;
+    using rpc_msg_queue = message_queue<rpc_msg_ptr>;
+    struct rpc_event {
+        rpc_msg_ptr              msg;
+        std::shared_future<void> sf;
+    };
+    rpc_msg_queue    queue;
+    socket_ptr       sock;
+    std::atomic_bool running;
+    std::thread      thread;
+};
+
+static void rpc_dispatcher_trampoline(rpc_dispatcher * dispatcher)
+{
+    dispatcher->work();
+}
+
+void rpc_dispatcher::send(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size) {
+    auto msg = std::make_shared<rpc_msg>();
+    msg->cmd = cmd;
+    msg->input = input;
+    msg->input_size = input_size;
+    msg->output = nullptr;
+    msg->output_size = 0;
+    GGML_ASSERT(queue.push(msg));
+    auto future = msg->completion.get_future();
+    future.wait();
+}
+
+void rpc_dispatcher::send_async(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size) {
+    auto msg = std::make_shared<rpc_msg>();
+    msg->cmd = cmd;
+    msg->input = input;
+    msg->input_size = input_size;
+    msg->output = nullptr;
+    msg->output_size = 0;
+    GGML_ASSERT(queue.push(msg));
+}
+
+void rpc_dispatcher::send(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size, void * output, size_t output_size) {
+    auto msg = std::make_shared<rpc_msg>();
+    msg->cmd = cmd;
+    msg->input = input;
+    msg->input_size = input_size;
+    msg->output = output;
+    msg->output_size = output_size;
+    GGML_ASSERT(queue.push(msg));
+    auto future = msg->completion.get_future();
+    future.wait();
+}
+
+void rpc_dispatcher::send_async(enum rpc_cmd cmd, std::shared_ptr<const void> input, size_t input_size, void * output, size_t output_size) {
+    auto msg = std::make_shared<rpc_msg>();
+    msg->cmd = cmd;
+    msg->input = input;
+    msg->input_size = input_size;
+    msg->output = output;
+    msg->output_size = output_size;
+    GGML_ASSERT(queue.push(msg));
+}
+
+ggml_backend_event_t rpc_dispatcher::event_new(ggml_backend_dev_t dev) {
+    rpc_event * ev = new rpc_event;
+    ev->msg = std::make_shared<rpc_msg>();
+    ev->msg->cmd = RPC_CMD_NONE;
+    ev->sf = ev->msg->completion.get_future().share();
+    GGML_ASSERT(queue.push(ev->msg));
+    return new ggml_backend_event {
+        /* .device  = */ dev,
+        /* .context = */ ev,
+    };
+}
+
+void rpc_dispatcher::event_free(ggml_backend_event_t event) {
+    rpc_event * ev = (rpc_event *)event->context;
+    delete ev;
+}
+
+void rpc_dispatcher::event_synchronize(ggml_backend_event_t event) {
+    rpc_event * ev = (rpc_event *)event->context;
+    ev->sf.wait();
+}
+
+void rpc_dispatcher::event_record(ggml_backend_event_t event) {
+    rpc_event * ev = (rpc_event *)event->context;
+    ev->msg = std::make_shared<rpc_msg>();
+    ev->msg->cmd = RPC_CMD_NONE;
+    ev->sf = ev->msg->completion.get_future().share();
+    GGML_ASSERT(queue.push(ev->msg));
+}
+
+void rpc_dispatcher::synchronize() {
+    // to ensure all messages are processed, submit dummy message and wait for it to complete
+    auto msg = std::make_shared<rpc_msg>();
+    msg->cmd = RPC_CMD_NONE;
+    GGML_ASSERT(queue.push(msg));
+    msg->completion.get_future().wait();
+}
+
+void rpc_dispatcher::start(const std::string & endpoint) {
    std::string host;
    int port;
    if (!parse_endpoint(endpoint, host, port)) {
-        GGML_LOG_ERROR("Failed to parse endpoint: %s\n", endpoint.c_str());
-        return nullptr;
+        GGML_ABORT("Failed to parse endpoint: %s\n", endpoint.c_str());
+    }
+    if (!rpc_transport_init()) {
+        GGML_ABORT("RPC transport initialization failed\n");
    }

-    if (!rpc_transport_init()) {
-        return nullptr;
-    }
-    auto sock = socket_t::connect(host.c_str(), port);
+    sock = socket_t::connect(host.c_str(), port);
    if (sock == nullptr) {
-        return nullptr;
+        GGML_ABORT("Failed to connect to %s\n", endpoint.c_str());
    }
    if (!negotiate_hello(sock)) {
-        return nullptr;
+        GGML_ABORT("RPC handshake failed for %s\n", endpoint.c_str());
    }
    LOG_DBG("[%s] connected to %s\n", __func__, endpoint.c_str());
-    sockets[endpoint] = sock;
-    return sock;
+    running = true;
+    thread = std::thread(rpc_dispatcher_trampoline, this);
+}
+
+void rpc_dispatcher::work() {
+    while (running) {
+        rpc_msg_ptr msg_ptr;
+        if (!queue.pop(&msg_ptr)) {
+            break;
+        }
+        if (msg_ptr->cmd != RPC_CMD_NONE) {
+            if (msg_ptr->output) {
+                bool status = send_rpc_cmd(sock, msg_ptr->cmd, msg_ptr->input.get(), msg_ptr->input_size, msg_ptr->output, msg_ptr->output_size);
+                RPC_STATUS_ASSERT(status);
+            } else {
+                bool status = send_rpc_cmd(sock, msg_ptr->cmd, msg_ptr->input.get(), msg_ptr->input_size);
+                RPC_STATUS_ASSERT(status);
+            }
+        }
+        msg_ptr->completion.set_value();
+    }
+}
+
+rpc_dispatcher::~rpc_dispatcher() {
+    running = false;
+    queue.interrupt();
+    sock = nullptr;
+    if (thread.joinable()) {
+        thread.join();
+    }
+}
+
+static std::shared_ptr<rpc_dispatcher> get_dispatcher(const std::string & endpoint) {
+    static std::mutex mutex;
+    std::lock_guard<std::mutex> lock(mutex);
+    static std::unordered_map<std::string, std::weak_ptr<rpc_dispatcher>> dispatchers;
+
+    auto it = dispatchers.find(endpoint);
+    if (it != dispatchers.end()) {
+        if (auto dispatcher = it->second.lock()) {
+            return dispatcher;
+        }
+    }
+
+    auto dispatcher = std::make_shared<rpc_dispatcher>();
+    dispatcher->start(endpoint);
+    dispatchers[endpoint] = dispatcher;
+    return dispatcher;
 }

 static void ggml_backend_rpc_buffer_free_buffer(ggml_backend_buffer_t buffer) {
    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
-    rpc_msg_free_buffer_req request = {ctx->remote_ptr};
-    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_FREE_BUFFER, &request, sizeof(request), nullptr, 0);
-    RPC_STATUS_ASSERT(status);
+    auto request = std::make_shared<rpc_msg_free_buffer_req>();
+    request->remote_ptr = ctx->remote_ptr;
+    ctx->dispatcher->send(RPC_CMD_FREE_BUFFER, request, sizeof(*request));
    delete ctx;
 }

@@ -403,10 +613,10 @@ static void * ggml_backend_rpc_buffer_get_base(ggml_backend_buffer_t buffer) {
    if (ctx->base_ptr != nullptr) {
        return ctx->base_ptr;
    }
-    rpc_msg_buffer_get_base_req request = {ctx->remote_ptr};
+    auto request = std::make_shared<rpc_msg_buffer_get_base_req>();
+    request->remote_ptr = ctx->remote_ptr;
    rpc_msg_buffer_get_base_rsp response;
-    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_BUFFER_GET_BASE, &request, sizeof(request), &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+    ctx->dispatcher->send(RPC_CMD_BUFFER_GET_BASE, request, sizeof(*request), &response, sizeof(response));
    ctx->base_ptr = reinterpret_cast<void *>(response.base_ptr);
    return ctx->base_ptr;
 }
@@ -463,12 +673,9 @@ static enum ggml_status ggml_backend_rpc_buffer_init_tensor(ggml_backend_buffer_
    // Due to bandwidth constraints, we only call the server init tensor functions if necessary.
    // In particular, only quantized tensors need padding
    if (ggml_is_quantized(tensor->type) && (tensor->ne[0] % 512 != 0) && (tensor->view_src == nullptr)) {
-        rpc_msg_init_tensor_req request;
-
-        request.tensor = serialize_tensor(tensor);
-
-        bool status = send_rpc_cmd(ctx->sock, RPC_CMD_INIT_TENSOR, &request, sizeof(request), nullptr, 0);
-        RPC_STATUS_ASSERT(status);
+        auto request = std::make_shared<rpc_msg_init_tensor_req>();
+        request->tensor = serialize_tensor(tensor);
+        ctx->dispatcher->send(RPC_CMD_INIT_TENSOR, request, sizeof(*request));
    }
    return GGML_STATUS_SUCCESS;
 }
@@ -476,27 +683,24 @@ static enum ggml_status ggml_backend_rpc_buffer_init_tensor(ggml_backend_buffer_
 static void ggml_backend_rpc_buffer_memset_tensor(
        ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
-    rpc_msg_memset_tensor_req request = {
-        /* .tensor = */ serialize_tensor(tensor),
-        /* .offset = */ offset,
-        /* .size   = */ size,
-        /* .value  = */ value,
-    };
-    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_MEMSET_TENSOR, &request, sizeof(request), nullptr, 0);
-    RPC_STATUS_ASSERT(status);
+    auto request = std::make_shared<rpc_msg_memset_tensor_req>();
+    request->tensor = serialize_tensor(tensor);
+    request->offset = offset;
+    request->size   = size;
+    request->value  = value;
+    ctx->dispatcher->send(RPC_CMD_MEMSET_TENSOR, request, sizeof(*request));
 }

 static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
    rpc_tensor rpc_tensor = serialize_tensor(tensor);
    if (size > HASH_THRESHOLD) {
-        rpc_msg_set_tensor_hash_req request;
-        request.tensor = rpc_tensor;
-        request.offset = offset;
-        request.hash = fnv_hash((const uint8_t*)data, size);
+        auto request = std::make_shared<rpc_msg_set_tensor_hash_req>();
+        request->tensor = rpc_tensor;
+        request->offset = offset;
+        request->hash = fnv_hash((const uint8_t*)data, size);
        rpc_msg_set_tensor_hash_rsp response;
-        bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR_HASH, &request, sizeof(request), &response, sizeof(response));
-        RPC_STATUS_ASSERT(status);
+        ctx->dispatcher->send(RPC_CMD_SET_TENSOR_HASH, request, sizeof(*request), &response, sizeof(response));
        if (response.result) {
            // the server has the same data, no need to send it
            return;
@@ -504,22 +708,21 @@ static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggm
    }
    // input serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes)
    size_t input_size = sizeof(rpc_tensor) + sizeof(uint64_t) + size;
-    std::vector<uint8_t> input(input_size, 0);
-    memcpy(input.data(), &rpc_tensor, sizeof(rpc_tensor));
-    memcpy(input.data() + sizeof(rpc_tensor), &offset, sizeof(offset));
-    memcpy(input.data() + sizeof(rpc_tensor) + sizeof(offset), data, size);
-    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR, input.data(), input.size());
-    RPC_STATUS_ASSERT(status);
+    uint8_t * input = new uint8_t[input_size]();
+    memcpy(input, &rpc_tensor, sizeof(rpc_tensor));
+    memcpy(input + sizeof(rpc_tensor), &offset, sizeof(offset));
+    memcpy(input + sizeof(rpc_tensor) + sizeof(offset), data, size);
+    std::shared_ptr<uint8_t> input_ptr(input, std::default_delete<uint8_t[]>());
+    ctx->dispatcher->send(RPC_CMD_SET_TENSOR, input_ptr, input_size);
 }

 static void ggml_backend_rpc_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
-    rpc_msg_get_tensor_req request;
-    request.tensor = serialize_tensor(tensor);
-    request.offset = offset;
-    request.size = size;
-    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_GET_TENSOR, &request, sizeof(request), data, size);
-    RPC_STATUS_ASSERT(status);
+    auto request = std::make_shared<rpc_msg_get_tensor_req>();
+    request->tensor = serialize_tensor(tensor);
+    request->offset = offset;
+    request->size = size;
+    ctx->dispatcher->send(RPC_CMD_GET_TENSOR, request, sizeof(*request), data, size);
 }

 static bool ggml_backend_rpc_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) {
@@ -529,16 +732,15 @@ static bool ggml_backend_rpc_buffer_cpy_tensor(ggml_backend_buffer_t buffer, con
        ggml_backend_rpc_buffer_context * src_ctx = (ggml_backend_rpc_buffer_context *)src_buffer->context;
        ggml_backend_buffer_t dst_buffer = dst->buffer;
        ggml_backend_rpc_buffer_context * dst_ctx = (ggml_backend_rpc_buffer_context *)dst_buffer->context;
-        if (src_ctx->sock != dst_ctx->sock) {
+        if (src_ctx->dispatcher != dst_ctx->dispatcher) {
            return false;
        }
        ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
-        rpc_msg_copy_tensor_req request;
-        request.src = serialize_tensor(src);
-        request.dst = serialize_tensor(dst);
+        auto request = std::make_shared<rpc_msg_copy_tensor_req>();
+        request->src = serialize_tensor(src);
+        request->dst = serialize_tensor(dst);
        rpc_msg_copy_tensor_rsp response;
-        bool status = send_rpc_cmd(ctx->sock, RPC_CMD_COPY_TENSOR, &request, sizeof(request), &response, sizeof(response));
-        RPC_STATUS_ASSERT(status);
+        ctx->dispatcher->send(RPC_CMD_COPY_TENSOR, request, sizeof(*request), &response, sizeof(response));
        return response.result;
    }
    return false;
@@ -546,9 +748,10 @@ static bool ggml_backend_rpc_buffer_cpy_tensor(ggml_backend_buffer_t buffer, con

 static void ggml_backend_rpc_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
-    rpc_msg_buffer_clear_req request = {ctx->remote_ptr, value};
-    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_BUFFER_CLEAR, &request, sizeof(request), nullptr, 0);
-    RPC_STATUS_ASSERT(status);
+    auto request = std::make_shared<rpc_msg_buffer_clear_req>();
+    request->remote_ptr = ctx->remote_ptr;
+    request->value = value;
+    ctx->dispatcher->send(RPC_CMD_BUFFER_CLEAR, request, sizeof(*request));
 }

 static ggml_backend_buffer_i ggml_backend_rpc_buffer_interface = {
@@ -572,15 +775,17 @@ static const char * ggml_backend_rpc_buffer_type_name(ggml_backend_buffer_type_t

 static ggml_backend_buffer_t ggml_backend_rpc_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
    ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context;
-    rpc_msg_alloc_buffer_req request = {buft_ctx->device, size};
+    auto request = std::make_shared<rpc_msg_alloc_buffer_req>();
+    request->device = buft_ctx->device;
+    request->size = size;
    rpc_msg_alloc_buffer_rsp response;
-    auto sock = get_socket(buft_ctx->endpoint);
-    bool status = send_rpc_cmd(sock, RPC_CMD_ALLOC_BUFFER, &request, sizeof(request), &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+
+    auto dispatcher = get_dispatcher(buft_ctx->endpoint);
+    dispatcher->send(RPC_CMD_ALLOC_BUFFER, request, sizeof(*request), &response, sizeof(response));
    if (response.remote_ptr != 0) {
        ggml_backend_buffer_t buffer = ggml_backend_buffer_init(buft,
            ggml_backend_rpc_buffer_interface,
-            new ggml_backend_rpc_buffer_context{sock, nullptr, response.remote_ptr},
+            new ggml_backend_rpc_buffer_context{dispatcher, nullptr, response.remote_ptr},
            response.remote_size);
        return buffer;
    } else {
@@ -588,11 +793,11 @@ static ggml_backend_buffer_t ggml_backend_rpc_buffer_type_alloc_buffer(ggml_back
    }
 }

-static size_t get_alignment(const std::shared_ptr<socket_t> & sock, uint32_t device) {
-    rpc_msg_get_alignment_req request = {device};
+static size_t get_alignment(const std::shared_ptr<rpc_dispatcher> & dispatcher, uint32_t device) {
+    auto request = std::make_shared<rpc_msg_get_alignment_req>();
+    request->device = device;
    rpc_msg_get_alignment_rsp response;
-    bool status = send_rpc_cmd(sock, RPC_CMD_GET_ALIGNMENT, &request, sizeof(request), &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+    dispatcher->send(RPC_CMD_GET_ALIGNMENT, request, sizeof(*request), &response, sizeof(response));
    return response.alignment;
 }

@@ -601,11 +806,11 @@ static size_t ggml_backend_rpc_buffer_type_get_alignment(ggml_backend_buffer_typ
    return buft_ctx->alignment;
 }

-static size_t get_max_size(const std::shared_ptr<socket_t> & sock, uint32_t device) {
-    rpc_msg_get_max_size_req request = {device};
+static size_t get_max_size(const std::shared_ptr<rpc_dispatcher> & dispatcher, uint32_t device) {
+    auto request = std::make_shared<rpc_msg_get_max_size_req>();
+    request->device = device;
    rpc_msg_get_max_size_rsp response;
-    bool status = send_rpc_cmd(sock, RPC_CMD_GET_MAX_SIZE, &request, sizeof(request), &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+    dispatcher->send(RPC_CMD_GET_MAX_SIZE, request, sizeof(*request), &response, sizeof(response));
    return response.max_size;
 }

@@ -628,23 +833,63 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty

    if (rpc_get) {
        ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context;
-        auto sock = get_socket(buft_ctx->endpoint);

-        rpc_msg_get_alloc_size_req request = {
-            /*.device =*/ buft_ctx->device,
-            /*.tensor =*/ serialize_tensor(tensor),
-            /*.srcs   =*/ {},
+        // Cache key for calls to read the alloc_size.
+        // We deliberately exclude src tensor dimensions from the key because:
+        // 1. For CPU backends, alloc_size = ggml_nbytes(output) regardless of src shapes
+        // 2. For GPU backends, the reservation graph uses max dimensions, so the
+        //    cached value from reservation is always >= any subsequent request
+        // 3. Including src dims causes cache misses per-ubatch (e.g. growing KV cache)
+        //    which blocks the main thread behind in-flight GRAPH_COMPUTE commands
+        struct alloc_size_cache_key {
+            uint32_t device;
+            uint32_t type;
+            uint32_t op;
+            int32_t  op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)];
+            uint32_t ne[GGML_MAX_DIMS];
        };

+        alloc_size_cache_key key = {};
+        key.device = buft_ctx->device;
+        key.type = tensor->type;
+        key.op = tensor->op;
+        memcpy(key.op_params, tensor->op_params, sizeof(key.op_params));
+        for (int i = 0; i < GGML_MAX_DIMS; i++) {
+            key.ne[i] = (uint32_t)tensor->ne[i];
+        }
+
+        uint64_t cache_hash = fnv_hash((const uint8_t *)&key, sizeof(key));
+        cache_hash = fnv_hash((const uint8_t *)buft_ctx->endpoint.data(), buft_ctx->endpoint.size(), cache_hash);
+
+        // alloc sizes are immutable for a given tensor configuration
+        static std::mutex cache_mutex;
+        static std::unordered_map<uint64_t, size_t> cache;
+
+        {
+            std::lock_guard<std::mutex> lock(cache_mutex);
+            auto it = cache.find(cache_hash);
+            if (it != cache.end()) {
+                return it->second;
+            }
+        }
+
+        auto request = std::make_shared<rpc_msg_get_alloc_size_req>();
+        request->device = buft_ctx->device;
+        request->tensor = serialize_tensor(tensor);
+
        // .get_alloc_size could be a function of the tensor's srcs, so we must serialize them as well
        for (int i = 0; i < GGML_MAX_SRC; i++) {
-            request.srcs[i] = serialize_tensor(tensor->src[i]);
+            request->srcs[i] = serialize_tensor(tensor->src[i]);
        }

-        // TODO: cache the alloc responses to avoid extra RPC calls?
        rpc_msg_get_alloc_size_rsp response;
-        bool status = send_rpc_cmd(sock, RPC_CMD_GET_ALLOC_SIZE, &request, sizeof(request), &response, sizeof(response));
-        RPC_STATUS_ASSERT(status);
+        auto dispatcher = get_dispatcher(buft_ctx->endpoint);
+        dispatcher->send(RPC_CMD_GET_ALLOC_SIZE, request, sizeof(*request), &response, sizeof(response));
+
+        {
+            std::lock_guard<std::mutex> lock(cache_mutex);
+            cache[cache_hash] = response.alloc_size;
+        }

        return response.alloc_size;
    }
@@ -673,9 +918,44 @@ static void ggml_backend_rpc_free(ggml_backend_t backend) {
    delete backend;
 }

+static void ggml_backend_rpc_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
+    ggml_backend_rpc_context * ctx = (ggml_backend_rpc_context *)backend->context;
+    rpc_tensor rpc_tensor = serialize_tensor(tensor);
+    if (size > HASH_THRESHOLD) {
+        auto request = std::make_shared<rpc_msg_set_tensor_hash_req>();
+        request->tensor = rpc_tensor;
+        request->offset = offset;
+        request->hash = fnv_hash((const uint8_t*)data, size);
+        rpc_msg_set_tensor_hash_rsp response;
+        // TODO: make this async
+        ctx->dispatcher->send(RPC_CMD_SET_TENSOR_HASH, request, sizeof(*request), &response, sizeof(response));
+        if (response.result) {
+            // the server has the same data, no need to send it
+            return;
+        }
+    }
+    // input serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes)
+    size_t input_size = sizeof(rpc_tensor) + sizeof(uint64_t) + size;
+    uint8_t * input = new uint8_t[input_size]();
+    memcpy(input, &rpc_tensor, sizeof(rpc_tensor));
+    memcpy(input + sizeof(rpc_tensor), &offset, sizeof(offset));
+    memcpy(input + sizeof(rpc_tensor) + sizeof(offset), data, size);
+    std::shared_ptr<uint8_t> input_ptr(input, std::default_delete<uint8_t[]>());
+    ctx->dispatcher->send_async(RPC_CMD_SET_TENSOR, input_ptr, input_size);
+}
+
+static void ggml_backend_rpc_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
+    ggml_backend_rpc_context * ctx = (ggml_backend_rpc_context *)backend->context;
+    auto request = std::make_shared<rpc_msg_get_tensor_req>();
+    request->tensor = serialize_tensor(tensor);
+    request->offset = offset;
+    request->size = size;
+    ctx->dispatcher->send_async(RPC_CMD_GET_TENSOR, request, sizeof(*request), data, size);
+}
+
 static void ggml_backend_rpc_synchronize(ggml_backend_t backend) {
-    GGML_UNUSED(backend);
-    // this is no-op because we don't have any async operations
+    ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context;
+    rpc_ctx->dispatcher->synchronize();
 }

 static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::vector<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) {
@@ -698,7 +978,7 @@ static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::ve
    tensors.push_back(result);
 }

-static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::vector<uint8_t> & output) {
+static uint8_t * serialize_graph(uint32_t device, const ggml_cgraph * cgraph, size_t * output_size) {
    uint32_t n_nodes = cgraph->n_nodes;
    std::vector<rpc_tensor> tensors;
    std::unordered_set<ggml_tensor*> visited;
@@ -708,9 +988,9 @@ static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::ve
    // serialization format:
    // | device (4 bytes) | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) |
    uint32_t n_tensors = tensors.size();
-    int output_size = 2*sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(rpc_tensor);
-    output.resize(output_size, 0);
-    uint8_t * dest = output.data();
+    *output_size = 2*sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(rpc_tensor);
+    uint8_t * output = new uint8_t[*output_size]();
+    uint8_t * dest = output;
    memcpy(dest, &device, sizeof(device));
    dest += sizeof(device);
    memcpy(dest, &n_nodes, sizeof(n_nodes));
@@ -723,6 +1003,7 @@ static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::ve
    dest += sizeof(n_tensors);
    rpc_tensor * out_tensors = (rpc_tensor *)dest;
    memcpy(out_tensors, tensors.data(), n_tensors * sizeof(rpc_tensor));
+    return output;
 }

 static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {
@@ -733,27 +1014,35 @@ static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, g
    GGML_ASSERT(cgraph->n_nodes > 0);
    bool reuse = cgraph->uid != 0 && rpc_dev_ctx->last_graph_uid == cgraph->uid;
    if (reuse) {
-        rpc_msg_graph_recompute_req request;
-        request.device = rpc_ctx->device;
-        auto sock = get_socket(rpc_ctx->endpoint);
-        bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_RECOMPUTE, &request, sizeof(request));
-        RPC_STATUS_ASSERT(status);
+        auto request = std::make_shared<rpc_msg_graph_recompute_req>();
+        request->device = rpc_ctx->device;
+        rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_RECOMPUTE, request, sizeof(*request));
    } else {
        rpc_dev_ctx->last_graph_uid = cgraph->uid;
-        std::vector<uint8_t> input;
-        serialize_graph(rpc_ctx->device, cgraph, input);
-        auto sock = get_socket(rpc_ctx->endpoint);
-        bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size());
-        RPC_STATUS_ASSERT(status);
+        size_t input_size = 0;
+        uint8_t * input = serialize_graph(rpc_ctx->device, cgraph, &input_size);
+        std::shared_ptr<uint8_t> input_ptr(input, std::default_delete<uint8_t[]>());
+        rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_COMPUTE, input_ptr, input_size);
    }
    return GGML_STATUS_SUCCESS;
 }

+static void ggml_backend_rpc_event_record(ggml_backend_t backend, ggml_backend_event_t event) {
+    ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context;
+    rpc_ctx->dispatcher->event_record(event);
+}
+
+static void ggml_backend_rpc_event_wait(ggml_backend_t backend, ggml_backend_event_t event) {
+    // this is noop for RPC as we have a single stream
+    GGML_UNUSED(backend);
+    GGML_UNUSED(event);
+}
+
 static ggml_backend_i ggml_backend_rpc_interface = {
    /* .get_name                = */ ggml_backend_rpc_name,
    /* .free                    = */ ggml_backend_rpc_free,
-    /* .set_tensor_async        = */ NULL,
-    /* .get_tensor_async        = */ NULL,
+    /* .set_tensor_async        = */ ggml_backend_rpc_set_tensor_async,
+    /* .get_tensor_async        = */ ggml_backend_rpc_get_tensor_async,
    /* .set_tensor_2d_async     = */ NULL,
    /* .get_tensor_2d_async     = */ NULL,
    /* .cpy_tensor_async        = */ NULL,
@@ -763,8 +1052,8 @@ static ggml_backend_i ggml_backend_rpc_interface = {
    /* .graph_plan_update       = */ NULL,
    /* .graph_plan_compute      = */ NULL,
    /* .graph_compute           = */ ggml_backend_rpc_graph_compute,
-    /* .event_record            = */ NULL,
-    /* .event_wait              = */ NULL,
+    /* .event_record            = */ ggml_backend_rpc_event_record,
+    /* .event_wait              = */ ggml_backend_rpc_event_wait,
    /* .graph_optimize          = */ NULL,
 };

@@ -778,13 +1067,9 @@ ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, u
    if (it != buft_map.end()) {
        return it->second;
    }
-    auto sock = get_socket(endpoint);
-    if (sock == nullptr) {
-        GGML_LOG_ERROR("Failed to connect to %s\n", endpoint);
-        return nullptr;
-    }
-    size_t alignment = get_alignment(sock, device);
-    size_t max_size = get_max_size(sock, device);
+    auto dispatcher = get_dispatcher(endpoint);
+    size_t alignment = get_alignment(dispatcher, device);
+    size_t max_size = get_max_size(dispatcher, device);
    ggml_backend_rpc_buffer_type_context * buft_ctx = new ggml_backend_rpc_buffer_type_context {
        /* .endpoint  = */ endpoint,
        /* .device    = */ device,
@@ -804,10 +1089,11 @@ ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, u

 ggml_backend_t ggml_backend_rpc_init(const char * endpoint, uint32_t device) {
    std::string dev_name = "RPC" + std::to_string(device) + "[" + std::string(endpoint) + "]";
+    auto dispatcher = get_dispatcher(endpoint);
    ggml_backend_rpc_context * ctx = new ggml_backend_rpc_context {
-        /* .endpoint       = */ endpoint,
-        /* .device         = */ device,
-        /* .name           = */ dev_name,
+        /* .dispatcher = */ dispatcher,
+        /* .device     = */ device,
+        /* .name       = */ dev_name,
    };
    auto reg = ggml_backend_rpc_add_server(endpoint);
    ggml_backend_t backend = new ggml_backend {
@@ -823,26 +1109,16 @@ bool ggml_backend_is_rpc(ggml_backend_t backend) {
    return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_rpc_guid());
 }

-static void get_device_memory(const std::shared_ptr<socket_t> & sock, uint32_t device, size_t * free, size_t * total) {
-    rpc_msg_get_device_memory_req request;
-    request.device = device;
+void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total) {
+    auto dispatcher = get_dispatcher(endpoint);
+    auto request = std::make_shared<rpc_msg_get_device_memory_req>();
+    request->device = device;
    rpc_msg_get_device_memory_rsp response;
-    bool status = send_rpc_cmd(sock, RPC_CMD_GET_DEVICE_MEMORY, &request, sizeof(request), &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+    dispatcher->send(RPC_CMD_GET_DEVICE_MEMORY, request, sizeof(*request), &response, sizeof(response));
    *free = response.free_mem;
    *total = response.total_mem;
 }

-void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total) {
-    auto sock = get_socket(endpoint);
-    if (sock == nullptr) {
-        *free = 0;
-        *total = 0;
-        return;
-    }
-    get_device_memory(sock, device, free, total);
-}
-
 // RPC server-side implementation

 class rpc_server {
@@ -1647,9 +1923,6 @@ static void rpc_serve_client(const std::vector<ggml_backend_t> & backends, const
                if (!server.free_buffer(request)) {
                    return;
                }
-                if (!send_msg(sock, nullptr, 0)) {
-                    return;
-                }
                break;
            }
            case RPC_CMD_BUFFER_CLEAR: {
@@ -1660,9 +1933,6 @@ static void rpc_serve_client(const std::vector<ggml_backend_t> & backends, const
                if (!server.buffer_clear(request)) {
                    return;
                }
-                if (!send_msg(sock, nullptr, 0)) {
-                    return;
-                }
                break;
            }
            case RPC_CMD_MEMSET_TENSOR: {
@@ -1673,9 +1943,6 @@ static void rpc_serve_client(const std::vector<ggml_backend_t> & backends, const
                if (!server.memset_tensor(request)) {
                    return;
                }
-                if (!send_msg(sock, nullptr, 0)) {
-                    return;
-                }
                break;
            }
            case RPC_CMD_SET_TENSOR: {
@@ -1710,9 +1977,6 @@ static void rpc_serve_client(const std::vector<ggml_backend_t> & backends, const
                if (!server.init_tensor(request)) {
                    return;
                }
-                if (!send_msg(sock, nullptr, 0)) {
-                    return;
-                }
                break;
            }
            case RPC_CMD_GET_TENSOR: {
@@ -1889,10 +2153,10 @@ static void ggml_backend_rpc_device_get_props(ggml_backend_dev_t dev, struct ggm
    props->type        = ggml_backend_rpc_device_get_type(dev);
    ggml_backend_rpc_device_get_memory(dev, &props->memory_free, &props->memory_total);
    props->caps = {
-        /* .async                 = */ false,
+        /* .async                 = */ true,
        /* .host_buffer           = */ false,
        /* .buffer_from_host_ptr  = */ false,
-        /* .events                = */ false,
+        /* .events                = */ true,
        /* .mmap_support          = */ true,
    };
 }
@@ -1929,6 +2193,24 @@ static bool ggml_backend_rpc_device_supports_buft(ggml_backend_dev_t dev, ggml_b
    return buft_ctx->endpoint == dev_ctx->endpoint && buft_ctx->device == dev_ctx->device;
 }

+static ggml_backend_event_t ggml_backend_rpc_device_event_new(ggml_backend_dev_t dev) {
+    ggml_backend_rpc_device_context * ctx = (ggml_backend_rpc_device_context *)dev->context;
+    auto dispatcher = get_dispatcher(ctx->endpoint);
+    return dispatcher->event_new(dev);
+}
+
+static void ggml_backend_rpc_device_event_free(ggml_backend_dev_t dev, ggml_backend_event_t event) {
+    ggml_backend_rpc_device_context * ctx = (ggml_backend_rpc_device_context *)dev->context;
+    auto dispatcher = get_dispatcher(ctx->endpoint);
+    dispatcher->event_free(event);
+}
+
+static void ggml_backend_rpc_device_event_synchronize(ggml_backend_dev_t dev, ggml_backend_event_t event) {
+    ggml_backend_rpc_device_context * ctx = (ggml_backend_rpc_device_context *)dev->context;
+    auto dispatcher = get_dispatcher(ctx->endpoint);
+    dispatcher->event_synchronize(event);
+}
+
 static const struct ggml_backend_device_i ggml_backend_rpc_device_i = {
    /* .get_name             = */ ggml_backend_rpc_device_get_name,
    /* .get_description      = */ ggml_backend_rpc_device_get_description,
@@ -1942,9 +2224,9 @@ static const struct ggml_backend_device_i ggml_backend_rpc_device_i = {
    /* .supports_op          = */ ggml_backend_rpc_device_supports_op,
    /* .supports_buft        = */ ggml_backend_rpc_device_supports_buft,
    /* .offload_op           = */ NULL,
-    /* .event_new            = */ NULL,
-    /* .event_free           = */ NULL,
-    /* .event_synchronize    = */ NULL,
+    /* .event_new            = */ ggml_backend_rpc_device_event_new,
+    /* .event_free           = */ ggml_backend_rpc_device_event_free,
+    /* .event_synchronize    = */ ggml_backend_rpc_device_event_synchronize,
 };

 // backend reg interface
@@ -2004,14 +2286,9 @@ ggml_backend_reg_t ggml_backend_rpc_reg(void) {
 }

 static uint32_t ggml_backend_rpc_get_device_count(const char * endpoint) {
-    auto sock = get_socket(endpoint);
-    if (sock == nullptr) {
-        GGML_LOG_ERROR("Failed to connect to %s\n", endpoint);
-        return 0;
-    }
+    auto dispatcher = get_dispatcher(endpoint);
    rpc_msg_device_count_rsp response;
-    bool status = send_rpc_cmd(sock, RPC_CMD_DEVICE_COUNT, nullptr, 0, &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+    dispatcher->send(RPC_CMD_DEVICE_COUNT, nullptr, 0, &response, sizeof(response));
    return response.device_count;
 }

@@ -1,3 +1,4 @@
+#include <array>
 #include <cstdint>
 #include <cstdio>
 #include <cstring>
@@ -150,7 +151,8 @@ struct sdpa_partition {

 // Build + compile the contiguous-input GQA SDPA graph (MatMul->Divide->Add->SoftMax->MatMul), f32 out.
 // Mirrors the hardware-verified scratch/onednn_sdpa_probe.cpp build_gqa (partitions=1, sdp_primitive_kernel_t).
-static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d) {
+static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d,
+                                 const std::array<int64_t, 5> & k_str, const std::array<int64_t, 5> & v_str) try {
    using ltype = logical_tensor::layout_type;
    using dt    = logical_tensor::data_type;
    using ldims = logical_tensor::dims;
@@ -158,11 +160,12 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int
    const int   rep = H / Hkv;
    const ldims q_sz = {1, Hkv, rep, q, d}, kv_sz = {1, Hkv, 1, seq, d}, s_sz = {1, Hkv, rep, q, seq},
                sc = {1, 1, 1, 1, 1}, msk = {1, 1, 1, q, seq}, o_sz = {1, Hkv, rep, q, d};
+    const ldims k_st(k_str.begin(), k_str.end()), v_st(v_str.begin(), v_str.end());
    int64_t        id = 0;
    sdpa_partition E;

    auto query  = logical_tensor(id++, t,  q_sz, ltype::strided);
-    auto key    = logical_tensor(id++, t,  kv_sz, ltype::strided);
+    auto key    = logical_tensor(id++, t,  kv_sz, k_st);
    auto score  = logical_tensor(id++, fi, s_sz, ltype::strided);
    auto bmm1   = op(id++, op::kind::MatMul, "bmm1");
    bmm1.set_attr<bool>(op::attr::transpose_b, true);          // key is [.., seq, d]
@@ -184,7 +187,7 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int
    smax.set_attr<std::string>(op::attr::mode, "inf_as_zero");
    smax.add_inputs({masked}); smax.add_outputs({probs});

-    auto value  = logical_tensor(id++, t,  kv_sz, ltype::strided);
+    auto value  = logical_tensor(id++, t,  kv_sz, v_st);
    // f16 output is REQUIRED to hit sdp_primitive_kernel_t (the systolic micro-kernel); an f32 output
    // falls to larger_partition_kernel_t which materializes N^2 (confirmed: scratch/onednn_sdpa_kernel_probe.cpp).
    // converted to the f32 ggml dst in the permute below.
@@ -198,6 +201,7 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int

    auto parts = g.get_partitions();
    if (parts.size() != 1 || !parts[0].is_supported()) {
+        GGML_LOG_WARN("%s: oneDNN did not fuse the SDPA graph; falling back to TILE kernel\n", __func__);
        return E;   // ok stays false -> caller falls back to TILE
    }
    E.ins      = parts[0].get_input_ports();
@@ -209,6 +213,12 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int
    E.ok       = true;
    return E;
 }
+catch (const std::exception & e) {
+    // compile() can reject a stride set the partitioner never inspects; memoise the failure so the
+    // fallback costs one build rather than one per call.
+    GGML_LOG_WARN("%s: oneDNN SDPA partition build failed (%s); falling back to TILE kernel\n", __func__, e.what());
+    return {};
+}

 void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tensor * dst) try {
    const ggml_tensor * Q    = dst->src[0];
@@ -234,13 +244,34 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
    ggml_sycl_pool_alloc<sycl::half> Qf(ctx.pool(), (size_t) H * q * d);
    cont_to_f16_sycl<float>((const char *) Q->data, Qf.get(), d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream);

-    // K/V: use pool-alloc for both F16 and dequant paths.
+    // K/V: bind the f16 cache in place. llama.cpp permutes it to [token][head][dim], so its head
+    // plane is strided rather than dense, which is what an explicit stride vector expresses.
+    // Quantized and f32 KV still stage a dense copy -- the layout the k_str/v_str defaults describe.
    sycl::half * K_ptr = nullptr;
    sycl::half * V_ptr = nullptr;
+    std::array<int64_t, 5> k_str{ Hkv * seq * d, seq * d, seq * d, d, 1 };
+    std::array<int64_t, 5> v_str = k_str;
    std::optional<ggml_sycl_pool_alloc<sycl::half>> Kf_pool;
    std::optional<ggml_sycl_pool_alloc<sycl::half>> Vf_pool;

-    if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) {
+    auto bindable = [](const ggml_tensor * t) {
+        return t->nb[0] == sizeof(sycl::half) && t->nb[1] % sizeof(sycl::half) == 0 &&
+               t->nb[2] % sizeof(sycl::half) == 0 && t->nb[3] % sizeof(sycl::half) == 0;
+    };
+    auto elem_strides = [](const ggml_tensor * t) {
+        const int64_t s1 = (int64_t) (t->nb[1] / t->nb[0]);
+        const int64_t s2 = (int64_t) (t->nb[2] / t->nb[0]);
+        const int64_t s3 = (int64_t) (t->nb[3] / t->nb[0]);
+        // dims are {mb=1, Hkv, rep=1, seq, d}; the size-1 dims at 0 and 2 never advance an address.
+        return std::array<int64_t, 5>{ s3, s2, s2, s1, 1 };
+    };
+
+    if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && bindable(K) && bindable(V)) {
+        K_ptr = (sycl::half *) K->data;
+        V_ptr = (sycl::half *) V->data;
+        k_str = elem_strides(K);
+        v_str = elem_strides(V);
+    } else if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) {
        Kf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d);
        Vf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d);
        cont_to_f16_sycl<sycl::half>((const char *) K->data, Kf_pool->get(), d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream);
@@ -341,19 +372,24 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso

    ggml_sycl_pool_alloc<sycl::half> outf(ctx.pool(), (size_t) H * q * d);   // f16 contiguous SDPA out [mb,H,q,d]

-    // compile once per (device, shape), reuse across layers/calls.
+    // compile once per (device, shape, KV strides), reuse across layers/calls. Stride 2 always
+    // repeats stride 1 and stride 4 is always 1, so the key covers every entry that can differ.
    static std::unordered_map<std::string, sdpa_partition> cache;
-    char keyb[96];
-    snprintf(keyb, sizeof(keyb), "%d:%lld:%lld:%lld:%lld:%lld", ggml_sycl_get_device(),
-             (long long) H, (long long) Hkv, (long long) q, (long long) seq, (long long) d);
+    char keyb[256];
+    snprintf(keyb, sizeof(keyb), "%d:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld", ggml_sycl_get_device(),
+             (long long) H, (long long) Hkv, (long long) q, (long long) seq, (long long) d,
+             (long long) k_str[0], (long long) k_str[1], (long long) k_str[3],
+             (long long) v_str[0], (long long) v_str[1], (long long) v_str[3]);
    auto it = cache.find(keyb);
    if (it == cache.end()) {
-        it = cache.emplace(keyb, build_sdpa(eng, (int) H, (int) Hkv, (int) q, (int) seq, (int) d)).first;
+        it = cache.emplace(keyb, build_sdpa(eng, (int) H, (int) Hkv, (int) q, (int) seq, (int) d, k_str, v_str)).first;
    }
    sdpa_partition & E = it->second;
-    // _supported() is authoritative: if it accepted this op the partition must build.
-    // A failure here is a gap in _supported() -- surface it, don't mask it with a fallback.
-    GGML_ASSERT(E.ok && "oneDNN SDPA partition failed to build for a _supported() shape");
+    if (!E.ok) {
+        // oneDNN can decline a shape or a stride set that _supported() never sees; build_sdpa warns per key.
+        ggml_sycl_flash_attn_ext_tile(ctx, dst);
+        return;
+    }

    auto id2ptr = [&](size_t r) -> void * {
        if (r == E.id_q)     return Qf.get();
@@ -104,7 +104,6 @@ enum best_fattn_kernel {


 static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const ggml_tensor * dst) {
-    GGML_UNUSED(device);
 #ifndef SYCL_FLASH_ATTN
    GGML_UNUSED(dst);
    return BEST_FATTN_KERNEL_NONE;
@@ -263,6 +262,11 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
            }
        } else {
            if (Q->ne[1] <= 2) {
+                // TILE is faster for quantized KV decode on Xe2 (BMG); keep VEC on untested archs
+                const gpu_arch arch = ggml_sycl_info().devices[device].hw_info.arch;
+                if (arch == gpu_arch::intel_gpu_bmg_g21 || arch == gpu_arch::intel_gpu_bmg_g31) {
+                    return BEST_FATTN_KERNEL_TILE;
+                }
                return BEST_FATTN_KERNEL_VEC;
            }
        }
@@ -767,6 +767,21 @@ static constexpr std::initializer_list<std::array<int, 3>> rms_norm_mul_rope_vie
    { 4, 0, 3 }, // set_rows->src[0] == view
 };

+static constexpr std::array<ggml_type, 9> lightning_indexer_k_types = {
+    GGML_TYPE_F32,
+    GGML_TYPE_F16,
+    GGML_TYPE_BF16,
+    GGML_TYPE_Q8_0,
+    GGML_TYPE_Q5_1,
+    GGML_TYPE_Q5_0,
+    GGML_TYPE_Q4_1,
+    GGML_TYPE_Q4_0,
+    GGML_TYPE_IQ4_NL,
+};
+
+static bool ggml_vk_lightning_indexer_k_type_supported(ggml_type type) {
+    return std::find(lightning_indexer_k_types.begin(), lightning_indexer_k_types.end(), type) != lightning_indexer_k_types.end();
+}

 struct vk_device_struct {
    std::recursive_mutex mutex;
@@ -1042,6 +1057,8 @@ struct vk_device_struct {
    vk_pipeline pipeline_argsort_large_f32[num_argsort_pipelines];
    vk_pipeline pipeline_topk_f32[num_topk_pipelines];
    vk_pipeline pipeline_sum_rows_f32;
+    vk_pipeline pipeline_cross_entropy_loss_f32, pipeline_cross_entropy_loss_f32_wg512;
+    vk_pipeline pipeline_cross_entropy_loss_back_f32, pipeline_cross_entropy_loss_back_f32_wg512;
    vk_pipeline pipeline_fwht_f32[4];
    vk_pipeline pipeline_cumsum_f32;
    vk_pipeline pipeline_cumsum_small_f32;
@@ -1066,6 +1083,7 @@ struct vk_device_struct {
    vk_pipeline pipeline_rwkv_wkv6_f32;
    vk_pipeline pipeline_rwkv_wkv7_f32;
    vk_pipeline pipeline_gated_linear_attn_f32;
+    vk_pipeline pipeline_lightning_indexer_f32[GGML_TYPE_COUNT];
    // [size_idx][kda] where size_idx: 0=d16, 1=d32, 2=d64, 3=d128
    vk_pipeline pipeline_gated_delta_net[4][2];
    vk_pipeline pipeline_ssm_scan_f32_d128;
@@ -1330,6 +1348,8 @@ struct vk_mat_mat_id_push_constants {
    uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d;
    uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11;
    uint32_t padded_N;
+    uint32_t n_experts;
+    uint32_t hoist_row_ids;
 };
 struct vk_mat_vec_id_push_constants {
    uint32_t ncols;
@@ -1410,6 +1430,10 @@ struct vk_op_count_experts_push_constants {
    uint32_t nb00;
    uint32_t nb01;
    uint32_t a_offset;
+    uint32_t n_experts;
+    uint32_t hoist_row_ids;
+    uint32_t ne00mp;
+    uint32_t ne00L;
 };

 struct vk_op_glu_push_constants {
@@ -1588,6 +1612,10 @@ template <> void init_pushconst_fastdiv(vk_op_glu_push_constants &p) {
    init_fastdiv_values(p.ne20,                p.ne2_0mp,      p.ne2_0L);
 }

+template <> void init_pushconst_fastdiv(vk_op_count_experts_push_constants &p) {
+    init_fastdiv_values(p.ne00, p.ne00mp, p.ne00L);
+}
+
 struct vk_op_binary_push_constants {
    uint32_t ne;
    uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03;
@@ -1846,6 +1874,26 @@ struct vk_op_gated_linear_attn_push_constants {
    uint32_t H;
    float scale;
 };
+struct vk_op_lightning_indexer_push_constants {
+    uint32_t n_kv;
+    uint32_t n_heads;
+    uint32_t n_tokens;
+    uint32_t n_streams;
+    uint32_t n_masks;
+    uint32_t dispatch_x;
+    uint32_t q_nb1;
+    uint32_t q_nb2;
+    uint32_t q_nb3;
+    uint32_t k_nb2;
+    uint32_t k_nb3;
+    uint32_t w_nb1;
+    uint32_t w_nb3;
+    uint32_t m_nb1;
+    uint32_t m_nb3;
+    uint32_t d_nb1;
+    uint32_t d_nb3;
+};
+static_assert(sizeof(vk_op_lightning_indexer_push_constants) <= 128);
 struct vk_op_gated_delta_net_push_constants {
    uint32_t H;
    uint32_t n_tokens;
@@ -3902,11 +3950,16 @@ static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const
    return vk_fa_pipeline_state{hsk, hsv, params.block_rows, params.block_cols, params.d_split, params.row_split, params.shmem_staging, params.path, params.workgroup_size, subgroup_size, aligned, f32acc, flags, params.limit_occupancy_shmem, k_type, v_type};
 }

+// Bytes per buffer block for the FaBlockBytesK/V spec constants. F32 is fed as
+// a vec4 "block" of 4 floats, everything else uses its ggml block size.
+static uint32_t fa_block_bytes(ggml_type t) {
+    if (t == GGML_TYPE_F32) {
+        return 16u;
+    }
+    return (uint32_t) ggml_type_size(t);
+}
+
 static std::vector<uint32_t> get_fa_spec_constants(const vk_fa_pipeline_state& state) {
-    const auto fa_block_bytes = [](ggml_type t) -> uint32_t {
-        if (t == GGML_TYPE_F32) return 16u;
-        return (uint32_t) ggml_type_size(t);
-    };
    return {
        /* 0 WorkGroupSize   */ state.workgroup_size,
        /* 1 Br              */ state.Br,
@@ -4169,10 +4222,16 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
    const uint32_t subgroup_size_16 = std::max(device->subgroup_size, 16u);
    const uint32_t subgroup_size_32 = std::max(device->subgroup_size, 32u);

+    // clamp WARP for l_/m_ warptiles so WM <= BM (breaks on subgroupSize > 64)
+    const uint32_t mm_warp_8  = std::min(subgroup_size_8,  64u);
+    const uint32_t mm_warp_16 = std::min(subgroup_size_16, 64u);
+
    const uint32_t mul_mat_subgroup_size = (device->vendor_id == VK_VENDOR_ID_INTEL && device->subgroup_size_control) ? device->subgroup_min_size : device->subgroup_size;
    const uint32_t mul_mat_subgroup_size_8 = std::max(mul_mat_subgroup_size, 8u);
    const uint32_t mul_mat_subgroup_size_16 = std::max(mul_mat_subgroup_size, 16u);
    const uint32_t mul_mat_subgroup_size_32 = std::max(mul_mat_subgroup_size, 32u);
+    const uint32_t mul_mat_mm_warp_8  = std::min(mul_mat_subgroup_size_8,  64u);
+    const uint32_t mul_mat_mm_warp_16 = std::min(mul_mat_subgroup_size_16, 64u);

    const bool subgroup_min_size_16 = (!device->subgroup_size_control && device->subgroup_size >= 16) ||
                                      (device->subgroup_size_control && device->subgroup_max_size >= 16);
@@ -4253,39 +4312,39 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {

        const uint32_t s_warptile_wm = device->subgroup_size == 8 ? 8 : 32;

-        l_warptile = { 128,             128, 128, 16, subgroup_size_8 * 2, 64, 2, tm_l, tn_l, tk_l, subgroup_size_8 };
-        m_warptile = { 128,              64,  64, 16, subgroup_size_8,     32, 2, tm_m, tn_m, tk_m, subgroup_size_8 };
-        s_warptile = { subgroup_size_32, 32,  32, 16, s_warptile_wm,       32, 2, tm_s, tn_s, tk_s, subgroup_size_8 };
+        l_warptile = { 128,             128, 128, 16, mm_warp_8 * 2, 64, 2, tm_l, tn_l, tk_l, mm_warp_8 };
+        m_warptile = { 128,              64,  64, 16, mm_warp_8,     32, 2, tm_m, tn_m, tk_m, mm_warp_8 };
+        s_warptile = { subgroup_size_32, 32,  32, 16, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 };

-        l_warptile_mmq = { 128,             128, 128, 32, subgroup_size_8 * 2, 64, 2, tm_l, tn_l, tk_l, subgroup_size_8 };
-        m_warptile_mmq = { 128,              64,  64, 32, subgroup_size_8,     32, 2, tm_m, tn_m, tk_m, subgroup_size_8 };
-        s_warptile_mmq = { subgroup_size_32, 32,  32, 32, s_warptile_wm,       32, 2, tm_s, tn_s, tk_s, subgroup_size_8 };
+        l_warptile_mmq = { 128,             128, 128, 32, mm_warp_8 * 2, 64, 2, tm_l, tn_l, tk_l, mm_warp_8 };
+        m_warptile_mmq = { 128,              64,  64, 32, mm_warp_8,     32, 2, tm_m, tn_m, tk_m, mm_warp_8 };
+        s_warptile_mmq = { subgroup_size_32, 32,  32, 32, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 };

        // Integer MMQ has a smaller shared memory profile, but heavier register use
-        l_warptile_mmq_int = { 128,             128, 128, 32, subgroup_size_8 * 2, 64, 2, 4, 4, 1, subgroup_size_8 };
-        m_warptile_mmq_int = { 128,              64,  64, 32, subgroup_size_8,     32, 2, 2, 2, 1, subgroup_size_8 };
-        s_warptile_mmq_int = { subgroup_size_32, 32,  32, 32, s_warptile_wm,       32, 2, 2, 1, 1, subgroup_size_8 };
+        l_warptile_mmq_int = { 128,             128, 128, 32, mm_warp_8 * 2, 64, 2, 4, 4, 1, mm_warp_8 };
+        m_warptile_mmq_int = { 128,              64,  64, 32, mm_warp_8,     32, 2, 2, 2, 1, mm_warp_8 };
+        s_warptile_mmq_int = { subgroup_size_32, 32,  32, 32, s_warptile_wm, 32, 2, 2, 1, 1, subgroup_size_8 };

        // K-quants use even more registers, mitigate by setting WMITER to 1
-        l_warptile_mmq_int_k = { 128,               128, 128, 32, subgroup_size_8 * 2, 64, 1, 4, 4, 1, subgroup_size_8 };
-        m_warptile_mmq_int_k = { 128,                64,  64, 32, subgroup_size_8,     32, 1, 2, 2, 1, subgroup_size_8 };
-        s_warptile_mmq_int_k = { subgroup_size_32,   32,  32, 32, s_warptile_wm,       32, 1, 2, 1, 1, subgroup_size_8 };
+        l_warptile_mmq_int_k = { 128,               128, 128, 32, mm_warp_8 * 2, 64, 1, 4, 4, 1, mm_warp_8 };
+        m_warptile_mmq_int_k = { 128,                64,  64, 32, mm_warp_8,     32, 1, 2, 2, 1, mm_warp_8 };
+        s_warptile_mmq_int_k = { subgroup_size_32,   32,  32, 32, s_warptile_wm, 32, 1, 2, 1, 1, subgroup_size_8 };

-        l_warptile_id = { 128,                      128, 128, 16, mul_mat_subgroup_size_16 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_subgroup_size_16 };
-        m_warptile_id = { 128,                       64,  64, 16, mul_mat_subgroup_size_16,     32, 2, tm_m, tn_m, tk_m, mul_mat_subgroup_size_16 };
-        s_warptile_id = { mul_mat_subgroup_size_16,  32,  32, 16, s_warptile_wm,                32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_16 };
+        l_warptile_id = { 128,                      128, 128, 16, mul_mat_mm_warp_16 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_mm_warp_16 };
+        m_warptile_id = { 128,                       64,  64, 16, mul_mat_mm_warp_16,     32, 2, tm_m, tn_m, tk_m, mul_mat_mm_warp_16 };
+        s_warptile_id = { mul_mat_subgroup_size_16,  32,  32, 16, s_warptile_wm,          32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_16 };

-        l_warptile_mmqid = { 128,                       128, 128, 32, mul_mat_subgroup_size_8 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_subgroup_size_8 };
-        m_warptile_mmqid = { 128,                        64,  64, 32, mul_mat_subgroup_size_8,     32, 2, tm_m, tn_m, tk_m, mul_mat_subgroup_size_8 };
-        s_warptile_mmqid = { mul_mat_subgroup_size_32,   32,  32, 32, s_warptile_wm,               32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_8 };
+        l_warptile_mmqid = { 128,                       128, 128, 32, mul_mat_mm_warp_8 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_mm_warp_8 };
+        m_warptile_mmqid = { 128,                        64,  64, 32, mul_mat_mm_warp_8,     32, 2, tm_m, tn_m, tk_m, mul_mat_mm_warp_8 };
+        s_warptile_mmqid = { mul_mat_subgroup_size_32,   32,  32, 32, s_warptile_wm,         32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_8 };

-        l_warptile_mmqid_int = { 128,                       128, 128, 32, mul_mat_subgroup_size_8 * 2, 64, 2, 4, 4, 1, mul_mat_subgroup_size_8 };
-        m_warptile_mmqid_int = { 128,                        64,  64, 32, mul_mat_subgroup_size_8,     32, 2, 2, 2, 1, mul_mat_subgroup_size_8 };
-        s_warptile_mmqid_int = { mul_mat_subgroup_size_32,   32,  32, 32, s_warptile_wm,               32, 2, 2, 1, 1, mul_mat_subgroup_size_8 };
+        l_warptile_mmqid_int = { 128,                       128, 128, 32, mul_mat_mm_warp_8 * 2, 64, 2, 4, 4, 1, mul_mat_mm_warp_8 };
+        m_warptile_mmqid_int = { 128,                        64,  64, 32, mul_mat_mm_warp_8,     32, 2, 2, 2, 1, mul_mat_mm_warp_8 };
+        s_warptile_mmqid_int = { mul_mat_subgroup_size_32,   32,  32, 32, s_warptile_wm,         32, 2, 2, 1, 1, mul_mat_subgroup_size_8 };

-        l_warptile_mmqid_int_k = { 128,                     128, 128, 32, mul_mat_subgroup_size_16 * 2, 64, 1, 4, 4, 1, mul_mat_subgroup_size_16 };
-        m_warptile_mmqid_int_k = { 128,                      64,  64, 32, mul_mat_subgroup_size_16,     32, 1, 2, 2, 1, mul_mat_subgroup_size_16 };
-        s_warptile_mmqid_int_k = { mul_mat_subgroup_size_32, 32,  32, 32, s_warptile_wm,                32, 1, 2, 1, 1, mul_mat_subgroup_size_16 };
+        l_warptile_mmqid_int_k = { 128,                     128, 128, 32, mul_mat_mm_warp_16 * 2, 64, 1, 4, 4, 1, mul_mat_mm_warp_16 };
+        m_warptile_mmqid_int_k = { 128,                      64,  64, 32, mul_mat_mm_warp_16,     32, 1, 2, 2, 1, mul_mat_mm_warp_16 };
+        s_warptile_mmqid_int_k = { mul_mat_subgroup_size_32, 32,  32, 32, s_warptile_wm,          32, 1, 2, 1, 1, mul_mat_subgroup_size_16 };

        // chip specific tuning
        if ((device->architecture == AMD_GCN) && (device->driver_id != vk::DriverId::eAmdProprietary)) {
@@ -4293,13 +4352,13 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
            m_warptile_mmqid = m_warptile_mmqid_int = { 256, 64, 64, 32, 16, 16, 2, 2, 2, 1, 16 };
        } else if (device->vendor_id == VK_VENDOR_ID_AMD && device->coopmat_support && device->driver_id != vk::DriverId::eAmdProprietary) {
            // This is intentionally using tx_m values, slight performance increase
-            l_warptile = { 256, 128, 128, 16, subgroup_size_8, 64, 2, tm_m, tn_m, tk_m, subgroup_size_8 };
-            l_warptile_mmq = l_warptile_mmq_int = { 256, 128, 128, 32, subgroup_size_8, 64, 2, tm_m, tn_m, tk_m, subgroup_size_8 };
-            l_warptile_mmq_int_k = { 256, 128, 128, 32, subgroup_size_16, 64, 1, 4, 2, 1, subgroup_size_16 };
+            l_warptile = { 256, 128, 128, 16, mm_warp_8, 64, 2, tm_m, tn_m, tk_m, mm_warp_8 };
+            l_warptile_mmq = l_warptile_mmq_int = { 256, 128, 128, 32, mm_warp_8, 64, 2, tm_m, tn_m, tk_m, mm_warp_8 };
+            l_warptile_mmq_int_k = { 256, 128, 128, 32, mm_warp_16, 64, 1, 4, 2, 1, mm_warp_16 };
        } else if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support) {
            // Xe2/Xe3 with coopmat enabled - warptile performance tuning
-            l_warptile = { 512, 128, 128, 16, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 };
-            l_warptile_mmq = { 512, 128, 128, 32, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 };
+            l_warptile = { 512, 128, 128, 16, mm_warp_8, 32, 2, tm_m, tn_m, tk_m, mm_warp_8 };
+            l_warptile_mmq = { 512, 128, 128, 32, mm_warp_8, 32, 2, tm_m, tn_m, tk_m, mm_warp_8 };
        }

        l_mmq_wg_denoms = l_wg_denoms = {128, 128, 1 };
@@ -5172,8 +5231,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
        const uint32_t s_warptile_wm = device->subgroup_size == 8 ? 8 : 32;

        // use scalar tile sizes
-        l_warptile = { 128, 128, 128, 16, subgroup_size_8 * 2, 64, 2, 4, 4, 1, subgroup_size_8 };
-        m_warptile = { 128,  64,  64, 16, subgroup_size_8, 32, 2, 4, 2, 1, subgroup_size_8 };
+        l_warptile = { 128, 128, 128, 16, mm_warp_8 * 2, 64, 2, 4, 4, 1, mm_warp_8 };
+        m_warptile = { 128,  64,  64, 16, mm_warp_8, 32, 2, 4, 2, 1, mm_warp_8 };
        s_warptile = { subgroup_size_32, 32, 32, 16, s_warptile_wm, 32, 2, 2, 2, 1, subgroup_size_8 };

        l_wg_denoms = {128, 128, 1 };
@@ -5758,6 +5817,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
    ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);

    ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_sum_rows_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_f32, "cross_entropy_loss_f32", cross_entropy_loss_f32_len, cross_entropy_loss_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_f32_wg512, "cross_entropy_loss_f32_wg512", cross_entropy_loss_f32_len, cross_entropy_loss_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, { 512 }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_back_f32, "cross_entropy_loss_back_f32", cross_entropy_loss_back_f32_len, cross_entropy_loss_back_f32_data, "main", 4, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_back_f32_wg512, "cross_entropy_loss_back_f32_wg512", cross_entropy_loss_back_f32_len, cross_entropy_loss_back_f32_data, "main", 4, sizeof(vk_op_push_constants), {1, 1, 1}, { 512 }, 1);
    // Intel Windows driver in range [32.0.101.8509, 32.0.101.8860) will crash when using fwht kernels so we gate that here
    const bool can_use_fwht = device->driver_id != vk::DriverId::eIntelProprietaryWindows ||
        !ggml_vk_intel_windows_driver_in_range(device->properties.driverVersion, 101, 8509, 101, 8860);
@@ -5786,7 +5849,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {

    ggml_vk_create_pipeline(device, device->pipeline_count_equal_i32, "count_equal_i32", count_equal_i32_len, count_equal_i32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, { device->subgroup_size }, 1);

-    ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true);
+    if (device->subgroup_arithmetic && device->subgroup_require_full_support) {
+        ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_subgroup_len, count_experts_subgroup_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true, true);
+    } else {
+        ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true);
+    }

    for (auto &s : device->pipeline_solve_tri_f32) {
        const vk_solve_tri_pipeline_state &state = s.first;
@@ -5835,6 +5902,17 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {

    ggml_vk_create_pipeline(device, device->pipeline_gated_linear_attn_f32, "gated_linear_attn_f32", gated_linear_attn_f32_len, gated_linear_attn_f32_data, "main", 6, sizeof(vk_op_gated_linear_attn_push_constants), {1, 1, 1}, {}, 1);

+    {
+        const bool li_subgroup = device->subgroup_arithmetic && device->subgroup_require_full_support;
+        const size_t li_len   = li_subgroup ? lightning_indexer_subgroup_f32_len  : lightning_indexer_f32_len;
+        const void * li_data  = li_subgroup ? (const void *)lightning_indexer_subgroup_f32_data : (const void *)lightning_indexer_f32_data;
+
+        for (ggml_type k_type : lightning_indexer_k_types) {
+            const std::string name = "lightning_indexer_" + std::string(ggml_type_name(k_type)) + "_k_f32";
+            ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f32[k_type], name.c_str(), li_len, li_data, "main", 5, sizeof(vk_op_lightning_indexer_push_constants), {1, 1, 1}, {(uint32_t)k_type, fa_block_bytes(k_type), device->subgroup_size}, 1, true, li_subgroup);
+        }
+    }
+
    {
        const uint32_t gdn_sizes[] = {16, 32, 64, 128};
        const char * gdn_names[][2] = {
@@ -8906,13 +8984,13 @@ static void ggml_vk_matmul_id(
        uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d,
        uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d,
        uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11,
-        uint32_t padded_n) {
+        uint32_t padded_n, bool hoist_row_ids) {
    VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), expert_count: (" << expert_count_buf.buffer->buffer << ", " << expert_count_buf.offset << ", " << expert_count_buf.size << "), " <<
        "m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " <<
        "batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " <<
        "n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")");
    const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d,
-                                              nei0, nei1, nbi1, ne11, padded_n };
+                                              nei0, nei1, nbi1, ne11, padded_n, n_as, uint32_t(hoist_row_ids) };
    ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf }, pc, { m, nei1, n_as });
 }

@@ -10098,6 +10176,12 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
    // const uint64_t ne23 = dst->ne[3];

    const uint64_t n_as = ne02;
+    // n_as counts, n_as offsets, one total, then one packed row id per (expert, token).
+    // Hoisting requires 16-bit indices for the packing and a table that fits one binding.
+    const uint64_t hoisted_row_id_words = 2 * n_as + 1 + nei0 * nei1;
+    const bool hoist_row_ids = n_as <= 256 && nei0 <= 0xffff && nei1 <= 0xffff &&
+                                hoisted_row_id_words * sizeof(uint32_t) <=
+                                    ctx->device->properties.limits.maxStorageBufferRange;

    ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context;
    ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context;
@@ -10238,7 +10322,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
    }
    vk_pipeline count_experts = ctx->device->pipeline_count_experts;

-    uint32_t expert_count_size = sizeof(uint32_t) * n_as;
+    const size_t expert_data_size = sizeof(uint32_t) *
+        (hoist_row_ids ? hoisted_row_id_words : n_as);

    {
        if (
@@ -10254,8 +10339,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
            ctx->prealloc_size_y = y_sz;
            ggml_vk_preallocate_buffers(ctx, subctx);
        }
-        if (ctx->prealloc_size_split_k < expert_count_size) {
-            ctx->prealloc_size_split_k = expert_count_size;
+        if (ctx->prealloc_size_split_k < expert_data_size) {
+            ctx->prealloc_size_split_k = expert_data_size;
            ggml_vk_preallocate_buffers(ctx, subctx);
        }

@@ -10321,18 +10406,23 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
        }
    }
    // Count how many times each expert is used
-    vk_subbuffer expert_count_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0);
+    vk_subbuffer expert_count_buf = { ctx->prealloc_split_k, 0, expert_data_size };
    if (ctx->prealloc_split_k_need_sync) {
        ggml_vk_sync_buffers(ctx, subctx);
    }
    {
-        const std::vector<uint32_t> pc = { (uint32_t)nei0,
+        vk_op_count_experts_push_constants pc = { (uint32_t)nei0,
                                           (uint32_t)nei1,
                                           (uint32_t)(nbi0 / ggml_type_size(ids->type)),
                                           (uint32_t)(nbi1 / ggml_type_size(ids->type)),
-                                           (uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)) };
+                                           (uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)),
+                                           (uint32_t)n_as,
+                                           uint32_t(hoist_row_ids),
+                                           0, 0 };
+        init_pushconst_fastdiv(pc);
        ggml_vk_dispatch_pipeline(ctx, subctx, count_experts,
-            { vk_subbuffer{ d_ids, ids_buf_offset, ids_sz }, expert_count_buf }, pc, { (uint32_t)n_as, 1, 1});
+            { vk_subbuffer{ d_ids, ids_buf_offset, ids_sz }, expert_count_buf }, pc,
+            { hoist_row_ids ? 1u : (uint32_t)n_as, 1, 1});
    }

    if (x_non_contig) {
@@ -10401,7 +10491,7 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
        { d_D, d_buf_offset, d_sz }, { d_ids, ids_buf_offset, ids_sz }, expert_count_buf,
        ne01, ne21, ne10, ne10, stride_b_y, ne01,
        stride_batch_x, stride_batch_y, ne20*ne21,
-        n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, padded_n
+        n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, padded_n, hoist_row_ids
    );  // NOLINT

    if (x_non_contig || qx_needs_dequant) {
@@ -11577,6 +11667,17 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
            return ctx->device->pipeline_sum_rows_f32;
        }
        return nullptr;
+    case GGML_OP_CROSS_ENTROPY_LOSS:
+        if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+            return src0->ne[0] > 1024 ? ctx->device->pipeline_cross_entropy_loss_f32_wg512 : ctx->device->pipeline_cross_entropy_loss_f32;
+        }
+        return nullptr;
+    case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
+        // src0 is the scalar grad; src1 is logits
+        if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && src2 && src2->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+            return src1->ne[0] > 1024 ? ctx->device->pipeline_cross_entropy_loss_back_f32_wg512 : ctx->device->pipeline_cross_entropy_loss_back_f32;
+        }
+        return nullptr;
    case GGML_OP_CUMSUM:
        if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
            if (src0->ne[0] <= 512) {
@@ -11674,6 +11775,12 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
            return ctx->device->pipeline_gated_linear_attn_f32;
        }
        return nullptr;
+    case GGML_OP_LIGHTNING_INDEXER:
+        // only the k type selects a pipeline, the other types are fixed by ggml_lightning_indexer()
+        if (ggml_vk_lightning_indexer_k_type_supported(src1->type)) {
+            return ctx->device->pipeline_lightning_indexer_f32[src1->type];
+        }
+        return nullptr;
    case GGML_OP_GATED_DELTA_NET:
        if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
            const uint32_t S_v = dst->src[2]->ne[0];
@@ -12749,6 +12856,55 @@ static void ggml_vk_gated_linear_attn(ggml_backend_vk_context * ctx, vk_context&
        pc, { (uint32_t)(n_seqs * n_heads), 1, 1 });
 }

+static void ggml_vk_lightning_indexer(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
+    const ggml_tensor * q = dst->src[0];
+    const ggml_tensor * k = dst->src[1];
+    const ggml_tensor * w = dst->src[2];
+    const ggml_tensor * m = dst->src[3];
+
+    vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, q, k, w, dst, dst->op);
+    GGML_ASSERT(pipeline != nullptr);
+
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+
+    const uint32_t n_kv      = k->ne[2];
+    const uint32_t n_heads   = q->ne[1];
+    const uint32_t n_tokens  = q->ne[2];
+    const uint32_t n_streams = q->ne[3];
+    const uint32_t n_masks   = m->ne[3];
+
+    const uint32_t n_outputs = (uint32_t)(dst->ne[0] * dst->ne[1] * dst->ne[3]);
+    const uint32_t dispatch_x = std::min(n_outputs, ctx->device->properties.limits.maxComputeWorkGroupCount[0]);
+    const uint32_t dispatch_y = CEIL_DIV(n_outputs, dispatch_x);
+
+    // q, w and dst are f32 and m is f16, so their strides are passed in elements;
+    // k may be quantized, so its strides stay in bytes
+    const uint32_t q_nb1 = q->nb[1] / sizeof(float);
+    const uint32_t q_nb2 = q->nb[2] / sizeof(float);
+    const uint32_t q_nb3 = q->nb[3] / sizeof(float);
+    const uint32_t k_nb2 = k->nb[2];
+    const uint32_t k_nb3 = k->nb[3];
+    const uint32_t w_nb1 = w->nb[1] / sizeof(float);
+    const uint32_t w_nb3 = w->nb[3] / sizeof(float);
+    const uint32_t m_nb1 = m->nb[1] / sizeof(ggml_fp16_t);
+    const uint32_t m_nb3 = m->nb[3] / sizeof(ggml_fp16_t);
+    const uint32_t d_nb1 = dst->nb[1] / sizeof(float);
+    const uint32_t d_nb3 = dst->nb[3] / sizeof(float);
+
+    const vk_op_lightning_indexer_push_constants pc = {
+        n_kv, n_heads, n_tokens, n_streams, n_masks, dispatch_x,
+        q_nb1, q_nb2, q_nb3,
+        k_nb2, k_nb3,
+        w_nb1, w_nb3,
+        m_nb1, m_nb3,
+        d_nb1, d_nb3,
+    };
+
+    ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
+        {ggml_vk_tensor_subbuffer(ctx, q), ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, w), ggml_vk_tensor_subbuffer(ctx, m), ggml_vk_tensor_subbuffer(ctx, dst)},
+        pc, {dispatch_x, dispatch_y, 1});
+}
+
 static void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
    const ggml_tensor * src_q     = dst->src[0];
    const ggml_tensor * src_v     = dst->src[2];
@@ -13942,6 +14098,103 @@ static void ggml_vk_cumsum(ggml_backend_vk_context * ctx, vk_context& subctx, co
    ctx->prealloc_split_k_need_sync = true;
 }

+static std::array<uint32_t, 3> ggml_vk_nrows_elements(uint32_t nr) {
+    if (nr > 262144) {
+        return { 512, 512, CEIL_DIV(nr, 262144) };
+    }
+    if (nr > 512) {
+        return { 512, CEIL_DIV(nr, 512), 1 };
+    }
+    return { nr, 1, 1 };
+}
+
+static void ggml_vk_cross_entropy_loss(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_contiguous(src0));
+    GGML_ASSERT(ggml_is_contiguous(src1));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+    GGML_ASSERT(ggml_are_same_shape(src0, src1));
+    GGML_ASSERT(ggml_is_scalar(dst));
+
+    const uint32_t nclasses = (uint32_t)src0->ne[0];
+    const uint32_t nrows    = (uint32_t)ggml_nrows(src0);
+
+    vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, nullptr, dst, GGML_OP_CROSS_ENTROPY_LOSS);
+    GGML_ASSERT(pipeline != nullptr);
+
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+    ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_sum_rows_f32, 1);
+
+    vk_subbuffer src0_buf = ggml_vk_tensor_subbuffer(ctx, src0);
+    vk_subbuffer src1_buf = ggml_vk_tensor_subbuffer(ctx, src1);
+    vk_subbuffer dst_buf  = ggml_vk_tensor_subbuffer(ctx, dst, true);
+
+    const vk_op_push_constants pc = { nclasses, nrows, 0.0f, 0.0f, 0.0f, 0.0f };
+
+    const size_t tmp_size = (size_t)nrows * sizeof(float);
+    if (ctx->prealloc_size_x < tmp_size) {
+        ctx->prealloc_size_x = tmp_size;
+        ggml_vk_preallocate_buffers(ctx, subctx);
+    }
+    if (ctx->prealloc_x_need_sync) {
+        ggml_vk_sync_buffers(ctx, subctx);
+    }
+
+    vk_subbuffer tmp_buf = { ctx->prealloc_x, 0, tmp_size };
+    ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src0_buf, src1_buf, tmp_buf }, pc, ggml_vk_nrows_elements(nrows));
+    ggml_vk_sync_buffers(ctx, subctx);
+
+    vk_op_sum_rows_push_constants sp = {};
+    sp.n_cols = nrows;
+    sp.ne01 = 1;
+    sp.ne02 = 1;
+    sp.weight = 1.0f;
+    init_pushconst_fastdiv(sp);
+    sp.misalign_offsets = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
+
+    ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_sum_rows_f32, { tmp_buf, dst_buf }, sp, { 1, 1, 1 });
+    ctx->prealloc_x_need_sync = true;
+}
+
+static void ggml_vk_cross_entropy_loss_back(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
+    const ggml_tensor * grad   = dst->src[0];
+    const ggml_tensor * logits = dst->src[1];
+    const ggml_tensor * labels = dst->src[2];
+
+    GGML_ASSERT(grad->type   == GGML_TYPE_F32);
+    GGML_ASSERT(logits->type == GGML_TYPE_F32);
+    GGML_ASSERT(labels->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type    == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_scalar(grad));
+    GGML_ASSERT(ggml_is_contiguous(grad));
+    GGML_ASSERT(ggml_is_contiguous(logits));
+    GGML_ASSERT(ggml_is_contiguous(labels));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+    GGML_ASSERT(ggml_are_same_shape(logits, labels));
+    GGML_ASSERT(ggml_are_same_shape(logits, dst));
+
+    const uint32_t nclasses = (uint32_t)logits->ne[0];
+    const uint32_t nrows    = (uint32_t)ggml_nrows(logits);
+
+    vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, grad, logits, labels, dst, GGML_OP_CROSS_ENTROPY_LOSS_BACK);
+    GGML_ASSERT(pipeline != nullptr);
+
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+
+    vk_subbuffer grad_buf   = ggml_vk_tensor_subbuffer(ctx, grad);
+    vk_subbuffer logits_buf = ggml_vk_tensor_subbuffer(ctx, logits);
+    vk_subbuffer labels_buf = ggml_vk_tensor_subbuffer(ctx, labels);
+    vk_subbuffer dst_buf    = ggml_vk_tensor_subbuffer(ctx, dst);
+
+    const vk_op_push_constants pc = { nclasses, nrows, 0.0f, 0.0f, 0.0f, 0.0f };
+    ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { grad_buf, logits_buf, labels_buf, dst_buf }, pc, ggml_vk_nrows_elements(nrows));
+}
+
 static void ggml_vk_argmax(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
    ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_ARGMAX, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], 0.0f, 0.0f, 0.0f, 0.0f });
 }
@@ -15687,6 +15940,14 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
    case GGML_OP_ARGMAX:
        ggml_vk_argmax(ctx, compute_ctx, src0, node);

+        break;
+    case GGML_OP_CROSS_ENTROPY_LOSS:
+        ggml_vk_cross_entropy_loss(ctx, compute_ctx, node);
+
+        break;
+    case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
+        ggml_vk_cross_entropy_loss_back(ctx, compute_ctx, node);
+
        break;
    case GGML_OP_COUNT_EQUAL:
        ggml_vk_count_equal(ctx, compute_ctx, src0, src1, node);
@@ -15770,6 +16031,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr

        break;

+    case GGML_OP_LIGHTNING_INDEXER:
+        ggml_vk_lightning_indexer(ctx, compute_ctx, node);
+
+        break;
+
    case GGML_OP_GATED_DELTA_NET:
        ggml_vk_gated_delta_net(ctx, compute_ctx, node);

@@ -17534,20 +17800,32 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
        return;
    }

-    auto const &is_empty = [](ggml_tensor * node) -> bool {
+    auto const &is_empty = [](const ggml_tensor * node) -> bool {
        return node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE;
    };

-    auto const &is_src_of = [](const ggml_tensor *dst, const ggml_tensor *src) -> bool {
+    auto const &is_src_of = [&is_empty](const ggml_tensor *dst, const ggml_tensor *src) -> bool {
+        auto const &base = [](const ggml_tensor * tensor) {
+            return tensor->view_src ? tensor->view_src : tensor;
+        };
        for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) {
            if (dst->src[s] == src) {
                return true;
            }
+            if (is_empty(dst) || is_empty(src)) {
+                continue;
+            }
+            // A source view of dst may read storage written through a different view by src.
+            if (dst->src[s] && base(dst->src[s]) == base(src)) {
+                return true;
+            }
+            // Moving dst forward may overwrite storage still read through a view by src.
+            if (src->src[s] && base(dst) == base(src->src[s])) {
+                return true;
+            }
        }
        // implicit dependency if they view the same tensor
-        const ggml_tensor *dst2 = dst->view_src ? dst->view_src : dst;
-        const ggml_tensor *src2 = src->view_src ? src->view_src : src;
-        if (dst2 == src2) {
+        if (base(dst) == base(src)) {
            return true;
        }
        return false;
@@ -18511,6 +18789,18 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
            }
        case GGML_OP_ARGMAX:
            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
+        case GGML_OP_CROSS_ENTROPY_LOSS:
+            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32
+                && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32
+                && ggml_are_same_shape(op->src[0], op->src[1])
+                && ggml_is_contiguous(op) && ggml_is_scalar(op) && op->type == GGML_TYPE_F32;
+        case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
+            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32 && ggml_is_scalar(op->src[0])
+                && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32
+                && ggml_is_contiguous(op->src[2]) && op->src[2]->type == GGML_TYPE_F32
+                && ggml_are_same_shape(op->src[1], op->src[2])
+                && ggml_are_same_shape(op->src[1], op)
+                && ggml_is_contiguous(op) && op->type == GGML_TYPE_F32;
        case GGML_OP_COUNT_EQUAL:
            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_I32
                && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_I32;
@@ -18536,6 +18826,40 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
        case GGML_OP_GATED_LINEAR_ATTN:
            // the shader block size is hardcoded to head_size 64
            return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && op->src[0]->ne[0] == 64;
+        case GGML_OP_LIGHTNING_INDEXER:
+            {
+                const ggml_tensor * q = op->src[0];
+                const ggml_tensor * k = op->src[1];
+                const ggml_tensor * w = op->src[2];
+                const ggml_tensor * m = op->src[3];
+
+                // the q/w/m types and the shape relationships between q, k, w, m and dst
+                // are already asserted in ggml_lightning_indexer()
+                if (!ggml_vk_lightning_indexer_k_type_supported(k->type) || !device->fp16) {
+                    return false;
+                }
+
+                // the shader block size is hardcoded to head size 128
+                if (q->ne[0] != 128) {
+                    return false;
+                }
+
+                // the shader indexes the buffers by element stride, and is dispatched
+                // without allow_misalign
+                for (const ggml_tensor * t : {q, k, w, m, op}) {
+                    if (t->nb[0] != ggml_type_size(t->type) ||
+                        (vk_tensor_offset(t) + t->view_offs) % device->properties.limits.minStorageBufferOffsetAlignment != 0) {
+                        return false;
+                    }
+                    // the strides get scaled down from bytes, so the division must be exact
+                    for (int i = 1; i < GGML_MAX_DIMS; ++i) {
+                        if (t->nb[i] % ggml_type_size(t->type) != 0) {
+                            return false;
+                        }
+                    }
+                }
+                return true;
+            }
        case GGML_OP_GATED_DELTA_NET:
            {
                const uint32_t S_v = op->src[2]->ne[0];
@@ -19437,6 +19761,10 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
            tensor_clone = ggml_mean(ggml_ctx, src_clone[0]);
        } else if (tensor->op == GGML_OP_ARGMAX) {
            tensor_clone = ggml_argmax(ggml_ctx, src_clone[0]);
+        } else if (tensor->op == GGML_OP_CROSS_ENTROPY_LOSS) {
+            tensor_clone = ggml_cross_entropy_loss(ggml_ctx, src_clone[0], src_clone[1]);
+        } else if (tensor->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK) {
+            tensor_clone = ggml_cross_entropy_loss_back(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]);
        } else if (tensor->op == GGML_OP_COUNT_EQUAL) {
            tensor_clone = ggml_count_equal(ggml_ctx, src_clone[0], src_clone[1]);
        } else if (tensor->op == GGML_OP_SOLVE_TRI) {
@@ -19541,6 +19869,8 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
            const float * op_params = (const float *)tensor->op_params;
            tensor_clone = ggml_gated_linear_attn(ggml_ctx, src_clone[0], src_clone[1],
            src_clone[2], src_clone[3], src_clone[4], op_params[0]);
+        } else if (tensor->op == GGML_OP_LIGHTNING_INDEXER) {
+            tensor_clone = ggml_lightning_indexer(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3]);
        } else if (tensor->op == GGML_OP_GATED_DELTA_NET) {
            tensor_clone = ggml_gated_delta_net(ggml_ctx, src_clone[0], src_clone[1],
            src_clone[2], src_clone[3], src_clone[4], src_clone[5],
@@ -2,6 +2,11 @@

 #extension GL_EXT_control_flow_attributes : enable

+#ifdef USE_SUBGROUPS
+#extension GL_KHR_shader_subgroup_basic : enable
+#extension GL_KHR_shader_subgroup_arithmetic : enable
+#endif
+
 #include "types.glsl"

 layout (push_constant) uniform parameter
@@ -11,6 +16,10 @@ layout (push_constant) uniform parameter
    uint32_t nb00;
    uint32_t nb01;
    uint32_t a_offset;
+    uint32_t n_experts;
+    uint32_t hoist_row_ids;
+    uint32_t ne00mp;
+    uint32_t ne00L;
 } p;

 #define BLOCK_SIZE 256
@@ -21,16 +30,98 @@ layout (binding = 0) readonly buffer A {uint data_a[];};
 layout (binding = 1) writeonly buffer D {uint data_d[];};

 shared uint vals[BLOCK_SIZE];
+shared uint offsets[BLOCK_SIZE];
+shared uint cursors[BLOCK_SIZE];

+// see init_fastdiv_values in ggml-vulkan.cpp
+uint fastdiv(uint n, uint mp, uint L) {
+    uint msbs, lsbs;
+    // msbs = mulhi(n, mp)
+    umulExtended(n, mp, msbs, lsbs);
+    return (msbs + n) >> L;
+}
+
+// data_d layout when p.hoist_row_ids is set:
+//   [0,              n_experts)   per-expert row count
+//   [n_experts,    2*n_experts)   per-expert start offset into the row id region
+//   [2*n_experts]                 total row count
+//   [2*n_experts + 1,         )   row ids grouped by expert, packed as (i01 << 16) | (i00 & 0xffff)
+// Otherwise only data_d[expert_id] is written, holding that expert's row count.
 void main() {
    const uint expert_id = gl_WorkGroupID.x;
    const uint num_elements = p.ne00 * p.ne01;
    const uint tid = gl_LocalInvocationID.x;

+    if (p.hoist_row_ids != 0) {
+        if (tid < p.n_experts) {
+            vals[tid] = 0;
+        }
+        barrier();
+
+        for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
+            const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
+            const uint i00 = idx - i01 * p.ne00;
+            const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
+            if (expert < p.n_experts) {
+                atomicAdd(vals[expert], 1);
+            }
+        }
+        barrier();
+
+#ifdef USE_SUBGROUPS
+        if (gl_SubgroupID == 0) {
+            // pad the trip count so the subgroup ops stay in uniform control flow
+            const uint n_experts_padded = (p.n_experts + gl_SubgroupSize - 1) & ~(gl_SubgroupSize - 1);
+            uint base = 0;
+            for (uint expert = gl_SubgroupInvocationID; expert < n_experts_padded; expert += gl_SubgroupSize) {
+                const bool in_range = expert < p.n_experts;
+                const uint count = in_range ? vals[expert] : 0;
+                const uint offset = base + subgroupExclusiveAdd(count);
+                if (in_range) {
+                    data_d[expert] = count;
+                    data_d[p.n_experts + expert] = offset;
+                    offsets[expert] = offset;
+                    cursors[expert] = 0;
+                }
+                base += subgroupAdd(count);
+            }
+            if (subgroupElect()) {
+                data_d[2 * p.n_experts] = base;
+            }
+        }
+#else
+        if (tid == 0) {
+            uint offset = 0;
+            for (uint expert = 0; expert < p.n_experts; ++expert) {
+                const uint count = vals[expert];
+                data_d[expert] = count;
+                data_d[p.n_experts + expert] = offset;
+                offsets[expert] = offset;
+                cursors[expert] = 0;
+                offset += count;
+            }
+            data_d[2 * p.n_experts] = offset;
+        }
+#endif
+        barrier();
+
+        for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
+            const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
+            const uint i00 = idx - i01 * p.ne00;
+            const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
+            if (expert < p.n_experts) {
+                const uint row = atomicAdd(cursors[expert], 1);
+                const uint packed_row_id = (i01 << 16) | (i00 & 0xffffu);
+                data_d[2 * p.n_experts + 1 + offsets[expert] + row] = packed_row_id;
+            }
+        }
+        return;
+    }
+
    uint count = 0;
    for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
-        const uint i01 = idx / p.ne00;
-        const uint i00 = idx % p.ne00;
+        const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
+        const uint i00 = idx - i01 * p.ne00;
        const uint a = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];

        count += uint(a == expert_id);
@@ -0,0 +1,78 @@
+#version 450
+
+#include "generic_head.glsl"
+#include "types.glsl"
+
+#extension GL_EXT_control_flow_attributes : enable
+
+layout(constant_id = 0) const uint BLOCK_SIZE = 32;
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
+layout (binding = 1) readonly buffer B {B_TYPE data_b[];};
+layout (binding = 2) writeonly buffer D {D_TYPE data_d[];};
+
+shared FLOAT_TYPE tmp[BLOCK_SIZE];
+
+FLOAT_TYPE wg_reduce_max(FLOAT_TYPE v) {
+    const uint tid = gl_LocalInvocationID.x;
+    tmp[tid] = v;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            tmp[tid] = max(tmp[tid], tmp[tid + s]);
+        }
+        barrier();
+    }
+    v = tmp[0];
+    barrier();
+    return v;
+}
+
+FLOAT_TYPE wg_reduce_sum(FLOAT_TYPE v) {
+    const uint tid = gl_LocalInvocationID.x;
+    tmp[tid] = v;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            tmp[tid] += tmp[tid + s];
+        }
+        barrier();
+    }
+    v = tmp[0];
+    barrier();
+    return v;
+}
+
+void main() {
+    const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
+    const uint tid = gl_LocalInvocationID.x;
+
+    if (row >= p.KY) {
+        return;
+    }
+
+    const uint off = row * p.KX;
+
+    FLOAT_TYPE max_logit = FLOAT_TYPE(uintBitsToFloat(0xFF800000));
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        max_logit = max(max_logit, FLOAT_TYPE(data_a[off + i]));
+    }
+    max_logit = wg_reduce_max(max_logit);
+
+    FLOAT_TYPE sum_exp = FLOAT_TYPE(0.0f);
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        sum_exp += exp(FLOAT_TYPE(data_a[off + i]) - max_logit);
+    }
+    const FLOAT_TYPE log_sum = log(wg_reduce_sum(sum_exp));
+
+    FLOAT_TYPE loss = FLOAT_TYPE(0.0f);
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        loss += (FLOAT_TYPE(data_a[off + i]) - max_logit - log_sum) * FLOAT_TYPE(data_b[off + i]);
+    }
+    loss = -wg_reduce_sum(loss) / FLOAT_TYPE(p.KY);
+
+    if (tid == 0) {
+        data_d[row] = D_TYPE(loss);
+    }
+}
@@ -0,0 +1,75 @@
+#version 450
+
+#include "generic_head.glsl"
+#include "types.glsl"
+
+#extension GL_EXT_control_flow_attributes : enable
+
+layout(constant_id = 0) const uint BLOCK_SIZE = 32;
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer G {A_TYPE data_g[];};
+layout (binding = 1) readonly buffer X {B_TYPE data_x[];};
+layout (binding = 2) readonly buffer Y {B_TYPE data_y[];};
+layout (binding = 3) writeonly buffer D {D_TYPE data_d[];};
+
+shared FLOAT_TYPE tmp[BLOCK_SIZE];
+
+FLOAT_TYPE wg_reduce_max(FLOAT_TYPE v) {
+    const uint tid = gl_LocalInvocationID.x;
+    tmp[tid] = v;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            tmp[tid] = max(tmp[tid], tmp[tid + s]);
+        }
+        barrier();
+    }
+    v = tmp[0];
+    barrier();
+    return v;
+}
+
+FLOAT_TYPE wg_reduce_sum(FLOAT_TYPE v) {
+    const uint tid = gl_LocalInvocationID.x;
+    tmp[tid] = v;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            tmp[tid] += tmp[tid + s];
+        }
+        barrier();
+    }
+    v = tmp[0];
+    barrier();
+    return v;
+}
+
+void main() {
+    const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
+    const uint tid = gl_LocalInvocationID.x;
+
+    if (row >= p.KY) {
+        return;
+    }
+
+    const uint off = row * p.KX;
+    const FLOAT_TYPE d_by_nrows = FLOAT_TYPE(data_g[0]) / FLOAT_TYPE(p.KY);
+
+    FLOAT_TYPE max_logit = FLOAT_TYPE(uintBitsToFloat(0xFF800000));
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        max_logit = max(max_logit, FLOAT_TYPE(data_x[off + i]));
+    }
+    max_logit = wg_reduce_max(max_logit);
+
+    FLOAT_TYPE sum_exp = FLOAT_TYPE(0.0f);
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        sum_exp += exp(FLOAT_TYPE(data_x[off + i]) - max_logit);
+    }
+    const FLOAT_TYPE inv_sum = FLOAT_TYPE(1.0f) / wg_reduce_sum(sum_exp);
+
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        const FLOAT_TYPE sm = exp(FLOAT_TYPE(data_x[off + i]) - max_logit) * inv_sum;
+        data_d[off + i] = D_TYPE((sm - FLOAT_TYPE(data_y[off + i])) * d_by_nrows);
+    }
+}
@@ -0,0 +1,55 @@
+#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
+
+// 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;
+        default:           return 1u;
+    }
+}
+
+// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte
+// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number
+// 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);
+        default:           return 1u;
+    }
+}
+
+bool fa_type_needs_shmem(uint ty) {
+    switch (ty) {
+        case FA_TYPE_IQ4_NL: return true;
+        default:             return false;
+    }
+}
+
+#endif // !defined(GGML_FA_TYPES_COMP)
@@ -88,17 +88,7 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
 #define BINDING_IDX_K 0
 #define BINDING_IDX_V 1

-// 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 "fa_types.glsl"

 #if defined(BFLOAT16)
 #define O_TYPE float
@@ -108,45 +98,6 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
 #define O_TYPEV4 FLOAT_TYPEV4
 #endif

-// 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;
-        default:           return 1u;
-    }
-}
-
-// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte
-// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number
-// 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);
-        default:           return 1u;
-    }
-}
-
-bool fa_type_needs_shmem(uint ty) {
-    switch (ty) {
-        case FA_TYPE_IQ4_NL: return true;
-        default:             return false;
-    }
-}
-
 // These can't be `const` globals because GLSL forbids function calls in global
 // const initializers, even when the spec constants would let the driver fold
 // them. Macros expand at the use site and fold after specialization.
@@ -0,0 +1,151 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : require
+#extension GL_EXT_shader_16bit_storage : require
+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
+#extension GL_KHR_shader_subgroup_basic : enable
+#if USE_SUBGROUP_ADD
+#extension GL_KHR_shader_subgroup_arithmetic : enable
+#endif
+
+#define BINDING_IDX_K 0u
+
+#include "types.glsl"
+#include "fa_types.glsl"
+#define FaTypeV FA_TYPE_F32
+
+layout(constant_id = 0) const uint FaTypeK = FA_TYPE_F32;
+layout(constant_id = 1) const uint FaBlockBytesK = 4;
+layout(constant_id = 2) const uint SUBGROUP_SIZE = 32;
+
+#include "flash_attn_dequant.glsl"
+
+// one workgroup computes one output element, one invocation per head element
+#define HEAD_SIZE 128
+
+layout(local_size_x = HEAD_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+layout(binding = 0) readonly buffer QBuf { float q[]; };
+layout(binding = 1) readonly buffer KBufF16 { float16_t k_f16[]; };
+layout(binding = 1) readonly buffer KBufF32 { float k_f32[]; };
+layout(binding = 1) readonly buffer KBufBF16 { uint16_t k_bf16[]; };
+layout(binding = 2) readonly buffer WBuf { float weights[]; };
+layout(binding = 3) readonly buffer MBuf { float16_t mask[]; };
+layout(binding = 4) writeonly buffer DstBuf { float dst[]; };
+
+layout(push_constant) uniform PushConstants {
+    uint n_kv;
+    uint n_heads;
+    uint n_tokens;
+    uint n_streams;
+    uint n_masks;
+    uint dispatch_x;
+    uint q_nb1;
+    uint q_nb2;
+    uint q_nb3;
+    uint k_nb2;
+    uint k_nb3;
+    uint w_nb1;
+    uint w_nb3;
+    uint m_nb1;
+    uint m_nb3;
+    uint d_nb1;
+    uint d_nb3;
+};
+
+shared float k_row[HEAD_SIZE];
+
+#if USE_SUBGROUP_ADD
+shared float sg_partials[HEAD_SIZE / SUBGROUP_SIZE];
+#else
+shared float partials[HEAD_SIZE];
+#endif
+
+void main() {
+    const uint tid = gl_LocalInvocationID.x;
+    const uint output_idx = gl_WorkGroupID.y * dispatch_x + gl_WorkGroupID.x;
+    const uint n_outputs = n_kv * n_tokens * n_streams;
+
+    if (fa_type_needs_shmem(FaTypeK)) {
+        init_iq_shmem(gl_WorkGroupSize);
+    }
+
+    if (output_idx >= n_outputs) {
+        return;
+    }
+
+    const uint ik = output_idx % n_kv;
+    const uint ts = output_idx / n_kv;
+    const uint t = ts % n_tokens;
+    const uint s = ts / n_tokens;
+    const uint k_offset = ik * k_nb2 + s * k_nb3;
+
+    // k strides come in as bytes, so scale them down to the view being indexed
+    const uint k_block_elems = fa_block_elems(FaTypeK);
+    const uint k_elem_bytes = FaBlockBytesK / k_block_elems;
+
+    if (FaTypeK == FA_TYPE_F16) {
+        k_row[tid] = float(k_f16[k_offset / k_elem_bytes + tid]);
+    } else if (FaTypeK == FA_TYPE_F32) {
+        k_row[tid] = k_f32[k_offset / k_elem_bytes + tid];
+    } else if (FaTypeK == FA_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;
+        const uint ib = coord / k_block_elems;
+        const uint iqs = coord % k_block_elems;
+        const vec4 values = dequantize4(ib, iqs, k_offset / FaBlockBytesK, BINDING_IDX_K);
+        k_row[coord + 0] = values.x;
+        k_row[coord + 1] = values.y;
+        k_row[coord + 2] = values.z;
+        k_row[coord + 3] = values.w;
+    }
+    barrier();
+
+    const float k_val = k_row[tid];
+
+    float score = 0.0;
+    for (uint h = 0; h < n_heads; ++h) {
+        const float prod = q[h * q_nb1 + t * q_nb2 + s * q_nb3 + tid] * k_val;
+
+#if USE_SUBGROUP_ADD
+        const float sg_sum = subgroupAdd(prod);
+        if (gl_SubgroupInvocationID == 0) {
+            sg_partials[gl_SubgroupID] = sg_sum;
+        }
+        barrier();
+
+        if (tid == 0) {
+            float sum = 0.0;
+            [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) {
+                sum += sg_partials[i];
+            }
+            score += max(sum, 0.0) * weights[h + t * w_nb1 + s * w_nb3];
+        }
+        // the reads above must complete before the next iteration overwrites sg_partials
+        barrier();
+#else
+        partials[tid] = prod;
+        barrier();
+
+        [[unroll]] for (uint stride = HEAD_SIZE / 2; stride > 0; stride >>= 1) {
+            if (tid < stride) {
+                partials[tid] += partials[tid + stride];
+            }
+            barrier();
+        }
+
+        if (tid == 0) {
+            score += max(partials[0], 0.0) * weights[h + t * w_nb1 + s * w_nb3];
+        }
+        // the read of partials[0] above must complete before the next iteration
+        // overwrites partials[tid]
+        barrier();
+#endif
+    }
+
+    if (tid == 0) {
+        const uint mask_offset = ik + t * m_nb1 + (s % n_masks) * m_nb3;
+        dst[ik + t * d_nb1 + s * d_nb3] = score + float(mask[mask_offset]);
+    }
+}
@@ -88,6 +88,9 @@ layout (push_constant) uniform parameter
    uint nei1;
    uint nbi1;
    uint ne11;
+    uint padded_N;
+    uint n_experts;
+    uint hoist_row_ids;
 #else
    uint base_work_group_z;
    uint num_batches;
@@ -214,27 +217,31 @@ void main() {
    const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B_EFF * LOAD_VEC_BATCH_B / BK;

 #ifdef MUL_MAT_ID
-#ifdef MUL_MAT_ID_USE_SUBGROUPS
-    if (bitCount(p.nei0) == 1) {
-        load_row_ids(expert_idx, true, ic);
+    if (p.hoist_row_ids != 0) {
+        load_row_ids_hoisted(expert_idx, ic);
    } else {
-        load_row_ids(expert_idx, false, ic);
-    }
+#ifdef MUL_MAT_ID_USE_SUBGROUPS
+        if (bitCount(p.nei0) == 1) {
+            load_row_ids(expert_idx, true, ic);
+        } else {
+            load_row_ids(expert_idx, false, ic);
+        }
 #else
-    _ne1 = 0;
-    for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
-        for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
-            if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
-                if (_ne1 >= ic * BN) {
-                    row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
+        _ne1 = 0;
+        for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
+            for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
+                if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
+                    if (_ne1 >= ic * BN) {
+                        row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
+                    }
+                    _ne1++;
                }
-                _ne1++;
            }
        }
-    }

-    barrier();
+        barrier();
 #endif
+    }

    // Workgroup has no work
    if (ic * BN >= _ne1) return;
--- a/Show More
+++ b/Show More