metal : fix memory leak in early return

ci : move more jobs to ccache-buckets (#28375 )
* move more jobs to ccache-buckets * add venv deps * also jq
2026-09-27 21:46:57 +02:00 · 2026-09-04 18:37:43 +02:00 · 2026-09-04 15:50:33 +02:00 · 2026-09-04 15:13:10 +02:00 · 2026-09-04 14:31:36 +02:00 · 2026-09-04 15:22:38 +03:00
269 changed files with 35951 additions and 1950 deletions
@@ -24,7 +24,7 @@ runs:

        write-host "Installing ROCm wheels for multi-arch support"
        # Install ROCm wheels for multi-arch support (this may take several minutes)
-        python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ inputs.version }}"
+        python -m pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ "rocm[libraries,devel]==${{ inputs.version }}"

        # Pre-expand the devel tree so it is included in the cache
        write-host "Initializing ROCm devel tree"
@@ -110,7 +110,7 @@ jobs:
      #   cache on:  https://github.com/ggerganov/tmp2/actions/runs/26534713799/job/78224189394
      #
      #- name: ccache
-      #  uses: ggml-org/[email protected]1
+      #  uses: ggml-org/[email protected]4
      #  with:
      #    key: android-ubuntu-arm64
      #    evict-old-files: 1d
@@ -47,11 +47,19 @@ jobs:
        uses: actions/checkout@v6

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: apple-arm64
-          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_CI }}
+        with:
+          key: apple-arm64
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache

      - name: Build
        id: cmake_build
@@ -66,7 +74,25 @@ jobs:
            -DGGML_RPC=ON \
            -DCMAKE_OSX_DEPLOYMENT_TARGET=13.3
          time cmake --build build --config Release -j $(sysctl -n hw.logicalcpu)
-          leaks -atExit -- ./build/bin/test-thread-safety -hf ggml-org/gemma-3-270m-qat-GGUF -ngl 99 -p "$(printf 'hello %.0s' {1..128})" -n 16 -c 512 -ub 32 -np 2 -t 2 -lv 1
+
+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: apple-arm64
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
+      - name: Check for leaks
+        run: |
+          cmd=(./build/bin/test-thread-safety -hf ggml-org/gemma-3-270m-qat-GGUF -ngl 99 -p "$(printf 'hello %.0s' {1..128})" -n 16 -c 512 -ub 32 -np 2 -t 2 -lv 1)
+          leaks -atExit -- "${cmd[@]}"
+          # Graphics devices are leaked by Metal in Apple code sometimes, so we ignore those leaks
+          OBJC_DEBUG_MISSING_POOLS=YES "${cmd[@]}" 2>&1 | awk '{ print } index($0, "autoreleased with no pool in place") && !/class [a-zA-Z0-9]+Device autoreleased/ { found = 1 } END { exit found }'

      - name: Test
        id: cmake_test
@@ -74,16 +100,6 @@ jobs:
          cd build
          ctest -L main -E "test-llama-archs" --verbose --timeout 900

-      - name: ccache-clear
-        uses: ./.github/actions/ccache-clear
-        env:
-          GH_TOKEN: ${{ github.token }}
-        with:
-          key: apple-arm64
-          older: 5m
-          min: 1
-          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
-
  macos-latest-x64:
    runs-on: macos-15-intel

@@ -93,11 +109,19 @@ jobs:
        uses: actions/checkout@v6

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: apple-x64
-          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_CI }}
+        with:
+          key: apple-x64
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache

      - name: Build
        id: cmake_build
@@ -114,22 +138,24 @@ jobs:
            -DCMAKE_OSX_DEPLOYMENT_TARGET=13.3
          time cmake --build build --config Release -j $(sysctl -n hw.logicalcpu)

+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: apple-x64
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
      - name: Test
        id: cmake_test
        run: |
          cd build
          ctest -L main --verbose --timeout 900

-      - name: ccache-clear
-        uses: ./.github/actions/ccache-clear
-        env:
-          GH_TOKEN: ${{ github.token }}
-        with:
-          key: apple-x64
-          older: 5m
-          min: 1
-          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
-
  macos-latest-ios-xcode:
    runs-on: macos-latest

@@ -62,11 +62,10 @@ jobs:
        uses: actions/checkout@v6

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: cpu-${{ matrix.os }}
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+          save: false

      - name: Build Dependencies
        id: build_depends
@@ -91,6 +90,15 @@ jobs:
          python3 -m pip install --upgrade pip setuptools
          pip3 install ./gguf-py

+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
+        with:
+          key: cpu-${{ matrix.os }}
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
+
      - name: Build
        id: cmake_build
        run: |
@@ -100,6 +108,18 @@ jobs:
            -DGGML_RPC=ON
          time cmake --build build --config Release -j $(nproc)

+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: cpu-${{ matrix.os }}
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
      - name: Test
        id: cmake_test
        run: |
@@ -117,18 +137,6 @@ jobs:
          ./bin/llama-convert-llama2c-to-ggml --copy-vocab-from-model ./tok512.bin --llama2c-model stories260K.bin --llama2c-output-model stories260K.gguf
          ./bin/llama-completion -m stories260K.gguf -p "One day, Lily met a Shoggoth" -n 500 -c 256

-      # note: real deletion only on push to master (same condition as the ccache save),
-      #       dry-run otherwise (the token is read-only on PRs from forks)
-      - name: ccache-clear
-        uses: ./.github/actions/ccache-clear
-        env:
-          GH_TOKEN: ${{ github.token }}
-        with:
-          key: cpu-${{ matrix.os }}
-          older: 5m
-          min: 1
-          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
-
  windows:
    name: windows / ${{ matrix.build }}
    runs-on: windows-2025
@@ -156,7 +164,7 @@ jobs:
        uses: actions/checkout@v6

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: cpu-windows-2025-${{ matrix.build }}
          variant: ccache
@@ -53,7 +53,7 @@ jobs:
          apt install -y cmake build-essential ninja-build libgomp1 git libssl-dev jq python3 python3-venv python3-pip

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: cuda-ubuntu-24.04-cuda
          save: false
@@ -61,7 +61,7 @@ jobs:
      - name: ccache-buckets-restore
        uses: ./.github/actions/ccache-buckets
        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
        with:
          key: cuda-ubuntu-24.04-cuda
          folder: llama.cpp
@@ -108,7 +108,7 @@ jobs:
          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/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: cuda-ubuntu-22.04-hip
          save: false
@@ -116,7 +116,7 @@ jobs:
      - name: ccache-buckets-restore
        uses: ./.github/actions/ccache-buckets
        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
        with:
          key: cuda-ubuntu-22.04-hip
          folder: llama.cpp
@@ -159,7 +159,7 @@ jobs:
          apt-get install -y build-essential git cmake libssl-dev jq

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: cuda-ubuntu-22.04-musa
          save: false
@@ -167,7 +167,7 @@ jobs:
      - name: ccache-buckets-restore
        uses: ./.github/actions/ccache-buckets
        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
        with:
          key: cuda-ubuntu-22.04-musa
          folder: llama.cpp
@@ -47,7 +47,7 @@ jobs:
        uses: actions/checkout@v6

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: release-windows-2022-${{ matrix.arch }}-cuda-${{ matrix.cuda }}

@@ -152,7 +152,7 @@ jobs:
          & "${env:HIP_PATH}\lib\llvm\bin\clang.exe" --version

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          # TODO: this build does not match the build in release.yml, so we use a different cache key
          #       ideally, the builds should match, similar to the CUDA build above so that we would be able
@@ -35,7 +35,7 @@ jobs:
        uses: actions/checkout@v6

      #- name: ccache
-      #  uses: ggml-org/[email protected].16
+      #  uses: ggml-org/[email protected].24
      #  with:
      #    key: msys-windows-2025-x64
      #    variant: ccache
@@ -44,7 +44,7 @@ jobs:
        uses: actions/checkout@v6

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: opencl-windows-2025-x64
          variant: ccache
@@ -32,8 +32,8 @@ env:
  LLAMA_ARG_LOG_COLORS: 1
  LLAMA_ARG_LOG_PREFIX: 1
  LLAMA_ARG_LOG_TIMESTAMPS: 1
-  # TODO: fix and re-enable the `test-llama-archs` and `test-recurrent-state-rollback`
-  CTEST_EXCLUDE: "test-llama-archs|^test-recurrent-state-rollback"
+  # TODO: fix failing tests on OpenVINO backend
+  CTEST_EXCLUDE: "test-llama-archs|^test-recurrent-state-|test-backend-ops|test-save-load-state"

 jobs:
  ubuntu-24-openvino:
@@ -105,7 +105,7 @@ jobs:
        uses: actions/checkout@v6

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: openvino-windows-2022
          variant: ccache
@@ -67,7 +67,7 @@ jobs:

      # note: sparing some ccache since these jobs run on dedicated runners that are not part of the organitzation
      #- name: ccache
-      #  uses: ggml-org/ccache-action@afde29e5b5422e5da23cb1f639e8baecadeadfc3 # https://github.com/ggml-org/ccache-action/pull/1
+      #  uses: ggml-org/ccache-action@v1.2.24
      #  with:
      #    key: riscv-ubuntu-native
      #    evict-old-files: 1d
@@ -137,7 +137,7 @@ jobs:

      # note: sparing some ccache since these jobs run on dedicated runners that are not part of the organitzation
      #- name: ccache
-      #  uses: ggml-org/ccache-action@afde29e5b5422e5da23cb1f639e8baecadeadfc3 # https://github.com/ggml-org/ccache-action/pull/1
+      #  uses: ggml-org/ccache-action@v1.2.24
      #  with:
      #    key: riscv-ubuntu-native-sanitizer-${{ matrix.sanitizer }}-${{ matrix.build_type }}
      #    evict-old-files: 1d
@@ -55,7 +55,7 @@ jobs:
        uses: actions/checkout@v6

      # - name: ccache
-      #   uses: ggml-org/[email protected]1
+      #   uses: ggml-org/[email protected]4
      #   if: ${{ matrix.sanitizer != 'UNDEFINED' }}
      #   with:
      #     key: ctest-${{ matrix.sanitizer }}-ubuntu-24.04
@@ -75,11 +75,19 @@ jobs:
          sudo apt-get install -y ./level-zero.deb ./level-zero-devel.deb

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: sycl-ubuntu-24-${{ matrix.build }}
-          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_CI }}
+        with:
+          key: sycl-ubuntu-24-${{ matrix.build }}
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache

      - name: Build
        id: cmake_build
@@ -96,15 +104,17 @@ jobs:
            -DGGML_SYCL_F16=${{ matrix.fp16 }}
          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: sycl-ubuntu-24-${{ matrix.build }}
-          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

  windows-latest-sycl:
    runs-on: windows-2022
@@ -137,7 +147,7 @@ jobs:
          "LEVEL_ZERO_V1_SDK_PATH=C:/level-zero-sdk" | Out-File -FilePath $env:GITHUB_ENV -Append

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: sycl-windows-latest
          variant: ccache
@@ -53,12 +53,20 @@ jobs:
          echo "CXX=g++-14" >> "$GITHUB_ENV"

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: vulkan-ubuntu-24.04-arm
          variant: ccache
-          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_CI }}
+        with:
+          key: vulkan-ubuntu-24.04-arm
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache

      - name: Configure
        id: cmake_configure
@@ -73,15 +81,17 @@ jobs:
        run: |
          time cmake --build build -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: vulkan-ubuntu-24.04-arm
-          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

  ubuntu-llvmpipe:
    runs-on: ubuntu-24.04
@@ -112,11 +122,19 @@ jobs:
          strip: 1

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: vulkan-ubuntu-24.04-llvmpipe
-          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_CI }}
+        with:
+          key: vulkan-ubuntu-24.04-llvmpipe
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache

      - name: Build
        id: cmake_build
@@ -127,6 +145,18 @@ jobs:
            -DGGML_VULKAN=ON
          cmake --build build --config Release -j $(nproc)

+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: vulkan-ubuntu-24.04-llvmpipe
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
      - name: Test
        id: cmake_test
        run: |
@@ -138,16 +168,6 @@ jobs:
          # test-backend-ops is too slow on llvmpipe, skip it
          ctest -L main -E test-backend-ops --verbose --timeout 900

-      - name: ccache-clear
-        uses: ./.github/actions/ccache-clear
-        env:
-          GH_TOKEN: ${{ github.token }}
-        with:
-          key: vulkan-ubuntu-24.04-llvmpipe
-          older: 5m
-          min: 1
-          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
-
  windows:
    runs-on: windows-2025

@@ -160,7 +180,7 @@ jobs:
        uses: actions/checkout@v6

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: cpu-windows-2025-x64-vulkan
          variant: ccache
@@ -54,11 +54,10 @@ jobs:
        uses: actions/checkout@v6

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: webgpu-ubuntu-24.04-arm-wasm
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+          save: false

      - name: Install Emscripten
        run: |
@@ -76,6 +75,15 @@ jobs:
            "https://github.com/google/dawn/releases/download/${DAWN_TAG}/${EMDAWN_PKG}"
          unzip emdawn.zip

+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
+        with:
+          key: webgpu-ubuntu-24.04-arm-wasm
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
+
      - name: Build WASM WebGPU
        run: |
          source emsdk/emsdk_env.sh
@@ -89,12 +97,14 @@ jobs:

          time cmake --build build-wasm --config Release --target test-backend-ops -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: webgpu-ubuntu-24.04-arm-wasm
-          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
@@ -69,11 +69,10 @@ jobs:
        uses: actions/checkout@v6

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: webgpu-macos-latest
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+          save: false

      - name: Dawn Dependency
        id: dawn-depends
@@ -88,6 +87,15 @@ jobs:
          mkdir dawn
          tar -xvf artifact.tar.gz -C dawn --strip-components=1

+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
+        with:
+          key: webgpu-macos-latest
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
+
      - name: Build
        id: cmake_build
        run: |
@@ -95,22 +103,24 @@ jobs:
          cmake -B build -G "Ninja" -DCMAKE_BUILD_TYPE=Release -DGGML_WEBGPU=ON -DGGML_METAL=OFF -DGGML_BLAS=OFF
          time cmake --build build --config Release -j $(sysctl -n hw.logicalcpu)

+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: webgpu-macos-latest
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
      - name: Test
        id: cmake_test
        run: |
          cd build
          ctest -L main --verbose --timeout 900

-      - name: ccache-clear
-        uses: ./.github/actions/ccache-clear
-        env:
-          GH_TOKEN: ${{ github.token }}
-        with:
-          key: webgpu-macos-latest
-          older: 5m
-          min: 1
-          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
-
  ubuntu:
    runs-on: ubuntu-24.04

@@ -120,11 +130,10 @@ jobs:
        uses: actions/checkout@v6

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: webgpu-ubuntu-24.04
-          evict-old-files: 1d
-          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+          save: false

      - name: Dependencies
        id: depends
@@ -148,6 +157,15 @@ jobs:
          mkdir dawn
          tar -xvf artifact.tar.gz -C dawn --strip-components=1

+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
+        with:
+          key: webgpu-ubuntu-24.04
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
+
      - name: Build
        id: cmake_build
        run: |
@@ -156,6 +174,18 @@ jobs:
            -DGGML_WEBGPU=ON
          time cmake --build build --config Release -j $(nproc)

+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: webgpu-ubuntu-24.04
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
      - name: Test
        id: cmake_test
        run: |
@@ -163,13 +193,3 @@ jobs:
          # This is using llvmpipe and runs slower than other backends
          # test-backend-ops is too slow on llvmpipe, skip it
          ctest -L main -E test-backend-ops --verbose --timeout 900
-
-      - name: ccache-clear
-        uses: ./.github/actions/ccache-clear
-        env:
-          GH_TOKEN: ${{ github.token }}
-        with:
-          key: webgpu-ubuntu-24.04
-          older: 5m
-          min: 1
-          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
@@ -29,7 +29,7 @@ jobs:
        uses: actions/checkout@v6

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: copilot-setup-steps
          evict-old-files: 1d
@@ -49,14 +49,22 @@ jobs:
        id: depends
        run: |
          sudo apt-get update
-          sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev libssl-dev python3
+          sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev libssl-dev python3 python3-venv python3-pip jq

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: hip-quality-check-ubuntu-22.04
-          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_CI }}
+        with:
+          key: hip-quality-check-ubuntu-22.04
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache

      - name: Build with Werror
        id: cmake_build
@@ -85,12 +93,14 @@ jobs:
          make -j $(nproc) 2>&1 | tee metrics.log | grep -v 'Rpass-analysis=kernel-resource-usage\|remark:\|^$'
          python3 ../scripts/hip/gcn-cdna-vgpr-check.py metrics.log

-      - 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: hip-quality-check-ubuntu-22.04
-          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
@@ -103,7 +103,7 @@ jobs:
          path: tools/ui/dist

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: release-${{ matrix.os }}-${{ matrix.arch }}

@@ -187,7 +187,7 @@ jobs:

      - name: ccache
        if: ${{ matrix.build != 's390x' }}
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: release-${{ matrix.os }}-cpu

@@ -272,7 +272,7 @@ jobs:
          fi

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: release-${{ matrix.os }}-vulkan

@@ -358,7 +358,7 @@ jobs:
      #   cache on:  https://github.com/ggerganov/tmp2/actions/runs/26534713799/job/78224189394
      #
      #- name: ccache
-      #  uses: ggml-org/[email protected]1
+      #  uses: ggml-org/[email protected]4
      #  with:
      #    key: release-android-arm64

@@ -436,7 +436,7 @@ jobs:
          path: tools/ui/dist

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: release-ubuntu-24.04-openvino-release-no-preset-v1

@@ -551,7 +551,7 @@ jobs:
          path: tools/ui/dist

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: release-windows-2022-openvino
          variant: ccache
@@ -679,7 +679,7 @@ jobs:
          choco install ninja

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: release-windows-2025-vs2026-${{ matrix.arch }}-cpu

@@ -725,7 +725,7 @@ jobs:
    strategy:
      matrix:
        include:
-          - ROCM_VERSION: "7.14.0"
+          - ROCM_VERSION: "10.0.0"
            gpu_targets: "gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
            build: x64

@@ -741,7 +741,7 @@ jobs:
          choco install ninja

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
          evict-old-files: 1d
@@ -923,7 +923,7 @@ jobs:

      # TODO: these jobs need to use llvm toolchain in order to utilize the ccache
      #- name: ccache
-      #  uses: ggml-org/[email protected]1
+      #  uses: ggml-org/[email protected]4
      #  with:
      #    key: release-windows-2025-${{ matrix.arch }}-${{ matrix.backend }}

@@ -1011,7 +1011,7 @@ jobs:
          choco install ninja

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: release-windows-2022-${{ matrix.arch }}-cuda-${{ matrix.cuda }}

@@ -1107,7 +1107,7 @@ jobs:
          "LEVEL_ZERO_V1_SDK_PATH=C:/level-zero-sdk" | Out-File -FilePath $env:GITHUB_ENV -Append

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: release-windows-2022-x64-sycl

@@ -1225,7 +1225,7 @@ jobs:
          path: tools/ui/dist

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: release-ubuntu-24.04-sycl-${{ matrix.build }}

@@ -1279,7 +1279,7 @@ jobs:
    strategy:
      matrix:
        include:
-          - ROCM_VERSION: "7.14.0"
+          - ROCM_VERSION: "10.0.0"
            gpu_targets: "gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
            build: 'x64'

@@ -1302,7 +1302,7 @@ jobs:
          tool-cache: true

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: release-ubuntu-24.04-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
          evict-old-files: 1d
@@ -1333,7 +1333,7 @@ jobs:
          # libraries = HIP runtime and CMake configs needed for linking
          # devel = compilers, headers, static libs
          python -m pip install --upgrade pip
-          python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ matrix.ROCM_VERSION }}"
+          python -m pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ "rocm[libraries,devel]==${{ matrix.ROCM_VERSION }}"

          # Get ROCm installation paths using the rocm-sdk CLI tool
          ROCM_PATH=$(rocm-sdk path --root)
@@ -1703,7 +1703,7 @@ jobs:
            - [Ubuntu s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz)
            - [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz)
            - [Ubuntu arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-arm64.tar.gz)
-            - [Ubuntu x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-7.14-x64.tar.gz)
+            - [Ubuntu x64 (ROCm 10.0)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-10.0-x64.tar.gz)
            - [Ubuntu x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-openvino-${{ needs.ubuntu-24-openvino.outputs.openvino_version }}-x64.tar.gz)
            - [Ubuntu x64 (SYCL FP32)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp32-x64.tar.gz)
            - [Ubuntu x64 (SYCL FP16)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp16-x64.tar.gz)
@@ -1721,7 +1721,7 @@ jobs:
            - [Windows x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-vulkan-x64.zip)
            - [Windows x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-openvino-${{ needs.windows-openvino.outputs.openvino_version }}-x64.zip)
            - [Windows x64 (SYCL)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-sycl-x64.zip)
-            - [Windows x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-rocm-7.14-x64.zip)
+            - [Windows x64 (ROCm 10.0)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-rocm-10.0-x64.zip)

            **openEuler:**
            - [DISABLED](https://github.com/ggml-org/llama.cpp/pull/23705)
@@ -103,7 +103,7 @@ jobs:
          source .venv/bin/activate
          cd tools/server/tests
          export ${{ matrix.extra_args }}
-          ./tests.sh
+          PYTEST_WORKERS=1 ./tests.sh

      - name: Slow tests
        id: server_integration_tests_slow
@@ -112,4 +112,4 @@ jobs:
          source .venv/bin/activate
          cd tools/server/tests
          export ${{ matrix.extra_args }}
-          SLOW_TESTS=1 ./tests.sh
+          PYTEST_WORKERS=1 SLOW_TESTS=1 ./tests.sh
@@ -102,7 +102,7 @@ jobs:
          ./tests.sh

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

    steps:
      - name: Clone
@@ -112,12 +112,42 @@ jobs:
          fetch-depth: 0
          ref: ${{ github.event.inputs.sha || github.event.pull_request.head.sha || github.sha || github.head_ref || github.ref_name }}

+      - name: Install dependencies
+        run: |
+          sudo apt update
+          sudo apt install -y cmake libssl-dev python3 python3-venv python3-pip
+
+      - name: ccache
+        uses: ggml-org/[email protected]
+        with:
+          restore: false
+          save: false
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        with:
+          key: self-hosted-server-cuda
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
+
      - name: Build
        id: cmake_build
        run: |
-          cmake -B build -DGGML_CUDA=ON -DGGML_SCHED_NO_REALLOC=ON
+          cmake -B build -DGGML_CUDA=ON -DGGML_SCHED_NO_REALLOC=ON -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc
          cmake --build build --config Release -j $(nproc) --target llama-server

+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: self-hosted-server-cuda
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
      - name: Python setup
        id: setup_python
        run: |
@@ -80,11 +80,19 @@ jobs:
          ref: ${{ github.event.inputs.sha || github.event.pull_request.head.sha || github.sha || github.head_ref || github.ref_name }}

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: server-ubuntu-24.04-arm
-          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_CI }}
+        with:
+          key: server-ubuntu-24.04-arm
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache

      - name: Build
        id: cmake_build
@@ -93,6 +101,18 @@ jobs:
            -DGGML_SCHED_NO_REALLOC=ON
          cmake --build build --config Release -j $(nproc) --target llama-server

+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: server-ubuntu-24.04-arm
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
      - name: Python setup
        id: setup_python
        uses: actions/setup-python@v6
@@ -128,16 +148,6 @@ jobs:
          export LLAMA_ARG_BACKEND_SAMPLING=1
          SLOW_TESTS=1 ./tests.sh

-      - name: ccache-clear
-        uses: ./.github/actions/ccache-clear
-        env:
-          GH_TOKEN: ${{ github.token }}
-        with:
-          key: server-ubuntu-24.04-arm
-          older: 5m
-          min: 1
-          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
-
  windows:
    runs-on: windows-2025

@@ -150,7 +160,7 @@ jobs:
          ref: ${{ github.event.inputs.sha || github.event.pull_request.head.sha || github.sha || github.head_ref || github.ref_name }}

      - name: ccache
-        uses: ggml-org/[email protected]1
+        uses: ggml-org/[email protected]4
        with:
          key: server-windows-2025-x64
          evict-old-files: 1d
@@ -1,4 +1,4 @@
-# date: Tue Aug 18 14:32:43 EEST 2026
+# date: Fri Sep  4 10:06:46 EEST 2026
 # this file is auto-generated by scripts/gen-authors.sh

 Нияз Гарифзянов <[email protected]>
@@ -46,6 +46,7 @@ Abhijit Ramesh <[email protected]>
 abhijitb11 <[email protected]>
 Abhilash Majumder <[email protected]>
 Abhinay Krishna <[email protected]>
+Abhiram <[email protected]>
 Abhishek Gopinath K <[email protected]>
 abotsis <[email protected]>
 Abraham Gonzalez <[email protected]>
@@ -87,6 +88,7 @@ akleine <[email protected]>
 Al G <[email protected]>
 Al Mochkin <[email protected]>
 Alan Gray <[email protected]>
+Alan Tseng <[email protected]>
 Alawode Oluwandabira <[email protected]>
 Albert Jin <[email protected]>
 Alberto <[email protected]>
@@ -136,7 +138,9 @@ alonfaraj <[email protected]>
 AlpinDale <[email protected]>
 alwqx <[email protected]>
 Aman <[email protected]>
+Aman Chadha(IVIXMMI) <[email protected]>
 Aman Gupta <[email protected]>
+Aman Karki <[email protected]>
 amd-dwang <[email protected]>
 amd-lalithnc <[email protected]>
 Amir <[email protected]>
@@ -187,6 +191,7 @@ Anton Mitkov <[email protected]>
 Antonis Makropoulos <[email protected]>
 Anudit Nagar <[email protected]>
 Anuj Attri <[email protected]>
+anujj <[email protected]>
 anzz1 <[email protected]>
 Aparna M P <[email protected]>
 Aparna M P <[email protected]>
@@ -196,6 +201,7 @@ arch-btw <[email protected]>
 arcrank <[email protected]>
 ardfork <[email protected]>
 Arik Poznanski <[email protected]>
+Aritro Bandyopadhyay <[email protected]>
 arlo-phoenix <[email protected]>
 Armen Kaleshian <[email protected]>
 Arsen Arutunan <[email protected]>
@@ -230,6 +236,7 @@ bandoti <[email protected]>
 Bar Haim <[email protected]>
 BarfingLemurs <[email protected]>
 Bart Louwers <[email protected]>
+Bartosz Taudul <[email protected]>
 Bartowski <[email protected]>
 Bartowski <[email protected]>
 Bas Nijholt <[email protected]>
@@ -277,6 +284,7 @@ Bono Lv <[email protected]>
 Borislav Stanimirov <[email protected]>
 Borislav Stanimirov <[email protected]>
 Bowen Han <[email protected]>
+Brad Smith <[email protected]>
 Branden Butler <[email protected]>
 Brandon Squizzato <[email protected]>
 Brian <[email protected]>
@@ -287,6 +295,7 @@ Bryan Honof <[email protected]>
 bryanSwk <[email protected]>
 bsilvereagle <[email protected]>
 bssrdf <[email protected]>
+Buğra Özgürsoy <[email protected]>
 byte-6174 <[email protected]>
 Caleb DeLeeuw <[email protected]>
 Calvin Laurenson <[email protected]>
@@ -326,6 +335,7 @@ Chenguang Li <[email protected]>
 Chenguang Li <[email protected]>
 Chipmunk <[email protected]>
 chiranko <[email protected]>
+Chris Danis <[email protected]>
 Chris Elrod <[email protected]>
 Chris Kuehl <[email protected]>
 Chris Lee <[email protected]>
@@ -356,6 +366,7 @@ clyang <[email protected]>
 cmdr2 <[email protected]>
 cmdr2 <[email protected]>
 cocktailpeanut <[email protected]>
+codemonkey <[email protected]>
 codezjx <[email protected]>
 coezbek <[email protected]>
 comex <[email protected]>
@@ -367,6 +378,8 @@ Copilot <[email protected]>
 Corentin REGAL <[email protected]>
 cphlipot <[email protected]>
 cpumaxx <[email protected]>
+cqderek <[email protected]>
+cqderek <[email protected]>
 crasm <[email protected]>
 crasm <[email protected]>
 crat0z <[email protected]>
@@ -427,6 +440,7 @@ DavidKorczynski <[email protected]>
 davidrhodus <[email protected]>
 Dawid Potocki <[email protected]>
 Dawid Wysocki <[email protected]>
+Daya Adianto <[email protected]>
 ddh0 <[email protected]>
 ddh0 <[email protected]>
 ddpasa <[email protected]>
@@ -463,6 +477,7 @@ Dmytro Romanov <[email protected]>
 Dobri Danchev <[email protected]>
 DocShotgun <[email protected]>
 Doctor Shotgun <[email protected]>
+Dominik Pantaleoni <[email protected]>
 Don Mahurin <[email protected]>
 Dong Won Kim <[email protected]>
 Donghyeon Jeong <[email protected]>
@@ -504,6 +519,7 @@ Emmanuel Ferdman <[email protected]>
 Emreerdog <[email protected]>
 Engininja2 <[email protected]>
 Equim <[email protected]>
+Eric A Stalee <[email protected]>
 Eric Curtin <[email protected]>
 Eric Curtin <[email protected]>
 Eric Curtin <[email protected]>
@@ -519,6 +535,7 @@ Esko Toivonen <[email protected]>
 Ethan Turner <[email protected]>
 Ettore Di Giacinto <[email protected]>
 EugeoSynthesisThirtyTwo <[email protected]>
+Eurekatic <[email protected]>
 Evan Huus <[email protected]>
 Evan Jones <[email protected]>
 Evan Miller <[email protected]>
@@ -677,6 +694,7 @@ HimariO <[email protected]>
 hipudding <[email protected]>
 Hitesh Chopra <[email protected]>
 hksdpc255 <[email protected]>
+hmirin <[email protected]>
 hmscider <[email protected]>
 Hoang Nguyen <[email protected]>
 hoangmit <[email protected]>
@@ -701,6 +719,7 @@ Huawei Lin <[email protected]>
 Hugo <[email protected]>
 Hugo Roussel <[email protected]>
 Huifeng Ou <[email protected]>
+HumerousGorgon <[email protected]>
 hutli <[email protected]>
 hutli <[email protected]>
 hutli <[email protected]>
@@ -738,12 +757,15 @@ intelmatt <[email protected]>
 iohub <[email protected]>
 Ionoclast Laboratories <[email protected]>
 iron <[email protected]>
+Isaac <[email protected]>
 Isaac McFadyen <[email protected]>
 IsaacDynamo <[email protected]>
 Ishaan Gandhi <[email protected]>
 iSma <[email protected]>
 Ismail <[email protected]>
 issixx <[email protected]>
+itsnotoger <[email protected]>
+itterative <[email protected]>
 Ivan <[email protected]>
 Ivan Chikish <[email protected]>
 Ivan Filipov <[email protected]>
@@ -768,6 +790,7 @@ Jakkala Mahesh <[email protected]>
 Jakub N <[email protected]>
 JamePeng <[email protected]>
 James A Capozzoli <[email protected]>
+James Francis <[email protected]>
 James O'Leary <[email protected]>
 James Reynolds <[email protected]>
 jameswu2014 <[email protected]>
@@ -798,6 +821,7 @@ Jed Fox <[email protected]>
 Jeff Bolz <[email protected]>
 Jeffrey Morgan <[email protected]>
 Jeffrey Quesnelle <[email protected]>
+Jeremie Miller <[email protected]>
 Jeremy Demeule <[email protected]>
 Jeremy Rand <[email protected]>
 Jeroen Mostert <[email protected]>
@@ -809,6 +833,7 @@ Jesse Jojo Johnson <[email protected]>
 Jesse LaRose <[email protected]>
 Jesse Posner <[email protected]>
 Jesus Talavera <[email protected]>
+Jetson Tan <[email protected]>
 Jett Janiak <[email protected]>
 Jeximo <[email protected]>
 JFLFY2255 <[email protected]>
@@ -825,6 +850,7 @@ Jie Fu (傅杰) <[email protected]>
 jiez <[email protected]>
 Jillis ter Hove <[email protected]>
 Jim Wu <[email protected]>
+Jingxin (Philip) Li <[email protected]>
 Jinwoo Jeong <[email protected]>
 Jinyang He <[email protected]>
 jinzihao <[email protected]>
@@ -850,11 +876,13 @@ John Balis <[email protected]>
 John Bean <[email protected]>
 John Eismeier <[email protected]>
 John Smith <[email protected]>
+John-Henry Lim <[email protected]>
 Johnathan Craig Maudlin <[email protected]>
 JohnnyB <[email protected]>
 johnson442 <[email protected]>
 jojorne <[email protected]>
 jon-chuang <[email protected]>
+Jonas J <[email protected]>
 Jonas Jankaitis <[email protected]>
 Jonas Wunderlich <[email protected]>
 Jonathan <[email protected]>
@@ -924,6 +952,7 @@ Karsten Weiss <[email protected]>
 Karthick <[email protected]>
 Karthik Kumar Viswanathan <[email protected]>
 Karthik Sethuraman <[email protected]>
+Kartik Gulia <[email protected]>
 Kartik Sirohi <[email protected]>
 Kashif Rasul <[email protected]>
 KASR <[email protected]>
@@ -931,6 +960,7 @@ Kasumi <[email protected]>
 Katostrofik <[email protected]>
 katsu560 <[email protected]>
 Kawrakow <[email protected]>
+kbenkhaled <[email protected]>
 kchro3 <[email protected]>
 kdkd <[email protected]>
 Keiichi Tabata <[email protected]>
@@ -939,6 +969,7 @@ Kenvix ⭐ <[email protected]>
 Kerfuffle <[email protected]>
 Kevin Gibbons <[email protected]>
 Kevin Hannon <[email protected]>
+Kevin Hopper <[email protected]>
 Kevin Ji <[email protected]>
 Kevin Kwok <[email protected]>
 Kevin Liu <[email protected]>
@@ -964,12 +995,14 @@ Konstantin Herud <[email protected]>
 Konstantin Zhuravlyov <[email protected]>
 Krishna Sridhar <[email protected]>
 krystiancha <[email protected]>
+krzsztf <[email protected]>
 kubawoo <[email protected]>
 kumaal <[email protected]>
 kunal-vaishnavi <[email protected]>
 kunnis <[email protected]>
 Kunshang Ji <[email protected]>
 kuronekosaiko <[email protected]>
+kurquhar <[email protected]>
 Kusha Gharahi <[email protected]>
 kustaaya <[email protected]>
 kuvaus <[email protected]>
@@ -981,6 +1014,7 @@ Kyle Liang <[email protected]>
 Kyle Mistele <[email protected]>
 KyleHagy <[email protected]>
 Kylin <[email protected]>
+Kyozzz <[email protected]>
 l-austenfeld <[email protected]>
 l3utterfly <[email protected]>
 l8bloom <[email protected]>
@@ -992,6 +1026,7 @@ Lars Sonchocky-Helldorf <[email protected]>
 las7 <[email protected]>
 Lasse Lauwerys <[email protected]>
 Laura <[email protected]>
+Laurent Zuijdwijk <[email protected]>
 Law Po Ying <[email protected]>
 lcy <[email protected]>
 ldwang <[email protected]>
@@ -1039,6 +1074,8 @@ Ludovic Henry <[email protected]>
 Ludovic Henry <[email protected]>
 Lukas Straub <[email protected]>
 Łukasz Ślusarczyk <[email protected]>
+Lukasz Stolcman <[email protected]>
+LunalFresh <[email protected]>
 Luo Tian <[email protected]>
 luoyu-intel <[email protected]>
 luyhcsu <[email protected]>
@@ -1054,6 +1091,7 @@ Maarten ter Huurne <[email protected]>
 Maciej Lisowski <[email protected]>
 Mack Straight <[email protected]>
 maddes8cht <[email protected]>
+Mads Marquart <[email protected]>
 Maël Kerbiriou <[email protected]>
 MaggotHATE <[email protected]>
 MagicExists <[email protected]>
@@ -1215,6 +1253,8 @@ Naco Siren <[email protected]>
 Nam D. Tran <[email protected]>
 nanahi <[email protected]>
 Nathan Epstein <[email protected]>
+Nathan Wilson <[email protected]>
+Nathanw1014 <[email protected]>
 Natsu <[email protected]>
 Nauful Shaikh <[email protected]>
 NawafAlansari <[email protected]>
@@ -1237,6 +1277,7 @@ niansa/tuxifan <[email protected]>
 Nicholai Tukanov <[email protected]>
 Nicholas Sparks <[email protected]>
 Nick <[email protected]>
+Nick Farrell <[email protected]>
 nick huang <[email protected]>
 Nick Lafleur <[email protected]>
 Nick Towle <[email protected]>
@@ -1259,6 +1300,7 @@ NikolaiLyssogor <[email protected]>
 Nikolaos Pothitos <[email protected]>
 Nikolas <[email protected]>
 Nikolay Popov <[email protected]>
+Nils Gladitz <[email protected]>
 Nindaleth <[email protected]>
 ningshanwutuobang <[email protected]>
 Noah <[email protected]>
@@ -1355,6 +1397,7 @@ Pop Flamingo <[email protected]>
 postmasters <[email protected]>
 Pouya <[email protected]>
 pqnet <[email protected]>
+Prabhsimran Singh <[email protected]>
 Prabod <[email protected]>
 Prajwal B Mehendarkar <[email protected]>
 Pranav Dhinakar <[email protected]>
@@ -1378,6 +1421,7 @@ qouoq <[email protected]>
 Qu Zongfu <[email protected]>
 quei <[email protected]>
 Quentin Bramas <[email protected]>
+QuintinShaw <[email protected]>
 QuintinShaw <[email protected]>
 qunash <[email protected]>
 quyentonndbs <[email protected]>
@@ -1462,6 +1506,7 @@ robertomeroni <[email protected]>
 Robey Holderith <[email protected]>
 Robin Davidsson <[email protected]>
 Robyn <[email protected]>
