models : support nanbeige4.2-3B (#27730 )

Co-authored-by: admin <lizongqiang@kanzhun.com>
hexagon: support for multi-NPU devices (IQ9, IQ10) and fully asynchronous backend (#26501 )
2026-08-27 10:37:33 +02:00 · 2026-08-27 07:55:31 +03:00 · 2026-08-26 18:46:50 -07:00 · 2026-08-26 23:34:28 +02:00 · 2026-08-26 21:07:24 +02:00 · 2026-08-26 20:19:19 +02:00
124 changed files with 7429 additions and 4334 deletions
@@ -58,34 +58,38 @@ runs:
      if: ${{ inputs.save == 'true' }}
      shell: bash
      run: |
-        set +e -uo pipefail
-        source .venv-hf/bin/activate
-        CCACHE_DIR=$(ccache -k cache_dir)
-        if [[ -d "$CCACHE_DIR" ]]; then
-            ccache -s
-            if [[ -n "${{ inputs.evict-old-files }}" ]]; then
-                ccache --evict-older-than "${{ inputs.evict-old-files }}"
+        if [[ -n "$HF_TOKEN" ]]; then
+            set +e -uo pipefail
+            source .venv-hf/bin/activate
+            CCACHE_DIR=$(ccache -k cache_dir)
+            if [[ -d "$CCACHE_DIR" ]]; then
+                ccache -s
+                if [[ -n "${{ inputs.evict-old-files }}" ]]; then
+                    ccache --evict-older-than "${{ inputs.evict-old-files }}"
+                fi
+                DATESTAMP=$(date -u +'%Y-%m-%dT%H:%M:%SZ')
+                CACHEFILE="${{ inputs.key }}-$DATESTAMP.tar.gz"
+                if tar -czf ccache_bucket.tar.gz -C "$CCACHE_DIR" .; then
+                    hf buckets cp ccache_bucket.tar.gz "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}/$CACHEFILE"
+                fi
+                rm ccache_bucket.tar.gz
+            else
+                echo "'$CCACHE_DIR' not found."
            fi
-            DATESTAMP=$(date -u +'%Y-%m-%dT%H:%M:%SZ')
-            CACHEFILE="${{ inputs.key }}-$DATESTAMP.tar.gz"
-            if tar -czf ccache_bucket.tar.gz -C "$CCACHE_DIR" .; then
-                hf buckets cp ccache_bucket.tar.gz "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}/$CACHEFILE"
-            fi
-            rm ccache_bucket.tar.gz
-        else
-            echo "'$CCACHE_DIR' not found."
        fi

    - name: Remove old ccache files from buckets
      if: ${{ inputs.save == 'true' }}
      shell: bash
      run: |
-        set +e -uo pipefail
-        source .venv-hf/bin/activate
-        CACHE_FILES=$(hf buckets list "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}" --json | jq -r '[.[] | select(.type == "file") | select((.uploaded_at | .[:19]+"Z" | fromdateiso8601) < (now - 5 * 60)) | select(.path | startswith("${{ inputs.folder }}/${{ inputs.key }}") and endswith(".tar.gz"))] | sort_by(.path)[:-1] | .[] | [.path // ""] | @tsv')
-        if [[ -n "$CACHE_FILES" ]]; then
-            echo "Removing old ccache files..."
-            while IFS=$'\t' read -r CACHE_PATH; do
-                hf buckets rm "hf://buckets/${{ inputs.hf_bucket }}/$CACHE_PATH" -y
-            done <<< "$CACHE_FILES"
+        if [[ -n "$HF_TOKEN" ]]; then
+            set +e -uo pipefail
+            source .venv-hf/bin/activate
+            CACHE_FILES=$(hf buckets list "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}" --json | jq -r '[.[] | select(.type == "file") | select((.uploaded_at | .[:19]+"Z" | fromdateiso8601) < (now - 5 * 60)) | select(.path | startswith("${{ inputs.folder }}/${{ inputs.key }}") and endswith(".tar.gz"))] | sort_by(.path)[:-1] | .[] | [.path // ""] | @tsv')
+            if [[ -n "$CACHE_FILES" ]]; then
+                echo "Removing old ccache files..."
+                while IFS=$'\t' read -r CACHE_PATH; do
+                    hf buckets rm "hf://buckets/${{ inputs.hf_bucket }}/$CACHE_PATH" -y
+                done <<< "$CACHE_FILES"
+            fi
        fi
@@ -65,7 +65,7 @@ jobs:
        with:
          key: cuda-ubuntu-24.04-cuda
          folder: llama.cpp
-          hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }}
+          hf_bucket: ggml-org/cache

      - name: Build with CMake
        # TODO: Remove GGML_CUDA_CUB_3DOT2 flag once CCCL 3.2 is bundled within CTK and that CTK version is used in this project
@@ -89,7 +89,7 @@ jobs:
          key: cuda-ubuntu-24.04-cuda
          folder: llama.cpp
          evict-old-files: 1d
-          hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }}
+          hf_bucket: ggml-org/cache
          save: true

  hip:
@@ -120,7 +120,7 @@ jobs:
        with:
          key: cuda-ubuntu-22.04-hip
          folder: llama.cpp
-          hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }}
+          hf_bucket: ggml-org/cache

      - name: Build with native CMake HIP support
        id: cmake_build
@@ -140,7 +140,7 @@ jobs:
          key: cuda-ubuntu-22.04-hip
          folder: llama.cpp
          evict-old-files: 1d
-          hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }}
+          hf_bucket: ggml-org/cache
          save: true

  musa:
@@ -171,7 +171,7 @@ jobs:
        with:
          key: cuda-ubuntu-22.04-musa
          folder: llama.cpp
-          hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }}
+          hf_bucket: ggml-org/cache

      - name: Build with native CMake MUSA support
        id: cmake_build
@@ -189,5 +189,5 @@ jobs:
          key: cuda-ubuntu-22.04-musa
          folder: llama.cpp
          evict-old-files: 1d
-          hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }}
+          hf_bucket: ggml-org/cache
          save: true
@@ -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
@@ -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
@@ -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,8 +714,10 @@ jobs:
        with:
          key: release-windows-2025-vs2026-${{ matrix.arch }}-cpu

+  # TODO: build only the ggml-hip backend like the other windows backend jobs
+  #       (windows-cuda, windows-sycl), then drop the ui-build dependency
  windows-rocm:
-    needs: [check-release]
+    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

    runs-on: windows-2022
@@ -769,6 +736,12 @@ jobs:
        with:
          fetch-depth: 0

+      - name: Download UI build
+        uses: actions/download-artifact@v7
+        with:
+          name: llama-ui.zip
+          path: tools/ui/dist
+
      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
        with:
@@ -879,6 +852,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 +884,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 +946,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 +976,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 +1047,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 +1083,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 +1153,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 +1195,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
@@ -1287,11 +1244,11 @@ jobs:
        with:
          key: release-ubuntu-24.04-sycl-${{ matrix.build }}

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

-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04

    permissions:
      actions: write
@@ -1310,12 +1267,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
@@ -1325,7 +1281,7 @@ jobs:
      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
        with:
-          key: release-ubuntu-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
+          key: release-ubuntu-24.04-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
          evict-old-files: 1d
          max-size: "1G"

@@ -1388,7 +1344,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)

@@ -1414,10 +1369,10 @@ jobs:
      - name: ccache-clear
        uses: ./.github/actions/ccache-clear
        with:
-          key: release-ubuntu-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
+          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 +1400,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 +1523,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,14 +1540,13 @@ jobs:
    runs-on: ubuntu-slim

    needs:
-      - get-version
      - windows
      - windows-cpu
      - windows-cuda
      - windows-sycl
      - windows-rocm
      - windows-openvino
-      - ubuntu-22-rocm
+      - ubuntu-24-rocm
      - ubuntu-cpu
      - ubuntu-vulkan
      - ubuntu-24-openvino
@@ -1628,24 +1579,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 +1623,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)
@@ -202,6 +202,10 @@ class NemotronHModel(GraniteHybridModel):
    is_moe: bool = False
    supports_mtp_export = True

+    _SSM_LAYER_TYPES = {"mamba", "linear_attention"}
+    _ATTN_LAYER_TYPES = {"attention", "full_attention"}
+    _MLP_LAYER_TYPES = {"moe"}
+
    def __init__(self, *args, **kwargs):
        # We have to determine the correct model architecture (MoE vs non-MoE) before
        # calling the parent __init__. This is because the parent constructor
@@ -242,8 +246,8 @@ class NemotronHModel(GraniteHybridModel):
            self._ssm_layers = [i for i, val in enumerate(pattern) if val == "M"]
            self._mlp_layers = [i for i, val in enumerate(pattern) if val == ("E" if self.is_moe else "-")]
        else:
-            self._ssm_layers = [i for i, val in enumerate(pattern) if val == "mamba"]
-            self._mlp_layers = [i for i, val in enumerate(pattern) if val == "moe"]
+            self._ssm_layers = [i for i, val in enumerate(pattern) if val in self._SSM_LAYER_TYPES]
+            self._mlp_layers = [i for i, val in enumerate(pattern) if val in self._MLP_LAYER_TYPES]

        # `--no-mtp` drops it entirely; `--mtp` exports only the MTP head
        self._mtp_bid: int | None = None
@@ -272,7 +276,7 @@ class NemotronHModel(GraniteHybridModel):
        if isinstance(pattern, str):
            return [i for i, val in enumerate(pattern) if val == "*"]

-        return [i for i, val in enumerate(pattern) if val == "attention"]
+        return [i for i, val in enumerate(pattern) if val in self._ATTN_LAYER_TYPES]

    @classmethod
    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
@@ -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"
@@ -110,6 +110,16 @@ set_and_check(GGML_INCLUDE_DIR "@PACKAGE_GGML_INCLUDE_INSTALL_DIR@")
 set_and_check(GGML_LIB_DIR "@PACKAGE_GGML_LIB_INSTALL_DIR@")
 #set_and_check(GGML_BIN_DIR "@PACKAGE_GGML_BIN_INSTALL_DIR@")

+if (NOT GGML_SHARED_LIB AND GGML_CPU_KLEIDIAI)
+    unset(KLEIDIAI_LIBRARY CACHE)
+    unset(KLEIDIAI_LIBRARY)
+    find_library(KLEIDIAI_LIBRARY kleidiai
+        REQUIRED
+        HINTS ${GGML_LIB_DIR}
+        NO_CMAKE_FIND_ROOT_PATH)
+    list(APPEND GGML_CPU_INTERFACE_LINK_LIBRARIES ${KLEIDIAI_LIBRARY})
+endif()
+
 if(NOT TARGET ggml::ggml)
    find_package(Threads REQUIRED)

@@ -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
@@ -83,6 +83,7 @@ extern "C" {
    GGML_API ggml_backend_buffer_t ggml_backend_multi_buffer_alloc_buffer(ggml_backend_buffer_t * buffers, size_t n_buffers);
    GGML_API bool                  ggml_backend_buffer_is_multi_buffer(ggml_backend_buffer_t buffer);
    GGML_API void                  ggml_backend_multi_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage);
+    GGML_API void                  ggml_backend_meta_buffer_set_usage (ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage);

    //
    // Backend (meta)
@@ -1168,7 +1168,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
 }

 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync) {
-    GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer));
    ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context;
    return ggml_backend_meta_get_split_state(buf_ctx->get_simple_tensor_container(tensor), tensor, assume_sync);
 }
@@ -1259,7 +1258,14 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m
            t_ij->data = (char *) ggml_backend_buffer_get_base(simple_buf)
                + size_t(tensor->data) - size_t(ggml_backend_buffer_get_base(tensor->buffer));
        }
-        t_ij->extra = tensor->extra;
+
+        if (simple_buf) {
+            // the backend that owns the buffer will set .extra
+            ggml_backend_buffer_init_tensor(simple_buf, t_ij);
+        } else {
+            t_ij->extra = tensor->extra;
+        }
+
        for (int i = 0; i < GGML_MAX_SRC; i++) {
            t_ij->src[i] = tensor->src[i];
            if (tensor->src[i] == tensor) {
@@ -1668,6 +1674,16 @@ bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf) {
    return buf != nullptr && buf->iface.free_buffer == ggml_backend_meta_buffer_iface.free_buffer;
 }

+void ggml_backend_meta_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage) {
+    GGML_ASSERT(ggml_backend_buffer_is_meta(buffer));
+    ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context;
+    for (size_t i = 0; i < buf_ctx->bufs.size(); i++) {
+        if (buf_ctx->bufs[i]) {
+            ggml_backend_buffer_set_usage(buf_ctx->bufs[i].get(), usage);
+        }
+    }
+}
+
 static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
    const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft);

@@ -182,6 +182,8 @@ void ggml_backend_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backe
    // FIXME: add a generic callback to the buffer interface
    if (ggml_backend_buffer_is_multi_buffer(buffer)) {
        ggml_backend_multi_buffer_set_usage(buffer, usage);
+    } else if (ggml_backend_buffer_is_meta(buffer)) {
+        ggml_backend_meta_buffer_set_usage(buffer, usage);
    }
 }

@@ -576,10 +576,25 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
    endif()

    if (GGML_CPU_KLEIDIAI)
-        message(STATUS "Using KleidiAI optimized kernels if applicable")
+        # upstream repo requires at least cmake 3.16
+        if (CMAKE_VERSION VERSION_LESS 3.16)
+            message(FATAL_ERROR "GGML_CPU_KLEIDIAI requires CMake >= 3.16")
+        endif()

-        # Disable the KleidiAI tests
-        set(KLEIDIAI_BUILD_TESTS  OFF)
+        set(GGML_CPU_KLEIDIAI_AARCH64 OFF)
+        if (GGML_SYSTEM_ARCH STREQUAL "ARM" AND
+            (APPLE OR WIN32 OR CMAKE_SYSTEM_NAME MATCHES "^(Linux|Android)$") AND
+            (CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64|ARM64|arm64-v8a)$" OR
+             CMAKE_OSX_ARCHITECTURES MATCHES "arm64" OR
+             CMAKE_GENERATOR_PLATFORM_LWR STREQUAL "arm64" OR
+             CMAKE_ANDROID_ARCH_ABI STREQUAL "arm64-v8a"))
+            set(GGML_CPU_KLEIDIAI_AARCH64 ON)
+        endif()
+        if (NOT GGML_CPU_KLEIDIAI_AARCH64)
+            message(FATAL_ERROR "GGML_CPU_KLEIDIAI requires a Linux, Android, Apple, or Windows AArch64/arm64 target")
+        endif()
+
+        message(STATUS "Using KleidiAI optimized kernels if applicable")

        # Fetch KleidiAI sources:
        include(FetchContent)
@@ -595,31 +610,49 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
            list(APPEND KLEIDIAI_FETCH_ARGS DOWNLOAD_EXTRACT_TIMESTAMP NEW)
        endif()

