model: add DSpark support for Nemotron3.5 (#27804 )

* model: add DSpark support for Nemotron3.5 * Update src/models/dflash.cpp Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> --------- Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
ggml-hexagon: add HTP unary ops for ABS and LOG (#27786 )
2026-08-28 02:57:42 +02:00 · 2026-08-28 01:49:27 +02:00 · 2026-08-27 15:05:57 -07:00 · 2026-08-27 14:38:02 -07:00 · 2026-08-27 22:39:14 +02:00 · 2026-08-27 22:18:34 +02:00
481 changed files with 16425 additions and 7493 deletions
@@ -90,6 +90,9 @@ RUN bash -c "source ${OpenVINO_DIR}/setupvars.sh && \
    cmake -B build/ReleaseOV -G Ninja \
        -DCMAKE_BUILD_TYPE=Release \
        -DLLAMA_BUILD_TESTS=OFF \
+        -DGGML_NATIVE=OFF \
+        -DGGML_BACKEND_DL=ON \
+        -DGGML_CPU_ALL_VARIANTS=ON \
        -DGGML_OPENVINO=ON && \
    cmake --build build/ReleaseOV --parallel "

@@ -0,0 +1,95 @@
+name: "ccache-buckets"
+description: "Save/restore latest GitHub Actions ccache matching a key prefix to/from HF buckets"
+inputs:
+  key:
+    description: "Cache key prefix to match and load"
+    required: true
+  folder:
+    description: "Bucket folder containing ccache files"
+    required: true
+  evict-old-files:
+    description: "Corresponds to the ccache --evict-older-than AGE option, where AGE is the number of seconds or days followed by the 's' or 'd' suffix respectively."
+    default: ''
+  save:
+    description: "Save ccache"
+    required: false
+    default: false
+    type: boolean
+  hf_bucket:
+    description: 'Hugging Face buckets path'
+    required: true
+
+runs:
+  using: "composite"
+  steps:
+    - name: Install Hugging Face Hub CLI
+      shell: bash
+      run: |
+        python3 -m venv .venv-hf
+        .venv-hf/bin/pip install -U huggingface_hub==1.28.0
+
+    - name: Restore ccache from buckets
+      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
+            CACHE_PATH=$(hf buckets list "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}" --json | jq -r '[.[] | select(.type == "file") | select(.path | startswith("${{ inputs.folder }}/${{ inputs.key }}") and endswith(".tar.gz"))] | sort_by(.path) | last | .path // ""')
+            if [[ -n "$CACHE_PATH" ]]; then
+                echo "Restoring ccache from '$CACHE_PATH'."
+                hf buckets cp "hf://buckets/${{ inputs.hf_bucket }}/$CACHE_PATH" ccache_bucket.tar.gz
+                mkdir -p ccache_bucket
+                if tar -xzf ccache_bucket.tar.gz -C ccache_bucket; then
+                    rm -rf "$CCACHE_DIR"
+                    mv ccache_bucket "$CCACHE_DIR"
+                    ccache -z
+                fi
+                rm ccache_bucket.tar.gz
+            else
+                echo "No ccache found."
+            fi
+        else
+            echo "'$CCACHE_DIR' not found."
+        fi
+
+    - name: Save ccache to buckets
+      if: ${{ inputs.save == 'true' }}
+      shell: bash
+      run: |
+        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
+        fi
+
+    - name: Remove old ccache files from buckets
+      if: ${{ inputs.save == 'true' }}
+      shell: bash
+      run: |
+        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
@@ -22,7 +22,8 @@ on:
    types: [opened, synchronize, reopened]
    paths: [
      '.github/workflows/build-apple.yml',
-      'ggml/src/ggml-metal/**'
+      'ggml/src/ggml-metal/**',
+      'ggml/src/ggml-rpc/**'
    ]

 concurrency:
@@ -50,14 +50,22 @@ jobs:
          DEBIAN_FRONTEND: noninteractive
        run: |
          apt update
-          apt install -y cmake build-essential ninja-build libgomp1 git libssl-dev
+          apt install -y cmake build-essential ninja-build libgomp1 git libssl-dev jq python3 python3-venv python3-pip

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
        with:
          key: cuda-ubuntu-24.04-cuda
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+          save: false
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: cuda-ubuntu-24.04-cuda
+          folder: llama.cpp
+          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
@@ -72,15 +80,17 @@ jobs:
            -DGGML_CUDA_CUB_3DOT2=ON
          cmake --build build

-      - name: ccache-clear
-        uses: ./.github/actions/ccache-clear
+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
        env:
-          GH_TOKEN: ${{ github.token }}
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
        with:
          key: cuda-ubuntu-24.04-cuda
-          older: 5m
-          min: 1
-          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true

  hip:
    runs-on: ubuntu-22.04
@@ -95,14 +105,22 @@ jobs:
        id: depends
        run: |
          sudo apt-get update
-          sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev libssl-dev rocwmma-dev
+          sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev libssl-dev rocwmma-dev jq python3-venv

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
        with:
          key: cuda-ubuntu-22.04-hip
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+          save: false
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: cuda-ubuntu-22.04-hip
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache

      - name: Build with native CMake HIP support
        id: cmake_build
@@ -113,15 +131,17 @@ jobs:
            -DGGML_HIP=ON
          cmake --build build --config Release -j $(nproc)

-      - name: ccache-clear
-        uses: ./.github/actions/ccache-clear
+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
        env:
-          GH_TOKEN: ${{ github.token }}
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
        with:
          key: cuda-ubuntu-22.04-hip
-          older: 5m
-          min: 1
-          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true

  musa:
    runs-on: ubuntu-22.04
@@ -136,14 +156,22 @@ jobs:
        id: depends
        run: |
          apt-get update
-          apt-get install -y build-essential git cmake libssl-dev
+          apt-get install -y build-essential git cmake libssl-dev jq

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
        with:
          key: cuda-ubuntu-22.04-musa
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+          save: false
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: cuda-ubuntu-22.04-musa
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache

      - name: Build with native CMake MUSA support
        id: cmake_build
@@ -152,12 +180,14 @@ jobs:
            -DGGML_MUSA=ON
          time cmake --build build --config Release -j $(nproc)

-      - name: ccache-clear
-        uses: ./.github/actions/ccache-clear
+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
        env:
-          GH_TOKEN: ${{ github.token }}
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
        with:
          key: cuda-ubuntu-22.04-musa
-          older: 5m
-          min: 1
-          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+          folder: llama.cpp
+          evict-old-files: 1d
+          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
@@ -84,11 +84,13 @@ jobs:

            New version has been released.

+            ## Assets
+
            ${{ steps.desc.outputs.nightly }}

-            **Web UI:** the `nightly-tag.txt` asset contains the tag of the corresponding nightly release
+            ## More info

-            **More info:** [dist : releases and versioning of ggml-org projects](https://github.com/ggml-org/ggml/discussions/1579)
+            - [Releases and versioning of `ggml-org` projects](https://github.com/ggml-org/ggml/discussions/1579)

            ## ${{ steps.desc.outputs.changelog_title }}

@@ -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,6 +714,8 @@ jobs:
        with:
          key: release-windows-2025-vs2026-${{ matrix.arch }}-cpu

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.21
@@ -1287,11 +1267,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 +1290,11 @@ jobs:
        with:
          fetch-depth: 0

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

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

@@ -1414,10 +1392,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 +1423,7 @@ jobs:
            -DLLAMA_BUILD_SERVER=OFF \
            -DCMAKE_SYSTEM_NAME=iOS \
            -DCMAKE_OSX_DEPLOYMENT_TARGET=16.0 \
-            -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml \
-            -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }}
+            -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml
          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu) -- CODE_SIGNING_ALLOWED=NO

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

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

  release:
    if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
@@ -1588,14 +1563,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 +1602,27 @@ jobs:
          path: ./artifact
          merge-multiple: true

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

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

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

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

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

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

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

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

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

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

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

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

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

      - name: Install Playwright browsers
@@ -4,7 +4,7 @@ include(CheckIncludeFileCXX)

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

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

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

    add_opt(common_arg(
        {"--spec-draft-p-split", "--draft-p-split"}, "P",
@@ -1177,6 +1177,8 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_
        "</tool_call>",
    };

+    auto is_qwen3_coder  = !supports_reasoning;
+
    if (supports_reasoning) {
        data.thinking_start_tag = "<think>";
        // Support both </think> and <tool_call> as reasoning end sequences.
@@ -1217,13 +1219,15 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_

    std::vector<std::string> tool_call_starts = { "<tool_call>" };

-    // Match complete <function=name> opener for Qwen3-Coder models that occasionally omit the
-    // starting <tool_call>. The model may hallucinate a tool name, but it is preferable over
-    // constraining on <function which may occur in valid content generation, e.g. #include <functional>
-    foreach_function(inputs.tools, [&](const json & tool) {
-        const std::string name = tool.at("function").at("name");
-        tool_call_starts.push_back("<function=" + name + ">");
-    });
+    if (is_qwen3_coder) {
+        // Match complete <function=name> opener for Qwen3-Coder models that occasionally omit the
+        // starting <tool_call>. The model may hallucinate a tool name, but it is preferable over
+        // constraining on <function which may occur in valid content generation, e.g. #include <functional>
+        foreach_function(inputs.tools, [&](const json & tool) {
+            const std::string name = tool.at("function").at("name");
+            tool_call_starts.push_back("<function=" + name + ">");
+        });
+    }

    auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
        auto generation_prompt = p.literal(GEN_PREFIX);
@@ -1288,10 +1292,13 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_

            auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0;

+            auto tool_call_body = tool_choice + "</tool_call>" + p.space();
+            auto tool_call      = p.rule("tool-call", "<tool_call>\n" + tool_call_body);
+
            // Qwen3-Coder models may occasionally omit the <tool_call> token.
-            auto tool_call_body  = tool_choice + "</tool_call>" + p.space();
-            auto tool_call_first = p.rule("tool-call-first", p.optional(p.literal("<tool_call>\n")) + tool_call_body);
-            auto tool_call       = p.rule("tool-call", "<tool_call>\n" + tool_call_body);
+            auto tool_call_first = is_qwen3_coder ?
+                p.rule("tool-call-first", p.optional(p.literal("<tool_call>\n")) + tool_call_body) :
+                tool_call;

            auto calls      = inputs.parallel_tool_calls ? tool_call_first + p.zero_or_more(tool_call) : tool_call_first;
            auto tool_calls = p.trigger_rule("tool-call-root", p.repeat(calls, min_calls, 1));
@@ -1688,6 +1688,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
    mparams.main_gpu        = params.main_gpu;
    mparams.split_mode      = params.split_mode;
    mparams.load_mode       = params.load_mode;
+    mparams.tensor_read_lazy = params.tensor_read_lazy;
    mparams.tensor_split    = params.tensor_split;
    mparams.check_tensors   = params.check_tensors;
    mparams.use_extra_bufts = !params.no_extra_bufts;
@@ -8,6 +8,7 @@
 #include "ggml.h"
 #include "llama.h"

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    virtual ~common_speculative_impl() = default;

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

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

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

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

    // dspark speculators
    bool sample_from_anchor = true;

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

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

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

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

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

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

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

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

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

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

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

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

            auto & result = *dp.result;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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


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

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

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

        # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd
@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+from typing import Iterable, cast
+
+import torch
+from torch import Tensor
+
+import gguf
+import numpy as np
+
+from .base import ModelBase
+from .qwen import _LinearAttentionVReorderBase, _Qwen35MRopeMixin
+from .qwen3vl import Qwen3VLVisionModel
+
+
+@ModelBase.register("Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM")
+@ModelBase.example("Qwen/Qwen3.8-Flash-Next")
+class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
+    """Qwen3.8-Flash-Next.
+
+    Shares the Qwen3.5 gated delta net and interleaved mrope, and adds three things:
+    hyper-connections in place of every layer norm, QSA sparse attention on the full
+    attention layers, and PLE n-gram hash embeddings on a single layer.
+    """
+
+    model_arch = gguf.MODEL_ARCH.QWEN4EXP
+
+    # the MTP block is a separate draft head; vLLM drops it too
+    supports_mtp_export = False
+    no_mtp = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # only the shard names, so the table itself is never held
+        self._ple_shards: dict[int, str] = {}
+        self._ple_row_dim: int | None = None
+
+    def _read_hash_constants(self, suffix: str) -> list[int]:
+        """Read an int64 PLE constant straight from the checkpoint.
+
+        prepare_tensors() casts every non-float dtype to float32 before
+        modify_tensors() sees it (base.py), which would silently round these
+        45-bit multipliers. Reading the lazy tensor here bypasses that.
+        """
+        for name, gen in self.model_tensors.items():
+            if name.endswith(suffix):
+                t = gen()
+                if t.dtype != torch.int64:
+                    t = t.to(torch.int64)
+                return [int(x) for x in t.tolist()]
+        raise ValueError(f"PLE constant {suffix!r} missing from the checkpoint")
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        hp = self.hparams
+
+        self.gguf_writer.add_hyper_connection_count(hp["hc_count"])
+        self.gguf_writer.add_hyper_connection_low_rank(hp["hc_lowrank"])
+
+        n_layer = hp["num_hidden_layers"]
+        self.gguf_writer.add_indexer_head_count(hp["indexer_n_heads"])
+        self.gguf_writer.add_indexer_key_length(hp["indexer_head_dim"])
+        self.gguf_writer.add_indexer_top_k(hp["indexer_budget"])
+        ratio = hp["indexer_compress_ratio"]
+        layer_types = hp["layer_types"]
+        self.gguf_writer.add_attention_compress_ratios(
+            [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)]
+        )
+
+        # ple_layer_ids is 1-based in the HF config; empty means no n-gram table,
+        # so emit no PLE keys rather than optional ones
+        ple_layers = [i - 1 for i in hp["ple_layer_ids"]]
+        if not ple_layers:
+            return
+        self.gguf_writer.add_ple_layers(ple_layers)
+        self.gguf_writer.add_ple_ngram_size(hp["ngram_size"])
+        self.gguf_writer.add_ple_heads_per_ngram(hp["heads_per_ngram"])
+        self.gguf_writer.add_ple_conv_kernel(hp["ple_conv_kernel_size"])
+        self.gguf_writer.add_ple_eos_token_id(self._eos_token_id())
+        # an image is decoded as an embeddings-only batch, so the graph has no placeholder
+        # ids to hash; carry the id and let it stand in for those positions
+        _img = self._image_token_id()
+        if _img is not None:
+            self.gguf_writer.add_ple_image_token_id(int(_img))
+        if self._ple_row_dim is not None:
+            self.gguf_writer.add_embedding_length_per_layer_input(self._ple_row_dim)
+
+        self.gguf_writer.add_ple_layer_multipliers(
+            self._read_hash_constants("ple_embedding.layer_multipliers"))
+        self.gguf_writer.add_ple_head_offsets(
+            self._read_hash_constants("ple_embedding.ngram_heads_offsets"))
+        self.gguf_writer.add_ple_head_vocab_sizes(
+            self._read_hash_constants("ple_embedding.ngram_heads_vocab_sizes"))
+
+    def _image_token_id(self) -> int | None:
+        img = self.hparams.get("image_token_id")
+        return None if img is None else int(img)
+
+    def _eos_token_id(self) -> int:
+        eos = self.hparams.get("eos_token_id")
+        if isinstance(eos, list):
+            # the PLE hash resets n-grams on the primary EOS
+            return int(eos[-1])
+        if eos is None:
+            raise ValueError("eos_token_id is required: the PLE hash resets its n-grams on it")
+        return int(eos)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # int64 hash constants must stay exact; 1-D tensors force F32, so use KV
+        if name.endswith("ple_embedding.layer_multipliers"):
+            self._ple_multipliers = [int(x) for x in data_torch.tolist()]
+            return []
+        if name.endswith("ple_embedding.ngram_heads_offsets"):
+            self._ple_head_offsets = [int(x) for x in data_torch.tolist()]
+            return []
+        if name.endswith("ple_embedding.ngram_heads_vocab_sizes"):
+            self._ple_head_vocab_sizes = [int(x) for x in data_torch.tolist()]
+            return []
+
+        if ".ngram_embedding.shard_" in name:
+            return self._place_ple_shard(data_torch, name)
+
+        # one projection feeds indexer q and k; split it, as minimax-m3 does
+        if ".indexer.index_qk_proj.weight" in name:
+            n_q = self.hparams["indexer_n_heads"] * self.hparams["indexer_head_dim"]
+            q = data_torch[:n_q]
+            k = data_torch[n_q:]
+            return [
+                (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid, ".weight"), q),
+                (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid, ".weight"), k),
+            ]
+
+        # Gemma zero-centred gammas the inherited norm.weight rule misses
+        if name.endswith((".ple.norm_key.weight", ".ple.norm_query.weight", ".ple.norm_conv.weight",
+                          ".indexer.q_layernorm.weight", ".indexer.k_layernorm.weight")):
+            return [(self.map_tensor_name(name), data_torch + 1)]
+
+        if name.endswith(".ple.conv1d.weight"):
+            return [(self.map_tensor_name(name), data_torch.squeeze())]
+
+        return super().modify_tensors(data_torch, name, bid)
+
+    # the shards concatenate into a tensor of well over 100 GB
+    # use LazyChunkedTensor here, a single shard resident at a time
+    def _place_ple_shard(self, data_torch: Tensor, name: str) -> Iterable[tuple[str, Tensor]]:
+
+        idx = int(name.rpartition(".shard_")[2].partition(".")[0])
+        n_parts = self.hparams["split_ngram_parts"]
+
+        self._ple_shards[idx] = name
+        self._ple_row_dim = int(data_torch.shape[-1])
+
+        if len(self._ple_shards) < n_parts:
+            return []
+
+        # the checkpoint may yield the shards in any order, the row order is by index
+        shards = [self._ple_shards[i] for i in sorted(self._ple_shards)]
+        rows = 0
+        for shard in shards:
+            shape = self.model_tensors[shard]().shape
+            if int(shape[-1]) != self._ple_row_dim:
+                raise ValueError(
+                    f"PLE shard {shard} has row dim {int(shape[-1])}, expected {self._ple_row_dim}")
+            rows += int(shape[0])
+
+        table = gguf.LazyChunkedTensor(
+            [self._load_ple_shard(shard) for shard in shards],
+            shape=(rows, self._ple_row_dim),
+            dtype=np.float32,
+        )
+        gguf_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.PER_LAYER_TOKEN_EMBD]
+        return [(gguf_name + ".weight", cast(Tensor, table))]
+
+    def _load_ple_shard(self, name: str):
+        def load() -> np.ndarray:
+            from .base import LazyTorchTensor
+
+            # a fresh lazy tensor every call, or to_eager() memoizes every shard
+            eager = LazyTorchTensor.to_eager(self.model_tensors[name]())
+            return eager.to(torch.float32).contiguous().numpy()
+        return load
+
+    def prepare_tensors(self):
+        super().prepare_tensors()
+        n_parts = self.hparams.get("split_ngram_parts", 0)
+        if self._ple_shards and len(self._ple_shards) != n_parts:
+            raise ValueError(
+                f"got {len(self._ple_shards)} PLE embedding shards, expected {n_parts}"
+            )
+
+
+@ModelBase.register("Qwen4ExpForConditionalGeneration")
+@ModelBase.example("Qwen/Qwen3.8-Flash-Next")
+class Qwen4ExpVisionModel(Qwen3VLVisionModel):
+    """The vision tower is an unmodified Qwen3-VL ViT."""
@@ -8,7 +8,7 @@
        "toolset":      { "value": "host=x86_64", "strategy": "external" },
        "cacheVariables": {
            "ANDROID_ABI":      "arm64-v8a",
-            "ANDROID_PLATFORM": "android-31",
+            "ANDROID_PLATFORM": "android-34",
            "CMAKE_TOOLCHAIN_FILE": "$env{ANDROID_NDK_ROOT}/build/cmake/android.toolchain.cmake",
            "CMAKE_C_FLAGS":   "-march=armv8.7a+fp16+dotprod+i8mm -fvectorize -ffp-model=fast -fno-finite-math-only -flto -D_GNU_SOURCE",
            "CMAKE_CXX_FLAGS": "-march=armv8.7a+fp16+dotprod+i8mm -fvectorize -ffp-model=fast -fno-finite-math-only -flto -D_GNU_SOURCE",
@@ -2,39 +2,47 @@

 ## Setup

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

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

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

 ## How to Build

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

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

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

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

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

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

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

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

 ### Windows

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

 ## How to Run

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

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

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

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

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

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

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

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

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

 ## Environment variables

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

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

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

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

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

  Examples:

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

 ## Large model handling

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

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

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

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

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

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

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


 ## How to Build

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

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

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

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

 ## How to Install

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

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

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

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

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

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

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

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

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

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

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

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

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

 ```
@@ -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,
@@ -70,6 +62,8 @@ enum htp_op_code {
    HTP_OP_UNARY_NEG,
    HTP_OP_UNARY_SOFTPLUS,
    HTP_OP_UNARY_TANH,
+    HTP_OP_UNARY_ABS,
+    HTP_OP_UNARY_LOG,
    HTP_OP_GLU_SWIGLU,
    HTP_OP_GLU_SWIGLU_OAI,
    HTP_OP_GLU_GEGLU,
@@ -99,12 +93,15 @@ enum htp_op_code {
    HTP_OP_CONCAT,
    HTP_OP_CLAMP,
    HTP_OP_IM2COL,
+    HTP_OP_FENCE,
+    HTP_OP_ALLREDUCE,
+    HTP_OP_ALLREDUCE_ADD,

    HTP_OP_INVALID
 };

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

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

 #define HTP_MMAP_MAX_VMEM  (2147483648u)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

-    res->d_queue = dispatch_queue_create("ggml-metal", DISPATCH_QUEUE_CONCURRENT);
-
-    res->use_fusion      = getenv("GGML_METAL_FUSION_DISABLE") == nil;
-    res->use_concurrency = getenv("GGML_METAL_CONCURRENCY_DISABLE") == nil;
-
-    {
-        const char * val = getenv("GGML_METAL_GRAPH_DEBUG");
-        res->debug_graph = val ? atoi(val) : 0;
-    }
-
-    {
-        const char * val = getenv("GGML_METAL_FUSION_DEBUG");
-        res->debug_fusion = val ? atoi(val) : 0;
-    }
-
-    res->use_graph_optimize = true;
-
-    if (getenv("GGML_METAL_GRAPH_OPTIMIZE_DISABLE") != NULL) {
-        res->use_graph_optimize = false;
-    }
-
-    memset(res->fuse_cnt, 0, sizeof(res->fuse_cnt));
-
-    GGML_LOG_INFO("%s: use fusion         = %s\n", __func__, res->use_fusion         ? "true" : "false");
-    GGML_LOG_INFO("%s: use concurrency    = %s\n", __func__, res->use_concurrency    ? "true" : "false");
-    GGML_LOG_INFO("%s: use graph optimize = %s\n", __func__, res->use_graph_optimize ? "true" : "false");
-
-    res->capture_compute = 0;
-    res->capture_started = false;
-    res->capture_scope = nil;
-
-    {
-        const char * val = getenv("GGML_METAL_CAPTURE_COMPUTE");
-        if (val) {
-            res->capture_compute = atoi(val);
+                return NULL;
+            }
        }
+
+        res->ev_cpy = ggml_metal_device_event_init(dev);
+
+        const struct ggml_metal_device_props * props_dev = ggml_metal_device_get_props(dev);
+
+        snprintf(res->name, sizeof(res->name), "%s", props_dev->name);
+
+        res->d_queue = dispatch_queue_create("ggml-metal", DISPATCH_QUEUE_CONCURRENT);
+
+        res->use_fusion      = getenv("GGML_METAL_FUSION_DISABLE") == nil;
+        res->use_concurrency = getenv("GGML_METAL_CONCURRENCY_DISABLE") == nil;
+
+        {
+            const char * val = getenv("GGML_METAL_GRAPH_DEBUG");
+            res->debug_graph = val ? atoi(val) : 0;
+        }
+
+        {
+            const char * val = getenv("GGML_METAL_FUSION_DEBUG");
+            res->debug_fusion = val ? atoi(val) : 0;
+        }
+
+        res->use_graph_optimize = true;
+
+        if (getenv("GGML_METAL_GRAPH_OPTIMIZE_DISABLE") != NULL) {
+            res->use_graph_optimize = false;
+        }
+
+        memset(res->fuse_cnt, 0, sizeof(res->fuse_cnt));
+
+        GGML_LOG_INFO("%s: use fusion         = %s\n", __func__, res->use_fusion         ? "true" : "false");
+        GGML_LOG_INFO("%s: use concurrency    = %s\n", __func__, res->use_concurrency    ? "true" : "false");
+        GGML_LOG_INFO("%s: use graph optimize = %s\n", __func__, res->use_graph_optimize ? "true" : "false");
+
+        res->capture_compute = 0;
+        res->capture_started = false;
+        res->capture_scope = nil;
+
+        {
+            const char * val = getenv("GGML_METAL_CAPTURE_COMPUTE");
+            if (val) {
+                res->capture_compute = atoi(val);
+            }
+        }
+
+        res->has_error = false;
+
+        res->gf = nil;
+        res->encode_async = nil;
+        for (int i = 0; i < GGML_METAL_MAX_COMMAND_BUFFERS; ++i) {
+            res->cmd_bufs[i].obj = nil;
+        }
+
+        res->cmd_bufs_ext = [[NSMutableArray alloc] init];
+
+        res->cmd_buf_last = nil;
+
+        res->pipelines_ext = ggml_metal_pipelines_init();
+
+        return res;
    }
-
-    res->has_error = false;
-
-    res->gf = nil;
-    res->encode_async = nil;
-    for (int i = 0; i < GGML_METAL_MAX_COMMAND_BUFFERS; ++i) {
-        res->cmd_bufs[i].obj = nil;
-    }
-
-    res->cmd_bufs_ext = [[NSMutableArray alloc] init];
-
-    res->cmd_buf_last = nil;
-
-    res->pipelines_ext = ggml_metal_pipelines_init();
-
-    return res;
 }

 void ggml_metal_free(ggml_metal_t ctx) {
@@ -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);
@@ -778,7 +778,9 @@ void ggml_metal_encoder_free(ggml_metal_encoder_t encoder) {
 }

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

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

    assert(dev != NULL);

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

-            snprintf(dev->props.name, sizeof(dev->props.name), "%s%d", "MTL", device);
-            const char * gpu_name = [[dev->mtl_device name] UTF8String];
-            if (n_devices > 1) {
-                snprintf(dev->props.desc, sizeof(dev->props.desc), "%s (dev p%d/v%d)",
-                         gpu_name, dev->props.device_phys, dev->props.device_virt);
-            } else {
-                snprintf(dev->props.desc, sizeof(dev->props.desc), "%s", gpu_name);
-            }
+                snprintf(dev->props.name, sizeof(dev->props.name), "%s%d", "MTL", device);
+                const char * gpu_name = [[dev->mtl_device name] UTF8String];
+                if (n_devices > 1) {
+                    snprintf(dev->props.desc, sizeof(dev->props.desc), "%s (dev p%d/v%d)",
+                             gpu_name, dev->props.device_phys, dev->props.device_virt);
+                } else {
+                    snprintf(dev->props.desc, sizeof(dev->props.desc), "%s", gpu_name);
+                }

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

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

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

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

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

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

@@ -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;
 }
@@ -204,6 +204,11 @@ static ggml_backend_buffer_t ggml_backend_metal_buffer_type_alloc_buffer(ggml_ba
    ggml_metal_device_t ctx_dev = (ggml_metal_device_t)buft->device->context;
    ggml_metal_buffer_t res = ggml_metal_buffer_init(ctx_dev, size, shared);

+    if (res == NULL) {
+        GGML_LOG_ERROR("%s: failed to allocate Metal buffer of %zu bytes (out of memory)\n", __func__, size);
+        return NULL;
+    }
+
    ggml_backend_buffer_i buf_i = ggml_metal_buffer_is_shared(res)
        ? ggml_backend_metal_buffer_shared_i
        : ggml_backend_metal_buffer_private_i;
@@ -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);
+    }
+}
@@ -903,6 +903,8 @@ struct ggml_backend_opencl_context {
    cl_kernel kernel_gemv_moe_mxfp4_f32_ns_wimg = nullptr;      // weight-as-texture MoE decode GEMV
    cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a = nullptr;   // dp4a (int8) mxfp4 MoE prefill GEMM
    cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a = nullptr;    // dp4a (int8) q4_0 MoE prefill GEMM
+    cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a_bin = nullptr;   // binary dp4a (int8) mxfp4 MoE prefill GEMM
+    cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a_bin = nullptr;    // binary dp4a (int8) q4_0 MoE prefill GEMM
    cl_kernel kernel_moe_reorder_b;
    cl_kernel kernel_moe_histogram, kernel_moe_scan, kernel_moe_fill, kernel_moe_scatter;
    cl_kernel kernel_moe_scatter_stable = nullptr;   // deterministic slot assignment
@@ -4248,6 +4250,24 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
        GGML_LOG_CONT(".");
    }

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

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

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

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

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

                        // dp4a GEMM
                        cl_kernel dk = backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a;
+                        if (backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin) {
+                            dk = backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin;
+                        }
+
                        int aidx = 0;
                        CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_mxfp4->q_img));
                        CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_mxfp4->e));
@@ -9,10 +9,18 @@ if (WIN32)
    target_link_libraries(ggml-rpc PRIVATE ws2_32)
 endif()

-# RDMA auto-detection (Linux only, requires libibverbs)
-if (NOT WIN32 AND NOT APPLE)
-    find_library(IBVERBS_LIB ibverbs)
-    if (IBVERBS_LIB)
+# RDMA auto-detection: Linux RoCE/IB via libibverbs, Apple RDMA-over-Thunderbolt via librdma
+if (APPLE)
+    set(RDMA_LIB_NAME rdma)
+    set(RDMA_DESC "Apple RDMA-over-Thunderbolt, UC")
+elseif (NOT WIN32)
+    set(RDMA_LIB_NAME ibverbs)
+    set(RDMA_DESC "auto-detected")
+endif()
+
+if (RDMA_LIB_NAME)
+    find_library(RDMA_LIB ${RDMA_LIB_NAME})
+    if (RDMA_LIB)
        option(GGML_RPC_RDMA "ggml: enable RDMA transport for RPC" ON)
    else()
        option(GGML_RPC_RDMA "ggml: enable RDMA transport for RPC" OFF)
@@ -22,12 +30,16 @@ else()
 endif()

 if (GGML_RPC_RDMA)
-    if (NOT IBVERBS_LIB)
-        find_library(IBVERBS_LIB ibverbs REQUIRED)
+    if (NOT RDMA_LIB)
+        find_library(RDMA_LIB ${RDMA_LIB_NAME} REQUIRED)
    endif()
    target_compile_definitions(ggml-rpc PRIVATE GGML_RPC_RDMA)
-    target_link_libraries(ggml-rpc PRIVATE ${IBVERBS_LIB})
-    message(STATUS "  RDMA transport enabled (auto-detected)")
+    target_link_libraries(ggml-rpc PRIVATE ${RDMA_LIB})
+    if (APPLE)
+        target_compile_definitions(ggml-rpc PRIVATE GGML_RPC_RDMA_APPLE)
+        target_sources(ggml-rpc PRIVATE transport-apple.cpp)
+    endif()
+    message(STATUS "  RDMA transport enabled (${RDMA_DESC})")
 else()
    message(STATUS "  RDMA transport disabled")
 endif()
@@ -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];
@@ -253,7 +259,10 @@ static bool send_msg(socket_ptr sock, const void * msg, size_t msg_size) {
    if (!sock->send_data(&msg_size, sizeof(msg_size))) {
        return false;
    }
-    return sock->send_data(msg, msg_size);
+    if (!sock->send_data(msg, msg_size)) {
+        return false;
+    }
+    return sock->flush();
 }

 static bool recv_msg(socket_ptr sock, void * msg, size_t msg_size) {
@@ -308,7 +317,7 @@ static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input,
    if (!sock->send_data(input, input_size)) {
        return false;
    }
-    return true;
+    return sock->flush();
 }

 // RPC request : | rpc_cmd (1 byte) | request_size (8 bytes) | request_data (request_size bytes) |
@@ -354,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;
 }

@@ -400,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;
 }
@@ -460,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;
 }
@@ -473,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;
@@ -501,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) {
@@ -526,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;
@@ -543,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 = {
@@ -569,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 {
@@ -585,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;
 }

@@ -598,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;
 }

@@ -625,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;
    }
@@ -670,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) {
@@ -695,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;
@@ -705,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));
@@ -720,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) {
@@ -730,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,
@@ -760,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,
 };

@@ -775,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,
@@ -801,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 {
@@ -820,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 {
@@ -1644,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: {
@@ -1657,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: {
@@ -1670,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: {
@@ -1707,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: {
@@ -1886,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,
    };
 }
@@ -1926,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,
@@ -1939,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
@@ -2001,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;
 }

@@ -0,0 +1,470 @@
+#include "transport-apple.h"
+#include "transport.h"
+#include "ggml-impl.h"
+
+#include <infiniband/verbs.h>
+
+#include <cerrno>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include <poll.h>
+#include <sys/socket.h>
+#include <unistd.h>
+
+// Apple RDMA-over-Thunderbolt (see Apple TN3205).
+//
+// Apple's RDMA is quite different from what's supported in Linux - deserving of its own transport implementation.
+// see https://developer.apple.com/documentation/technotes/tn3205-low-latency-communication-with-rdma-over-thunderbolt for details
+// at a high level the main differences are:
+// UC(unreliable connection) on Apple vs RC(reliable connection) QP transport types on Linux (though in practice UC on Apple is still lossless)
+// fixed 128KiB stride on Apple vs variable chunk size on Linux
+// relying on Apple's hardware credit based flow control vs RNR NAKs + retries on Linux
+//
+// on Apple a SEND and its corresponding RECV must cover the same number of 4 KiB Thunderbolt frames,
+// so every SEND posts a whole 128KiB stride over the wire, even when partially filled.
+// (In testing 128KiB was the best performing among 32, 64, 128, 256)
+
+static constexpr uint32_t RDMA_SEG_MAGIC   = 0x52534547u; // "RSEG"
+static constexpr int      RDMA_NBUF        = 16;          // ring depth (frames per direction)
+static constexpr size_t   RDMA_FRAME       = 4096;        // Thunderbolt frame (fixed on Apple)
+static constexpr size_t   RDMA_STRIDE      = 128 * 1024;  // 32 Thunderbolt frames; NBUF x this = 2 MiB pinned per direction
+static constexpr uint32_t RDMA_PSN         = 0;           // any value works if both sides match: UC has no retransmit
+static constexpr size_t   RDMA_GID_SIZE    = 16;
+
+static_assert(RDMA_STRIDE % RDMA_FRAME == 0, "RDMA_STRIDE must be a whole number of frames");
+// TN3205 counts queue depth in Thunderbolt frames, not work requests.
+static constexpr uint32_t RDMA_QP_WR       = (uint32_t)RDMA_NBUF * (RDMA_STRIDE / RDMA_FRAME);
+static constexpr uint64_t RDMA_RECV_WR     = 1ull << 20;  // wr_id bit tagging recv completions
+static constexpr uint64_t RDMA_WR_IDX_MASK = 0xffff;      // buffer index in the low bits of wr_id
+static constexpr uint8_t  RDMA_SYNC_READY  = 0x2A;        // readiness-handshake byte (peer activated)
+
+struct rdma_seg_hdr {
+    uint32_t magic; // RDMA_SEG_MAGIC; a mismatch means the stream desynced
+    uint32_t len;   // payload bytes in this frame; the rest of the stride is padding
+};
+static constexpr size_t RDMA_PAYLOAD = RDMA_STRIDE - sizeof(rdma_seg_hdr);
+
+struct apple_rdma_caps {
+    uint32_t qpn;
+    uint16_t lid;
+    uint16_t reserved;
+    uint8_t  gid[RDMA_GID_SIZE];
+};
+
+static_assert(sizeof(apple_rdma_caps) == RPC_CONN_CAPS_SIZE, "apple_rdma_caps must match conn_caps size");
+
+struct apple_rdma::impl {
+    int fd = -1;                      // bootstrap TCP socket, kept as the liveness anchor
+
+    struct ibv_context * ctx = nullptr;
+    struct ibv_pd * pd = nullptr;
+    struct ibv_cq * cq = nullptr;     // one CQ for both directions; RDMA_RECV_WR tags recv completions
+    struct ibv_qp * qp = nullptr;
+
+    uint8_t       * send_mem = nullptr;
+    struct ibv_mr * send_mr  = nullptr;
+    uint8_t       * recv_mem = nullptr;
+    struct ibv_mr * recv_mr  = nullptr;
+
+    int      send_busy[RDMA_NBUF] = {};  // 1 while this buffer has a send in flight
+    // completed recv frames, oldest first: ring index, bytes already handed to
+    // the reader, and total payload length
+    struct { int buf; uint32_t off; uint32_t len; } inq[RDMA_NBUF] = {};
+    int      inq_head = 0;
+    int      inq_count = 0;
+    int      pend_buf = -1;
+    uint32_t pend_len = 0;
+    bool     broken = false;
+
+    uint32_t     qpn = 0;
+    uint8_t      port = 0;
+    int          gid_idx = 0;
+    enum ibv_mtu path_mtu = IBV_MTU_1024;
+
+    int  progress();
+    bool acquire_pending();
+    bool post_pending();
+
+    bool post_recv(int i) {
+        struct ibv_sge sge = {};
+        sge.addr   = (uintptr_t)(recv_mem + (size_t)i * RDMA_STRIDE);
+        sge.length = (uint32_t)RDMA_STRIDE;
+        sge.lkey   = recv_mr->lkey;
+        struct ibv_recv_wr wr = {}, * bad = nullptr;
+        wr.wr_id   = RDMA_RECV_WR | (uint64_t)i;
+        wr.sg_list = &sge;
+        wr.num_sge = 1;
+        return ibv_post_recv(qp, &wr, &bad) == 0;
+    }
+
+    bool post_send(int i, size_t len) {
+        struct ibv_sge sge = {};
+        sge.addr   = (uintptr_t)(send_mem + (size_t)i * RDMA_STRIDE);
+        sge.length = (uint32_t)len;
+        sge.lkey   = send_mr->lkey;
+        struct ibv_send_wr wr = {}, * bad = nullptr;
+        wr.wr_id   = (uint64_t)i;
+        wr.sg_list = &sge;
+        wr.num_sge = 1;
+        wr.opcode  = IBV_WR_SEND;
+        wr.send_flags = IBV_SEND_SIGNALED;
+        return ibv_post_send(qp, &wr, &bad) == 0;
+    }
+
+    ~impl() {
+        broken = true;
+        // the QP must be destroyed before the memory it can still write to is
+        // deregistered and freed: ERR only starts flushing the posted WQEs
+        if (qp) {
+            struct ibv_qp_attr a = {};
+            a.qp_state = IBV_QPS_ERR;
+            ibv_modify_qp(qp, &a, IBV_QP_STATE);
+            struct ibv_wc wc[RDMA_NBUF * 2];
+            while (ibv_poll_cq(cq, RDMA_NBUF * 2, wc) > 0) {}
+            ibv_destroy_qp(qp);
+        }
+        if (send_mr) ibv_dereg_mr(send_mr);
+        if (recv_mr) ibv_dereg_mr(recv_mr);
+        free(send_mem);
+        free(recv_mem);
+        if (cq)  ibv_destroy_cq(cq);
+        if (pd)  ibv_dealloc_pd(pd);
+        if (ctx) ibv_close_device(ctx);
+    }
+};
+
+apple_rdma::apple_rdma(std::unique_ptr<impl> p) : pimpl(std::move(p)) {}
+
+apple_rdma::~apple_rdma() = default;
+
+bool apple_rdma::broken() const {
+    return pimpl->broken;
+}
+
+// The readiness handshake below still runs over the bootstrap socket, one byte
+// each way, before the transport is declared live.
+static bool tcp_send_byte(int fd, uint8_t b) {
+    ssize_t n;
+    do { n = ::send(fd, &b, sizeof(b), 0); } while (n < 0 && errno == EINTR);
+    return n == sizeof(b);
+}
+
+static bool tcp_recv_byte(int fd, uint8_t * b) {
+    ssize_t n;
+    do { n = ::recv(fd, b, sizeof(*b), 0); } while (n < 0 && errno == EINTR);
+    return n == (ssize_t)sizeof(*b);
+}
+
+// Index of the GID on this port equal to the target, or -1. Thunderbolt GIDs are
+// RoCEv2 IPv4-mapped (::ffff:a.b.c.d), so this matches the local TCP address.
+static int rdma_match_gid(struct ibv_context * ctx, uint8_t port, int gid_tbl_len,
+                          const uint8_t * target, union ibv_gid * out) {
+    for (int i = 0; i < gid_tbl_len; i++) {
+        union ibv_gid g;
+        if (ibv_query_gid(ctx, port, i, &g) != 0) continue;
+        if (memcmp(g.raw, target, RDMA_GID_SIZE) != 0) continue;
+        if (out) *out = g;
+        return i;
+    }
+    return -1;
+}
+
+// First ACTIVE port on the device. Only a cabled, up Thunderbolt link reports
+// ACTIVE, and it is not always port 1, so the port cannot be hardcoded the way
+// the Linux path does. Returns 0 if none.
+static uint8_t rdma_first_active_port(struct ibv_context * ctx, struct ibv_port_attr * out) {
+    struct ibv_device_attr da;
+    if (ibv_query_device(ctx, &da) != 0) return 0;
+    for (uint8_t p = 1; p <= da.phys_port_cnt; p++) {
+        struct ibv_port_attr pa;
+        if (ibv_query_port(ctx, p, &pa) != 0) continue;
+        if (pa.state == IBV_PORT_ACTIVE) { if (out) *out = pa; return p; }
+    }
+    return 0;
+}
+
+// Called before the endpoints are exchanged: pick the local device facing this
+// peer, create a UC QP and register the frame rings. RDMA is point-to-point, so
+// the device is the one whose GID equals the bootstrap connection's local
+// address, i.e. the one cabled to the peer.
+std::unique_ptr<apple_rdma> apple_rdma::probe(int fd, const uint8_t * target_gid, uint8_t * caps) {
+    int ndev = 0;
+    ibv_device ** devs = ibv_get_device_list(&ndev);
+    if (!devs) return nullptr;
+
+    ibv_context * ctx = nullptr;
+    uint8_t port = 0;
+    struct ibv_port_attr pa = {};
+    union ibv_gid gid = {};
+    int gid_idx = -1;
+    std::string matched;
+    for (int d = 0; d < ndev; d++) {
+        ibv_context * c = ibv_open_device(devs[d]);
+        if (!c) continue;
+        struct ibv_port_attr p = {};
+        uint8_t pt = rdma_first_active_port(c, &p);
+        int gi = pt ? rdma_match_gid(c, pt, p.gid_tbl_len, target_gid, &gid) : -1;
+        if (gi < 0) { ibv_close_device(c); continue; }
+        ctx = c; port = pt; pa = p; gid_idx = gi;
+        const char * name = ibv_get_device_name(devs[d]);
+        matched = name ? name : "";
+        break;
+    }
+    ibv_free_device_list(devs);
+    if (!ctx) return nullptr;
+
+    std::unique_ptr<impl> c(new impl());
+    c->fd  = fd;
+    c->ctx = ctx;
+    c->port = port;
+    c->gid_idx = gid_idx;
+    c->path_mtu = pa.active_mtu;
+
+    c->pd = ibv_alloc_pd(ctx);
+    if (!c->pd) return nullptr;
+
+    c->cq = ibv_create_cq(ctx, 2 * RDMA_QP_WR + 1, nullptr, nullptr, 0);
+    if (!c->cq) return nullptr;
+
+    ibv_qp_init_attr qia = {};
+    qia.send_cq = c->cq;
+    qia.recv_cq = c->cq;
+    qia.qp_type = IBV_QPT_UC;
+    qia.cap.max_send_wr  = RDMA_QP_WR;
+    qia.cap.max_recv_wr  = RDMA_QP_WR;
+    qia.cap.max_send_sge = 1;
+    qia.cap.max_recv_sge = 1;
+    c->qp = ibv_create_qp(c->pd, &qia);
+    if (!c->qp) return nullptr;
+
+    {
+        ibv_qp_attr a = {};
+        a.qp_state = IBV_QPS_INIT;
+        a.pkey_index = 0;
+        a.port_num = port;
+        a.qp_access_flags = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE;
+        if (ibv_modify_qp(c->qp, &a,
+                IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | IBV_QP_ACCESS_FLAGS) != 0) {
+            return nullptr;
+        }
+    }
+
+    long page = sysconf(_SC_PAGESIZE);
+    if (page <= 0) page = 4096;
+    const size_t ring_bytes = (size_t)RDMA_NBUF * RDMA_STRIDE;
+    if (posix_memalign((void **)&c->send_mem, (size_t)page, ring_bytes) != 0) c->send_mem = nullptr;
+    if (posix_memalign((void **)&c->recv_mem, (size_t)page, ring_bytes) != 0) c->recv_mem = nullptr;
+    if (!c->send_mem || !c->recv_mem) return nullptr;
+
+    // Apple's provider rejects LOCAL_WRITE-only MRs even for two-sided SEND/RECV.
+    const int mr_flags = IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ | IBV_ACCESS_REMOTE_WRITE;
+    c->send_mr = ibv_reg_mr(c->pd, c->send_mem, ring_bytes, mr_flags);
+    c->recv_mr = ibv_reg_mr(c->pd, c->recv_mem, ring_bytes, mr_flags);
+    if (!c->send_mr || !c->recv_mr) return nullptr;
+
+    // Recvs are posted in activate() after the RTS transition, not here: Apple's
+    // provider rejects ibv_post_recv on a QP that has not reached RTS.
+
+    c->qpn = c->qp->qp_num;
+
+    apple_rdma_caps rc = {};
+    rc.qpn = c->qpn;
+    rc.lid = pa.lid;
+    memcpy(rc.gid, gid.raw, RDMA_GID_SIZE);
+    memcpy(caps, &rc, sizeof(rc));
+
+    GGML_LOG_INFO("RDMA(Apple/UC) probed: dev=%s port=%u gid=%d qpn=%u lid=%u mtu=%d ring=%d x %zu KiB\n",
+                  matched.c_str(), port, gid_idx, c->qpn, (unsigned)pa.lid, 128 << c->path_mtu,
+                  RDMA_NBUF, RDMA_STRIDE / 1024);
+    return std::unique_ptr<apple_rdma>(new apple_rdma(std::move(c)));
+}
+
+// Called once the peer's endpoint has arrived: INIT -> RTR -> RTS (UC: GID/GRH
+// addressing, no timeout/retry/rnr/rd_atomic), then the readiness handshake.
+bool apple_rdma::activate(const uint8_t * caps) {
+    impl * c = pimpl.get();
+
+    apple_rdma_caps rc = {};
+    memcpy(&rc, caps, sizeof(rc));
+
+    bool ok = true;
+    {
+        ibv_qp_attr a = {};
+        a.qp_state   = IBV_QPS_RTR;
+        a.path_mtu   = c->path_mtu;
+        a.rq_psn     = RDMA_PSN;
+        a.dest_qp_num = rc.qpn;
+        a.ah_attr.is_global     = 1;
+        a.ah_attr.port_num      = c->port;
+        a.ah_attr.sl            = 0;
+        a.ah_attr.src_path_bits = 0;
+        a.ah_attr.dlid          = rc.lid;
+        a.ah_attr.grh.hop_limit  = 1;
+        a.ah_attr.grh.sgid_index = (uint8_t)c->gid_idx;
+        memcpy(&a.ah_attr.grh.dgid, rc.gid, RDMA_GID_SIZE);
+        if (ibv_modify_qp(c->qp, &a,
+                IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN | IBV_QP_RQ_PSN) != 0) {
+            GGML_LOG_ERROR("RDMA(Apple/UC) RTR failed: %s\n", strerror(errno));
+            ok = false;
+        }
+    }
+    if (ok) {
+        ibv_qp_attr a = {};
+        a.qp_state = IBV_QPS_RTS;
+        a.sq_psn   = RDMA_PSN;
+        if (ibv_modify_qp(c->qp, &a, IBV_QP_STATE | IBV_QP_SQ_PSN) != 0) {
+            GGML_LOG_ERROR("RDMA(Apple/UC) RTS failed: %s\n", strerror(errno));
+            ok = false;
+        }
+    }
+
+    // Recvs are posted only now: the controller starts processing them at RTR.
+    for (int i = 0; ok && i < RDMA_NBUF; i++) {
+        if (!c->post_recv(i)) {
+            GGML_LOG_ERROR("RDMA(Apple/UC) post_recv %d/%d failed\n", i, RDMA_NBUF);
+            ok = false;
+        }
+    }
+
+    // A queue pair processes receives only after RTR and the transitions above can
+    // fail on one side alone, so neither peer sends a frame until both report their
+    // recvs posted.
+    uint8_t peer_ready = 0;
+    if (!tcp_send_byte(c->fd, ok ? RDMA_SYNC_READY : 0) || !tcp_recv_byte(c->fd, &peer_ready)) {
+        return false;
+    }
+    if (!ok || peer_ready != RDMA_SYNC_READY) {
+        return false;
+    }
+
+    GGML_LOG_INFO("RDMA(Apple/UC) activated: qpn=%u->%u mtu=%d rx_depth=%d\n",
+                  c->qpn, rc.qpn, 128 << c->path_mtu, RDMA_NBUF);
+    return true;
+}
+
+// Drain the CQ: release completed send buffers, queue completed recv frames for
+// the reader. Returns the number of completions reaped, or -1 on error.
+int apple_rdma::impl::progress() {
+    struct ibv_wc wc[RDMA_NBUF * 2];
+    int n = ibv_poll_cq(cq, RDMA_NBUF * 2, wc);
+    if (n < 0) { GGML_LOG_ERROR("RDMA(Apple/UC) poll_cq failed\n"); broken = true; return -1; }
+    for (int j = 0; j < n; j++) {
+        uint64_t id = wc[j].wr_id;
+        bool is_recv = (id & RDMA_RECV_WR) != 0;
+        if (wc[j].status != IBV_WC_SUCCESS) {
+            GGML_LOG_ERROR("RDMA(Apple/UC) %s wc error: status=%d\n", is_recv ? "recv" : "send", wc[j].status);
+            broken = true;
+            return -1;
+        }
+        if (is_recv) {
+            int b = (int)(id & RDMA_WR_IDX_MASK);
+            const rdma_seg_hdr * h = (const rdma_seg_hdr *)(recv_mem + (size_t)b * RDMA_STRIDE);
+            if (h->magic != RDMA_SEG_MAGIC) { GGML_LOG_ERROR("RDMA(Apple/UC) bad frame magic\n"); broken = true; return -1; }
+            if (h->len > RDMA_PAYLOAD) { GGML_LOG_ERROR("RDMA(Apple/UC) frame len %u exceeds payload\n", h->len); broken = true; return -1; }
+            int slot = (inq_head + inq_count) % RDMA_NBUF;
+            inq[slot].buf  = b;
+            inq[slot].off  = 0;
+            inq[slot].len  = h->len;
+            inq_count++;
+        } else {
+            send_busy[(int)(id & RDMA_WR_IDX_MASK)] = 0;
+        }
+    }
+    return n;
+}
+
+// Reserve a free send buffer to coalesce into, waiting on progress if none free.
+bool apple_rdma::impl::acquire_pending() {
+    if (pend_buf >= 0) return true;
+    for (;;) {
+        if (broken) return false;
+        for (int k = 0; k < RDMA_NBUF; k++) if (!send_busy[k]) { pend_buf = k; pend_len = 0; return true; }
+        if (progress() < 0) return false;
+    }
+}
+
+// Post the pending frame. The whole STRIDE goes out even when only partly filled:
+// TN3205 requires a SEND and its matching RECV to cover the same number of
+// Thunderbolt frames, so a short send would fail the peer's receive.
+bool apple_rdma::impl::post_pending() {
+    if (pend_buf < 0) return true;
+    int i = pend_buf;
+    rdma_seg_hdr * h = (rdma_seg_hdr *)(send_mem + (size_t)i * RDMA_STRIDE);
+    h->magic = RDMA_SEG_MAGIC;
+    h->len   = pend_len;
+    if (!post_send(i, RDMA_STRIDE)) { broken = true; return false; }
+    send_busy[i] = 1;
+    pend_buf = -1;
+    pend_len = 0;
+    return true;
+}
+
+// Coalescing write: append into the pending frame, posting a full frame when it
+// fills. The trailing partial is posted by flush() at each message boundary.
+bool apple_rdma::send(const void * data, size_t size) {
+    impl * c = pimpl.get();
+    const uint8_t * p = (const uint8_t *)data;
+    while (size > 0) {
+        if (c->broken) return false;
+        if (!c->acquire_pending()) return false;
+        uint8_t * sb = c->send_mem + (size_t)c->pend_buf * RDMA_STRIDE;
+        size_t space = RDMA_PAYLOAD - c->pend_len;
+        size_t chunk = size < space ? size : space;
+        memcpy(sb + sizeof(rdma_seg_hdr) + c->pend_len, p, chunk);
+        c->pend_len += (uint32_t)chunk;
+        p += chunk;
+        size -= chunk;
+        if (c->pend_len == RDMA_PAYLOAD) { if (!c->post_pending()) return false; }
+    }
+    return true;
+}
+
+bool apple_rdma::recv(void * data, size_t size) {
+    impl * c = pimpl.get();
+    uint8_t * p = (uint8_t *)data;
+    if (!c->post_pending()) return false;   // turnaround: flush the coalesced request
+    unsigned idle = 0;
+    while (size > 0) {
+        if (c->inq_count == 0) {
+            if (c->broken) return false;
+            int n = c->progress();
+            if (n < 0) return false;
+            if (n == 0) {
+                // UC gives no disconnect notification, so the bootstrap TCP fd is
+                // the liveness anchor: nothing crosses it once RDMA is up, so any
+                // readability means the peer's FIN (macOS has no POLLRDHUP).
+                // Same idle interval as the Linux path.
+                if ((++idle & 0xFFFFF) == 0) {
+                    struct pollfd pfd = { c->fd, POLLIN, 0 };
+                    if (poll(&pfd, 1, 0) > 0 &&
+                        (pfd.revents & (POLLIN | POLLHUP | POLLERR | POLLNVAL))) {
+                        return false;
+                    }
+                }
+            } else {
+                idle = 0;
+            }
+            continue;
+        }
+        idle = 0;
+        int slot = c->inq_head;
+        int b = c->inq[slot].buf;
+        uint32_t avail = c->inq[slot].len - c->inq[slot].off;
+        uint32_t take = (size < (size_t)avail) ? (uint32_t)size : avail;
+        memcpy(p, c->recv_mem + (size_t)b * RDMA_STRIDE + sizeof(rdma_seg_hdr) + c->inq[slot].off, take);
+        p += take;
+        size -= take;
+        c->inq[slot].off += take;
+        if (c->inq[slot].off == c->inq[slot].len) {
+            if (!c->post_recv(b)) { c->broken = true; return false; }
+            c->inq_head = (c->inq_head + 1) % RDMA_NBUF;
+            c->inq_count--;
+        }
+    }
+    return true;
+}
+
+bool apple_rdma::flush() {
+    return pimpl->post_pending();
+}
@@ -0,0 +1,27 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+struct apple_rdma {
+    // target_gid is 16 bytes in, caps is RPC_CONN_CAPS_SIZE bytes out.
+    static std::unique_ptr<apple_rdma> probe(int fd, const uint8_t * target_gid, uint8_t * caps);
+    ~apple_rdma();
+
+    // Peer endpoint from its caps, which must be non-zero: this blocks on a
+    // readiness handshake over fd that the peer only joins if it also has RDMA.
+    bool activate(const uint8_t * caps);
+
+    bool send(const void * data, size_t size);
+    bool recv(void * data, size_t size);
+    // Post the trailing partial frame; must be called at every message boundary.
+    bool flush();
+    // True once the connection has failed; the caller should drop the socket.
+    bool broken() const;
+
+private:
+    struct impl;
+    explicit apple_rdma(std::unique_ptr<impl> p);
+    std::unique_ptr<impl> pimpl;
+};
@@ -18,15 +18,20 @@
 #  include <unistd.h>
 #endif
 #include <cstdlib>
+#include <cstring>
 #include <mutex>
 #include <optional>

 #ifdef GGML_RPC_RDMA
 #  include <infiniband/verbs.h>
+#  include <array>
 #  include <time.h>
 #  ifndef _WIN32
 #    include <poll.h>
 #  endif
+#  ifdef GGML_RPC_RDMA_APPLE
+#    include "transport-apple.h"
+#  endif
 #endif // GGML_RPC_RDMA

 #ifdef _WIN32
@@ -42,10 +47,13 @@ static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG");
    do { if (RPC_DEBUG) GGML_LOG_DEBUG(__VA_ARGS__); } while (0)

 #ifdef GGML_RPC_RDMA
-static constexpr size_t RDMA_CHUNK    = 256 * 1024;   // 256 KiB per send/recv (fits default 8 MiB memlock)
-static constexpr int    RDMA_RX_DEPTH = 24;            // pre-posted recv ring: 24 × 256 KiB = 6 MiB
 static constexpr size_t RDMA_GID_SIZE = 16;            // RoCE GID / IB GID is always 16 bytes
 using rdma_gid_t = std::array<uint8_t, RDMA_GID_SIZE>;
+#endif // GGML_RPC_RDMA
+
+#if defined(GGML_RPC_RDMA) && !defined(GGML_RPC_RDMA_APPLE)
+static constexpr size_t RDMA_CHUNK    = 256 * 1024;   // 256 KiB per send/recv (fits default 8 MiB memlock)
+static constexpr int    RDMA_RX_DEPTH = 24;            // pre-posted recv ring: 24 × 256 KiB = 6 MiB

 struct rdma_conn {
    struct ibv_context * ctx = nullptr;
@@ -111,27 +119,33 @@ struct rdma_caps {

 static_assert(sizeof(rdma_caps) == RPC_CONN_CAPS_SIZE, "rdma_caps must match conn_caps size");

-#endif // GGML_RPC_RDMA
+#endif // GGML_RPC_RDMA && !GGML_RPC_RDMA_APPLE

 struct socket_t::impl {
    impl(sockfd_t fd) : use_rdma(false), fd(fd) {}
    ~impl();
    bool send_data(const void * data, size_t size);
    bool recv_data(void * data, size_t size);
+    bool flush();
    void get_caps(uint8_t * local_caps);
    void update_caps(const uint8_t * remote_caps);

 #ifdef GGML_RPC_RDMA
-    bool tcp_peer_closed();
    std::optional<rdma_gid_t> rdma_build_target_gid();
+
+#  ifdef GGML_RPC_RDMA_APPLE
+    std::unique_ptr<apple_rdma> rdma;
+#  else
    bool rdma_probe();
-    bool rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid);
-    bool rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc);
    bool rdma_send(const void * data, size_t size);
    bool rdma_recv(void * data, size_t size);
+    bool tcp_peer_closed();
+    bool rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid);
+    bool rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc);

    std::unique_ptr<rdma_conn> rdma;
    rdma_local_info            rdma_local = {};
+#  endif
 #endif // GGML_RPC_RDMA
    bool     use_rdma;
    sockfd_t fd;
@@ -151,17 +165,6 @@ socket_t::impl::~impl() {

 #ifdef GGML_RPC_RDMA

-bool socket_t::impl::tcp_peer_closed() {
-    if (fd < 0) return false;
-#ifndef _WIN32
-    struct pollfd pfd = { fd, POLLIN | POLLRDHUP, 0 };
-    int r = poll(&pfd, 1, 0);
-    return r > 0 && (pfd.revents & (POLLHUP | POLLERR | POLLRDHUP));
-#else
-    return false;
-#endif
-}
-
 // Build a RoCE GID-shaped 16-byte target from a TCP socket's local address.
 // Used to match the socket's local IP against the kernel's GID table so that
 // a single memcmp handles IPv4, IPv4-mapped IPv6, and native IPv6 uniformly:
@@ -191,6 +194,19 @@ std::optional<rdma_gid_t> socket_t::impl::rdma_build_target_gid() {
    return std::nullopt;
 }

+#ifndef GGML_RPC_RDMA_APPLE
+
+bool socket_t::impl::tcp_peer_closed() {
+    if (fd < 0) return false;
+#ifndef _WIN32
+    struct pollfd pfd = { fd, POLLIN | POLLRDHUP, 0 };
+    int r = poll(&pfd, 1, 0);
+    return r > 0 && (pfd.revents & (POLLHUP | POLLERR | POLLRDHUP));
+#else
+    return false;
+#endif
+}
+
 bool socket_t::impl::rdma_probe() {
    const char * dev_env = std::getenv("GGML_RDMA_DEV");
    const char * gid_env = std::getenv("GGML_RDMA_GID");
@@ -457,10 +473,16 @@ bool socket_t::impl::rdma_recv(void * data, size_t size) {
    return true;
 }

+#endif // !GGML_RPC_RDMA_APPLE (Linux RC transport)
+
 #endif // GGML_RPC_RDMA

 bool socket_t::impl::send_data(const void * data, size_t size) {
-#ifdef GGML_RPC_RDMA
+#ifdef GGML_RPC_RDMA_APPLE
+    if (use_rdma) {
+        return rdma->send(data, size);
+    }
+#elif defined(GGML_RPC_RDMA)
    if (use_rdma) {
        return rdma_send(data, size);
    }
@@ -480,7 +502,11 @@ bool socket_t::impl::send_data(const void * data, size_t size) {
 }

 bool socket_t::impl::recv_data(void * data, size_t size) {
-#ifdef GGML_RPC_RDMA
+#ifdef GGML_RPC_RDMA_APPLE
+    if (use_rdma) {
+        return rdma->recv(data, size);
+    }
+#elif defined(GGML_RPC_RDMA)
    if (use_rdma) {
        return rdma_recv(data, size);
    }
@@ -506,6 +532,15 @@ bool socket_t::impl::recv_data(void * data, size_t size) {
 void socket_t::impl::get_caps(uint8_t * local_caps) {
    memset(local_caps, 0, RPC_CONN_CAPS_SIZE);
 #ifdef GGML_RPC_RDMA
+    if (std::getenv("GGML_RPC_NO_RDMA")) {
+        return;
+    }
+#  ifdef GGML_RPC_RDMA_APPLE
+    auto target_gid = rdma_build_target_gid();
+    if (target_gid) {
+        rdma = apple_rdma::probe(fd, target_gid->data(), local_caps);
+    }
+#  else
    rdma_local = {};
    if (rdma_probe()) {
        rdma_caps rc = {};
@@ -516,21 +551,30 @@ void socket_t::impl::get_caps(uint8_t * local_caps) {
    } else {
        rdma.reset();
    }
+#  endif
 #endif // GGML_RPC_RDMA
 }

 void socket_t::impl::update_caps(const uint8_t * remote_caps) {
 #ifdef GGML_RPC_RDMA
-    if (!rdma) {
-        return;
+    // a peer that has no RDMA advertises all-zero caps and takes no further part
+    // in the negotiation, so drop to TCP without reporting a failure
+    bool remote_rdma = false;
+    for (size_t i = 0; i < RPC_CONN_CAPS_SIZE; i++) {
+        remote_rdma |= remote_caps[i] != 0;
    }
-    rdma_caps rc = {};
-    memcpy(&rc, remote_caps, sizeof(rc));
-    if (rc.qpn == 0) {
+    if (!rdma || !remote_rdma) {
        rdma.reset();
        return;
    }
-    if (rdma_activate(rc.qpn, rc.psn, rc.gid)) {
+#  ifdef GGML_RPC_RDMA_APPLE
+    bool activated = rdma->activate(remote_caps);
+#  else
+    rdma_caps rc = {};
+    memcpy(&rc, remote_caps, sizeof(rc));
+    bool activated = rdma_activate(rc.qpn, rc.psn, rc.gid);
+#  endif
+    if (activated) {
        use_rdma = true;
    } else {
        GGML_LOG_ERROR("RDMA activate failed, staying on TCP\n");
@@ -541,6 +585,14 @@ void socket_t::impl::update_caps(const uint8_t * remote_caps) {
 #endif // GGML_RPC_RDMA
 }

+bool socket_t::impl::flush() {
+#ifdef GGML_RPC_RDMA_APPLE
+    if (use_rdma) {
+        return rdma->flush();
+    }
+#endif
+    return true;
+}

 /////////////////////////////////////////////////////////////////////////////

@@ -556,6 +608,10 @@ bool socket_t::recv_data(void * data, size_t size) {
    return pimpl->recv_data(data, size);
 }

+bool socket_t::flush() {
+    return pimpl->flush();
+}
+
 void socket_t::get_caps(uint8_t * local_caps) {
    return pimpl->get_caps(local_caps);
 }
@@ -15,6 +15,10 @@ struct socket_t {

    bool send_data(const void * data, size_t size);
    bool recv_data(void * data, size_t size);
+    // Must be called at every message boundary: the RDMA transport coalesces
+    // writes into fixed-size frames and posts the trailing partial frame only
+    // here. No-op on TCP.
+    bool flush();

    socket_ptr accept();

@@ -767,6 +767,21 @@ static constexpr std::initializer_list<std::array<int, 3>> rms_norm_mul_rope_vie
    { 4, 0, 3 }, // set_rows->src[0] == view
 };

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

 struct vk_device_struct {
    std::recursive_mutex mutex;
@@ -1042,6 +1057,8 @@ struct vk_device_struct {
    vk_pipeline pipeline_argsort_large_f32[num_argsort_pipelines];
    vk_pipeline pipeline_topk_f32[num_topk_pipelines];
    vk_pipeline pipeline_sum_rows_f32;
+    vk_pipeline pipeline_cross_entropy_loss_f32, pipeline_cross_entropy_loss_f32_wg512;
+    vk_pipeline pipeline_cross_entropy_loss_back_f32, pipeline_cross_entropy_loss_back_f32_wg512;
    vk_pipeline pipeline_fwht_f32[4];
    vk_pipeline pipeline_cumsum_f32;
    vk_pipeline pipeline_cumsum_small_f32;
@@ -1066,6 +1083,7 @@ struct vk_device_struct {
    vk_pipeline pipeline_rwkv_wkv6_f32;
    vk_pipeline pipeline_rwkv_wkv7_f32;
    vk_pipeline pipeline_gated_linear_attn_f32;
+    vk_pipeline pipeline_lightning_indexer_f32[GGML_TYPE_COUNT];
    // [size_idx][kda] where size_idx: 0=d16, 1=d32, 2=d64, 3=d128
    vk_pipeline pipeline_gated_delta_net[4][2];
    vk_pipeline pipeline_ssm_scan_f32_d128;
@@ -1846,6 +1864,26 @@ struct vk_op_gated_linear_attn_push_constants {
    uint32_t H;
    float scale;
 };
+struct vk_op_lightning_indexer_push_constants {
+    uint32_t n_kv;
+    uint32_t n_heads;
+    uint32_t n_tokens;
+    uint32_t n_streams;
+    uint32_t n_masks;
+    uint32_t dispatch_x;
+    uint32_t q_nb1;
+    uint32_t q_nb2;
+    uint32_t q_nb3;
+    uint32_t k_nb2;
+    uint32_t k_nb3;
+    uint32_t w_nb1;
+    uint32_t w_nb3;
+    uint32_t m_nb1;
+    uint32_t m_nb3;
+    uint32_t d_nb1;
+    uint32_t d_nb3;
+};
+static_assert(sizeof(vk_op_lightning_indexer_push_constants) <= 128);
 struct vk_op_gated_delta_net_push_constants {
    uint32_t H;
    uint32_t n_tokens;
@@ -3902,11 +3940,16 @@ static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const
    return vk_fa_pipeline_state{hsk, hsv, params.block_rows, params.block_cols, params.d_split, params.row_split, params.shmem_staging, params.path, params.workgroup_size, subgroup_size, aligned, f32acc, flags, params.limit_occupancy_shmem, k_type, v_type};
 }

+// Bytes per buffer block for the FaBlockBytesK/V spec constants. F32 is fed as
+// a vec4 "block" of 4 floats, everything else uses its ggml block size.
+static uint32_t fa_block_bytes(ggml_type t) {
+    if (t == GGML_TYPE_F32) {
+        return 16u;
+    }
+    return (uint32_t) ggml_type_size(t);
+}
+
 static std::vector<uint32_t> get_fa_spec_constants(const vk_fa_pipeline_state& state) {
-    const auto fa_block_bytes = [](ggml_type t) -> uint32_t {
-        if (t == GGML_TYPE_F32) return 16u;
-        return (uint32_t) ggml_type_size(t);
-    };
    return {
        /* 0 WorkGroupSize   */ state.workgroup_size,
        /* 1 Br              */ state.Br,
@@ -4169,10 +4212,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 +4302,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 +4342,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 +5221,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 +5807,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);
@@ -5835,6 +5888,17 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {

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

+    {
+        const bool li_subgroup = device->subgroup_arithmetic && device->subgroup_require_full_support;
+        const size_t li_len   = li_subgroup ? lightning_indexer_subgroup_f32_len  : lightning_indexer_f32_len;
+        const void * li_data  = li_subgroup ? (const void *)lightning_indexer_subgroup_f32_data : (const void *)lightning_indexer_f32_data;
+
+        for (ggml_type k_type : lightning_indexer_k_types) {
+            const std::string name = "lightning_indexer_" + std::string(ggml_type_name(k_type)) + "_k_f32";
+            ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f32[k_type], name.c_str(), li_len, li_data, "main", 5, sizeof(vk_op_lightning_indexer_push_constants), {1, 1, 1}, {(uint32_t)k_type, fa_block_bytes(k_type), device->subgroup_size}, 1, true, li_subgroup);
+        }
+    }
+
    {
        const uint32_t gdn_sizes[] = {16, 32, 64, 128};
        const char * gdn_names[][2] = {
@@ -11577,6 +11641,17 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
            return ctx->device->pipeline_sum_rows_f32;
        }
        return nullptr;
+    case GGML_OP_CROSS_ENTROPY_LOSS:
+        if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+            return src0->ne[0] > 1024 ? ctx->device->pipeline_cross_entropy_loss_f32_wg512 : ctx->device->pipeline_cross_entropy_loss_f32;
+        }
+        return nullptr;
+    case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
+        // src0 is the scalar grad; src1 is logits
+        if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && src2 && src2->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+            return src1->ne[0] > 1024 ? ctx->device->pipeline_cross_entropy_loss_back_f32_wg512 : ctx->device->pipeline_cross_entropy_loss_back_f32;
+        }
+        return nullptr;
    case GGML_OP_CUMSUM:
        if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
            if (src0->ne[0] <= 512) {
@@ -11674,6 +11749,12 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
            return ctx->device->pipeline_gated_linear_attn_f32;
        }
        return nullptr;
+    case GGML_OP_LIGHTNING_INDEXER:
+        // only the k type selects a pipeline, the other types are fixed by ggml_lightning_indexer()
+        if (ggml_vk_lightning_indexer_k_type_supported(src1->type)) {
+            return ctx->device->pipeline_lightning_indexer_f32[src1->type];
+        }
+        return nullptr;
    case GGML_OP_GATED_DELTA_NET:
        if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
            const uint32_t S_v = dst->src[2]->ne[0];
@@ -12749,6 +12830,55 @@ static void ggml_vk_gated_linear_attn(ggml_backend_vk_context * ctx, vk_context&
        pc, { (uint32_t)(n_seqs * n_heads), 1, 1 });
 }

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

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

        break;

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

@@ -18511,6 +18751,18 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
            }
        case GGML_OP_ARGMAX:
            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
+        case GGML_OP_CROSS_ENTROPY_LOSS:
+            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32
+                && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32
+                && ggml_are_same_shape(op->src[0], op->src[1])
+                && ggml_is_contiguous(op) && ggml_is_scalar(op) && op->type == GGML_TYPE_F32;
+        case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
+            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32 && ggml_is_scalar(op->src[0])
+                && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32
+                && ggml_is_contiguous(op->src[2]) && op->src[2]->type == GGML_TYPE_F32
+                && ggml_are_same_shape(op->src[1], op->src[2])
+                && ggml_are_same_shape(op->src[1], op)
+                && ggml_is_contiguous(op) && op->type == GGML_TYPE_F32;
        case GGML_OP_COUNT_EQUAL:
            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_I32
                && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_I32;
@@ -18536,6 +18788,40 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
        case GGML_OP_GATED_LINEAR_ATTN:
            // the shader block size is hardcoded to head_size 64
            return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && op->src[0]->ne[0] == 64;
+        case GGML_OP_LIGHTNING_INDEXER:
+            {
+                const ggml_tensor * q = op->src[0];
+                const ggml_tensor * k = op->src[1];
+                const ggml_tensor * w = op->src[2];
+                const ggml_tensor * m = op->src[3];
+
+                // the q/w/m types and the shape relationships between q, k, w, m and dst
+                // are already asserted in ggml_lightning_indexer()
+                if (!ggml_vk_lightning_indexer_k_type_supported(k->type) || !device->fp16) {
+                    return false;
+                }
+
+                // the shader block size is hardcoded to head size 128
+                if (q->ne[0] != 128) {
+                    return false;
+                }
+
+                // the shader indexes the buffers by element stride, and is dispatched
+                // without allow_misalign
+                for (const ggml_tensor * t : {q, k, w, m, op}) {
+                    if (t->nb[0] != ggml_type_size(t->type) ||
+                        (vk_tensor_offset(t) + t->view_offs) % device->properties.limits.minStorageBufferOffsetAlignment != 0) {
+                        return false;
+                    }
+                    // the strides get scaled down from bytes, so the division must be exact
+                    for (int i = 1; i < GGML_MAX_DIMS; ++i) {
+                        if (t->nb[i] % ggml_type_size(t->type) != 0) {
+                            return false;
+                        }
+                    }
+                }
+                return true;
+            }
        case GGML_OP_GATED_DELTA_NET:
            {
                const uint32_t S_v = op->src[2]->ne[0];
@@ -19437,6 +19723,10 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
            tensor_clone = ggml_mean(ggml_ctx, src_clone[0]);
        } else if (tensor->op == GGML_OP_ARGMAX) {
            tensor_clone = ggml_argmax(ggml_ctx, src_clone[0]);
+        } else if (tensor->op == GGML_OP_CROSS_ENTROPY_LOSS) {
+            tensor_clone = ggml_cross_entropy_loss(ggml_ctx, src_clone[0], src_clone[1]);
+        } else if (tensor->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK) {
+            tensor_clone = ggml_cross_entropy_loss_back(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]);
        } else if (tensor->op == GGML_OP_COUNT_EQUAL) {
            tensor_clone = ggml_count_equal(ggml_ctx, src_clone[0], src_clone[1]);
        } else if (tensor->op == GGML_OP_SOLVE_TRI) {
@@ -19541,6 +19831,8 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
            const float * op_params = (const float *)tensor->op_params;
            tensor_clone = ggml_gated_linear_attn(ggml_ctx, src_clone[0], src_clone[1],
            src_clone[2], src_clone[3], src_clone[4], op_params[0]);
+        } else if (tensor->op == GGML_OP_LIGHTNING_INDEXER) {
+            tensor_clone = ggml_lightning_indexer(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3]);
        } else if (tensor->op == GGML_OP_GATED_DELTA_NET) {
            tensor_clone = ggml_gated_delta_net(ggml_ctx, src_clone[0], src_clone[1],
            src_clone[2], src_clone[3], src_clone[4], src_clone[5],
@@ -0,0 +1,78 @@
+#version 450
+
+#include "generic_head.glsl"
+#include "types.glsl"
+
+#extension GL_EXT_control_flow_attributes : enable
+
+layout(constant_id = 0) const uint BLOCK_SIZE = 32;
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
+layout (binding = 1) readonly buffer B {B_TYPE data_b[];};
+layout (binding = 2) writeonly buffer D {D_TYPE data_d[];};
+
+shared FLOAT_TYPE tmp[BLOCK_SIZE];
+
+FLOAT_TYPE wg_reduce_max(FLOAT_TYPE v) {
+    const uint tid = gl_LocalInvocationID.x;
+    tmp[tid] = v;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            tmp[tid] = max(tmp[tid], tmp[tid + s]);
+        }
+        barrier();
+    }
+    v = tmp[0];
+    barrier();
+    return v;
+}
+
+FLOAT_TYPE wg_reduce_sum(FLOAT_TYPE v) {
+    const uint tid = gl_LocalInvocationID.x;
+    tmp[tid] = v;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            tmp[tid] += tmp[tid + s];
+        }
+        barrier();
+    }
+    v = tmp[0];
+    barrier();
+    return v;
+}
+
+void main() {
+    const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
+    const uint tid = gl_LocalInvocationID.x;
+
+    if (row >= p.KY) {
+        return;
+    }
+
+    const uint off = row * p.KX;
+
+    FLOAT_TYPE max_logit = FLOAT_TYPE(uintBitsToFloat(0xFF800000));
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        max_logit = max(max_logit, FLOAT_TYPE(data_a[off + i]));
+    }
+    max_logit = wg_reduce_max(max_logit);
+
+    FLOAT_TYPE sum_exp = FLOAT_TYPE(0.0f);
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        sum_exp += exp(FLOAT_TYPE(data_a[off + i]) - max_logit);
+    }
+    const FLOAT_TYPE log_sum = log(wg_reduce_sum(sum_exp));
+
+    FLOAT_TYPE loss = FLOAT_TYPE(0.0f);
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        loss += (FLOAT_TYPE(data_a[off + i]) - max_logit - log_sum) * FLOAT_TYPE(data_b[off + i]);
+    }
+    loss = -wg_reduce_sum(loss) / FLOAT_TYPE(p.KY);
+
+    if (tid == 0) {
+        data_d[row] = D_TYPE(loss);
+    }
+}
@@ -0,0 +1,75 @@
+#version 450
+
+#include "generic_head.glsl"
+#include "types.glsl"
+
+#extension GL_EXT_control_flow_attributes : enable
+
+layout(constant_id = 0) const uint BLOCK_SIZE = 32;
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer G {A_TYPE data_g[];};
+layout (binding = 1) readonly buffer X {B_TYPE data_x[];};
+layout (binding = 2) readonly buffer Y {B_TYPE data_y[];};
+layout (binding = 3) writeonly buffer D {D_TYPE data_d[];};
+
+shared FLOAT_TYPE tmp[BLOCK_SIZE];
+
+FLOAT_TYPE wg_reduce_max(FLOAT_TYPE v) {
+    const uint tid = gl_LocalInvocationID.x;
+    tmp[tid] = v;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            tmp[tid] = max(tmp[tid], tmp[tid + s]);
+        }
+        barrier();
+    }
+    v = tmp[0];
+    barrier();
+    return v;
+}
+
+FLOAT_TYPE wg_reduce_sum(FLOAT_TYPE v) {
+    const uint tid = gl_LocalInvocationID.x;
+    tmp[tid] = v;
+    barrier();
+    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            tmp[tid] += tmp[tid + s];
+        }
+        barrier();
+    }
+    v = tmp[0];
+    barrier();
+    return v;
+}
+
+void main() {
+    const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
+    const uint tid = gl_LocalInvocationID.x;
+
+    if (row >= p.KY) {
+        return;
+    }
+
+    const uint off = row * p.KX;
+    const FLOAT_TYPE d_by_nrows = FLOAT_TYPE(data_g[0]) / FLOAT_TYPE(p.KY);
+
+    FLOAT_TYPE max_logit = FLOAT_TYPE(uintBitsToFloat(0xFF800000));
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        max_logit = max(max_logit, FLOAT_TYPE(data_x[off + i]));
+    }
+    max_logit = wg_reduce_max(max_logit);
+
+    FLOAT_TYPE sum_exp = FLOAT_TYPE(0.0f);
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        sum_exp += exp(FLOAT_TYPE(data_x[off + i]) - max_logit);
+    }
+    const FLOAT_TYPE inv_sum = FLOAT_TYPE(1.0f) / wg_reduce_sum(sum_exp);
+
+    for (uint i = tid; i < p.KX; i += BLOCK_SIZE) {
+        const FLOAT_TYPE sm = exp(FLOAT_TYPE(data_x[off + i]) - max_logit) * inv_sum;
+        data_d[off + i] = D_TYPE((sm - FLOAT_TYPE(data_y[off + i])) * d_by_nrows);
+    }
+}
@@ -0,0 +1,55 @@
+#if !defined(GGML_FA_TYPES_COMP)
+#define GGML_FA_TYPES_COMP
+
+// FaTypeK / FaTypeV spec constant values. These mirror enum ggml_type so the
+// host can pass the type directly. Keep in sync with ggml.h.
+#define FA_TYPE_F32   0u
+#define FA_TYPE_F16   1u
+#define FA_TYPE_Q4_0  2u
+#define FA_TYPE_Q4_1  3u
+#define FA_TYPE_Q5_0  6u
+#define FA_TYPE_Q5_1  7u
+#define FA_TYPE_Q8_0  8u
+#define FA_TYPE_IQ4_NL 20u
+#define FA_TYPE_BF16 30u
+
+// Number of matrix elements per buffer block, derived from the K/V type spec
+// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1
+// and bypasses the dequant path entirely. Quants follow their ggml block sizes.
+uint fa_block_elems(uint ty) {
+    switch (ty) {
+        case FA_TYPE_F32:  return 4u;
+        case FA_TYPE_F16:  return 1u;
+        case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0);
+        case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1);
+        case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0);
+        case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1);
+        case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0);
+        case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL);
+        case FA_TYPE_BF16: return 1u;
+        default:           return 1u;
+    }
+}
+
+// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte
+// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number
+// of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R.
+uint fa_quant_r_mmq(uint ty) {
+    switch (ty) {
+        case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0);
+        case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1);
+        case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0);
+        case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1);
+        case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0);
+        default:           return 1u;
+    }
+}
+
+bool fa_type_needs_shmem(uint ty) {
+    switch (ty) {
+        case FA_TYPE_IQ4_NL: return true;
+        default:             return false;
+    }
+}
+
+#endif // !defined(GGML_FA_TYPES_COMP)
@@ -88,17 +88,7 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
 #define BINDING_IDX_K 0
 #define BINDING_IDX_V 1

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

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

-// Number of matrix elements per buffer block, derived from the K/V type spec
-// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1
-// and bypasses the dequant path entirely. Quants follow their ggml block sizes.
-uint fa_block_elems(uint ty) {
-    switch (ty) {
-        case FA_TYPE_F32:  return 4u;
-        case FA_TYPE_F16:  return 1u;
-        case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0);
-        case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1);
-        case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0);
-        case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1);
-        case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0);
-        case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL);
-        case FA_TYPE_BF16: return 1u;
-        default:           return 1u;
-    }
-}
-
-// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte
-// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number
-// of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R.
-uint fa_quant_r_mmq(uint ty) {
-    switch (ty) {
-        case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0);
-        case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1);
-        case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0);
-        case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1);
-        case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0);
-        default:           return 1u;
-    }
-}
-
-bool fa_type_needs_shmem(uint ty) {
-    switch (ty) {
-        case FA_TYPE_IQ4_NL: return true;
-        default:             return false;
-    }
-}
-
 // These can't be `const` globals because GLSL forbids function calls in global
 // const initializers, even when the spec constants would let the driver fold
 // them. Macros expand at the use site and fold after specialization.
@@ -0,0 +1,151 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : require
+#extension GL_EXT_shader_16bit_storage : require
+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
+#extension GL_KHR_shader_subgroup_basic : enable
+#if USE_SUBGROUP_ADD
+#extension GL_KHR_shader_subgroup_arithmetic : enable
+#endif
+
+#define BINDING_IDX_K 0u
+
+#include "types.glsl"
+#include "fa_types.glsl"
+#define FaTypeV FA_TYPE_F32
+
+layout(constant_id = 0) const uint FaTypeK = FA_TYPE_F32;
+layout(constant_id = 1) const uint FaBlockBytesK = 4;
+layout(constant_id = 2) const uint SUBGROUP_SIZE = 32;
+
+#include "flash_attn_dequant.glsl"
+
+// one workgroup computes one output element, one invocation per head element
+#define HEAD_SIZE 128
+
+layout(local_size_x = HEAD_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+layout(binding = 0) readonly buffer QBuf { float q[]; };
+layout(binding = 1) readonly buffer KBufF16 { float16_t k_f16[]; };
+layout(binding = 1) readonly buffer KBufF32 { float k_f32[]; };
+layout(binding = 1) readonly buffer KBufBF16 { uint16_t k_bf16[]; };
+layout(binding = 2) readonly buffer WBuf { float weights[]; };
+layout(binding = 3) readonly buffer MBuf { float16_t mask[]; };
+layout(binding = 4) writeonly buffer DstBuf { float dst[]; };
+
+layout(push_constant) uniform PushConstants {
+    uint n_kv;
+    uint n_heads;
+    uint n_tokens;
+    uint n_streams;
+    uint n_masks;
+    uint dispatch_x;
+    uint q_nb1;
+    uint q_nb2;
+    uint q_nb3;
+    uint k_nb2;
+    uint k_nb3;
+    uint w_nb1;
+    uint w_nb3;
+    uint m_nb1;
+    uint m_nb3;
+    uint d_nb1;
+    uint d_nb3;
+};
+
+shared float k_row[HEAD_SIZE];
+
+#if USE_SUBGROUP_ADD
+shared float sg_partials[HEAD_SIZE / SUBGROUP_SIZE];
+#else
+shared float partials[HEAD_SIZE];
+#endif
+
+void main() {
+    const uint tid = gl_LocalInvocationID.x;
+    const uint output_idx = gl_WorkGroupID.y * dispatch_x + gl_WorkGroupID.x;
+    const uint n_outputs = n_kv * n_tokens * n_streams;
+
+    if (fa_type_needs_shmem(FaTypeK)) {
+        init_iq_shmem(gl_WorkGroupSize);
+    }
+
+    if (output_idx >= n_outputs) {
+        return;
+    }
+
+    const uint ik = output_idx % n_kv;
+    const uint ts = output_idx / n_kv;
+    const uint t = ts % n_tokens;
+    const uint s = ts / n_tokens;
+    const uint k_offset = ik * k_nb2 + s * k_nb3;
+
+    // k strides come in as bytes, so scale them down to the view being indexed
+    const uint k_block_elems = fa_block_elems(FaTypeK);
+    const uint k_elem_bytes = FaBlockBytesK / k_block_elems;
+
+    if (FaTypeK == FA_TYPE_F16) {
+        k_row[tid] = float(k_f16[k_offset / k_elem_bytes + tid]);
+    } else if (FaTypeK == FA_TYPE_F32) {
+        k_row[tid] = k_f32[k_offset / k_elem_bytes + tid];
+    } else if (FaTypeK == FA_TYPE_BF16) {
+        k_row[tid] = bf16_to_fp32(uint(k_bf16[k_offset / k_elem_bytes + tid]));
+    } else if (4 * tid < HEAD_SIZE) {
+        const uint coord = 4 * tid;
+        const uint ib = coord / k_block_elems;
+        const uint iqs = coord % k_block_elems;
+        const vec4 values = dequantize4(ib, iqs, k_offset / FaBlockBytesK, BINDING_IDX_K);
+        k_row[coord + 0] = values.x;
+        k_row[coord + 1] = values.y;
+        k_row[coord + 2] = values.z;
+        k_row[coord + 3] = values.w;
+    }
+    barrier();
+
+    const float k_val = k_row[tid];
+
+    float score = 0.0;
+    for (uint h = 0; h < n_heads; ++h) {
+        const float prod = q[h * q_nb1 + t * q_nb2 + s * q_nb3 + tid] * k_val;
+
+#if USE_SUBGROUP_ADD
+        const float sg_sum = subgroupAdd(prod);
+        if (gl_SubgroupInvocationID == 0) {
+            sg_partials[gl_SubgroupID] = sg_sum;
+        }
+        barrier();
+
+        if (tid == 0) {
+            float sum = 0.0;
+            [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) {
+                sum += sg_partials[i];
+            }
+            score += max(sum, 0.0) * weights[h + t * w_nb1 + s * w_nb3];
+        }
+        // the reads above must complete before the next iteration overwrites sg_partials
+        barrier();
+#else
+        partials[tid] = prod;
+        barrier();
+
+        [[unroll]] for (uint stride = HEAD_SIZE / 2; stride > 0; stride >>= 1) {
+            if (tid < stride) {
+                partials[tid] += partials[tid + stride];
+            }
+            barrier();
+        }
+
+        if (tid == 0) {
+            score += max(partials[0], 0.0) * weights[h + t * w_nb1 + s * w_nb3];
+        }
+        // the read of partials[0] above must complete before the next iteration
+        // overwrites partials[tid]
+        barrier();
+#endif
+    }
+
+    if (tid == 0) {
+        const uint mask_offset = ik + t * m_nb1 + (s % n_masks) * m_nb3;
+        dst[ik + t * d_nb1 + s * d_nb3] = score + float(mask[mask_offset]);
+    }
+}
@@ -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"}}));
@@ -1067,6 +1069,12 @@ void process_shaders() {

    string_to_spv("gated_linear_attn_f32", "gla.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));

+    // Compile IQ4_NL support in so its shared LUT is available when K uses it.
+    // K quant type is selected at runtime via the FaTypeK spec constant.
+    std::map<std::string, std::string> li_dict = {{"FLOAT_TYPE", "float"}, {"FLOAT_TYPEV4", "vec4"}, {"DATA_A_IQ4_NL", "1"}};
+    string_to_spv("lightning_indexer_f32", "lightning_indexer.comp", li_dict);
+    string_to_spv("lightning_indexer_subgroup_f32", "lightning_indexer.comp", merge_maps(li_dict, {{"USE_SUBGROUP_ADD", "1"}}));
+
    string_to_spv("rwkv_wkv7_f32", "wkv7.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));

    string_to_spv("gated_delta_net_f32", "gated_delta_net.comp", merge_maps(base_dict, {{"FLOAT_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}, {"USE_SUBGROUP_CLUSTERED", "1"}}));
@@ -162,7 +162,12 @@ class Keys:
        TARGET_LAYERS                     = "{arch}.target_layers"
        TARGET_HIDDEN_SIZE                = "{arch}.target_hidden_size"
        BLOCK_SIZE                        = "{arch}.block_size"
+        CONV_KERNEL_SIZE                  = "{arch}.conv_kernel_size"
+        CONV_GROUP_SIZE                   = "{arch}.conv_group_size"
+        SELECTOR_RANK                     = "{arch}.selector_rank"
+        SELECTOR_TOP_K                    = "{arch}.selector_top_k"
        SAMPLE_FROM_ANCHOR                = "{arch}.sample_from_anchor"
+        HAS_CONFIDENCE_HEAD               = "{arch}.has_confidence_head"
        NORM_BEFORE_RESIDUAL              = "{arch}.norm_before_residual"
        NORM_BEFORE_FC                    = "{arch}.norm_before_fc"

@@ -225,6 +230,19 @@ class Keys:
        COUNT                = "{arch}.hyper_connection.count"
        SINKHORN_ITERATIONS  = "{arch}.hyper_connection.sinkhorn_iterations"
        EPSILON              = "{arch}.hyper_connection.epsilon"
+        # absent means the mix projection is full rank (DeepSeek-V4 behaviour)
+        LOW_RANK             = "{arch}.hyper_connection.low_rank"
+
+    class PerLayerEmbedding:
+        LAYERS             = "{arch}.ple.layers"
+        NGRAM_SIZE         = "{arch}.ple.ngram_size"
+        HEADS_PER_NGRAM    = "{arch}.ple.heads_per_ngram"
+        CONV_KERNEL        = "{arch}.ple.conv_kernel"
+        LAYER_MULTIPLIERS  = "{arch}.ple.layer_multipliers"
+        HEAD_OFFSETS       = "{arch}.ple.head_offsets"
+        HEAD_VOCAB_SIZES   = "{arch}.ple.head_vocab_sizes"
+        EOS_TOKEN_ID       = "{arch}.ple.eos_token_id"
+        IMAGE_TOKEN_ID     = "{arch}.ple.image_token_id"

    class Rope:
        DIMENSION_COUNT           = "{arch}.rope.dimension_count"
@@ -494,6 +512,7 @@ class MODEL_ARCH(IntEnum):
    QWEN3VLMOE       = auto()
    QWEN35           = auto()
    QWEN35MOE        = auto()
+    QWEN4EXP         = auto()
    PHI2             = auto()
    PHI3             = auto()
    PHIMOE           = auto()
@@ -636,6 +655,9 @@ class MODEL_TENSOR(IntEnum):
    HC_HEAD_FN           = auto()
    HC_HEAD_BASE         = auto()
    HC_HEAD_SCALE        = auto()
+    HC_HEAD_NORM         = auto() # qwen4exp
+    HC_HEAD_DOWN         = auto() # qwen4exp
+    HC_HEAD_UP           = auto() # qwen4exp
    ROPE_FREQS           = auto()
    ROPE_FACTORS_LONG    = auto()
    ROPE_FACTORS_SHORT   = auto()
@@ -780,6 +802,20 @@ class MODEL_TENSOR(IntEnum):
    HC_FFN_FN            = auto()
    HC_FFN_BASE          = auto()
    HC_FFN_SCALE         = auto()
+    HC_ATTN_NORM         = auto() # qwen4exp
+    HC_ATTN_DOWN         = auto() # qwen4exp
+    HC_ATTN_UP           = auto() # qwen4exp
+    HC_ATTN_INJECT       = auto() # qwen4exp
+    HC_FFN_NORM          = auto() # qwen4exp
+    HC_FFN_DOWN          = auto() # qwen4exp
+    HC_FFN_UP            = auto() # qwen4exp
+    HC_FFN_INJECT        = auto() # qwen4exp
+    PLE_KEY              = auto() # qwen4exp
+    PLE_VALUE            = auto() # qwen4exp
+    PLE_NORM_KEY         = auto() # qwen4exp
+    PLE_NORM_QUERY       = auto() # qwen4exp
+    PLE_NORM_CONV        = auto() # qwen4exp
+    PLE_CONV1D           = auto() # qwen4exp
    ATTN_COMPRESSOR_WKV  = auto()
    ATTN_COMPRESSOR_WGATE = auto()
    ATTN_COMPRESSOR_APE  = auto()
@@ -1146,6 +1182,13 @@ class MODEL_TENSOR(IntEnum):
    DSPARK_MARKOV_W1       = auto()  # markov head: prev-token embed
    DSPARK_MARKOV_W2       = auto()  # markov head: bias projection
    DSPARK_CONF_PROJ       = auto()  # confidence head
+    DFLASH_ATTN_CONV_BASE  = auto()
+    DFLASH_ATTN_CONV_PROJ  = auto()
+    DFLASH_FFN_CONV_BASE   = auto()
+    DFLASH_FFN_CONV_PROJ   = auto()
+    DFLASH_SELECTOR_PREV   = auto()
+    DFLASH_SELECTOR_NEXT   = auto()
+    DFLASH_SELECTOR_HIDDEN = auto()
    # lfm2 audio
    A_ENC_NORM_CONV        = auto()
    A_ENC_LINEAR_POS       = auto()
@@ -1217,6 +1260,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
    MODEL_ARCH.QWEN3VLMOE:       "qwen3vlmoe",
    MODEL_ARCH.QWEN35:           "qwen35",
    MODEL_ARCH.QWEN35MOE:        "qwen35moe",
+    MODEL_ARCH.QWEN4EXP:         "qwen4exp",
    MODEL_ARCH.PHI2:             "phi2",
    MODEL_ARCH.PHI3:             "phi3",
    MODEL_ARCH.PHIMOE:           "phimoe",
@@ -1358,6 +1402,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
    MODEL_TENSOR.HC_HEAD_FN:                "output_hc_fn",
    MODEL_TENSOR.HC_HEAD_BASE:              "output_hc_base",
    MODEL_TENSOR.HC_HEAD_SCALE:             "output_hc_scale",
+    MODEL_TENSOR.HC_HEAD_NORM:              "output_hc_norm",                 # qwen4exp
+    MODEL_TENSOR.HC_HEAD_DOWN:              "output_hc_down",                 # qwen4exp
+    MODEL_TENSOR.HC_HEAD_UP:                "output_hc_up",                   # qwen4exp
    MODEL_TENSOR.ROPE_FREQS:                "rope_freqs",
    MODEL_TENSOR.ROPE_FACTORS_LONG:         "rope_factors_long",
    MODEL_TENSOR.ROPE_FACTORS_SHORT:        "rope_factors_short",
@@ -1502,6 +1549,20 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
    MODEL_TENSOR.HC_FFN_FN:                 "blk.{bid}.hc_ffn_fn",
    MODEL_TENSOR.HC_FFN_BASE:               "blk.{bid}.hc_ffn_base",
    MODEL_TENSOR.HC_FFN_SCALE:              "blk.{bid}.hc_ffn_scale",
+    MODEL_TENSOR.HC_ATTN_NORM:              "blk.{bid}.hc_attn_norm",         # qwen4exp
+    MODEL_TENSOR.HC_ATTN_DOWN:              "blk.{bid}.hc_attn_down",         # qwen4exp
+    MODEL_TENSOR.HC_ATTN_UP:                "blk.{bid}.hc_attn_up",           # qwen4exp
+    MODEL_TENSOR.HC_ATTN_INJECT:            "blk.{bid}.hc_attn_inject",       # qwen4exp
+    MODEL_TENSOR.HC_FFN_NORM:               "blk.{bid}.hc_ffn_norm",          # qwen4exp
+    MODEL_TENSOR.HC_FFN_DOWN:               "blk.{bid}.hc_ffn_down",          # qwen4exp
+    MODEL_TENSOR.HC_FFN_UP:                 "blk.{bid}.hc_ffn_up",            # qwen4exp
+    MODEL_TENSOR.HC_FFN_INJECT:             "blk.{bid}.hc_ffn_inject",        # qwen4exp
+    MODEL_TENSOR.PLE_KEY:                   "blk.{bid}.ple_key",              # qwen4exp
+    MODEL_TENSOR.PLE_VALUE:                 "blk.{bid}.ple_value",            # qwen4exp
+    MODEL_TENSOR.PLE_NORM_KEY:              "blk.{bid}.ple_norm_key",         # qwen4exp
+    MODEL_TENSOR.PLE_NORM_QUERY:            "blk.{bid}.ple_norm_query",       # qwen4exp
+    MODEL_TENSOR.PLE_NORM_CONV:             "blk.{bid}.ple_norm_conv",        # qwen4exp
+    MODEL_TENSOR.PLE_CONV1D:                "blk.{bid}.ple_conv1d",           # qwen4exp
    MODEL_TENSOR.ATTN_COMPRESSOR_WKV:       "blk.{bid}.attn_compressor_kv",
    MODEL_TENSOR.ATTN_COMPRESSOR_WGATE:     "blk.{bid}.attn_compressor_gate",
    MODEL_TENSOR.ATTN_COMPRESSOR_APE:       "blk.{bid}.attn_compressor_ape",
@@ -1895,6 +1956,13 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
    MODEL_TENSOR.DSPARK_MARKOV_W1:          "markov_w1",
    MODEL_TENSOR.DSPARK_MARKOV_W2:          "markov_w2",
    MODEL_TENSOR.DSPARK_CONF_PROJ:          "conf_proj",
+    MODEL_TENSOR.DFLASH_ATTN_CONV_BASE:     "blk.{bid}.attn_conv_base",
+    MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ:     "blk.{bid}.attn_conv_proj",
+    MODEL_TENSOR.DFLASH_FFN_CONV_BASE:      "blk.{bid}.ffn_conv_base",
+    MODEL_TENSOR.DFLASH_FFN_CONV_PROJ:      "blk.{bid}.ffn_conv_proj",
+    MODEL_TENSOR.DFLASH_SELECTOR_PREV:      "selector_predecessor",
+    MODEL_TENSOR.DFLASH_SELECTOR_NEXT:      "selector_successor",
+    MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN:    "selector_hidden",
    MODEL_TENSOR.D2T:                       "d2t",
 }

@@ -2795,6 +2863,58 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
        MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
    ],
+    MODEL_ARCH.QWEN4EXP: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT,
+        # no OUTPUT_NORM / ATTN_NORM / ATTN_POST_NORM: hyper-connections replace every layer norm
+        MODEL_TENSOR.HC_HEAD_NORM,
+        MODEL_TENSOR.HC_HEAD_DOWN,
+        MODEL_TENSOR.HC_HEAD_UP,
+        MODEL_TENSOR.HC_ATTN_NORM,
+        MODEL_TENSOR.HC_ATTN_DOWN,
+        MODEL_TENSOR.HC_ATTN_UP,
+        MODEL_TENSOR.HC_ATTN_INJECT,
+        MODEL_TENSOR.HC_FFN_NORM,
+        MODEL_TENSOR.HC_FFN_DOWN,
+        MODEL_TENSOR.HC_FFN_UP,
+        MODEL_TENSOR.HC_FFN_INJECT,
+        # full attention layers: ATTN_Q holds [q|gate] interleaved per head
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.INDEXER_Q_PROJ,
+        MODEL_TENSOR.INDEXER_K_PROJ,
+        MODEL_TENSOR.INDEXER_Q_NORM,
+        MODEL_TENSOR.INDEXER_K_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.SSM_BETA,
+        MODEL_TENSOR.SSM_ALPHA,
+        MODEL_TENSOR.SSM_OUT,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_INP_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_GATE_UP_EXP,
+        MODEL_TENSOR.PER_LAYER_TOKEN_EMBD,
+        MODEL_TENSOR.PLE_KEY,
+        MODEL_TENSOR.PLE_VALUE,
+        MODEL_TENSOR.PLE_NORM_KEY,
+        MODEL_TENSOR.PLE_NORM_QUERY,
+        MODEL_TENSOR.PLE_NORM_CONV,
+        MODEL_TENSOR.PLE_CONV1D,
+    ],
    MODEL_ARCH.PLAMO: [
        MODEL_TENSOR.TOKEN_EMBD,
        MODEL_TENSOR.OUTPUT_NORM,
@@ -4953,6 +5073,13 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
        MODEL_TENSOR.DSPARK_MARKOV_W1,
        MODEL_TENSOR.DSPARK_MARKOV_W2,
        MODEL_TENSOR.DSPARK_CONF_PROJ,
+        MODEL_TENSOR.DFLASH_ATTN_CONV_BASE,
+        MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ,
+        MODEL_TENSOR.DFLASH_FFN_CONV_BASE,
+        MODEL_TENSOR.DFLASH_FFN_CONV_PROJ,
+        MODEL_TENSOR.DFLASH_SELECTOR_PREV,
+        MODEL_TENSOR.DFLASH_SELECTOR_NEXT,
+        MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN,
    ],
    MODEL_ARCH.MISTRAL4: [
        MODEL_TENSOR.TOKEN_EMBD,
@@ -993,9 +993,24 @@ class GGUFWriter:
    def add_block_size(self, value: int) -> None:
        self.add_uint32(Keys.LLM.BLOCK_SIZE.format(arch=self.arch), value)

+    def add_conv_kernel_size(self, value: int) -> None:
+        self.add_uint32(Keys.LLM.CONV_KERNEL_SIZE.format(arch=self.arch), value)
+
+    def add_conv_group_size(self, value: int) -> None:
+        self.add_uint32(Keys.LLM.CONV_GROUP_SIZE.format(arch=self.arch), value)
+
+    def add_selector_rank(self, value: int) -> None:
+        self.add_uint32(Keys.LLM.SELECTOR_RANK.format(arch=self.arch), value)
+
+    def add_selector_top_k(self, value: int) -> None:
+        self.add_uint32(Keys.LLM.SELECTOR_TOP_K.format(arch=self.arch), value)
+
    def add_sample_from_anchor(self, value: bool) -> None:
        self.add_bool(Keys.LLM.SAMPLE_FROM_ANCHOR.format(arch=self.arch), value)

+    def add_has_confidence_head(self, value: bool) -> None:
+        self.add_bool(Keys.LLM.HAS_CONFIDENCE_HEAD.format(arch=self.arch), value)
+
    def add_target_layers(self, value: Sequence[int]) -> None:
        self.add_array(Keys.LLM.TARGET_LAYERS.format(arch=self.arch), value)

@@ -1029,6 +1044,40 @@ class GGUFWriter:
    def add_hyper_connection_epsilon(self, value: float) -> None:
        self.add_float32(Keys.HyperConnection.EPSILON.format(arch=self.arch), value)

+    def add_hyper_connection_low_rank(self, value: int) -> None:
+        self.add_uint32(Keys.HyperConnection.LOW_RANK.format(arch=self.arch), value)
+
+    def add_ple_layers(self, values: Sequence[int]) -> None:
+        self.add_array(Keys.PerLayerEmbedding.LAYERS.format(arch=self.arch), values)
+
+    def add_ple_ngram_size(self, value: int) -> None:
+        self.add_uint32(Keys.PerLayerEmbedding.NGRAM_SIZE.format(arch=self.arch), value)
+
+    def add_ple_heads_per_ngram(self, value: int) -> None:
+        self.add_uint32(Keys.PerLayerEmbedding.HEADS_PER_NGRAM.format(arch=self.arch), value)
+
+    def add_ple_conv_kernel(self, value: int) -> None:
+        self.add_uint32(Keys.PerLayerEmbedding.CONV_KERNEL.format(arch=self.arch), value)
+
+    # multipliers reach ~2.4e13; default INT32 inference would truncate them
+    def _add_u64_array(self, key: str, values: Sequence[int]) -> None:
+        self.add_key_value(key, list(values), GGUFValueType.ARRAY, GGUFValueType.UINT64)
+
+    def add_ple_layer_multipliers(self, values: Sequence[int]) -> None:
+        self._add_u64_array(Keys.PerLayerEmbedding.LAYER_MULTIPLIERS.format(arch=self.arch), values)
+
+    def add_ple_head_offsets(self, values: Sequence[int]) -> None:
+        self._add_u64_array(Keys.PerLayerEmbedding.HEAD_OFFSETS.format(arch=self.arch), values)
+
+    def add_ple_head_vocab_sizes(self, values: Sequence[int]) -> None:
+        self._add_u64_array(Keys.PerLayerEmbedding.HEAD_VOCAB_SIZES.format(arch=self.arch), values)
+
+    def add_ple_eos_token_id(self, value: int) -> None:
+        self.add_uint32(Keys.PerLayerEmbedding.EOS_TOKEN_ID.format(arch=self.arch), value)
+
+    def add_ple_image_token_id(self, value: int) -> None:
+        self.add_uint32(Keys.PerLayerEmbedding.IMAGE_TOKEN_ID.format(arch=self.arch), value)
+
    def add_attention_scale(self, value: float) -> None:
        self.add_float32(Keys.Attention.SCALE.format(arch=self.arch), value)

@@ -226,3 +226,64 @@ class LazyNumpyTensor(LazyBase):
        return eager.tofile(*args, **kwargs)

    # TODO: __array_function__
+
+
+# Tensor written to file one row-chunk at a time
+class LazyChunkedTensor:
+
+    def __init__(
+        self, chunks: list[Callable[[], np.ndarray]], shape: tuple[int, ...], dtype: DTypeLike,
+        qtype: Any = None, byteswap: bool = False,
+    ):
+        self._chunks = chunks
+        self._qtype = qtype
+        self._byteswap = byteswap
+        self.shape = tuple(shape)
+        self.dtype = np.dtype(dtype)
+
+    @property
+    def nbytes(self) -> int:
+        n = self.dtype.itemsize
+        for d in self.shape:
+            n *= d
+        return n
+
+    def numpy(self) -> LazyChunkedTensor:
+        return self
+
+    def quantize(self, qtype: Any) -> LazyChunkedTensor:
+        from .constants import GGMLQuantizationType
+        from .quants import QuantError, quant_shape_to_byte_shape
+
+        if qtype == GGMLQuantizationType.F32:
+            shape, dtype = self.shape, np.dtype(np.float32)
+        elif qtype == GGMLQuantizationType.F16:
+            shape, dtype = self.shape, np.dtype(np.float16)
+        else:
+            try:
+                shape, dtype = quant_shape_to_byte_shape(self.shape, qtype), np.dtype(np.uint8)
+            except ValueError as e:
+                # raised here and not per chunk, so callers can still fall back to F16
+                raise QuantError(str(e)) from e
+        return LazyChunkedTensor(self._chunks, shape, dtype, qtype, self._byteswap)
+
+    def byteswap(self, inplace: bool = False) -> LazyChunkedTensor:
+        if inplace:
+            raise NotImplementedError("a chunked tensor cannot be byteswapped in place")
+        return LazyChunkedTensor(self._chunks, self.shape, self.dtype, self._qtype, not self._byteswap)
+
+    def tofile(self, *args, **kwargs) -> None:
+        from .quants import quantize
+
+        written = 0
+        for load_chunk in self._chunks:
+            chunk = load_chunk()
+            if self._qtype is not None:
+                # exact only because chunks split on rows, and blocks never cross one
+                chunk = quantize(chunk, self._qtype)
+            if self._byteswap:
+                chunk = chunk.byteswap(inplace=False)
+            chunk.tofile(*args, **kwargs)
+            written += chunk.nbytes
+            del chunk
+        assert written == self.nbytes, f"chunked tensor wrote {written} bytes, expected {self.nbytes}"
@@ -1355,6 +1355,34 @@ class TensorNameMap:
            "model.confidence_head.proj", # dspark
        ),

+        MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: (
+            "model.layers.{bid}.attention_conv.base_kernel",
+        ),
+
+        MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: (
+            "model.layers.{bid}.attention_conv.kernel_projection",
+        ),
+
+        MODEL_TENSOR.DFLASH_FFN_CONV_BASE: (
+            "model.layers.{bid}.mlp_conv.base_kernel",
+        ),
+
+        MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: (
+            "model.layers.{bid}.mlp_conv.kernel_projection",
+        ),
+
+        MODEL_TENSOR.DFLASH_SELECTOR_PREV: (
+            "model.candidate_selector.predecessor_codebook",
+        ),
+
+        MODEL_TENSOR.DFLASH_SELECTOR_NEXT: (
+            "model.candidate_selector.successor_codebook",
+        ),
+
+        MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: (
+            "model.candidate_selector.hidden_projection",
+        ),
+
        MODEL_TENSOR.CLS: (
            "classifier",       # jina
            "classifier.dense", # roberta
@@ -2680,6 +2708,65 @@ class TensorNameMap:
                "model.layers.{bid}.post_attention_layernorm",
            ),
        },
+        MODEL_ARCH.QWEN4EXP: {
+            MODEL_TENSOR.HC_ATTN_NORM: (
+                "model.layers.{bid}.attn_hyper_connection.hc_norm",
+            ),
+            MODEL_TENSOR.HC_ATTN_DOWN: (
+                "model.layers.{bid}.attn_hyper_connection.input_mix_weight_down",
+            ),
+            MODEL_TENSOR.HC_ATTN_UP: (
+                "model.layers.{bid}.attn_hyper_connection.input_mix_weight_up",
+            ),
+            MODEL_TENSOR.HC_ATTN_INJECT: (
+                "model.layers.{bid}.attn_hyper_connection.block_inject_weight",
+            ),
+            MODEL_TENSOR.HC_FFN_NORM: (
+                "model.layers.{bid}.mlp_hyper_connection.hc_norm",
+            ),
+            MODEL_TENSOR.HC_FFN_DOWN: (
+                "model.layers.{bid}.mlp_hyper_connection.input_mix_weight_down",
+            ),
+            MODEL_TENSOR.HC_FFN_UP: (
+                "model.layers.{bid}.mlp_hyper_connection.input_mix_weight_up",
+            ),
+            MODEL_TENSOR.HC_FFN_INJECT: (
+                "model.layers.{bid}.mlp_hyper_connection.block_inject_weight",
+            ),
+            MODEL_TENSOR.HC_HEAD_NORM: (
+                "model.hyper_connection_mixer.hc_norm",
+            ),
+            MODEL_TENSOR.HC_HEAD_DOWN: (
+                "model.hyper_connection_mixer.input_mix_weight_down",
+            ),
+            MODEL_TENSOR.HC_HEAD_UP: (
+                "model.hyper_connection_mixer.input_mix_weight_up",
+            ),
+            MODEL_TENSOR.INDEXER_Q_NORM: (
+                "model.layers.{bid}.self_attn.indexer.q_layernorm",
+            ),
+            MODEL_TENSOR.INDEXER_K_NORM: (
+                "model.layers.{bid}.self_attn.indexer.k_layernorm",
+            ),
+            MODEL_TENSOR.PLE_KEY: (
+                "model.layers.{bid}.ple.key_proj",
+            ),
+            MODEL_TENSOR.PLE_VALUE: (
+                "model.layers.{bid}.ple.value_proj",
+            ),
+            MODEL_TENSOR.PLE_NORM_KEY: (
+                "model.layers.{bid}.ple.norm_key",
+            ),
+            MODEL_TENSOR.PLE_NORM_QUERY: (
+                "model.layers.{bid}.ple.norm_query",
+            ),
+            MODEL_TENSOR.PLE_NORM_CONV: (
+                "model.layers.{bid}.ple.norm_conv",
+            ),
+            MODEL_TENSOR.PLE_CONV1D: (
+                "model.layers.{bid}.ple.conv1d",
+            ),
+        },
    }

    mapping: dict[str, tuple[MODEL_TENSOR, str]]
@@ -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" {
@@ -214,6 +214,12 @@ extern "C" {
    LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode);
    LLAMA_API enum llama_load_mode llama_load_mode_from_str(const char * str);

+    enum llama_tensor_read_lazy {
+        LLAMA_TENSOR_READ_LAZY_OFF  = 0, // always read the whole tensor up front
+        LLAMA_TENSOR_READ_LAZY_AUTO = 1, // lazy only for marked tensors larger than 4 GiB (requires mmap)
+        LLAMA_TENSOR_READ_LAZY_ON   = 2, // read the rows of tensors marked by the arch on demand (requires mmap)
+    };
+
    enum llama_context_type {
        LLAMA_CONTEXT_TYPE_DEFAULT = 0,
        LLAMA_CONTEXT_TYPE_MTP     = 1,
@@ -315,6 +321,8 @@ extern "C" {
        enum llama_split_mode split_mode; // how to split the model across multiple GPUs
        enum llama_load_mode  load_mode;  // how to load the model

+        enum llama_tensor_read_lazy tensor_read_lazy; // on-demand reading of tensors marked by the arch
+
        // the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE
        int32_t main_gpu;

@@ -437,6 +445,7 @@ extern "C" {
        const struct llama_model_kv_override * kv_overrides;        // pointer to kv overrides
        const struct llama_model_tensor_override * tt_overrides;    // pointer to tensor overrides
        const int32_t * prune_layers;                               // pointer to layer indices to prune
+        size_t max_buf_size;                                        // max bytes of tensor rows kept in memory at once, 0 = default (8 GiB)
    } llama_model_quantize_params;

    typedef struct llama_logit_bias {
--- a/Show More
+++ b/Show More