+Rock Chen <[email protected]>
 Rőczey Barnabás <[email protected]>
 RodriMora <[email protected]>
 Roger Chen <[email protected]>
@@ -1499,17 +1544,21 @@ runfuture <[email protected]>
 RunningLeon <[email protected]>
 RunningLeon <[email protected]>
 Russyyds <[email protected]>
+Ryan C <[email protected]>
 Ryan Goulden <[email protected]>
 Ryan Landay <[email protected]>
 Ryan Mangeno <[email protected]>
 Ryder Wishart <[email protected]>
 Ryuei <[email protected]>
 s-goto-11 <[email protected]>
+s0mecode <[email protected]>
 s8322 <[email protected]>
+Saad Ali <[email protected]>
 Saba Fallah <[email protected]>
 Saba Fallah <[email protected]>
 Sachin Desai <[email protected]>
 Sachin Sharma <[email protected]>
+Safi Ullah <[email protected]>
 safranowith <[email protected]>
 SakuraUmi <[email protected]>
 Salvador E. Tropea <[email protected]>
@@ -1552,6 +1601,7 @@ Sergey Alirzaev <[email protected]>
 Sergey Alirzaev <[email protected]>
 Sergey Fedorov <[email protected]>
 Sergey Malinin <[email protected]>
+Sergey Sklyarov <[email protected]>
 Sergio López <[email protected]>
 Sergio López <[email protected]>
 Sergiu <[email protected]>
@@ -1582,11 +1632,13 @@ Shawn Gu <[email protected]>
 Shawn yang <[email protected]>
 Shelby Jenkins <[email protected]>
 Sheldon Robinson <[email protected]>
+Shenghan Yang <[email protected]>
 shibe2 <[email protected]>
 Shijie <[email protected]>
 Shin-myoung-serp <[email protected]>
 Shintarou Okada <[email protected]>
 shivamkumard-ctrl <[email protected]>
+Shobhit <[email protected]>
 Shouyu <[email protected]>
 Shouzheng Liu <[email protected]>
 Shouzheng Liu <[email protected]>
@@ -1607,6 +1659,7 @@ Sigbjørn Skjæret <[email protected]>
 Sigbjørn Skjæret <[email protected]>
 simevo <[email protected]>
 Simon Redman <[email protected]>
+Simon Teixidor <[email protected]>
 Simon Willison <[email protected]>
 simon886212 <[email protected]>
 Simranjeet Singh <[email protected]>
@@ -1663,6 +1716,7 @@ stevenkuang <[email protected]>
 Steward Garcia <[email protected]>
 StrangeBytesDev <[email protected]>
 strawberrymelonpanda <[email protected]>
+Strongtut <[email protected]>
 Suaj Carrot <[email protected]>
 sudhiarm <[email protected]>
 Sukriti Sharma <[email protected]>
@@ -1687,6 +1741,7 @@ Tamar <[email protected]>
 tamarPal <[email protected]>
 Tameem <[email protected]>
 Tamotsu Takahashi <[email protected]>
+Tanner Bruhn <[email protected]>
 tarcey <[email protected]>
 Tarek Dakhran <[email protected]>
 Tarek Dakhran <[email protected]>
@@ -1696,6 +1751,7 @@ Taylor <[email protected]>
 tc-mb <[email protected]>
 TecJesh <[email protected]>
 Tei Home <[email protected]>
+Tekin Ertekin <[email protected]>
 Tekin Ertekin <[email protected]>
 tempstudio <[email protected]>
 teo <[email protected]>
@@ -1737,6 +1793,7 @@ Ting Lou <[email protected]>
 Ting Lou <[email protected]>
 Ting Sun <[email protected]>
 Titaniumtown <[email protected]>
+Tiwei Bie <[email protected]>
 tjohnman <[email protected]>
 Tobias Lütke <[email protected]>
 Toby <[email protected]>
@@ -1813,6 +1870,7 @@ Vishal Agarwal <[email protected]>
 Vishal Singh <[email protected]>
 Vitali Lovich <[email protected]>
 Vivian <[email protected]>
+vk <[email protected]>
 Vlad <[email protected]>
 Vladimir <[email protected]>
 Vladimir Malyutin <[email protected]>
@@ -1897,6 +1955,7 @@ Yaiko <[email protected]>
 Yakine Tahtah <[email protected]>
 YangLe <[email protected]>
 yangli2 <[email protected]>
+Yaniss Amazouz <[email protected]>
 Yann Follet <[email protected]>
 Yanzhao Wang <[email protected]>
 Yarden Tal <[email protected]>
@@ -4,7 +4,7 @@ include(CheckIncludeFileCXX)

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

@@ -57,6 +57,7 @@
 /ggml/src/ggml-cann/                    @ggml-org/ggml-cann
 /ggml/src/ggml-common.h                 @ggerganov
 /ggml/src/ggml-cpu/                     @ggerganov
+/ggml/src/ggml-cpu/iqp.*                @bartowski1182
 /ggml/src/ggml-cpu/spacemit/            @alex-spacemit
 /ggml/src/ggml-cuda/                    @ggml-org/ggml-cuda
 /ggml/src/ggml-cuda/vendors/hip.h       @IMbackK