-        if (CMAKE_VERSION VERSION_GREATER_EQUAL "3.28")
-            FetchContent_Declare(KleidiAI_Download
-                ${KLEIDIAI_FETCH_ARGS}
+        FetchContent_Declare(kleidiai
+            ${KLEIDIAI_FETCH_ARGS}
+        )
+
+        # Disable tests and benchmark building
+        set(KLEIDIAI_BUILD_TESTS OFF CACHE BOOL "" FORCE)
+        set(KLEIDIAI_BUILD_BENCHMARK OFF CACHE BOOL "" FORCE)
+
+        # Use the Populate/add_subdirectory flow for compatibility with CMake 3.16.
+        FetchContent_GetProperties(kleidiai
+            SOURCE_DIR KLEIDIAI_SRC
+            BINARY_DIR KLEIDIAI_BIN
+            POPULATED KLEIDIAI_POPULATED
+        )
+        if (NOT KLEIDIAI_POPULATED)
+            FetchContent_Populate(kleidiai)
+            FetchContent_GetProperties(kleidiai
+                SOURCE_DIR KLEIDIAI_SRC
+                BINARY_DIR KLEIDIAI_BIN
+            )
+        endif()
+
+        if (NOT TARGET kleidiai)
+            add_subdirectory(
+                "${CMAKE_CURRENT_SOURCE_DIR}/ggml-cpu/kleidiai"
+                "${CMAKE_CURRENT_BINARY_DIR}/kleidiai-wrapper"
                EXCLUDE_FROM_ALL
            )
-
-            FetchContent_MakeAvailable(KleidiAI_Download)
-            FetchContent_GetProperties(KleidiAI_Download SOURCE_DIR KLEIDIAI_SRC)
-        else()
-            FetchContent_Declare(KleidiAI_Download
-                ${KLEIDIAI_FETCH_ARGS}
-            )
-
-            FetchContent_GetProperties(KleidiAI_Download
-                SOURCE_DIR KLEIDIAI_SRC
-                POPULATED KLEIDIAI_POPULATED
-            )
-
-            if (NOT KLEIDIAI_POPULATED)
-                FetchContent_Populate(KleidiAI_Download)
-                FetchContent_GetProperties(KleidiAI_Download SOURCE_DIR KLEIDIAI_SRC)
+            if (NOT CMAKE_SKIP_INSTALL_RULES AND
+                (NOT DEFINED BUILD_SHARED_LIBS OR NOT BUILD_SHARED_LIBS))
+                install(TARGETS kleidiai ARCHIVE)
            endif()
        endif()

-        add_compile_definitions(GGML_USE_CPU_KLEIDIAI)
+        if (NOT TARGET kleidiai)
+            message(FATAL_ERROR "KleidiAI target was not created")
+        endif()
+
+        set_target_properties(kleidiai PROPERTIES POSITION_INDEPENDENT_CODE ON)
+
+        target_link_libraries(${GGML_CPU_NAME} PRIVATE kleidiai)
+
+        target_compile_definitions(${GGML_CPU_NAME} PRIVATE GGML_USE_CPU_KLEIDIAI)

        list(APPEND GGML_CPU_SOURCES
            ggml-cpu/kleidiai/kleidiai.cpp
@@ -627,108 +660,6 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
            ggml-cpu/kleidiai/kleidiai.h
            ggml-cpu/kleidiai/kernels.h
            )
-
-        # KleidiAI
-        include_directories(
-            ${KLEIDIAI_SRC}/
-            ${KLEIDIAI_SRC}/kai/
-            ${KLEIDIAI_SRC}/kai/ukernels/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/)
-
-        set(ARCH_FLAGS_TEMP "${ARCH_FLAGS}")
-        if (NOT ARCH_FLAGS_TEMP)
-            string(REGEX MATCH "-march=[^ ]+" ARCH_FLAGS_TEMP "${CMAKE_C_FLAGS}")
-        endif()
-        string(FIND "${ARCH_FLAGS_TEMP}" "+dotprod" DOTPROD_ENABLED)
-        string(FIND "${ARCH_FLAGS_TEMP}" "+i8mm" I8MM_ENABLED)
-        string(FIND "${ARCH_FLAGS_TEMP}" "+sme" SME_ENABLED)
-        string(FIND "${ARCH_FLAGS_TEMP}" "+sve" SVE_ENABLED)
-
-        set(PRIVATE_ARCH_FLAGS ${ARCH_FLAGS_TEMP})
-
-        list(APPEND GGML_KLEIDIAI_SOURCES
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p_f32.c
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p4x8sb_f32_neon.c
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon.c
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p_f32_neon.c
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0.c
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_quant_pack_qai8dxp_f32.c
-            ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi8cxp_qsi8cx_neon.c)
-
-        if (NOT DOTPROD_ENABLED MATCHES -1)
-            list(APPEND GGML_KLEIDIAI_SOURCES
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4x4_1x4_neon_dotprod.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x4_qsi4c32p4x4_16x4_neon_dotprod.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp4x4_qsi8cxp4x4_16x4_neon_dotprod.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4x4_1x4_neon_dotprod.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x8_qsi8cxp4x8_1x4_neon_dotprod.c)
-        endif()
-
-        if (NOT I8MM_ENABLED MATCHES -1)
-            list(APPEND GGML_KLEIDIAI_SOURCES
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p4x8_16x4_neon_i8mm.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp4x8_qsi8cxp4x8_16x4_neon_i8mm.c)
-        endif()
-
-        if (NOT SME_ENABLED MATCHES -1)
-            list(APPEND GGML_KLEIDIAI_SME_SOURCES
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa_asm.S)
-            set_source_files_properties(${GGML_KLEIDIAI_SME_SOURCES}
-                PROPERTIES COMPILE_OPTIONS "-fno-tree-vectorize;${ARCH_FLAGS_TEMP}+sve+sve2+sme")
-            list(APPEND GGML_CPU_SOURCES ${GGML_KLEIDIAI_SME_SOURCES})
-
-            list(APPEND GGML_KLEIDIAI_SME2_SOURCES
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/kai_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/kai_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_bf16p2vlx2_f32_sme.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f16pmrx2_f32_neon.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f32p2vlx1_f32_sme.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f32p2vlx1_f32_sme_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme_asm.S
-                ${KLEIDIAI_SRC}/kai/kai_common_sme_asm.S)
-            set_source_files_properties(${GGML_KLEIDIAI_SME2_SOURCES}
-                PROPERTIES COMPILE_OPTIONS "-fno-tree-vectorize;${ARCH_FLAGS_TEMP}+sve+sve2+sme2+fp16")
-            list(APPEND GGML_CPU_SOURCES ${GGML_KLEIDIAI_SME2_SOURCES})
-            set(PRIVATE_ARCH_FLAGS "-fno-tree-vectorize;${PRIVATE_ARCH_FLAGS}")
-        endif()
-
-        if (NOT SVE_ENABLED MATCHES -1)
-            list(APPEND GGML_KLEIDIAI_SOURCES
-                ${KLEIDIAI_SRC}/kai/kai_common_sve_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod.c
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p8x8_16x8_sve_i8mm_asm.S
-                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p8x8_16x8_sve_i8mm.c)
-        endif()
-
-        set_source_files_properties(${GGML_KLEIDIAI_SOURCES} PROPERTIES COMPILE_OPTIONS "${PRIVATE_ARCH_FLAGS}")
-        list(APPEND GGML_CPU_SOURCES ${GGML_KLEIDIAI_SOURCES})
    endif()

    message(STATUS "Adding CPU backend variant ${GGML_CPU_NAME}: ${ARCH_FLAGS} ${ARCH_DEFINITIONS}")
@@ -0,0 +1,14 @@
+set(BUILD_SHARED_LIBS OFF)
+set(CMAKE_SKIP_INSTALL_RULES TRUE)
+
+add_subdirectory("${KLEIDIAI_SRC}" "${KLEIDIAI_BIN}" EXCLUDE_FROM_ALL)
+
+if (NOT TARGET kleidiai)
+    message(FATAL_ERROR "KleidiAI target was not created")
+endif()
+
+if (MSVC)
+    target_compile_options(kleidiai PRIVATE $<$<COMPILE_LANGUAGE:C,CXX>:/WX->)
+else()
+    target_compile_options(kleidiai PRIVATE $<$<COMPILE_LANGUAGE:C,CXX>:-Wno-error>)
+endif()
@@ -3,44 +3,44 @@
 //

 // KleidiAI micro-kernels
-#include "kai_matmul_clamp_f32_qsi8d32p_qsi4c32p_interface.h"
-#include "kai_matmul_clamp_f32_qai8dxp_qsi8cxp_interface.h"
-#include "kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod.h"
-#include "kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4x4_1x4_neon_dotprod.h"
-#include "kai_matmul_clamp_f32_qsi8d32p4x4_qsi4c32p4x4_16x4_neon_dotprod.h"
-#include "kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p4x8_16x4_neon_i8mm.h"
-#include "kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot.h"
-#include "kai_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa.h"
-#include "kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa.h"
-#include "kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot.h"
-#include "kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa.h"
-#include "kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot.h"
-#include "kai_matmul_clamp_f32_qai8dxp1x8_qsi8cxp4x8_1x4_neon_dotprod.h"
-#include "kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4x4_1x4_neon_dotprod.h"
-#include "kai_matmul_clamp_f32_qai8dxp4x4_qsi8cxp4x4_16x4_neon_dotprod.h"
-#include "kai_matmul_clamp_f32_qai8dxp4x8_qsi8cxp4x8_16x4_neon_i8mm.h"
-#include "kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p8x8_16x8_sve_i8mm.h"
-#include "kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod.h"
-#include "kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.h"
-#include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.h"
-#include "kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.h"
-#include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p_qsi4c32p_interface.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp_qsi8cxp_interface.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4x4_1x4_neon_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x4_qsi4c32p4x4_16x4_neon_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p4x8_16x4_neon_i8mm.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot.h"
+#include "kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/kai_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x8_qsi8cxp4x8_1x4_neon_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4x4_1x4_neon_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp4x4_qsi8cxp4x4_16x4_neon_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp4x8_qsi8cxp4x8_16x4_neon_i8mm.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p8x8_16x8_sve_i8mm.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.h"
+#include "kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.h"

-#include "kai_lhs_pack_bf16p2vlx2_f32_sme.h"
-#include "kai_lhs_pack_f32p2vlx1_f32_sme.h"
-#include "kai_lhs_quant_pack_qsi8d32p_f32.h"
-#include "kai_lhs_quant_pack_qsi8d32p4x8sb_f32_neon.h"
-#include "kai_lhs_quant_pack_qsi8d32p_f32_neon.h"
-#include "kai_lhs_quant_pack_qai8dxp_f32.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_pack_bf16p2vlx2_f32_sme.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_pack_f32p2vlx1_f32_sme.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p_f32.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p4x8sb_f32_neon.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p_f32_neon.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_quant_pack_qai8dxp_f32.h"

-#include "kai_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme.h"
-#include "kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.h"
-#include "kai_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0.h"
-#include "kai_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon.h"
-#include "kai_rhs_pack_nxk_qsi8cxp_qsi8cx_neon.h"
-#include "kai_lhs_pack_f16pmrx2_f32_neon.h"
+#include "kai/ukernels/matmul/pack/kai_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme.h"
+#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.h"
+#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0.h"
+#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon.h"
+#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi8cxp_qsi8cx_neon.h"
+#include "kai/ukernels/matmul/pack/kai_lhs_pack_f16pmrx2_f32_neon.h"

-#include "kai_common.h"
+#include "kai/kai_common.h"

 #include "simd-mappings.h"

@@ -328,9 +328,8 @@ static void dequantize_row_qsi8cxp(
 }

 static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
-#if defined(__ARM_FEATURE_SME)
    {
-        /* SME GEMM */
+        /* SME2 GEMM */
        /* .kern_info = */ {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa,
@@ -351,7 +350,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_size_ex        = */ &lhs_ps_fn6<kai_get_lhs_packed_size_lhs_pack_f16pmrx2_f32_neon>,
            /* .pack_func_ex          = */ &lhs_pack_void_fn10<kai_run_lhs_pack_f16pmrx2_f32_neon>,
        },
-        /* SME GEMV */
+        /* SME2 GEMV */
        /* .kern_info = */ {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot,
@@ -378,13 +377,13 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_stride_ex      = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon>,
            /* .pack_func_ex          = */ &rhs_pack_fn12<kai_run_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon>,
        },
-        /* .required_cpu       = */ CPU_FEATURE_SME2,
+        /* .required_cpu       = */ CPU_FEATURE_SME2 | CPU_FEATURE_FP16,
        /* .lhs_type           = */ GGML_TYPE_F32,
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
    {
-        /* SME GEMM */
+        /* SME2 GEMM */
        /* .kern_info = */ {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa,
@@ -404,7 +403,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_pack_bf16p2vlx2_f32_sme>,
            /* .pack_func_ex          = */ &lhs_pack_void_fn9<kai_run_lhs_pack_bf16p2vlx2_f32_sme>,
        },
-        /* SME GEMV */
+        /* SME2 GEMV */
        /* .kern_info = */ {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa,
@@ -436,9 +435,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
        /* .rhs_type           = */ GGML_TYPE_F16,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
 #if defined(__APPLE__)
-#if defined(__ARM_FEATURE_DOTPROD)
    {
        /* DOTPROD GEMM */
        /* .kern_info = */ {
@@ -492,8 +489,6 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
-#if defined(__ARM_FEATURE_MATMUL_INT8)
    {
        /* i8mm GEMM */
        /* .kern_info = */ {
@@ -515,7 +510,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_size_ex        = */ &lhs_ps_fn6<kai_get_lhs_packed_size_lhs_quant_pack_qsi8d32p4x8sb_f32_neon>,
            /* .pack_func_ex          = */ &lhs_pack_float_fn10<kai_run_lhs_quant_pack_qsi8d32p4x8sb_f32_neon>,
        },
-        /* i8mm GEMV */
+        /* DOTPROD GEMV */
        /* .kern_info = */ {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod,
@@ -542,14 +537,12 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_stride_ex      = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
            /* .pack_func_ex          = */ &rhs_pack_fn12<kai_run_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
        },
-        /* .required_cpu       = */ CPU_FEATURE_I8MM,
+        /* .required_cpu       = */ CPU_FEATURE_I8MM | CPU_FEATURE_DOTPROD,
        /* .lhs_type           = */ GGML_TYPE_F32,
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
 #else
-#if defined(__ARM_FEATURE_SVE)
    {
        /* SVE i8mm GEMM */
        /* .kern_info = */ {
@@ -603,8 +596,6 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
-#if defined(__ARM_FEATURE_MATMUL_INT8)
    {
        /* i8mm GEMM */
        /* .kern_info = */ {
@@ -626,7 +617,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_size_ex        = */ &lhs_ps_fn6<kai_get_lhs_packed_size_lhs_quant_pack_qsi8d32p4x8sb_f32_neon>,
            /* .pack_func_ex          = */ &lhs_pack_float_fn10<kai_run_lhs_quant_pack_qsi8d32p4x8sb_f32_neon>,
        },
-        /* i8mm GEMV */
+        /* DOTPROD GEMV */
        /* .kern_info = */ {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod,
@@ -653,13 +644,11 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_stride_ex      = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
            /* .pack_func_ex          = */ &rhs_pack_fn12<kai_run_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
        },
-        /* .required_cpu       = */ CPU_FEATURE_I8MM,
+        /* .required_cpu       = */ CPU_FEATURE_I8MM | CPU_FEATURE_DOTPROD,
        /* .lhs_type           = */ GGML_TYPE_F32,
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif // __ARM_FEATURE_MATMUL_INT8
-#if defined(__ARM_FEATURE_DOTPROD)
    {
        /* DOTPROD GEMM */
        /* .kern_info = */ {
@@ -713,15 +702,13 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
 #endif
    { /* Sentinel */ }
 };

 static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
-#if defined(__ARM_FEATURE_SME)
    {
-        /* SME GEMM */
+        /* SME2 GEMM */
        {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa,
@@ -741,7 +728,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
            /* .pack_func_ex          = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
        },
-        /* SME GEMV */
+        /* SME2 GEMV */
        {
            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot,
            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot,
@@ -826,8 +813,6 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
        /* .rhs_type           = */ GGML_TYPE_Q8_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
-#if defined(__ARM_FEATURE_MATMUL_INT8)
    {
        /* I8MM GEMM */
        {
@@ -876,13 +861,11 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
            /* .packed_stride_ex      = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
            /* .pack_func_ex          = */ &rhs_pack_scale_fn12<kai_run_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
        },
-        /* .required_cpu       = */ CPU_FEATURE_I8MM,
+        /* .required_cpu       = */ CPU_FEATURE_I8MM | CPU_FEATURE_DOTPROD,
        /* .lhs_type           = */ GGML_TYPE_F32,
        /* .rhs_type           = */ GGML_TYPE_Q8_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
-#if defined(__ARM_FEATURE_DOTPROD)
    {
        /* DOTPROD GEMM */
        {
@@ -936,12 +919,10 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
        /* .rhs_type           = */ GGML_TYPE_Q8_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
    { /* Sentinel */ }
 };

 static ggml_kleidiai_kernels ggml_kleidiai_kernels_f32[] = {
-#if defined(__ARM_FEATURE_SME)
    {
        /* SME2 GEMM */
        {
@@ -1048,7 +1029,6 @@ static ggml_kleidiai_kernels ggml_kleidiai_kernels_f32[] = {
        /* .rhs_type           = */ GGML_TYPE_F32,
        /* .op_type            = */ GGML_TYPE_F32,
    },
-#endif
    { /* Sentinel */ }
 };

@@ -1056,10 +1036,6 @@ ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, c
    ggml_kleidiai_kernels * kernel = nullptr;

    if (tensor->op == GGML_OP_MUL_MAT && tensor->src[0] != nullptr && tensor->src[1] != nullptr) {
-#if defined(__ARM_FEATURE_SME)          ||  \
-    defined(__ARM_FEATURE_DOTPROD)      ||  \
-    defined(__ARM_FEATURE_MATMUL_INT8)  ||  \
-    defined(__ARM_FEATURE_SVE)
        auto try_table = [&](auto & table) {
            for (size_t i = 0; i < NELEMS(table) - 1; ++i) {
                if ((cpu_features & table[i].required_cpu) == table[i].required_cpu &&
@@ -1080,12 +1056,6 @@ ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, c
        } else {
            try_table(gemm_gemv_kernels);
        }
-#else
-    GGML_UNUSED(gemm_gemv_kernels);
-    GGML_UNUSED(gemm_gemv_kernels_q8);
-    GGML_UNUSED(ggml_kleidiai_kernels_f32);
-    GGML_UNUSED(cpu_features);
-#endif
    }

    return kernel;
@@ -1094,19 +1064,13 @@ ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, c
 ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q4_0(cpu_feature features) {
    ggml_kleidiai_kernels * kernels = nullptr;

-#if defined(__ARM_FEATURE_SME)          ||  \
-    defined(__ARM_FEATURE_DOTPROD)      ||  \
-    defined(__ARM_FEATURE_MATMUL_INT8)  ||  \
-    defined(__ARM_FEATURE_SVE)
    for (size_t i = 0; i < NELEMS(gemm_gemv_kernels) - 1; ++i) {
-        if ((features & gemm_gemv_kernels[i].required_cpu) == gemm_gemv_kernels[i].required_cpu) {
+        if ((features & gemm_gemv_kernels[i].required_cpu) == gemm_gemv_kernels[i].required_cpu &&
+            gemm_gemv_kernels[i].rhs_type == GGML_TYPE_Q4_0) {
            kernels = &gemm_gemv_kernels[i];
            break;
        }
    }
-#else
-    GGML_UNUSED(features);
-#endif

    return kernels;
 }
@@ -1114,16 +1078,12 @@ ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q4_0(cpu_feature features)
 ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q8_0(cpu_feature features) {
    ggml_kleidiai_kernels * kernels = nullptr;

-#if defined(__ARM_FEATURE_SME) || defined(__ARM_FEATURE_DOTPROD) || defined(__ARM_FEATURE_MATMUL_INT8)
    for (size_t i = 0; i < NELEMS(gemm_gemv_kernels_q8) - 1; ++i) {
        if ((features & gemm_gemv_kernels_q8[i].required_cpu) == gemm_gemv_kernels_q8[i].required_cpu) {
            kernels = &gemm_gemv_kernels_q8[i];
            break;
        }
    }
-#else
-    GGML_UNUSED(features);
-#endif

    return kernels;
 }
@@ -1131,16 +1091,11 @@ ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q8_0(cpu_feature features)
 ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_f32(cpu_feature features) {
    ggml_kleidiai_kernels * kernels = nullptr;

-#if defined(__ARM_FEATURE_SME)
    for (size_t i = 0; i < NELEMS(ggml_kleidiai_kernels_f32) - 1; ++i) {
        if ((features & ggml_kleidiai_kernels_f32[i].required_cpu) == ggml_kleidiai_kernels_f32[i].required_cpu) {
            kernels = &ggml_kleidiai_kernels_f32[i];
            break;
        }
    }
-#else
-    GGML_UNUSED(features);
-#endif
-
    return kernels;
 }
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: Copyright 2025 Arm Limited and/or its affiliates <open-source-office@arm.com>
+// SPDX-FileCopyrightText: Copyright 2025-2026 Arm Limited and/or its affiliates <open-source-office@arm.com>
 // SPDX-License-Identifier: MIT
 //

@@ -12,7 +12,8 @@ enum cpu_feature {
    CPU_FEATURE_I8MM    = 2,
    CPU_FEATURE_SVE     = 4,
    CPU_FEATURE_SME     = 8,
-    CPU_FEATURE_SME2    = 16
+    CPU_FEATURE_SME2    = 16,
+    CPU_FEATURE_FP16    = 32
 };

 inline cpu_feature& operator|=(cpu_feature& lhs, cpu_feature rhs) {
@@ -48,7 +48,7 @@

 #include "kernels.h"

-#include "kai_common.h"
+#include "kai/kai_common.h"

 #define GGML_COMMON_DECL_CPP
 #include "ggml-common.h"
@@ -316,6 +316,7 @@ static void init_kleidiai_context(void) {

        ctx.features  = (runtime_feat.has_dotprod  ? CPU_FEATURE_DOTPROD : CPU_FEATURE_NONE) |
                        (runtime_feat.has_i8mm     ? CPU_FEATURE_I8MM    : CPU_FEATURE_NONE) |
+                        (runtime_feat.has_fp16     ? CPU_FEATURE_FP16    : CPU_FEATURE_NONE) |
                        (runtime_feat.sve_cnt == QK8_0 ? CPU_FEATURE_SVE : CPU_FEATURE_NONE);

        if (env_threads) {
@@ -1,4 +1,4 @@
-static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal(ggml_type type, int J, bool fallback) {
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal_dp4a(ggml_type type, int J, bool fallback) {
    CASE(GGML_TYPE_Q1_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
@@ -0,0 +1,273 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal_older(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 1, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
@@ -314,7 +314,9 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
    }

    if (ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_DP4A) {
-        return false;
+        // for MoE, mmq is faster even without native dp4a
+        // TODO: check if cards older than pascal might benefit from this as well
+        return cc >= GGML_CUDA_CC_PASCAL && n_experts > 0;
    }

 #ifdef GGML_CUDA_FORCE_MMQ
@@ -213,7 +213,8 @@ struct ggml_cuda_mmq_config {
        return ggml_cuda_mmq_config((type_), (nthreads_), (occupancy_), (I_), (J_), (sram_layout_), (K_vram_), (stream_k_), (fallback_)); \
    }                                                                                                                                     \

-#include "mmq-config-pascal.cuh"
+#include "mmq-config-pascal-older.cuh"
+#include "mmq-config-pascal-dp4a.cuh"
 #include "mmq-config-ampere.cuh"
 #include "mmq-config-blackwell.cuh"

@@ -247,7 +248,10 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty
    if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) {
        return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
    }
-    return ggml_cuda_mmq_get_config_pascal(type, J, fallback);
+    if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_DP4A) {
+        return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback);
+    }
+    return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback);
 }

 static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback) {
@@ -268,8 +272,10 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t
    return ggml_cuda_mmq_get_config_blackwell(type, J, fallback);
 #elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
    return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
+#elif __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A
+    return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback);
 #else
-    return ggml_cuda_mmq_get_config_pascal(type, J, fallback);
+    return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback);
 #endif // BLACKWELL_MMA_AVAILABLE
 #endif // GGML_USE_HIP
    GGML_UNUSED_VARS(type, J, fallback);
@@ -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,
@@ -99,12 +91,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 +107,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 +173,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 +214,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 +231,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); \
 } \

@@ -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:
@@ -818,12 +867,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 +882,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 +892,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 +1120,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
@@ -572,7 +572,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched
    return res;
 }

-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_metal_library_t lib, const ggml_tensor * op)  {
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_metal_library_t lib, const ggml_tensor * op, bool tail)  {
    GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);

    char base[256];
@@ -580,7 +580,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_me

    const int nsg = (ne00 + 31)/32;

-    snprintf(base, 256, "kernel_ssm_scan_%s", ggml_type_name(op->src[0]->type));
+    snprintf(base, 256, "kernel_ssm_scan_%s%s", ggml_type_name(op->src[0]->type), tail ? "_tail" : "");
    snprintf(name, 256, "%s_nsg=%d", base, nsg);

    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
@@ -598,6 +598,27 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_me
    return res;
 }

+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan_ssd_mma(ggml_metal_library_t lib, const ggml_tensor * op)  {
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_ssm_scan_ssd_mma_%s", ggml_type_name(op->src[0]->type));
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+    }
+
+    // acs/exp(acs)/state-decay vectors + dtX + SAM rows + two 8x8 tiles per simdgroup
+    res.smem = (3*OP_SSM_SCAN_SSD_CS +
+                OP_SSM_SCAN_SSD_CS*OP_SSM_SCAN_SSD_HD +
+                OP_SSM_SCAN_SSD_NSG*8*OP_SSM_SCAN_SSD_CS +
+                OP_SSM_SCAN_SSD_NSG*2*8*8)*sizeof(float);
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv(ggml_metal_library_t lib, const ggml_tensor * op) {
    char base[256];
    char name[256];
@@ -129,7 +129,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc           (ggml_metal_library_t lib, enum ggml_op op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv          (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched  (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs);
-struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan          (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan          (ggml_metal_library_t lib, const struct ggml_tensor * op, bool tail);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan_ssd_mma  (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv              (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net   (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri         (ggml_metal_library_t lib, const struct ggml_tensor * op);
@@ -158,6 +158,10 @@
 #define OP_SUM_ROWS_NUM_SUM_ROWS 10
 #define OP_SUM_ROWS_NUM_MEAN     11

+#define OP_SSM_SCAN_SSD_CS  64 // Metal-specific; Chunk Size; 64 is largest multiple of 8 (simdgroup tile) fitting into 32 KiB Metal threadgroup mem limit (~26.75 KiB shared mem; see smem layout comment in kernel_ssm_scan_ssd_mma_f32)
+#define OP_SSM_SCAN_SSD_HD  64 // Metal-specific; Head Dim the MMA kernel is specialized for (Mamba-2); use_mma gates on d_inner == this
+#define OP_SSM_SCAN_SSD_NSG 4  // Metal-specific; Number of SimdGroups per threadgroup; NSG*32 == threads dispatched per threadgroup
+
 // kernel argument structs
 //
 // - element counters (e.g. ne00) typically use int32_t to reduce register usage
@@ -893,6 +897,8 @@ typedef struct {
    int64_t  n_head;
    int64_t  n_group;
    int64_t  n_seq_tokens;
+    int64_t  n_seq_tokens_total;
+    int64_t  token_offset;
    int64_t  n_seqs;
    int64_t  K;
    uint64_t s_off;
@@ -1677,6 +1677,7 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {

    ggml_metal_library_t lib = ctx->lib;
    ggml_metal_encoder_t enc = ctx->enc;
+    const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev);

    GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
    GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
@@ -1722,6 +1723,8 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
        /*.n_head       =*/ n_head,
        /*.n_group      =*/ n_group,
        /*.n_seq_tokens =*/ n_seq_tokens,
+        /*.n_seq_tokens_total =*/ n_seq_tokens,
+        /*.token_offset =*/ 0,
        /*.n_seqs       =*/ n_seqs,
        /*.K            =*/ K,
        /*.s_off        =*/ ggml_nelements(op->src[1]) * sizeof(float),
@@ -1751,26 +1754,53 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
        /*.nb0          =*/ nb0,
    };

-    auto pipeline = ggml_metal_library_get_pipeline_ssm_scan(lib, op);
+    constexpr int64_t CHUNK = OP_SSM_SCAN_SSD_CS;

-    GGML_ASSERT(d_state <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
+    const int64_t snap_reserve = K > 1 ? K : 0; // tokens reserved for sequential kernel rollback snapshots
+    const int64_t mma_tokens = ((n_seq_tokens - snap_reserve) / CHUNK) * CHUNK; // largest multiple of CHUNK that leaves snap_reserve for the tail
+    const bool use_mma =
+        mma_tokens > 0 &&
+        ne30 == 1 && // checks that A tensor is set to scalar decay per head (A shape {1, n_head})
+        props_dev->has_simdgroup_mm && // hardware check for M1 or newer
+        d_state % 8 == 0 && // d_state must be multiple of 8 to align with simdgroup_float 8x8 tiles
+        d_inner == OP_SSM_SCAN_SSD_HD; // mma kernel is specialized for the Mamba-2 head dim; this checks it

-    const size_t smem = pipeline.smem;
+    const auto dispatch = [&](ggml_metal_pipeline_with_params pipeline, int64_t nth, int64_t n_tg_x) {
+        GGML_ASSERT(nth <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
+        GGML_ASSERT(pipeline.smem <= props_dev->max_theadgroup_memory_size);

-    ggml_metal_encoder_set_pipeline(enc, pipeline);
-    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[1]), 2);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[2]), 3);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[3]), 4);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[4]), 5);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[5]), 6);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[6]), 7);
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         8);
+        ggml_metal_encoder_set_pipeline(enc, pipeline);
+        ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[1]), 2);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[2]), 3);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[3]), 4);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[4]), 5);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[5]), 6);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[6]), 7);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         8);
+        ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0);
+        ggml_metal_encoder_dispatch_threadgroups(enc, n_tg_x, n_head, n_seqs, nth, 1, 1);
+    };

-    ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
+    if (!use_mma) {
+        dispatch(ggml_metal_library_get_pipeline_ssm_scan(lib, op, false), d_state, d_inner);
+        return 1;
+    }

-    ggml_metal_encoder_dispatch_threadgroups(enc, d_inner, n_head, n_seqs, d_state, 1, 1);
+    args.n_seq_tokens = mma_tokens;
+    dispatch(
+        ggml_metal_library_get_pipeline_ssm_scan_ssd_mma(lib, op),
+        OP_SSM_SCAN_SSD_NSG*32,
+        1);
+
+    if (mma_tokens < n_seq_tokens) {
+        ggml_metal_op_concurrency_reset(ctx);
+
+        args.n_seq_tokens = n_seq_tokens - mma_tokens;
+        args.token_offset = mma_tokens;
+        dispatch(ggml_metal_library_get_pipeline_ssm_scan(lib, op, true), d_state, d_inner);
+    }

    return 1;
 }