@@ -13,7 +13,7 @@
 [![Docker](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/docker.yml?label=Docker)](https://github.com/ggml-org/llama.cpp/actions/workflows/docker.yml)
 [![Winget](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/winget.yml?label=Winget)](https://github.com/ggml-org/llama.cpp/actions/workflows/winget.yml)

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

 </div>

@@ -74,7 +74,7 @@ The `llama.cpp` project is build on top of the [ggml](https://github.com/ggml-or
 | [CANN](docs/build.md#cann) | Ascend NPU |
 | [CUDA](docs/build.md#cuda) | Nvidia GPU |
 | [HIP](docs/build.md#hip) | AMD GPU |
-| [Hexagon [In Progress]](docs/backend/snapdragon/README.md) | Snapdragon |
+| [Hexagon](docs/backend/snapdragon/README.md) | Snapdragon |
 | [IBM zDNN](docs/backend/zDNN.md) | IBM Z & LinuxONE |
 | [MUSA](docs/build.md#musa) | Moore Threads GPU |
 | [Metal](docs/build.md#metal-build) | Apple Silicon |
@@ -80,7 +80,7 @@ static const command cmds[] = {
 #undef UPDATE_HIDDEN

 static int version(int /*argc*/, char ** /*argv*/) {
-    llama_print_build_info(llama_version());
+    llama_print_build_info(llama_version(), stdout);
    return 0;
 }

@@ -18,7 +18,7 @@ LLAMA_BUILD_TESTS=OFF
 LLAMA_BUILD_SERVER=OFF
 LLAMA_BUILD_MTMD=ON
 GGML_METAL=ON
-GGML_METAL_EMBED_LIBRARY=ON
+GGML_METAL_EMBED_LIBRARY=${GGML_METAL_EMBED_LIBRARY:-ON}
 GGML_BLAS_DEFAULT=ON
 GGML_OPENMP=OFF

@@ -169,6 +169,14 @@ setup_framework_structure() {
    cp tools/mtmd/mtmd.h           ${header_path}
    cp tools/mtmd/mtmd-helper.h    ${header_path}

+    if [[ "$GGML_METAL_EMBED_LIBRARY" == "OFF" ]]; then
+        if [[ "$platform" == "macos" ]]; then
+            cp ${build_dir}/bin/*.metallib ${build_dir}/framework/${framework_name}.framework/Versions/A/Resources/
+        else
+            cp ${build_dir}/bin/*.metallib ${build_dir}/framework/${framework_name}.framework/
+        fi
+    fi
+
    # Create module map (common for all platforms)
    cat > ${module_path}module.modulemap << EOF
 framework module llama {
@@ -450,6 +458,7 @@ build_ios_sim() {
        -DIOS=ON \
        -DCMAKE_SYSTEM_NAME=iOS \
        -DCMAKE_OSX_SYSROOT=iphonesimulator \
+        -DGGML_METAL_TARGET_OS=ios \
        -DCMAKE_OSX_ARCHITECTURES="arm64;x86_64" \
        -DCMAKE_XCODE_ATTRIBUTE_SUPPORTED_PLATFORMS=iphonesimulator \
        -DCMAKE_C_FLAGS="${COMMON_C_FLAGS}" \
@@ -467,6 +476,7 @@ build_ios_device() {
        -DCMAKE_OSX_DEPLOYMENT_TARGET=${IOS_MIN_OS_VERSION} \
        -DCMAKE_SYSTEM_NAME=iOS \
        -DCMAKE_OSX_SYSROOT=iphoneos \
+        -DGGML_METAL_TARGET_OS=ios \
        -DCMAKE_OSX_ARCHITECTURES="arm64" \
        -DCMAKE_XCODE_ATTRIBUTE_SUPPORTED_PLATFORMS=iphoneos \
        -DCMAKE_C_FLAGS="${COMMON_C_FLAGS}" \
@@ -498,6 +508,7 @@ build_visionos() {
        -DCMAKE_OSX_ARCHITECTURES="arm64" \
        -DCMAKE_SYSTEM_NAME=visionOS \
        -DCMAKE_OSX_SYSROOT=xros \
+        -DGGML_METAL_TARGET_OS=xros \
        -DCMAKE_XCODE_ATTRIBUTE_SUPPORTED_PLATFORMS=xros \
        -DCMAKE_C_FLAGS="${COMMON_C_FLAGS}" \
        -DCMAKE_CXX_FLAGS="${COMMON_CXX_FLAGS}" \
@@ -516,6 +527,7 @@ build_visionos_sim() {
        -DCMAKE_OSX_ARCHITECTURES="arm64;x86_64" \
        -DCMAKE_SYSTEM_NAME=visionOS \
        -DCMAKE_OSX_SYSROOT=xrsimulator \
+        -DGGML_METAL_TARGET_OS=xros \
        -DCMAKE_XCODE_ATTRIBUTE_SUPPORTED_PLATFORMS=xrsimulator \
        -DCMAKE_C_FLAGS="${COMMON_C_FLAGS}" \
        -DCMAKE_CXX_FLAGS="${COMMON_CXX_FLAGS}" \
@@ -534,6 +546,7 @@ build_tvos_sim() {
        -DCMAKE_OSX_DEPLOYMENT_TARGET=${TVOS_MIN_OS_VERSION} \
        -DCMAKE_SYSTEM_NAME=tvOS \
        -DCMAKE_OSX_SYSROOT=appletvsimulator \
+        -DGGML_METAL_TARGET_OS=tvos \
        -DCMAKE_OSX_ARCHITECTURES="arm64;x86_64" \
        -DGGML_METAL=ON \
        -DCMAKE_XCODE_ATTRIBUTE_SUPPORTED_PLATFORMS=appletvsimulator \
@@ -552,6 +565,7 @@ build_tvos_device() {
        -DCMAKE_OSX_DEPLOYMENT_TARGET=${TVOS_MIN_OS_VERSION} \
        -DCMAKE_SYSTEM_NAME=tvOS \
        -DCMAKE_OSX_SYSROOT=appletvos \
+        -DGGML_METAL_TARGET_OS=tvos \
        -DCMAKE_OSX_ARCHITECTURES="arm64" \
        -DGGML_METAL=ON \
        -DCMAKE_XCODE_ATTRIBUTE_SUPPORTED_PLATFORMS=appletvos \
@@ -189,8 +189,8 @@ if [ ! -z ${GG_BUILD_OPENVINO} ]; then
    fi
    CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_OPENVINO=ON"

-    # TODO: fix and re-enable the `test-llama-archs` and `test-recurrent-state-rollback*`
-    CTEST_EXTRA="-E test-llama-archs|^test-recurrent-state-rollback"
+    # TODO: fix failing tests on OpenVINO backend
+    CTEST_EXTRA="-E test-llama-archs|^test-recurrent-state-|test-backend-ops|test-save-load-state"
 fi

 ## helpers
@@ -960,6 +960,11 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
        ));
    }

+    // if the preserve_reasoning kwarg was not specified explicitly, enable it by default
+    if (!params.default_template_kwargs.count("preserve_reasoning")) {
+        params.default_template_kwargs["preserve_reasoning"] = "true";
+    }
+
    return true;
 }

@@ -3553,6 +3558,10 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
                    LOG_WRN("Setting 'enable_thinking' via --chat-template-kwargs is deprecated. "
                            "Use --reasoning on / --reasoning off instead.\n");
                }
+                if (item.key() == "preserve_reasoning") {
+                    LOG_WRN("Setting 'preserve_reasoning' via --chat-template-kwargs is deprecated. "
+                            "Use --reasoning-preserve / --no-reasoning-preserve instead.\n");
+                }
                params.default_template_kwargs[item.key()] = item.value().dump();
            }
        }
@@ -3743,7 +3752,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
    add_opt(common_arg(
        {"--reasoning-preserve"},
        {"--no-reasoning-preserve"},
-        "preserve reasoning trace in the full history, not just the last assistant message (default: template default)\n"
+        "preserve reasoning trace in the full history, not just the last assistant message (default: enabled)\n"
        "compatible with certain templates having 'supports_preserve_reasoning' capability\n"
        "example: https://docs.z.ai/guides/capabilities/thinking-mode#preserved-thinking",
        [](common_params & params, bool value) {
@@ -3752,6 +3761,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            } else {
                params.default_template_kwargs["preserve_reasoning"] = "false";
            }
+            params.preserve_reasoning_specified = true;
        }
    ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_REASONING_PRESERVE"));
    add_opt(common_arg(
@@ -29,7 +29,7 @@ const char * llama_build_info(void) {
    return s.c_str();
 }

-void llama_print_build_info(const char * llama_version) {
-    fprintf(stderr, "version: %s (build %d, commit %s)\n", llama_version, llama_build_number(), llama_commit());
-    fprintf(stderr, "built with %s for %s\n", llama_compiler(), llama_build_target());
+void llama_print_build_info(const char * llama_version, FILE * stream) {
+    fprintf(stream, "version: %s (build %d, commit %s)\n", llama_version, llama_build_number(), llama_commit());
+    fprintf(stream, "built with %s for %s\n", llama_compiler(), llama_build_target());
 }
@@ -1,5 +1,7 @@
 #pragma once

+#include <cstdio>
+
 int llama_build_number(void);

 const char * llama_commit(void);
@@ -8,4 +10,4 @@ const char * llama_compiler(void);
 const char * llama_build_target(void);
 const char * llama_build_info(void);

-void llama_print_build_info(const char *);
+void llama_print_build_info(const char *, FILE * = stderr);
@@ -270,7 +270,7 @@ struct common_params_sampling {
        COMMON_SAMPLER_TYPE_TEMPERATURE,
    };

-    common_grammar              grammar;      // optional grammar constraint (user / output-format / tool-calls)
+    common_grammar                      grammar;          // optional grammar constraint (user / output-format / tool-calls)
    bool                                grammar_lazy = false;
    std::vector<common_grammar_trigger> grammar_triggers; // optional triggers (for lazy grammars)
    std::set<llama_token>               preserved_tokens;
@@ -657,6 +657,7 @@ struct common_params {
    std::string ssl_file_cert = "";                                                                         // NOLINT

    std::map<std::string, std::string> default_template_kwargs;
+    bool preserve_reasoning_specified = false;

    // CLI params
    std::string server_base; // if set, connect to this server instead of starting a new one
@@ -748,6 +748,10 @@ private:
            optional_props.push_back("*");
        }

+        if (required_props.empty() && optional_props.empty()) {
+            return "\"{\" space \"}\"";
+        }
+
        std::string rule = "\"{\" space ";
        for (size_t i = 0; i < required_props.size(); i++) {
            if (i > 0) {
@@ -438,7 +438,7 @@ void common_log_flush(struct common_log * log) {
    log->resume();
 }

-static int common_get_verbosity(enum ggml_log_level level) {
+int common_log_get_verbosity(enum ggml_log_level level) {
    switch (level) {
        case GGML_LOG_LEVEL_DEBUG: return LOG_LEVEL_DEBUG;
        case GGML_LOG_LEVEL_INFO:  return LOG_LEVEL_TRACE;
@@ -452,7 +452,7 @@ static int common_get_verbosity(enum ggml_log_level level) {
 }

 void common_log_default_callback(enum ggml_log_level level, const char * text, void * /*user_data*/) {
-    auto verbosity = common_get_verbosity(level);
+    auto verbosity = common_log_get_verbosity(level);
    if (verbosity <= common_log_verbosity_thold) {
        common_log_add(common_log_main(), level, "%s", text);
    }
@@ -43,6 +43,8 @@ int  common_log_get_verbosity_thold(void);

 void common_log_set_verbosity_thold(int verbosity); // not thread-safe

+int common_log_get_verbosity(enum ggml_log_level level);
+
 void common_log_default_callback(enum ggml_log_level level, const char * text, void * user_data);

 // the common_log uses an internal worker thread to print/write log messages
@@ -124,6 +124,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "HunYuanMoEV1ForCausalLM": "hunyuan",
    "HunYuanVLForConditionalGeneration": "hunyuan",
    "HYV3ForCausalLM": "hunyuan",
+    "HYV4ForCausalLM": "hy_v4",
    "IQuestCoderForCausalLM": "llama",
    "InternLM2ForCausalLM": "internlm",
    "InternLM3ForCausalLM": "internlm",
@@ -188,6 +189,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "NanbeigeForCausalLM": "nanbeige",
    "NemotronForCausalLM": "nemotron",
    "NemotronHForCausalLM": "nemotron",
+    "NemotronHPuzzleForCausalLM": "nemotron",
    "NeoBERT": "bert",
    "NeoBERTForSequenceClassification": "bert",
    "NeoBERTLMHead": "bert",
@@ -286,6 +288,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "CogVLMForCausalLM": "cogvlm",
    "DeepseekOCR2ForCausalLM": "deepseek",
    "DeepseekOCRForCausalLM": "deepseek",
+    "DeepseekV4ForCausalLM": "deepseek",
    "Dots3NoteForCausalLM": "dots3",
    "Dots3NoteForConditionalGeneration": "dots3",
    "DotsOCRForCausalLM": "dotsocr",
@@ -1507,6 +1507,9 @@ class TextModel(ModelBase):
        if chkhsh == "bba3b3366b646dbdded5dbc42d59598b849371afc42f7beafa914afaa5b70aa6":
            # ref: https://huggingface.co/tencent/Hunyuan-4B-Instruct
            res = "hunyuan-dense"
+        if chkhsh == "e6ddf9c6686791c12d698d34c31ab9be1fea9af5a3d9a6909783ab382198ae1c":
+            # ref: https://huggingface.co/tencent/Hy4-preview
+            res = "hy_v4"
        if chkhsh == "a6b57017d60e6edb4d88ecc2845188e0eb333a70357e45dcc9b53964a73bbae6":
            # ref: https://huggingface.co/tiiuae/Falcon-H1-0.5B-Base
            res = "falcon-h1"
@@ -578,6 +578,8 @@ class DeepseekV4Model(TextModel):
    @classmethod
    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
        name, gen = item
+        if name.startswith(("aligner.", "image_")):
+            return None
        if name.startswith("mtp."):
            if not cls.mtp_only:
                cls._skipped_mtp_tensors += 1
@@ -853,6 +855,7 @@ class DeepseekV4Model(TextModel):
            "ffn_norm.weight": (gguf.MODEL_TENSOR.FFN_NORM, ".weight"),
            "ffn.gate.weight": (gguf.MODEL_TENSOR.FFN_GATE_INP, ".weight"),
            "ffn.gate.bias": (gguf.MODEL_TENSOR.FFN_EXP_PROBS_B, ".bias"),
+            "ffn.gate.bias_vl": (gguf.MODEL_TENSOR.FFN_EXP_PROBS_B_VL, ".bias"),
            "ffn.gate.tid2eid": (gguf.MODEL_TENSOR.FFN_GATE_TID2EID, ".weight"),
            "ffn.shared_experts.w1.weight": (gguf.MODEL_TENSOR.FFN_GATE_SHEXP, ".weight"),
            "ffn.shared_experts.w2.weight": (gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, ".weight"),
@@ -878,6 +881,10 @@ class DeepseekV4Model(TextModel):
        if re.match(r"layers\.\d+\.ffn\.experts\.\d+\.w[123]\.(weight|scale)$", name):
            return []

+        # hash layers route text tokens via tid2eid and image tokens via bias_vl; gate.bias is unused
+        if name.endswith(".ffn.gate.bias") and bid is not None and bid < self.hparams["num_hash_layers"]:
+            return []
+
        tensor_key, suffix = self._map_dsv4_tensor_name(name, bid)
        if tensor_key == gguf.MODEL_TENSOR.FFN_GATE_TID2EID:
            return []
@@ -1000,6 +1007,13 @@ class DeepseekV4DSparkModel(DeepseekV4Model):
            return self._DSPARK_ROOT_MAP[name]
        return super()._map_dsv4_tensor_name(name, bid)

+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # the DFlash draft uses the plain exp-probs bias (ffn.gate.bias -> FFN_EXP_PROBS_B);
+        # the mtmd-only hash routing tensors (bias_vl, tid2eid) are not part of the DFLASH arch
+        if name.endswith(".ffn.gate.bias_vl"):
+            return
+        yield from super().modify_tensors(data_torch, name, bid)
+
    def set_vocab(self):
        if self.target_model_dir is None:
            raise ValueError("DeepSeek-V4 DSpark requires --target-model-dir with the target tokenizer")
@@ -1018,3 +1032,73 @@ class DeepseekV4DSparkModel(DeepseekV4Model):

        self.gguf_writer.add_block_size(self.hparams["dspark_block_size"])
        self.gguf_writer.add_target_layers([layer + 1 for layer in self.hparams["dspark_target_layer_ids"]])
+
+
+@ModelBase.register("DeepseekV4ForCausalLM")
+@ModelBase.example("deepseek-ai/DeepSeek-V4-Flash-Vision-Exp")
+class DeepseekV4FlashVisionModel(MmprojModel):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.hparams_vision is not None
+        # no preprocessor_config.json in the repo; normalization is (x/255 - 0.5) / 0.5
+        # ref: inference/image_processor.py (load_image)
+        self.preprocessor_config = {
+            "image_mean": [0.5, 0.5, 0.5],
+            "image_std":  [0.5, 0.5, 0.5],
+            **self.preprocessor_config,
+        }
+
+    def get_vision_config(self) -> dict[str, Any] | None:
+        cfg = self.global_config
+        if cfg.get("vision_n_layers", 0) == 0:
+            raise ValueError("DeepseekV4FlashVisionModel requires vision_n_layers > 0 in the model config")
+        return {
+            "num_hidden_layers":   cfg["vision_n_layers"],
+            "hidden_size":         cfg["vision_dim"],
+            "num_attention_heads": cfg["vision_n_heads"],
+            "intermediate_size":   cfg["vision_inter_dim"],
+            "patch_size":          cfg["vision_patch_size"],
+            # dynamic resolution; only used for compat / warmup
+            "image_size":          cfg["vision_patch_size"] * cfg["vision_downsample_ratio"] * 16,
+            "rope_theta":          cfg.get("vision_rope_theta", 10000.0),
+            "downsample_ratio":    cfg["vision_downsample_ratio"],
+            "min_pixels":          cfg["vision_min_pixels"],
+        }
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        assert self.hparams_vision is not None
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.DEEPSEEK4V)
+        # vision RMSNorm eps is the pytorch default, NOT the LLM's rms_norm_eps (1e-20)
+        # ref: inference/vision.py (RMSNorm)
+        self.gguf_writer.add_vision_attention_layernorm_eps(1e-6)
+        self.gguf_writer.add_vision_use_silu(True) # SwiGLU MLP
+        self.gguf_writer.add_vision_projector_scale_factor(self.hparams_vision["downsample_ratio"])
+        self.gguf_writer.add_vision_min_pixels(self.hparams_vision["min_pixels"])
+        # hardcoded on the C++ side (see PROJECTOR_TYPE_DEEPSEEK4V in clip.cpp)
+        # if future models use different values, add GGUF keys for those
+        assert self.global_config["vision_max_n_token"] == 384
+        assert self.global_config["vision_max_wh_ratio"] == 8
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, _ = item
+        if not (name.startswith(("vision.", "aligner.", "image_"))):
+            return None
+        return super().filter_tensors(item)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        assert self.hparams_vision is not None
+        if name == "vision.patch_embed.proj.weight":
+            # nn.Linear over flattened (3, p, p) patches == conv2d weight
+            p = self.hparams_vision["patch_size"]
+            data_torch = data_torch.reshape(data_torch.shape[0], 3, p, p)
+
+        if ".mlp.w1." in name:
+            # fused SwiGLU gate+up
+            gate, up = data_torch.chunk(2, dim=0)
+            yield from super().modify_tensors(gate, name.replace("w1", "w1_gate"), bid)
+            yield from super().modify_tensors(up, name.replace("w1", "w1_up"), bid)
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)
@@ -0,0 +1,311 @@
+from __future__ import annotations
+
+import re
+from typing import Iterable
+
+import torch
+
+from .base import ModelBase, gguf, logger
+from .deepseek import DeepseekV2Model
+
+
+def split_kv_b_proj(weight: torch.Tensor, n_head: int, qk_nope: int, v_head_dim: int):
+    """Split kv_b_proj into k_b (transposed) and v_b, matching DeepSeek MLA absorption.
+
+    weight: [n_head*(qk_nope+v_head_dim), kv_lora_rank].
+    Returns (k_b, v_b): k_b [n_head, kv_lora_rank, qk_nope], v_b [n_head, v_head_dim, kv_lora_rank].
+    """
+    kv_lora = weight.shape[-1]
+    assert weight.shape[0] == n_head * (qk_nope + v_head_dim)
+    kv_b = weight.view(n_head, qk_nope + v_head_dim, kv_lora)
+    k_b, v_b = torch.split(kv_b, [qk_nope, v_head_dim], dim=1)
+    k_b = k_b.transpose(1, 2).contiguous()  # [n_head, kv_lora, qk_nope]
+    return k_b, v_b.contiguous()
+
+
+def split_gate_up(weight: torch.Tensor, moe_intermediate_size: int):
+    """Split a fused stacked gate_up expert tensor into (gate, up).
+
+    weight: [n_expert, 2*moe_intermediate_size, hidden] (gate first, up second).
+    Returns (gate, up) each [n_expert, moe_intermediate_size, hidden].
+    """
+    assert weight.shape[1] == 2 * moe_intermediate_size, f"{weight.shape[1]} != 2*{moe_intermediate_size}"
+    gate = weight[:, :moe_intermediate_size, :].contiguous()
+    up = weight[:, moe_intermediate_size:, :].contiguous()
+    return gate, up
+
+
+@ModelBase.register("HYV4ForCausalLM")
+class HYV4Model(DeepseekV2Model):
+    """HY_V4: DeepSeek-V3 style MLA + MoE with iHC, a gated MLA output and a learnable sink.
+
+    Reuses DeepseekV2Model for the vocab and the MLA metadata, but overrides the tensor mapping
+    because HY_V4 ships pre-stacked / fused experts plus extra iHC, gate and sink tensors. The
+    rope rows are mapped straight through (no permute) - the graph rotates consecutive pairs.
+
+    DSA is supported: indexer weights are exported for the layers marked "full" in indexer_types.
+    "shared" layers reuse the top-k of the last preceding full layer at inference time, so they
+    carry no indexer weights.
+
+    MTP (num_nextn_predict_layers) is dropped, so the GGUF cannot be used for speculative
+    decoding. The reference only runs the MTP layers while training or while speculating, so they
+    cannot change single-token logits.
+    """
+
+    model_arch = gguf.MODEL_ARCH.HY_V4
+
+    # tensors a "full" indexer layer must carry
+    INDEXER_SUFFIXES = frozenset({
+        "self_attn.indexer.wq_b.weight",
+        "self_attn.indexer.wk.weight",
+        "self_attn.indexer.k_norm.weight",
+        "self_attn.indexer.k_norm.bias",
+        "self_attn.indexer.weights_proj.weight",
+    })
+
+    @classmethod
+    def filter_tensors(cls, item):
+        # drop MTP here, not in modify_tensors, so the weights are never read
+        if item[0].startswith("model.mtp_layers."):
+            return None
+        return super().filter_tensors(item)
+
+    def _check_indexer_hparams(self):
+        for key in ("index_n_heads", "index_head_dim", "index_topk"):
+            if key not in self.hparams:
+                raise ValueError(f"HY_V4 has DSA layers but no {key}")
+
+    def indexer_is_full(self) -> list[bool] | None:
+        """Per-layer indexer ownership, or None when the checkpoint has no DSA.
+
+        indexer_types entries are "full" (owns an indexer) or "shared" (reuses the preceding
+        full layer's top-k). Missing indexer_types with sparse layers means every sparse layer
+        owns one.
+        """
+        hparams = self.hparams
+        n_layer = hparams["num_hidden_layers"]
+        indexer_types = hparams.get("indexer_types")
+
+        # the reference drives DSA off indexer_types alone; layer_types is only a fallback for
+        # checkpoints predating it (it was renamed to deepseek_sparse_attention upstream)
+        if indexer_types is None:
+            layer_types = hparams.get("layer_types") or []
+            sparse = {"sparse_attention", "deepseek_sparse_attention"}
+            if not any(t in sparse for t in layer_types):
+                return None
+            if len(layer_types) < n_layer:
+                raise ValueError(f"HY_V4 layer_types has {len(layer_types)} entries, need {n_layer}")
+            self._check_indexer_hparams()
+            return [t in sparse for t in layer_types[:n_layer]]
+
+        self._check_indexer_hparams()
+
+        if len(indexer_types) < n_layer:
+            raise ValueError(f"HY_V4 indexer_types has {len(indexer_types)} entries, need {n_layer}")
+        unknown = {t for t in indexer_types[:n_layer]} - {"full", "shared"}
+        if unknown:
+            raise ValueError(f"HY_V4 unknown indexer_types values: {sorted(unknown)}")
+        is_full = [t == "full" for t in indexer_types[:n_layer]]
+        if is_full and not is_full[0]:
+            raise ValueError("HY_V4 layer 0 must be indexer_types 'full' (nothing precedes it to share)")
+        return is_full
+
+    def set_gguf_parameters(self):
+        hparams = self.hparams
+
+        # HY4 has n_group == topk_group == 1 (no group routing). Drop the keys so the base does
+        # not emit expert_group_count/used; llama.cpp then takes the ungrouped MoE path.
+        if hparams.get("n_group") == 1 and hparams.get("topk_group") == 1:
+            hparams.pop("n_group", None)
+            hparams.pop("topk_group", None)
+
+        # HY_V4 config expresses dense/sparse layers via mlp_layer_types, but DeepseekV2Model
+        # needs first_k_dense_replace. Derive it as the contiguous leading "dense" block
+        # (the real config.json also carries first_k_dense_replace; prefer it when present,
+        # but assert the two agree so a mismatch fails loudly).
+        mlp_types = hparams.get("mlp_layer_types")
+        explicit = hparams.get("first_k_dense_replace")
+        derived = None
+        if mlp_types is not None:
+            lead = 0
+            for t in mlp_types:
+                if t == "dense":
+                    lead += 1
+                else:
+                    break
+            if any(t == "dense" for t in mlp_types[lead:]):
+                raise NotImplementedError("HY_V4 converter expects a contiguous leading dense block")
+            derived = lead
+        if explicit is not None and derived is not None and explicit != derived:
+            raise ValueError(
+                f"HY_V4 first_k_dense_replace ({explicit}) disagrees with mlp_layer_types "
+                f"leading-dense count ({derived})"
+            )
+        if explicit is None:
+            if derived is None:
+                raise ValueError("HY_V4 needs first_k_dense_replace or mlp_layer_types to place dense layers")
+            hparams["first_k_dense_replace"] = derived
+
+        # reuse DeepseekV2 MLA + MoE metadata (forces num_key_value_heads=1, writes q/kv lora,
+        # key/value lengths, expert counts, weights scale/norm, rope dims, etc.)
+        super().set_gguf_parameters()
+
+        # HY4 uses DeepSeek-V3 sigmoid routing with e_score_correction_bias. The config has no
+        # scoring_func key, so the base does not write a gating func; set it explicitly.
+        self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
+
+        # routed-expert SwiGLU logits clamp (only routed experts; shared/dense are not clamped,
+        # so swiglu_clamp_shexp is intentionally not written). 0.0 disables the clamp.
+        swiglu_limit = float(hparams.get("swiglu_limit", 0.0) or 0.0)
+        if swiglu_limit > 0.0:
+            self.gguf_writer.add_swiglu_clamp_exp([swiglu_limit] * self.block_count)
+
+        # iHC (independent Hyper-Connections)
+        self.gguf_writer.add_hyper_connection_count(hparams["hc_mult"])
+        self.gguf_writer.add_hyper_connection_epsilon(hparams["hc_eps"])
+        self.gguf_writer.add_hyper_connection_magnitude(hparams["hc_magnitude"])
+
+        # is_full is written explicitly; the graph must not infer it from tensor presence
+        is_full = self.indexer_is_full()
+        if is_full is not None:
+            self.gguf_writer.add_indexer_head_count(hparams["index_n_heads"])
+            self.gguf_writer.add_indexer_key_length(hparams["index_head_dim"])
+            self.gguf_writer.add_indexer_top_k(hparams["index_topk"])
+            self.gguf_writer.add_indexer_types(is_full)
+            logger.info(
+                "HY_V4 DSA: %d/%d layers own an indexer (top_k=%d, n_heads=%d, head_dim=%d)",
+                sum(is_full), len(is_full), hparams["index_topk"],
+                hparams["index_n_heads"], hparams["index_head_dim"],
+            )
+
+        if hparams.get("num_nextn_predict_layers", 0):
+            logger.warning(
+                "HY_V4: dropping %d MTP (nextn) layer(s) - the reference runs them only under "
+                "training / speculative decoding. This GGUF cannot be used for speculative decoding.",
+                hparams["num_nextn_predict_layers"],
+            )
+
+    def prepare_tensors(self):
+        # validate before the base materializes tensors, so a mismatch fails early
+        is_full = self.indexer_is_full()
+        if is_full is not None:
+            present: dict[int, set[str]] = {}
+            for name in self.model_tensors:
+                m = re.match(r"model\.layers\.(\d+)\.(self_attn\.indexer\..+)$", name)
+                if m:
+                    present.setdefault(int(m.group(1)), set()).add(m.group(2))
+            for il, expect_full in enumerate(is_full):
+                seen = present.get(il, set())
+                if expect_full and seen != self.INDEXER_SUFFIXES:
+                    raise ValueError(
+                        f"HY_V4 layer {il} is indexer_types 'full' but is missing indexer tensors: "
+                        f"{sorted(self.INDEXER_SUFFIXES - seen)}"
+                    )
+                if not expect_full and seen:
+                    raise ValueError(
+                        f"HY_V4 layer {il} is indexer_types 'shared' but carries indexer tensors: "
+                        f"{sorted(seen)}"
+                    )
+
+        super().prepare_tensors()
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        # iHC mixing matrices are 2D .weight tensors that the reference keeps in fp32
+        # (_keep_in_fp32_modules_strict). 1D tensors (hc_base/scale, attn_sinks,
+        # e_score_correction_bias) and the router (FFN_GATE_INP) are already forced F32 by the
+        # base rules. Force the HC *_fn matrices here.
+        if new_name.endswith(("hc_attn_fn.weight", "hc_ffn_fn.weight", "output_hc_fn.weight")):
+            return gguf.GGMLQuantizationType.F32
+        # indexer k_norm is fp32 in the reference; the base rules already cover
+        # *_norm.weight and INDEXER_PROJ, but not this bias
+        if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.INDEXER_K_NORM, bid, suffix=".bias"):
+            return gguf.GGMLQuantizationType.F32
+        # enable_lm_head_fp32: mirror the reference fp32 LM-head matmul by keeping output F32.
+        if new_name == "output.weight" and self.hparams.get("enable_lm_head_fp32", False):
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    def modify_tensors(self, data_torch: torch.Tensor, name: str, bid: int | None) -> Iterable[tuple[str, torch.Tensor]]:
+        hparams = self.hparams
+        n_head = hparams["num_attention_heads"]
+        qk_nope = hparams["qk_nope_head_dim"]
+        v_head_dim = hparams["v_head_dim"]
+        moe_inter = hparams["moe_intermediate_size"]
+
+        tn = self.format_tensor_name
+
+        # ---- global (non per-layer) ----
+        if name == "model.embed_tokens.weight":
+            return [(tn(gguf.MODEL_TENSOR.TOKEN_EMBD), data_torch)]
+        if name == "model.norm.weight":
+            return [(tn(gguf.MODEL_TENSOR.OUTPUT_NORM), data_torch)]
+        if name == "lm_head.weight":
+            return [(tn(gguf.MODEL_TENSOR.OUTPUT), data_torch)]
+        if name == "model.hc_head.hc_head_fn":
+            return [(tn(gguf.MODEL_TENSOR.HC_HEAD_FN), data_torch)]
+        if name == "model.hc_head.hc_head_base":
+            return [(tn(gguf.MODEL_TENSOR.HC_HEAD_BASE), data_torch)]
+        if name == "model.hc_head.hc_head_scale":
+            return [(tn(gguf.MODEL_TENSOR.HC_HEAD_SCALE), data_torch)]
+
+        assert bid is not None, f"expected a per-layer tensor, got {name!r}"
+
+        # ---- per-layer, keyed by suffix after 'model.layers.{bid}.' ----
+        suffix = name.split(f"model.layers.{bid}.", 1)[-1]
+
+        # note: q_b_proj and kv_a_proj_with_mqa are mapped straight through (no RoPE permute),
+        # the graph rotates consecutive pairs so the rows need no reordering
+        simple = {
+            "input_layernorm.weight":          (gguf.MODEL_TENSOR.ATTN_NORM, ".weight"),
+            "post_attention_layernorm.weight": (gguf.MODEL_TENSOR.FFN_NORM,  ".weight"),
+            "self_attn.q_a_proj.weight":       (gguf.MODEL_TENSOR.ATTN_Q_A, ".weight"),
+            "self_attn.q_a_layernorm.weight":  (gguf.MODEL_TENSOR.ATTN_Q_A_NORM, ".weight"),
+            "self_attn.q_b_proj.weight":       (gguf.MODEL_TENSOR.ATTN_Q_B, ".weight"),
+            "self_attn.kv_a_proj_with_mqa.weight": (gguf.MODEL_TENSOR.ATTN_KV_A_MQA, ".weight"),
+            "self_attn.kv_a_layernorm.weight": (gguf.MODEL_TENSOR.ATTN_KV_A_NORM, ".weight"),
+            "self_attn.o_proj.weight":         (gguf.MODEL_TENSOR.ATTN_OUT, ".weight"),
+            "self_attn.linear_gate.weight":    (gguf.MODEL_TENSOR.ATTN_GATE, ".weight"),
+            "self_attn.learnable_sink_param":  (gguf.MODEL_TENSOR.ATTN_SINKS, ".weight"),
+            "self_attn.indexer.wq_b.weight":   (gguf.MODEL_TENSOR.INDEXER_ATTN_Q_B, ".weight"),
+            "self_attn.indexer.wk.weight":     (gguf.MODEL_TENSOR.INDEXER_ATTN_K, ".weight"),
+            "self_attn.indexer.k_norm.weight": (gguf.MODEL_TENSOR.INDEXER_K_NORM, ".weight"),
+            "self_attn.indexer.k_norm.bias":   (gguf.MODEL_TENSOR.INDEXER_K_NORM, ".bias"),
+            "self_attn.indexer.weights_proj.weight": (gguf.MODEL_TENSOR.INDEXER_PROJ, ".weight"),
+            "hc_attn_layer.hc_pre.hc_fn":      (gguf.MODEL_TENSOR.HC_ATTN_FN, ".weight"),
+            "hc_attn_layer.hc_pre.hc_base":    (gguf.MODEL_TENSOR.HC_ATTN_BASE, ".weight"),
+            "hc_attn_layer.hc_pre.hc_scale":   (gguf.MODEL_TENSOR.HC_ATTN_SCALE, ".weight"),
+            "hc_mlp_layer.hc_pre.hc_fn":       (gguf.MODEL_TENSOR.HC_FFN_FN, ".weight"),
+            "hc_mlp_layer.hc_pre.hc_base":     (gguf.MODEL_TENSOR.HC_FFN_BASE, ".weight"),
+            "hc_mlp_layer.hc_pre.hc_scale":    (gguf.MODEL_TENSOR.HC_FFN_SCALE, ".weight"),
+            "mlp.gate.weight":                 (gguf.MODEL_TENSOR.FFN_GATE_INP, ".weight"),
+            "mlp.gate.e_score_correction.bias":(gguf.MODEL_TENSOR.FFN_EXP_PROBS_B, ".bias"),
+            "mlp.gate_proj.weight":            (gguf.MODEL_TENSOR.FFN_GATE, ".weight"),
+            "mlp.up_proj.weight":              (gguf.MODEL_TENSOR.FFN_UP, ".weight"),
+            "mlp.down_proj.weight":            (gguf.MODEL_TENSOR.FFN_DOWN, ".weight"),
+            "mlp.shared_experts.gate_proj.weight": (gguf.MODEL_TENSOR.FFN_GATE_SHEXP, ".weight"),
+            "mlp.shared_experts.up_proj.weight":   (gguf.MODEL_TENSOR.FFN_UP_SHEXP, ".weight"),
+            "mlp.shared_experts.down_proj.weight": (gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, ".weight"),
+        }
+        if suffix in simple:
+            key, sfx = simple[suffix]
+            return [(tn(key, bid, sfx), data_torch)]
+
+        # kv_b_proj: split into k_b (transposed) and v_b
+        if suffix == "self_attn.kv_b_proj.weight":
+            k_b, v_b = split_kv_b_proj(data_torch, n_head, qk_nope, v_head_dim)
+            return [
+                (tn(gguf.MODEL_TENSOR.ATTN_K_B, bid), k_b),
+                (tn(gguf.MODEL_TENSOR.ATTN_V_B, bid), v_b),
+            ]
+
+        # fused stacked experts: split gate_up into gate/up
+        if suffix == "mlp.experts.gate_up_proj":
+            gate, up = split_gate_up(data_torch, moe_inter)
+            return [
+                (tn(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid), gate),
+                (tn(gguf.MODEL_TENSOR.FFN_UP_EXP, bid), up),
+            ]
+        if suffix == "mlp.experts.down_proj":
+            return [(tn(gguf.MODEL_TENSOR.FFN_DOWN_EXP, bid), data_torch)]
+
+        raise ValueError(f"Unsupported HY_V4 tensor {name!r} (suffix {suffix!r})")
@@ -5,6 +5,7 @@ from typing import Any, Callable, Iterable, TYPE_CHECKING
 import torch

 if TYPE_CHECKING:
+    from pathlib import Path
    from torch import Tensor

 from .base import MmprojModel, ModelBase, TextModel, gguf, logger
@@ -201,6 +202,7 @@ class NemotronHModel(GraniteHybridModel):
    model_arch = gguf.MODEL_ARCH.NEMOTRON_H
    is_moe: bool = False
    supports_mtp_export = True
+    _experts: list[dict[str, Tensor]] | None = None

    _SSM_LAYER_TYPES = {"mamba", "linear_attention"}
    _ATTN_LAYER_TYPES = {"attention", "full_attention"}
@@ -513,3 +515,88 @@ class NemotronHModel(GraniteHybridModel):
            experts = [k for d in self._experts for k in d.keys()]
            if len(experts) > 0:
                raise ValueError(f"Unprocessed experts: {experts}")
+
+
+@ModelBase.register("NemotronHPuzzleForCausalLM")
+@ModelBase.example("nvidia/NVIDIA-Nemotron-Labs-3-Puzzle-75B-A9B-BF16")
+class NemotronHPuzzleModel(NemotronHModel):
+    """NVIDIA Puzzle: NemotronH with a per-block MoE config (block_configs).
+
+    The checkpoint also ships an MTP draft head (mtp.safetensors). It is skipped
+    here: there is no Puzzle MTP inference path in tree, and the head is laid out
+    by mtp_block_configs rather than the mtp.layers.* form NemotronHModel maps."""
+
+    model_arch = gguf.MODEL_ARCH.NEMOTRON_H_MOE
+    is_moe: bool = True
+    supports_mtp_export = False
+
+    def __init__(self, dir_model: "Path", *args, **kwargs):
+        hparams = dict(kwargs.pop("hparams", None) or ModelBase.load_hparams(dir_model, self.is_mistral_format))
+
+        self.block_configs: list[dict] = hparams["block_configs"]
+        self.n_layer_trunk = len(self.block_configs)
+
+        # block_configs carries the per-block MoE shape, and is the authority on the
+        # block pattern too: the layers_block_type the HF config wrapper computes is
+        # not sized to it.
+        hparams["num_hidden_layers"] = self.n_layer_trunk
+        hparams["layers_block_type"] = [bc["block_type"] for bc in self.block_configs]
+
+        self.model_arch = gguf.MODEL_ARCH.NEMOTRON_H_MOE
+
+        # Bypass NemotronHModel.__init__: it assumes a flat num_experts_per_tok /
+        # moe_intermediate_size and a layers_block_type sized to block_count, neither
+        # of which hold for Puzzle's per-block config.
+        GraniteHybridModel.__init__(self, dir_model, *args, hparams=hparams, **kwargs)
+
+        self.head_dim = self.find_hparam(["head_dim", "attention_head_dim"])
+        self.d_inner = self.find_hparam(["num_heads"]) * self.d_model
+
+        # NemotronHModel.__init__ folds an MTP block into block_count when the
+        # config carries num_nextn_predict_layers; Puzzle's config does, but its
+        # head has a different layout and no inference path, so stay opted out.
+        self._mtp_bid = None
+
+    def set_gguf_parameters(self):
+        GraniteHybridModel.set_gguf_parameters(self)
+
+        head_dim = self.head_dim
+        if head_dim is None:
+            raise ValueError("Could not find the attention head dim in config")
+        self.gguf_writer.add_key_length(head_dim)
+        self.gguf_writer.add_value_length(head_dim)
+
+        ffn_lengths = [bc.get("moe_intermediate_size") or 0 for bc in self.block_configs]
+        experts_used = [bc.get("num_experts_per_tok") or 0 for bc in self.block_configs]
+
+        self.gguf_writer.add_feed_forward_length(ffn_lengths)
+        self.gguf_writer.add_expert_feed_forward_length(ffn_lengths)
+        self.gguf_writer.add_expert_used_count(experts_used)
+
+        self.gguf_writer.add_expert_shared_feed_forward_length(self.hparams["moe_shared_expert_intermediate_size"])
+        self.gguf_writer.add_expert_count(self.hparams["n_routed_experts"])
+        self.gguf_writer.add_expert_shared_count(self.hparams["n_shared_experts"])
+        self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])
+        self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])
+        self.gguf_writer.add_expert_group_count(self.hparams["n_group"])
+        self.gguf_writer.add_moe_latent_size(self.hparams["moe_latent_size"])
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # The official BF16 checkpoint (NVIDIA-Nemotron-Labs-3-Puzzle-75B-A9B-BF16)
+        # names the trunk "model.*" (model.layers.*, model.embeddings, model.norm_f)
+        # where the original release used the NemotronH-style "backbone.*", and spells
+        # the router bias "e_score_correction_bias" instead of "e_score_correction.bias";
+        # normalize so both convert identically.
+        if name.startswith("model."):
+            name = "backbone." + name[len("model."):]
+        if name.endswith("mixer.gate.e_score_correction_bias"):
+            name = name[: -len("e_score_correction_bias")] + "e_score_correction.bias"
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        # Drop the MTP head unconditionally; see the class docstring.
+        if item[0].startswith("mtp."):
+            return None
+        return super().filter_tensors(item)
@@ -276,6 +276,10 @@ class Qwen3TTSSpeakerEncoderModel(MmprojModel):
        # ConvTranspose1d kernels: only F16/F32 are implemented, no BF16
        if new_name.endswith(".conv.weight") and (".up.blk." in new_name or ".dac.blk." in new_name):
            return gguf.GGMLQuantizationType.F32
+        # the code predictor FFN intermediate peaks around 1.5e5, above the F16 range, and mul_mat
+        # casts its input to the weight type
+        if new_name.startswith("a.gen.code.blk.") and new_name.endswith(".ffn_down.weight"):
+            return gguf.GGMLQuantizationType.F32
        return super().tensor_force_quant(name, new_name, bid, n_dims)

    @classmethod
@@ -176,6 +176,7 @@ pre_computed_hashes = [
    {"name": "minerva-7b", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/sapienzanlp/Minerva-7B-base-v1.0", "chkhsh": "1431a23e583c97432bc230bff598d103ddb5a1f89960c8f1d1051aaa944d0b35"},
    {"name": "hunyuan", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tencent/Hunyuan-A13B-Instruct", "chkhsh": "7e57df22b1fe23a7b1e1c7f3dc4e3f96d43a4eb0836d0c6bdc3436d7b2f1c664"},
    {"name": "hunyuan-dense", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tencent/Hunyuan-4B-Instruct", "chkhsh": "bba3b3366b646dbdded5dbc42d59598b849371afc42f7beafa914afaa5b70aa6"},
+    {"name": "hy_v4", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tencent/Hy4-preview", "chkhsh": "e6ddf9c6686791c12d698d34c31ab9be1fea9af5a3d9a6909783ab382198ae1c"},
    # falcon-h1 series uses 4 different tokenizers across model sizes (0.5b - 34b), hence we need to define 4 different hashes
    {"name": "falcon-h1", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tiiuae/Falcon-H1-0.5B-Base", "chkhsh": "a6b57017d60e6edb4d88ecc2845188e0eb333a70357e45dcc9b53964a73bbae6"},
    {"name": "falcon-h1", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tiiuae/Falcon-H1-1B-Base", "chkhsh": "60476e1243776c4fb1b993dbd7a5f15ac22f83c80afdf425fa5ae01c8d44ef86"},
@@ -53,7 +53,7 @@ To see what it might look like visually, here's an old demo of an interactive se
 https://user-images.githubusercontent.com/271616/225014776-1d567049-ad71-4ef2-b050-55b0b3b9274c.mp4

 ## Cross-compile CLI using Android NDK
-It's possible to build `llama.cpp` for Android on your host system via CMake and the Android NDK. If you are interested in this path, ensure you already have an environment prepared to cross-compile programs for Android (i.e., install the Android SDK). Note that, unlike desktop environments, the Android environment ships with a limited set of native libraries, and so only those libraries are available to CMake when building with the Android NDK (see: https://developer.android.com/ndk/guides/stable_apis.)
+It's possible to build `llama.cpp` for Android on your host system via CMake and the Android NDK. If you are interested in this path, ensure you already have an environment prepared to cross-compile programs for Android (i.e., install the Android SDK/NDK and set `ANDROID_NDK` to the NDK root). Note that, unlike desktop environments, the Android environment ships with a limited set of native libraries, and so only those libraries are available to CMake when building with the Android NDK (see: https://developer.android.com/ndk/guides/stable_apis.)

 Once you're ready and have cloned `llama.cpp`, invoke the following in the project directory:

@@ -62,18 +62,22 @@ $ cmake \
  -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \
  -DANDROID_ABI=arm64-v8a \
  -DANDROID_PLATFORM=android-28 \
-  -DCMAKE_C_FLAGS="-march=armv8.7a" \
-  -DCMAKE_CXX_FLAGS="-march=armv8.7a" \
+  -DGGML_NATIVE=OFF \
  -DGGML_OPENMP=OFF \
  -DGGML_LLAMAFILE=OFF \
+  -DLLAMA_OPENSSL=OFF \
  -B build-android
 ```

 Notes:
+  - `GGML_NATIVE=OFF` is required for cross-compilation because the host CPU is not the Android target CPU
  - While later versions of Android NDK ship with OpenMP, it must still be installed by CMake as a dependency, which is not supported at this time
  - `llamafile` does not appear to support Android devices (see: https://github.com/Mozilla-Ocho/llamafile/issues/325)
+  - `LLAMA_OPENSSL=OFF` avoids depending on OpenSSL, which is not part of the Android NDK stable native API set

-The above command should configure `llama.cpp` with the most performant options for modern devices. Even if your device is not running `armv8.7a`, `llama.cpp` includes runtime checks for available CPU features it can use.
+The above command configures a portable Android `arm64-v8a` build. Do not add a global `-march` flag unless you intentionally want to raise the baseline instruction set for every compiled source.
+
+For optional KleidiAI acceleration on Android `arm64-v8a`, see the [Arm KleidiAI section in build.md](./build.md#arm-kleidiai).

 Feel free to adjust the Android ABI for your target. Once the project is configured:

@@ -795,7 +795,8 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_ENABLE_FLASH_ATTN | 1 (default) or 0| Enable Flash-Attention. It can reduce memory usage. The performance impact depends on the LLM.|
 | GGML_SYCL_ENABLE_OPT | 0 or 1 (default)| Enable optimize features for Intel GPUs. (Recommended to 0 for Intel devices older than Gen 10) |
 | GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. |
-| GGML_SYCL_ENABLE_HOST_PINNED_MEM | 0 or 1 (default) | Enable host pinned memory to speed up copy data from host to device. When disable it, host memory will common malloc() on CPU.|
+| GGML_SYCL_ENABLE_HOST_PINNED_MEM | 0 or 1 (default) | Enable host pinned memory to speed up copy data from host to device. When disable it, host memory will common malloc() on CPU. Disable it when use `--load-model mlock`.|
+| GGML_SYCL_HOST_PINNED_MEM_2G | 0 (default) or 1 | Limit the max memory allocation to be no more than 2GB when enable host pinned memory. USM allocations above 2 GiB take the relaxed/large-allocation path, which serializes H2D copies with compute and prevents copy/compute overlap. It will impact the startup time. Need more test. Depend on `GGML_SYCL_ENABLE_HOST_PINNED_MEM=1`.|
 | GGML_SYCL_GET_MEM_API | 0 (default) or 1  | Set to get memory info (free, total) by Level Zero or SYCL API:<br>0 - Level Zero API: support more GPUs, only run on Level Zero running time. When there is an error, fallback to call SYCL API. Depend on GGML_SYCL_SUPPORT_LEVEL_ZERO_API.<br>1 - SYCL API: legacy, support more running time, it can't get the free size of some GPUs (like Arc770). In such case, return total size for free size.|
 | GGML_SYCL_USE_LEVEL_ZERO_API | 1 (default) or 0 | Use Level Zero API for device memory allocation instead of SYCL. Reduces system RAM usage on Intel dGPUs by avoiding DMA-buf/TTM host memory staging. Requires GGML_SYCL_SUPPORT_LEVEL_ZERO_API=ON at build time. SYCL backend always runs on Level Zero running time even if it's set as OFF (The SYCL api will be usage for memory allocation).|
 | GGML_SYCL_ENABLE_DNN | 0 or 1 (default)| Enable running computations through oneDNN and always use oneMKL. |
@@ -805,7 +806,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_ENABLE_MKL_FA | 1 (default) or 0 | Enable oneMKL GEMM flash attention for XMX-accelerated prompt processing with quantized KV cache. Automatically activates during prefill (prompt processing) when all conditions are met: (1) flash-attn enabled (`-fa` or `--flash-attn on`), (2) KV cache quantized (`--cache-type-k q8_0 --cache-type-v q8_0` or other `*_0/*_1` types), (3) batch size ≥ 1024 (`--batch-size 1024`), (4) prompt length ≥ 1024 tokens. Set to 0 to force the TILE kernel for A/B testing. Example minimum command: `llama-cli -m model.gguf -fa -ngl 99 --cache-type-k q8_0 --cache-type-v q8_0 --batch-size 1024 -p "your prompt"` |
 | GGML_SYCL_MKL_FA_DEBUG | 0 (default) or 1 | Enable per-call diagnostic logging for MKL flash attention: GEMM/softmax timings, interleaved-head detection, and buffer memory usage. |
 | GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. |
-| GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. |
+| GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. Unsupported types and layouts fall back to the standalone op kernels. See `ggml_sycl_can_fuse()`. |
 | GGML_SYCL_ENABLE_ESIMD | 0 or 1 (default)| Enable ESIMD kernels when available. |
 | ZES_ENABLE_SYSMAN | 0 (default) or 1 | Support to get free memory of GPU by sycl::aspect::ext_intel_free_memory.<br>Recommended to use when --split-mode = layer |
 | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. |
@@ -27,6 +27,7 @@ The following sections describe how to build with different backends and options
 * [OpenCL](#opencl)
 * [Android](#android-1)
 * [OpenVINO](#openvino)
+* [Hexagon](#hexagon)
 * [Notes about GPU-accelerated backends](#notes-about-gpu-accelerated-backends)

 ## CPU Build
@@ -299,7 +300,6 @@ The following compilation options are also available to tweak performance:
 |-------------------------------|------------------------|---------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | GGML_CUDA_FORCE_MMQ           | Boolean                | false   | Force the use of custom matrix multiplication kernels for quantized models instead of FP16 cuBLAS even if there is no int8 tensor core implementation available (affects V100, CDNA and RDNA3+). MMQ kernels are enabled by default on GPUs with int8 tensor core support. With MMQ force enabled, speed for large batch sizes will be worse but VRAM consumption will be lower. |
 | GGML_CUDA_FORCE_CUBLAS        | Boolean                | false   | Force the use of FP16 cuBLAS instead of custom matrix multiplication kernels for quantized models. There may be issues with numerical overflows (except for V100, CDNA and RDNA4 which use FP32 compute type by default) and memory use will be higher. Prompt processing may become faster on recent datacenter GPUs (the custom kernels were tuned primarily for RTX 3000/4000).   |
-| GGML_CUDA_PEER_MAX_BATCH_SIZE | Positive integer       | 128     | Maximum batch size for which to enable peer access between multiple GPUs. Peer access requires either Linux or NVLink. When using NVLink enabling peer access for larger batch sizes is potentially beneficial.                                                                                                                                                                  |
 | GGML_CUDA_FA_ALL_QUANTS       | Boolean                | false   | Compile support for all KV cache quantization type (combinations) for the FlashAttention CUDA kernels. More fine-grained control over KV cache size but compilation takes much longer.                                                                                                                                                                                           |

 ## MUSA
@@ -614,30 +614,100 @@ You can test with:
 For detailed information about hardware support, setup instructions, and performance optimization, refer to [llama.cpp for ZenDNN](./backend/ZenDNN.md).

 ## Arm® KleidiAI™
-KleidiAI is a library of optimized microkernels for AI workloads, specifically designed for Arm CPUs. These microkernels enhance performance and can be enabled for use by the CPU backend.
+KleidiAI provides optimized Arm CPU microkernels used by the ggml CPU backend. Enabling it at build time makes those kernels available; it does not force every operation to use KleidiAI. At runtime, llama.cpp selects the best compatible CPU kernel from the detected CPU features, tensor type, operation shape, and active backend priority.
+
+Supported targets:
+
+| Platform | Supported ABI / architecture | Notes |
+| --- | --- | --- |
+| Linux | AArch64 / arm64 | Runtime CPU feature detection is automatic. |
+| Android | `arm64-v8a` | Use the Android NDK command below for a portable build. |
+| Apple | arm64 | Runtime CPU feature detection is automatic. Non-streaming SVE vector length is treated as unavailable. |
+| Windows | arm64 | Runtime CPU feature detection is automatic. SMCU count is treated as unknown until a detection path is verified. |
+
+`GGML_CPU_KLEIDIAI=ON` is valid only for AArch64/arm64 builds. Do not enable it for x86, 32-bit Arm, or Android ABIs other than `arm64-v8a`.
+
+### Native AArch64/arm64 build
+
+From the llama.cpp source directory:

-To enable KleidiAI, go to the llama.cpp directory and build using CMake
 ```bash
-cmake -B build -DGGML_CPU_KLEIDIAI=ON
+cmake -S . -B build -DGGML_CPU_KLEIDIAI=ON
 cmake --build build --config Release
 ```
-You can verify that KleidiAI is being used by running
+
+### Android arm64-v8a NDK build
+
+Set `ANDROID_NDK` to the Android NDK root, then run the following from the llama.cpp source directory. This command configures a portable Android `arm64-v8a` build with KleidiAI enabled and avoids Android dependencies that are not part of the NDK stable native API set.
+
+```bash
+cmake -S . -B build-android \
+  -DCMAKE_BUILD_TYPE=Release \
+  -DCMAKE_TOOLCHAIN_FILE="$ANDROID_NDK/build/cmake/android.toolchain.cmake" \
+  -DANDROID_ABI=arm64-v8a \
+  -DANDROID_PLATFORM=android-28 \
+  -DGGML_CPU_KLEIDIAI=ON \
+  -DGGML_NATIVE=OFF \
+  -DGGML_OPENMP=OFF \
+  -DGGML_LLAMAFILE=OFF \
+  -DLLAMA_OPENSSL=OFF
+cmake --build build-android --config Release --parallel
+cmake --install build-android --prefix {install-dir} --config Release
+```
+
+Important Android options:
+
+- `GGML_CPU_KLEIDIAI=ON` enables KleidiAI for Android `arm64-v8a`.
+- `GGML_NATIVE=OFF` is required for cross-compilation because the build host CPU is not the Android target CPU.
+- `GGML_OPENMP=OFF` avoids adding an OpenMP runtime dependency to this NDK command-line build.
+- `GGML_LLAMAFILE=OFF` avoids the llamafile backend, which is not supported on Android.
+- `LLAMA_OPENSSL=OFF` avoids depending on OpenSSL, which is not part of the Android NDK stable native API set.
+
+The Android Studio project under `examples/llama.android` enables KleidiAI automatically for `arm64-v8a`. For Android command-line CMake builds on `arm64-v8a`, pass `-DGGML_CPU_KLEIDIAI=ON` explicitly.
+
+Global -march flags such as `-march=armv8.7a` flag are not required for a portable Android `arm64-v8a` build. Global `-march` flags raise the baseline instruction set for generic code. No manual architecture-specific source selection is required; llama.cpp selects compatible KleidiAI kernels at runtime. The KleidiAI libraries internal CMake handles the -march flags for each particular kernel.
+
+### Verifying the build
+
+Run an installed or in-tree binary:
+
 ```bash
 ./build/bin/llama-cli -m PATH_TO_MODEL -p "What is a car?"
 ```
-If KleidiAI is enabled, the output will contain a line similar to:
+
+If KleidiAI is enabled, the output contains a line similar to:
+
 ```
 load_tensors: CPU_KLEIDIAI model buffer size =  3474.00 MiB
 ```
-KleidiAI’s microkernels implement optimized tensor operations using Arm CPU features such as dotprod, int8mm, SVE, and SME. Llama.cpp selects the most efficient kernels at runtime based on detected CPU capabilities.
-On CPUs that support SME, SME microkernels are enabled automatically using runtime detection.
-The environment variable GGML_KLEIDIAI_SME can be used to control SME behavior:
- Not set: enable SME automatically if supported and detected.
- 0: disable SME.
- <n> > 0: enable SME and assume <n> available SME units (override auto detection).
-If SME is not supported by the CPU, SME microkernels are always disabled.

-Depending on your build target, other higher priority backends may be enabled by default. To ensure the CPU backend is used, you must disable the higher priority backends either at compile time, e.g. -DGGML_METAL=OFF, or during run-time using the command line option `--device none`.
+This confirms that the model has tensors allocated through the KleidiAI CPU buffer. It does not prove that every operation, or any specific SME-family operation, used a KleidiAI microkernel. Runtime CPU features, tensor type, operation shape, and backend priority still control dispatch.
+
+Depending on the build target, another backend may have higher priority than the CPU backend. To force CPU execution for a run, disable higher priority backends at build time, for example `-DGGML_METAL=OFF`, or use a runtime device option such as `--device none` where supported.
+
+### Runtime dispatch
+
+KleidiAI microkernels use Arm CPU features such as dotprod, i8mm, SVE, and SME/SME2. Build-time configuration makes the kernels available. Runtime dispatch selects a compatible kernel for the detected CPU and operation. Older or lower-feature CPUs fall back automatically to compatible kernels.
+
+KleidiAI accelerates selected `GGML_OP_MUL_MAT` paths for F32 and common quantized formats. Exact coverage depends on the bundled KleidiAI version and the llama.cpp runtime selector, so unsupported tensor types, unsupported operation shapes, or higher priority backends may bypass KleidiAI even when the CPU supports the required Arm feature. This is also why a model may not use SME-family kernels on SME-capable hardware.
+
+The current llama.cpp KleidiAI SVE selector only enables SVE kernels when the runtime SVE vector length is known to be QK8_0 bytes, currently 32 bytes. Linux and Android query this at runtime. Apple reports SVE capability separately from userspace non-streaming SVE availability, so llama.cpp treats the SVE vector length as unknown there. Windows exposes SVE feature presence but not the runtime SVE vector length used by this selector, so that value is also treated as unknown. Windows arm64 also treats SMCU count as unknown until a detection mechanism is verified.
+
+The set of available SME-family kernels depends on the bundled KleidiAI version and the detected CPU capabilities. Production configuration does not require any KleidiAI runtime environment variables.
+
+### Diagnostics and debug overrides
+
+KleidiAI runtime environment variables are diagnostics/debug overrides, not production configuration. Leave them unset for normal use.
+
+`GGML_KLEIDIAI_SME` controls SME-family kernel selection and overrides the maximum number of threads assigned to selected quantized SME-family kernels:
+
+- Not set: use automatic runtime detection.
+- `0`: disable SME-family kernels.
+- `<n> > 0`: enable compatible SME-family kernels and allow up to `<n>` threads for quantized SME-family kernels.
+
+On Windows arm64, use `GGML_KLEIDIAI_SME=<n>` as the temporary diagnostics/debug override for SME thread-cap calibration until automatic SMCU count detection is verified.
+
+If the CPU does not support the required SME-family capability for a bundled kernel, that kernel is disabled regardless of the environment variable.

 ## OpenCL

@@ -760,6 +830,9 @@ To read documentation for how to build on IBM Z & LinuxONE, [click here](./build

 For build instructions and usage examples, refer to [OPENVINO.md](backend/OPENVINO.md).

+### Hexagon
+
+Check [README.md](./backend/snapdragon/README.md) for target specific build and run info.

 ---
 ## Notes about GPU-accelerated backends
@@ -12,116 +12,117 @@ Legend:
 - 🟡 Partially supported by this backend
 - ❌ Not supported by this backend

-| Operation | BLAS | CANN | CPU | CUDA | ET | MTL | OpenCL | SYCL | Vulkan | WebGPU | ZenDNN | zDNN |
-|-----------|------|------|------|------|------|------|------|------|------|------|------|------|
-|                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              ACC | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             ADD1 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                           ADD_ID | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                           ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                            CLAMP | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        COL2IM_1D | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                          CONV_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
-|                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                          CONV_3D | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                CONV_TRANSPOSE_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                CONV_TRANSPOSE_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                              COS | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                      COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                              CPY | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
-|               CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                           CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                              DIV | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                     DSV4_HC_COMB | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                     DSV4_HC_POST | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                      DSV4_HC_PRE | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                              DUP | ❌ | ✅ | ✅ | 🟡 | ❌ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                              ELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              EXP | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                            EXPM1 | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             FILL | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                   FLASH_ATTN_EXT | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
-|                            FLOOR | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                  GATED_DELTA_NET | ❌ | ❌ | ✅ | ❌ | ✅ | 🟡 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                GATED_LINEAR_ATTN | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            GEGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        GEGLU_ERF | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                      GEGLU_QUICK | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             GELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                         GELU_ERF | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                       GELU_QUICK | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                         GET_ROWS | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ❌ | ❌ |
-|                    GET_ROWS_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ |
-|                       GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                LIGHTNING_INDEXER | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
-|                 MUL_MAT_HADAMARD | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ❌ |
-|                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                     OPT_STEP_SGD | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                         OUT_PROD | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | 🟡 |
-|                              PAD | ❌ | 🟡 | ✅ | 🟡 | ❌ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ |
-|                   PAD_REFLECT_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                          POOL_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                          POOL_2D | ❌ | 🟡 | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             RELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                      REPEAT_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                         RMS_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                    RMS_NORM_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                             ROLL | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                             ROPE | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        ROPE_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            ROUND | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        RWKV_WKV6 | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                        RWKV_WKV7 | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            SCALE | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              SET | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ |
-|                         SET_ROWS | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | 🟡 | ❌ | ❌ |
-|                              SGN | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                          SIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             SILU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        SILU_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
-|                              SIN | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                         SOFTPLUS | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                         SOFT_MAX | ❌ | 🟡 | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                    SOFT_MAX_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ |
-|                        SOLVE_TRI | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ |
-|                              SQR | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             SQRT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                         SSM_CONV | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                         SSM_SCAN | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ❌ | ❌ |
-|                             STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              SUB | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              SUM | ❌ | 🟡 | ✅ | 🟡 | ❌ | 🟡 | ❌ | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
-|                         SUM_ROWS | ❌ | ✅ | ✅ | 🟡 | ❌ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ |
-|                           SWIGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                       SWIGLU_OAI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|               TIMESTEP_EMBEDDING | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ❌ | ❌ |
-|                              TRI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                            TRUNC | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                          UPSCALE | ❌ | 🟡 | ✅ | ✅ | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                            XIELU | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+| Operation | BLAS | CANN | CPU | CUDA | ET | HTP | MTL | OpenCL | SYCL | Vulkan | WebGPU | ZenDNN | zDNN |
+|-----------|------|------|------|------|------|------|------|------|------|------|------|------|------|
+|                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              ACC | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             ADD1 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                           ADD_ID | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            CLAMP | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                        COL2IM_1D | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
+|                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
+|                          CONV_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
+|                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                          CONV_3D | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                CONV_TRANSPOSE_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                CONV_TRANSPOSE_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                              COS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                      COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                              CPY | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
+|               CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                           CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                              DIV | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                     DSV4_HC_COMB | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     DSV4_HC_POST | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                      DSV4_HC_PRE | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                              DUP | ❌ | ✅ | ✅ | 🟡 | ❌ | ❌ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                              ELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              EXP | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            EXPM1 | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             FILL | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                   FLASH_ATTN_EXT | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
+|                            FLOOR | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                  GATED_DELTA_NET | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                GATED_LINEAR_ATTN | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                            GEGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                        GEGLU_ERF | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                      GEGLU_QUICK | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             GELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         GELU_ERF | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                       GELU_QUICK | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         GET_ROWS | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ❌ | ❌ |
+|                    GET_ROWS_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ |
+|                       GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ❌ |
+|                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                LIGHTNING_INDEXER | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
+|                 MUL_MAT_HADAMARD | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ❌ |
+|                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ |
+|                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                     OPT_STEP_SGD | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                         OUT_PROD | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | 🟡 |
+|                              PAD | ❌ | 🟡 | ✅ | 🟡 | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ |
+|                   PAD_REFLECT_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                          POOL_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                          POOL_2D | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             RELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
+|                      REPEAT_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                         RMS_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                    RMS_NORM_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                             ROLL | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                             ROPE | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                        ROPE_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                            ROUND | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                        RWKV_WKV6 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                        RWKV_WKV7 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                            SCALE | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              SET | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ |
+|                         SET_ROWS | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | 🟡 | ❌ | ❌ |
+|                              SGN | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                          SIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             SILU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                        SILU_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
+|                              SIN | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         SOFTPLUS | ❌ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         SOFT_MAX | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                    SOFT_MAX_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ |
+|                        SOLVE_TRI | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ |
+|                              SQR | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             SQRT | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         SSM_CONV | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         SSM_SCAN | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ❌ | ❌ |
+|                             STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              SUB | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              SUM | ❌ | 🟡 | ✅ | 🟡 | ❌ | ❌ | 🟡 | ❌ | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
+|                         SUM_ROWS | ❌ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ |
+|                           SWIGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                     SWIGLU_CLAMP | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                       SWIGLU_OAI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|               TIMESTEP_EMBEDDING | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ❌ | ❌ |
+|                              TRI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            TRUNC | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                          UPSCALE | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            XIELU | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -734,6 +734,9 @@ class SchemaConverter:
            )
            optional_props.append("*")

+        if not required_props and not optional_props:
+            return '"{" space "}"'
+
        rule = '"{" space '
        rule += ' "," space '.join(prop_kv_rule_names[k] for k in required_props)

@@ -6,6 +6,8 @@ Finetuning of Stories 260K and LLaMA 3.2 1b seems to work with 24 GB of memory.
 **For CPU training, compile llama.cpp without any additional backends such as CUDA.**
 **For CUDA training, use the maximum number of GPU layers.**

+Flash attention is disabled during training because `FLASH_ATTN_EXT` has no backward pass.
+
 Proof of concept:

 ``` sh
@@ -4,7 +4,7 @@ project("ggml" C CXX ASM)

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

@@ -200,8 +200,6 @@ option(GGML_CUDA                            "ggml: use CUDA"
 option(GGML_MUSA                            "ggml: use MUSA"                                  OFF)
 option(GGML_CUDA_FORCE_MMQ                  "ggml: use mmq kernels instead of cuBLAS"         OFF)
 option(GGML_CUDA_FORCE_CUBLAS               "ggml: always use cuBLAS instead of mmq kernels"  OFF)
-set   (GGML_CUDA_PEER_MAX_BATCH_SIZE "128" CACHE STRING
-                                            "ggml: max. batch size for using peer access")
 option(GGML_CUDA_NO_PEER_COPY               "ggml: do not use peer to peer copies"            OFF)
 option(GGML_CUDA_NO_VMM                     "ggml: do not try to use CUDA VMM"                OFF)
 option(GGML_CUDA_FA                         "ggml: compile ggml FlashAttention CUDA kernels"  ON)
@@ -242,6 +240,8 @@ option(GGML_METAL_EMBED_LIBRARY             "ggml: embed Metal library"
 set   (GGML_METAL_MACOSX_VERSION_MIN "" CACHE STRING
                                            "ggml: metal minimum macOS version")
 set   (GGML_METAL_STD "" CACHE STRING       "ggml: metal standard version (-std flag)")
+set   (GGML_METAL_TARGET_OS "macos" CACHE STRING
+                                            "ggml: metal -mtargetos OS name (macos, ios, xros, tvos)")
 option(GGML_OPENMP                          "ggml: use OpenMP"                                ON)
 option(GGML_OPENMP_FETCH                    "ggml: fetch LLVM OpenMP"                         OFF)
 option(GGML_RPC                             "ggml: use RPC"                                   OFF)
@@ -404,10 +404,6 @@ write_basic_package_version_file(
    VERSION ${GGML_INSTALL_VERSION}
    COMPATIBILITY SameMajorVersion)

-target_compile_definitions(ggml-base PRIVATE
-    GGML_VERSION="${GGML_INSTALL_VERSION}"
-    GGML_COMMIT="${GGML_BUILD_COMMIT}"
-)
 message(STATUS "ggml version: ${GGML_INSTALL_VERSION}")
 message(STATUS "ggml commit:  ${GGML_BUILD_COMMIT}")

@@ -424,10 +424,6 @@ extern "C" {
    // Compare the output of two backends
    GGML_API bool ggml_backend_compare_graph_backend(ggml_backend_t backend1, ggml_backend_t backend2, struct ggml_cgraph * graph, ggml_backend_eval_callback callback, void * user_data, struct ggml_tensor const * const * test_nodes, size_t num_test_nodes);

-    // returns true for ops that may require additional memory for fleeting data on some backends,
-    // i.e. the backend's get_alloc_size may return more than ggml_nbytes for the output tensor
-    GGML_API bool ggml_backend_op_alloc_size_may_expand(enum ggml_op op);
-
    // Tensor initialization
    GGML_API enum ggml_status ggml_backend_tensor_alloc(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, void * addr);
    GGML_API enum ggml_status ggml_backend_view_init(struct ggml_tensor * tensor);
@@ -2453,6 +2453,12 @@ extern "C" {
    GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec(
            const struct ggml_tensor * a);

+    // Use finite mask entries as a sparse K/V set. Set 0 to disable.
+    // n_kv_max must bound the number of finite entries in every mask row.
+    GGML_API void ggml_flash_attn_ext_set_n_kv_max(
+            struct ggml_tensor * a,
+            int32_t              n_kv_max);
+
    GGML_API void ggml_flash_attn_ext_add_sinks(
            struct ggml_tensor * a,
            struct ggml_tensor * sinks);
@@ -213,7 +213,9 @@ set_target_properties(ggml-base PROPERTIES
    SOVERSION ${GGML_VERSION_MAJOR}
 )

-target_include_directories(ggml-base PRIVATE .)
+configure_file(ggml-version.h.in ${CMAKE_CURRENT_BINARY_DIR}/ggml-version.h @ONLY)
+
+target_include_directories(ggml-base PRIVATE . ${CMAKE_CURRENT_BINARY_DIR})
 if (GGML_BACKEND_DL)
    target_compile_definitions(ggml-base PUBLIC GGML_BACKEND_DL)
 endif()
@@ -34,6 +34,11 @@ extern "C" {
        void * context;
    };

+    // [TAG_ALLOC_SIZE_EXPAND]
+    // returns true for ops that may require additional memory for fleeting data on some backends,
+    // i.e. the backend buffer type's get_alloc_size may return more than ggml_nbytes for the output tensor
+    GGML_API bool ggml_op_alloc_size_may_expand(enum ggml_op op);
+
    //
    // Backend buffer
    //
@@ -490,7 +490,13 @@ static ggml_backend_reg_t ggml_backend_load_best(const char * name, bool silent,
 #endif
        // default search paths: executable directory, current directory
        search_paths.push_back(get_executable_path());
-        search_paths.push_back(fs::current_path());
+        std::error_code cwd_ec;
+        const fs::path cwd = fs::current_path(cwd_ec);
+        if (cwd_ec) {
+            GGML_LOG_DEBUG("%s: current_path() failure, error-message: %s\n", __func__, cwd_ec.message().c_str());
+        } else {
+            search_paths.push_back(cwd);
+        }
    } else {
        search_paths.push_back(fs::u8path(user_search_path));
    }
@@ -508,8 +514,14 @@ static ggml_backend_reg_t ggml_backend_load_best(const char * name, bool silent,
            }
            continue;
        }
-        fs::directory_iterator dir_it(search_path, fs::directory_options::skip_permission_denied);
-        for (const auto & entry : dir_it) {
+        std::error_code dir_ec;
+        fs::directory_iterator dir_it(search_path, fs::directory_options::skip_permission_denied, dir_ec);
+        if (dir_ec) {
+            GGML_LOG_DEBUG("%s: failed to enumerate %s: %s\n", __func__, path_str(search_path).c_str(), dir_ec.message().c_str());
+            continue;
+        }
+        for (const fs::directory_iterator end; dir_it != end; dir_it.increment(dir_ec)) {
+            const auto & entry = *dir_it;
            if (entry.is_regular_file(ec)) {
                auto filename = entry.path().filename();
                auto ext = entry.path().extension();
@@ -71,7 +71,7 @@ size_t ggml_backend_buft_get_alloc_size(ggml_backend_buffer_type_t buft, const s
        GGML_ASSERT(size <= ggml_nbytes(tensor) ||
                    ggml_op_is_empty(tensor->op) ||
                    ggml_is_quantized(tensor->type) || // [TAG_ALLOC_SIZE_EXPAND]
-                    ggml_backend_op_alloc_size_may_expand(tensor->op));
+                    ggml_op_alloc_size_may_expand(tensor->op));

        return size;
    }
@@ -2109,10 +2109,7 @@ ggml_backend_t ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched,

 // utils

-// [TAG_ALLOC_SIZE_EXPAND]
-// returns true for ops that may require additional memory for fleeting data on some backends,
-// i.e. the backend's get_alloc_size may return more than ggml_nbytes for the output tensor
-bool ggml_backend_op_alloc_size_may_expand(enum ggml_op op) {
+bool ggml_op_alloc_size_may_expand(enum ggml_op op) {
    switch (op) {
        case GGML_OP_FLASH_ATTN_EXT:
        case GGML_OP_MUL_MAT:
@@ -1131,7 +1131,7 @@ GGML_TABLE_END()
 #define NGRID_IQ1S 2048
 #define IQ1S_DELTA 0.125f
 #define IQ1M_DELTA 0.125f
-#if defined(GGML_COMMON_IMPL_C)
+#if defined(GGML_COMMON_IMPL_C) || defined(GGML_COMMON_IMPL_CPP)
 GGML_TABLE_BEGIN(uint64_t, iq1s_grid, NGRID_IQ1S)
    0xffffffffffffffff, 0xffffffffffffff01, 0xffffffffffff0000, 0xffffffffffff01ff,
    0xffffffffffff0101, 0xffffffffff00ff00, 0xffffffffff000000, 0xffffffffff01ffff,
@@ -31,6 +31,8 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
        ggml-cpu/ggml-cpu.cpp
        ggml-cpu/repack.cpp
        ggml-cpu/repack.h
+        ggml-cpu/iqp.cpp
+        ggml-cpu/iqp.h
        ggml-cpu/hbm.cpp
        ggml-cpu/hbm.h
        ggml-cpu/quants.c
@@ -453,12 +455,16 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
                ggml-cpu/spacemit/repack.h
                ggml-cpu/spacemit/ime_env.cpp
                ggml-cpu/spacemit/ime_env.h
-                ggml-cpu/spacemit/ime1_kernels.cpp
-                ggml-cpu/spacemit/ime2_kernels.cpp
                ggml-cpu/spacemit/ime_kernels.h
                ggml-cpu/spacemit/rvv_kernels.cpp
                ggml-cpu/spacemit/rvv_kernels.h
            )
+            if ("RISCV64_SPACEMIT_IME1" IN_LIST RISCV64_SPACEMIT_IME_SPEC)
+                list(APPEND GGML_CPU_SOURCES ggml-cpu/spacemit/ime1_kernels.cpp)
+            endif()
+            if ("RISCV64_SPACEMIT_IME2" IN_LIST RISCV64_SPACEMIT_IME_SPEC)
+                list(APPEND GGML_CPU_SOURCES ggml-cpu/spacemit/ime2_kernels.cpp)
+            endif()
        endif()
        if(NOT GGML_CPU_ALL_VARIANTS)
            set(MARCH_STR "rv64gc")
@@ -636,7 +636,7 @@ void ggml_vec_dot_q5_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const voi
        const float32x4_t v_xyf = vec_float(v_xy);

        const float32x4_t v_d = vec_splats(GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d));
-        const float32x4_t v_acc = vec_madd(v_xyf, v_d, v_acc);
+        const float32x4_t v_acc = vec_madd(v_xyf, v_d, vec_splats(0.0f));

        sumf += vec_hsum_f32x4(v_acc) + summs;
    }
@@ -4,6 +4,7 @@
 #include "ggml-backend-impl.h"
 #include "ggml-backend.h"
 #include "traits.h"
+#include "iqp.h"
 #include "ggml-cpu-impl.h"
 #include "ggml-impl.h"
 #include "quants.h"
@@ -1363,6 +1364,13 @@ UseGgmlGemm1:;

    ggml_barrier(params->threadpool);

+    // IQ panel gemm (see iqp.h) - must come after the barrier above, it consumes the q8_K rows
+    // of src1 from the work buffer
+    if (ggml_cpu_iqp_supports_mul_mat(dst) && !params->use_ref) {
+        ggml_compute_forward_mul_mat_iqp(params, dst);
+        return;
+    }
+
 #if GGML_USE_LLAMAFILE
    if (src1->type != vec_dot_type) {
        const void* wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata;
@@ -1580,6 +1588,16 @@ static void ggml_compute_forward_mul_mat_id(
    char (*atomic_current_chunk)[CACHE_LINE_SIZE] = // [n_as]
        incr_ptr_aligned(&wdata_cur, CACHE_LINE_SIZE * n_as, CACHE_LINE_SIZE);

+    // IQ panel gemm (see iqp.h); per expert eligibility is decided below, but the work buffer is
+    // reserved for the whole node (ggml_graph_plan sizes it without params, use_ref only skips the dispatch)
+    const bool iqp = ggml_cpu_iqp_supports_mul_mat_id(dst) && !params->use_ref;
+
+    char * iqp_panels = NULL;
+
+    if (iqp) {
+        iqp_panels = incr_ptr_aligned(&wdata_cur, nth * ggml_cpu_iqp_scratch_size(dst), 64);
+    }
+
    GGML_ASSERT(params->wsize >= (size_t)((char *) wdata_cur - (char *) params->wdata));

    if (src1->type != vec_dot_type) {
@@ -1651,6 +1669,13 @@ static void ggml_compute_forward_mul_mat_id(
            continue;
        }

+        if (iqp && ggml_cpu_iqp_mul_mat_id_min_batch(cne1)) {
+            ggml_compute_forward_mul_mat_id_iqp(params, dst, cur_a, cne1, (const int32_t *) &MMID_MATRIX_ROW(cur_a, 0),
+                                                iqp_panels);
+
+            continue;
+        }
+
        const char * src0_cur = (const char *) src0->data + cur_a * nb02;
        const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata;
        const size_t row_size = ggml_row_size(vec_dot_type, ne10);
@@ -2858,6 +2883,11 @@ struct ggml_cplan ggml_graph_plan(
                        if (node->src[1]->type != vec_dot_type) {
                            cur = ggml_row_size(vec_dot_type, ggml_nelements(node->src[1]));
                        }
+
+                        // the IQ panel path needs one scratch panel per thread past the q8_K rows
+                        if (ggml_cpu_iqp_supports_mul_mat(node)) {
+                            cur = GGML_PAD(cur, 64) + n_tasks * ggml_cpu_iqp_scratch_size(node);
+                        }
                    } break;
                case GGML_OP_MUL_MAT_ID:
                    {
@@ -2877,6 +2907,10 @@ struct ggml_cplan ggml_graph_plan(
                        cur += n_as*ids->ne[0]*ids->ne[1]*sizeof(struct mmid_row_mapping) + sizeof(int64_t);
                        // atomic_current_chunk
                        cur += CACHE_LINE_SIZE*n_as + CACHE_LINE_SIZE;
+                        // the IQ panel path needs one scratch panel per thread on top of that
+                        if (ggml_cpu_iqp_supports_mul_mat_id(node)) {
+                            cur += n_tasks * ggml_cpu_iqp_scratch_size(node) + 64;
+                        }
                    } break;
                case GGML_OP_OUT_PROD:
                    {
@@ -0,0 +1,39 @@
+#pragma once
+
+#include "ggml-cpu-impl.h"
+#include "ggml.h"
+
+// GGML internal header
+
+// batched mul_mat path for the grid based IQ types: decode 8 src0 rows at a time into per thread scratch
+// (block_iqp_x8, see iqp.cpp) and run an integer gemm over them against all src1 columns
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// whether cne1 rows of src1 are enough for the decode to pay for itself, per expert, for MUL_MAT_ID
+bool ggml_cpu_iqp_mul_mat_id_min_batch(int64_t cne1);
+
+bool ggml_cpu_iqp_supports_mul_mat(const struct ggml_tensor * dst);
+
+// node level test only - per expert eligibility is decided with ggml_cpu_iqp_mul_mat_id_min_batch
+bool ggml_cpu_iqp_supports_mul_mat_id(const struct ggml_tensor * dst);
+
+// per thread panel scratch bytes, padded
+size_t ggml_cpu_iqp_scratch_size(const struct ggml_tensor * dst);
+
+// must be called after src1 has been converted to q8_K into params->wdata and the threads have synchronized on it
+void ggml_compute_forward_mul_mat_iqp(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+
+// one expert: expert_rows points at its row of the matrix_rows table of (i1, i2) int32 pairs, panels at the base of the per thread panel scratches
+void ggml_compute_forward_mul_mat_id_iqp(const struct ggml_compute_params * params,
+                                         struct ggml_tensor *               dst,
+                                         int64_t                            cur_a,
+                                         int64_t                            cne1,
+                                         const int32_t *                    expert_rows,
+                                         void *                             panels);
+
+#ifdef __cplusplus
+}
+#endif
@@ -1823,7 +1823,7 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
        const bool src0_is_kleidiai =
            op->src[0]->buffer &&
            (ggml_n_dims(op->src[0]) == 2) &&
-            op->src[0]->buffer->buft == ggml_backend_cpu_kleidiai_buffer_type() &&
+            op->src[0]->buffer->buft->context == this &&
            slot_total > 0;

        if ((op->op == GGML_OP_MUL_MAT || op->op == GGML_OP_GET_ROWS) &&
@@ -1862,7 +1862,7 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {

    ggml::cpu::tensor_traits * get_tensor_traits(const struct ggml_tensor * op) override {
        if (op->op == GGML_OP_MUL_MAT || op->op == GGML_OP_GET_ROWS) {
-            if (op->src[0]->buffer && op->src[0]->buffer->buft == ggml_backend_cpu_kleidiai_buffer_type()) {
+            if (op->src[0]->buffer && op->src[0]->buffer->buft->context == this) {
                return (ggml::cpu::tensor_traits *) op->src[0]->extra;
            } else {
                // KleidiAI only has kernels for Q4_0 and Q8_0. For a quantized weight of any
@@ -129,8 +129,6 @@ if (CUDAToolkit_FOUND)
                             ${GGML_SOURCES_CUDA}
                            )

-    add_compile_definitions(GGML_CUDA_PEER_MAX_BATCH_SIZE=${GGML_CUDA_PEER_MAX_BATCH_SIZE})
-
    if (GGML_CUDA_GRAPHS)
        add_compile_definitions(GGML_CUDA_USE_GRAPHS)
    endif()
@@ -52,6 +52,7 @@
 #define GGML_CUDA_CC_VOLTA           700
 #define GGML_CUDA_CC_TURING          750
 #define GGML_CUDA_CC_AMPERE          800
+#define GGML_CUDA_CC_ORIN            870
 #define GGML_CUDA_CC_ADA_LOVELACE    890
 #define GGML_CUDA_CC_HOPPER          900
 // While BW spans CC 1000, 1100 & 1200, we are integrating Tensor Core instructions available to 1200 family, see
@@ -718,6 +718,9 @@ static __global__ void flash_attn_mask_to_KV_max(
    KV_max[sequence*ne31 + jt] = KV_max_sj;
 }

+void ggml_cuda_flash_attn_ext_compact_mask(
+        const ggml_tensor * mask, int32_t * indices, int32_t n_kv_max, cudaStream_t stream);
+
 template<int D, int ncols1, int ncols2> // D == head size
 __launch_bounds__(D, 1)
 static __global__ void flash_attn_stream_k_fixup_uniform(
@@ -972,7 +975,8 @@ static __global__ void flash_attn_combine_results(
 template <int DV, int ncols1, int ncols2>
 void launch_fattn(
    ggml_backend_cuda_context & ctx, ggml_tensor * dst, fattn_kernel_t fattn_kernel, const int nwarps, const size_t nbytes_shared,
-    const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const int warp_size = WARP_SIZE
+    const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const bool use_sparse,
+    const int warp_size = WARP_SIZE
 ) {
    constexpr int ncols = ncols1 * ncols2;

@@ -1088,10 +1092,20 @@ void launch_fattn(
    const int ntiles_z_gqa = ((gqa_ratio + ncols2 - 1) / ncols2);
    const int ntiles_dst   = ntiles_x * ntiles_z_gqa * K->ne[2] * Q->ne[3];

+    const int32_t n_kv_max = use_sparse ? ggml_get_op_params_i32(KQV, 4) : 0;
+    if (use_sparse) {
+        GGML_ASSERT(mask != nullptr);
+        GGML_ASSERT(n_kv_max > 0);
+        const size_t mask_rows = size_t(mask->ne[1]) * mask->ne[3];
+
+        KV_max.alloc(size_t(n_kv_max) * mask_rows);
+        ggml_cuda_flash_attn_ext_compact_mask(mask, KV_max.ptr, n_kv_max, main_stream);
+    }
+
    // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped.
    // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or
    //     multiple sequences of possibly different lengths.
-    if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) {
+    if (!use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) {
        const int64_t s31 = mask->nb[1] / sizeof(half2);
        const int64_t s33 = mask->nb[3] / sizeof(half2);

@@ -1114,7 +1128,8 @@ void launch_fattn(
    GGML_ASSERT(max_blocks_per_sm > 0);
    int parallel_blocks = max_blocks_per_sm;

-    const int ntiles_KV = (K->ne[1] + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length.
+    const int64_t n_kv = use_sparse ? n_kv_max : K->ne[1];
+    const int ntiles_KV = (n_kv + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length.

    dim3 blocks_num;
    if (stream_k) {
@@ -1218,7 +1233,7 @@ void launch_fattn(
        !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr,
        scale, max_bias, m0, m1, n_head_log2, logit_softcap,
        Q->ne[0], ne01,     Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3],
-        K->ne[0], K->ne[1], K->ne[2], K->ne[3], nb11, nb12, nb13,
+        K->ne[0], n_kv, K->ne[2], K->ne[3], nb11, nb12, nb13,
        nb21, nb22, nb23,
        mask ? mask->ne[1] : 0, mask ? mask->ne[2] : 0, mask ? mask->ne[3] : 0,
        mask ? mask->nb[1] : 0, mask ? mask->nb[2] : 0, mask ? mask->nb[3] : 0
@@ -2,6 +2,7 @@
 #include "cp-async.cuh"
 #include "mma.cuh"
 #include "fattn-common.cuh"
+#include "fattn-swizzle.cuh"

 using namespace ggml_cuda_mma;

@@ -66,7 +67,7 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co
    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 32, 128, 2,  32,  96,  64,  64, 2, true);
    GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 64, 128, 2,  32,  96,  64,  64, 2, true);

-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8,  64, 4,  64, 128, 128, 128, 2, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8, 128, 2,  64, 128, 128, 128, 2, true);
    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16,  64, 4,  32, 128, 128, 128, 2, true);
    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 128, 2,  32, 128, 128, 128, 2, true);
    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2,  32, 128, 128, 128, 2, true);
@@ -349,20 +350,24 @@ static __host__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV,
    return cp_async_available(cc) && ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2, cc) : 0;
 }

-static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, const int ncols1, const int ncols2) {
+static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(
+        const int DKQ, const int DV, const int ncols1, const int ncols2, const bool use_sparse) {
 #ifdef CP_ASYNC_AVAILABLE
-    return ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2) : 0;
+    const int nstages_target = ncols2 >= 2 ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols1*ncols2) : 0;
+    // sparse gather is not implemented for multi-stage loading
+    return use_sparse && nstages_target > 1 ? 1 : nstages_target;
 #else
-    GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2);
+    GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2, use_sparse);
    return 0;
 #endif // CP_ASYNC_AVAILABLE
 }

 // ------------------------------------------------------------------------------------------------------------------

-template<int stride_tile, int nwarps, int nbatch_fa, bool use_cp_async, bool oob_check>
+template<int stride_tile, bool swz, int nwarps, int nbatch_fa, bool use_cp_async, bool oob_check, bool use_sparse>
 static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
-        const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV, const int i_sup) {
+        const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV,
+        const int k_VKQ_0, const int i_sup, const int32_t * const __restrict__ indices) {
    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
    // K/V data is loaded with decreasing granularity for D for better memory bandwidth.
    // The minimum granularity is 16 bytes.
@@ -370,7 +375,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
    const int chunks_per_row = D2 / h2_per_chunk;
    if constexpr (use_cp_async) {
        static_assert(warp_size == 32, "bad warp_size");
-        static_assert(!oob_check, "OOB check not compatible with cp_async");
+        static_assert(!oob_check || use_sparse, "OOB check not compatible with cp_async");
        constexpr int preload = 64;

        const unsigned int tile_KV_32 = ggml_cuda_cvta_generic_to_shared(tile_KV);
@@ -393,11 +398,25 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
                    break;
                }

+                int64_t i_KV;
+                if constexpr (use_sparse) {
+                    // padded slots gather row 0, the -inf mask removes their contribution
+                    const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : 0;
+                    i_KV = index >= 0 ? index : 0;
+                } else {
+                    i_KV = k_VKQ_0 + i;
+                }
+
 #pragma unroll
                for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) {
                    const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k);

-                    cp_async_cg_16<preload>(tile_KV_32 + i*(stride_tile*sizeof(half2)) + k*16, KV + i*stride_KV + k*h2_per_chunk);
+                    if constexpr (swz) {
+                        const int smem_offs_b = ggml_cuda_fattn_smem_swizzle::bytes_rc<stride_tile>(i, k*h2_per_chunk);
+                        cp_async_cg_16<preload>(tile_KV_32 + smem_offs_b, KV + i_KV*stride_KV + k*h2_per_chunk);
+                    } else {
+                        cp_async_cg_16<preload>(tile_KV_32 + i*(stride_tile*sizeof(half2)) + k*16, KV + i_KV*stride_KV + k*h2_per_chunk);
+                    }
                }
            }
        };
@@ -432,8 +451,18 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
                for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) {
                    const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k);

-                    ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4,
-                        !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero);
+                    const half2 * src;
+                    if constexpr (use_sparse) {
+                        const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : -1;
+                        src = index >= 0 ? KV + int64_t(index)*stride_KV + k*h2_per_chunk : zero;
+                    } else {
+                        src = !oob_check || i < i_sup ? KV + int64_t(k_VKQ_0 + i)*stride_KV + k*h2_per_chunk : zero;
+                    }
+                    if constexpr (swz) {
+                        ggml_cuda_memcpy_1<16>((char *) tile_KV + ggml_cuda_fattn_smem_swizzle::bytes_rc<stride_tile>(i, k*h2_per_chunk), src);
+                    } else {
+                        ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4, src);
+                    }
                }
            }
        };
@@ -447,14 +476,16 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
    }
 }

-template<int ncols1, int nwarps, int nbatch_fa, bool use_cp_async, bool oob_check>
+template<int ncols1, int nwarps, int nbatch_fa, bool use_cp_async, bool oob_check, bool use_sparse>
 static __device__ __forceinline__ void flash_attn_ext_f16_load_mask(
        const half * const __restrict__ mask_h, half * const __restrict__ tile_mask,
-        const int stride_mask, const int i_sup, const int j0, const uint3 ne01) {
+        const int stride_mask, const int k_VKQ_0, const int i_sup, const int j0, const uint3 ne01,
+        const int32_t * const __restrict__ indices) {
    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
    if constexpr (use_cp_async) {
        static_assert(nbatch_fa <= 8*warp_size && nbatch_fa % 8 == 0, "bad nbatch_fa");
        static_assert(!oob_check, "OOB check incompatible with cp_async");
+        static_assert(!use_sparse, "sparse gather incompatible with cp_async");
        constexpr int preload = nbatch_fa >= 32 ? nbatch_fa * sizeof(half) : 64;
        constexpr int cols_per_warp = 8*warp_size/nbatch_fa;
        constexpr int stride_j = nwarps * cols_per_warp;
@@ -472,9 +503,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask(

            const int i = 8 * (threadIdx.x % (nbatch_fa/8));

-            cp_async_cg_16<preload>(tile_mask_32 + j_sram*(nbatch_fa*sizeof(half) + 16) + i*sizeof(half), mask_h + int64_t(j_vram)*stride_mask + i);
+            cp_async_cg_16<preload>(tile_mask_32 + j_sram*(nbatch_fa*sizeof(half) + 16) + i*sizeof(half), mask_h + int64_t(j_vram)*stride_mask + k_VKQ_0 + i);
        }
-    } else if constexpr (oob_check) {
+    } else if constexpr (oob_check || use_sparse) {
 #pragma unroll
        for (int j1 = 0; j1 < ncols1; j1 += nwarps) {
            const int j_sram = j1 + threadIdx.y;
@@ -488,7 +519,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask(
            for (int i0 = 0; i0 < nbatch_fa; i0 += warp_size) {
                const int i = i0 + threadIdx.x;

-                tile_mask[j_sram*(nbatch_fa + 8) + i] = i < i_sup ? mask_h[int64_t(j_vram)*stride_mask + i] : half(0.0f);
+                if constexpr (use_sparse) {
+                    const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : -1;
+                    tile_mask[j_sram*(nbatch_fa + 8) + i] = index >= 0 ? mask_h[int64_t(j_vram)*stride_mask + index] : half(-INFINITY);
+                } else {
+                    tile_mask[j_sram*(nbatch_fa + 8) + i] = i < i_sup ? mask_h[int64_t(j_vram)*stride_mask + k_VKQ_0 + i] : half(0.0f);
+                }
            }
        }
    } else if constexpr (nbatch_fa < 2*warp_size) {
@@ -505,7 +541,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask(

            const int i = threadIdx.x % (warp_size/cols_per_warp);

-            ggml_cuda_memcpy_1<sizeof(half2)>(tile_mask + j_sram*(nbatch_fa + 8) + 2*i, mask_h + int64_t(j_vram)*stride_mask + 2*i);
+            ggml_cuda_memcpy_1<sizeof(half2)>(tile_mask + j_sram*(nbatch_fa + 8) + 2*i, mask_h + int64_t(j_vram)*stride_mask + k_VKQ_0 + 2*i);
        }
    } else {
 #pragma unroll
@@ -521,20 +557,21 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask(
            for (int i0 = 0; i0 < nbatch_fa; i0 += 2*warp_size) {
                const int i = i0 + 2*threadIdx.x;

-                ggml_cuda_memcpy_1<sizeof(half2)>(tile_mask + j_sram*(nbatch_fa + 8) + i, mask_h + int64_t(j_vram)*stride_mask + i);
+                ggml_cuda_memcpy_1<sizeof(half2)>(tile_mask + j_sram*(nbatch_fa + 8) + i, mask_h + int64_t(j_vram)*stride_mask + k_VKQ_0 + i);
            }
        }
    }
 }

 template<int DKQ, int DV, int ncols1, int ncols2, int nwarps,
-    bool use_logit_softcap, bool V_is_K_view, bool needs_fixup, bool is_fixup, bool last_iter, bool oob_check,
+    bool use_logit_softcap, bool V_is_K_view, bool use_sparse, bool needs_fixup, bool is_fixup, bool last_iter, bool oob_check,
    typename T_A_KQ, typename T_B_KQ, typename T_C_KQ, typename T_A_VKQ, typename T_B_VKQ, typename T_C_VKQ>
 static __device__ __forceinline__ void flash_attn_ext_f16_iter(
        const float2 * const __restrict__ Q_f2,
        const half2  * const __restrict__ K_h2,
        const half2  * const __restrict__ V_h2,
        const half   * const __restrict__ mask_h,
+        const int32_t * const __restrict__ indices,
        float2       * const __restrict__ dstk,
        float2       * const __restrict__ dstk_fixup,
        const float scale,
@@ -566,11 +603,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
    constexpr int  nbatch_K2       = ggml_cuda_fattn_mma_get_nbatch_K2(DKQ, DV, ncols);
    constexpr int  nbatch_V2       = ggml_cuda_fattn_mma_get_nbatch_V2(DKQ, DV, ncols);
    constexpr bool Q_in_reg        = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols);
-    constexpr int  nstages         = ggml_cuda_fattn_mma_get_nstages  (DKQ, DV, ncols1, ncols2);
+    constexpr int  nstages         = ggml_cuda_fattn_mma_get_nstages  (DKQ, DV, ncols1, ncols2, use_sparse);

-    constexpr int stride_tile_K = nbatch_K2 + 4;
-
-    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : nbatch_V2 + 4;
+    // swizzle the tile stride for K and V based on the batch size.
+    constexpr int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2);
+    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2);
+    constexpr bool swz_K = ggml_cuda_fattn_smem_swizzle::enabled(nbatch_K2);
+    constexpr bool swz_V = V_is_K_view ? swz_K : ggml_cuda_fattn_smem_swizzle::enabled(nbatch_V2);

    const int k_VKQ_0 = kb0 * nbatch_fa;
 #if defined(TURING_MMA_AVAILABLE)
@@ -588,13 +627,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
        constexpr bool use_cp_async = true;
        cp_async_wait_all();
        __syncthreads();
-        flash_attn_ext_f16_load_tile<stride_tile_V, nwarps, nbatch_fa, use_cp_async, oob_check>
-            (V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, nbatch_V2, stride_V, k_VKQ_sup);
+        flash_attn_ext_f16_load_tile<stride_tile_V, swz_V, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+            (V_h2, tile_V, nbatch_V2, stride_V, k_VKQ_0, k_VKQ_sup, nullptr);
    } else {
-        constexpr bool use_cp_async = nstages == 1;
+        // the sparse mask values are gathered per element, always load them synchronously
+        constexpr bool use_cp_async = nstages == 1 && !use_sparse;
        if (ncols2 > 1 || mask_h) {
-            flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, use_cp_async, oob_check>
-                (mask_h + k_VKQ_0, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01);
+            flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+                (mask_h, tile_mask, stride_mask, k_VKQ_0, k_VKQ_sup, jt*ncols1, ne01, indices);
        }
    }

@@ -607,8 +647,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
        if constexpr (nstages <= 1) {
            const int k0_diff = k0_stop - k0_start;
            constexpr bool use_cp_async = nstages == 1;
-            flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check>
-                (K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup);
+            flash_attn_ext_f16_load_tile<stride_tile_K, swz_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+                (K_h2 + k0_start, tile_K, k0_diff, stride_K, k_VKQ_0, k_VKQ_sup, indices);
            if (use_cp_async) {
                cp_async_wait_all();
            }
@@ -623,7 +663,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
 #pragma unroll
                for (int k_KQ_0 = k0_start; k_KQ_0 < k0_stop; k_KQ_0 += T_A_KQ::J) {
                    T_A_KQ K_A;
-                    load_ldmatrix(K_A, tile_K + i_KQ_0*stride_tile_K + (k_KQ_0 - k0_start), stride_tile_K);
+                    ggml_cuda_fattn_smem_swizzle::load_ldmatrix<stride_tile_K, swz_K>(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start);
                    if constexpr (cols_per_warp == 8) {
                        mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[k_KQ_0/T_A_KQ::J]);
                    } else {
@@ -649,7 +689,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
                    const int i_KQ_0 = i_KQ_00 + (threadIdx.y % np)*T_A_KQ::I;

                    T_A_KQ K_A;
-                    load_ldmatrix(K_A, tile_K + i_KQ_0*stride_tile_K + (k_KQ_0 - k0_start), stride_tile_K);
+                    ggml_cuda_fattn_smem_swizzle::load_ldmatrix<stride_tile_K, swz_K>(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start);

                    if constexpr (cols_per_warp == 8) {
                        mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[0]);
@@ -933,6 +973,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
    }

    if constexpr (nstages > 1) {
+        static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading");
        static_assert(!V_is_K_view, "K data reuse not implemented multi-stage loading");
        // Preload K tile for next iteration:
        constexpr bool use_cp_async = true;
@@ -940,11 +981,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
        __syncthreads();
        if (!last_iter) {
            if (ncols2 > 1 || mask_h) {
-                flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, use_cp_async, oob_check>
-                    (mask_h + k_VKQ_0 + nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01);
+                flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+                    (mask_h, tile_mask, stride_mask, k_VKQ_0 + nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr);
            }
-            flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check>
-                (K_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup);
+            flash_attn_ext_f16_load_tile<stride_tile_K, swz_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+                (K_h2, tile_K, nbatch_K2, stride_K, k_VKQ_0 + nbatch_fa, k_VKQ_sup, nullptr);
        }
    }

@@ -959,8 +1000,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
            const int i0_diff = i0_stop - i0_start;
            if (!V_is_K_view || i0_stop > 2*nbatch_K2) {
                constexpr bool use_cp_async = nstages == 1;
-                flash_attn_ext_f16_load_tile<stride_tile_V, nwarps, nbatch_fa, use_cp_async, oob_check>
-                    (V_h2 + int64_t(k_VKQ_0)*stride_V + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup);
+                flash_attn_ext_f16_load_tile<stride_tile_V, swz_V, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+                    (V_h2 + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_0, k_VKQ_sup, indices);
                if (use_cp_async) {
                    cp_async_wait_all();
                }
@@ -978,7 +1019,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
                const int k0 = k00 + (threadIdx.y % np)*T_A_VKQ::J;

                T_A_VKQ A; // Transposed in SRAM but not in registers, gets transposed on load.
-                load_ldmatrix_trans(A, tile_V_i + 2*k0*stride_tile_V + (i_VKQ_0 - i0_start)/2, stride_tile_V);
+                ggml_cuda_fattn_smem_swizzle::load_ldmatrix_trans<stride_tile_V, swz_V>(A, tile_V, (int)(tile_V_i - tile_V) + 2*k0*stride_tile_V + (i_VKQ_0 - i0_start)/2);
                if constexpr (T_B_KQ::I == 8) {
                    mma(VKQ_C[i_VKQ_0/T_A_VKQ::I], A, B[k00/(np*T_A_VKQ::J)]);
                } else {
@@ -1004,7 +1045,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
                const int k0 = k00 + (threadIdx.y % np)*T_A_VKQ::I;

                T_A_VKQ A; // Transposed in both SRAM and registers, load normally.
-                load_ldmatrix(A, tile_V_i + k0*stride_tile_V + (i_VKQ_0 - i0_start)/2, stride_tile_V);
+                ggml_cuda_fattn_smem_swizzle::load_ldmatrix<stride_tile_V, swz_V>(A, tile_V, (int)(tile_V_i - tile_V) + k0*stride_tile_V + (i_VKQ_0 - i0_start)/2);
                mma(VKQ_C[i_VKQ_0/i0_stride], B[k00/(np*T_A_VKQ::I)], A);
            }
        }
@@ -1015,7 +1056,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
        }
    }
 #else
-    GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup,
+    GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup,
        scale, slope, logit_softcap, ne01, ne02,
        stride_K, stride_V, stride_mask,
        tile_Q, tile_K, tile_V, tile_mask,
@@ -1113,12 +1154,13 @@ template<int DV, int ncols> struct mma_tile_sizes {
 };
 #endif // defined(TURING_MMA_AVAILABLE)

-template<int DKQ, int DV, int ncols1, int ncols2, int nwarps, bool use_logit_softcap, bool V_is_K_view, bool needs_fixup, bool is_fixup>
+template<int DKQ, int DV, int ncols1, int ncols2, int nwarps, bool use_logit_softcap, bool V_is_K_view, bool use_sparse, bool needs_fixup, bool is_fixup>
 static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
        const float2 * const __restrict__ Q_f2,
        const half2  * const __restrict__ K_h2,
        const half2  * const __restrict__ V_h2,
        const half   * const __restrict__ mask_h,
+        const int32_t * const __restrict__ indices,
        const float  * const __restrict__ sinks_f,
        float2       * const __restrict__ dstk,
        float2       * const __restrict__ dstk_fixup,
@@ -1158,7 +1200,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
    constexpr int  nbatch_V2       = ggml_cuda_fattn_mma_get_nbatch_V2     (DKQ, DV, ncols);
    constexpr int  nbatch_combine  = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols);
    constexpr bool Q_in_reg        = ggml_cuda_fattn_mma_get_Q_in_reg      (DKQ, DV, ncols);
-    constexpr int  nstages         = ggml_cuda_fattn_mma_get_nstages       (DKQ, DV, ncols1, ncols2);
+    constexpr int  nstages         = ggml_cuda_fattn_mma_get_nstages       (DKQ, DV, ncols1, ncols2, use_sparse);

    if (cols_per_warp > ncols) {
        NO_DEVICE_CODE;
@@ -1168,10 +1210,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
    static_assert(nwarps * (cols_per_warp/ncols2) % ncols1 == 0, "bad nwarps");

    constexpr int stride_tile_Q = DKQ/2     + 4;
-    constexpr int stride_tile_K = nbatch_K2 + 4;
-
-    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : nbatch_V2 + 4;
+    // swizzle the tile stride for K and V based on the batch size.
+    constexpr int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2);
+    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2);
    constexpr int stride_tile_KV_max = stride_tile_K > stride_tile_V ? stride_tile_K : stride_tile_V;
+    constexpr bool swz_K = ggml_cuda_fattn_smem_swizzle::enabled(nbatch_K2);
+    constexpr bool swz_V = V_is_K_view ? swz_K : ggml_cuda_fattn_smem_swizzle::enabled(nbatch_V2);

    extern __shared__ half2 tile_Q[];
    half2 * tile_K    = Q_in_reg              ? tile_Q                             : tile_Q + ncols     * stride_tile_Q;
@@ -1257,37 +1301,38 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(

    // Preload mask and K data for first iteration when using cp_async with multiple stages:
    if constexpr (nstages > 1) {
+        static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading");
        static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi-stage pipeline");
        constexpr bool use_cp_async = true;
        constexpr bool oob_check    = false;
        constexpr int  k_VKQ_sup    = nbatch_fa;
        if (ncols2 > 1 || mask_h) {
-            flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, use_cp_async, oob_check>
-                (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01);
+            flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+                (mask_h, tile_mask, stride_mask, kb0*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr);
        }
-        flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check>
-            (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup);
+        flash_attn_ext_f16_load_tile<stride_tile_K, swz_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+            (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, nullptr);
    }

    // kb0_start is always < kb0_stop so the last iter can be executed unconditionally.
-    if constexpr (ncols2 == 1) {
+    if constexpr (ncols2 == 1 || use_sparse) {
        constexpr bool oob_check = true;
        for (; kb0 < kb0_stop-1; ++kb0) {
            constexpr bool last_iter = false;
            constexpr int  k_VKQ_sup = nbatch_fa;
            flash_attn_ext_f16_iter
-                <DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, needs_fixup, is_fixup, last_iter, oob_check,
+                <DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup, last_iter, oob_check,
                 T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ>
-                (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap,
+                (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap,
                 ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C,
                 KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup);
        }
        constexpr bool last_iter = true;
        const     int  k_VKQ_sup = ne11 - kb0*nbatch_fa;
        flash_attn_ext_f16_iter
-            <DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, needs_fixup, is_fixup, last_iter, oob_check,
+            <DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup, last_iter, oob_check,
              T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ>
-            (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap,
+            (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap,
             ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C,
             KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup);
    } else {
@@ -1296,18 +1341,18 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
            constexpr bool last_iter = false;
            constexpr int  k_VKQ_sup = nbatch_fa;
            flash_attn_ext_f16_iter
-                <DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, needs_fixup, is_fixup, last_iter, oob_check,
+                <DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup, last_iter, oob_check,
                 T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ>
-                (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap,
+                (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap,
                 ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C,
                 KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup);
        }
        constexpr bool last_iter = true;
        constexpr int  k_VKQ_sup = nbatch_fa;
        flash_attn_ext_f16_iter
-            <DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, needs_fixup, is_fixup, last_iter, oob_check,
+            <DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup, last_iter, oob_check,
             T_A_KQ, T_B_KQ, T_C_KQ, T_A_VKQ, T_B_VKQ, T_C_VKQ>
-            (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap,
+            (Q_f2, K_h2, V_h2, mask_h, indices, dstk, dstk_fixup, scale, slope, logit_softcap,
             ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C,
             KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup);
    }
@@ -1430,11 +1475,17 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
    constexpr int tile_stride = nbatch_combine + 4;
    static_assert((DV/2) % nbatch_combine == 0, "bad nbatch_combine");

+    constexpr bool combine_needs_sync = swz_K || swz_V;
+
    if constexpr (cols_per_warp == 8) {
        const int jc_cwmo = (threadIdx.x % (2*T_C_VKQ::J)) / T_C_VKQ::J; // jc combine write meta offset
        const int jc_cwm = threadIdx.y*(2*T_C_VKQ::J) + 2*T_C_VKQ::get_j(-1) + jc_cwmo; // jc combine write meta
        const float2 KQ_cmr = make_float2(KQ_max[jc_cwmo], KQ_rowsum[jc_cwmo]); // KQ combine max rowsum

+        if constexpr (combine_needs_sync) {
+            __syncthreads();
+        }
+
        if (((!needs_fixup && !is_fixup) || np > 1) && threadIdx.x < 2*T_C_VKQ::J) {
            // Use the 16 bytes of padding in each row to store the meta data: KQ max, KQ rowsum, KQ max scale.
            ((float2 *) tile_Q)[jc_cwm*(tile_stride/2) + nbatch_combine/2] = KQ_cmr;
@@ -1471,6 +1522,10 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
        const bool thread_should_write = T_C_KQ::J == 8 || T_C_KQ::get_j(threadIdx.x & 2) < 8;
 #endif // defined(TURING_MMA_AVAILABLE)

+        if constexpr (combine_needs_sync) {
+            __syncthreads();
+        }
+
        if (((!needs_fixup && !is_fixup) || np > 1) && thread_should_write) {
            ((float2 *) tile_Q)[jc_cwm*(tile_stride/2) + nbatch_combine/2] = KQ_cmr;
        }
@@ -1692,7 +1747,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
        }
    }
 #else
-    GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dstk_fixup,
+    GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dstk_fixup,
        scale, slope, logit_softcap, ne01, ne02, gqa_ratio,
        stride_Q1, stride_Q2, stride_K, stride_V, stride_mask,
        jt, kb0_start, kb0_stop);
@@ -1700,7 +1755,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
 #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
 }

-template<int DKQ, int DV, int ncols1, int ncols2, bool use_logit_softcap, bool V_is_K_view>
+static constexpr __host__ __device__ bool ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(
+        const int DKQ, const int DV, const int ncols1, const int ncols2) {
+    return (DKQ == 512 && DV == 512 && ncols1 == 1 && ncols2 == 8) ||
+           (DKQ == 576 && DV == 512 && ncols1 == 1 && ncols2 == 16);
+}
+
+template<int DKQ, int DV, int ncols1, int ncols2, bool use_logit_softcap, bool V_is_K_view, bool use_sparse>
 __launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2))
 static __global__ void flash_attn_ext_f16(
        const char * Q_ptr,
@@ -1726,14 +1787,15 @@ static __global__ void flash_attn_ext_f16(
                            const int32_t nb31, const int32_t nb32, const int64_t nb33) {
    ggml_cuda_pdl_sync(); // TODO optimize placement
 #if defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE))
-    const char * GGML_CUDA_RESTRICT Q        = Q_ptr;
-    const char * GGML_CUDA_RESTRICT K        = K_ptr;
-    const char * GGML_CUDA_RESTRICT V        = V_ptr;
-    const char * GGML_CUDA_RESTRICT mask     = mask_ptr;
-    const char * GGML_CUDA_RESTRICT sinks    = sinks_ptr;
-    const int  * GGML_CUDA_RESTRICT KV_max   = KV_max_ptr;
-    float      * GGML_CUDA_RESTRICT dst      = dst_ptr;
-    float2     * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr;
+    const char * GGML_CUDA_RESTRICT Q              = Q_ptr;
+    const char * GGML_CUDA_RESTRICT K              = K_ptr;
+    const char * GGML_CUDA_RESTRICT V              = V_ptr;
+    const char * GGML_CUDA_RESTRICT mask           = mask_ptr;
+    const char * GGML_CUDA_RESTRICT sinks          = sinks_ptr;
+    const int  * GGML_CUDA_RESTRICT KV_max         = use_sparse ? nullptr : KV_max_ptr;
+    const int  * GGML_CUDA_RESTRICT sparse_indices = use_sparse ? KV_max_ptr : nullptr;
+    float      * GGML_CUDA_RESTRICT dst            = dst_ptr;
+    float2     * GGML_CUDA_RESTRICT dst_meta       = dst_meta_ptr;

    // Skip unused kernel variants for faster compilation:
    if (use_logit_softcap && !(DKQ == 128 || DKQ == 256 || DKQ == 512)) {
@@ -1744,6 +1806,11 @@ static __global__ void flash_attn_ext_f16(
        NO_DEVICE_CODE;
        return;
    }
+
+    if (!ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2) && use_sparse) {
+        NO_DEVICE_CODE;
+        return;
+    }
 #ifdef VOLTA_MMA_AVAILABLE
    if (ncols1*ncols2 < 32) {
        NO_DEVICE_CODE;
@@ -1820,6 +1887,7 @@ static __global__ void flash_attn_ext_f16(

        const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV);
        const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr;
+        const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt*ncols1)*ne11 : nullptr;

        const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f;

@@ -1829,13 +1897,13 @@ static __global__ void flash_attn_ext_f16(
        constexpr bool is_fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
        if (kb0_start == 0) {
            constexpr bool needs_fixup = false; // CUDA block is working on an entire tile.
-            flash_attn_ext_f16_process_tile<DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, needs_fixup, is_fixup>
-                (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
+            flash_attn_ext_f16_process_tile<DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup>
+                (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
                 ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop);
        } else {
            constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile.
-            flash_attn_ext_f16_process_tile<DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, needs_fixup, is_fixup>
-                (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
+            flash_attn_ext_f16_process_tile<DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup>
+                (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
                 ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop);
        }

@@ -1866,6 +1934,7 @@ static __global__ void flash_attn_ext_f16(

    const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV);
    const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr;
+    const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*ne31 + jt*ncols1)*ne11 : nullptr;

    const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f;

@@ -1875,8 +1944,8 @@ static __global__ void flash_attn_ext_f16(

    constexpr bool is_fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
    constexpr bool needs_fixup = false;
-    flash_attn_ext_f16_process_tile<DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, needs_fixup, is_fixup>
-        (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
+    flash_attn_ext_f16_process_tile<DKQ, DV, ncols1, ncols2, nwarps, use_logit_softcap, V_is_K_view, use_sparse, needs_fixup, is_fixup>
+        (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
         ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop);
 #else
    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
@@ -1892,6 +1961,8 @@ static __global__ void flash_attn_ext_f16(
 #endif // defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE))
 }

+bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
 template <int DKQ, int DV, int ncols1, int ncols2>
 void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
    const ggml_tensor * KQV = dst;
@@ -1914,8 +1985,11 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml

    constexpr bool V_is_K_view = DKQ == 576; // Guaranteed by the kernel selection logic in fattn.cu

-    const size_t nbytes_shared_KV_1stage = nbatch_fa            * std::max(nbatch_K2 + 4,  nbatch_V2 + 4) * sizeof(half2);
-    const size_t nbytes_shared_KV_2stage = nbatch_fa            *         (nbatch_K2 + 4 + nbatch_V2 + 4) * sizeof(half2);
+    // KV tile strides must match flash_attn_ext_f16_iter / _process_tile.
+    const int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2, cc);
+    const int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2, cc);
+    const size_t nbytes_shared_KV_1stage = nbatch_fa            * std::max(stride_tile_K,  stride_tile_V) * sizeof(half2);
+    const size_t nbytes_shared_KV_2stage = nbatch_fa            *         (stride_tile_K + stride_tile_V) * sizeof(half2);
    const size_t nbytes_shared_Q         = ncols                * (DKQ/2 + 4)                             * sizeof(half2);
    const size_t nbytes_shared_mask      = ncols1               * (nbatch_fa/2 + 4)                       * sizeof(half2);
    const size_t nbytes_shared_combine   = nwarps*cols_per_warp * (nbatch_combine + 4)                    * sizeof(half2);
@@ -1935,20 +2009,49 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
    using fattn_kernel_ptr_t = fattn_kernel_t;
 #endif // defined(GGML_USE_HIP)
    fattn_kernel_t fattn_kernel;
+    bool use_sparse = false;
    if (logit_softcap == 0.0f) {
        constexpr bool use_logit_softcap = false;
-        fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view>;
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+        if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2)) {
+            if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst)) {
+                constexpr bool use_sparse_kernel = true;
+                fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;
+                use_sparse = true;
+
+                static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false};
+                if (!shared_memory_limit_raised[id]) {
+                    CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast<fattn_kernel_ptr_t>(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total));
+                    shared_memory_limit_raised[id] = true;
+                }
+            } else {
+                constexpr bool use_sparse_kernel = false;
+                fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;
+
+                static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false};
+                if (!shared_memory_limit_raised[id]) {
+                    CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast<fattn_kernel_ptr_t>(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total));
+                    shared_memory_limit_raised[id] = true;
+                }
+            }
+        } else
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+        {
+            constexpr bool use_sparse_kernel = false;
+            fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;

 #if !defined(GGML_USE_MUSA)
-        static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false};
-        if (!shared_memory_limit_raised[id]) {
-            CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast<fattn_kernel_ptr_t>(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total));
-            shared_memory_limit_raised[id] = true;
-        }
+            static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false};
+            if (!shared_memory_limit_raised[id]) {
+                CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast<fattn_kernel_ptr_t>(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total));
+                shared_memory_limit_raised[id] = true;
+            }
 #endif // !defined(GGML_USE_MUSA)
+        }
    } else {
        constexpr bool use_logit_softcap = true;
-        fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view>;
+        constexpr bool use_sparse_kernel = false;
+        fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;

 #if !defined(GGML_USE_MUSA)
        static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false};
@@ -1960,7 +2063,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
    }

    launch_fattn<DV, ncols1, ncols2>
-        (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host);
+        (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host);
 }


@@ -0,0 +1,126 @@
+#pragma once
+
+#include "common.cuh"
+#include "mma.cuh"
+
+// XOR swizzle for K/V SMEM tiles to avoid bank conflicts without row padding (Turing+ only).
+// Stride must be a multiple of 32 half2 columns, otherwise we keep +4 row padding.
+
+namespace ggml_cuda_fattn_smem_swizzle {
+
+static __host__ __device__ constexpr bool bank_aligned(const int nbatch_2) {
+    return nbatch_2 >= 32 && nbatch_2 % 32 == 0;
+}
+
+static __device__ constexpr bool enabled(const int nbatch_2) {
+#if defined(TURING_MMA_AVAILABLE)
+    return bank_aligned(nbatch_2);
+#else
+    GGML_UNUSED(nbatch_2);
+    return false;
+#endif // defined(TURING_MMA_AVAILABLE)
+}
+
+static __host__ bool enabled(const int nbatch_2, const int cc) {
+#ifdef GGML_USE_HIP
+    GGML_UNUSED(nbatch_2);
+    GGML_UNUSED(cc);
+    return false;
+#else
+    return turing_mma_available(cc) && bank_aligned(nbatch_2);
+#endif // GGML_USE_HIP
+}
+
+static __device__ constexpr int tile_stride(const int nbatch_2) {
+    return enabled(nbatch_2) ? nbatch_2 : nbatch_2 + 4;
+}
+
+static __host__ int tile_stride(const int nbatch_2, const int cc) {
+    return enabled(nbatch_2, cc) ? nbatch_2 : nbatch_2 + 4;
+}
+
+// Swizzled byte offset for tile element (row, col_h2), same map used for writes and reads.
+template<int stride_h2>
+static __device__ __forceinline__ int bytes_rc(const int row, const int col_h2) {
+    static_assert(bank_aligned(stride_h2), "swizzled tile needs a stride that is a multiple of 32");
+    return ((row * stride_h2 + col_h2) * (int) sizeof(half2)) ^ ((row & 7) << 4);
+}
+
+// ldmatrix.x4 via 64-bit generic pointer.
+static __device__ __forceinline__ void ldmatrix_x4(int * xi, const half2 * addr) {
+#if defined(TURING_MMA_AVAILABLE)
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3])
+        : "l"(addr));
+#else
+    GGML_UNUSED_VARS(xi, addr);
+    NO_DEVICE_CODE;
+#endif // defined(TURING_MMA_AVAILABLE)
+}
+
+static __device__ __forceinline__ void ldmatrix_x4_trans(int * xi, const half2 * addr) {
+#if defined(TURING_MMA_AVAILABLE)
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3])
+        : "l"(addr));
+#else
+    GGML_UNUSED_VARS(xi, addr);
+    NO_DEVICE_CODE;
+#endif // defined(TURING_MMA_AVAILABLE)
+}
+
+// Per-lane swizzled address for one tile<16, 8, half2> ldmatrix: 16 rows, 4 half2 columns per lane.
+template<int stride_h2>
+static __device__ __forceinline__ const half2 * lane_addr(
+        const half2 * tile_base, const int base_row, const int base_col_h2, const int I, const int J) {
+    static_assert(bank_aligned(stride_h2), "swizzled tile needs a stride that is a multiple of 32");
+    const int lane_row = threadIdx.x % I;
+    const int lane_col = (threadIdx.x / I) * (J / 2);
+    uint32_t byte_off = (uint32_t) ((base_row + lane_row)*stride_h2 + base_col_h2 + lane_col) * (uint32_t) sizeof(half2);
+    byte_off ^= (uint32_t) (((base_row + lane_row) & 7) << 4);
+    return (const half2 *) ((const char *) tile_base + byte_off);
+}
+
+template<int stride_h2, bool swz, typename TileT>
+static __device__ __forceinline__ void load_ldmatrix(
+        TileT & t, const half2 * tile_base, const int base_row, const int base_col_h2) {
+    if constexpr (swz) {
+        static_assert(std::is_same_v<TileT, ggml_cuda_mma::tile<16, 8, half2>>,
+            "the swizzled layout is only supported for tile<16, 8, half2>");
+        ldmatrix_x4((int *) t.x, lane_addr<stride_h2>(tile_base, base_row, base_col_h2, TileT::I, TileT::J));
+    } else {
+        ggml_cuda_mma::load_ldmatrix(t, tile_base + base_row*stride_h2 + base_col_h2, stride_h2);
+    }
+}
+
+template<int stride_h2, bool swz, typename TileT>
+static __device__ __forceinline__ void load_ldmatrix(TileT & t, const half2 * tile_base, const int off_h2) {
+    if constexpr (swz) {
+        load_ldmatrix<stride_h2, swz>(t, tile_base, off_h2 / stride_h2, off_h2 % stride_h2);
+    } else {
+        ggml_cuda_mma::load_ldmatrix(t, tile_base + off_h2, stride_h2);
+    }
+}
+
+template<int stride_h2, bool swz, typename TileT>
+static __device__ __forceinline__ void load_ldmatrix_trans(
+        TileT & t, const half2 * tile_base, const int base_row, const int base_col_h2) {
+    if constexpr (swz) {
+        static_assert(std::is_same_v<TileT, ggml_cuda_mma::tile<16, 8, half2>>,
+            "the swizzled layout is only supported for tile<16, 8, half2>");
+        ldmatrix_x4_trans((int *) t.x, lane_addr<stride_h2>(tile_base, base_row, base_col_h2, TileT::I, TileT::J));
+    } else {
+        ggml_cuda_mma::load_ldmatrix_trans(t, tile_base + base_row*stride_h2 + base_col_h2, stride_h2);
+    }
+}
+
+template<int stride_h2, bool swz, typename TileT>
+static __device__ __forceinline__ void load_ldmatrix_trans(TileT & t, const half2 * tile_base, const int off_h2) {
+    if constexpr (swz) {
+        load_ldmatrix_trans<stride_h2, swz>(t, tile_base, off_h2 / stride_h2, off_h2 % stride_h2);
+    } else {
+        ggml_cuda_mma::load_ldmatrix_trans(t, tile_base + off_h2, stride_h2);
+    }
+}
+
+} // namespace ggml_cuda_fattn_smem_swizzle
@@ -1163,7 +1163,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm
            const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc);
            fattn_kernel_t fattn_kernel = flash_attn_tile<DKQ, DV, cols_per_block/ncols2, ncols2, use_logit_softcap>;
            launch_fattn<DV, cols_per_block/ncols2, ncols2>
-                (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size);
+                (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size);
            return;
        }
    }
@@ -1179,7 +1179,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm
            const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc);
            fattn_kernel_t fattn_kernel = flash_attn_tile<DKQ, DV, cols_per_block/ncols2, ncols2, use_logit_softcap>;
            launch_fattn<DV, cols_per_block/ncols2, ncols2>
-                (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size);
+                (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size);
            return;
        }
    }
@@ -1191,7 +1191,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm
            const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc);
            fattn_kernel_t fattn_kernel = flash_attn_tile<DKQ, DV, cols_per_block/ncols2, ncols2, use_logit_softcap>;
            launch_fattn<DV, cols_per_block/ncols2, ncols2>
-                (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size);
+                (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size);
            return;
        }
    }
@@ -1203,7 +1203,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm
            const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc);
            fattn_kernel_t fattn_kernel = flash_attn_tile<DKQ, DV, cols_per_block/ncols2, ncols2, use_logit_softcap>;
            launch_fattn<DV, cols_per_block/ncols2, ncols2>
-                (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size);
+                (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size);
            return;
        }
    }
@@ -1215,7 +1215,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm
            const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc);
            fattn_kernel_t fattn_kernel = flash_attn_tile<DKQ, DV, cols_per_block/ncols2, ncols2, use_logit_softcap>;
            launch_fattn<DV, cols_per_block/ncols2, ncols2>
-                (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size);
+                (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size);
            return;
        }
    }
@@ -1226,7 +1226,7 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm
        const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc);
        fattn_kernel_t fattn_kernel = flash_attn_tile<DKQ, DV, cols_per_block/ncols2, ncols2, use_logit_softcap>;
        launch_fattn<DV, cols_per_block/ncols2, ncols2>
-            (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, warp_size);
+            (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size);
        return;
    }

@@ -540,7 +540,7 @@ void ggml_cuda_flash_attn_ext_vec_case_impl(ggml_backend_cuda_context & ctx, ggm
    const bool need_f16_K = type_K == GGML_TYPE_F16;
    const bool need_f16_V = type_V == GGML_TYPE_F16;
    constexpr size_t nbytes_shared = 0;
-    launch_fattn<D, cols_per_block, 1>(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
+    launch_fattn<D, cols_per_block, 1>(ctx, dst, fattn_kernel, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false, false);
 }

 template <int D, ggml_type type_K, ggml_type type_V>
@@ -5,11 +5,144 @@
 #include "fattn-vec.cuh"
 #include "fattn.cuh"

+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+__launch_bounds__(256, 1)
+static __global__ void flash_attn_mask_to_sparse_indices(
+        const half * mask_ptr, int32_t * indices_ptr, const int ne30, const int n_kv_max,
+        const int64_t s31, const int64_t s33) {
+    ggml_cuda_pdl_sync();
+
+    constexpr int values_per_lane = 8;
+    const int tid      = threadIdx.x;
+    const int warp     = tid / WARP_SIZE;
+    const int lane     = tid % WARP_SIZE;
+    const int sequence = blockIdx.y;
+    const int query    = blockIdx.x;
+
+    const half * mask = mask_ptr + sequence*s33 + query*s31;
+    int32_t * indices = indices_ptr + (int64_t(sequence)*gridDim.x + query)*n_kv_max;
+
+    __shared__ int warp_offsets[256/WARP_SIZE];
+    __shared__ int row_count;
+    __shared__ int chunk_count;
+
+    if (tid == 0) {
+        row_count = 0;
+    }
+    __syncthreads();
+
+    for (int i0 = 0; i0 < ne30; i0 += blockDim.x*values_per_lane) {
+        uint32_t selected_warp[values_per_lane];
+        int warp_count = 0;
+#pragma unroll
+        for (int item = 0; item < values_per_lane; ++item) {
+            const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane;
+            const bool selected = i < ne30 && isfinite(__half2float(mask[i]));
+            selected_warp[item] = __ballot_sync(0xFFFFFFFF, selected);
+            warp_count += __popc(selected_warp[item]);
+        }
+
+        if (lane == 0) {
+            warp_offsets[warp] = warp_count;
+        }
+        __syncthreads();
+
+        if (tid == 0) {
+            int offset = 0;
+#pragma unroll
+            for (int iw = 0; iw < 256/WARP_SIZE; ++iw) {
+                const int count = warp_offsets[iw];
+                warp_offsets[iw] = offset;
+                offset += count;
+            }
+            chunk_count = offset;
+        }
+        __syncthreads();
+
+        const uint32_t lane_mask = lane == 0 ? 0 : (1u << lane) - 1;
+        int warp_item_offset = 0;
+#pragma unroll
+        for (int item = 0; item < values_per_lane; ++item) {
+            const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane;
+            const int dst = row_count + warp_offsets[warp] + warp_item_offset + __popc(selected_warp[item] & lane_mask);
+            if ((selected_warp[item] & (uint32_t(1) << lane)) && dst < n_kv_max) {
+                indices[dst] = i;
+            }
+            warp_item_offset += __popc(selected_warp[item]);
+        }
+        __syncthreads();
+
+        if (tid == 0) {
+            row_count += chunk_count;
+        }
+        __syncthreads();
+    }
+
+    const int count = row_count;
+    for (int i = count + tid; i < n_kv_max; i += blockDim.x) {
+        indices[i] = -1;
+    }
+    __syncthreads();
+
+    // the dependent grid reads indices, signal once the row is complete
+    ggml_cuda_pdl_lc();
+}
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+
+void ggml_cuda_flash_attn_ext_compact_mask(
+        const ggml_tensor * mask, int32_t * indices, int32_t n_kv_max, cudaStream_t stream) {
+#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+    GGML_UNUSED_VARS(mask, indices, n_kv_max, stream);
+    GGML_ABORT("sparse flash attention is only supported on NVIDIA CUDA");
+#else
+    const int64_t s31 = mask->nb[1] / sizeof(half);
+    const int64_t s33 = mask->nb[3] / sizeof(half);
+    const dim3 blocks_num(mask->ne[1], mask->ne[3], 1);
+    const dim3 block_dim(256, 1, 1);
+    const ggml_cuda_kernel_launch_params launch_params(blocks_num, block_dim, 0, stream);
+    ggml_cuda_kernel_launch(flash_attn_mask_to_sparse_indices, launch_params,
+        (const half *) mask->data, indices, int(mask->ne[0]), n_kv_max, s31, s33);
+    CUDA_CHECK(cudaGetLastError());
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+}
+
+bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+    GGML_UNUSED_VARS(ctx, dst);
+    return false;
+#else
+    const ggml_tensor * Q    = dst->src[0];
+    const ggml_tensor * K    = dst->src[1];
+    const ggml_tensor * mask = dst->src[3];
+    const int cc = ggml_cuda_info().devices[ctx.device].cc;
+
+    float max_bias = 0.0f;
+    float logit_softcap = 0.0f;
+    memcpy(&max_bias,      (const float *) dst->op_params + 1, sizeof(float));
+    memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float));
+
+    const int32_t n_kv_max = ggml_get_op_params_i32(dst, 4);
+    return GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) &&
+        mask != nullptr && n_kv_max > 0 && max_bias == 0.0f && logit_softcap == 0.0f &&
+        mask->ne[0] == K->ne[1] && mask->ne[1] >= Q->ne[1] && mask->ne[2] == 1 &&
+        K->ne[1] >= std::max<int64_t>(4096, 2LL*n_kv_max);
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+}
+
 template <int DKQ, int DV, int ncols2>
 static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
    const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
    const ggml_tensor * Q = dst->src[0];

+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+    if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, 1, ncols2)) {
+        if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(ctx, dst)) {
+            ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 1, ncols2>(ctx, dst);
+            return;
+        }
+    }
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+
    if constexpr (ncols2 <= 8) {
        if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) {
            ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 8/ncols2, ncols2>(ctx, dst);
@@ -32,6 +32,7 @@
 #include "ggml-cuda/mmq.cuh"
 #include "ggml-cuda/mmvf.cuh"
 #include "ggml-cuda/mmvq.cuh"
+#include "ggml-cuda/moe-weighted-reduction.cuh"
 #include "ggml-cuda/norm.cuh"
 #include "ggml-cuda/opt-step-adamw.cuh"
 #include "ggml-cuda/opt-step-sgd.cuh"
@@ -3026,6 +3027,150 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph,
    return is_ok;
 }

+// The long form spans 2*k + 1 nodes. ggml_can_fuse_subgraph() accepts at most
+// 31 nodes, so k <= 15; larger values use the per-operation path.
+static constexpr int MOE_WEIGHTED_REDUCTION_MAX_EXPERTS = 15;
+
+struct ggml_cuda_moe_weighted_reduction_match {
+    const ggml_tensor * experts      = nullptr;
+    const ggml_tensor * expert_scale = nullptr;
+    const ggml_tensor * weights      = nullptr;
+    ggml_tensor *       dst          = nullptr;
+    int                 node_count   = 0;
+};
+
+static bool ggml_cuda_match_moe_weighted_reduction(
+        const ggml_cgraph * cgraph,
+        int node_idx,
+        ggml_cuda_moe_weighted_reduction_match & match) {
+    const ggml_tensor * first = cgraph->nodes[node_idx];
+    if (first->op != GGML_OP_MUL || first->type != GGML_TYPE_F32 || !ggml_is_contiguous(first)) {
+        return false;
+    }
+
+    auto split_mul = [](const ggml_tensor * mul, const ggml_tensor *& full, const ggml_tensor *& broadcast) {
+        auto is_weights = [mul](const ggml_tensor * tensor) {
+            return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && tensor->ne[0] == 1 &&
+                tensor->ne[1] == mul->ne[1] && tensor->ne[2] == mul->ne[2] && tensor->ne[3] == mul->ne[3];
+        };
+        auto is_experts = [mul](const ggml_tensor * tensor) {
+            return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) &&
+                ggml_are_same_shape(tensor, mul);
+        };
+
+        if (is_experts(mul->src[0]) && is_weights(mul->src[1])) {
+            full      = mul->src[0];
+            broadcast = mul->src[1];
+            return true;
+        }
+        if (is_experts(mul->src[1]) && is_weights(mul->src[0])) {
+            full      = mul->src[1];
+            broadcast = mul->src[0];
+            return true;
+        }
+        return false;
+    };
+
+    const ggml_tensor * weighted     = first;
+    const ggml_tensor * experts      = nullptr;
+    const ggml_tensor * expert_scale = nullptr;
+    const ggml_tensor * weights      = nullptr;
+    int                 mul_count    = 1;
+
+    // Match both structural forms:
+    //   (experts * expert_scale) * router_weight
+    //   experts * router_weight
+    // The matcher does not depend on the model or quantization type.
+    if (node_idx + 1 < cgraph->n_nodes) {
+        const ggml_tensor * second = cgraph->nodes[node_idx + 1];
+        const ggml_tensor * scaled = nullptr;
+        const ggml_tensor * route  = nullptr;
+        const ggml_tensor * raw    = nullptr;
+        const ggml_tensor * scale  = nullptr;
+        if (second->op == GGML_OP_MUL && second->type == GGML_TYPE_F32 && ggml_is_contiguous(second) &&
+                split_mul(second, scaled, route) && scaled == first && split_mul(first, raw, scale)) {
+            weighted     = second;
+            experts      = raw;
+            expert_scale = scale;
+            weights      = route;
+            mul_count    = 2;
+        }
+    }
+
+    if (experts == nullptr && !split_mul(first, experts, weights)) {
+        return false;
+    }
+
+    const int     n_expert_used = (int) weighted->ne[1];
+    const int64_t n_tokens      = weighted->ne[2] * weighted->ne[3];
+    if (n_expert_used < 2 || n_expert_used > MOE_WEIGHTED_REDUCTION_MAX_EXPERTS || n_tokens <= 0) {
+        return false;
+    }
+
+    const int node_count = 2 * n_expert_used + mul_count - 1;
+    if (node_idx + node_count > cgraph->n_nodes) {
+        return false;
+    }
+
+    std::vector<ggml_op> ops(node_count, GGML_OP_VIEW);
+    ops[0] = GGML_OP_MUL;
+    if (mul_count == 2) {
+        ops[1] = GGML_OP_MUL;
+    }
+    std::vector<const ggml_tensor *> views;
+    views.reserve(n_expert_used);
+    const ggml_tensor * previous = nullptr;
+    int n_adds = 0;
+    for (int offset = mul_count; offset < node_count; ++offset) {
+        const ggml_tensor * candidate = cgraph->nodes[node_idx + offset];
+        ops[offset] = candidate->op;
+
+        if (candidate->op == GGML_OP_VIEW) {
+            const int expert = (int) views.size();
+            if (expert >= n_expert_used || candidate->src[0] != weighted || candidate->view_src != weighted ||
+                    candidate->type != GGML_TYPE_F32 || candidate->ne[0] != weighted->ne[0] ||
+                    candidate->ne[1] != n_tokens || candidate->ne[2] != 1 || candidate->ne[3] != 1 ||
+                    candidate->nb[0] != weighted->nb[0] || candidate->nb[1] != weighted->nb[2] ||
+                    candidate->view_offs != (size_t) expert * weighted->nb[1]) {
+                return false;
+            }
+            views.push_back(candidate);
+            continue;
+        }
+
+        if (candidate->op != GGML_OP_ADD || views.size() < 2 || n_adds + 1 >= (int) views.size()) {
+            return false;
+        }
+        const ggml_tensor * lhs = n_adds == 0 ? views[0] : previous;
+        const ggml_tensor * rhs = views[n_adds + 1];
+        if (candidate->src[0] != lhs || candidate->src[1] != rhs || candidate->type != GGML_TYPE_F32) {
+            return false;
+        }
+        previous = candidate;
+        ++n_adds;
+    }
+
+    if ((int) views.size() != n_expert_used || n_adds != n_expert_used - 1 || previous == nullptr) {
+        return false;
+    }
+    if (!ggml_is_contiguous(previous) || previous->ne[0] != weighted->ne[0] ||
+            previous->ne[1] != n_tokens || previous->ne[2] != 1 || previous->ne[3] != 1) {
+        return false;
+    }
+
+    const int output_idx = node_idx + node_count - 1;
+    if (!ggml_can_fuse_subgraph(cgraph, node_idx, node_count, ops.data(), &output_idx, 1)) {
+        return false;
+    }
+
+    match.experts      = experts;
+    match.expert_scale = expert_scale;
+    match.weights      = weights;
+    match.dst          = cgraph->nodes[output_idx];
+    match.node_count   = node_count;
+    return true;
+}
+

 static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
                               int                                       node_idx,
@@ -3288,6 +3433,18 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph

    ggml_tensor * node = cgraph->nodes[i];

+    if (node->op == GGML_OP_MUL) {
+        ggml_cuda_moe_weighted_reduction_match match;
+        if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) {
+            const int output_idx = i + match.node_count - 1;
+            if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, match.node_count, &output_idx, 1)) {
+                ggml_cuda_op_moe_weighted_reduction(
+                    *cuda_ctx, match.experts, match.expert_scale, match.weights, match.dst);
+                return match.node_count - 1;
+            }
+        }
+    }
+
    // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache
    if (node->op == GGML_OP_GATED_DELTA_NET) {
        ggml_cuda_gated_delta_net_fused_cache fused_state_cpy;
@@ -4340,10 +4497,30 @@ static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_ev
 }

 static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) {
-    GGML_UNUSED(params);
-
    ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;

+    static const bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION"));
+    if (!disable_fusion) {
+        for (int i = 0; i < cgraph->n_nodes; ++i) {
+            if (cgraph->nodes[i]->op != GGML_OP_MUL) {
+                continue;
+            }
+
+            ggml_cuda_moe_weighted_reduction_match match;
+            if (!ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) {
+                continue;
+            }
+
+            params->add_alloc_dep(params->user_data, const_cast<ggml_tensor *>(match.experts), match.dst);
+            params->add_alloc_dep(params->user_data, const_cast<ggml_tensor *>(match.weights), match.dst);
+            if (match.expert_scale != nullptr) {
+                params->add_alloc_dep(
+                    params->user_data, const_cast<ggml_tensor *>(match.expert_scale), match.dst);
+            }
+            i += match.node_count - 1;
+        }
+    }
+
 #ifdef USE_CUDA_GRAPH
    const void * graph_key = ggml_cuda_graph_get_key(cgraph);
    const bool use_cuda_graph = ggml_cuda_graph_set_enabled(cuda_ctx, graph_key);
@@ -4365,10 +4542,12 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph
    ggml_cuda_stream_context & stream_context = cuda_ctx->stream_context();
    stream_context.reset();

-    if (!use_cuda_graph || ggml_backend_cuda_get_device_count() != 1) {
+    if (!use_cuda_graph) {
        return;
    }

+    ggml_cuda_set_device(cuda_ctx->device);
+
    // number of out-degrees for a particular node
    std::unordered_map<const ggml_tensor *, int> fan_out;
    // reverse mapping of node to index in the cgraph
@@ -148,7 +148,6 @@ static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma(
    typedef tile<16,  8, int, input_layout>        tile_B;
    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;

-    constexpr int I             = ggml_cuda_mmq_get_I(type, J, fallback);
    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -204,7 +203,6 @@ static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma(
    typedef tile< 8, 8, int> tile_B;
    typedef tile<16, 8, int> tile_C;

-    constexpr int I             = ggml_cuda_mmq_get_I(type, J, fallback);
    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -320,7 +318,6 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    typedef tile<16,  8, int, input_layout>        tile_B;
    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;

-    constexpr int I             = ggml_cuda_mmq_get_I(type, J, fallback);
    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -371,7 +368,6 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    typedef tile< 8,  8, int> tile_B;
    typedef tile<16,  8, int> tile_C;

-    constexpr int I             = ggml_cuda_mmq_get_I(type, J, fallback);
    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -486,7 +482,6 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    typedef tile<16,  4, int, input_layout>        tile_B;
    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;

-    constexpr int I             = ggml_cuda_mmq_get_I(type, J, fallback);
    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -537,7 +532,6 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    typedef tile< 8, 4, int> tile_B;
    typedef tile<16, 8, int> tile_C;

-    constexpr int I             = ggml_cuda_mmq_get_I(type, J, fallback);
    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -686,7 +680,6 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    typedef tile<16,  4, int, input_layout>        tile_B;
    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;

-    constexpr int I             = ggml_cuda_mmq_get_I(type, J, fallback);
    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -756,7 +749,6 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    typedef tile< 8, 4, int> tile_B;
    typedef tile<16, 8, int> tile_C;

-    constexpr int I             = ggml_cuda_mmq_get_I(type, J, fallback);
    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -1023,7 +1015,6 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    typedef tile<16,  4, int, input_layout>        tile_B;
    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;

-    constexpr int I             = ggml_cuda_mmq_get_I(type, J, fallback);
    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -1075,7 +1066,6 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    typedef tile< 8, 4, int> tile_B;
    typedef tile<16, 8, int> tile_C;

-    constexpr int I             = ggml_cuda_mmq_get_I(type, J, fallback);
    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -1190,7 +1180,6 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    typedef tile<8,  8, int>   tile_B;
    typedef tile<16, 8, float> tile_C;

-    constexpr int I             = ggml_cuda_mmq_get_I(type, J, fallback);
    constexpr int sram_stride   = ggml_cuda_mmq_get_sram_stride(type, J, fallback);
    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
    constexpr int ntx           = rows_per_warp / tile_C::I;
@@ -481,9 +481,6 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
    typedef tile<16,  8, int> tile_C;
 #endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)

-    constexpr int warp_size     = ggml_cuda_get_physical_warp_size();
-    constexpr int nwarps        = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
-    constexpr int I             = ggml_cuda_mmq_get_I(type, J, fallback);
    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback);
    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.

@@ -540,8 +537,6 @@ struct ggml_cuda_mmq_util_funcs {

 template <ggml_type type, int J, bool fallback>
 static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_funcs() {
-    constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
-
    if (!ggml_cuda_mmq_get_config(type, J, fallback).use_mma_data_layout()) {
        switch (type) {
            case GGML_TYPE_Q1_0:
@@ -326,6 +326,18 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) {
                return ne11 <= MMVQ_MAX_BATCH_SIZE;
        }
    }
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_ORIN) {
+        switch (type) { // tuned for Jetson Orin
+            case GGML_TYPE_Q2_K:
+            case GGML_TYPE_Q3_K:
+            case GGML_TYPE_Q4_K:
+            case GGML_TYPE_Q5_K:
+            case GGML_TYPE_Q6_K:
+                return ne11 <= 1;
+            default:
+                return ne11 <= MMVQ_MAX_BATCH_SIZE;
+        }
+    }
    if (GGML_CUDA_CC_IS_CDNA(cc)) {
        if (GGML_CUDA_CC_IS_CDNA1(cc)) {
            switch (type) {
@@ -0,0 +1,65 @@
+#include "moe-weighted-reduction.cuh"
+
+static __global__ void moe_weighted_reduction_f32(const float * __restrict__ experts,
+                                                  const float * __restrict__ expert_scale,
+                                                  const float * __restrict__ weights,
+                                                  float * __restrict__ dst,
+                                                  const int64_t n_embd,
+                                                  const int     n_expert_used) {
+    const int64_t token = blockIdx.x;
+    const int64_t col   = (int64_t) blockIdx.y * blockDim.x + threadIdx.x;
+    if (col >= n_embd) {
+        return;
+    }
+
+    const uint64_t first_row   = (uint64_t) token * n_expert_used;
+    const float    first_scale = expert_scale != nullptr ? expert_scale[first_row] : 1.0f;
+    float          sum         = (experts[first_row * n_embd + col] * first_scale) * weights[first_row];
+
+    for (int expert = 1; expert < n_expert_used; ++expert) {
+        const uint64_t row   = first_row + expert;
+        const float   scale = expert_scale != nullptr ? expert_scale[row] : 1.0f;
+        sum += (experts[row * n_embd + col] * scale) * weights[row];
+    }
+    dst[token * n_embd + col] = sum;
+}
+
+static void launch_moe_weighted_reduction(const float * experts,
+                                          const float * expert_scale,
+                                          const float * weights,
+                                          float *       dst,
+                                          int64_t       n_embd,
+                                          int64_t       n_tokens,
+                                          int           n_expert_used,
+                                          cudaStream_t  stream) {
+    constexpr int threads = 256;
+    const dim3 blocks(n_tokens, (n_embd + threads - 1) / threads, 1);
+    moe_weighted_reduction_f32
+        <<<blocks, threads, 0, stream>>>(experts, expert_scale, weights, dst, n_embd, n_expert_used);
+}
+
+void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx,
+                                         const ggml_tensor *         experts,
+                                         const ggml_tensor *         expert_scale,
+                                         const ggml_tensor *         weights,
+                                         ggml_tensor *               dst) {
+    GGML_ASSERT(experts->type == GGML_TYPE_F32);
+    GGML_ASSERT(weights->type == GGML_TYPE_F32);
+    GGML_ASSERT(expert_scale == nullptr || expert_scale->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_contiguous(experts));
+    GGML_ASSERT(ggml_is_contiguous(weights));
+    GGML_ASSERT(expert_scale == nullptr || ggml_is_contiguous(expert_scale));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+
+    const int64_t n_embd        = experts->ne[0];
+    const int64_t n_expert_used = experts->ne[1];
+    const int64_t n_tokens      = experts->ne[2] * experts->ne[3];
+    cudaStream_t  stream        = ctx.stream();
+
+    launch_moe_weighted_reduction((const float *) experts->data,
+                                  expert_scale ? (const float *) expert_scale->data : nullptr,
+                                  (const float *) weights->data,
+                                  (float *) dst->data, n_embd, n_tokens, (int) n_expert_used, stream);
+    CUDA_CHECK(cudaGetLastError());
+}
@@ -0,0 +1,7 @@
+#include "common.cuh"
+
+void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx,
+                                         const ggml_tensor *         experts,
+                                         const ggml_tensor *         expert_scale,
+                                         const ggml_tensor *         weights,
+                                         ggml_tensor *               dst);
@@ -98,12 +98,26 @@ static int    opt_ar_select = 2; // 2 = fused ALLREDUCE+ADD (DMA, default), 1 =
 //     https://docs.qualcomm.com/doc/80-N2040-61/topic/hvx-pmu-events.html
 static u32vec opt_pmu_evt { 0x3, 0x111, 0x100, 0x105, 0x240, 0x256, 0x7D, 0x8C };

-static int opt_opbatch  = 1024; // max number of ops in a batch
-static int opt_opqueue  = 64;   // max number of pending batches
+static int opt_opbatch  = 1280; // max number of ops in a batch
+static int opt_opqueue  = 32;   // max number of pending batches
 static int opt_optrace  = 0;    // trace buffer size per thread (0 means default)
 static int opt_oppoll   = 0;    // polling for batch completions
 static int opt_opfusion = 1;    // enable/disable op fusion

+enum ggml_hexagon_fusion_flags {
+    GGML_HEXAGON_FUSE_ALLREDUCE_ADD = (1 << 1), // 2
+    GGML_HEXAGON_FUSE_RMS_NORM_MUL  = (1 << 2), // 4
+    GGML_HEXAGON_FUSE_MUL_MAT_ADD   = (1 << 3), // 8
+    GGML_HEXAGON_FUSE_MUL_MAT_NX    = (1 << 4), // 16
+    GGML_HEXAGON_FUSE_MUL_MAT_ID_NX = (1 << 5), // 32
+};
+
+static inline bool ggml_hexagon_is_fusion_enabled(int flag) {
+    if (opt_opfusion <= 0) return false;
+    if (opt_opfusion == 1) return true; // 1 enables all
+    return (opt_opfusion & flag) != 0;
+}
+
 static std::regex* opt_opfilter = NULL; // regex of ops to not claim

 #define HEX_VERBOSE(...) \
@@ -293,6 +307,15 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
    struct htp_mm_kernel_params * kparams
 );

+static void ggml_hexagon_precompute_fused_mmidnx_params(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    int32_t n_weights,
+    struct htp_mm_kernel_params * kparams
+);
+
 static bool ggml_hexagon_precompute_allreduce_params(
    const struct ggml_hexagon_session * sess,
    const struct ggml_tensor * dst,
@@ -304,8 +327,12 @@ static bool ggml_hexagon_precompute_allreduce_params(
 );

 static bool mm_is_hmx_eligible(const ggml_tensor * t);
+static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams);
+static bool is_supported_mul_mat_id_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams);
 static bool is_mergeable_mul_mat(const ggml_tensor * t);
 static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor * n2);
+static bool is_mergeable_mul_mat_id(const ggml_tensor * t);
+static bool is_mergeable_mul_mat_id_pair(const ggml_tensor * n1, const ggml_tensor * n2);

 // ** backend sessions

@@ -1832,6 +1859,42 @@ struct ggml_hexagon_opbatch {
        }
    }

+    void sort_buffers() {
+        if (n_bufs <= 1) return;
+
+        std::vector<int> order(n_bufs);
+        for (unsigned int i = 0; i < n_bufs; i++) { order[i] = (int) i; }
+
+        std::stable_sort(order.begin(), order.end(), [&](int a, int b) {
+            return h_bufs[a].size > h_bufs[b].size;
+        });
+
+        bool already_sorted = true;
+        for (unsigned int i = 0; i < n_bufs; i++) {
+            if (order[i] != (int) i) {
+                already_sorted = false;
+                break;
+            }
+        }
+        if (already_sorted) return;
+
+        std::vector<uint16_t> remap(n_bufs);
+        std::vector<htp_buf_desc> sorted_bufs(n_bufs);
+        for (unsigned int new_bi = 0; new_bi < n_bufs; new_bi++) {
+            int old_bi = order[new_bi];
+            remap[old_bi] = (uint16_t) new_bi;
+            sorted_bufs[new_bi] = h_bufs[old_bi];
+        }
+
+        for (unsigned int i = 0; i < n_bufs; i++) {
+            h_bufs[i] = sorted_bufs[i];
+        }
+
+        for (unsigned int i = 0; i < n_tens; i++) {
+            h_tens[i].bi = remap[h_tens[i].bi];
+        }
+    }
+
    bool try_fuse_allreduce_add(const htp_opnode & node) {
        if (n_ops == 0 || opt_ar_select != 2) return false;
        if (node.opcode != HTP_OP_ADD) return false;
@@ -2144,9 +2207,15 @@ struct ggml_hexagon_opbatch {
            if (x_in != x || w_in->type != w0->type || w_in->ne[0] != w0->ne[0]) {
                return false;
            }
+            if (!last_node.fused.empty() && (mm_is_hmx_eligible(last_node.fused[0]) != mm_is_hmx_eligible(node.node))) {
+                return false;
+            }

            struct htp_mm_kernel_params kparams;
            ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, curr_n + 1, &kparams);
+            if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) {
+                return false;
+            }
            if ((size_t) kparams.vtcm_size > sess->vtcm_size) {
                HEX_VERBOSE("ggml-hex: %s skip NX fusion: VTCM needed (%d) > budget (%zu)\n",
                            sess->c_name(), kparams.vtcm_size, sess->vtcm_size);
@@ -2210,6 +2279,9 @@ struct ggml_hexagon_opbatch {

            struct htp_mm_kernel_params kparams;
            ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, 2, &kparams);
+            if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) {
+                return false;
+            }
            if ((size_t) kparams.vtcm_size > sess->vtcm_size) {
                HEX_VERBOSE("ggml-hex: %s skip NX fusion: VTCM needed (%d) > budget (%zu)\n",
                            sess->c_name(), kparams.vtcm_size, sess->vtcm_size);
@@ -2272,18 +2344,172 @@ struct ggml_hexagon_opbatch {
        return false;
    }

-enum ggml_hexagon_fusion_flags {
-    GGML_HEXAGON_FUSE_ALLREDUCE_ADD = (1 << 1), // 2
-    GGML_HEXAGON_FUSE_RMS_NORM_MUL  = (1 << 2), // 4
-    GGML_HEXAGON_FUSE_MUL_MAT_ADD   = (1 << 3), // 8
-    GGML_HEXAGON_FUSE_MUL_MAT_NX    = (1 << 4), // 16
-};
+    bool try_fuse_mul_mat_id_nx(const htp_opnode & node) {
+        if (n_ops == 0 || node.opcode != HTP_OP_MUL_MAT_ID) return false;
+        if (!is_mergeable_mul_mat_id(node.node)) return false;

-static inline bool ggml_hexagon_is_fusion_enabled(int flag) {
-    if (opt_opfusion <= 0) return false;
-    if (opt_opfusion == 1) return true; // 1 enables all
-    return (opt_opfusion & flag) != 0;
-}
+        const ggml_tensor * w_in   = node.src0();
+        const ggml_tensor * x_in   = node.src1();
+        const ggml_tensor * ids_in = node.node->src[2];
+        const ggml_tensor * d_in   = node.dst();
+        if (!w_in || !x_in || !ids_in || !d_in) return false;
+
+        htp_opnode & last_node = ops[n_ops - 1];
+
+        // Case 1: last_node is already MUL_MAT_ID_NX
+        if (last_node.opcode == HTP_OP_MUL_MAT_ID_NX) {
+            const uint32_t curr_n = (uint32_t) last_node.outputs.size();
+            if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 2 >= HTP_OP_MAX_INPUTS) {
+                return false;
+            }
+
+            const ggml_tensor * w0  = last_node.inputs[0];
+            const ggml_tensor * x   = last_node.inputs[curr_n];
+            const ggml_tensor * ids = last_node.inputs[curr_n + 1];
+
+            if (x_in != x || ids_in != ids || w_in->type != w0->type || w_in->ne[0] != w0->ne[0] || w_in->ne[2] != w0->ne[2]) {
+                return false;
+            }
+            if (!last_node.fused.empty() && (mm_is_hmx_eligible(last_node.fused[0]) != mm_is_hmx_eligible(node.node))) {
+                return false;
+            }
+
+            struct htp_mm_kernel_params kparams;
+            ggml_hexagon_precompute_fused_mmidnx_params(sess, w0, x, d_in, curr_n + 1, &kparams);
+            if (!is_supported_mul_mat_id_nx_kernel(w0, &kparams)) {
+                return false;
+            }
+            if ((size_t) kparams.vtcm_size > sess->vtcm_size) {
+                HEX_VERBOSE("ggml-hex: %s skip ID NX fusion: VTCM needed (%d) > budget (%zu)\n",
+                            sess->c_name(), kparams.vtcm_size, sess->vtcm_size);
+                return false;
+            }
+
+            size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0;
+            auto fit_t = [&](const ggml_tensor * t) {
+                if (!t) return;
+                if (!t_map.count(t)) {
+                    extra_tens++;
+                    auto sbuf = static_cast<ggml_hexagon_shared_buffer *>(t->buffer->context);
+                    if (!b_map.count(sbuf->fd())) {
+                        extra_vmem += sbuf->size();
+                        extra_bufs += 1;
+                    }
+                }
+            };
+            fit_t(w_in);
+            fit_t(d_in);
+            if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) {
+                return false;
+            }
+
+            last_node.inputs[curr_n] = w_in;
+            last_node.inputs[curr_n + 1] = x;
+            last_node.inputs.push_back(ids);
+            last_node.outputs.push_back(d_in);
+            last_node.fused.push_back(node.node);
+            memcpy(last_node.kernel_params, &kparams, sizeof(kparams));
+
+            htp_op_desc & o = h_ops[n_ops - 1];
+            memcpy(o.kernel_params, &kparams, sizeof(kparams));
+
+            for (uint32_t s = 0; s <= curr_n + 2; s++) {
+                o.src[s] = add_tensor(last_node.inputs[s]);
+            }
+            for (uint32_t s = curr_n + 3; s < HTP_OP_MAX_INPUTS; s++) {
+                o.src[s] = 0xffff;
+            }
+            for (uint32_t d = 0; d <= curr_n; d++) {
+                o.dst[d] = add_tensor(last_node.outputs[d]);
+            }
+            for (uint32_t d = curr_n + 1; d < HTP_OP_MAX_OUTPUTS; d++) {
+                o.dst[d] = 0xffff;
+            }
+
+            HEX_VERBOSE("ggml-hex: %s fused MUL_MAT_ID_NX (N=%u, #%u)\n", sess->c_name(), curr_n + 1, n_ops - 1);
+            return true;
+        }
+
+        // Case 2: last_node is single MUL_MAT_ID
+        if (last_node.opcode == HTP_OP_MUL_MAT_ID) {
+            if (!is_mergeable_mul_mat_id_pair(last_node.node, node.node)) {
+                return false;
+            }
+
+            const ggml_tensor * w0  = last_node.src0();
+            const ggml_tensor * x   = last_node.src1();
+            const ggml_tensor * ids = last_node.node->src[2];
+            const ggml_tensor * w1  = node.src0();
+            if (!w0 || !x || !ids || !w1) return false;
+
+            struct htp_mm_kernel_params kparams;
+            ggml_hexagon_precompute_fused_mmidnx_params(sess, w0, x, node.dst(), 2, &kparams);
+            if (!is_supported_mul_mat_id_nx_kernel(w0, &kparams)) {
+                return false;
+            }
+            if ((size_t) kparams.vtcm_size > sess->vtcm_size) {
+                HEX_VERBOSE("ggml-hex: %s skip ID NX fusion: VTCM needed (%d) > budget (%zu)\n",
+                            sess->c_name(), kparams.vtcm_size, sess->vtcm_size);
+                return false;
+            }
+
+            size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0;
+            auto fit_t = [&](const ggml_tensor * t) {
+                if (!t) return;
+                if (!t_map.count(t)) {
+                    extra_tens++;
+                    auto sbuf = static_cast<ggml_hexagon_shared_buffer *>(t->buffer->context);
+                    if (!b_map.count(sbuf->fd())) {
+                        extra_vmem += sbuf->size();
+                        extra_bufs += 1;
+                    }
+                }
+            };
+            fit_t(w1);
+            fit_t(node.dst());
+            if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) {
+                return false;
+            }
+
+            const ggml_tensor * dst_0 = last_node.dst();
+            const ggml_tensor * dst_1 = node.dst();
+
+            last_node.opcode = HTP_OP_MUL_MAT_ID_NX;
+            last_node.name   = "MUL_MAT_ID_NX";
+            last_node.inputs.clear();
+            last_node.inputs.push_back(w0);
+            last_node.inputs.push_back(w1);
+            last_node.inputs.push_back(x);
+            last_node.inputs.push_back(ids);
+            last_node.outputs.clear();
+            last_node.outputs.push_back(dst_0);
+            last_node.outputs.push_back(dst_1);
+            last_node.fused.push_back(node.node);
+            memcpy(last_node.kernel_params, &kparams, sizeof(kparams));
+
+            htp_op_desc & o = h_ops[n_ops - 1];
+            o.opcode = HTP_OP_MUL_MAT_ID_NX;
+            memcpy(o.kernel_params, &kparams, sizeof(kparams));
+
+            o.src[0] = add_tensor(w0);
+            o.src[1] = add_tensor(w1);
+            o.src[2] = add_tensor(x);
+            o.src[3] = add_tensor(ids);
+            for (uint32_t s = 4; s < HTP_OP_MAX_INPUTS; s++) {
+                o.src[s] = 0xffff;
+            }
+            o.dst[0] = add_tensor(dst_0);
+            o.dst[1] = add_tensor(dst_1);
+            for (uint32_t d = 2; d < HTP_OP_MAX_OUTPUTS; d++) {
+                o.dst[d] = 0xffff;
+            }
+
+            HEX_VERBOSE("ggml-hex: %s fused MUL_MAT_ID_NX (N=2, #%u)\n", sess->c_name(), n_ops - 1);
+            return true;
+        }
+
+        return false;
+    }

    bool try_fuse(const htp_opnode & node) {
        if (!opt_opfusion) return false;
@@ -2291,6 +2517,7 @@ static inline bool ggml_hexagon_is_fusion_enabled(int flag) {
        if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_RMS_NORM_MUL)  && try_fuse_rms_norm_mul(node))  return true;
        if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_ADD)   && try_fuse_mul_mat_add(node))   return true;
        if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_NX)    && try_fuse_mul_mat_nx(node))    return true;
+        if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_ID_NX) && try_fuse_mul_mat_id_nx(node)) return true;
        return false;
    }
 };
@@ -2350,6 +2577,8 @@ struct ggml_hexagon_opqueue {
        delete shm_buf;
    }

+    size_t shm_size() const { return shm_buf ? shm_buf->size() : 0; }
+
    // push new batch
    bool push(htp_opbatch_req& req, dspqueue_buffer& dbuf, ggml_hexagon_opbatch* op_batch) {
        static_assert(sizeof(htp_opbatch_req) % 8 == 0, "sizeof(htp_opbatch_req) must be multiple of 8");
@@ -2396,6 +2625,8 @@ struct ggml_hexagon_opqueue {
        uint8_t * t_ptr = m_ptr; m_ptr += t_size;
        uint8_t * o_ptr = m_ptr;

+        op_batch->sort_buffers();
+
        memcpy(b_ptr, (void *) op_batch->h_bufs.data(), b_size);
        memcpy(t_ptr, (void *) op_batch->h_tens.data(), t_size);
        memcpy(o_ptr, (void *) op_batch->h_ops.data(),  o_size);
@@ -3018,7 +3249,8 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
        opt_vmem = ggml_hexagon_measure_max_vmem(this);
        GGML_LOG_INFO("ggml-hex: %s measured max vmem %zu\n", this->c_name(), opt_vmem);
    }
-    this->max_vmem = opt_vmem;
+    const size_t shm_size = this->op_queue->shm_size();
+    this->max_vmem = (opt_vmem > shm_size) ? (opt_vmem - shm_size) : opt_vmem;

    this->op_batch = new ggml_hexagon_opbatch(this, opt_opbatch, this->max_vmem);

@@ -3378,6 +3610,10 @@ static bool ggml_hexagon_matmul_is_hmx_eligible(
    bool is_matmul_id,
    bool is_batched
 ) {
+    if (src1->type != GGML_TYPE_F32) {
+        return false;
+    }
+
    const int ne00  = src0->ne[0];
    const int ne11  = src1->ne[1];
    const int ne12  = src1->ne[2];
@@ -3408,7 +3644,8 @@ static bool ggml_hexagon_matmul_is_hmx_eligible(
        return false;
    }

-    // M alignment: Use HMX when M > HTP_MM_HMX_MIN_NROWS
+    // M alignment: Use HMX when M > HTP_MM_HMX_MIN_NROWS.
+    // For MUL_MAT_ID, src1 shape is [K, n_expert_used, n_tokens, 1], so n_tokens is ne12.
    const int m = is_matmul_id ? ne12 : ne11;
    if (m <= HTP_MM_HMX_MIN_NROWS) {
        return false;
@@ -3460,7 +3697,7 @@ static bool ggml_hexagon_precompute_hmx_mm_params(

    if (!use_grouped) {
        // Fallback to simple 2D path (group_size = 1)
-        const int m_id_rows = (int) ((size_t) dst->ne[1] * dst->ne[2]);
+        const int m_id_rows = (dst && is_matmul_id) ? (int) ((size_t) dst->ne[1] * dst->ne[2]) : 0;
        if (!htp_mm_hmx_solve_2d_params(wtype, ne00_padded, m_id_rows, ne01_padded, ne11_padded, ne11, n_threads, pipeline, is_matmul_id, aligned_tile_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
            return false;
        }
@@ -3768,8 +4005,10 @@ static void ggml_hexagon_precompute_unary_params(

    kparams->n_threads = n_threads;

-    const size_t src0_data_row_size = src0->ne[0] * sizeof(float);
-    const size_t dst_data_row_size  = dst->ne[0]  * sizeof(float);
+    const size_t elem_size = ggml_type_size(src0->type);
+
+    const size_t src0_data_row_size = src0->ne[0] * elem_size;
+    const size_t dst_data_row_size  = dst->ne[0]  * ggml_type_size(dst->type);

    const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, 128);
    const size_t dst_row_size_aligned  = hex_round_up(dst_data_row_size,  128);
@@ -3783,7 +4022,7 @@ static void ggml_hexagon_precompute_unary_params(

    if (op == HTP_OP_RMS_NORM_MUL) {
        GGML_ASSERT(src1 != nullptr);
-        src1_data_row_size = src1->ne[0] * sizeof(float);
+        src1_data_row_size = src1->ne[0] * ggml_type_size(src1->type);
        src1_row_size_aligned = hex_round_up(src1_data_row_size, 128);
        broadcast_weight = (src1->ne[1] * src1->ne[2] * src1->ne[3] == 1);
    }
@@ -3797,7 +4036,7 @@ static void ggml_hexagon_precompute_unary_params(

    htp_unary_vtcm_layout_build(&L, op, src0->ne[0], dst->ne[0],
                                op == HTP_OP_RMS_NORM_MUL ? src1->ne[0] : 0,
-                                broadcast_weight, n_threads, sess->vtcm_size,
+                                broadcast_weight, n_threads, sess->vtcm_size, elem_size,
                                &col_tile, &vtcm_row_per_thread);

    kparams->col_tile = col_tile;
@@ -3918,64 +4157,113 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
 ) {
    memset(kparams, 0, sizeof(*kparams));

-    const int wtype = src0->type;
-    const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype);
+    const int ne00 = src0->ne[0];
+    const int ne01 = src0->ne[1];
+    const int ne02 = src0->ne[2];
+    const int ne03 = src0->ne[3];

    const int ne10 = src1->ne[0];
-    const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3];
-    const size_t src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
-    const size_t src0_row_size = src0->nb[1];
+    const int ne11 = src1->ne[1];
+    const int ne12 = src1->ne[2];
+    const int ne13 = src1->ne[3];

-    uint32_t best_n_prefetch = 16;
+    const int wtype = src0->type;
+    const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype);
+    const int ne00_padded = is_repack ? hex_round_up(ne00, 32) : ne00;
+    const int ne01_padded = is_repack ? hex_round_up(ne01, 32) : ne01;
+    const int ne11_padded = hex_round_up(ne11, 32);

-    if (is_repack) {
-        const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
-        best_n_prefetch = 2;
-        for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
-            struct htp_mm_hvx_vtcm_layout L;
-            htp_mm_hvx_vtcm_layout_build(
-                &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
-                0, src0_row_size, src1_row_size, 0, d, false, true
-            );
-            if (L.total_bytes <= sess->vtcm_size) {
-                best_n_prefetch = d;
-                break;
-            }
+    const size_t vtcm_budget = sess->vtcm_size;
+    const bool is_batched = (ne02 * ne03 > 1 || ne12 * ne13 > 1);
+
+    bool hmx_enabled = (sess->n_hmx > 0) && (opt_mm_select >= 3);
+    if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, src1, nullptr, ne01_padded, false, is_batched)) {
+        if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, nullptr, wtype, ne00_padded, ne01_padded, ne02, ne11, ne12, ne11_padded, false, is_batched, vtcm_budget, kparams)) {
+            kparams->n_weights = n_weights;
+            goto finalize;
        }
    }

-    struct htp_mm_hvx_vtcm_layout L;
-    bool try_tiled = (opt_mm_select >= 2);
-
-    // Test tiled first
-    htp_mm_hvx_vtcm_layout_build(
-        &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
-        0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true
-    );
-
-    if (try_tiled && L.total_bytes <= sess->vtcm_size) {
-        kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
-        kparams->vtcm_src0_size = L.src0_bytes;
-        kparams->vtcm_src1_size = L.src1_bytes;
-        kparams->vtcm_dst_size  = L.dst_bytes;
-        kparams->vtcm_size      = L.total_bytes;
-        kparams->n_prefetch     = best_n_prefetch;
-        kparams->n_weights      = n_weights;
-    } else {
-        kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT;
-        size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
-
-        htp_mm_hvx_vtcm_layout_build(
-            &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
-            0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true
-        );
-        kparams->vtcm_src0_size = L.src0_bytes;
-        kparams->vtcm_src1_size = L.src1_bytes;
-        kparams->vtcm_dst_size  = L.dst_bytes;
-        kparams->vtcm_size      = L.total_bytes;
-        kparams->n_prefetch     = best_n_prefetch;
-        kparams->n_weights      = n_weights;
+    if (!is_repack) {
+        kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED;
+        return;
    }
+
+    {
+        const int src1_nrows = ne11 * ne12 * ne13;
+        const size_t src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
+        const size_t src0_row_size = src0->nb[1];
+
+        uint32_t best_n_prefetch = 16;
+
+        if (is_repack) {
+            const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
+            best_n_prefetch = 2;
+            for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
+                struct htp_mm_hvx_vtcm_layout L;
+                htp_mm_hvx_vtcm_layout_build(
+                    &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
+                    0, src0_row_size, src1_row_size, 0, d, false, true
+                );
+                if (L.total_bytes <= sess->vtcm_size) {
+                    best_n_prefetch = d;
+                    break;
+                }
+            }
+        }
+
+        struct htp_mm_hvx_vtcm_layout L;
+        bool try_tiled = (opt_mm_select >= 2);
+
+        // Test tiled first
+        htp_mm_hvx_vtcm_layout_build(
+            &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
+            0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true
+        );
+
+        if (try_tiled && L.total_bytes <= sess->vtcm_size) {
+            kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
+            kparams->vtcm_src0_size = L.src0_bytes;
+            kparams->vtcm_src1_size = L.src1_bytes;
+            kparams->vtcm_dst_size  = L.dst_bytes;
+            kparams->vtcm_size      = L.total_bytes;
+            kparams->n_prefetch     = best_n_prefetch;
+            kparams->n_weights      = n_weights;
+        } else {
+            kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT;
+            size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
+
+            htp_mm_hvx_vtcm_layout_build(
+                &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
+                0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true
+            );
+            kparams->vtcm_src0_size = L.src0_bytes;
+            kparams->vtcm_src1_size = L.src1_bytes;
+            kparams->vtcm_dst_size  = L.dst_bytes;
+            kparams->vtcm_size      = L.total_bytes;
+            kparams->n_prefetch     = best_n_prefetch;
+            kparams->n_weights      = n_weights;
+        }
+    }
+
+finalize:
+    kparams->div_ne12_ne1 = init_fastdiv_values(ne12 * ne11);
+    kparams->div_ne1      = init_fastdiv_values(ne11);
+    kparams->div_r2       = init_fastdiv_values(ne02 > 0 ? ne12 / ne02 : 1);
+    kparams->div_r3       = init_fastdiv_values(ne03 > 0 ? ne13 / ne03 : 1);
+    kparams->div_ne11     = init_fastdiv_values(ne11);
+}
+
+static void ggml_hexagon_precompute_fused_mmidnx_params(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0, // W0
+    const struct ggml_tensor * src1, // x
+    const struct ggml_tensor * dst,  // dst0
+    int32_t n_weights,
+    struct htp_mm_kernel_params * kparams
+) {
+    ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, kparams);
+    kparams->n_weights = n_weights;
 }

 static bool ggml_hexagon_tensor_is_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) {
@@ -4010,11 +4298,6 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s
                return false;
            }

-            // hardcoded limit to refuse the lm-head for now
-            if (src0->ne[1] > 32768) {
-                return false;
-            }
-
            if (src1->ne[2] != 1 || src1->ne[3] != 1) {
                return false;  // no broadcasting (for now)
            }
@@ -4170,15 +4453,39 @@ static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * ses
    const struct ggml_tensor * src0 = op->src[0];
    const struct ggml_tensor * dst  = op;

-    if (src0->type != GGML_TYPE_F32) {
+    if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16) {
        return false;
    }
-    if (dst->type != GGML_TYPE_F32) {
+    if (dst->type != src0->type) {
        return false;
    }
    if (!ggml_is_contiguous_rows(src0)) {
        return false;
    }
+
+    // F16 device kernels only cover this explicit whitelist (must stay in sync with
+    // the is_f16 whitelist in execute_op_unary(), unary-ops.c).
+    if (src0->type == GGML_TYPE_F16) {
+        switch (op->op) {
+            case GGML_OP_NORM:
+            case GGML_OP_RMS_NORM:
+            case GGML_OP_L2_NORM:
+            case GGML_OP_SCALE:
+            case GGML_OP_CLAMP:
+            case GGML_OP_SQR:
+            case GGML_OP_SQRT:
+            case GGML_OP_LOG:
+                break;
+            case GGML_OP_UNARY:
+                if (ggml_get_unary_op(op) != GGML_UNARY_OP_ABS) {
+                    return false;
+                }
+                break;
+            default:
+                return false;
+        }
+    }
+
    if (!ggml_are_same_shape(src0, dst)) {
        return false;
    }
@@ -4348,6 +4655,13 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session *
    const struct ggml_tensor * src1 = op->src[1]; // indices
    const struct ggml_tensor * dst  = op;

+    if (src0->extra) {
+        const auto * extra = (const ggml_hexagon_tensor_extra *) src0->extra;
+        if (extra->flags & GGML_HEXAGON_TENSOR_REPACK) {
+            return false;
+        }
+    }
+
    if (src0->type != GGML_TYPE_F32 && src0->ne[0] < 32) {
        return false;
    }
@@ -4734,10 +5048,43 @@ static bool mm_is_hmx_eligible(const ggml_tensor * t) {
    return ggml_hexagon_matmul_is_hmx_eligible(src0, src1, t, ne01_padded, is_matmul_id, is_batched);
 }

+static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams) {
+    if (kparams->n_hmx) {
+        return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D;
+    }
+
+    if (!ggml_hexagon_is_repack_type(src0->type)) {
+        return false;
+    }
+
+    return kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT;
+}
+
+static bool is_supported_mul_mat_id_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams) {
+    if (kparams->n_hmx) {
+        return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D;
+    }
+
+    if (!ggml_hexagon_is_repack_type(src0->type)) {
+        return false;
+    }
+
+    return kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK;
+}
+
 static bool is_mergeable_mul_mat(const ggml_tensor * t) {
-    if (!t || t->op != GGML_OP_MUL_MAT)   return false;
-    if (t->src[1]->type != GGML_TYPE_F32) return false;
-    return ggml_is_quantized(t->src[0]->type) && !mm_is_hmx_eligible(t);
+    if (!t || t->op != GGML_OP_MUL_MAT) return false;
+
+    const ggml_tensor * src0 = t->src[0];
+    const ggml_tensor * src1 = t->src[1];
+    if (src1->type != GGML_TYPE_F32) return false;
+    if (src0->ne[2] != 1 || src0->ne[3] != 1) return false;
+
+    if (mm_is_hmx_eligible(t)) {
+        return ggml_hexagon_is_hmx_weight_type(src0->type);
+    }
+
+    return ggml_hexagon_is_repack_type(src0->type);
 }

 static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor * n2) {
@@ -4753,6 +5100,41 @@ static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor
    if (n1->src[0]->type != n2->src[0]->type) {
        return false;
    }
+    if (mm_is_hmx_eligible(n1) != mm_is_hmx_eligible(n2)) {
+        return false;
+    }
+    return true;
+}
+
+static bool is_mergeable_mul_mat_id(const ggml_tensor * t) {
+    if (!t || t->op != GGML_OP_MUL_MAT_ID) return false;
+
+    const ggml_tensor * src0 = t->src[0];
+    return ggml_hexagon_is_repack_type(src0->type);
+}
+
+static bool is_mergeable_mul_mat_id_pair(const ggml_tensor * n1, const ggml_tensor * n2) {
+    if (!is_mergeable_mul_mat_id(n1) || !is_mergeable_mul_mat_id(n2)) {
+        return false;
+    }
+    if (n1->src[1] != n2->src[1]) {
+        return false;
+    }
+    if (n1->src[2] != n2->src[2]) {
+        return false;
+    }
+    if (n1->src[0]->ne[0] != n2->src[0]->ne[0]) {
+        return false;
+    }
+    if (n1->src[0]->ne[2] != n2->src[0]->ne[2]) {
+        return false;
+    }
+    if (n1->src[0]->type != n2->src[0]->type) {
+        return false;
+    }
+    if (mm_is_hmx_eligible(n1) != mm_is_hmx_eligible(n2)) {
+        return false;
+    }
    return true;
 }

@@ -4776,8 +5158,8 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg

            if (graph->nodes[i]->op == GGML_OP_RMS_NORM && ggml_can_fuse(graph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) {
                extra->flags |= GGML_HEXAGON_TENSOR_FUSEABLE;
-            } else if (graph->nodes[i]->op == GGML_OP_MUL_MAT) {
-                if ((i + 1 < graph->n_nodes && graph->nodes[i + 1]->op == GGML_OP_ADD && ggml_can_fuse(graph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD })) ||
+            } else if (graph->nodes[i]->op == GGML_OP_MUL_MAT || graph->nodes[i]->op == GGML_OP_MUL_MAT_ID) {
+                if ((i + 1 < graph->n_nodes && graph->nodes[i + 1]->op == GGML_OP_ADD && ggml_can_fuse(graph, i, { graph->nodes[i]->op, GGML_OP_ADD })) ||
                    ggml_node_has_n_uses(graph, i, 1)) {
                    extra->flags |= GGML_HEXAGON_TENSOR_FUSEABLE;
                }
@@ -315,7 +315,8 @@ 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_NX || node.opcode == HTP_OP_MUL_MAT_ADD) {
+            node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ID_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;
@@ -323,6 +323,10 @@ int op_cpy(struct htp_ops_context * octx) {
        return HTP_STATUS_NO_SUPPORT;
    }

+    FARF(HIGH, "cpy-%s-%s: (%ux%ux%ux%u) -> (%ux%ux%ux%u) : use_dma=%d n_threads %u\n",
+         src0->type == HTP_TYPE_F32 ? "f32" : "f16", dst->type == HTP_TYPE_F32 ? "f32" : "f16",
+         ne00, ne01, ne02, ne03, ne0, ne1, ne2, ne3, use_dma, 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 {
@@ -1138,6 +1138,15 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
    gctx.vtcm_base = octx->ctx->vtcm_base;
    gctx.vtcm_per_thread = 2 * state_aligned;

+    FARF(HIGH, "gated-delta-net-f32: q(%ux%ux%ux%u) k(%ux%ux%ux%u) v(%ux%ux%ux%u) state(%ux%ux%ux%u) -> (%ux%ux%ux%u) : "
+         "vtcm-size %zu n_threads %u\n",
+         q->ne[0], q->ne[1], q->ne[2], q->ne[3],
+         k->ne[0], k->ne[1], k->ne[2], k->ne[3],
+         v->ne[0], v->ne[1], v->ne[2], v->ne[3],
+         state->ne[0], state->ne[1], state->ne[2], state->ne[3],
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
+         gctx.vtcm_per_thread * octx->n_threads, octx->n_threads);
+
    if (n_tokens == 1) {
        worker_pool_run_func(octx->ctx->worker_pool, gated_delta_net_f32_tg_thread, &gctx, octx->n_threads);
    } else {
@@ -247,6 +247,14 @@ int op_get_rows(struct htp_ops_context * octx) {
        }
    }

+    FARF(HIGH, "get-rows: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu dst-vtcm-size %zu use_dma=%d n_threads %d\n",
+         octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
+         octx->src[1]->ne[0], octx->src[1]->ne[1], octx->src[1]->ne[2], octx->src[1]->ne[3],
+         octx->dst->ne[0], octx->dst->ne[1], octx->dst->ne[2], octx->dst->ne[3],
+         grctx.vtcm_layout.src0_bytes_per_thread * kparams->n_threads,
+         grctx.vtcm_layout.dst_bytes_per_thread  * kparams->n_threads,
+         kparams->use_dma, kparams->n_threads);
+
    work_queue_run(octx->ctx->work_queue, q_func, &grctx, kparams->n_threads);
    return HTP_STATUS_OK;
 }
@@ -118,6 +118,7 @@ struct htp_context {
 int op_matmul(struct htp_ops_context * octx);
 int op_matmul_id(struct htp_ops_context * octx);
 int op_matmul_nx(struct htp_ops_context * octx);
+int op_matmul_id_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);
@@ -52,6 +52,7 @@ enum htp_op_code {
    HTP_OP_MUL_MAT,
    HTP_OP_MUL_MAT_ID,
    HTP_OP_MUL_MAT_NX,
+    HTP_OP_MUL_MAT_ID_NX,
    HTP_OP_MUL_MAT_ADD,
    HTP_OP_RMS_NORM,
    HTP_OP_RMS_NORM_MUL,
@@ -358,6 +358,54 @@ static inline void hvx_clamp_scalar_f32(uint8_t * restrict dst, const uint8_t *
    }
 }

+#define HVX_OP_CLAMP_SCALAR_F16(v)                                       \
+    ({                                                                   \
+        HVX_VectorPred pred_cap_right = Q6_Q_vcmp_gt_VhfVhf(v, max_vec); \
+        HVX_VectorPred pred_cap_left  = Q6_Q_vcmp_gt_VhfVhf(min_vec, v); \
+        HVX_Vector tmp = Q6_V_vmux_QVV(pred_cap_right, max_vec, v);      \
+        Q6_V_vmux_QVV(pred_cap_left, min_vec, tmp);                      \
+    })
+
+static inline void hvx_clamp_scalar_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, const _Float16 min, const _Float16 max, uint32_t n) {
+    const HVX_Vector min_vec = hvx_vec_splat_f16(min);
+    const HVX_Vector max_vec = hvx_vec_splat_f16(max);
+    assert((unsigned long) dst % 128 == 0);
+    assert((unsigned long) src % 128 == 0);
+    hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(_Float16), hvx_vec_store_a, HVX_OP_CLAMP_SCALAR_F16);
+}
+
+static inline void hvx_clamp_scalar_f16_au(uint8_t * restrict dst, const uint8_t * restrict src, const _Float16 min, const _Float16 max, uint32_t n) {
+    const HVX_Vector min_vec = hvx_vec_splat_f16(min);
+    const HVX_Vector max_vec = hvx_vec_splat_f16(max);
+    assert((unsigned long) dst % 128 == 0);
+    hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(_Float16), hvx_vec_store_a, HVX_OP_CLAMP_SCALAR_F16);
+}
+
+static inline void hvx_clamp_scalar_f16_ua(uint8_t * restrict dst, const uint8_t * restrict src, const _Float16 min, const _Float16 max, uint32_t n) {
+    const HVX_Vector min_vec = hvx_vec_splat_f16(min);
+    const HVX_Vector max_vec = hvx_vec_splat_f16(max);
+    assert((unsigned long) src % 128 == 0);
+    hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(_Float16), hvx_vec_store_u, HVX_OP_CLAMP_SCALAR_F16);
+}
+
+static inline void hvx_clamp_scalar_f16_uu(uint8_t * restrict dst, const uint8_t * restrict src, const _Float16 min, const _Float16 max, uint32_t n) {
+    const HVX_Vector min_vec = hvx_vec_splat_f16(min);
+    const HVX_Vector max_vec = hvx_vec_splat_f16(max);
+    hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(_Float16), hvx_vec_store_u, HVX_OP_CLAMP_SCALAR_F16);
+}
+
+static inline void hvx_clamp_scalar_f16(uint8_t * restrict dst, const uint8_t * restrict src, const _Float16 min, const _Float16 max, const int num_elems) {
+    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) {
+        hvx_clamp_scalar_f16_aa(dst, src, min, max, num_elems);
+    } else if (hex_is_aligned((void *) dst, 128)) {
+        hvx_clamp_scalar_f16_au(dst, src, min, max, num_elems);
+    } else if (hex_is_aligned((void *) src, 128)) {
+        hvx_clamp_scalar_f16_ua(dst, src, min, max, num_elems);
+    } else {
+        hvx_clamp_scalar_f16_uu(dst, src, min, max, num_elems);
+    }
+}
+
 //
 // Abs
 //
@@ -386,11 +434,69 @@ static inline void hvx_abs_f32_aa(uint8_t * restrict dst, const uint8_t * restri
    }
 }

+#define hvx_abs_f16_loop_body(dst_type, src_type, vec_store)               \
+    do {                                                                   \
+        dst_type * restrict vdst  = (dst_type *) dst;                      \
+        src_type * restrict vsrc = (src_type *) src;                       \
+                                                                           \
+        const uint32_t elem_size = sizeof(_Float16);                       \
+        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_f16(vsrc[i]);                            \
+        }                                                                  \
+        if (nloe) {                                                        \
+            HVX_Vector v = hvx_vec_abs_f16(vsrc[i]);                       \
+            vec_store((void *) &vdst[i], nloe * elem_size, v);             \
+        }                                                                  \
+    } while(0)
+
+static inline void hvx_abs_f16_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_abs_f16_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
+}
+
+static inline void hvx_abs_f16_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    hvx_abs_f16_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a);
+}
+
+static inline void hvx_abs_f16_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) src % 128 == 0);
+    hvx_abs_f16_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u);
+}
+
+static inline void hvx_abs_f16_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    hvx_abs_f16_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u);
+}
+
+static inline void hvx_abs_f16(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t num_elems) {
+    if (hex_is_aligned((void *) dst, 128)) {
+        if (hex_is_aligned((void *) src, 128)) {
+            hvx_abs_f16_aa(dst, src, num_elems);
+        } else {
+            hvx_abs_f16_au(dst, src, num_elems);
+        }
+    } else {
+        if (hex_is_aligned((void *) src, 128)) {
+            hvx_abs_f16_ua(dst, src, num_elems);
+        } else {
+            hvx_abs_f16_uu(dst, src, num_elems);
+        }
+    }
+}
+
 //
 // Square
 //

-#define hvx_sqr_f32_loop_body(dst_type, src_type, vec_store)           \
+#define hvx_sqr_f32_loop_body(dst_type, src_type, vec_store)               \
    do {                                                                   \
        dst_type * restrict vdst  = (dst_type *) dst;                      \
        src_type * restrict vsrc = (src_type *) src;                       \
@@ -404,10 +510,10 @@ static inline void hvx_abs_f32_aa(uint8_t * restrict dst, const uint8_t * restri
                                                                           \
        _Pragma("unroll(4)")                                               \
        for (; i < nvec; i++) {                                            \
-            vdst[i] = HVX_OP_MUL_F32(vsrc[i], vsrc[i]);                        \
+            vdst[i] = HVX_OP_MUL_F32(vsrc[i], vsrc[i]);                    \
        }                                                                  \
        if (nloe) {                                                        \
-            HVX_Vector v = HVX_OP_MUL_F32(vsrc[i], vsrc[i]);                   \
+            HVX_Vector v = HVX_OP_MUL_F32(vsrc[i], vsrc[i]);               \
            vec_store((void *) &vdst[i], nloe * elem_size, v);             \
        }                                                                  \
    } while(0)
@@ -448,6 +554,64 @@ static inline void hvx_sqr_f32(uint8_t * restrict dst, const uint8_t * restrict
    }
 }

+#define hvx_sqr_f16_loop_body(dst_type, src_type, vec_store)               \
+    do {                                                                   \
+        dst_type * restrict vdst  = (dst_type *) dst;                      \
+        src_type * restrict vsrc = (src_type *) src;                       \
+                                                                           \
+        const uint32_t elem_size = sizeof(_Float16);                       \
+        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_OP_MUL_F16(vsrc[i], vsrc[i]);                    \
+        }                                                                  \
+        if (nloe) {                                                        \
+            HVX_Vector v = HVX_OP_MUL_F16(vsrc[i], vsrc[i]);               \
+            vec_store((void *) &vdst[i], nloe * elem_size, v);             \
+        }                                                                  \
+    } while(0)
+
+static inline void hvx_sqr_f16_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_sqr_f16_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
+}
+
+static inline void hvx_sqr_f16_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    hvx_sqr_f16_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a);
+}
+
+static inline void hvx_sqr_f16_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) src % 128 == 0);
+    hvx_sqr_f16_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u);
+}
+
+static inline void hvx_sqr_f16_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    hvx_sqr_f16_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u);
+}
+
+static inline void hvx_sqr_f16(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t num_elems) {
+    if (hex_is_aligned((void *) dst, 128)) {
+        if (hex_is_aligned((void *) src, 128)) {
+            hvx_sqr_f16_aa(dst, src, num_elems);
+        } else {
+            hvx_sqr_f16_au(dst, src, num_elems);
+        }
+    } else {
+        if (hex_is_aligned((void *) src, 128)) {
+            hvx_sqr_f16_ua(dst, src, num_elems);
+        } else {
+            hvx_sqr_f16_uu(dst, src, num_elems);
+        }
+    }
+}
+
 #undef HVX_OP_ADD_F32
 #undef HVX_OP_SUB_F32
 #undef HVX_OP_MUL_F32
@@ -464,6 +628,7 @@ static inline void hvx_sqr_f32(uint8_t * restrict dst, const uint8_t * restrict
 #undef hvx_scalar_loop_body
 #undef HVX_OP_MIN_SCALAR
 #undef HVX_OP_CLAMP_SCALAR
+#undef HVX_OP_CLAMP_SCALAR_F16
 #undef DEFINE_HVX_BINARY_OP_VARIANTS
 #undef HVX_BINARY_DISPATCHER
 #undef UNUSED
@@ -86,4 +86,33 @@ static inline void hvx_log_f32_aa(uint8_t * restrict dst, const uint8_t * restri
    }
 }

+// Compute log(x) for f16 by promoting to f32, applying hvx_vec_log_f32, and narrowing back.
+static inline void hvx_log_f16_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 nvec = n / VLEN_FP16;
+    const uint32_t nloe = n % VLEN_FP16;
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        HVX_VectorPair p = hvx_vec_f16_to_f32(vsrc[i]);
+        HVX_Vector r0 = hvx_vec_log_f32(Q6_V_lo_W(p));
+        HVX_Vector r1 = hvx_vec_log_f32(Q6_V_hi_W(p));
+        vdst[i] = hvx_vec_f32_to_f16(r0, r1);
+    }
+    if (nloe) {
+        HVX_VectorPair p = hvx_vec_f16_to_f32(vsrc[i]);
+        HVX_Vector r0 = hvx_vec_log_f32(Q6_V_lo_W(p));
+        HVX_Vector r1 = hvx_vec_log_f32(Q6_V_hi_W(p));
+        HVX_Vector v = hvx_vec_f32_to_f16(r0, r1);
+        hvx_vec_store_a((void *) &vdst[i], nloe * SIZEOF_FP16, v);
+    }
+}
+
 #endif /* HVX_LOG_H */
@@ -254,4 +254,201 @@ static inline void hvx_fast_l2_norm_f32(const uint8_t * restrict src,
    }
 }

+// F16 norm kernels: reduce and scale in f32 (via promote/narrow), matching the
+// precision-preserving pattern used by the flash-attn f16 kernels.
+
+static inline void hvx_fast_rms_norm_f16(const uint8_t * restrict src,
+                                         uint8_t * restrict dst,
+                                         const int num_elems,
+                                         float     epsilon) {
+
+    const HVX_Vector * restrict v_src = (HVX_Vector *) src;
+    HVX_Vector * restrict v_dst       = (HVX_Vector *) dst;
+
+    const int nvec = num_elems / VLEN_FP16;    // number of full f16 vectors
+    const int nloe = num_elems % VLEN_FP16;    // leftover elements
+
+    HVX_Vector sum_v = Q6_V_vsplat_R(0x00000000);
+    HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
+
+    #pragma unroll(4)
+    for (int i = 0; i < nvec; i++) {
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v_src[i]);
+        HVX_Vector p0 = Q6_V_lo_W(p);
+        HVX_Vector p1 = Q6_V_hi_W(p);
+        sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, Q6_Vqf32_vmpy_VsfVsf(p0, p0));
+        sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, Q6_Vqf32_vmpy_VsfVsf(p1, p1));
+    }
+
+    if (nloe > 0) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * SIZEOF_FP16);
+        HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v1);
+        HVX_Vector p0 = Q6_V_lo_W(p);
+        HVX_Vector p1 = Q6_V_hi_W(p);
+        sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, Q6_Vqf32_vmpy_VsfVsf(p0, p0));
+        sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, Q6_Vqf32_vmpy_VsfVsf(p1, p1));
+    }
+
+    sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
+
+    HVX_Vector t_v            = hvx_vec_splat_f32((float) num_elems);
+    HVX_Vector denom_v        = hvx_vec_inverse_f32(t_v);
+    HVX_Vector mean_v         = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v);
+    HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v);
+
+    HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
+
+    #pragma unroll(4)
+    for (int i = 0; i < nvec; i++) {
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v_src[i]);
+        HVX_Vector r0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(p), scale_v));
+        HVX_Vector r1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(p), scale_v));
+        v_dst[i] = hvx_vec_f32_to_f16(r0, r1);
+    }
+
+    if (nloe > 0) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * SIZEOF_FP16);
+        HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v1);
+        HVX_Vector r0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(p), scale_v));
+        HVX_Vector r1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(p), scale_v));
+        HVX_Vector result = hvx_vec_f32_to_f16(r0, r1);
+        hvx_vec_store_a(&v_dst[nvec], nloe * SIZEOF_FP16, result);
+    }
+}
+
+static inline void hvx_fast_norm_f16(const uint8_t * restrict src,
+                                     uint8_t * restrict dst,
+                                     const int num_elems,
+                                     float     epsilon) {
+
+    const HVX_Vector * restrict v_src = (HVX_Vector *) src;
+    HVX_Vector * restrict v_dst       = (HVX_Vector *) dst;
+
+    const int nvec = num_elems / VLEN_FP16;
+    const int nloe = num_elems % VLEN_FP16;
+
+    HVX_Vector sum_sq_v = Q6_V_vsplat_R(0x00000000);
+    HVX_Vector sum_x_v  = Q6_V_vsplat_R(0x00000000);
+    HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
+
+    #pragma unroll(4)
+    for (int i = 0; i < nvec; i++) {
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v_src[i]);
+        HVX_Vector p0 = Q6_V_lo_W(p);
+        HVX_Vector p1 = Q6_V_hi_W(p);
+        sum_sq_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_sq_v, Q6_Vqf32_vmpy_VsfVsf(p0, p0));
+        sum_sq_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_sq_v, Q6_Vqf32_vmpy_VsfVsf(p1, p1));
+        sum_x_v  = Q6_Vqf32_vadd_Vqf32Vqf32(sum_x_v,  Q6_Vqf32_vadd_VsfVsf(p0, Q6_V_vzero()));
+        sum_x_v  = Q6_Vqf32_vadd_Vqf32Vqf32(sum_x_v,  Q6_Vqf32_vadd_VsfVsf(p1, Q6_V_vzero()));
+    }
+
+    if (nloe > 0) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * SIZEOF_FP16);
+        HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v1);
+        HVX_Vector p0 = Q6_V_lo_W(p);
+        HVX_Vector p1 = Q6_V_hi_W(p);
+        sum_sq_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_sq_v, Q6_Vqf32_vmpy_VsfVsf(p0, p0));
+        sum_sq_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_sq_v, Q6_Vqf32_vmpy_VsfVsf(p1, p1));
+        sum_x_v  = Q6_Vqf32_vadd_Vqf32Vqf32(sum_x_v,  Q6_Vqf32_vadd_VsfVsf(p0, Q6_V_vzero()));
+        sum_x_v  = Q6_Vqf32_vadd_Vqf32Vqf32(sum_x_v,  Q6_Vqf32_vadd_VsfVsf(p1, Q6_V_vzero()));
+    }
+
+    sum_sq_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_sq_v));
+    sum_x_v  = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_x_v));
+
+    HVX_Vector t_v            = hvx_vec_splat_f32((float) num_elems);
+    HVX_Vector denom_v        = hvx_vec_inverse_f32(t_v);
+    HVX_Vector mean_sq_v      = Q6_Vqf32_vmpy_VsfVsf(sum_sq_v, denom_v);
+    HVX_Vector mean_x_v       = Q6_Vqf32_vmpy_VsfVsf(sum_x_v,  denom_v);
+    HVX_Vector mean_x_sq_v    = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(mean_x_v), Q6_Vsf_equals_Vqf32(mean_x_v));
+    HVX_Vector var_v          = Q6_Vqf32_vsub_Vqf32Vqf32(mean_sq_v, mean_x_sq_v);
+    HVX_Vector var_epsilon_v  = Q6_Vqf32_vadd_Vqf32Vsf(var_v, epsilon_v);
+
+    HVX_Vector scale_v  = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(var_epsilon_v));
+    HVX_Vector mean_x_b = hvx_vec_repl_f32(Q6_Vsf_equals_Vqf32(mean_x_v));
+
+    #pragma unroll(4)
+    for (int i = 0; i < nvec; i++) {
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v_src[i]);
+        HVX_Vector d0 = Q6_Vqf32_vsub_VsfVsf(Q6_V_lo_W(p), mean_x_b);
+        HVX_Vector d1 = Q6_Vqf32_vsub_VsfVsf(Q6_V_hi_W(p), mean_x_b);
+        HVX_Vector r0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(d0), scale_v));
+        HVX_Vector r1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(d1), scale_v));
+        v_dst[i] = hvx_vec_f32_to_f16(r0, r1);
+    }
+
+    if (nloe > 0) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * SIZEOF_FP16);
+        HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v1);
+        HVX_Vector d0 = Q6_Vqf32_vsub_VsfVsf(Q6_V_lo_W(p), mean_x_b);
+        HVX_Vector d1 = Q6_Vqf32_vsub_VsfVsf(Q6_V_hi_W(p), mean_x_b);
+        HVX_Vector r0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(d0), scale_v));
+        HVX_Vector r1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(d1), scale_v));
+        HVX_Vector result = hvx_vec_f32_to_f16(r0, r1);
+        hvx_vec_store_a(&v_dst[nvec], nloe * SIZEOF_FP16, result);
+    }
+}
+
+static inline void hvx_fast_l2_norm_f16(const uint8_t * restrict src,
+                                        uint8_t * restrict dst,
+                                        const int num_elems,
+                                        float     epsilon) {
+
+    const HVX_Vector * restrict v_src = (HVX_Vector *) src;
+    HVX_Vector * restrict v_dst       = (HVX_Vector *) dst;
+
+    const int nvec = num_elems / VLEN_FP16;
+    const int nloe = num_elems % VLEN_FP16;
+
+    HVX_Vector sum_v = hvx_vec_splat_f32(0.0f);
+
+    #pragma unroll(4)
+    for (int i = 0; i < nvec; i++) {
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v_src[i]);
+        HVX_Vector p0 = Q6_V_lo_W(p);
+        HVX_Vector p1 = Q6_V_hi_W(p);
+        sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, Q6_Vqf32_vmpy_VsfVsf(p0, p0));
+        sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, Q6_Vqf32_vmpy_VsfVsf(p1, p1));
+    }
+
+    if (nloe > 0) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * SIZEOF_FP16);
+        HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v1);
+        HVX_Vector p0 = Q6_V_lo_W(p);
+        HVX_Vector p1 = Q6_V_hi_W(p);
+        sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, Q6_Vqf32_vmpy_VsfVsf(p0, p0));
+        sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, Q6_Vqf32_vmpy_VsfVsf(p1, p1));
+    }
+
+    HVX_Vector sum_sf    = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
+    HVX_Vector rsqrt_v   = hvx_vec_rsqrt_f32(sum_sf);
+    HVX_Vector sqrt_v    = hvx_vec_inverse_f32(rsqrt_v);
+    HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
+    HVX_Vector denom_v   = Q6_Vsf_vmax_VsfVsf(sqrt_v, epsilon_v);
+    HVX_Vector scale_v   = hvx_vec_inverse_f32(denom_v);
+
+    #pragma unroll(4)
+    for (int i = 0; i < nvec; i++) {
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v_src[i]);
+        HVX_Vector r0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(p), scale_v));
+        HVX_Vector r1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(p), scale_v));
+        v_dst[i] = hvx_vec_f32_to_f16(r0, r1);
+    }
+
+    if (nloe > 0) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * SIZEOF_FP16);
+        HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+        HVX_VectorPair p = hvx_vec_f16_to_f32(v1);
+        HVX_Vector r0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(p), scale_v));
+        HVX_Vector r1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(p), scale_v));
+        HVX_Vector result = hvx_vec_f32_to_f16(r0, r1);
+        hvx_vec_store_a(&v_dst[nvec], nloe * SIZEOF_FP16, result);
+    }
+}
+
 #endif // HVX_NORM_H
@@ -130,4 +130,70 @@ static inline void hvx_scale_offset_f32(uint8_t * restrict dst, const uint8_t *
    }
 }

+// Scale+offset computed by promoting f16 -> f32, then narrowing the result back to f16.
+#define hvx_scale_offset_f16_loop_body(dst_type, src_type, vec_store)                                                \
+    do {                                                                                                             \
+        dst_type * restrict vdst = (dst_type *) dst;                                                                 \
+        src_type * restrict vsrc = (src_type *) src;                                                                 \
+                                                                                                                     \
+        HVX_Vector vs = hvx_vec_splat_f32(scale);                                                                    \
+        HVX_Vector vo = hvx_vec_splat_f32(offset);                                                                   \
+                                                                                                                     \
+        const uint32_t nvec = n / VLEN_FP16;                                                                         \
+        const uint32_t nloe = n % VLEN_FP16;                                                                         \
+                                                                                                                     \
+        uint32_t i = 0;                                                                                              \
+                                                                                                                     \
+        _Pragma("unroll(4)")                                                                                         \
+        for (; i < nvec; ++i) {                                                                                      \
+            HVX_VectorPair p = hvx_vec_f16_to_f32(vsrc[i]);                                                          \
+            HVX_Vector r0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vsf(Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(p), vs), vo)); \
+            HVX_Vector r1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vsf(Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(p), vs), vo)); \
+            vdst[i] = hvx_vec_f32_to_f16(r0, r1);                                                                    \
+        }                                                                                                            \
+        if (nloe) {                                                                                                  \
+            HVX_VectorPair p = hvx_vec_f16_to_f32(vsrc[i]);                                                          \
+            HVX_Vector r0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vsf(Q6_Vqf32_vmpy_VsfVsf(Q6_V_lo_W(p), vs), vo)); \
+            HVX_Vector r1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vsf(Q6_Vqf32_vmpy_VsfVsf(Q6_V_hi_W(p), vs), vo)); \
+            HVX_Vector v = hvx_vec_f32_to_f16(r0, r1);                                                               \
+            vec_store((void *) &vdst[i], nloe * SIZEOF_FP16, v);                                                     \
+        }                                                                                                            \
+    } while(0)
+
+static inline void hvx_scale_offset_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) {
+    assert((size_t) dst % 128 == 0);
+    assert((size_t) src % 128 == 0);
+    hvx_scale_offset_f16_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
+}
+
+static inline void hvx_scale_offset_f16_au(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) {
+    assert((size_t) dst % 128 == 0);
+    hvx_scale_offset_f16_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a);
+}
+
+static inline void hvx_scale_offset_f16_ua(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) {
+    assert((size_t) src % 128 == 0);
+    hvx_scale_offset_f16_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u);
+}
+
+static inline void hvx_scale_offset_f16_uu(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) {
+    hvx_scale_offset_f16_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u);
+}
+
+static inline void hvx_scale_offset_f16(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) {
+    if (((size_t) dst & 127) == 0) {
+        if (((size_t) src & 127) == 0) {
+            hvx_scale_offset_f16_aa(dst, src, n, scale, offset);
+        } else {
+            hvx_scale_offset_f16_au(dst, src, n, scale, offset);
+        }
+    } else {
+        if (((size_t) src & 127) == 0) {
+            hvx_scale_offset_f16_ua(dst, src, n, scale, offset);
+        } else {
+            hvx_scale_offset_f16_uu(dst, src, n, scale, offset);
+        }
+    }
+}
+
 #endif // HVX_SCALE_H