@@ -159,7 +159,9 @@ kernel void kernel_ssm_conv_f32_f32_batched_4(

 // ref: ggml.c:ggml_compute_forward_ssm_scan_f32, Mamba-2 part
 // Optimized version: reduces redundant memory loads by having one thread load shared values
-kernel void kernel_ssm_scan_f32(
+// TAIL == false is the whole-sequence / decode path: token_offset folds away at compile time.
+template<bool TAIL>
+kernel void kernel_ssm_scan_impl(
        constant ggml_metal_kargs_ssm_scan & args,
        device const void * src0,
        device const void * src1,
@@ -200,13 +202,17 @@ kernel void kernel_ssm_scan_f32(
    const int32_t n_t = args.n_seq_tokens;
    const int32_t n_s = args.n_seqs;
    const int32_t K   = args.K;
+    const int32_t n_t_total = TAIL ? args.n_seq_tokens_total : n_t;
+    const int32_t t_off     = TAIL ? args.token_offset       : 0;

    const int32_t s_off = args.s_off;

    device const int32_t * ids = (device const int32_t *) src6;

-    device const float * s0_buff = (device const float *) ((device const char *) src0 + ir*args.nb02 + ids[i3]*args.nb03);
    device       float * s_buff  = (device       float *) ((device       char *) dst  + ir*args.nb02 +      i3*args.nb03 + s_off);
+    device const float * s0_buff = t_off != 0 ?
+        s_buff :
+        (device const float *) ((device const char *) src0 + ir*args.nb02 + ids[i3]*args.nb03);

    const int32_t i = i0 + i1*nc;
    const int32_t g = ir / (nh / ng); // repeat_interleave
@@ -218,12 +224,12 @@ kernel void kernel_ssm_scan_f32(

    const float A0 = A[i0%args.ne30];

-    device const float * x  = (device const float *)((device const char *) src1 + i1*args.nb10  + ir*args.nb11 + i3*args.nb13); // {dim, nh, nt, ns}
-    device const float * dt = (device const float *)((device const char *) src2 + ir*args.nb20  + i3*args.nb22);                // {nh, nt, ns}
-    device const float * B  = (device const float *)((device const char *) src4 +  g*args.nb41  + i3*args.nb43);                // {d_state, ng, nt, ns}
-    device const float * C  = (device const float *)((device const char *) src5 +  g*args.nb51  + i3*args.nb53);                // {d_state, ng, nt, ns}
+    device const float * x  = (device const float *)((device const char *) src1 + i1*args.nb10 + ir*args.nb11 + t_off*args.nb12 + i3*args.nb13); // {dim, nh, nt, ns}
+    device const float * dt = (device const float *)((device const char *) src2 + ir*args.nb20 + t_off*args.nb21 + i3*args.nb22);                 // {nh, nt, ns}
+    device const float * B  = (device const float *)((device const char *) src4 + g*args.nb41 + t_off*args.nb42 + i3*args.nb43);                  // {d_state, ng, nt, ns}
+    device const float * C  = (device const float *)((device const char *) src5 + g*args.nb51 + t_off*args.nb52 + i3*args.nb53);                  // {d_state, ng, nt, ns}

-    device float * y = dst + (i1 + ir*(nr) + i3*(n_t*nh*nr)); // {dim, nh, nt, ns}
+    device float * y = dst + (i1 + ir*nr + t_off*nh*nr + i3*(n_t_total*nh*nr)); // {dim, nh, nt, ns}

    for (int i2 = 0; i2 < n_t; i2 += sgptg) {
        threadgroup_barrier(mem_flags::mem_threadgroup);
@@ -285,3 +291,183 @@ kernel void kernel_ssm_scan_f32(

    s_buff[i] = s;
 }
+
+typedef decltype(kernel_ssm_scan_impl<false>) kernel_ssm_scan_t;
+
+template [[host_name("kernel_ssm_scan_f32")]]      kernel kernel_ssm_scan_t kernel_ssm_scan_impl<false>;
+template [[host_name("kernel_ssm_scan_f32_tail")]] kernel kernel_ssm_scan_t kernel_ssm_scan_impl<true>;
+
+// Chunked SSD SSM scan via Metal simdgroup MMatrix Multiply-Accumulate (simdgroup_float8x8) fast path.
+// One threadgroup per (head, sequence) and tokens are processed in chunks.
+// C*B^T computed in each chunk one time and reused across the head_dim channel tiles.
+kernel void kernel_ssm_scan_ssd_mma_f32(
+        constant ggml_metal_kargs_ssm_scan & args,
+        device const void * src0,
+        device const void * src1,
+        device const void * src2,
+        device const void * src3,
+        device const void * src4,
+        device const void * src5,
+        device const void * src6,
+        device      float * dst,
+        threadgroup float * shared [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort  tiitg[[thread_index_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]]) {
+    constexpr short CS  = OP_SSM_SCAN_SSD_CS;
+    constexpr short TC  = 8; // Tile Count of each edge in a simdgroup 8x8 tile
+    constexpr short HD  = OP_SSM_SCAN_SSD_HD;
+    constexpr short NSG = OP_SSM_SCAN_SSD_NSG;
+
+    // acs/exp(acs)/state-decay vectors, dtX[CS][HD], four private SAM row tiles [8][CS],
+    // and two 8x8 scratch tiles per simdgroup. Total: 26.75 KiB.
+    threadgroup float * shared_acs         = shared;
+    threadgroup float * shared_exp_acs     = shared + CS;
+    threadgroup float * shared_state_decay = shared + 2*CS;
+    threadgroup float * shared_dtx         = shared + 3*CS;
+    threadgroup float * shared_sam         = shared + 3*CS + CS*HD;
+    threadgroup float * sam_rows    = shared_sam + sgitg*TC*CS;
+    threadgroup float * shared_tile = shared_sam + NSG*TC*CS;
+    threadgroup float * tile0       = shared_tile + sgitg*2*TC*TC;
+    threadgroup float * tile1       = tile0 + TC*TC;
+
+    const int32_t ir = tgpig.y; // current head
+    const int32_t i3 = tgpig.z; // current seq
+
+    const int32_t nc  = args.d_state;
+    const int32_t nr  = args.d_inner;
+    const int32_t nh  = args.n_head;
+    const int32_t ng  = args.n_group;
+    const int32_t n_t = args.n_seq_tokens;
+    const int32_t n_t_total = args.n_seq_tokens_total;
+    const int32_t g   = ir / (nh / ng);
+
+    device const int32_t * ids = (device const int32_t *) src6;
+
+    device const float * s0_buff = (device const float *) ((device const char *) src0 + ir*args.nb02 + ids[i3]*args.nb03);
+    device       float * s_buff  = (device       float *) ((device       char *) dst  + ir*args.nb02 +      i3*args.nb03 + args.s_off);
+
+    device const float * A  = (device const float *) ((device const char *) src3 + ir*args.nb31);
+    device const float * x  = (device const float *) ((device const char *) src1 + ir*args.nb11 + i3*args.nb13);
+    device const float * dt = (device const float *) ((device const char *) src2 + ir*args.nb20 + i3*args.nb22);
+    device const float * B  = (device const float *) ((device const char *) src4 +  g*args.nb41 + i3*args.nb43);
+    device const float * C  = (device const float *) ((device const char *) src5 +  g*args.nb51 + i3*args.nb53);
+
+    device float * y = dst + (ir*nr + i3*(n_t_total*nh*nr));
+
+    for (int32_t t0 = 0; t0 < n_t; t0 += CS) {
+        for (int32_t idx = tiitg; idx < CS*HD; idx += NSG*N_SIMDWIDTH) {
+            const int32_t t = idx / HD;
+            const int32_t c = idx % HD;
+            const float dt0  = dt[(t0 + t) * (int32_t) args.ns21];
+            const float dtsp = dt0 <= 20.0f ? log(1.0f + exp(dt0)) : dt0;
+            shared_dtx[idx] = x[(t0 + t) * (int32_t) args.ns12 + c] * dtsp;
+        }
+        if (tiitg < CS) {
+            const float dt0  = dt[(t0 + tiitg) * (int32_t) args.ns21];
+            const float dtsp = dt0 <= 20.0f ? log(1.0f + exp(dt0)) : dt0;
+            shared_acs[tiitg] = dtsp * A[0];
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiitg == 0) {
+            float acc = 0.0f;
+            for (short t = 0; t < CS; ++t) {
+                acc += shared_acs[t];
+                shared_acs[t] = acc;
+            }
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        if (tiitg < CS) {
+            shared_exp_acs[tiitg] = exp(shared_acs[tiitg]);
+            shared_state_decay[tiitg] = exp(shared_acs[CS - 1] - shared_acs[tiitg]);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        device const float * state = t0 == 0 ? s0_buff : s_buff;
+
+        // Build one 8x64 row tile of SAM per simdgroup, then reuse it across every channel tile.
+        for (short ib = sgitg; ib < CS/TC; ib += NSG) {
+            for (short jb = 0; jb <= ib; ++jb) {
+                simdgroup_float8x8 cb = make_filled_simdgroup_matrix<float, 8>(0.0f);
+
+                for (int32_t k0 = 0; k0 < nc; k0 += TC) {
+                    simdgroup_float8x8 mc;
+                    simdgroup_float8x8 mb;
+                    simdgroup_load(mc, C + (t0 + ib*TC)*(int32_t) args.ns52 + k0, args.ns52);
+                    simdgroup_load(mb, B + (t0 + jb*TC)*(int32_t) args.ns42 + k0, args.ns42, 0, true);
+                    simdgroup_multiply_accumulate(cb, mc, mb, cb);
+                }
+
+                threadgroup float * sam = sam_rows + jb*TC;
+                simdgroup_store(cb, sam, CS);
+                simdgroup_barrier(mem_flags::mem_threadgroup);
+                for (short e = tiisg; e < TC*TC; e += N_SIMDWIDTH) {
+                    const short ri = e / TC;
+                    const short rj = e % TC;
+                    const short i  = ib*TC + ri;
+                    const short j  = jb*TC + rj;
+                    sam[ri*CS + rj] = j <= i ?
+                        sam[ri*CS + rj] * exp(shared_acs[i] - shared_acs[j]) : 0.0f;
+                }
+                simdgroup_barrier(mem_flags::mem_threadgroup);
+            }
+
+            for (short ch = 0; ch < HD/TC; ++ch) {
+                simdgroup_float8x8 y_diag  = make_filled_simdgroup_matrix<float, 8>(0.0f);
+                simdgroup_float8x8 y_inter = make_filled_simdgroup_matrix<float, 8>(0.0f);
+
+                for (short jb = 0; jb <= ib; ++jb) {
+                    simdgroup_float8x8 sam;
+                    simdgroup_float8x8 mdtx;
+                    simdgroup_load(sam,  sam_rows + jb*TC,                   CS);
+                    simdgroup_load(mdtx, shared_dtx + jb*TC*HD + ch*TC,     HD);
+                    simdgroup_multiply_accumulate(y_diag, sam, mdtx, y_diag);
+                }
+
+                for (int32_t k0 = 0; k0 < nc; k0 += TC) {
+                    simdgroup_float8x8 mc;
+                    simdgroup_float8x8 ms;
+                    simdgroup_load(mc, C + (t0 + ib*TC)*(int32_t) args.ns52 + k0, args.ns52);
+                    simdgroup_load(ms, state + ch*TC*nc + k0, nc, 0, true);
+                    simdgroup_multiply_accumulate(y_inter, mc, ms, y_inter);
+                }
+
+                simdgroup_store(y_diag,  tile0, TC);
+                simdgroup_store(y_inter, tile1, TC);
+                simdgroup_barrier(mem_flags::mem_threadgroup);
+                for (short e = tiisg; e < TC*TC; e += N_SIMDWIDTH) {
+                    const short ri = e / TC;
+                    const short ci = e % TC;
+                    const int32_t token = t0 + ib*TC + ri;
+                    y[token*nh*nr + ch*TC + ci] =
+                        tile0[e] + shared_exp_acs[ib*TC + ri] * tile1[e];
+                }
+                simdgroup_barrier(mem_flags::mem_threadgroup);
+            }
+        }
+
+        // All simdgroups must finish reading s_buff before any thread overwrites it.
+        threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup);
+
+        // Keep the carried-state reduction in token order. Reassociating this particular product
+        // with MMA compounds rounding differences at every chunk boundary; CB, y_diag, and C*S
+        // remain on the matrix unit.
+        const float chunk_decay = exp(shared_acs[CS - 1]);
+        for (int32_t idx = tiitg; idx < nc*HD; idx += NSG*N_SIMDWIDTH) {
+            const int32_t ci = idx / nc;
+            const int32_t si = idx % nc;
+            float state_c = 0.0f;
+            for (short t = 0; t < CS; ++t) {
+                state_c += shared_state_decay[t] *
+                    B[(t0 + t)*(int32_t) args.ns42 + si] *
+                    shared_dtx[t*HD + ci];
+            }
+            s_buff[idx] = chunk_decay * state[idx] + state_c;
+        }
+
+        // All state tiles must be visible before the next chunk consumes s_buff as S_prev.
+        threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup);
+    }
+}
@@ -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;
 }

@@ -1042,6 +1042,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;
@@ -4169,10 +4171,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 +4261,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 +4301,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 +5180,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 +5766,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);
@@ -11577,6 +11589,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) {
@@ -13942,6 +13965,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 +15807,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);
@@ -18511,6 +18639,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;
@@ -19437,6 +19577,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) {
@@ -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);
+    }
+}
@@ -1029,6 +1029,8 @@ void process_shaders() {

    string_to_spv("argmax_f32", "argmax.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "int"}}));
    string_to_spv("sum_rows_f32", "sum_rows.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
+    string_to_spv("cross_entropy_loss_f32", "cross_entropy_loss.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
+    string_to_spv("cross_entropy_loss_back_f32", "cross_entropy_loss_back.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
    string_to_spv("fwht_f32", "fwht.comp", {});
    string_to_spv("fwht_shmem_f32", "fwht.comp", {{"FWHT_SHMEM", "1"}});
    string_to_spv("count_equal_i32", "count_equal.comp", merge_maps(base_dict, {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}}));
@@ -43,10 +43,10 @@
 #define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq'

 #define LLAMA_SESSION_MAGIC   LLAMA_FILE_MAGIC_GGSN
-#define LLAMA_SESSION_VERSION 9
+#define LLAMA_SESSION_VERSION 10

 #define LLAMA_STATE_SEQ_MAGIC   LLAMA_FILE_MAGIC_GGSQ
-#define LLAMA_STATE_SEQ_VERSION 2
+#define LLAMA_STATE_SEQ_VERSION 3

 #ifdef __cplusplus
 extern "C" {
@@ -1,49 +0,0 @@
-#!/bin/sh
-#
-
-# Basedir on device
-basedir=/data/local/tmp/llama.cpp
-
-branch=.
-[ "$B" != "" ] && branch=$B
-
-adbserial=
-[ "$S" != "" ] && adbserial="-s $S"
-
-adbhost=
-[ "$H" != "" ] && adbhost="-H $H"
-
-model="Llama-3.2-3B-Instruct-Q4_0.gguf"
-[ "$M" != "" ] && model="$M"
-
-device="HTP0"
-[ "$D" != "" ] && device="$D"
-
-verbose=
-[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" cli_opts="$cli_opts -v"
-
-profile=
-[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" cli_opts="$cli_opts -v"
-
-opmask=
-[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE"
-
-nhvx=
-[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX"
-
-ndev=
-[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV"
-
-hb=
-[ "$HB" != "" ] && hb="GGML_HEXAGON_HOSTBUF=$HB"
-
-set -x
-
-adb $adbserial $adbhost shell " \
-  cd $basedir;         \
-  LD_LIBRARY_PATH=$basedir/$branch/lib   \
-  ADSP_LIBRARY_PATH=$basedir/$branch/lib \
-    $ndev $nhvx $opmask $verbose $profile $hb ./$branch/bin/llama-bench --device $device --load-mode none -m $basedir/../gguf/$model \
-        --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 \
-        --ubatch-size 1024 -fa 1 -ngl 99 $cli_opts $@   \
-"
@@ -1,78 +0,0 @@
-#!/bin/sh
-#
-
-# Basedir on device
-basedir=/data/local/tmp/llama.cpp
-
-cli_opts=
-
-branch=.
-[ "$B" != "" ] && branch=$B
-
-adbserial=
-[ "$S" != "" ] && adbserial="-s $S"
-
-adbhost=
-[ "$H" != "" ] && adbhost="-H $H"
-
-model="Llama-3.2-3B-Instruct-Q4_0.gguf"
-[ "$M" != "" ] && model="$M"
-
-device="HTP0"
-[ "$D" != "" ] && device="$D"
-
-verbose=
-[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" cli_opts="$cli_opts -v"
-
-sched=
-[ "$SCHED" != "" ] && sched="GGML_SCHED_DEBUG=2" cli_opts="$cli_opts -v"
-
-profile=
-[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" cli_opts="$cli_opts -v"
-
-opmask=
-[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE"
-
-nhvx=
-[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX"
-
-hmx=
-[ "$HMX" != "" ] && hmx="GGML_HEXAGON_USE_HMX=$HMX"
-
-ndev=
-[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV"
-
-hb=
-[ "$HB" != "" ] && hb="GGML_HEXAGON_HOSTBUF=$HB"
-
-opbatch=
-[ "$OB" != "" ] && opbatch="GGML_HEXAGON_OPBATCH=$OB"
-
-opqueue=
-[ "$OQ" != "" ] && opqueue="GGML_HEXAGON_OPQUEUE=$OQ"
-
-opflt=
-[ "$OF" != "" ] && opflt="GGML_HEXAGON_OPFILTER=$OF"
-
-vmem=
-[ "$VM" != "" ] && opflt="GGML_HEXAGON_VMEM=$VM"
-
-mbuf=
-[ "$MB" != "" ] && opflt="GGML_HEXAGON_MBUF=$MB"
-vmem=
-[ "$VM" != "" ] && vmem="GGML_HEXAGON_VMEM=$VM"
-
-mbuf=
-[ "$MB" != "" ] && mbuf="GGML_HEXAGON_MBUF=$MB"
-set -x
-
-adb $adbserial $adbhost shell " \
-  cd $basedir; ulimit -c unlimited;        \
-    LD_LIBRARY_PATH=$basedir/$branch/lib   \
-    ADSP_LIBRARY_PATH=$basedir/$branch/lib \
-    $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $opflt $vmem $mbuf \
-      ./$branch/bin/llama-cli --load-mode none -m $basedir/../gguf/$model \
-         --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1           \
-         --ctx-size 8192 --ubatch-size 1024 -fa on                 \
-         -ngl 99 --device $device $cli_opts $@                     \
-"
@@ -1,86 +0,0 @@
-#!/bin/sh
-#
-
-# Basedir on device
-basedir=/data/local/tmp/llama.cpp
-
-cli_opts=
-
-branch=.
-[ "$B" != "" ] && branch=$B
-
-adbserial=
-[ "$S" != "" ] && adbserial="-s $S"
-
-adbhost=
-[ "$H" != "" ] && adbhost="-H $H"
-
-model="Llama-3.2-3B-Instruct-Q4_0.gguf"
-[ "$M" != "" ] && model="$M"
-
-device="HTP0"
-[ "$D" != "" ] && device="$D"
-
-verbose=
-[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" cli_opts="$cli_opts -v"
-
-sched=
-[ "$SCHED" != "" ] && sched="GGML_SCHED_DEBUG=2" cli_opts="$cli_opts -v"
-
-profile=
-[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" cli_opts="$cli_opts -v"
-
-opmask=
-[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE"
-
-nhvx=
-[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX"
-
-hmx=
-[ "$HMX" != "" ] && hmx="GGML_HEXAGON_USE_HMX=$HMX"
-
-ndev=
-[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV"
-
-hb=
-[ "$HB" != "" ] && hb="GGML_HEXAGON_HOSTBUF=$HB"
-
-opbatch=
-[ "$OB" != "" ] && opbatch="GGML_HEXAGON_OPBATCH=$OB"
-
-opqueue=
-[ "$OQ" != "" ] && opqueue="GGML_HEXAGON_OPQUEUE=$OQ"
-
-oppoll=
-[ "$OP" != "" ] && oppoll="GGML_HEXAGON_OPPOLL=$OP"
-
-opflt=
-[ "$OF" != "" ] && opflt="GGML_HEXAGON_OPFILTER=$OF"
-
-opfuse=
-[ "$OC" != "" ] && opfuse="GGML_HEXAGON_OPFUSION=$OC"
-
-vmem=
-[ "$VM" != "" ] && vmem="GGML_HEXAGON_VMEM=$VM"
-
-mbuf=
-[ "$MB" != "" ] && mbuf="GGML_HEXAGON_MBUF=$MB"
-
-mmsel=
-[ "$MM" != "" ] && mmsel="GGML_HEXAGON_MM_SELECT=$MM"
-
-fasel=
-[ "$FA" != "" ] && fasel="GGML_HEXAGON_FA_SELECT=$FA"
-
-set -x
-
-adb $adbserial $adbhost shell " \
-  cd $basedir; ulimit -c unlimited;        \
-    LD_LIBRARY_PATH=$basedir/$branch/lib   \
-    ADSP_LIBRARY_PATH=$basedir/$branch/lib \
-    $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $oppoll $opflt $opfuse $vmem $mbuf $mmsel $fasel \
-      ./$branch/bin/llama-completion --load-mode none -m $basedir/../gguf/$model \
-         --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1                  \
-         --ctx-size 8192 --ubatch-size 1024 -fa on                        \
-         -ngl 99 --device $device $cli_opts $@                            \
-"
@@ -1,71 +0,0 @@
-#!/bin/sh
-#
-
-# Basedir on device
-basedir=/data/local/tmp/llama.cpp
-
-cli_opts=
-
-branch=.
-[ "$B" != "" ] && branch=$B
-
-adbserial=
-[ "$S" != "" ] && adbserial="-s $S"
-
-adbhost=
-[ "$H" != "" ] && adbhost="-H $H"
-
-model="gemma-3-4b-it-Q4_0.gguf"
-[ "$M" != "" ] && model="$M"
-
-mmproj="mmproj-F16.gguf"
-[ "$MMPROJ" != "" ] && mmproj="$MMPROJ"
-
-image=
-[ "$IMG" != "" ] && image="$IMG"
-
-device="HTP0"
-[ "$D" != "" ] && device="$D"
-
-verbose=
-[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V"
-
-experimental="GGML_HEXAGON_EXPERIMENTAL=1"
-[ "$E" != "" ] && experimental="GGML_HEXAGON_EXPERIMENTAL=$E"
-
-sched=
-[ "$SCHED" != "" ] && sched="GGML_SCHED_DEBUG=2" cli_opts="$cli_opts -v"
-
-profile=
-[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF"
-
-opmask=
-[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE"
-
-nhvx=
-[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX"
-
-hmx=
-[ "$HMX" != "" ] && hmx="GGML_HEXAGON_USE_HMX=$HMX"
-
-ndev=
-[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV"
-
-# MTMD backend device for vision model (defaults to CPU if not set)
-mtmd_backend=
-[ "$MTMD_DEVICE" != "" ] && mtmd_backend="MTMD_BACKEND_DEVICE=$MTMD_DEVICE"
-
-set -x
-
-adb $adbserial $adbhost shell " \
-  cd $basedir; ulimit -c unlimited;        \
-    LD_LIBRARY_PATH=$basedir/$branch/lib   \
-    ADSP_LIBRARY_PATH=$basedir/$branch/lib \
-    $verbose $experimental $sched $opmask $profile $hmx $nhvx $ndev $mtmd_backend \
-      ./$branch/bin/llama-mtmd-cli --load-mode none -m $basedir/../gguf/$model      \
-         --mmproj $basedir/../gguf/$mmproj                                   \
-         --image $basedir/../gguf/$image                                     \
-         --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1                     \
-         --ctx-size 8192 --ubatch-size 1024 -fa on                           \
-         -ngl 99 --device $device -v $cli_opts $@                            \
-"
@@ -1,72 +0,0 @@
-#!/bin/sh
-#
-
-# Basedir on device
-basedir=/data/local/tmp/llama.cpp
-
-cli_opts=
-
-branch=.
-[ "$B" != "" ] && branch=$B
-
-adbserial=
-[ "$S" != "" ] && adbserial="-s $S"
-
-adbhost=
-[ "$H" != "" ] && adbhost="-H $H"
-
-device="HTP0"
-[ "$D" != "" ] && device="$D"
-
-verbose=
-[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V"
-
-sched=
-[ "$SCHED" != "" ] && sched="GGML_SCHED_DEBUG=2" cli_opts="$cli_opts -v"
-
-profile=
-[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF"
-
-opmask=
-[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE"
-
-nhvx=
-[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX"
-
-hmx=
-[ "$HMX" != "" ] && hmx="GGML_HEXAGON_USE_HMX=$HMX"
-
-ndev=
-[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV"
-
-hb=
-[ "$HB" != "" ] && hb="GGML_HEXAGON_HOSTBUF=$HB"
-
-opbatch=
-[ "$OB" != "" ] && opbatch="GGML_HEXAGON_OPBATCH=$OB"
-
-opqueue=
-[ "$OQ" != "" ] && opqueue="GGML_HEXAGON_OPQUEUE=$OQ"
-
-oppoll=
-[ "$OP" != "" ] && oppoll="GGML_HEXAGON_OPPOLL=$OP"
-
-opfuse=
-[ "$OC" != "" ] && opfuse="GGML_HEXAGON_OPFUSION=$OC"
-
-mmsel=
-[ "$MM" != "" ] && mmsel="GGML_HEXAGON_MM_SELECT=$MM"
-
-fasel=
-[ "$FA" != "" ] && fasel="GGML_HEXAGON_FA_SELECT=$FA"
-
-set -x
-
-tool=$1; shift
-
-adb $adbserial $adbhost shell " \
-  cd $basedir; ulimit -c unlimited;        \
-    LD_LIBRARY_PATH=$basedir/$branch/lib   \
-    ADSP_LIBRARY_PATH=$basedir/$branch/lib \
-    $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $oppoll $opfuse $mmsel $fasel ./$branch/bin/$tool $@ \
-"
@@ -0,0 +1,260 @@
+#!/usr/bin/env python3
+#
+# Build llama.cpp for Snapdragon (via Docker or natively) and push to device.
+#
+
+import sys
+import os
+import argparse
+import subprocess
+import platform
+import shutil
+import logging
+
+logger = logging.getLogger("build")
+
+
+def parse_target(target_str):
+    if not target_str:
+        return None, None
+    if target_str.startswith("adb") or target_str.startswith("android"):
+        parts = target_str.split(":", 1)
+        serial = parts[1] if len(parts) > 1 else None
+        return "android", serial
+    elif target_str.startswith("lnx") or target_str.startswith("linux") or target_str.startswith("ubuntu"):
+        parts = target_str.split(":", 1)
+        host = parts[1] if len(parts) > 1 else None
+        return "linux", host
+    elif target_str in ("wos", "windows"):
+        return "windows", None
+    else:
+        return None, None
+
+
+def get_uid_gid():
+    if platform.system() != "Windows":
+        return [f"{os.getuid()}:{os.getgid()}"]
+    return []
+
+
+def main():
+    logging.basicConfig(level=logging.INFO, format='%(message)s')
+    parser = argparse.ArgumentParser(
+        description="Build llama.cpp for Snapdragon using cross-compilation docker containers or natively."
+    )
+    parser.add_argument("--target", default="android", help="Compilation target and deployment definition (e.g. android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, windows/wos) (default: android)")
+    parser.add_argument("--build-dir", help="Build directory name (defaults to build-TARGET[-dbg], e.g. build-android)")
+    parser.add_argument("--install-dir", help="Install directory name (defaults to pkg-TARGET[-dbg], e.g. pkg-android)")
+    parser.add_argument("--jobs", "-j", type=int, help="Number of build jobs (defaults to CPU thread count)")
+    parser.add_argument("--no-docker", action="store_true", help="Build natively on the host instead of in a docker container")
+    parser.add_argument("--preset", help="Override the CMake preset to use")
+    parser.add_argument("--debug", action="store_true", help="Build in debug mode (uses -debug presets instead of -release)")
+
+    # Push options
+    parser.add_argument("--push", action="store_true", help="Push built package to the target device via ADB or SSH/SCP")
+    parser.add_argument("--target-dir", help="Target directory on the device (default: /data/local/tmp/llama.cpp for Android, ~/llama.cpp for Linux)")
+
+    # Toolchain options
+    parser.add_argument("--toolchain-version", default="v0.7", help="Docker toolchain image version/tag (default: v0.7)")
+    parser.add_argument("--toolchain-url", default="ghcr.io/snapdragon-toolchain", help="Docker toolchain registry URL/namespace (default: ghcr.io/snapdragon-toolchain)")
+
+    args = parser.parse_args()
+
+    target_type, target_val = parse_target(args.target)
+    if not target_type:
+        logger.error(f"Error: Invalid target format '{args.target}'. Must be android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, or windows/wos.")
+        sys.exit(1)
+
+    # Determine preset and check if it's debug
+    preset = args.preset
+    if preset:
+        is_debug = args.debug or ("debug" in preset.lower())
+    else:
+        is_debug = args.debug
+        config_type = "debug" if is_debug else "release"
+        if args.no_docker:
+            if target_type == "windows" or platform.system() == "Windows":
+                preset = f"arm64-windows-snapdragon-{config_type}"
+            elif target_type == "linux":
+                preset = f"arm64-linux-snapdragon-{config_type}"
+            else:
+                preset = f"arm64-android-snapdragon-{config_type}"
+        else:
+            preset = f"arm64-linux-snapdragon-{config_type}" if target_type == "linux" else f"arm64-android-snapdragon-{config_type}"
+
+    target_prefix = args.target.split(":", 1)[0]
+    suffix = "-dbg" if is_debug else ""
+
+    build_dir = args.build_dir
+    if not build_dir:
+        build_dir = f"build-{target_prefix}{suffix}"
+
+    install_dir = args.install_dir
+    if not install_dir:
+        install_dir = f"pkg-{target_prefix}{suffix}"
+
+    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+
+    # Ensure CMakeUserPresets.json is in the workspace root, update if docs version is newer
+    preset_src = os.path.join(repo_root, "docs", "backend", "snapdragon", "CMakeUserPresets.json")
+    preset_dst = os.path.join(repo_root, "CMakeUserPresets.json")
+    if os.path.exists(preset_src):
+        should_copy = False
+        if not os.path.exists(preset_dst):
+            should_copy = True
+        else:
+            # Check modification times
+            src_mtime = os.path.getmtime(preset_src)
+            dst_mtime = os.path.getmtime(preset_dst)
+            if src_mtime > dst_mtime:
+                preset_bak = preset_dst + ".bak"
+                logger.info(f"Docs CMakeUserPresets.json is newer. Backing up existing {preset_dst} to {preset_bak}")
+                shutil.copy2(preset_dst, preset_bak)
+                should_copy = True
+
+        if should_copy:
+            logger.info(f"Copying CMakeUserPresets.json from {preset_src} to {preset_dst}")
+            shutil.copy2(preset_src, preset_dst)
+    else:
+        logger.warning("Warning: CMakeUserPresets.json not found in docs/backend/snapdragon/.")
+
+    jobs = args.jobs if args.jobs else os.cpu_count() or 4
+
+    if target_type == "windows":
+        logger.info("Windows target selected. Forcing native compilation...")
+        args.no_docker = True
+        if platform.system() != "Windows":
+            logger.warning("Warning: Windows compilation is intended to run on Windows arm64 hosts.")
+
+    if args.no_docker:
+        # Native/local host build
+        logger.info("Running native/local CMake build...")
+        install_prefix = os.path.join(repo_root, install_dir, "llama.cpp")
+
+        # Configure
+        configure_cmd = ["cmake", f"--preset={preset}", "-B", build_dir]
+        logger.info(f"+ {' '.join(configure_cmd)}")
+        res = subprocess.run(configure_cmd, cwd=repo_root)
+        if res.returncode != 0:
+            logger.error("CMake configuration failed.")
+            sys.exit(res.returncode)
+
+        # Build
+        build_cmd = ["cmake", "--build", build_dir, "-j", str(jobs)]
+        logger.info(f"+ {' '.join(build_cmd)}")
+        res = subprocess.run(build_cmd, cwd=repo_root)
+        if res.returncode != 0:
+            logger.error("CMake build failed.")
+            sys.exit(res.returncode)
+
+        # Install
+        install_cmd = ["cmake", "--install", build_dir, "--prefix", install_prefix]
+        logger.info(f"+ {' '.join(install_cmd)}")
+        res = subprocess.run(install_cmd, cwd=repo_root)
+        if res.returncode != 0:
+            logger.error("CMake install failed.")
+            sys.exit(res.returncode)
+    else:
+        # Docker-based build
+        logger.info("Running Docker-based cross-compilation build...")
+        image_name = "arm64-linux" if target_type == "linux" else "arm64-android"
+        image = f"{args.toolchain_url}/{image_name}:{args.toolchain_version}"
+
+        install_prefix_container = f"/workspace/{install_dir}/llama.cpp"
+
+        build_sh_cmd = (
+            f"cmake --preset {preset} -B /workspace/{build_dir} && "
+            f"cmake --build /workspace/{build_dir} -j {jobs} && "
+            f"cmake --install /workspace/{build_dir} --prefix {install_prefix_container}"
+        )
+
+        docker_cmd = [
+            "docker", "run", "--rm",
+            "--volume", f"{repo_root}:/workspace",
+            "--workdir", "/workspace",
+            "--platform", "linux/amd64"
+        ]
+        uid_gid = get_uid_gid()
+        if uid_gid:
+            docker_cmd += ["-u", uid_gid[0]]
+
+        docker_cmd += [image, "bash", "-c", build_sh_cmd]
+
+        logger.info(f"+ {' '.join(docker_cmd)}")
+        res = subprocess.run(docker_cmd, cwd=repo_root)
+        if res.returncode != 0:
+            logger.error("Docker-based build failed.")
+            sys.exit(res.returncode)
+
+    logger.info("\nBuild and installation completed successfully!")
+
+    # Push/deploy if requested
+    if args.push:
+        src_path = os.path.join(repo_root, install_dir, "llama.cpp")
+        if not os.path.exists(src_path):
+            logger.error(f"Error: installation directory {src_path} does not exist. Cannot deploy.")
+            sys.exit(1)
+
+        # Resolve target directory on device
+        target_dir = args.target_dir
+        if not target_dir:
+            target_dir = "/data/local/tmp/llama.cpp" if target_type == "android" else "~/llama.cpp"
+        target_dir = target_dir.rstrip("/")
+
+        sub_items = [item for item in os.listdir(src_path) if not item.startswith(".")]
+
+        if target_type == "android":
+            logger.info("\nPushing built artifacts to Android device via ADB...")
+            adb_cmd = ["adb"]
+            if target_val: # serial
+                adb_cmd += ["-s", target_val]
+
+            # Clean stale package files on device
+            if sub_items:
+                clean_paths = " ".join(f"{target_dir}/{item}" for item in sub_items)
+                clean_cmd = adb_cmd + ["shell", f"rm -rf {clean_paths}"]
+                logger.info(f"+ {' '.join(clean_cmd)}")
+                subprocess.run(clean_cmd)
+
+            # Android destination directory is target_dir
+            push_cmd = adb_cmd + ["push", os.path.join(src_path, "."), target_dir]
+            logger.info(f"+ {' '.join(push_cmd)}")
+            res = subprocess.run(push_cmd)
+            if res.returncode != 0:
+                logger.error("ADB push failed.")
+                sys.exit(res.returncode)
+            logger.info("ADB push completed successfully!")
+
+        elif target_type == "linux":
+            ssh_host = target_val
+            if not ssh_host:
+                logger.error("Error: SSH host not specified in target (e.g. use linux:user@host, lnx:user@host, or ubuntu:user@host). Cannot deploy.")
+                sys.exit(1)
+            logger.info(f"\nDeploying built artifacts to Linux device {ssh_host} via SSH/SCP...")
+
+            # Clean stale package files on remote host
+            if sub_items:
+                clean_paths = " ".join(f"{target_dir}/{item}" for item in sub_items)
+                clean_cmd = ["ssh", ssh_host, f"rm -rf {clean_paths}"]
+                logger.info(f"+ {' '.join(clean_cmd)}")
+                subprocess.run(clean_cmd)
+
+            # Deploy to target_dir
+            deploy_cmd = ["scp", "-r", os.path.join(src_path, "."), f"{ssh_host}:{target_dir}"]
+            logger.info(f"+ {' '.join(deploy_cmd)}")
+            res = subprocess.run(deploy_cmd)
+            if res.returncode != 0:
+                logger.error("SSH/SCP deploy failed.")
+                sys.exit(res.returncode)
+            logger.info("SSH/SCP deploy completed successfully!")
+
+        elif target_type == "windows":
+            logger.info("\nPush for Windows on Snapdragon (windows) target is currently a stub.")
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        logger.info("\nInterrupted by user.")
+        sys.exit(130)
@@ -34,6 +34,26 @@ trace_pattern = re.compile(
    r"trace-evt\s+(?P<event>[A-Z_0-9\-]+):\s+thread\s+(?P<thread>\d+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)"
 )

+device_pattern = re.compile(r"\b(HTP\d+(?::\d+)?)\s+(?:profile-op|trace-evt)\b")
+
+
+def extract_device(line):
+    m = device_pattern.search(line)
+    if m:
+        return m.group(1)
+    return "HTP0"
+
+
+def device_matches(record_device, target_device):
+    targets = [t.strip() for t in target_device.split(',')]
+    for target in targets:
+        if record_device == target:
+            return True
+        if record_device.startswith(target + ":"):
+            return True
+    return False
+
+
 logger = logging.getLogger("ggml-hexagon-profile")


@@ -72,7 +92,7 @@ class CycleUnwrapper:
        return raw + self.high_part


-def parse_log(file_path, pmu_index=None):
+def parse_log(file_path, pmu_index=None, limit=None, device_filter=None, op_filter_re=None):
    try:
        if file_path != "-":
            f = open(file_path, 'r', encoding='utf-8', errors='ignore')
@@ -85,13 +105,22 @@ def parse_log(file_path, pmu_index=None):
    all_ops: List[Dict[str, Any]] = []
    all_traces: List[Dict[str, Any]] = []
    current_op: Optional[Dict[str, Any]] = None
+    ops_count_per_device = {}
+    if device_filter is not None:
+        for target in device_filter.split(','):
+            ops_count_per_device[target.strip()] = 0
+    limit_reached = False

-    timestamp_pattern = re.compile(r"^(?P<min>\d+)\.(?P<sec>\d+)\.(?P<ms>\d+)\.(?P<us>\d+)\s+[A-Z]\s+")
-    unwrapper = None
-    trace_unwrapper = None
+    timestamp_pattern = re.compile(r"(?P<min>\d+)\.(?P<sec>\d+)\.(?P<ms>\d+)\.(?P<us>\d+)\s+[A-Z]\s+")
+    unwrappers = {}
+    last_batch_start = {}
+    trace_unwrappers = {}

    for line in f:
-        ts_match = timestamp_pattern.match(line)
+        if "profile-op" not in line and "trace-evt" not in line:
+            continue
+
+        ts_match = timestamp_pattern.search(line)
        abs_usec = 0
        if ts_match:
            abs_usec = (
@@ -100,8 +129,11 @@ def parse_log(file_path, pmu_index=None):
                + int(ts_match.group('us'))
            )

-        if "|" in line and "profile-op" in line:
-            parts = [p.strip() for p in line.split("|")]
+        device = extract_device(line)
+
+        idx = line.find("profile-op")
+        if idx != -1 and "|" in line[idx:]:
+            parts = [p.strip() for p in line[idx:].split("|")]
            prefix = parts[0]
            prefix_match = re.search(r"profile-op\s+(?P<op_name>[A-Z_0-9+]+)", prefix)
            if not prefix_match:
@@ -145,7 +177,6 @@ def parse_log(file_path, pmu_index=None):
                except (ValueError, IndexError):
                    pmu_val = None

-            evt_val = None
            evt_val = None
            if types.startswith("evt-cnt "):
                try:
@@ -158,14 +189,18 @@ def parse_log(file_path, pmu_index=None):
            if op_name == "OPBATCH":
                if cycles_start_raw:
                    unwrapped_cycles_start = int(cycles_start_raw)
-                    unwrapper = CycleUnwrapper(unwrapped_cycles_start)
-                    trace_unwrapper = CycleUnwrapper(unwrapped_cycles_start)
+                    unwrappers[device] = CycleUnwrapper(unwrapped_cycles_start)
+                    last_batch_start[device] = unwrapped_cycles_start
+                    for k in list(trace_unwrappers.keys()):
+                        if k[0] == device:
+                            del trace_unwrappers[k]
            else:
-                if cycles_start_raw and unwrapper is not None:
-                    unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw))
+                if cycles_start_raw:
+                    device_unwrapper = unwrappers.get(device)
+                    if device_unwrapper is not None:
+                        unwrapped_cycles_start = device_unwrapper.unwrap(int(cycles_start_raw))

-            idx = line.find("profile-op ")
-            op_text = line[idx + 11:].strip() if idx != -1 else line.strip()
+            op_text = re.sub(r"^profile-op\s+", "", line[idx:]).strip() if idx != -1 else line.strip()

            current_op = {
                'name':         op_name,
@@ -180,24 +215,58 @@ def parse_log(file_path, pmu_index=None):
                'pmu_val':      pmu_val,
                'evt_val':      evt_val,
                'abs_usec':     abs_usec,
-                'trace_events': []
+                'trace_events': [],
+                'device':       device
            }
            all_ops.append(current_op)
+
+            # Check if matching early exit criteria
+            matched = False
+            matched_target = None
+            if device_filter is not None:
+                targets = [t.strip() for t in device_filter.split(',')]
+                for target in targets:
+                    if device == target or device.startswith(target + ":"):
+                        matched = True
+                        matched_target = target
+                        break
+            else:
+                matched = True
+                matched_target = device
+
+            if op_filter_re is not None and not op_filter_re.search(op_text):
+                matched = False
+
+            if matched:
+                if matched_target not in ops_count_per_device:
+                    ops_count_per_device[matched_target] = 0
+                ops_count_per_device[matched_target] += 1
+
+            if limit is not None and len(ops_count_per_device) > 0 and all(count >= limit for count in ops_count_per_device.values()):
+                limit_reached = True
+
+            if limit_reached and op_name == "OPBATCH":
+                break
            continue

        trace_match = trace_pattern.search(line)
        if trace_match:
+            thread = int(trace_match.group('thread'))
            raw_cyc = int(trace_match.group('cycles'))
            unwrapped_cyc = None
-            if trace_unwrapper is not None:
-                unwrapped_cyc = trace_unwrapper.unwrap(raw_cyc)
+            th_key = (device, thread)
+            if th_key not in trace_unwrappers:
+                batch_start = last_batch_start.get(device)
+                trace_unwrappers[th_key] = CycleUnwrapper(batch_start)
+            unwrapped_cyc = trace_unwrappers[th_key].unwrap(raw_cyc)
            all_traces.append({
-                'thread': int(trace_match.group('thread')),
+                'thread': thread,
                'event':  trace_match.group('event'),
                'info':   int(trace_match.group('info')),
                'cycles': raw_cyc,
                'unwrapped_cycles': unwrapped_cyc,
-                'state':  trace_match.group('state')
+                'state':  trace_match.group('state'),
+                'device': device
            })

    f.close()
@@ -207,39 +276,45 @@ def parse_log(file_path, pmu_index=None):
        op['start_cycles'] = op['unwrapped_cycles_start']
        op['end_cycles'] = op['start_cycles'] + op['cycles'] if op['start_cycles'] is not None else None

-    # Filter ops with valid start_cycles
-    valid_ops = [op for op in all_ops if op['start_cycles'] is not None and op['end_cycles'] is not None]
+    # Group ops by device
+    valid_ops_by_dev = defaultdict(list)
+    for op in all_ops:
+        if op['start_cycles'] is not None and op['end_cycles'] is not None:
+            valid_ops_by_dev[op['device']].append(op)

-    # Separate OPBATCH ops from other ops
-    opbatch_ops = [op for op in valid_ops if op['name'] == "OPBATCH"]
-    other_ops = [op for op in valid_ops if op['name'] != "OPBATCH"]
-
-    # Sort them by start_cycles to enable binary search
-    opbatch_ops.sort(key=lambda op: op['start_cycles'])
-    other_ops.sort(key=lambda op: op['start_cycles'])
-
-    opbatch_starts = [op['start_cycles'] for op in opbatch_ops]
-    other_starts = [op['start_cycles'] for op in other_ops]
-
-    # Map trace events to any operator whose cycles contain them
+    # Group trace events by device
+    traces_by_dev = defaultdict(list)
    for e in all_traces:
-        cyc = e['unwrapped_cycles']
-        if cyc is None:
-            continue
+        if e['unwrapped_cycles'] is not None:
+            traces_by_dev[e['device']].append(e)

-        # Map to OPBATCH
-        idx = bisect.bisect_right(opbatch_starts, cyc) - 1
-        if idx >= 0:
-            op = opbatch_ops[idx]
-            if op['start_cycles'] <= cyc <= op['end_cycles']:
-                op['trace_events'].append(e)
+    for device, dev_ops in valid_ops_by_dev.items():
+        opbatch_ops = [op for op in dev_ops if op['name'] == "OPBATCH"]
+        other_ops = [op for op in dev_ops if op['name'] != "OPBATCH"]

-        # Map to other ops
-        idx = bisect.bisect_right(other_starts, cyc) - 1
-        if idx >= 0:
-            op = other_ops[idx]
-            if op['start_cycles'] <= cyc <= op['end_cycles']:
-                op['trace_events'].append(e)
+        opbatch_ops.sort(key=lambda op: op['start_cycles'])
+        other_ops.sort(key=lambda op: op['start_cycles'])
+
+        opbatch_starts = [op['start_cycles'] for op in opbatch_ops]
+        other_starts = [op['start_cycles'] for op in other_ops]
+
+        dev_traces = traces_by_dev.get(device, [])
+        for e in dev_traces:
+            cyc = e['unwrapped_cycles']
+
+            # Map to OPBATCH
+            idx = bisect.bisect_right(opbatch_starts, cyc) - 1
+            if idx >= 0:
+                op = opbatch_ops[idx]
+                if op['start_cycles'] <= cyc <= op['end_cycles']:
+                    op['trace_events'].append(e)
+
+            # Map to other ops
+            idx = bisect.bisect_right(other_starts, cyc) - 1
+            if idx >= 0:
+                op = other_ops[idx]
+                if op['start_cycles'] <= cyc <= op['end_cycles']:
+                    op['trace_events'].append(e)

    return all_ops

@@ -563,6 +638,7 @@ def main():
    parser.add_argument("--timeline", type=str, nargs='?', const='summary', choices=["summary", "bubbles"],
                        help="Output ASCII art event summary or thread idle bubble analysis (default: summary)")
    parser.add_argument("--filter", type=str, help="Regex filter matching against the original profile-op line")
+    parser.add_argument("--device", type=str, help="Device to filter by (e.g. HTP0, HTP0:0) or 'split' to generate separate reports per device")

    group = parser.add_mutually_exclusive_group()
    group.add_argument("--head", type=int, help="Limit to first N ops")
@@ -586,29 +662,84 @@ def main():
                logger.warning(f"Invalid width format '{w}'")

    final_pmu_name = (args.pmu_name or f"#{args.pmu_index}") if args.pmu_index is not None else None
-    ops = parse_log(args.logfile, pmu_index=args.pmu_index)

+    op_filter_re = None
    if args.filter:
        try:
-            filter_re = re.compile(args.filter)
+            op_filter_re = re.compile(args.filter)
        except re.error as e:
            logger.error(f"Invalid regex filter: {e}")
            sys.exit(1)
-        ops = [op for op in ops if filter_re.search(op['op_text'])]

-    if args.head is not None:
-        ops = ops[:args.head]
-    elif args.tail is not None:
-        ops = ops[-args.tail:]
+    limit = args.head if args.head is not None else None
+    device_filter = args.device if (args.device and args.device != "split") else None
+    ops = parse_log(args.logfile, pmu_index=args.pmu_index, limit=limit, device_filter=device_filter, op_filter_re=op_filter_re)

-    if args.timeline:
-        for op in ops:
-            if args.timeline == "summary":
-                print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'])
-            elif args.timeline == "bubbles":
-                print_bubbles_timeline(op)
+    if args.device and args.device != "split":
+        ops = [op for op in ops if device_matches(op['device'], args.device)]
+
+    if args.device == "split":
+        unique_devices = sorted(list(set(op['device'] for op in ops)))
+        for dev in unique_devices:
+            dev_ops = [op for op in ops if device_matches(op['device'], dev)]
+
+            if args.filter:
+                try:
+                    filter_re = re.compile(args.filter)
+                except re.error as e:
+                    logger.error(f"Invalid regex filter: {e}")
+                    sys.exit(1)
+                dev_ops = [op for op in dev_ops if filter_re.search(op['op_text'])]
+
+            if args.head is not None:
+                dev_ops = dev_ops[:args.head]
+            elif args.tail is not None:
+                dev_ops = dev_ops[-args.tail:]
+
+            logger.info("\n=========================================")
+            logger.info(f" Device: {dev}")
+            logger.info("=========================================")
+
+            if args.timeline:
+                for op in dev_ops:
+                    if args.timeline == "summary":
+                        print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'])
+                    elif args.timeline == "bubbles":
+                        print_bubbles_timeline(op)
+            else:
+                generate_report(dev_ops, args.top, overrides, args.sort, pmu_name=final_pmu_name)
    else:
-        generate_report(ops, args.top, overrides, args.sort, pmu_name=final_pmu_name)
+        if args.filter:
+            try:
+                filter_re = re.compile(args.filter)
+            except re.error as e:
+                logger.error(f"Invalid regex filter: {e}")
+                sys.exit(1)
+            ops = [op for op in ops if filter_re.search(op['op_text'])]
+
+        if args.head is not None or args.tail is not None:
+            ops_by_dev = defaultdict(list)
+            for op in ops:
+                ops_by_dev[op['device']].append(op)
+
+            filtered_ops = []
+            for dev in sorted(ops_by_dev.keys()):
+                dev_ops = ops_by_dev[dev]
+                if args.head is not None:
+                    dev_ops = dev_ops[:args.head]
+                elif args.tail is not None:
+                    dev_ops = dev_ops[-args.tail:]
+                filtered_ops.extend(dev_ops)
+            ops = filtered_ops
+
+        if args.timeline:
+            for op in ops:
+                if args.timeline == "summary":
+                    print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'])
+                elif args.timeline == "bubbles":
+                    print_bubbles_timeline(op)
+        else:
+            generate_report(ops, args.top, overrides, args.sort, pmu_name=final_pmu_name)


 if __name__ == "__main__":
@@ -20,6 +20,31 @@ trace_pattern = re.compile(
    r"trace-evt\s+(?P<event>[A-Z_0-9\-]+):\s+thread\s+(?P<thread>\d+)\s+info\s+(?P<info>\d+)\s+(?P<state>start|stop)\s+(?P<cycles>\d+)"
 )

+device_pattern = re.compile(r"\b(HTP\d+(?::\d+)?)\s+(?:profile-op|trace-evt)\b")
+
+
+def extract_device(line):
+    m = device_pattern.search(line)
+    if m:
+        return m.group(1)
+    return "HTP0"
+
+
+def device_matches(record_device, target_device):
+    targets = [t.strip() for t in target_device.split(',')]
+    for target in targets:
+        if record_device == target:
+            return True
+        if record_device.startswith(target + ":"):
+            return True
+    return False
+
+
+def get_split_output_path(base_path, device_name):
+    safe_device = device_name.replace(':', '_')
+    root, ext = os.path.splitext(base_path)
+    return f"{root}-{safe_device}{ext}"
+

 def normalize_event_name(evt_type, info=0):
    if evt_type == "HVX_COMP":
@@ -54,7 +79,79 @@ class CycleUnwrapper:
        return raw + self.high_part


-def parse_log(file_path):
+class DeviceTimeMapper:
+    def __init__(self, dev, ops):
+        self.dev = dev
+        self.batches = []
+        for op in ops:
+            if op.get('device') == dev and op.get('name') == 'OPBATCH' and op.get('unwrapped_cycles_start') is not None:
+                cycles = op.get('cycles', 0)
+                usec = op.get('usec', 0)
+                start_cyc = op['unwrapped_cycles_start']
+                freq = (cycles / usec) if usec > 0 and cycles > 0 else 1000.0
+                if freq <= 0:
+                    freq = 1000.0
+                self.batches.append({
+                    'start_cycles': start_cyc,
+                    'cycles': cycles,
+                    'end_cycles': start_cyc + cycles,
+                    'usec': usec,
+                    'dur_ns': usec * 1000,
+                    'freq_mhz': freq,
+                })
+
+        self.batches.sort(key=lambda b: b['start_cycles'])
+
+        for i, b in enumerate(self.batches):
+            if i == 0:
+                b['start_time_ns'] = 0
+            else:
+                prev = self.batches[i - 1]
+                idle_cyc = max(0, b['start_cycles'] - prev['end_cycles'])
+                idle_ns = int(round((idle_cyc / prev['freq_mhz']) * 1000))
+                b['start_time_ns'] = prev['start_time_ns'] + prev['dur_ns'] + idle_ns
+
+        self.batch_starts = [b['start_cycles'] for b in self.batches]
+
+        valid_starts = [op['unwrapped_cycles_start'] for op in ops if op.get('device') == dev and op.get('unwrapped_cycles_start') is not None]
+        self.min_cyc = min(valid_starts) if valid_starts else 0
+        if self.batches:
+            self.default_freq = self.batches[0]['freq_mhz']
+        else:
+            freqs = [op['cycles'] / op['usec'] for op in ops if op.get('device') == dev and op.get('usec', 0) > 0 and op.get('cycles', 0) > 0]
+            self.default_freq = statistics.mean(freqs) if freqs else 1000.0
+
+    def get_batch(self, cyc):
+        if not self.batches:
+            return None
+        idx = bisect.bisect_right(self.batch_starts, cyc) - 1
+        if idx >= 0:
+            return self.batches[idx]
+        return self.batches[0]
+
+    def get_freq(self, cyc=None):
+        if cyc is not None:
+            b = self.get_batch(cyc)
+            if b is not None:
+                return b['freq_mhz']
+        return self.default_freq
+
+    def cycle_to_ns(self, cyc):
+        if cyc is None:
+            return 0
+        b = self.get_batch(cyc)
+        if b is not None:
+            return b['start_time_ns'] + int(round(((cyc - b['start_cycles']) / b['freq_mhz']) * 1000))
+        return int(round(((cyc - self.min_cyc) / self.default_freq) * 1000))
+
+    def dur_cycles_to_ns(self, cyc_start, cyc_dur):
+        if cyc_dur is None:
+            return 0
+        freq = self.get_freq(cyc_start)
+        return int(round((cyc_dur / freq) * 1000))
+
+
+def parse_log(file_path, limit=None, device_filter=None, op_filter_re=None):
    try:
        if file_path != "-":
            f = open(file_path, 'r', encoding='utf-8', errors='ignore')
@@ -67,14 +164,25 @@ def parse_log(file_path):
    all_ops: List[Dict[str, Any]] = []
    all_traces: List[Dict[str, Any]] = []
    current_op: Optional[Dict[str, Any]] = None
-    unwrapper = None
-    trace_unwrapper = None
+    ops_count_per_device = {}
+    if device_filter is not None:
+        for target in device_filter.split(','):
+            ops_count_per_device[target.strip()] = 0
+    limit_reached = False
+    unwrappers = {}
+    last_batch_start = {}
+    trace_unwrappers = {}
    line_idx = 0

    for line in f:
        line_idx += 1
-        if "|" in line and "profile-op" in line:
-            parts = [p.strip() for p in line.split("|")]
+        if "profile-op" not in line and "trace-evt" not in line:
+            continue
+        device = extract_device(line)
+
+        idx = line.find("profile-op")
+        if idx != -1 and "|" in line[idx:]:
+            parts = [p.strip() for p in line[idx:].split("|")]
            prefix = parts[0]
            prefix_match = re.search(r"profile-op\s+(?P<op_name>[A-Z_0-9+]+)", prefix)
            if not prefix_match:
@@ -115,14 +223,18 @@ def parse_log(file_path):
            if op_name == "OPBATCH":
                if cycles_start_raw:
                    unwrapped_cycles_start = int(cycles_start_raw)
-                    unwrapper = CycleUnwrapper(unwrapped_cycles_start)
-                    trace_unwrapper = CycleUnwrapper(unwrapped_cycles_start)
+                    unwrappers[device] = CycleUnwrapper(unwrapped_cycles_start)
+                    last_batch_start[device] = unwrapped_cycles_start
+                    for k in list(trace_unwrappers.keys()):
+                        if k[0] == device:
+                            del trace_unwrappers[k]
            else:
-                if cycles_start_raw and unwrapper is not None:
-                    unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw))
+                if cycles_start_raw:
+                    device_unwrapper = unwrappers.get(device)
+                    if device_unwrapper is not None:
+                        unwrapped_cycles_start = device_unwrapper.unwrap(int(cycles_start_raw))

-            idx = line.find("profile-op ")
-            op_text = line[idx + 11:].strip() if idx != -1 else line.strip()
+            op_text = re.sub(r"^profile-op\s+", "", line[idx:]).strip() if idx != -1 else line.strip()

            evt_str = None
            if types.startswith("evt-cnt "):
@@ -142,24 +254,59 @@ def parse_log(file_path):
                'cycles_start': int(cycles_start_raw) if cycles_start_raw else None,
                'unwrapped_cycles_start': unwrapped_cycles_start,
                'trace_events': [],
-                'line_num':     line_idx
+                'line_num':     line_idx,
+                'device':       device
            }
            all_ops.append(current_op)
+
+            # Check if matching early exit criteria
+            matched = False
+            matched_target = None
+            if device_filter is not None:
+                targets = [t.strip() for t in device_filter.split(',')]
+                for target in targets:
+                    if device == target or device.startswith(target + ":"):
+                        matched = True
+                        matched_target = target
+                        break
+            else:
+                matched = True
+                matched_target = device
+
+            if op_filter_re is not None and not op_filter_re.search(op_text):
+                matched = False
+
+            if matched:
+                if matched_target not in ops_count_per_device:
+                    ops_count_per_device[matched_target] = 0
+                ops_count_per_device[matched_target] += 1
+
+            if limit is not None and len(ops_count_per_device) > 0 and all(count >= limit for count in ops_count_per_device.values()):
+                limit_reached = True
+
+            if limit_reached and op_name == "OPBATCH":
+                break
            continue

        trace_match = trace_pattern.search(line)
        if trace_match:
+            thread = int(trace_match.group('thread'))
            raw_cyc = int(trace_match.group('cycles'))
            unwrapped_cyc = None
-            if trace_unwrapper is not None:
-                unwrapped_cyc = trace_unwrapper.unwrap(raw_cyc)
+            th_key = (device, thread)
+            if th_key not in trace_unwrappers:
+                batch_start = last_batch_start.get(device)
+                trace_unwrappers[th_key] = CycleUnwrapper(batch_start)
+            unwrapped_cyc = trace_unwrappers[th_key].unwrap(raw_cyc)
            all_traces.append({
-                'thread': int(trace_match.group('thread')),
+                'thread': thread,
                'event':  trace_match.group('event'),
                'info':   int(trace_match.group('info')),
                'cycles': raw_cyc,
                'unwrapped_cycles': unwrapped_cyc,
-                'state':  trace_match.group('state')
+                'state':  trace_match.group('state'),
+                'line_num': line_idx,
+                'device': device
            })

    f.close()
@@ -274,27 +421,24 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path):
        logger.warning("No operators found after filtering.")
        return

-    # Compute average frequency
-    frequencies = []
-    for op in filtered_ops:
-        if op['usec'] > 0 and op['cycles'] > 0:
-            frequencies.append(op['cycles'] / op['usec'])
-    avg_freq_mhz = statistics.mean(frequencies) if frequencies else 1000.0
-    if avg_freq_mhz <= 0:
-        avg_freq_mhz = 1000.0
-
    # Assign start and end cycles to each operator
    for op in filtered_ops:
        op['start_cycles'] = op['unwrapped_cycles_start']
-        op['end_cycles'] = op['start_cycles'] + op['cycles']
+        op['end_cycles'] = op['start_cycles'] + op['cycles'] if op['start_cycles'] is not None else None

-    global_min_cyc = min(op['start_cycles'] for op in filtered_ops if op['start_cycles'] is not None)
+    # Get list of unique devices present in the operations
+    unique_devices = sorted(list(set(op['device'] for op in filtered_ops)))
+    device_to_idx = {dev: idx for idx, dev in enumerate(unique_devices)}
+    time_mappers = {dev: DeviceTimeMapper(dev, filtered_ops) for dev in unique_devices}

    # Process events
    completed_events = []
    if trace_events:
        trace_events = sorted(trace_events, key=lambda e: e['unwrapped_cycles'])
-        one_usec_cycles = max(avg_freq_mhz, 1.0)
+
+        one_usec_cycles = {}
+        for dev in unique_devices:
+            one_usec_cycles[dev] = max(time_mappers[dev].get_freq(), 1.0)

        active_starts = {}
        for e in trace_events:
@@ -303,31 +447,36 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path):
            info = e['info']
            state = e['state']
            cyc = e['unwrapped_cycles']
+            dev = e['device']

-            key = (t, evt, info)
+            key = (dev, t, evt, info)
            if state == 'start':
                # Handle missing stop (start followed by another start)
                if key in active_starts:
-                    prev_start = active_starts[key]
+                    prev_e = active_starts[key]
                    completed_events.append({
                        'thread': t,
                        'event': evt,
                        'info': info,
-                        'start_cyc': prev_start,
-                        'end_cyc': prev_start + one_usec_cycles,
+                        'start_cyc': prev_e['unwrapped_cycles'],
+                        'end_cyc': prev_e['unwrapped_cycles'] + one_usec_cycles.get(dev, 1000.0),
+                        'line_num': prev_e.get('line_num'),
                        'missing_stop': True,
+                        'device': dev
                    })
-                active_starts[key] = cyc
+                active_starts[key] = e
            elif state == 'stop':
                if key in active_starts:
-                    start_cyc = active_starts[key]
+                    prev_e = active_starts[key]
                    del active_starts[key]
                    completed_events.append({
                        'thread': t,
                        'event': evt,
                        'info': info,
-                        'start_cyc': start_cyc,
+                        'start_cyc': prev_e['unwrapped_cycles'],
                        'end_cyc': cyc,
+                        'line_num': prev_e.get('line_num'),
+                        'device': dev
                    })
                else:
                    # Handle missing start (stop without start)
@@ -335,31 +484,36 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path):
                        'thread': t,
                        'event': evt,
                        'info': info,
-                        'start_cyc': cyc - one_usec_cycles,
+                        'start_cyc': cyc - one_usec_cycles.get(dev, 1000.0),
                        'end_cyc': cyc,
+                        'line_num': e.get('line_num'),
                        'missing_start': True,
+                        'device': dev
                    })

        # Clear remaining unmatched starts
-        for key, start_cyc in active_starts.items():
-            t, evt, info = key
+        for key, prev_e in active_starts.items():
+            dev, t, evt, info = key
            completed_events.append({
                'thread': t,
                'event': evt,
                'info': info,
-                'start_cyc': start_cyc,
-                'end_cyc': start_cyc + one_usec_cycles,
+                'start_cyc': prev_e['unwrapped_cycles'],
+                'end_cyc': prev_e['unwrapped_cycles'] + one_usec_cycles.get(dev, 1000.0),
+                'line_num': prev_e.get('line_num'),
                'missing_stop': True,
+                'device': dev
            })

    completed_events.sort(key=lambda e: e['start_cyc'])

-    # Convert event times to microseconds and apply clamp rounded to 1ns resolution (3 decimals)
+    # Convert event times to nanoseconds using per-device / per-batch time mapper
    for e in completed_events:
-        start_us = (e['start_cyc'] - global_min_cyc) / avg_freq_mhz
-        dur_us = (e['end_cyc'] - e['start_cyc']) / avg_freq_mhz
-        e['ts_ns'] = int(round(start_us * 1000))
-        e['dur_ns'] = int(round(max(dur_us, 0.1) * 1000))
+        dev = e['device']
+        tm = time_mappers[dev]
+        e['ts_ns'] = tm.cycle_to_ns(e['start_cyc'])
+        dur_ns = tm.dur_cycles_to_ns(e['start_cyc'], e['end_cyc'] - e['start_cyc'])
+        e['dur_ns'] = max(dur_ns, 100)

    # Allocate slots (sub-tracks) to prevent overlaps on same virtual track
    active_slots = defaultdict(list)
@@ -368,14 +522,15 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path):
        evt = e['event']
        ts = e['ts_ns']
        dur = e['dur_ns']
+        dev = e['device']

        norm_evt = normalize_event_name(evt, e['info'])
        if norm_evt == "DMA":
-            track_key = (t, "DMA")
+            track_key = (dev, t, "DMA")
        elif t == 10:
-            track_key = (t, "HMX")
+            track_key = (dev, t, "HMX")
        else:
-            track_key = (t, "HVX")
+            track_key = (dev, t, "HVX")

        slots = active_slots[track_key]
        allocated_slot = -1
@@ -395,6 +550,7 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path):
        t = e['thread']
        evt = e['event']
        slot = e['slot']
+        dev = e['device']

        norm_evt = normalize_event_name(evt, e['info'])
        if norm_evt == "DMA":
@@ -408,56 +564,69 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path):
            evt_id = 2

        t_sort = 1 if t == 10 else t + 2
+        dev_idx = device_to_idx[dev]
+
        # Unique UUID for each sub-track
        if t == 10:
-            uuid = 20  # HMX thread track UUID
+            uuid = dev_idx * 10000000 + 20  # HMX thread track UUID
        else:
-            uuid = int(t_sort * 1000000 + evt_id * 1000 + slot)
+            uuid = int(dev_idx * 10000000 + t_sort * 1000000 + evt_id * 1000 + slot)
        e['uuid'] = uuid
-        used_tracks[uuid] = (t, track_evt, slot)
+        used_tracks[uuid] = (dev, t, track_evt, slot)

    with open(output_path, "wb") as f:
-        # Define Process with EXPLICIT child sorting
-        proc_desc = make_process_descriptor(1, "HTP NPU")
-        proc_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(1, process=proc_desc, child_ordering=3))
-        write_trace_packet_to_file(f, proc_packet)
+        for dev in unique_devices:
+            dev_idx = device_to_idx[dev]
+            pid = dev_idx + 1
+            proc_uuid = dev_idx * 10000000 + 1

-        # Define Operators Track (UUID = 2) as a thread track at rank 1, tid 8
-        op_thread_desc = make_thread_descriptor(1, 8, "Ops", sort_index=1)
-        op_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(2, parent_uuid=1, thread=op_thread_desc))
-        write_trace_packet_to_file(f, op_packet)
+            # Define Process with EXPLICIT child sorting
+            proc_name = dev
+            proc_desc = make_process_descriptor(pid, proc_name)
+            proc_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(proc_uuid, process=proc_desc, child_ordering=3))
+            write_trace_packet_to_file(f, proc_packet)

-        # Define HMX Thread Track (UUID = 20) at rank 2, tid 9
-        hmx_thread_desc = make_thread_descriptor(1, 9, "HMX", sort_index=2)
-        hmx_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(20, parent_uuid=1, thread=hmx_thread_desc))
-        write_trace_packet_to_file(f, hmx_packet)
+            # Define Operators Track as a thread track
+            op_track_uuid = dev_idx * 10000000 + 2
+            op_tid = pid * 100 + 8
+            op_thread_desc = make_thread_descriptor(pid, op_tid, "Ops", sort_index=1)
+            op_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(op_track_uuid, parent_uuid=proc_uuid, thread=op_thread_desc))
+            write_trace_packet_to_file(f, op_packet)

-        # Define Thread Tracks (T0, T1, ..., T9)
-        unique_threads = sorted(list(set(t for (t, _, _) in used_tracks.values() if t != 10)))
-        for t in unique_threads:
-            thread_uuid = 10 + t
-            thread_name = f"T{t}"
-            # Sort order starts from index 3 (T0 -> 3, T1 -> 4, etc.)
-            sort_index = 3 + t
-            tid = 10 + t
-            thread_desc = make_thread_descriptor(1, tid, thread_name, sort_index=sort_index)
-            thread_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(
-                thread_uuid,
-                parent_uuid=1,
-                thread=thread_desc,
-                sibling_order_rank=sort_index,
-                child_ordering=3  # Explicit child sorting for sub-tracks
-            ))
-            write_trace_packet_to_file(f, thread_packet)
+            # Define HMX Thread Track at rank 2
+            hmx_track_uuid = dev_idx * 10000000 + 20
+            hmx_tid = pid * 100 + 9
+            hmx_thread_desc = make_thread_descriptor(pid, hmx_tid, "HMX", sort_index=2)
+            hmx_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(hmx_track_uuid, parent_uuid=proc_uuid, thread=hmx_thread_desc))
+            write_trace_packet_to_file(f, hmx_packet)
+
+            # Define Thread Tracks (T0, T1, ..., T9) for this device
+            dev_used_tracks = {uuid: val for uuid, val in used_tracks.items() if val[0] == dev}
+            unique_threads = sorted(list(set(t for (_, t, _, _) in dev_used_tracks.values() if t != 10)))
+            for t in unique_threads:
+                thread_uuid = dev_idx * 10000000 + 10 + t
+                thread_name = f"T{t}"
+                sort_index = 3 + t
+                tid = pid * 100 + 10 + t
+                thread_desc = make_thread_descriptor(pid, tid, thread_name, sort_index=sort_index)
+                thread_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(
+                    thread_uuid,
+                    parent_uuid=proc_uuid,
+                    thread=thread_desc,
+                    sibling_order_rank=sort_index,
+                    child_ordering=3  # Explicit child sorting for sub-tracks
+                ))
+                write_trace_packet_to_file(f, thread_packet)

        # Define Track descriptors for sub-tracks parented to thread tracks
        for uuid in sorted(used_tracks.keys()):
-            if uuid == 20:
+            dev, t, evt, slot = used_tracks[uuid]
+            dev_idx = device_to_idx[dev]
+            if t == 10:
                continue
-            t, evt, slot = used_tracks[uuid]
            name = f"T{t} {evt}"
            rank = 0 if evt == "HVX" else 1
-            parent_thread_uuid = 10 + t
+            parent_thread_uuid = dev_idx * 10000000 + 10 + t
            # Sibling merge behavior: 1 (SIBLING_MERGE_BEHAVIOR_BY_TRACK_NAME)
            track_desc = make_track_descriptor(
                uuid=uuid,
@@ -470,15 +639,18 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path):
            write_trace_packet_to_file(f, track_packet)

        # Emit Operators
-        last_op_end_ns = 0
+        last_op_end_ns = defaultdict(int)
        for op in filtered_ops:
-            op_start_ns = int(round(((op['start_cycles'] - global_min_cyc) / avg_freq_mhz) * 1000))
-            op_dur_ns = int(round((op['cycles'] / avg_freq_mhz) * 1000))
+            dev = op['device']
+            dev_idx = device_to_idx[dev]
+            tm = time_mappers[dev]
+            op_start_ns = tm.cycle_to_ns(op['start_cycles'])
+            op_dur_ns = tm.dur_cycles_to_ns(op['start_cycles'], op['cycles'])
            if op['name'] != "OPBATCH":
-                if op_start_ns < last_op_end_ns:
-                    op_start_ns = last_op_end_ns
+                if op_start_ns < last_op_end_ns[dev]:
+                    op_start_ns = last_op_end_ns[dev]
                clamped_dur = max(op_dur_ns, 100) # Clamp to 100ns (0.1us)
-                last_op_end_ns = op_start_ns + clamped_dur
+                last_op_end_ns[dev] = op_start_ns + clamped_dur
            else:
                clamped_dur = max(op_dur_ns, 100)

@@ -495,24 +667,41 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path):
            if 'evt' in op and op['evt']:
                debug_annots.append(make_debug_annotation("evt", string_val=op['evt']))

+            op_track_uuid = dev_idx * 10000000 + 2
+
            # Slice Begin
-            evt_begin = make_track_event(1, 2, name=f"{op['name']} ({op['dims']})", category="operator", debug_annotations=debug_annots)
+            evt_begin = make_track_event(1, op_track_uuid, name=f"{op['name']} ({op['dims']})", category="operator", debug_annotations=debug_annots)
            packet_begin = make_trace_packet(op_start_ns, track_event=evt_begin)
            write_trace_packet_to_file(f, packet_begin)

            # Slice End
-            evt_end = make_track_event(2, 2)
+            evt_end = make_track_event(2, op_track_uuid)
            packet_end = make_trace_packet(op_start_ns + clamped_dur, track_event=evt_end)
            write_trace_packet_to_file(f, packet_end)

        # Emit Thread Trace Events
        for e in completed_events:
            norm_name = normalize_event_name(e['event'], e['info'])
-            name = f"DMA {e['info']}" if norm_name == "DMA" else norm_name
+            if norm_name == "DMA":
+                name = f"DMA {e['info']}"
+            elif norm_name == "FENCE":
+                name = f"FENCE {e['info']}" if e.get('info') is not None and e['info'] != 0 else "FENCE"
+            else:
+                name = norm_name
+
            if e.get('missing_start') or e.get('missing_stop'):
                name += "!"

            debug_annots = []
+            if 'line_num' in e and e['line_num'] is not None:
+                debug_annots.append(make_debug_annotation("line", int_val=e['line_num']))
+            if norm_name == "FENCE" and e.get('info') is not None:
+                debug_annots.append(make_debug_annotation("seq", int_val=e['info']))
+            elif norm_name == "DMA" and e.get('info') is not None:
+                debug_annots.append(make_debug_annotation("channel", int_val=e['info']))
+            elif e.get('info') is not None and e['info'] != 0:
+                debug_annots.append(make_debug_annotation("info", int_val=e['info']))
+
            if e.get('missing_start'):
                debug_annots.append(make_debug_annotation("missing_start", string_val="true"))
            if e.get('missing_stop'):
@@ -536,6 +725,7 @@ def main():
    parser.add_argument("logfile", help="Path to hex-log profile file")
    parser.add_argument("-o", "--output", default="optrace.perfetto-trace", help="Output trace file path (default: optrace.perfetto-trace)")
    parser.add_argument("--filter", type=str, help="Regex filter matching against the original profile-op line")
+    parser.add_argument("--device", type=str, help="Device to filter by (e.g. HTP0, HTP0:0) or 'split' to generate separate files per device")

    group = parser.add_mutually_exclusive_group()
    group.add_argument("--head", type=int, help="Limit to first N ops")
@@ -544,7 +734,21 @@ def main():
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO, format='%(message)s')

-    ops, traces = parse_log(args.logfile)
+    op_filter_re = None
+    if args.filter:
+        try:
+            op_filter_re = re.compile(args.filter)
+        except re.error as e:
+            logger.error(f"Invalid regex filter: {e}")
+            sys.exit(1)
+
+    limit = args.head if args.head is not None else None
+    device_filter = args.device if (args.device and args.device != "split") else None
+    ops, traces = parse_log(args.logfile, limit=limit, device_filter=device_filter, op_filter_re=op_filter_re)
+
+    if args.device and args.device != "split":
+        ops = [op for op in ops if device_matches(op['device'], args.device)]
+        traces = [t for t in traces if device_matches(t['device'], args.device)]

    if args.filter:
        try:
@@ -554,35 +758,60 @@ def main():
            sys.exit(1)
        ops = [op for op in ops if filter_re.search(op['op_text'])]

-    if args.head is not None:
-        ops = ops[:args.head]
-    elif args.tail is not None:
-        ops = ops[-args.tail:]
+    if args.head is not None or args.tail is not None:
+        ops_by_dev = defaultdict(list)
+        for op in ops:
+            ops_by_dev[op['device']].append(op)
+
+        filtered_ops = []
+        for dev in sorted(ops_by_dev.keys()):
+            dev_ops = ops_by_dev[dev]
+            if args.head is not None:
+                dev_ops = dev_ops[:args.head]
+            elif args.tail is not None:
+                dev_ops = dev_ops[-args.tail:]
+            filtered_ops.extend(dev_ops)
+        ops = filtered_ops

    if args.filter or args.head is not None or args.tail is not None:
-        valid_ranges = []
+        # Group valid ranges by device
+        valid_ranges_by_dev = defaultdict(list)
        for op in ops:
            start_cyc = op['unwrapped_cycles_start']
            end_cyc = start_cyc + op['cycles'] if start_cyc is not None else None
            if start_cyc is not None and end_cyc is not None:
-                valid_ranges.append((start_cyc, end_cyc))
+                valid_ranges_by_dev[op['device']].append((start_cyc, end_cyc))

-        valid_ranges.sort(key=lambda r: r[0])
-        range_starts = [r[0] for r in valid_ranges]
+        for dev in valid_ranges_by_dev:
+            valid_ranges_by_dev[dev].sort(key=lambda r: r[0])
+
+        range_starts_by_dev = {dev: [r[0] for r in ranges] for dev, ranges in valid_ranges_by_dev.items()}

        filtered_traces = []
        for e in traces:
            cyc = e['unwrapped_cycles']
            if cyc is None:
                continue
+            dev = e['device']
+            range_starts = range_starts_by_dev.get(dev)
+            if not range_starts:
+                continue
            idx = bisect.bisect_right(range_starts, cyc) - 1
            if idx >= 0:
-                start, end = valid_ranges[idx]
+                start, end = valid_ranges_by_dev[dev][idx]
                if start <= cyc <= end:
                    filtered_traces.append(e)
        traces = filtered_traces

-    generate_perfetto_trace(ops, traces, args.output)
+    if args.device == "split":
+        unique_devices = sorted(list(set(op['device'] for op in ops)))
+        for dev in unique_devices:
+            dev_ops = [op for op in ops if device_matches(op['device'], dev)]
+            dev_traces = [t for t in traces if device_matches(t['device'], dev)]
+            out_path = get_split_output_path(args.output, dev)
+            generate_perfetto_trace(dev_ops, dev_traces, out_path)
+    else:
+        generate_perfetto_trace(ops, traces, args.output)


 if __name__ == "__main__":
@@ -0,0 +1,405 @@
+#!/usr/bin/env python3
+#
+# Run llama.cpp tools on Snapdragon devices (natively, via ADB, or SSH).
+#
+
+import sys
+import os
+import argparse
+import subprocess
+import platform
+import shlex
+import logging
+
+logger = logging.getLogger("run")
+
+
+def parse_target(target_str):
+    if not target_str:
+        return None, None
+    if target_str.startswith("adb") or target_str.startswith("android"):
+        parts = target_str.split(":", 1)
+        serial = parts[1] if len(parts) > 1 else None
+        return "android", serial
+    elif target_str.startswith("lnx") or target_str.startswith("linux") or target_str.startswith("ubuntu"):
+        parts = target_str.split(":", 1)
+        host = parts[1] if len(parts) > 1 else None
+        return "linux", host
+    elif target_str in ("wos", "windows"):
+        return "windows", None
+    else:
+        return None, None
+
+
+def shlex_join(args_list):
+    if hasattr(shlex, 'join'):
+        return shlex.join(args_list)
+    import pipes
+    return " ".join(pipes.quote(x) for x in args_list)
+
+
+def main():
+    logging.basicConfig(level=logging.INFO, format='%(message)s')
+    # Split arguments at '--'
+    if '--' in sys.argv:
+        idx = sys.argv.index('--')
+        run_args = sys.argv[1:idx]
+        cmd_args = sys.argv[idx + 1:]
+    else:
+        run_args = sys.argv[1:]
+        cmd_args = []
+
+    parser = argparse.ArgumentParser(
+        description="Unified runner for llama.cpp tools on Snapdragon (natively, via ADB, or via SSH)."
+    )
+    parser.add_argument("--target", help="Execution target (e.g. android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, windows/wos) (default: local run)")
+    parser.add_argument("--target-dir", help="Target directory on the device (default: /data/local/tmp/llama.cpp for Android, ~/llama.cpp for Linux)")
+    parser.add_argument("--install-dir", help="Install directory name (defaults to pkg-TARGET or pkg-TARGET-dbg prefix based on target)")
+    parser.add_argument("--debug", action="store_true", help="Use debug build (defaults to pkg-TARGET-dbg folder)")
+    parser.add_argument("--devices", "--device", "-d", help="Select execution devices (split into NPU and OpenCL GPUs automatically, default: HTP0)")
+    parser.add_argument("--verbose", help="Verbose level (enables both Hexagon and OpenCL kernel cache debugging)")
+    parser.add_argument("--profile", help="Profiling flag (enables Hexagon profiling and OpenCL autotuning)")
+    parser.add_argument("--sched-debug", action="store_true", help="Enable GGML/llama.cpp scheduler debug output (GGML_SCHED_DEBUG=2)")
+    parser.add_argument("--mtmd-device", help="Specify the backend device ID for Multi-Threaded Multi-Device setup (MTMD_BACKEND_DEVICE)")
+
+    # Hexagon specific parameters
+    parser.add_argument("--hex-verbose", help="Enable verbose logging (GGML_HEXAGON_VERBOSE)")
+    parser.add_argument("--hex-profile", help="Enable NPU/Hexagon profiling and performance metrics print (GGML_HEXAGON_PROFILE)")
+    parser.add_argument("--hex-nhvx", help="Number of HVX units to use (GGML_HEXAGON_NHVX)")
+    parser.add_argument("--hex-nhmx", help="Number of HMX units to use. 0 disables HMX power-up (GGML_HEXAGON_NHMX)")
+    parser.add_argument("--hex-hostbuf", help="Enable host buffers (GGML_HEXAGON_HOSTBUF)")
+    parser.add_argument("--hex-opbatch", help="Maximum number of operations to batch into a single HTP execution (GGML_HEXAGON_OPBATCH)")
+    parser.add_argument("--hex-opqueue", help="Size of the asynchronous NPU operation queue (GGML_HEXAGON_OPQUEUE)")
+    parser.add_argument("--hex-oppoll", default="1", help="Enable (1) or Disable (0) polling for NPU opbatch completion (GGML_HEXAGON_OPPOLL) (default: 1)")
+    parser.add_argument("--hex-opfilter", help="Regex pattern to filter/select which operators are offloaded to NPU (GGML_HEXAGON_OPFILTER)")
+    parser.add_argument("--hex-opfusion", help="NPU graph node fusion optimization level (0: disabled, 1: enabled) (GGML_HEXAGON_OPFUSION)")
+    parser.add_argument("--hex-vmem", help="Maximum NPU VMEM size limit in MB to allocate (GGML_HEXAGON_VMEM)")
+    parser.add_argument("--hex-mbuf", help="Maximum host buffer size limit in MB to allocate (GGML_HEXAGON_MBUF)")
+    parser.add_argument("--hex-mm-select", help="Select MUL_MAT and MUL_MAT_ID kernel (GGML_HEXAGON_MM_SELECT) 3:HMX,2:HVX-tiled,1:HVX-flat,0:disable")
+    parser.add_argument("--hex-fa-select", help="Select Flash Attention kernel (GGML_HEXAGON_FA_SELECT) 2:HMX,1:HVX,0:disable")
+    parser.add_argument("--hex-ar-select", help="Select All-Reduce kernel (GGML_HEXAGON_AR_SELECT) 1:enable,0:disable")
+    parser.add_argument("--hex-etm", help="Enable Embedded Trace Macrocell hardware tracing / trace logging (GGML_HEXAGON_ETM)")
+    parser.add_argument("--hex-arch", help="Target Hexagon NPU architecture version override (v73, v75, v79, v81, etc.) (GGML_HEXAGON_ARCH)")
+    parser.add_argument("--hex-optrace", help="Trace buffer size in number of records (GGML_HEXAGON_OPTRACE)")
+
+    # OpenCL specific parameters
+    parser.add_argument("--cl-platform", help="Select OpenCL platform name/regex (e.g. Qualified Qualcomm OpenCL platform) (GGML_OPENCL_PLATFORM)")
+    parser.add_argument("--cl-device", help="Select OpenCL device name/regex (e.g. Adreno GPU) (GGML_OPENCL_DEVICE)")
+    parser.add_argument("--cl-opfilter", help="Regex pattern to filter/select which operators are offloaded to OpenCL (GGML_OPENCL_OPFILTER)")
+    parser.add_argument("--cl-disable-fusion", action="store_true", help="Disable OpenCL kernel fusion optimizations (GGML_OPENCL_DISABLE_FUSION)")
+    parser.add_argument("--cl-cache-dir", help="Directory path to store compiled OpenCL program binaries (GGML_OPENCL_KERNEL_CACHE_DIR)")
+    parser.add_argument("--cl-cache-debug", help="Enable verbose debugging logs for the kernel caching system (GGML_OPENCL_KERNEL_CACHE_DEBUG)")
+    parser.add_argument("--cl-fa-tune", action="store_true", help="Enable automatic Flash Attention kernel autotuning (GGML_OPENCL_FA_TUNE)")
+    parser.add_argument("--cl-adreno-xmem", action="store_true", help="Enforce matmul using texture/image (xmem) memory paths on Adreno GPUs (GGML_OPENCL_ADRENO_XMEM_GEMM)")
+    parser.add_argument("--cl-adreno-large-buffer", action="store_true", help="Allow allocating larger buffer sizes on Adreno GPUs (GGML_OPENCL_ADRENO_USE_LARGE_BUFFER)")
+
+    args = parser.parse_args(run_args)
+
+    if not cmd_args:
+        parser.print_help()
+        logger.error("\nError: No command specified after '--'")
+        sys.exit(1)
+
+    target_type = None
+    target_val = None
+    target_prefix = None
+    if args.target:
+        target_type, target_val = parse_target(args.target)
+        if not target_type:
+            logger.error(f"Error: Invalid target format '{args.target}'. Must be android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, or windows/wos.")
+            sys.exit(1)
+        target_prefix = args.target.split(":", 1)[0]
+
+    # Resolve install directory
+    install_dir = args.install_dir
+    if not install_dir:
+        if target_prefix:
+            suffix = "-dbg" if args.debug else ""
+            install_dir = f"pkg-{target_prefix}{suffix}"
+        else:
+            # Smart branch folder detection for local run if default is not set
+            prefixes = ("wos", "windows", "lnx", "linux", "ubuntu", "adb", "android")
+            suffixes = ("-dbg", "") if args.debug else ("", "-dbg")
+            found = False
+            for suffix in suffixes:
+                for prefix in prefixes:
+                    test_path = f"./pkg-{prefix}{suffix}/llama.cpp"
+                    if os.path.exists(test_path):
+                        install_dir = f"pkg-{prefix}{suffix}"
+                        found = True
+                        break
+                if found:
+                    break
+            if not install_dir:
+                install_dir = "pkg-android"  # Fallback default
+
+    # Host side package path
+    package_path = os.path.join(install_dir, "llama.cpp")
+
+    # Environment variables to map
+    env_vars = {}
+
+    def set_env(env_name, opt_val):
+        if opt_val is not None:
+            env_vars[env_name] = str(opt_val)
+        elif env_name in os.environ:
+            env_vars[env_name] = os.environ[env_name]
+
+    # Resolve and filter devices (HTP vs OpenCL)
+    devices_val = args.devices if args.devices is not None else "HTP0"
+    if devices_val.isdigit():
+        hex_devices = devices_val
+        cl_device = ""
+    else:
+        parts = [p.strip() for p in devices_val.split(",")]
+        # Any device containing "htp" is Hexagon, rest is OpenCL
+        hex_parts = [p for p in parts if "htp" in p.lower()]
+        cl_parts = [p for p in parts if "htp" not in p.lower()]
+        hex_devices = ",".join(hex_parts)
+        cl_device = ",".join(cl_parts)
+
+    # Set Hexagon devices
+    if hex_devices:
+        env_vars["GGML_HEXAGON_DEVICES"] = hex_devices
+    elif "GGML_HEXAGON_DEVICES" in os.environ:
+        env_vars["GGML_HEXAGON_DEVICES"] = os.environ["GGML_HEXAGON_DEVICES"]
+
+    # Set OpenCL device (unless overridden by --cl-device)
+    final_cl_device = args.cl_device if args.cl_device is not None else cl_device
+    if final_cl_device:
+        env_vars["GGML_OPENCL_DEVICE"] = final_cl_device
+    elif "GGML_OPENCL_DEVICE" in os.environ:
+        env_vars["GGML_OPENCL_DEVICE"] = os.environ["GGML_OPENCL_DEVICE"]
+
+    # Map shared & backend-specific parameters with correct overrides
+
+    # Verbose logging mapping
+    hex_verbose_val = args.hex_verbose if args.hex_verbose is not None else args.verbose
+    set_env("GGML_HEXAGON_VERBOSE", hex_verbose_val)
+
+    cl_cache_debug_val = args.cl_cache_debug if args.cl_cache_debug is not None else args.verbose
+    set_env("GGML_OPENCL_KERNEL_CACHE_DEBUG", cl_cache_debug_val)
+
+    # Profiling mapping
+    hex_profile_val = args.hex_profile if args.hex_profile is not None else args.profile
+    set_env("GGML_HEXAGON_PROFILE", hex_profile_val)
+
+    if args.cl_fa_tune or args.profile is not None:
+        env_vars["GGML_OPENCL_FA_TUNE"] = "1"
+    elif "GGML_OPENCL_FA_TUNE" in os.environ:
+        env_vars["GGML_OPENCL_FA_TUNE"] = os.environ["GGML_OPENCL_FA_TUNE"]
+
+    # Other Hexagon environment variables
+    set_env("GGML_HEXAGON_NHVX", args.hex_nhvx)
+    set_env("GGML_HEXAGON_NHMX", args.hex_nhmx)
+    set_env("GGML_HEXAGON_HOSTBUF", args.hex_hostbuf)
+    set_env("GGML_HEXAGON_OPBATCH", args.hex_opbatch)
+    set_env("GGML_HEXAGON_OPQUEUE", args.hex_opqueue)
+    set_env("GGML_HEXAGON_OPPOLL", args.hex_oppoll)
+    set_env("GGML_HEXAGON_OPFILTER", args.hex_opfilter)
+    set_env("GGML_HEXAGON_OPFUSION", args.hex_opfusion)
+    set_env("GGML_HEXAGON_VMEM", args.hex_vmem)
+    set_env("GGML_HEXAGON_MBUF", args.hex_mbuf)
+    set_env("GGML_HEXAGON_MM_SELECT", args.hex_mm_select)
+    set_env("GGML_HEXAGON_FA_SELECT", args.hex_fa_select)
+    set_env("GGML_HEXAGON_AR_SELECT", args.hex_ar_select)
+    set_env("GGML_HEXAGON_ETM", args.hex_etm)
+    set_env("GGML_HEXAGON_ARCH", args.hex_arch)
+    set_env("GGML_HEXAGON_OPTRACE", args.hex_optrace)
+    set_env("MTMD_BACKEND_DEVICE", args.mtmd_device)
+
+    # OpenCL environment variables
+    set_env("GGML_OPENCL_PLATFORM", args.cl_platform)
+    set_env("GGML_OPENCL_OPFILTER", args.cl_opfilter)
+    set_env("GGML_OPENCL_KERNEL_CACHE_DIR", args.cl_cache_dir)
+
+    if args.cl_disable_fusion:
+        env_vars["GGML_OPENCL_DISABLE_FUSION"] = "1"
+    elif "GGML_OPENCL_DISABLE_FUSION" in os.environ:
+        env_vars["GGML_OPENCL_DISABLE_FUSION"] = os.environ["GGML_OPENCL_DISABLE_FUSION"]
+
+    if args.cl_adreno_xmem:
+        env_vars["GGML_OPENCL_ADRENO_XMEM_GEMM"] = "1"
+    elif "GGML_OPENCL_ADRENO_XMEM_GEMM" in os.environ:
+        env_vars["GGML_OPENCL_ADRENO_XMEM_GEMM"] = os.environ["GGML_OPENCL_ADRENO_XMEM_GEMM"]
+
+    if args.cl_adreno_large_buffer:
+        env_vars["GGML_OPENCL_ADRENO_USE_LARGE_BUFFER"] = "1"
+    elif "GGML_OPENCL_ADRENO_USE_LARGE_BUFFER" in os.environ:
+        env_vars["GGML_OPENCL_ADRENO_USE_LARGE_BUFFER"] = os.environ["GGML_OPENCL_ADRENO_USE_LARGE_BUFFER"]
+
+    if args.sched_debug:
+        env_vars["GGML_SCHED_DEBUG"] = "2"
+
+    # Resolve executable path
+    executable = cmd_args[0]
+    known_binaries = ["llama-cli", "llama-bench", "llama-completion", "llama-mtmd-cli", "test-backend-ops"]
+    if executable in known_binaries:
+        if target_type in ("android", "linux"):
+            resolved_exec = f"./bin/{executable}"
+        else:
+            if platform.system() == "Windows":
+                resolved_exec = os.path.normpath(os.path.join(package_path, "bin", f"{executable}.exe"))
+            else:
+                resolved_exec = os.path.normpath(os.path.join(package_path, "bin", executable))
+        cmd_args[0] = resolved_exec
+
+    # Infer device string to pass to the tool
+    basename = os.path.basename(executable)
+    if basename.endswith(".exe"):
+        basename = basename[:-4]
+
+    device_val = None
+    if basename == "test-backend-ops":
+        for i in range(len(cmd_args)):
+            if cmd_args[i] in ("-p", "--params") and i + 1 < len(cmd_args):
+                val = cmd_args[i + 1]
+                new_val = ""
+                for j, char in enumerate(val):
+                    if char in ('[', ']'):
+                        if j > 0 and val[j - 1] == '\\':
+                            new_val += char
+                        else:
+                            new_val += '\\' + char
+                    else:
+                        new_val += char
+                cmd_args[i + 1] = new_val
+
+        has_b = any(arg == "-b" for arg in cmd_args)
+        if not has_b:
+            if args.devices:
+                if args.devices.isdigit():
+                    n = int(args.devices)
+                    device_val = ",".join(f"HTP{i}" for i in range(n))
+                else:
+                    device_val = args.devices
+            elif "D" in os.environ:
+                device_val = os.environ["D"]
+            elif "DEVICE" in os.environ:
+                device_val = os.environ["DEVICE"]
+            else:
+                device_val = "HTP0"
+            if device_val:
+                cmd_args += ["-b", device_val]
+    else:
+        has_device = any(arg.startswith("--device") for arg in cmd_args)
+        if not has_device:
+            if args.devices:
+                if args.devices.isdigit():
+                    n = int(args.devices)
+                    device_val = ",".join(f"HTP{i}" for i in range(n))
+                else:
+                    device_val = args.devices
+            elif "D" in os.environ:
+                device_val = os.environ["D"]
+            elif "DEVICE" in os.environ:
+                device_val = os.environ["DEVICE"]
+            else:
+                device_val = "HTP0"
+            if device_val:
+                cmd_args += ["--device", device_val]
+
+    # Automatically add -v to known llama tools if sched-debug, verbose, or profile are set
+    verbose_trigger = (
+        args.sched_debug
+        or args.verbose is not None
+        or args.profile is not None
+        or args.hex_verbose is not None
+        or args.hex_profile is not None
+        or args.hex_optrace is not None
+    )
+    if verbose_trigger and basename in ("llama-cli", "llama-completion", "llama-bench", "llama-server", "llama-mtmd-cli"):
+        if "-v" not in cmd_args and "--verbose" not in cmd_args:
+            cmd_args.append("-v")
+
+    # Inject defaults for llama-cli, llama-completion, and llama-server if not overridden by the user
+    if basename in ("llama-cli", "llama-completion", "llama-server"):
+        if "-ngl" not in cmd_args and "--n-gpu-layers" not in cmd_args:
+            cmd_args += ["-ngl", "99"]
+        if "--ubatch-size" not in cmd_args and "-ub" not in cmd_args:
+            cmd_args += ["--ubatch-size", "1024"]
+        if "-fa" not in cmd_args and "--flash-attn" not in cmd_args:
+            cmd_args += ["-fa", "on"]
+
+    if basename in ("llama-cli", "llama-completion", "llama-server", "llama-bench"):
+        if "-t" not in cmd_args and "--threads" not in cmd_args:
+            cmd_args += ["-t", "6"]
+
+    # Resolve target directory on device
+    target_dir = args.target_dir
+    if not target_dir:
+        target_dir = "/data/local/tmp/llama.cpp" if target_type == "android" else "~/llama.cpp"
+
+    if target_type == "android":
+        # Run via ADB
+        adb_base = ["adb"]
+        if target_val: # serial
+            adb_base += ["-s", target_val]
+
+        env_parts = [
+            "LD_LIBRARY_PATH=./lib",
+            "ADSP_LIBRARY_PATH=./lib"
+        ]
+        for k, v in env_vars.items():
+            env_parts.append(f"{k}={v}")
+        env_str = " ".join(env_parts)
+
+        cmd_str = shlex_join(cmd_args)
+        adb_shell_cmd = f"cd {target_dir} && ulimit -c unlimited && {env_str} {cmd_str}"
+        full_cmd = adb_base + ["shell", adb_shell_cmd]
+
+        logger.info(f"+ {' '.join(full_cmd)}")
+        res = subprocess.run(full_cmd)
+        sys.exit(res.returncode)
+
+    elif target_type == "linux":
+        ssh_host = target_val
+        if not ssh_host:
+            logger.error("Error: SSH host not specified in target (e.g. use linux:user@host, lnx:user@host, or ubuntu:user@host). Cannot execute.")
+            sys.exit(1)
+
+        # Linux remote run via SSH
+        env_parts = [
+            "LD_LIBRARY_PATH=./lib",
+            "ADSP_LIBRARY_PATH=./lib"
+        ]
+        for k, v in env_vars.items():
+            env_parts.append(f"{k}={v}")
+        env_str = " ".join(env_parts)
+
+        cmd_str = shlex_join(cmd_args)
+        ssh_shell_cmd = f"cd {target_dir} && ulimit -c unlimited && {env_str} {cmd_str}"
+        full_cmd = ["ssh", ssh_host, ssh_shell_cmd]
+
+        logger.info(f"+ {' '.join(full_cmd)}")
+        res = subprocess.run(full_cmd)
+        sys.exit(res.returncode)
+
+    elif target_type == "windows":
+        logger.info("Windows target execution is currently a stub.")
+        sys.exit(0)
+
+    else:
+        # Run locally
+        local_env = os.environ.copy()
+        lib_dir = os.path.normpath(os.path.join(package_path, "lib"))
+        local_env["ADSP_LIBRARY_PATH"] = lib_dir
+        if platform.system() == "Windows":
+            local_env["PATH"] = lib_dir + os.path.pathsep + local_env.get("PATH", "")
+        else:
+            local_env["LD_LIBRARY_PATH"] = lib_dir + os.path.pathsep + local_env.get("LD_LIBRARY_PATH", "")
+
+        for k, v in env_vars.items():
+            local_env[k] = v
+
+        logger.info(f"+ {shlex_join(cmd_args)}")
+        res = subprocess.run(cmd_args, env=local_env)
+        sys.exit(res.returncode)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        logger.info("\nInterrupted by user.")
+        sys.exit(130)
@@ -1,48 +0,0 @@
-
-#!/usr/bin/env pwsh
-
-# Basedir on device
-$basedir=".\pkg-snapdragon"
-
-$cli_opts=$args
-
-$model="Llama-3.2-3B-Instruct-Q4_0.gguf"
-if ($null -ne $env:M) {
-    $model=$env:M
-}
-
-$device="HTP0"
-if ($null -ne $env:D) {
-    $device=$env:D
-}
-
-if ($null -ne $env:V) {
-    $env:GGML_HEXAGON_VERBOSE=$env:V
-}
-
-if ($null -ne $env:PROF) {
-    $env:GGML_HEXAGON_PROFILE=$env:PROF
-}
-
-if ($null -ne $env:OPSTAGE) {
-    $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE
-}
-
-if ($null -ne $env:NHVX) {
-    $env:GGML_HEXAGON_NHVX=$env:NHVX
-}
-
-if ($null -ne $env:NDEV) {
-    $env:GGML_HEXAGON_NDEV=$env:NDEV
-}
-
-if ($null -ne $env:HB) {
-    $env:GGML_HEXAGON_HOSTBUF=$env:HB
-}
-
-$env:ADSP_LIBRARY_PATH="$basedir\lib"
-
-& "$basedir\bin\llama-bench.exe" `
-    --load-mode none -m $basedir\..\..\gguf\$model `
-    --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 `
-    --ubatch-size 1024 -ngl 99 --device $device $cli_opts
@@ -1,53 +0,0 @@
-
-#!/usr/bin/env pwsh
-
-# Basedir on device
-$basedir=".\pkg-snapdragon"
-
-$cli_opts=$args
-
-$model="Llama-3.2-3B-Instruct-Q4_0.gguf"
-if ($null -ne $env:M) {
-    $model=$env:M
-}
-
-$device="HTP0"
-if ($null -ne $env:D) {
-    $device=$env:D
-}
-
-if ($null -ne $env:V) {
-    $env:GGML_HEXAGON_VERBOSE=$env:V
-}
-
-if ($null -ne $env:SCHED) {
-    $env:GGML_SCHED_DEBUG=$env:SCHED; $cli_opts="$cli_opts -v"
-}
-
-if ($null -ne $env:PROF) {
-    $env:GGML_HEXAGON_PROFILE=$env:PROF
-}
-
-if ($null -ne $env:OPSTAGE) {
-    $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE
-}
-
-if ($null -ne $env:NHVX) {
-    $env:GGML_HEXAGON_NHVX=$env:NHVX
-}
-
-if ($null -ne $env:NDEV) {
-    $env:GGML_HEXAGON_NDEV=$env:NDEV
-}
-
-if ($null -ne $env:HB) {
-    $env:GGML_HEXAGON_HOSTBUF=$env:HB
-}
-
-$env:ADSP_LIBRARY_PATH="$basedir\lib"
-
-& "$basedir\bin\llama-cli.exe" `
-    --load-mode none -m $basedir\..\..\gguf\$model `
-    --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 `
-    --ctx-size 8192 --ubatch-size 1024 -fa on `
-    -ngl 99 --device $device $cli_opts
@@ -1,53 +0,0 @@
-
-#!/usr/bin/env pwsh
-
-# Basedir on device
-$basedir=".\pkg-snapdragon"
-
-$cli_opts=$args
-
-$model="Llama-3.2-3B-Instruct-Q4_0.gguf"
-if ($null -ne $env:M) {
-    $model=$env:M
-}
-
-$device="HTP0"
-if ($null -ne $env:D) {
-    $device=$env:D
-}
-
-if ($null -ne $env:V) {
-    $env:GGML_HEXAGON_VERBOSE=$env:V
-}
-
-if ($null -ne $env:SCHED) {
-    $env:GGML_SCHED_DEBUG=$env:SCHED; $cli_opts="$cli_opts -v"
-}
-
-if ($null -ne $env:PROF) {
-    $env:GGML_HEXAGON_PROFILE=$env:PROF
-}
-
-if ($null -ne $env:OPSTAGE) {
-    $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE
-}
-
-if ($null -ne $env:NHVX) {
-    $env:GGML_HEXAGON_NHVX=$env:NHVX
-}
-
-if ($null -ne $env:NDEV) {
-    $env:GGML_HEXAGON_NDEV=$env:NDEV
-}
-
-if ($null -ne $env:HB) {
-    $env:GGML_HEXAGON_HOSTBUF=$env:HB
-}
-
-$env:ADSP_LIBRARY_PATH="$basedir\lib"
-
-& "$basedir\bin\llama-completion.exe" `
-    --load-mode none -m $basedir\..\..\gguf\$model `
-    --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 `
-    --ctx-size 8192 --ubatch-size 1024 -fa on `
-    -ngl 99 -no-cnv --device $device $cli_opts
@@ -1,68 +0,0 @@
-#!/usr/bin/env pwsh
-
-# Basedir on device
-$basedir=".\pkg-snapdragon"
-
-$cli_opts=$args
-
-$model="gemma-3-4b-it-Q4_0.gguf"
-if ($null -ne $env:M) {
-    $model=$env:M
-}
-
-$mmproj="mmproj-F16.gguf"
-if ($null -ne $env:MMPROJ) {
-    $mmproj=$env:MMPROJ
-}
-
-$image=""
-if ($null -ne $env:IMG) {
-    $image=$env:IMG
-}
-
-$device="HTP0"
-if ($null -ne $env:D) {
-    $device=$env:D
-}
-
-if ($null -ne $env:V) {
-    $env:GGML_HEXAGON_VERBOSE=$env:V
-}
-
-if ($null -ne $env:SCHED) {
-    $env:GGML_SCHED_DEBUG=$env:SCHED; $cli_opts="$cli_opts -v"
-}
-
-if ($null -ne $env:PROF) {
-    $env:GGML_HEXAGON_PROFILE=$env:PROF
-}
-
-if ($null -ne $env:OPSTAGE) {
-    $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE
-}
-
-if ($null -ne $env:NHVX) {
-    $env:GGML_HEXAGON_NHVX=$env:NHVX
-}
-
-if ($null -ne $env:NDEV) {
-    $env:GGML_HEXAGON_NDEV=$env:NDEV
-}
-
-if ($null -ne $env:HB) {
-    $env:GGML_HEXAGON_HOSTBUF=$env:HB
-}
-
-if ($null -ne $env:MTMD_DEVICE) {
-    $env:MTMD_BACKEND_DEVICE=$env:MTMD_DEVICE
-}
-
-$env:ADSP_LIBRARY_PATH="$basedir\lib"
-
-& "$basedir\bin\llama-mtmd-cli.exe" `
-    --load-mode none -m $basedir\..\..\gguf\$model `
-    --mmproj $basedir\..\..\gguf\$mmproj `
-    --image $basedir\..\..\gguf\$image `
-    --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 `
-    --ctx-size 8192 --ubatch-size 1024 -fa on `
-    -ngl 99 --device $device -v $cli_opts
@@ -1,56 +0,0 @@
-
-#!/usr/bin/env pwsh
-
-# Basedir on device
-$basedir=".\pkg-snapdragon"
-
-if ($args.Count -eq 0) {
-    Write-Host "No arguments provided.Expected the tool and argument to run."
-    exit -1
-}
-
-$tool=$args[0]
-$cli_opts=@()
-
-if ($args.Count -gt 1) {
-    $cli_opts=$args[1..($args.Count - 1)]
-    $remainingArgs = $args[1..($args.Count - 1)]
-}
-
-$device="HTP0"
-if ($null -ne $env:D) {
-    $device=$env:D
-}
-
-if ($null -ne $env:V) {
-    $env:GGML_HEXAGON_VERBOSE=$env:V
-}
-
-if ($null -ne $env:SCHED) {
-    $env:GGML_SCHED_DEBUG=$env:SCHED; $cli_opts="$cli_opts -v"
-}
-
-if ($null -ne $env:PROF) {
-    $env:GGML_HEXAGON_PROFILE=$env:PROF
-}
-
-if ($null -ne $env:OPSTAGE) {
-    $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE
-}
-
-if ($null -ne $env:NHVX) {
-    $env:GGML_HEXAGON_NHVX=$env:NHVX
-}
-
-if ($null -ne $env:NDEV) {
-    $env:GGML_HEXAGON_NDEV=$env:NDEV
-}
-
-if ($null -ne $env:HB) {
-    $env:GGML_HEXAGON_HOSTBUF=$env:HB
-}
-
-$env:ADSP_LIBRARY_PATH="$basedir\lib"
-
-& "$basedir\bin\$tool" `
-    $cli_opts
@@ -1,105 +0,0 @@
-# Requires Run as Administrator is NOT strictly necessary for User-scope env vars,
-# but recommended for creating directories in C:\ root if permissions are restricted.
-
-$ErrorActionPreference = "Stop"
-
-# --- Configuration ---
-$BaseDir = "C:\Qualcomm"
-
-# SDK 1: Hexagon
-$HexagonUrl     = "https://github.com/snapdragon-toolchain/hexagon-sdk/releases/download/v6.6.0.0/hexagon-sdk-v6.6.0.0-arm64-wos.tar.xz"
-$HexagonParent  = Join-Path $BaseDir "Hexagon_SDK"
-$HexagonSdkVersion   = "6.6.0.0"
-$HexagonToolsVersion = "19.0.07"
-$HexagonSdkTarget    = Join-Path $HexagonParent $HexagonSdkVersion
-$HexagonToolsTarget  = Join-Path $HexagonSdkTarget "\tools\HEXAGON_Tools\$HexagonToolsVersion"
-
-# SDK 2: OpenCL
-$OpenCLUrl      = "https://github.com/snapdragon-toolchain/opencl-sdk/releases/download/v2.3.2/adreno-opencl-sdk-v2.3.2-arm64-wos.tar.xz"
-$OpenCLParent   = Join-Path $BaseDir "OpenCL_SDK"
-$OpenCLVersion  = "2.3.2"
-$OpenCLTarget   = Join-Path $OpenCLParent $OpenCLVersion
-
-# --- Helper Function ---
-function Install-QualcommSDK {
-    param (
-        [string]$Url,
-        [string]$ParentDir,
-        [string]$TargetDir,
-        [string]$Name
-    )
-
-    # 1. Create Parent Directory
-    if (-not (Test-Path -Path $ParentDir)) {
-        Write-Host "Creating directory: $ParentDir" -ForegroundColor Cyan
-        New-Item -Path $ParentDir -ItemType Directory -Force | Out-Null
-    }
-
-    # 2. Check for Specific Version Directory
-    if (Test-Path -Path $TargetDir) {
-        Write-Host "$Name ($TargetDir) already exists. Skipping download." -ForegroundColor Green
-    }
-    else {
-        Write-Host "$Name not found. preparing to download..." -ForegroundColor Yellow
-
-        # Create the target directory to extract into
-        New-Item -Path $TargetDir -ItemType Directory -Force | Out-Null
-
-        # Define temporary archive path
-        $TempFile = Join-Path $ParentDir "temp_sdk.tar.xz"
-
-        try {
-            # Download
-            Write-Host "Downloading from: $Url"
-            Invoke-WebRequest -Uri $Url -OutFile $TempFile
-
-            # Untar
-            # Note: We assume Windows includes tar.exe (Win 10 build 17063+)
-            Write-Host "Extracting archive to $TargetDir..."
-
-            # We use -C to extract contents INTO the target directory created above
-            tar -xJvf $TempFile -C $TargetDir\..
-
-            Write-Host "Extraction complete." -ForegroundColor Green
-        }
-        catch {
-            Write-Error "Failed to download or extract $Name. Error: $_"
-            # Cleanup target dir if failed so script tries again next time
-            Remove-Item -Path $TargetDir -Recurse -Force -ErrorAction SilentlyContinue
-        }
-        finally {
-            # Cleanup Archive
-            if (Test-Path $TempFile) { Remove-Item $TempFile -Force }
-        }
-    }
-}
-
-# --- Execution ---
-
-# 1. Ensure Base C:\Qualcomm exists
-if (-not (Test-Path $BaseDir)) {
-    New-Item -Path $BaseDir -ItemType Directory -Force | Out-Null
-}
-
-# 2. Run Install Logic
-Install-QualcommSDK -Url $HexagonUrl -ParentDir $HexagonParent -TargetDir $HexagonSdkTarget -Name "Hexagon SDK"
-Install-QualcommSDK -Url $OpenCLUrl -ParentDir $OpenCLParent -TargetDir $OpenCLTarget -Name "OpenCL SDK"
-
-# --- Environment Variables ---
-
-Write-Host "`nSetting Environment Variables..." -ForegroundColor Cyan
-
-# Set OPENCL_SDK_ROOT
-[System.Environment]::SetEnvironmentVariable('OPENCL_SDK_ROOT', $OpenCLTarget, [System.EnvironmentVariableTarget]::User)
-$env:OPENCL_SDK_ROOT = $OpenCLTarget # Set for current session as well
-Write-Host "OPENCL_SDK_ROOT set to:  $OpenCLTarget"
-
-# Set HEXAGON_SDK_ROOT
-[System.Environment]::SetEnvironmentVariable('HEXAGON_SDK_ROOT', $HexagonSdkTarget, [System.EnvironmentVariableTarget]::User)
-$env:HEXAGON_SDK_ROOT = $HexagonSdkTarget # Set for current session as well
-Write-Host "HEXAGON_SDK_ROOT set to: $HexagonSdkTarget"
-
-# Set HEXAGON_SDK_ROOT
-[System.Environment]::SetEnvironmentVariable('HEXAGON_TOOLS_ROOT', $HexagonToolsTarget, [System.EnvironmentVariableTarget]::User)
-$env:HEXAGON_TOOLS_ROOT = $HexagonToolsTarget # Set for current session as well
-Write-Host "HEXAGON_TOOLS_ROOT set to: $HexagonToolsTarget"
@@ -12,6 +12,7 @@
 #include <limits>
 #include <map>
 #include <stdexcept>
+#include <unordered_map>

 static bool ggml_is_power_of_2(int n) {
    return (n & (n - 1)) == 0;
@@ -1128,11 +1129,18 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &

            cells.pos_set(idx, ubatch.pos[i]);

-            if (ubatch.is_pos_2d()) {
-                llama_kv_cell_ext ext {
-                    /*.x =*/ ubatch.pos[i + ubatch.n_tokens*2],
-                    /*.y =*/ ubatch.pos[i + ubatch.n_tokens],
-                };
+            if (ubatch.is_pos_2d() || ubatch.token) {
+                llama_kv_cell_ext ext;
+
+                if (ubatch.is_pos_2d()) {
+                    ext.x = ubatch.pos[i + ubatch.n_tokens*2];
+                    ext.y = ubatch.pos[i + ubatch.n_tokens];
+                }
+
+                if (ubatch.token) {
+                    ext.tok = ubatch.token[i];
+                }
+
                cells.ext_set(idx, ext);
            }

@@ -1805,6 +1813,69 @@ void llama_kv_cache::set_input_v_rot(ggml_tensor * dst) const {
    memcpy(dst->data, attn_rot_hadamard.at(n_rot).data(), ggml_nbytes(dst));
 }

+bool llama_kv_cache::has_cell_ext() const {
+    return hparams.n_pos_per_embd() > 1;
+}
+
+void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const {
+    const uint32_t n_tokens = ubatch.n_tokens;
+
+    res.clear();
+    res.resize(n_tokens*n, LLAMA_TOKEN_NULL);
+
+    if (n == 0) {
+        return;
+    }
+
+    // note: apply_ubatch() has already stored the current ubatch
+    //       the window below thus covers tokens of this very ubatch as well, which is what we want
+    llama_pos p_min = std::numeric_limits<llama_pos>::max();
+    llama_pos p_max = std::numeric_limits<llama_pos>::min();
+
+    std::bitset<LLAMA_MAX_SEQ> seqs;
+
+    for (uint32_t i = 0; i < n_tokens; ++i) {
+        p_min = std::min(p_min, ubatch.pos[i]);
+        p_max = std::max(p_max, ubatch.pos[i]);
+    }
+
+    for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
+        seqs.set(ubatch.seq_id_unq[s]);
+    }
+
+    // (seq_id, pos) -> token, for every cell that could be a predecessor of a ubatch token
+    std::unordered_map<uint64_t, llama_token> hist;
+
+    const auto key = [](llama_seq_id seq_id, llama_pos pos) {
+        return ((uint64_t) seq_id << 32) | (uint32_t) pos;
+    };
+
+    for (uint32_t s = 0; s < n_stream; ++s) {
+        v_cells[s].for_each_token_in(seqs, p_min - (llama_pos) n, p_max,
+            [&](llama_seq_id seq_id, llama_pos pos, llama_token tok) {
+                hist[key(seq_id, pos)] = tok;
+            });
+    }
+
+    for (uint32_t i = 0; i < n_tokens; ++i) {
+        // TODO: a token that belongs to more than one sequence has an ambiguous history.
+        //       the n-gram architectures have to reject such batches
+        const llama_seq_id seq_id = ubatch.seq_id[i][0];
+
+        for (uint32_t j = 0; j < n; ++j) {
+            const llama_pos p = ubatch.pos[i] - (llama_pos) (n - j);
+            if (p < 0) {
+                continue;
+            }
+
+            const auto it = hist.find(key(seq_id, p));
+            if (it != hist.end()) {
+                res[i*n + j] = it->second;
+            }
+        }
+    }
+}
+
 size_t llama_kv_cache::total_size() const {
    size_t size = 0;

@@ -2106,7 +2177,7 @@ void llama_kv_cache::state_write_meta(llama_io_write_i & io, const cell_ranges_t
            io.write(&pos,      sizeof(pos));
            io.write(&n_seq_id, sizeof(n_seq_id));

-            if (hparams.n_pos_per_embd() > 1) {
+            if (has_cell_ext()) {
                const llama_kv_cell_ext ext = cells.ext_get(i);
                io.write(&ext, sizeof(ext));
            }
@@ -2243,12 +2314,17 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
                return false;
            }

-            if (hparams.n_pos_per_embd() > 1) {
+            if (has_cell_ext()) {
                llama_kv_cell_ext ext;
                io.read(&ext, sizeof(ext));

-                ubatch.pos[i + ubatch.n_tokens]   = ext.y;
-                ubatch.pos[i + ubatch.n_tokens*2] = ext.x;
+                if (hparams.n_pos_per_embd() > 1) {
+                    ubatch.pos[i + ubatch.n_tokens]   = ext.y;
+                    ubatch.pos[i + ubatch.n_tokens*2] = ext.x;
+                }
+
+                // apply_ubatch() below restores ext.tok from the ubatch tokens
+                ubatch.token[i] = ext.tok;
            }

            // read the sequence id, but directly discard it - we will use dest_seq_id instead
@@ -2268,7 +2344,8 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
            return false;
        }

-        // TODO: we cannot yet restore llama_kv_cell_ext as the apply_ubatch() does not support it yet
+        // note: apply_ubatch() rebuilds llama_kv_cell_ext from the ubatch
+        //       only ext.tok and the M-RoPE 2D position round-trip through it
        //       see: https://github.com/ggml-org/llama.cpp/pull/16825#issuecomment-3460868350
        apply_ubatch(sinfo, ubatch);

@@ -2301,7 +2378,7 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32

            cells.pos_set(i, pos);

-            if (hparams.n_pos_per_embd() > 1) {
+            if (has_cell_ext()) {
                llama_kv_cell_ext ext;
                io.read(&ext, sizeof(ext));
                cells.ext_set(i, ext);
@@ -2652,3 +2729,7 @@ void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const {
 void llama_kv_cache_context::set_input_v_rot(ggml_tensor * dst) const {
    kv->set_input_v_rot(dst);
 }
+
+void llama_kv_cache_context::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const {
+    kv->get_prev_tokens(ubatch, n, res);
+}
@@ -219,6 +219,14 @@ public:
    void set_input_k_rot(ggml_tensor * dst) const;
    void set_input_v_rot(ggml_tensor * dst) const;

+    // true if llama_kv_cell_ext holds information that has to survive a state save/restore
+    bool has_cell_ext() const;
+
+    // for every token of the ubatch, the ids of the n tokens that precede it in its sequence
+    // entries with no matching cell are set to LLAMA_TOKEN_NULL
+    // note: used by n-gram input embeddings
+    void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const;
+
 private:
    const llama_model & model;
    const llama_hparams & hparams;
@@ -401,6 +409,9 @@ public:
    void set_input_k_rot(ggml_tensor * dst) const;
    void set_input_v_rot(ggml_tensor * dst) const;

+    // see llama_kv_cache::get_prev_tokens()
+    void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const;
+
 private:
    llama_memory_status status;

@@ -15,6 +15,10 @@ struct llama_kv_cell_ext {
    llama_pos x = 0;
    llama_pos y = 0;

+    // when tok = LLAMA_TOKEN_NULL when the cell is produced by embedding input (i.e. multimodal)
+    // use case: n-gram embeddings hash
+    llama_token tok = LLAMA_TOKEN_NULL;
+
    // return true if the current 2D spatial position is greater than other
    bool is_2d_gt(llama_pos ox, llama_pos oy) const {
        return (y > oy) || (y == oy && x > ox);
@@ -23,7 +27,7 @@ struct llama_kv_cell_ext {
    void reset() {
        static_assert(std::is_trivially_copyable_v<llama_kv_cell_ext>);

-        memset(this, 0, sizeof(*this));
+        *this = llama_kv_cell_ext{};
    }
 };

@@ -305,6 +309,29 @@ public:
        return seq[i].test(seq_id);
    }

+    // gather the token ids of the cells in `seqs` with position in [p0, p1)
+    // the callback receives (seq_id, pos, token) for every such (cell, seq) pair
+    // note: used by n-gram input embeddings to recover the tokens preceding a ubatch
+    template<typename F>
+    void for_each_token_in(const std::bitset<LLAMA_MAX_SEQ> & seqs, llama_pos p0, llama_pos p1, F && f) const {
+        for (const auto & i : used) {
+            if (pos[i] < p0 || pos[i] >= p1) {
+                continue;
+            }
+
+            const auto m = seq[i] & seqs;
+            if (m.none()) {
+                continue;
+            }
+
+            for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
+                if (m.test(s)) {
+                    f(s, pos[i], ext[i].tok);
+                }
+            }
+        }
+    }
+
    // note: call only if the cell is not empty and the seq_id is not in the cell
    void seq_add(uint32_t i, llama_seq_id seq_id) {
        assert(i < pos.size());
@@ -103,6 +103,7 @@ llama_model_nanbeige::graph::graph(const llama_model & model, const llm_graph_pa
    ggml_tensor * inp_out_ids = build_inp_out_ids();

    for (int il = 0; il < n_layer; ++il) {
+        res->t_layer_inp[il] = inpL;
        ggml_tensor * inpSA = inpL;

        cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
@@ -2469,8 +2469,13 @@ struct test_set_rows : public test_case {
    // See dicussion here: https://github.com/ggml-org/llama.cpp/pull/23760#issuecomment-4566312209
    double max_nmse_err(ggml_backend_t backend) override {
        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
-        if (type_dst == GGML_TYPE_Q8_0 && strcmp(ggml_backend_reg_name(reg), "WebGPU") == 0) {
-            return std::max(test_case::max_nmse_err(backend), 2e-7);
+        if (type_dst == GGML_TYPE_Q8_0) {
+            if (strcmp(ggml_backend_reg_name(reg), "WebGPU") == 0) {
+                return std::max(test_case::max_nmse_err(backend), 2e-7);
+            }
+            if (strcmp(ggml_backend_reg_name(reg), "HTP") == 0) {
+                return std::max(test_case::max_nmse_err(backend), 5e-6);
+            }
        }
        return test_case::max_nmse_err(backend);
    }
@@ -4120,9 +4125,10 @@ struct test_ssm_scan : public test_case {
    const int64_t n_seqs;
    const bool    xbc_overlap;
    const int64_t K;
+    const bool    weak_decay;

    std::string vars() override {
-        return VARS_TO_STR9(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap, K);
+        return VARS_TO_STR10(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap, K, weak_decay);
    }

    test_ssm_scan(ggml_type type = GGML_TYPE_F32,
@@ -4133,8 +4139,9 @@ struct test_ssm_scan : public test_case {
            int64_t n_seq_tokens = 32,
            int64_t n_seqs = 32,
            bool xbc_overlap = false,
-            int64_t K = 1)
-        : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap), K(K) {}
+            int64_t K = 1,
+            bool weak_decay = false)
+        : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap), K(K), weak_decay(weak_decay) {}

    double max_nmse_err() override {
        // SSD path (head_dim > 1) uses FP16 intermediates (M matrix, X_dt); Mamba-1 is pure FP32.
@@ -4187,7 +4194,7 @@ struct test_ssm_scan : public test_case {
                continue;
            } else if (t->ne[1] == n_head && t->ne[2] == 1) {
                // A {1 or d_state, n_head}: negative decay (2-D tensor, ne[2]==1 distinguishes from 3-D/4-D tensors)
-                init_tensor_uniform(t, -1.0f, -0.5f);
+                init_tensor_uniform(t, weak_decay ? -0.02f : -1.0f, weak_decay ? -0.005f : -0.5f);
            } else {
                init_tensor_uniform(t);
            }
@@ -9111,6 +9118,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 4, 2, false, /*K=*/4)); // Mamba-2 rollback snapshots
    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 8, 2, false, /*K=*/3)); // Mamba-2 rollback overflow
    test_cases.emplace_back(new test_ssm_scan_rollback(GGML_TYPE_F32, 128, 64, 16, 2, 8, 2, /*K=*/3)); // rollback snapshots match prefix states
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 64, 4)); // Metal SSD one chunk MMA only, no seq tail
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 65, 2)); // SSD one chunk + 1-token sequential tail
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 128, 2)); // SSD multi-chunk, no tail (exercises the chunk-to-chunk state handoff)
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 128, 2, false, /*K=*/1, /*weak_decay=*/true)); // SSD multi-chunk, carried state not numerically negligible

    test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 1, 1));
    test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 1));
@@ -8,7 +8,8 @@
 		ChatFormInputFileInputInvisible,
 		ChatFormMcpResourcesList,
 		ChatFormPickers,
-		DialogMcpResourcesBrowser
+		DialogMcpResourcesBrowser,
+		DialogMcpServers
 	} from '$lib/components/app';
 	import {
 		CLIPBOARD_CONTENT_QUOTE_PREFIX,
@@ -183,6 +184,9 @@
 	let isResourceDialogOpen = $state(false);
 	let preSelectedResourceUri = $state<string | undefined>(undefined);

+	// MCP Servers Dialog State
+	let isMcpServersDialogOpen = $state(false);
+
 	let currentConfig = $derived(settingsStore.config);

 	let pasteLongTextToFileLength = $derived.by(() => {
@@ -618,6 +622,7 @@
 				onFileUpload={handleFileUpload}
 				onMcpPromptClick={showMcpPromptButton ? () => pickers.openPromptPicker() : undefined}
 				onMcpResourcesClick={() => (isResourceDialogOpen = true)}
+				onMcpSettingsClick={() => (isMcpServersDialogOpen = true)}
 				onMicClick={handleMicClick}
 				{onStop}
 				onSystemPromptClick={() => onSystemPromptClick?.({ files: uploadedFiles, message: value })}
@@ -656,3 +661,5 @@
 	}}
 	preSelectedUri={preSelectedResourceUri}
 />
+
+<DialogMcpServers bind:open={isMcpServersDialogOpen} />
@@ -1,10 +1,6 @@
 <script lang="ts">
-	import { File, FolderOpen, MessageSquare, Plus, Zap } from '@lucide/svelte';
-	import {
-		ChatFormActionAddMcpServersSubmenu,
-		ChatFormActionAddReasoningSubmenu,
-		ChatFormActionAddToolsSubmenu
-	} from '$lib/components/app';
+	import { File, MessageSquare, Plus } from '@lucide/svelte';
+	import { ChatFormActionAddToolsSubmenu, McpLogo } from '$lib/components/app';
 	import { buttonVariants } from '$lib/components/ui/button';
 	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
 	import * as Tooltip from '$lib/components/ui/tooltip';
@@ -31,11 +27,6 @@
 	// must not restore focus to the trigger on close
 	let suppressCloseAutoFocus = false;

-	function handleMcpSettingsClick() {
-		dropdownOpen = false;
-		chatFormActions.onMcpSettingsClick?.();
-	}
-
 	const attachmentMenu = useAttachmentMenu(
 		() => ({
 			hasAudioModality: chatFormActions.hasAudioModality,
@@ -93,10 +84,6 @@
 				}
 			}}
 		>
-			<ChatFormActionAddReasoningSubmenu />
-
-			<DropdownMenu.Separator />
-
 			<DropdownMenu.Sub>
 				<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
 					<File class={ICON_CLASS_DEFAULT} />
@@ -156,31 +143,14 @@

 			<ChatFormActionAddToolsSubmenu />

-			<ChatFormActionAddMcpServersSubmenu onMcpSettingsClick={handleMcpSettingsClick} />
+			<DropdownMenu.Item
+				class="flex cursor-pointer items-center gap-2"
+				onclick={chatFormActions.onMcpSettingsClick}
+			>
+				<McpLogo class={ICON_CLASS_DEFAULT} />

-			{#if chatFormActions.hasMcpPromptsSupport}
-				<DropdownMenu.Separator />
-
-				<DropdownMenu.Item
-					class="flex cursor-pointer items-center gap-2"
-					onclick={chatFormActions.onMcpPromptClick}
-				>
-					<Zap class={ICON_CLASS_DEFAULT} />
-
-					<span>MCP Prompt</span>
-				</DropdownMenu.Item>
-			{/if}
-
-			{#if chatFormActions.hasMcpResourcesSupport}
-				<DropdownMenu.Item
-					class="flex cursor-pointer items-center gap-2"
-					onclick={chatFormActions.onMcpResourcesClick}
-				>
-					<FolderOpen class={ICON_CLASS_DEFAULT} />
-
-					<span>MCP Resources</span>
-				</DropdownMenu.Item>
-			{/if}
+				<span>MCP Servers</span>
+			</DropdownMenu.Item>
 		</DropdownMenu.Content>
 	</DropdownMenu.Root>
 </div>
@@ -1,152 +0,0 @@
-<script lang="ts">
-	import { Plus, Settings } from '@lucide/svelte';
-	import { goto } from '$app/navigation';
-	import { DropdownMenuSearchable, McpLogo, McpServerIdentity } from '$lib/components/app';
-	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
-	import { Switch } from '$lib/components/ui/switch';
-	import { ICON_CLASS_DEFAULT, ROUTES } from '$lib/constants';
-	import { HealthCheckStatus } from '$lib/enums';
-	import { conversationsStore, mcpStore } from '$lib/stores';
-	import type { MCPServerSettingsEntry } from '$lib/types';
-
-	interface Props {
-		onMcpSettingsClick?: () => void;
-	}
-
-	let { onMcpSettingsClick }: Props = $props();
-
-	let mcpSearchQuery = $state('');
-	// Every configured server is listed; `enabled` is an on/off state,
-	// not a visibility filter, so a disabled server stays toggleable.
-	let mcpServers = $derived(mcpStore.getServers());
-	let hasMcpServers = $derived(mcpServers.length > 0);
-	let filteredMcpServers = $derived.by(() => {
-		const query = mcpSearchQuery.toLowerCase().trim();
-
-		if (!query) return mcpServers;
-
-		return mcpServers.filter((s) => {
-			const name = getServerLabel(s).toLowerCase();
-			const url = s.url.toLowerCase();
-
-			return name.includes(query) || url.includes(query);
-		});
-	});
-
-	function getServerLabel(server: MCPServerSettingsEntry): string {
-		return mcpStore.getServerLabel(server);
-	}
-
-	function isServerEnabledForChat(serverId: string): boolean {
-		return conversationsStore.preferences.isMcpServerEnabledForChat(serverId);
-	}
-
-	async function toggleServerForChat(serverId: string) {
-		await conversationsStore.preferences.toggleMcpServerForChat(serverId);
-	}
-
-	function handleMcpSubMenuOpen(open: boolean) {
-		if (open) {
-			mcpSearchQuery = '';
-			mcpStore.runHealthChecksForServers(mcpServers);
-		}
-	}
-
-	function handleMcpSettingsClick() {
-		onMcpSettingsClick?.();
-
-		goto(`${hasMcpServers ? '' : '?add'}${ROUTES.MCP_SERVERS}`);
-	}
-</script>
-
-<DropdownMenu.Root>
-	<DropdownMenu.Sub onOpenChange={handleMcpSubMenuOpen}>
-		<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
-			<McpLogo class={ICON_CLASS_DEFAULT} />
-
-			<span>MCP Servers</span>
-		</DropdownMenu.SubTrigger>
-
-		<DropdownMenu.SubContent class="w-72 pt-0">
-			{#if hasMcpServers}
-				<DropdownMenuSearchable
-					bind:searchValue={mcpSearchQuery}
-					emptyMessage="No servers found"
-					isEmpty={filteredMcpServers.length === 0}
-					placeholder="Search servers..."
-				>
-					<div class="max-h-64 overflow-y-auto">
-						{#each filteredMcpServers as server (server.id)}
-							{@const healthState = mcpStore.getHealthCheckState(server.id)}
-							{@const hasError = healthState.status === HealthCheckStatus.ERROR}
-							{@const isEnabledForChat = isServerEnabledForChat(server.id)}
-							{@const displayName = getServerLabel(server)}
-							{@const faviconUrl = mcpStore.getServerFavicon(server.id)}
-
-							<button
-								class="flex w-full items-center justify-between gap-2 rounded-sm px-2 py-2 text-left transition-colors hover:bg-accent disabled:cursor-not-allowed disabled:opacity-50"
-								disabled={hasError}
-								onclick={() => !hasError && toggleServerForChat(server.id)}
-								type="button"
-							>
-								<div class="flex min-w-0 flex-1 items-center gap-2">
-									<div class="min-w-0 flex-1">
-										<McpServerIdentity
-											{displayName}
-											{faviconUrl}
-											iconClass={ICON_CLASS_DEFAULT}
-											iconRounded="rounded-sm"
-											nameClass="text-sm"
-											showVersion={false}
-										/>
-									</div>
-
-									{#if hasError}
-										<span
-											class="shrink-0 rounded bg-destructive/15 px-1.5 py-0.5 text-xs text-destructive"
-										>
-											Error
-										</span>
-									{/if}
-								</div>
-
-								<Switch
-									checked={isEnabledForChat}
-									disabled={hasError}
-									onCheckedChange={() => toggleServerForChat(server.id)}
-									onclick={(e) => e.stopPropagation()}
-								/>
-							</button>
-						{/each}
-					</div>
-
-					{#snippet footer()}
-						<DropdownMenu.Item
-							class="flex cursor-pointer items-center gap-2"
-							onclick={handleMcpSettingsClick}
-						>
-							<Settings class={ICON_CLASS_DEFAULT} />
-
-							<span>Manage MCP Servers</span>
-						</DropdownMenu.Item>
-					{/snippet}
-				</DropdownMenuSearchable>
-			{:else}
-				<div class="px-2 py-3 text-center text-sm text-muted-foreground">
-					No MCP servers configured
-				</div>
-
-				<DropdownMenu.Separator />
-
-				<DropdownMenu.Item
-					class="flex cursor-pointer items-center gap-2"
-					onclick={handleMcpSettingsClick}
-				>
-					<Plus class={ICON_CLASS_DEFAULT} />
-
-					<span>Add MCP Servers</span>
-				</DropdownMenu.Item>
-			{/if}
-		</DropdownMenu.SubContent>
-	</DropdownMenu.Sub>
-</DropdownMenu.Root>
@@ -0,0 +1,51 @@
+<script lang="ts">
+	import { FolderOpen, Server, Zap } from '@lucide/svelte';
+	import { McpLogo } from '$lib/components/app';
+	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants';
+	import { getChatFormActionsContext } from '$lib/contexts';
+
+	const chatFormActions = getChatFormActionsContext();
+
+	function handleServersClick() {
+		chatFormActions.onMcpSettingsClick?.();
+	}
+</script>
+
+<DropdownMenu.Sub>
+	<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
+		<McpLogo class={ICON_CLASS_DEFAULT} />
+
+		<span>MCP</span>
+	</DropdownMenu.SubTrigger>
+
+	<DropdownMenu.SubContent class="w-48">
+		<DropdownMenu.Item class="flex cursor-pointer items-center gap-2" onclick={handleServersClick}>
+			<Server class={ICON_CLASS_DEFAULT} />
+
+			<span>Servers</span>
+		</DropdownMenu.Item>
+
+		{#if chatFormActions.hasMcpPromptsSupport}
+			<DropdownMenu.Item
+				class="flex cursor-pointer items-center gap-2"
+				onclick={chatFormActions.onMcpPromptClick}
+			>
+				<Zap class={ICON_CLASS_DEFAULT} />
+
+				<span>Prompts</span>
+			</DropdownMenu.Item>
+		{/if}
+
+		{#if chatFormActions.hasMcpResourcesSupport}
+			<DropdownMenu.Item
+				class="flex cursor-pointer items-center gap-2"
+				onclick={chatFormActions.onMcpResourcesClick}
+			>
+				<FolderOpen class={ICON_CLASS_DEFAULT} />
+
+				<span>Resources</span>
+			</DropdownMenu.Item>
+		{/if}
+	</DropdownMenu.SubContent>
+</DropdownMenu.Sub>
@@ -1,6 +1,5 @@
 <script lang="ts">
 	import { SkipForward, Square } from '@lucide/svelte';
-	import { goto } from '$app/navigation';
 	import { page } from '$app/state';
 	import {
 		ChatFormActionModels,
@@ -10,7 +9,7 @@
 		ChatFormContextGauge
 	} from '$lib/components/app';
 	import { Button } from '$lib/components/ui/button';
-	import { ICON_CLASS_DEFAULT, ROUTES } from '$lib/constants';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants';
 	import { setChatFormActionsContext } from '$lib/contexts';
 	import { FileTypeCategory, MessageRole } from '$lib/enums';
 	import { ChatService } from '$lib/services';
@@ -34,6 +33,7 @@
 		onSystemPromptClick?: () => void;
 		onMcpPromptClick?: () => void;
 		onMcpResourcesClick?: () => void;
+		onMcpSettingsClick?: () => void;
 	}

 	let {
@@ -47,6 +47,7 @@
 		onFileUpload,
 		onMcpPromptClick,
 		onMcpResourcesClick,
+		onMcpSettingsClick,
 		onMicClick,
 		onStop,
 		onSystemPromptClick,
@@ -163,7 +164,7 @@
 			return onMcpResourcesClick;
 		},
 		get onMcpSettingsClick() {
-			return () => goto(ROUTES.MCP_SERVERS);
+			return onMcpSettingsClick;
 		},
 		get onSystemPromptClick() {
 			return onSystemPromptClick;
@@ -221,25 +221,17 @@ export { default as ChatFormActionModels } from './ChatForm/ChatFormActions/Chat
 export { default as ChatFormActionAddToolsSubmenu } from './ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte';

 /**
- * Dropdown submenu for managing MCP servers in the chat form.
+ * Dropdown submenu for MCP prompts and resources in the chat form.
 *
- * Displays a searchable list of enabled MCP servers with toggle switches
- * to enable/disable each server for chat. Shows server favicon, health status,
- * and a "Manage MCP Servers" settings link.
- *
- * Features:
- * - Search/filter servers by name or URL
- * - Per-server toggle to enable/disable for chat
- * - Health check indicator (shows "Error" badge for failed servers)
- * - Server favicon display
- * - Settings link to manage MCP server configuration
+ * Shows an "MCP" sub-menu item with entries for MCP Prompts and MCP
+ * Resources. Only visible when the server supports them.
 *
 * @example
 * ```svelte
- * <ChatFormActionAddMcpServersSubmenu onMcpSettingsClick={handleMcpSettingsClick} />
+ * <ChatFormActionAddMcpSubmenu />
 * ```
 */
-export { default as ChatFormActionAddMcpServersSubmenu } from './ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpServersSubmenu.svelte';
+export { default as ChatFormActionAddMcpSubmenu } from './ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpSubmenu.svelte';

 /**
 * Dropdown submenu for selecting reasoning effort level.
@@ -253,7 +253,7 @@
 </script>

 <Dialog.Root onOpenChange={handleOpenChange} {open}>
-	<Dialog.Content class="max-h-[80vh] !max-w-4xl overflow-hidden p-0">
+	<Dialog.Content class="max-h-[80vh] md:max-w-4xl! w-full! overflow-hidden p-0">
 		<Dialog.Header class="border-b border-border/30 px-6 py-4">
 			<Dialog.Title class="flex items-center gap-2">
 				<FolderOpen class="h-5 w-5" />
@@ -246,7 +246,7 @@
 </script>

 <Dialog.Root onOpenChange={handleOpenChange} {open}>
-	<Dialog.Content class="sm:max-w-2xl">
+	<Dialog.Content class="max-w-2xl!">
 		<Dialog.Header>
 			<Dialog.Title class="select-none">Add New MCP Server</Dialog.Title>
 		</Dialog.Header>
@@ -0,0 +1,33 @@
+<script lang="ts">
+	import { McpLogo } from '$lib/components/app';
+	import { SettingsMcpServers } from '$lib/components/app/settings';
+	import * as Dialog from '$lib/components/ui/dialog';
+
+	interface Props {
+		open?: boolean;
+		onOpenChange?: (open: boolean) => void;
+	}
+
+	let { onOpenChange, open = $bindable(false) }: Props = $props();
+
+	function handleOpenChange(value: boolean) {
+		open = value;
+		onOpenChange?.(value);
+	}
+</script>
+
+<Dialog.Root onOpenChange={handleOpenChange} {open}>
+	<Dialog.Content
+		class="md:h-[calc(100vh-4rem)]! md:max-h-240! md:w-[calc(100vw-4rem)]! md:max-w-360! flex flex-col"
+	>
+		<Dialog.Header>
+			<Dialog.Title class="flex items-center gap-2">
+				<McpLogo class="h-5 w-5" />
+
+				<span>MCP Servers</span>
+			</Dialog.Title>
+		</Dialog.Header>
+
+		<SettingsMcpServers class="mt-4" />
+	</Dialog.Content>
+</Dialog.Root>
@@ -14,6 +14,7 @@
 <Dialog.Root bind:open {onOpenChange}>
 	<Dialog.Content
 		class="z-999999 grid max-h-full max-w-full! grid-rows-[1fr_auto] overflow-hidden p-0 md:h-[90vh] md:max-w-[90vw]!"
+		showCloseButton
 	>
 		<MermaidPreview {svgHtml} />
 	</Dialog.Content>
@@ -0,0 +1,34 @@
+<script lang="ts">
+	import { Settings } from '@lucide/svelte';
+	import { SettingsChat } from '$lib/components/app/settings';
+	import * as Dialog from '$lib/components/ui/dialog';
+
+	interface Props {
+		open?: boolean;
+		onOpenChange?: (open: boolean) => void;
+		initialSection?: string;
+	}
+
+	let { initialSection, onOpenChange, open = $bindable(false) }: Props = $props();
+
+	function handleOpenChange(value: boolean) {
+		open = value;
+		onOpenChange?.(value);
+	}
+</script>
+
+<Dialog.Root onOpenChange={handleOpenChange} {open}>
+	<Dialog.Content
+		class="md:h-[calc(100vh-4rem)]! md:max-h-240! md:w-[calc(100vw-4rem)]! md:max-w-6xl! flex flex-col p-0 md:p-6 gap-0"
+	>
+		<Dialog.Header class="md:p-0 p-4">
+			<Dialog.Title class="flex items-center gap-2">
+				<Settings class="h-5 w-5" />
+
+				<span>Settings</span>
+			</Dialog.Title>
+		</Dialog.Header>
+
+		<SettingsChat {initialSection} onClose={() => (open = false)} onSectionChange={() => {}} />
+	</Dialog.Content>
+</Dialog.Root>
@@ -18,6 +18,23 @@
 */
 export { default as DialogMcpServerAddNew } from './DialogMcpServerAddNew.svelte';

+/**
+ * **DialogMcpServers** - MCP servers dialog shown from the chat form
+ *
+ * Shows the same MCP server list as the `/mcp-servers` route inside a modal
+ * dialog.
+ */
+export { default as DialogMcpServers } from './DialogMcpServers.svelte';
+
+/**
+ * **DialogSettingsChat** - Chat settings shown in a modal dialog
+ *
+ * Wraps the full SettingsChat layout (sidebar, mobile header, fields, footer)
+ * inside a ShadCN Dialog instead of a dedicated route. Section switching is
+ * handled in-app via `onSectionChange` rather than URL navigation.
+ */
+export { default as DialogSettingsChat } from './DialogSettingsChat.svelte';
+
 /**
 * **DialogExportSettings** - Settings export dialog with sensitive data warning
 *
@@ -1,13 +1,22 @@
 <script lang="ts">
-	import { Database, FileText, ListChecks, MessageSquare, Sparkles, Wrench } from '@lucide/svelte';
+	import {
+		Database,
+		ExternalLink,
+		FileText,
+		ListChecks,
+		MessageSquare,
+		Sparkles,
+		Wrench
+	} from '@lucide/svelte';
 	import { Badge } from '$lib/components/ui/badge';
 	import type { MCPCapabilitiesInfo } from '$lib/types';

 	interface Props {
 		capabilities?: MCPCapabilitiesInfo;
+		onBrowseResources?: () => void;
 	}

-	let { capabilities }: Props = $props();
+	let { capabilities, onBrowseResources }: Props = $props();
 </script>

 {#if capabilities}
@@ -20,10 +29,24 @@
 	{/if}

 	{#if capabilities.server.resources}
-		<Badge class="h-5 gap-1 bg-blue-50 px-1.5 text-[10px] dark:bg-blue-950" variant="outline">
+		<Badge
+			class="h-5 cursor-pointer gap-1 bg-blue-50 px-1.5 text-[10px] transition-colors hover:bg-blue-100 dark:bg-blue-950 dark:hover:bg-blue-900"
+			onclick={onBrowseResources}
+			onkeydown={(e) => {
+				if (e.key === 'Enter' || e.key === ' ') {
+					e.preventDefault();
+					onBrowseResources?.();
+				}
+			}}
+			role="button"
+			tabindex={0}
+			variant="outline"
+		>
 			<Database class="h-3 w-3 text-blue-600 dark:text-blue-400" />

 			Resources
+
+			<ExternalLink class="h-3 w-3 text-blue-600 dark:text-blue-400" />
 		</Badge>
 	{/if}

--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
zql GitHub admin	d7a2074112	models : support nanbeige4.2-3B (#27730 ) Co-authored-by: admin <lizongqiang@kanzhun.com>	2026-08-27 07:55:31 +03:00
Max KrasnyanskyandGitHub	192067b72d	hexagon: support for multi-NPU devices (IQ9, IQ10) and fully asynchronous backend (#26501 ) * hexagon: use non-host bufs by default and make the backend fully async * hex-hb: remove optional hostbuf support and fix async copy * hex-unary: relax supported unary check * hex-bufs: use same get_alignment for host bufs * snapdragon: bump android_platform to 34 * hex-rows: super hacky get/set rows for q8_0 * hex-get-rows: fix q8_0 * hex-get-rows: supprot for f16 and cleanup for q8_0 * hex-get-rows: generic macros and specialized thread funcs * hex-get-rows: add DMA pipeline, vtcm_layout and kernel params * hex-set-rows: fix q8_0 support, add dma and tracing * hex-tests: override nmse threshold for HTP of Q8_0 quants * hex-fa: add support for Q8_0 with inplace dequantizers * hex-get-rows: simplify type dispatch * hex-rows: simplify GET/SET_ROWS DMA pipeline * hex-async: add events, set/get-tensor-async and rest of the async api support * hex-repack: use slice instead of expert in repack functions * hex-cpy: update event/async-cpy logging * hex-set-rows: optimize smaller tensors * hex-geglu: fix perf regression with larger tensors * hex-get-rows: add missing header * hex-set-rows: add missing header * hex-bufs: ressurect GGML_HEXAGON_HOSTBUF but disable it by default * hexagon: do not reject ops with non-heaxon buffers * hex-get-rows: apply >=32 restriction only for q8_0 * hex-res: bump vtcm acquire timeout to 10 seconds * hex-bufs: add support for cloning buffers between sessions to speed up tensor copies * hex-async: rework event recording and batch flushing and integrate with meta backend * hex-bufs: improved handling of repacked tensors * hex-repack: handle get_tensor_2d offsets * hex-dev: add support for devices with multiple NPUs * hex-sync: add support for sync tokens to synchronize npu devices for async splits * hex-mmap: cleanup mmap calls and add a retry for robustness * hex-sync: add failsafe if sync wait gets stuck * hex-sync: use sync_seq to check for completed events * hex-sync: rotate tokens for extra robustness * hex-devs: add supprot for legacy device names for now * hex-bufs: add support for auto-cloning buffers from diff sessions * hex-fusion: simplify and optimize htp-opnode fusion handling * hex-sync: override opnode name so that it shows up in the profiles * hex-trace: update scripts to handle multiple devices * hex-sync: bump the size of the opbatch queue and number of sync tokens * hex-cpy-sync: do not explicitly flush opbatches in cpy_tensor_async and add support for cpy-dma * hex-sync: add graph-flush threshold to avoid single op batches * hex-sync: add sync_peer so that we can flush peers we depend on during cross-device ops * hex-bufs: introduce tensor->extra and shadow_bufs for repacking * hex-l2: flush tiny tensors inline * hex-sync: use explicit l2flush for sync tokens * hex-extra: track weight flags via tensor extra * hex-fence: rename sync to fence * hex-repack: proper handling of set-tensor-2d in the shadow_buf * hex-trace: remove obsolete opstage mask that we used for profiling * hex-env: remove obsolete use_hmx variable * hexagon: new unified run.py and build.py and updated docs * snapdragon: update run script to auto-escapt test-backend-op -p argument * hex-scripts: fix trailing spaces * hex-scripts: fix flake8 warnings * snapdragon: cleanup dst lib/bin dirs before copying new build * hex-ops: add support for allreduce * hex-ar: improved allreduce with dma pipeline * hex-ar: align macros * hex-ar: consistent use of fence_seq * hex-ar: add AR_SELECT env var to select ALLREDUCE kernel or fallback * hex-ar: add proper synchronize handling for ALLREDUCE * hex-opbatch: looks like we now just rely on backend.synchronise to flush the batches, no need to flush them by threshold * hex-ar: bump block size to improve dma efficiency * hex-ar: fused ALLREDUCE+ADD * hex-ar: cleaner fence buffer management * hex-ar: futher allreduce tweaking to remove race conditions * hex-ar: add simple solver and remove non-dma kernels * hex-ar: add row-broadcast to fuse with bias ADD * hex-fence: pass seq numbers via op_params * hex-ar: allow for both entry/exit seq for completing entry wait * hex-ar: align macros * hex-ar: do not refetch broadcast row * hex-fusion: move all fusion into opbatch::add_op for consistency with ALLREDUCE and things * hex-fusion: fix incorrect MUL_MAT reordering * hex-mm: make fused 2x and 3x matmuls more generic * hex-fusion: move tensor fusion tagging to graph_compute * hexagon: make sure to copy tensor->extra by value * hex-get-rows: fix offset calc with row-chunking * hex-repack: get_tensor_2d fixes for non-zero offsets * snapdragon: make profile/trace scripts more robust and donot mix stdout/stderr by default * hex-devices: use legacy device nameing by default to ease the transition * hex-devices: hardcode CDSP domain IDs for current devices for now * hex-optrace: improve multi-NPU timestamp alignment and overall handling of cycle values * hex-optrace: more robust handling of the fence events	2026-08-26 18:46:50 -07:00
Xuan-Son NguyenandGitHub	925e117994	llama: add token ID tracking to KV cell (#27762 ) * kv: track token id * rm get_prev_tokens, move it to the main pr * nits * add get_prev_tokens	2026-08-26 23:34:28 +02:00
Aleksander GrygierandGitHub	539f24529b	ui: Move Settings and MCP Servers routes to dialog-based views (#27744 ) * ui : open MCP servers in a dialog from the chat form Replace the MCP servers submenu with a single "MCP Servers" item that opens a new DialogMcpServers dialog instead of navigating to the /mcp-servers route. Assisted-by: pi * ui : browse MCP resources from the server card Make the Resources capability badge clickable so it opens the MCP resources browser dialog, and drop the page-only chrome from SettingsMcpServers. Assisted-by: pi * ui : remove mcp-servers route and sidebar entry MCP servers are now managed in a dialog, so drop the dedicated route and the sidebar icon that navigated to it. Assisted-by: pi * ui : remove unused MCP servers submenu component The submenu was replaced by the MCP servers dialog, so delete the component and its export. Assisted-by: pi * feat(ui): add DialogSettingsChat dialog * refactor(ui): switch SettingsChat to in-app section navigation * feat(ui): open settings as dialog from sidebar * refactor(ui): remove settings route and URL-based settings navigation * fix(ui): adjust MCP dialogs for new base sizing * chore: Formatting & linting	2026-08-26 21:07:24 +02:00
Aleksander GrygierandGitHub	0379a19f09	ui: Update Dialog component styling (#27743 ) * feat(ui): make base dialog responsive and support sticky headers * ui: move dialog close button to the sticky header Assisted-by: pi * chore: Formatting & linting	2026-08-26 20:19:19 +02:00
Ruben OrtlamandGitHub	5e6a37cb11	vulkan: warptiles currently assume warp sizes <= 64, clamp to work around larger warps (#27726 )	2026-08-26 19:02:06 +03:00
Pranav UttarkarandGitHub	bf94216469	Implemented vulkan cross_entropy_loss and cross_entropy_loss_back (#27216 )	2026-08-26 16:49:32 +02:00
Radoslav GerganovandGitHub	d0132a680a	rpc : implement event and async backend APIs (#18626 ) * rpc : implement event and async backend APIs * cache responses from RPC_CMD_GET_ALLOC_SIZE	2026-08-26 17:34:46 +03:00
Aleksander GrygierandGitHub	4d19b28769	ci: Clean up UI builds from releases (#27706 ) * ci : inline UI version resolution into ui-build.yml * ci : build UI once and reuse the artifact in release jobs Server jobs now extract the ui-build artifact into tools/ui/dist instead of npm-building the UI. Also removes the get-version job and the no-op -DHF_UI_VERSION flags. Assisted-by: pi:Kimi-K3 * ui : disable the npm UI build by default (LLAMA_BUILD_UI=OFF) The flag now only controls building the UI from source via npm. The UI is still embedded by default from local tools/ui/dist or the prebuilt download (LLAMA_USE_PREBUILT_UI=ON). CI jobs no longer npm-build the UI; server-sanitize does not need node anymore. Assisted-by: pi:Kimi-K3 * ci : rename the ui-build artifact to llama-ui.zip Consistent with the other artifact names in the Actions summary. Assisted-by: pi:Kimi-K3 * ci : clarify the windows artifact merge in release.yml The windows-cuda/vulkan/sycl jobs build only the backend library; llama-server (with the embedded UI) is injected into their zips from the windows-cpu package during the release. State this in the job comments and use accurate wording in the merge step. Assisted-by: pi:Kimi-K3	2026-08-26 14:12:09 +02:00
David FriehsandGitHub	fc35562ba4	cuda: unblock mmq for MoE on sm_60 (#26264 ) * cuda: unblock mmq for MoE on sm_60 * cuda: duplicate mmq-config-pascal for dp4a and older * cuda: reduce occupancy on non-dp4a pascal for Q2_K, Q4_K, Q5_K, Q6_K	2026-08-26 18:35:54 +08:00
Sigbjørn SkjæretandGitHub	da9b5d68c3	ci : make cache bucket public (#27728 ) * check for hf token * make bucket public	2026-08-26 12:08:23 +02:00
Daniel BeveniusandGitHub	dac869b0a0	conversion : fix Nemotron 3.5 Lightning layers (#27729 ) This commit contains a fix for the conversion of NVIDIA Nemotron 3.5 Lightning which currently incorrectly converts when using a transformers version later than 5.5.1. When converting using [convert](https://github.com/ggml-org/convert) the transformers version is 5.13.1 and this produces the following: ```console WARNING:gguf.gguf_writer:Duplicated key name 'nemotron_h_moe.attention.head_count_kv', overwriting it with new value [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] of type ARRAY ``` This does not happen with transformers 5.5.1. The reason seems to be that the configuration is different in later versions, for example when using 5.13.1 the configuration block looks like this: ```console transformers 5.13.1 raw has layers_block_type: True autoconfig has layers_block_type: True autoconfig layers_block_type: [ 'linear_attention', 'moe', 'linear_attention', 'moe', 'linear_attention', 'full_attention', 'moe', ... ] ``` And with 5.5.1 we get: ```console transformers 5.5.1 raw has layers_block_type: True autoconfig has layers_block_type: True autoconfig layers_block_type: [ 'mamba', 'moe', 'mamba', 'moe', 'mamba', 'attention', 'moe' ... ] ``` In our conversion script we only match for attention, not full attention which is causing this issue. With the changes in this commit the output with transformers 5.13.1 will be: ```console (venv) $ gguf-dump models/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16.gguf \| grep head_count_kv INFO:gguf-dump:* Loading: models/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16.gguf 29: [INT32] \| 52 \| nemotron_h_moe.attention.head_count_kv = [0, 0, 0, 0, 0, 2, ...] ``` Resolves: https://github.com/ggml-org/llama.cpp/issues/27718 Refs: https://github.com/ggml-org/convert/actions/runs/32949047680/job/98116096069#step:5:2391	2026-08-26 12:05:31 +02:00
Dominik Pantaleoni GitHub dpantaleoni forforever73	11cd988428	ggml-metal: add chunked SSD MMA for Mamba-2 prefill optimization (#26647 ) * metal: WIP chunked SSD SSM_SCAN kernels for multi-token prefill * metal: drop scalar SSD path; MMA + sequential tail * drop WIP ssm scan test noise * remove state_from_dst and rename CS and NSG constants * remove unrelated added whitespace padding * added clarity to mma_tokens calculation * added clarity to use_mma bool checks * added comments to metal ssd op constants for clarity * reserve K tokens for sequential kernel rollback snapshots * reset concurrency between mma and seq tail * remove print args no longer used * fixed comment to no longer point to specific line * add FC_SSM_SCAN so seq path skips token offlset unless it's mma tail * added changes to new ssm.metal for rebase after ggml-metal.metal refactor * specialize ssm_scan tail with a template instead of a function constant --------- Co-authored-by: dpantaleoni <dominikpantaleoni@gmail.com> Co-authored-by: forforever73 <690105611@qq.com>	2026-08-26 11:57:07 +03:00
Max KrasnyanskyandGitHub	5d5cb4c3a4	ggml-meta: propagate buffer usage and call init on the new tensors (#27586 )	2026-08-26 08:27:51 +03:00
Jonathan ClohessyandGitHub	d222767c7a	kleidiai: Rework KleidiAI Build System/Integration (#26077 ) * Rework KleidiAI Build System/Integration Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com> * Add fp16 guard, and fix cmake caching issue Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com> * Fix formatting, and rebase issue Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com> --------- Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>	2026-08-25 14:07:29 -07:00
Mario LimoncielloandGitHub	eab8ee41f8	ci : update OS used for ROCM to Ubuntu 24.04 (#27681 ) This matches what other build targets use and also what AMD advertises wheels as supporting.	2026-08-25 20:14:16 +03:00