@@ -123,4 +123,67 @@ static inline void hvx_sqrt_f32(uint8_t * restrict dst, const uint8_t * restrict
    }
 }

+// Compute sqrt(x) for f16 by promoting to f32, applying hvx_vec_rsqrt_f32, and narrowing back.
+#define hvx_sqrt_f16_loop_body(dst_type, src_type, vec_store)                          \
+    do {                                                                               \
+        dst_type * restrict vdst = (dst_type *) dst;                                   \
+        src_type * restrict vsrc = (src_type *) src;                                   \
+                                                                                       \
+        const uint32_t nvec = n / VLEN_FP16;                                           \
+        const uint32_t nloe = n % VLEN_FP16;                                           \
+                                                                                       \
+        uint32_t i = 0;                                                                \
+                                                                                       \
+        _Pragma("unroll(4)")                                                           \
+        for (; i < nvec; i++) {                                                        \
+            HVX_VectorPair p = hvx_vec_f16_to_f32(vsrc[i]);                            \
+            HVX_Vector r0 = HVX_OP_MUL(hvx_vec_rsqrt_f32(Q6_V_lo_W(p)), Q6_V_lo_W(p)); \
+            HVX_Vector r1 = HVX_OP_MUL(hvx_vec_rsqrt_f32(Q6_V_hi_W(p)), Q6_V_hi_W(p)); \
+            vdst[i] = hvx_vec_f32_to_f16(r0, r1);                                      \
+        }                                                                              \
+        if (nloe) {                                                                    \
+            HVX_VectorPair p = hvx_vec_f16_to_f32(vsrc[i]);                            \
+            HVX_Vector r0 = HVX_OP_MUL(hvx_vec_rsqrt_f32(Q6_V_lo_W(p)), Q6_V_lo_W(p)); \
+            HVX_Vector r1 = HVX_OP_MUL(hvx_vec_rsqrt_f32(Q6_V_hi_W(p)), Q6_V_hi_W(p)); \
+            HVX_Vector v = hvx_vec_f32_to_f16(r0, r1);                                 \
+            vec_store((void *) &vdst[i], nloe * SIZEOF_FP16, v);                       \
+        }                                                                              \
+    } while(0)
+
+static inline void hvx_sqrt_f16_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_sqrt_f16_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
+}
+
+static inline void hvx_sqrt_f16_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    hvx_sqrt_f16_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a);
+}
+
+static inline void hvx_sqrt_f16_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) src % 128 == 0);
+    hvx_sqrt_f16_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u);
+}
+
+static inline void hvx_sqrt_f16_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    hvx_sqrt_f16_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u);
+}
+
+static inline void hvx_sqrt_f16(uint8_t * restrict dst, const uint8_t * restrict src, const int num_elems) {
+    if ((unsigned long) dst % 128 == 0) {
+        if ((unsigned long) src % 128 == 0) {
+            hvx_sqrt_f16_aa(dst, src, num_elems);
+        } else {
+            hvx_sqrt_f16_au(dst, src, num_elems);
+        }
+    } else {
+        if ((unsigned long) src % 128 == 0) {
+            hvx_sqrt_f16_ua(dst, src, num_elems);
+        } else {
+            hvx_sqrt_f16_uu(dst, src, num_elems);
+        }
+    }
+}
+
 #endif /* HVX_SQRT_H */
@@ -753,6 +753,9 @@ 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_ID_NX:
+            return op_matmul_id_nx(octx);
+
        case HTP_OP_MUL_MAT_NX:
            return op_matmul_nx(octx);

@@ -878,8 +881,8 @@ static inline void drop_mmap(struct htp_context *ctx, struct htp_mmap *m) {
    }
 }

-static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) {
-    if (b->base) return; // already mapped
+static inline bool mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) {
+    if (b->base) return true; // already mapped

    // find unused mapping
    for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) {
@@ -887,8 +890,8 @@ static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) {
        if (!m->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
+                FARF(HIGH, "mmap failed (will attempt defrag) : fd %u size %u", b->fd, (uint32_t) b->size);
+                return false;
            }

            m->base   = b->base = (uint64_t) va;
@@ -896,12 +899,12 @@ static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) {
            m->size   = b->size;

            FARF(ALWAYS, "mmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size);
-            return;
+            return true;
        }
    }

    FARF(ERROR, "mmap failed : exceeded mapping capacity limit of %u", HTP_MAX_MMAPS);
-    abort();
+    return false;
 }

 static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uint32_t n_bufs) {
@@ -934,12 +937,32 @@ static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uin
        }
    }

-    // Create missing mappings
+    // Create missing mappings (pass 1)
+    bool mmap_ok = true;
    for (uint32_t i=0; i < n_bufs; i++) {
        struct htp_buf_desc *b = bufs + i;
-        mmap_buf(ctx, b);
+        if (!mmap_buf(ctx, b)) {
+            mmap_ok = false;
+            break;
+        }
        FARF(HIGH, "prep-buf #%u : pass1 fd %u base %p size %u flags 0x%x", i, b->fd, (void*) b->base, (uint32_t) b->size, b->flags);
    }
+
+    if (!mmap_ok) {
+        // Attempt clean defragmentation: drop all mappings and remap (pass 2)
+        FARF(HIGH, "prep-bufs : dropping all mappings to defragment address space");
+        for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) { drop_mmap(ctx, ctx->mmap + i); }
+
+        for (uint32_t i=0; i < n_bufs; i++) {
+            struct htp_buf_desc *b = bufs + i;
+            b->base = 0;
+            if (!mmap_buf(ctx, b)) {
+                FARF(ERROR, "prep-bufs : mmap failed after defragmentation (fd %u size %u)", b->fd, (uint32_t) b->size);
+                abort();
+            }
+            FARF(HIGH, "prep-buf #%u : pass2 fd %u base %p size %u flags 0x%x", i, b->fd, (void*) b->base, (uint32_t) b->size, b->flags);
+        }
+    }
 }

 static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t idx, struct htp_tensor *t) {
@@ -981,7 +1004,7 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, u
        octx->src_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma;

        FARF(HIGH, "prep-src #%u: data %p size %u : %u:%u:%u:%u", op->src[i], (void*) src->data, src->size,
-            src->ne[0], src->ne[1], src->ne[3], src->ne[3]);
+            src->ne[0], src->ne[1], src->ne[2], src->ne[3]);
    }

    htp_tensor_flush_all(octx->ctx, octx->src, HTP_OP_MAX_INPUTS);
@@ -134,7 +134,8 @@ static inline int htp_mm_hmx_compute_chunks(size_t   vtcm_total,
    size_t best_mn   = 0;
    size_t best_m = 0, best_n = 0;

-    const size_t n_max = hex_align_down((size_t)n, HTP_MM_HMX_TILE_N_COLS);
+    const size_t max_nc_budget = (usable / per_n_cost);
+    const size_t n_max = hex_align_down(hex_smin((size_t)n, max_nc_budget), HTP_MM_HMX_TILE_N_COLS);
    for (size_t nc = n_max; nc >= HTP_MM_HMX_TILE_N_COLS; nc -= HTP_MM_HMX_TILE_N_COLS) {
        size_t n_fixed = 0, ncmn = 0, mc_denom = 0;
        if (hex_mul_overflow(nc, per_n_cost, &n_fixed)) continue;
@@ -299,6 +300,15 @@ static inline void htp_mm_hmx_get_batched_chunk_costs(
    *size_per_mn_out = sizeof(uint16_t);
 }

+static inline size_t htp_mm_hmx_get_2d_overhead(bool pipeline, bool is_matmul_id) {
+    size_t num_regions = pipeline ? 7 : (is_matmul_id ? 4 : 5);
+    return num_regions * HTP_MM_HMX_TILE_SIZE + 256;
+}
+
+static inline size_t htp_mm_hmx_get_batched_overhead(void) {
+    return 5 * HTP_MM_HMX_TILE_SIZE + 256;
+}
+
 struct htp_mm_hmx_vtcm_layout {
    // Byte offsets from vtcm_base for each region
    size_t off_weight[2];     // [1] is only used when pipelined
@@ -568,10 +578,8 @@ static inline void htp_mm_hvx_vtcm_layout_build(
                }

                size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
-                size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
-                if (dst_size_per_thread < quant_scratch_size_per_thread) {
-                    dst_size_per_thread = quant_scratch_size_per_thread;
-                }
+                size_t dst_slice_per_thread = (dst_nrows > 0 && src1_nrows == 1) ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0;
+                size_t dst_size_per_thread = (dst_slice_per_thread > quant_scratch_size_per_thread) ? dst_slice_per_thread : quant_scratch_size_per_thread;
                dst_sz = dst_size_per_thread * n_threads;
                break;
            }
@@ -592,10 +600,8 @@ static inline void htp_mm_hvx_vtcm_layout_build(
                }

                size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
-                size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
-                if (dst_size_per_thread < quant_scratch_size_per_thread) {
-                    dst_size_per_thread = quant_scratch_size_per_thread;
-                }
+                size_t dst_slice_per_thread = dst_nrows > 0 ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0;
+                size_t dst_size_per_thread = (dst_slice_per_thread > quant_scratch_size_per_thread) ? dst_slice_per_thread : quant_scratch_size_per_thread;
                dst_sz = dst_size_per_thread * n_threads;
                break;
            }
@@ -658,7 +664,7 @@ static inline bool htp_mm_hmx_solve_batched_params(

    int act_threads = n_threads;
    while (act_threads >= 1) {
-        size_t group_overhead = 256;
+        size_t group_overhead = htp_mm_hmx_get_batched_overhead();
        size_t group_size_per_n, group_size_per_m, group_size_per_mn;
        htp_mm_hmx_get_batched_chunk_costs(k, group_size, &group_size_per_n, &group_size_per_m, &group_size_per_mn);

@@ -725,7 +731,7 @@ static inline bool htp_mm_hmx_solve_2d_params(

    int act_threads = n_threads;
    while (act_threads >= 1) {
-        size_t simple_2d_overhead = 256;
+        size_t simple_2d_overhead = htp_mm_hmx_get_2d_overhead(pipeline, is_matmul_id);
        size_t simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn;
        htp_mm_hmx_get_2d_chunk_costs(wtype, k, pipeline, aligned_tile_size, &simple_2d_size_per_n, &simple_2d_size_per_m, &simple_2d_size_per_mn);

@@ -216,6 +216,14 @@ int op_set_rows(struct htp_ops_context * octx) {
        default:            return HTP_STATUS_NO_SUPPORT;
    }

+    FARF(HIGH, "set-rows: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu dst-vtcm-size %zu n_threads %d\n",
+         octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
+         octx->src[1]->ne[0], octx->src[1]->ne[1], octx->src[1]->ne[2], octx->src[1]->ne[3],
+         octx->dst->ne[0], octx->dst->ne[1], octx->dst->ne[2], octx->dst->ne[3],
+         srctx.vtcm_layout.src0_bytes_per_thread * kparams->n_threads,
+         srctx.vtcm_layout.dst_bytes_per_thread  * kparams->n_threads,
+         kparams->n_threads);
+
    work_queue_run(octx->ctx->work_queue, q_func, &srctx, kparams->n_threads);

    return HTP_STATUS_OK;
@@ -234,6 +234,146 @@ static void sqrt_f32(const float * restrict src,
    }
 }

+static void scale_f16(const _Float16 * restrict src,
+                      _Float16 * restrict dst,
+                      const uint32_t num_rows,
+                      const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+    float scale = 0.f;
+    float bias  = 0.f;
+    memcpy(&scale, &op_params[0], sizeof(float));
+    memcpy(&bias,  &op_params[1], sizeof(float));
+
+    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_scale_offset_f16_aa((uint8_t *) dst_local, (const uint8_t *) src_local, ne0, scale, bias);
+    }
+}
+
+static void clamp_f16(const _Float16 * restrict src,
+                      _Float16 * restrict dst,
+                      const uint32_t num_rows,
+                      const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+    float min = 0.f;
+    float max = 0.f;
+    memcpy(&min, &op_params[0], sizeof(float));
+    memcpy(&max, &op_params[1], sizeof(float));
+
+    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_clamp_scalar_f16(dst_local, src_local, (_Float16) min, (_Float16) max, ne0);
+    }
+}
+
+static void rms_norm_f16(const _Float16 * restrict src,
+                         _Float16 * restrict dst,
+                         const uint32_t num_rows,
+                         const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+    float epsilon = 0.f;
+    memcpy(&epsilon, op_params, sizeof(float));
+
+    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_fast_rms_norm_f16((const uint8_t *) src_local, (uint8_t *) dst_local, ne0, epsilon);
+    }
+}
+
+static void norm_f16(const _Float16 * restrict src,
+                     _Float16 * restrict dst,
+                     const uint32_t num_rows,
+                     const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+    float epsilon = 0.f;
+    memcpy(&epsilon, op_params, sizeof(float));
+
+    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_fast_norm_f16((const uint8_t *) src_local, (uint8_t *) dst_local, ne0, epsilon);
+    }
+}
+
+static void sqr_f16(const _Float16 * restrict src,
+                    _Float16 * 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_sqr_f16_aa((uint8_t *) dst_local, (const uint8_t *) src_local, ne0);
+    }
+}
+
+static void sqrt_f16(const _Float16 * restrict src,
+                     _Float16 * 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_sqrt_f16_aa((uint8_t *) dst_local, (const uint8_t *) src_local, ne0);
+    }
+}
+
+static void abs_f16(const _Float16 * restrict src,
+                    _Float16 * 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_f16_aa((uint8_t *) dst_local, (const uint8_t *) src_local, ne0);
+    }
+}
+
+static void log_f16(const _Float16 * restrict src,
+                    _Float16 * 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_f16_aa((uint8_t *) dst_local, (const uint8_t *) src_local, ne0);
+    }
+}
+
+static void l2_norm_f16(const _Float16 * restrict src,
+                        _Float16 * restrict dst,
+                        const uint32_t num_rows,
+                        const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+    float epsilon = 0.f;
+    memcpy(&epsilon, op_params, sizeof(float));
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_f = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_f       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_fast_l2_norm_f16((const uint8_t *)src_f, (uint8_t *)dst_f, ne0, epsilon);
+    }
+}
+
 static void neg_f32(const float * restrict src,
                    float * restrict dst,
                    const uint32_t num_rows,
@@ -471,8 +611,8 @@ static void log_f32(const float * restrict src,
    }
 }

-#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) {                                \
+#define DEFINE_UNARY_TASK_IMPL(NAME, TYPE, SUFFIX, IS_RMS_NORM_MUL, IS_TRI, CORE_EXPR)                              \
+static void unary_task_##SUFFIX##_##NAME(unsigned int nth, unsigned int ith, void * data) {                         \
    const struct htp_unary_context * uctx = (const struct htp_unary_context *) data;                                \
    struct htp_ops_context * octx = uctx->octx;                                                                     \
    const struct htp_tensor * src = octx->src[0];                                                                   \
@@ -536,7 +676,7 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
    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",                 \
+        FARF(ERROR, "unary-" #SUFFIX " : current VTCM reservation %zu is too small, needed at least %zu\n",         \
             uctx->vtcm_src0_size_per_thread, src0_row_size_aligned);                                               \
        return;                                                                                                     \
    }                                                                                                               \
@@ -578,11 +718,11 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
        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;                                                 \
-        float * src1_vtcm = NULL;                                                                                   \
+        TYPE * dst_vtcm  = (TYPE *) dma_queue_pop(dma_queue).src;                                                   \
+        TYPE * src0_vtcm = (TYPE *) dma_queue_pop(dma_queue).dst;                                                   \
+        TYPE * src1_vtcm = NULL;                                                                                    \
        if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) {                                                         \
-            src1_vtcm = (float *) dma_queue_pop(dma_queue).dst;                                                     \
+            src1_vtcm = (TYPE *) dma_queue_pop(dma_queue).dst;                                                      \
        }                                                                                                           \
                                                                                                                    \
        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                      \
@@ -625,6 +765,10 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
    dma_queue_flush(dma_queue);                                                                                     \
 }

+// F32 unary task: row-block DMA/VTCM plumbing, float-typed VTCM buffers.
+#define DEFINE_UNARY_TASK(NAME, IS_RMS_NORM_MUL, IS_TRI, CORE_EXPR) \
+    DEFINE_UNARY_TASK_IMPL(NAME, float, f32, IS_RMS_NORM_MUL, IS_TRI, CORE_EXPR)
+
 DEFINE_UNARY_TASK(norm,           false, false, norm_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(rms_norm,       false, false, rms_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(rms_norm_mul,   true,  false, rms_norm_mul_f32(src0_vtcm, uctx->broadcast_weight ? (const float *) src1_vtcm_data : src1_vtcm, dst_vtcm, block_size, uctx))
@@ -644,6 +788,18 @@ DEFINE_UNARY_TASK(unary_log,      false, false, log_f32(src0_vtcm, dst_vtcm, blo
 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))

+// F16 unary tasks: same DMA/VTCM plumbing as DEFINE_UNARY_TASK, but VTCM buffers are
+// _Float16-typed. None of the current F16 ops need RMS_NORM_MUL or TRI support.
+DEFINE_UNARY_TASK_IMPL(norm,      _Float16, f16, false, false, norm_f16(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK_IMPL(rms_norm,  _Float16, f16, false, false, rms_norm_f16(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK_IMPL(scale,     _Float16, f16, false, false, scale_f16(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK_IMPL(clamp,     _Float16, f16, false, false, clamp_f16(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK_IMPL(sqr,       _Float16, f16, false, false, sqr_f16(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK_IMPL(sqrt,      _Float16, f16, false, false, sqrt_f16(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK_IMPL(l2_norm,   _Float16, f16, false, false, l2_norm_f16(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK_IMPL(unary_abs, _Float16, f16, false, false, abs_f16(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK_IMPL(unary_log, _Float16, f16, false, false, log_f16(src0_vtcm, dst_vtcm, block_size, uctx))
+
 // Apply a pointwise unary op to one column tile that is already in VTCM.
 #define DEFINE_UNARY_TILED_TASK(NAME, IS_TRI, CORE_TILE_EXPR)                                                       \
 static void unary_task_f32_tiled_##NAME(unsigned int nth, unsigned int ith, void * data) {                          \
@@ -892,50 +1048,76 @@ DEFINE_UNARY_TILED_TASK(unary_abs,      false, hvx_abs_f32_aa(dst_vtcm, src_vtcm
 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) {
+static int execute_op_unary(struct htp_ops_context * octx) {
    int err = HTP_STATUS_OK;

    const struct htp_tensor * src0 = octx->src[0];
    const struct htp_tensor * dst  = octx->dst;

+    const bool is_f16 = (src0->type == HTP_TYPE_F16);
+
    const char * op_type = NULL;

    switch (octx->op) {
-        case HTP_OP_NORM:            op_type = "norm-f32";         break;
-        case HTP_OP_RMS_NORM:        op_type = "rmsnorm-f32";      break;
-        case HTP_OP_RMS_NORM_MUL:    op_type = "rmsnorm-mul-f32";  break;
-        case HTP_OP_SCALE:           op_type = "scale-f32";        break;
-        case HTP_OP_CLAMP:           op_type = "clamp-f32";        break;
-        case HTP_OP_SQR:             op_type = "sqr-f32";          break;
-        case HTP_OP_SQRT:            op_type = "sqrt-f32";         break;
-        case HTP_OP_UNARY_NEG:       op_type = "neg-f32";          break;
-        case HTP_OP_UNARY_EXP:       op_type = "exp-f32";          break;
-        case HTP_OP_UNARY_SIGMOID:   op_type = "sigmoid-f32";      break;
-        case HTP_OP_UNARY_SILU:      op_type = "silu-f32";         break;
-        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;
+        case HTP_OP_NORM:            op_type = is_f16 ? "norm-f16"     : "norm-f32";         break;
+        case HTP_OP_RMS_NORM:        op_type = is_f16 ? "rmsnorm-f16"  : "rmsnorm-f32";      break;
+        case HTP_OP_RMS_NORM_MUL:    op_type = "rmsnorm-mul-f32";                            break;
+        case HTP_OP_SCALE:           op_type = is_f16 ? "scale-f16"    : "scale-f32";        break;
+        case HTP_OP_CLAMP:           op_type = is_f16 ? "clamp-f16"    : "clamp-f32";        break;
+        case HTP_OP_SQR:             op_type = is_f16 ? "sqr-f16"      : "sqr-f32";          break;
+        case HTP_OP_SQRT:            op_type = is_f16 ? "sqrt-f16"     : "sqrt-f32";         break;
+        case HTP_OP_UNARY_NEG:       op_type = "neg-f32";                                    break;
+        case HTP_OP_UNARY_EXP:       op_type = "exp-f32";                                    break;
+        case HTP_OP_UNARY_SIGMOID:   op_type = "sigmoid-f32";                                break;
+        case HTP_OP_UNARY_SILU:      op_type = "silu-f32";                                   break;
+        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 = is_f16 ? "abs-f16"      : "abs-f32";          break;
+        case HTP_OP_UNARY_LOG:       op_type = is_f16 ? "log-f16"      : "log-f32";          break;
+        case HTP_OP_L2_NORM:         op_type = is_f16 ? "l2norm-f16"   : "l2norm-f32";       break;
+        case HTP_OP_TRI:             op_type = "tri-f32";                                    break;

        default:
            FARF(ERROR, "Unsupported unary Op %u\n", octx->op);
            return HTP_STATUS_NO_SUPPORT;
    }

+    // F16 only has row-block kernels for this subset of ops (see the dispatch switch
+    // below) - reject everything else up front, before touching kparams/VTCM.
+    if (is_f16) {
+        switch (octx->op) {
+            case HTP_OP_NORM:
+            case HTP_OP_RMS_NORM:
+            case HTP_OP_SCALE:
+            case HTP_OP_CLAMP:
+            case HTP_OP_SQR:
+            case HTP_OP_SQRT:
+            case HTP_OP_L2_NORM:
+            case HTP_OP_UNARY_ABS:
+            case HTP_OP_UNARY_LOG:
+                break;
+            default:
+                FARF(ERROR, "unary-%s: not supported for F16\n", op_type);
+                return HTP_STATUS_NO_SUPPORT;
+        }
+    }
+
    const struct htp_unary_kernel_params * kparams = (const struct htp_unary_kernel_params *) octx->kernel_params;

    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
    const uint32_t n_threads  = kparams->n_threads;

-    const size_t src0_data_row_size = src0->ne[0] * sizeof(float);
-    const size_t dst_data_row_size  = dst->ne[0]  * sizeof(float);
+    const size_t elem_size = is_f16 ? sizeof(_Float16) : sizeof(float);
+
+    const size_t src0_data_row_size = src0->ne[0] * elem_size;
+    const size_t dst_data_row_size  = dst->ne[0]  * elem_size;

    const size_t src0_row_size_aligned = kparams->src0_row_size_aligned;
    const size_t dst_row_size_aligned  = kparams->dst_row_size_aligned;

+    // Always 0 for F16 - htp_unary_vtcm_layout_build() keeps F16 on the row-block path,
+    // since only F32 has unary_task_f32_tiled_* kernels.
    const uint32_t col_tile = kparams->col_tile;

    size_t src1_data_row_size = 0;
@@ -943,6 +1125,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
    bool broadcast_weight = kparams->broadcast_weight;
    const struct htp_tensor * src1 = NULL;

+    // RMS_NORM_MUL fusion is F32-only (its weight tensor is always F32; see
+    // try_fuse_node()'s type guard), so this never triggers when is_f16 is true.
    if (octx->op == HTP_OP_RMS_NORM_MUL) {
        src1 = octx->src[1];
        src1_data_row_size = src1->ne[0] * sizeof(float);
@@ -987,7 +1171,7 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {

            .block                 = kparams->block,
            .nc                    = src0->ne[0],
-            .col_tile              = (uint32_t) kparams->col_tile,
+            .col_tile              = col_tile,
            .broadcast_weight      = broadcast_weight,

            .vtcm_src0             = VTCM_LAYOUT_PTR(uint8_t, base, 0),
@@ -1020,6 +1204,19 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
                case HTP_OP_TRI:             task_func = unary_task_f32_tiled_tri;            break;
                default:                     break;
            }
+        } else if (is_f16) {
+            switch (octx->op) {
+                case HTP_OP_NORM:            task_func = unary_task_f16_norm;                 break;
+                case HTP_OP_RMS_NORM:        task_func = unary_task_f16_rms_norm;             break;
+                case HTP_OP_SCALE:           task_func = unary_task_f16_scale;                break;
+                case HTP_OP_CLAMP:           task_func = unary_task_f16_clamp;                break;
+                case HTP_OP_SQR:             task_func = unary_task_f16_sqr;                  break;
+                case HTP_OP_SQRT:            task_func = unary_task_f16_sqrt;                 break;
+                case HTP_OP_L2_NORM:         task_func = unary_task_f16_l2_norm;              break;
+                case HTP_OP_UNARY_ABS:       task_func = unary_task_f16_unary_abs;            break;
+                case HTP_OP_UNARY_LOG:       task_func = unary_task_f16_unary_log;            break;
+                default:                     break;
+            }
        } else {
            switch (octx->op) {
                case HTP_OP_NORM:            task_func = unary_task_f32_norm;                 break;
@@ -1047,7 +1244,7 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
        if (task_func) {
            worker_pool_run_func(octx->ctx->worker_pool, task_func, &uctx, n_threads);
        } else {
-            FARF(ERROR, "execute_op_unary_f32: task function is NULL for op %d\n", octx->op);
+            FARF(ERROR, "execute_op_unary: task function is NULL for op %d\n", octx->op);
            err = HTP_STATUS_NO_SUPPORT;
        }
    }
@@ -1058,7 +1255,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
 int op_unary(struct htp_ops_context * octx) {
    switch (octx->src[0]->type) {
        case HTP_TYPE_F32:
-            return execute_op_unary_f32(octx);
+        case HTP_TYPE_F16:
+            return execute_op_unary(octx);

        default:
            return HTP_STATUS_NO_SUPPORT;
@@ -85,17 +85,19 @@ static inline void htp_unary_vtcm_layout_build(
    bool broadcast_weight,
    uint32_t n_threads,
    size_t vtcm_size,
+    size_t elem_size,
    uint32_t * out_col_tile,
    uint32_t * out_vtcm_row_per_thread
 ) {
-    const size_t src0_data_row_size = ne00 * sizeof(float);
-    const size_t dst_data_row_size  = ne10 * sizeof(float);
+    const size_t src0_data_row_size = ne00 * elem_size;
+    const size_t dst_data_row_size  = ne10 * elem_size;

    const size_t src0_row_size_aligned = hex_round_up(src0_data_row_size, 128);
    const size_t dst_row_size_aligned  = hex_round_up(dst_data_row_size,  128);

    size_t src1_row_size_aligned = 0;
    if (op == HTP_OP_RMS_NORM_MUL) {
+        // RMS_NORM_MUL fusion is F32-only; its weight tensor is always F32.
        const size_t src1_data_row_size = ne11 * sizeof(float);
        src1_row_size_aligned = hex_round_up(src1_data_row_size, 128);
    }
@@ -125,12 +127,19 @@ static inline void htp_unary_vtcm_layout_build(

    const bool is_reduction = (op == HTP_OP_NORM || op == HTP_OP_RMS_NORM ||
                               op == HTP_OP_RMS_NORM_MUL || op == HTP_OP_L2_NORM);
+    // The tiled fallback path below only has F32 task functions (unary_task_f32_tiled_*);
+    // F16 has no tiled kernels, so it must stay on the row-block path like reduction ops.
+    // NOTE: if F16 ends up with vtcm_row_per_thread == 0 here (row too large for the VTCM
+    // budget), execute_op_unary() will see BLOCK == 0 and skip computation for that op
+    // (logged via FARF(ERROR, ...)) since there is no F16 tiled fallback. This is a known
+    // limitation; supporting it would require adding F16 tiled kernels.
+    const bool is_f16 = (elem_size == sizeof(_Float16));
    uint32_t col_tile = 0;

-    if (vtcm_row_per_thread == 0 && !is_reduction) {
+    if (vtcm_row_per_thread == 0 && !is_reduction && !is_f16) {
        const size_t per_thread_budget = vtcm_size / n_threads;
        const size_t col_tile_bytes = hex_align_down(per_thread_budget / 4, 128);
-        col_tile = (uint32_t) (col_tile_bytes / sizeof(float));
+        col_tile = (uint32_t) (col_tile_bytes / elem_size);

        L->src0_bytes = col_tile_bytes * 2;
        L->dst_bytes  = col_tile_bytes * 2;
@@ -127,6 +127,18 @@ else()
        configure_file(${src} ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${src} COPYONLY)
    endforeach()

+    # CMAKE_OSX_SYSROOT is an SDK name or path - xcrun accepts both
+    set(METAL_SDK ${CMAKE_OSX_SYSROOT})
+    if (NOT METAL_SDK)
+        set(METAL_SDK macosx)
+    endif()
+
+    if (CMAKE_OSX_SYSROOT MATCHES "[Ss]imulator")
+        set(METAL_TARGET_SIM "-simulator")
+    else()
+        set(METAL_TARGET_SIM "")
+    endif()
+
    if (GGML_METAL_SHADER_DEBUG)
        # note: disabling fast math is needed in order to pass tests/test-backend-ops
        # note: adding -fno-inline fixes the tests when using MTL_SHADER_VALIDATION=1
@@ -138,9 +150,19 @@ else()
        set(XC_FLAGS -O3)
    endif()

+    execute_process(COMMAND xcrun -sdk ${METAL_SDK} --show-sdk-version OUTPUT_VARIABLE METAL_SDK_VERSION OUTPUT_STRIP_TRAILING_WHITESPACE)
+    if (METAL_SDK_VERSION VERSION_GREATER_EQUAL 26.0)
+        set(GGML_METAL_HAS_TENSOR_LIB ON)
+    else()
+        message(STATUS "Metal SDK ${METAL_SDK_VERSION} does not support the tensor API, skipping ggml-tensor.metallib")
+    endif()
+
    if (GGML_METAL_MACOSX_VERSION_MIN)
        message(STATUS "Adding  -mmacosx-version-min=${GGML_METAL_MACOSX_VERSION_MIN} flag to metal compilation")
        list   (APPEND XC_FLAGS -mmacosx-version-min=${GGML_METAL_MACOSX_VERSION_MIN})
+    elseif (NOT GGML_METAL_TARGET_OS STREQUAL "macos" AND CMAKE_OSX_DEPLOYMENT_TARGET)
+        message(STATUS "Adding  -mtargetos=${GGML_METAL_TARGET_OS}${CMAKE_OSX_DEPLOYMENT_TARGET}${METAL_TARGET_SIM} flag to metal compilation")
+        list   (APPEND XC_FLAGS -mtargetos=${GGML_METAL_TARGET_OS}${CMAKE_OSX_DEPLOYMENT_TARGET}${METAL_TARGET_SIM})
    endif()

    if (GGML_METAL_STD)
@@ -156,26 +178,51 @@ else()
        list(APPEND AIR_FILES ${AIR})
        add_custom_command(
            OUTPUT ${AIR}
-            COMMAND xcrun -sdk macosx metal ${XC_FLAGS} -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${src} -o ${AIR}
+            COMMAND xcrun -sdk ${METAL_SDK} metal ${XC_FLAGS} -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${src} -o ${AIR}
            DEPENDS ${src} kernels/common.h kernels/dequantize.h kernels/quantize.h ${METALLIB_COMMON} ggml-metal-impl.h
            COMMENT "Compiling ${src}"
            VERBATIM
        )
    endforeach()

+    set(METALLIB_FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib)
+
+    # the tensor API kernels go in a separate metallib, loaded only where supported
+    if (GGML_METAL_HAS_TENSOR_LIB)
+        set(AIR_MM_TENSOR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/mul_mm_tensor.air")
+        # the tensor API needs OS 26+
+        set(XC_FLAGS_TENSOR ${XC_FLAGS} -mtargetos=${GGML_METAL_TARGET_OS}26.0${METAL_TARGET_SIM})
+        add_custom_command(
+            OUTPUT ${AIR_MM_TENSOR}
+            COMMAND xcrun -sdk ${METAL_SDK} metal ${XC_FLAGS_TENSOR} -DGGML_METAL_HAS_TENSOR -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/mul_mm.metal -o ${AIR_MM_TENSOR}
+            DEPENDS kernels/mul_mm.metal kernels/common.h kernels/dequantize.h ${METALLIB_COMMON} ggml-metal-impl.h
+            COMMENT "Compiling kernels/mul_mm.metal (tensor API)"
+            VERBATIM
+        )
+
+        add_custom_command(
+            OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-tensor.metallib
+            COMMAND xcrun -sdk ${METAL_SDK} metallib ${AIR_MM_TENSOR} -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-tensor.metallib
+            DEPENDS ${AIR_MM_TENSOR}
+            COMMENT "Linking tensor API Metal kernels into ggml-tensor.metallib"
+        )
+
+        list(APPEND METALLIB_FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-tensor.metallib)
+    endif()
+
    add_custom_command(
        OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
-        COMMAND xcrun -sdk macosx metallib ${AIR_FILES} -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
+        COMMAND xcrun -sdk ${METAL_SDK} metallib ${AIR_FILES} -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
        COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h
        COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal-impl.h
        COMMAND rm -rf ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels
-        DEPENDS ${AIR_FILES}
+        DEPENDS ${AIR_FILES} ${AIR_MM_TENSOR}
        COMMENT "Linking Metal kernels into default.metallib"
    )

    add_custom_target(
        ggml-metal-lib ALL
-        DEPENDS ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
+        DEPENDS ${METALLIB_FILES}
    )
 endif() # GGML_METAL_EMBED_LIBRARY

@@ -187,7 +234,7 @@ if (NOT GGML_METAL_EMBED_LIBRARY)
    )

    install(
-        FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
+        FILES ${METALLIB_FILES}
        DESTINATION ${CMAKE_INSTALL_BINDIR}
    )
 endif()
@@ -69,6 +69,10 @@ struct ggml_metal {
    // extra command buffers for things like getting, setting and copying tensors
    NSMutableArray * cmd_bufs_ext;

+    // buffers to release after async Metal operations complete
+    // if Metal released them, it would do so on a Metal-internal thread without an autorelease pool, which could cause leaks
+    NSMutableArray * buf_refs;
+
    // the last command buffer queued into the Metal queue with operations relevant to the current Metal backend
    id<MTLCommandBuffer> cmd_buf_last;

@@ -107,6 +111,7 @@ ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) {
        id<MTLCommandQueue> queue = ggml_metal_device_get_queue(dev);
        if (queue == nil) {
            GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__);
+            free(res);
            return NULL;
        }

@@ -179,6 +184,7 @@ ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) {
        }

        res->cmd_bufs_ext = [[NSMutableArray alloc] init];
+        res->buf_refs     = [[NSMutableArray alloc] init];

        res->cmd_buf_last = nil;

@@ -206,6 +212,11 @@ void ggml_metal_free(ggml_metal_t ctx) {
    [ctx->cmd_bufs_ext removeAllObjects];
    [ctx->cmd_bufs_ext release];

+    @autoreleasepool {
+        [ctx->buf_refs removeAllObjects];
+        [ctx->buf_refs release];
+    }
+
    if (ctx->pipelines_ext) {
        ggml_metal_pipelines_free(ctx->pipelines_ext);
        ctx->pipelines_ext = nil;
@@ -294,6 +305,10 @@ void ggml_metal_synchronize(ggml_metal_t ctx) {

        [ctx->cmd_bufs_ext removeAllObjects];
    }
+
+    @autoreleasepool {
+        [ctx->buf_refs removeAllObjects];
+    }
 }

 static struct ggml_metal_buffer_id ggml_metal_get_buffer_id(const struct ggml_tensor * t) {
@@ -337,6 +352,8 @@ void ggml_metal_set_tensor_async(ggml_metal_t ctx, struct ggml_tensor * tensor,

        [encoder endEncoding];
        [cmd_buf commit];
+
+        [ctx->buf_refs addObject:buf_src];
        [buf_src release];

        // do not wait here for completion
@@ -381,6 +398,8 @@ void ggml_metal_get_tensor_async(ggml_metal_t ctx, const struct ggml_tensor * te

        [encoder endEncoding];
        [cmd_buf commit];
+
+        [ctx->buf_refs addObject:buf_dst];
        [buf_dst release];

        // do not wait here for completion
@@ -1336,7 +1336,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht(ggml_metal_
    return res;
 }

-// note: reuse the argsort kernel for top_k
+// note: reuse the argsort kernel for the bitonic top_k fallback
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal_library_t lib, const ggml_tensor * op) {
    assert(op->op == GGML_OP_TOP_K);

@@ -1364,6 +1364,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal
    return res;
 }

+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix(ggml_metal_library_t lib, const ggml_tensor * op) {
+    assert(op->op == GGML_OP_TOP_K);
+
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_top_k_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->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);
+    }
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge(ggml_metal_library_t lib, const ggml_tensor * op) {
    assert(op->op == GGML_OP_TOP_K);

@@ -1560,6 +1577,26 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext(
    return res;
 }

+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(
+        ggml_metal_library_t lib,
+        const ggml_tensor * op) {
+    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
+    assert(op->src[3]);
+
+    char name[256];
+
+    snprintf(name, 256, "kernel_flash_attn_ext_vec_idx");
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr);
+    }
+
+    GGML_UNUSED(op);
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec(
        ggml_metal_library_t lib,
        const ggml_tensor * op,
@@ -1568,6 +1605,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
        bool    has_bias,
        bool    has_scap,
        bool    has_kvpad,
+        bool    has_sparse,
        int32_t nqpsg,
        int32_t ne,
        int32_t nsg,
@@ -1597,13 +1635,14 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
            dv,
            qne_suffix);

-    snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
+    snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_sparse=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
            base,
            has_mask,
            has_sinks,
            has_bias,
            has_scap,
            has_kvpad,
+            has_sparse,
            ns10,
            ns20,
            nsg, nwg);
@@ -1616,7 +1655,8 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
        ggml_metal_cv_set_bool(cv, has_sinks, FC_FLASH_ATTN_EXT_VEC + 1);
        ggml_metal_cv_set_bool(cv, has_bias,  FC_FLASH_ATTN_EXT_VEC + 2);
        ggml_metal_cv_set_bool(cv, has_scap,  FC_FLASH_ATTN_EXT_VEC + 3);
-        ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT_VEC + 4);
+        ggml_metal_cv_set_bool(cv, has_kvpad,  FC_FLASH_ATTN_EXT_VEC + 4);
+        ggml_metal_cv_set_bool(cv, has_sparse, FC_FLASH_ATTN_EXT_VEC + 5);

        ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT_VEC + 20);
        ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT_VEC + 21);
@@ -145,6 +145,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge     (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht              (ggml_metal_library_t lib, int n);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k             (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix       (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge       (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin               (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse );
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin_one           (ggml_metal_library_t lib, enum ggml_op op);
@@ -200,6 +201,10 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att
        int32_t ns10,
        int32_t ns20);

+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(
+        ggml_metal_library_t lib,
+        const struct ggml_tensor * op);
+
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec(
        ggml_metal_library_t lib,
        const struct ggml_tensor * op,
@@ -208,6 +213,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att
        bool    has_bias,
        bool    has_scap,
        bool    has_kvpad,
+        bool    has_sparse,
        int32_t nqpsg,
        int32_t ne,
        int32_t nsg,
@@ -258,6 +264,7 @@ enum ggml_metal_device_id {
    GGML_METAL_DEVICE_M5_PRO,
    GGML_METAL_DEVICE_M5_MAX,
    GGML_METAL_DEVICE_M5_ULTRA,
+    GGML_METAL_DEVICE_A18_PRO,
 };

 const char * ggml_metal_device_id_token(enum ggml_metal_device_id id);
@@ -27,6 +27,9 @@
 static const NSInteger MTLGPUFamilyMetal3_GGML = 5001;
 static const NSInteger MTLGPUFamilyMetal4_GGML = 5002;

+// MTLLanguageVersion4_0 is not present in older SDKs
+static const NSUInteger MTLLanguageVersion4_0_GGML = 4 << 16;
+
 #if !GGML_METAL_EMBED_LIBRARY
 // Here to assist with NSBundle Path Hack
@interface GGMLMetalClass : NSObject
@@ -154,6 +157,9 @@ struct ggml_metal_library {
    // nil in single_library mode (everything resolves to objs[0]).
    NSMutableDictionary<NSString *, NSNumber *> * fn_to_lib;

+    // kernels from a second metallib, resolved ahead of the combined library
+    NSSet<NSString *> * override_fns;
+
    ggml_metal_device_t dev;
    ggml_metal_pipelines_t pipelines; // cache of compiled pipelines

@@ -174,6 +180,18 @@ static void ggml_metal_library_build_index(ggml_metal_library_t lib) {
    }
 }

+// note: defined below, after struct ggml_metal_device
+static void ggml_metal_device_disable_tensor(ggml_metal_device_t dev);
+
+// the tensor API headers are exposed to the shader compiler only at Metal language version 4.0
+static void ggml_metal_compile_options_set_lang(MTLCompileOptions * options, bool has_tensor) {
+    if (!has_tensor) {
+        return;
+    }
+
+    options.languageVersion = (MTLLanguageVersion) MTLLanguageVersion4_0_GGML;
+}
+
 // Parse a `#include "name"` line. Returns the quoted name in *include_name on
 // success. Whitespace-tolerant; ignores `#include <...>` (system headers).
 static bool ggml_metal_library_parse_quoted_include(NSString * line, NSString ** include_name) {
@@ -313,6 +331,7 @@ static bool ggml_metal_library_compile_all(
            @autoreleasepool {
                MTLCompileOptions * options = [MTLCompileOptions new];
                options.preprocessorMacros = prep;
+                ggml_metal_compile_options_set_lang(options, ggml_metal_device_get_props(res->dev)->has_tensor);

                lib = [device newLibraryWithSource:src options:options error:&error];

@@ -369,6 +388,46 @@ static bool ggml_metal_library_compile_all(
    return ok;
 }

+// look for <name>.metallib as a bundle resource, then next to the running binary
+static NSString * ggml_metal_find_metallib(NSBundle * bundle, NSString * name) {
+    NSError * error = nil;
+
+    NSString * path_lib = [bundle pathForResource:name ofType:@"metallib"];
+    if (path_lib == nil) {
+        // Try to find the resource in the directory where the current binary located.
+        NSString * bin_cur = [[NSProcessInfo processInfo] arguments][0];
+        NSString * bin_dir = [bin_cur stringByDeletingLastPathComponent];
+
+        NSString * path_lib_default = [NSString pathWithComponents:@[bin_dir, [name stringByAppendingPathExtension:@"metallib"]]];
+        if ([[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
+            GGML_LOG_INFO("%s: found '%s'\n", __func__, [path_lib_default UTF8String]);
+
+            NSDictionary * atts = [[NSFileManager defaultManager] attributesOfItemAtPath:path_lib_default error:&error];
+            if (atts && atts[NSFileType] == NSFileTypeSymbolicLink) {
+                // Optionally, if this is a symlink, try to resolve it.
+                path_lib_default = [[NSFileManager defaultManager] destinationOfSymbolicLinkAtPath:path_lib_default error:&error];
+                if (path_lib_default && [path_lib_default length] > 0 && ![[path_lib_default substringToIndex:1] isEqualToString:@"/"]) {
+                    // It is a relative path, adding the binary directory as directory prefix.
+                    path_lib_default = [NSString pathWithComponents:@[bin_dir, path_lib_default]];
+                }
+                if (!path_lib_default || ![[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
+                    // Link to the resource could not be resolved.
+                    path_lib_default = nil;
+                } else {
+                    GGML_LOG_INFO("%s: symlink resolved '%s'\n", __func__, [path_lib_default UTF8String]);
+                }
+            }
+        } else {
+            // The resource couldn't be found in the binary's directory.
+            path_lib_default = nil;
+        }
+
+        path_lib = path_lib_default;
+    }
+
+    return path_lib;
+}
+
 ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
    id<MTLDevice> device = ggml_metal_device_get_obj(dev);

@@ -432,38 +491,7 @@ ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
    const int64_t t_start = ggml_time_us();

    NSError * error = nil;
-    NSString * path_lib = [bundle pathForResource:@"default" ofType:@"metallib"];
-    if (path_lib == nil) {
-        // Try to find the resource in the directory where the current binary located.
-        NSString * bin_cur = [[NSProcessInfo processInfo] arguments][0];
-        NSString * bin_dir = [bin_cur stringByDeletingLastPathComponent];
-
-        NSString * path_lib_default = [NSString pathWithComponents:@[bin_dir, @"default.metallib"]];
-        if ([[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
-            GGML_LOG_INFO("%s: found '%s'\n", __func__, [path_lib_default UTF8String]);
-
-            NSDictionary * atts = [[NSFileManager defaultManager] attributesOfItemAtPath:path_lib_default error:&error];
-            if (atts && atts[NSFileType] == NSFileTypeSymbolicLink) {
-                // Optionally, if this is a symlink, try to resolve it.
-                path_lib_default = [[NSFileManager defaultManager] destinationOfSymbolicLinkAtPath:path_lib_default error:&error];
-                if (path_lib_default && [path_lib_default length] > 0 && ![[path_lib_default substringToIndex:1] isEqualToString:@"/"]) {
-                    // It is a relative path, adding the binary directory as directory prefix.
-                    path_lib_default = [NSString pathWithComponents:@[bin_dir, path_lib_default]];
-                }
-                if (!path_lib_default || ![[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
-                    // Link to the resource could not be resolved.
-                    path_lib_default = nil;
-                } else {
-                    GGML_LOG_INFO("%s: symlink resolved '%s'\n", __func__, [path_lib_default UTF8String]);
-                }
-            }
-        } else {
-            // The resource couldn't be found in the binary's directory.
-            path_lib_default = nil;
-        }
-
-        path_lib = path_lib_default;
-    }
+    NSString * path_lib = ggml_metal_find_metallib(bundle, @"default");

    if (path_lib != nil) {
        // pre-compiled library found: a single combined default.metallib
@@ -478,6 +506,30 @@ ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
            return NULL;
        }

+        // the tensor API kernels are built into a separate metallib
+        if (ggml_metal_device_get_props(dev)->has_tensor) {
+            NSString * path_mm = ggml_metal_find_metallib(bundle, @"ggml-tensor");
+
+            id<MTLLibrary> lib_mm = nil;
+            if (path_mm != nil) {
+                lib_mm = [device newLibraryWithURL:[NSURL fileURLWithPath:path_mm] error:&error];
+                if (!lib_mm && error) {
+                    GGML_LOG_ERROR("%s: %s\n", __func__, [[error description] UTF8String]);
+                }
+            }
+
+            if (lib_mm) {
+                GGML_LOG_INFO("%s: loaded '%s'\n", __func__, [path_mm UTF8String]);
+
+                res->objs[GGML_METAL_LIB_MUL_MM] = [lib_mm retain];
+                res->override_fns                = [[NSSet setWithArray:[lib_mm functionNames]] retain];
+            } else {
+                GGML_LOG_INFO("%s: ggml-tensor.metallib not found - disabling the tensor API\n", __func__);
+
+                ggml_metal_device_disable_tensor(dev);
+            }
+        }
+
        GGML_LOG_INFO("%s: loaded in %.3f sec\n", __func__, (ggml_time_us() - t_start) / 1e6);
        return res;
    }
@@ -557,6 +609,7 @@ ggml_metal_library_t ggml_metal_library_init_from_source(ggml_metal_device_t dev

        MTLCompileOptions * options = [MTLCompileOptions new];
        options.preprocessorMacros = prep;
+        ggml_metal_compile_options_set_lang(options, ggml_metal_device_get_props(dev)->has_tensor);

        library = [device newLibraryWithSource:src options:options error:&error];
        if (error) {
@@ -615,6 +668,10 @@ void ggml_metal_library_free(ggml_metal_library_t lib) {
        [lib->fn_to_lib release];
    }

+    if (lib->override_fns) {
+        [lib->override_fns release];
+    }
+
    ggml_metal_pipelines_free(lib->pipelines);

    [lib->lock release];
@@ -676,7 +733,9 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_compile_pipeline(ggml_
        // route to the library that actually defines this kernel; fn_to_lib is
        // built from -[MTLLibrary functionNames] so it's always in sync
        int lib_idx = 0;
-        if (!lib->single_library) {
+        if (lib->override_fns && [lib->override_fns containsObject:base_func]) {
+            lib_idx = GGML_METAL_LIB_MUL_MM;
+        } else if (!lib->single_library) {
            NSNumber * idx = lib->fn_to_lib[base_func];
            if (!idx) {
                [lib->lock unlock];
@@ -997,6 +1056,7 @@ static const struct {
    DEV("M5 Pro",   GGML_METAL_DEVICE_M5_PRO),
    DEV("M5 Max",   GGML_METAL_DEVICE_M5_MAX),
    DEV("M5 Ultra", GGML_METAL_DEVICE_M5_ULTRA),
+    DEV("A18 Pro",  GGML_METAL_DEVICE_A18_PRO),
 #undef DEV
 };

@@ -1287,19 +1347,21 @@ ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) {
 void ggml_metal_device_free(ggml_metal_device_t dev) {
    assert(dev != NULL);

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

-    ggml_metal_library_free(dev->library);
-    dev->library = NULL;
+        ggml_metal_library_free(dev->library);
+        dev->library = NULL;

-    if (dev->mtl_queue) {
-        [dev->mtl_queue release];
-        dev->mtl_queue = nil;
-    }
+        if (dev->mtl_queue) {
+            [dev->mtl_queue release];
+            dev->mtl_queue = nil;
+        }

-    if (dev->mtl_device) {
-        [dev->mtl_device release];
-        dev->mtl_device = nil;
+        if (dev->mtl_device) {
+            [dev->mtl_device release];
+            dev->mtl_device = nil;
+        }
    }

    free(dev);
@@ -1387,12 +1449,14 @@ ggml_metal_event_t ggml_metal_device_event_init(ggml_metal_device_t dev) {
 }

 void ggml_metal_device_event_free(ggml_metal_device_t dev, ggml_metal_event_t ev) {
-    id<MTLSharedEvent> event = ev->obj;
-    [event release];
+    @autoreleasepool {
+        id<MTLSharedEvent> event = ev->obj;
+        [event release];

-    free(ev);
+        free(ev);

-    GGML_UNUSED(dev);
+        GGML_UNUSED(dev);
+    }
 }

 void ggml_metal_device_event_synchronize(ggml_metal_device_t dev, ggml_metal_event_t ev) {
@@ -1407,8 +1471,10 @@ void ggml_metal_device_event_synchronize(ggml_metal_device_t dev, ggml_metal_eve

 void ggml_metal_device_get_memory(ggml_metal_device_t dev, size_t * free, size_t * total) {
    if (@available(macOS 10.12, iOS 16.0, *)) {
-        *total = dev->mtl_device.recommendedMaxWorkingSetSize;
-        *free  = *total - dev->mtl_device.currentAllocatedSize;
+        *total     = dev->mtl_device.recommendedMaxWorkingSetSize;
+        size_t cur = dev->mtl_device.currentAllocatedSize;
+        // it's possible to allocate more than `recommendedMaxWorkingSetSize`
+        *free      = *total > cur ? *total - cur : 0;
    } else {
        *free = 0;
        *total = 0;
@@ -1538,6 +1604,12 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                        return true;
                    case GGML_TYPE_BF16:
                        return has_bfloat;
+                    case GGML_TYPE_Q4_0:
+                    case GGML_TYPE_Q4_1:
+                    case GGML_TYPE_Q5_0:
+                    case GGML_TYPE_Q5_1:
+                    case GGML_TYPE_Q8_0:
+                        return true;
                    default:
                        return false;
                }
@@ -1856,6 +1928,10 @@ const struct ggml_metal_device_props * ggml_metal_device_get_props(ggml_metal_de
    return &dev->props;
 }

+static void ggml_metal_device_disable_tensor(ggml_metal_device_t dev) {
+    dev->props.has_tensor = false;
+}
+
 //
 // device buffers
 //
@@ -2157,14 +2233,16 @@ ggml_metal_buffer_t ggml_metal_buffer_map(ggml_metal_device_t dev, void * ptr, s
 }

 void ggml_metal_buffer_free(ggml_metal_buffer_t buf) {
-    ggml_metal_device_rsets_rm(buf->dev, buf->rset);
+    @autoreleasepool {
+        ggml_metal_device_rsets_rm(buf->dev, buf->rset);

-    for (int i = 0; i < buf->n_buffers; i++) {
-        [buf->buffers[i].metal release];
+        for (int i = 0; i < buf->n_buffers; i++) {
+            [buf->buffers[i].metal release];
+        }
+
+        ggml_metal_buffer_rset_free(buf);
    }

-    ggml_metal_buffer_rset_free(buf);
-
    if (buf->is_shared && buf->owned) {
 #if TARGET_OS_OSX
        vm_deallocate((vm_map_t)mach_task_self(), (vm_address_t)buf->all_data, buf->all_size);
--- a/Show More
+++ b/Show More