ui : add cache (#28802 )

Signed-off-by: Adrien Gallouët <angt@huggingface.co>
server : allow model downloads at model limit fix issue #26809 (#28530 )
2026-09-12 17:46:56 +02:00 · 2026-09-12 16:09:46 +02:00 · 2026-09-12 11:50:35 +02:00 · 2026-09-12 11:26:53 +02:00 · 2026-09-12 11:23:54 +02:00 · 2026-09-12 10:15:08 +03:00
184 changed files with 12713 additions and 5879 deletions
@@ -137,19 +137,19 @@ runs:
      run: |
          mkdir -p "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4"
          choco install unzip -y
-          curl -O "https://packages.nvidia.com/bin-archive/pool/windows-x86_64/5B515474-7E78-11F1-8656-C51E4F4B317F/cccl-windows-x86_64-13.3.4.1.2-archive.zip"
-          curl -O "https://packages.nvidia.com/bin-archive/pool/windows-x86_64/5B515474-7E78-11F1-8656-C51E4F4B317F/cuda_crt-windows-x86_64-13.4.46-archive.zip"
-          curl -O "https://packages.nvidia.com/bin-archive/pool/windows-x86_64/5B515474-7E78-11F1-8656-C51E4F4B317F/cuda_nvcc-windows-x86_64-13.4.46-archive.zip"
-          curl -O "https://packages.nvidia.com/bin-archive/pool/windows-x86_64/5B515474-7E78-11F1-8656-C51E4F4B317F/libnvvm-windows-x86_64-13.4.46-archive.zip"
-          curl -O "https://packages.nvidia.com/bin-archive/pool/windows-arm64/5B515474-7E78-11F1-8656-C51E4F4B317F/cuda_cudart-windows-arm64-13.4.46-archive.zip"
-          curl -O "https://packages.nvidia.com/bin-archive/pool/windows-arm64/5B515474-7E78-11F1-8656-C51E4F4B317F/libcublas-windows-arm64-13.7.0.10-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cccl/windows-x86_64/cccl-windows-x86_64-13.3.4.2.1-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_crt/windows-x86_64/cuda_crt-windows-x86_64-13.4.59-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/windows-x86_64/cuda_nvcc-windows-x86_64-13.4.59-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/libnvvm/windows-x86_64/libnvvm-windows-x86_64-13.4.59-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/windows-arm64/cuda_cudart-windows-arm64-13.4.49-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/libcublas/windows-arm64/libcublas-windows-arm64-13.7.0.27-archive.zip"
          unzip '*.zip' -d "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4"
-          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cccl-windows-x86_64-13.3.4.1.2-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
-          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cuda_crt-windows-x86_64-13.4.46-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
-          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cuda_nvcc-windows-x86_64-13.4.46-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
-          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\libnvvm-windows-x86_64-13.4.46-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
-          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cuda_cudart-windows-arm64-13.4.46-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
-          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\libcublas-windows-arm64-13.7.0.10-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cccl-windows-x86_64-13.3.4.2.1-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cuda_crt-windows-x86_64-13.4.59-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cuda_nvcc-windows-x86_64-13.4.59-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\libnvvm-windows-x86_64-13.4.59-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\cuda_cudart-windows-arm64-13.4.49-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\libcublas-windows-arm64-13.7.0.27-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" /E /I /H /Y
          echo "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4\bin" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append
          echo "CUDA_PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
          echo "CUDA_PATH_V13_4=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
@@ -33,7 +33,7 @@ env:
  LLAMA_ARG_LOG_PREFIX: 1
  LLAMA_ARG_LOG_TIMESTAMPS: 1
  # TODO: fix failing tests on OpenVINO backend
-  CTEST_EXCLUDE: "test-llama-archs|^test-recurrent-state-|test-backend-ops|test-save-load-state"
+  CTEST_EXCLUDE: "test-llama-archs|^test-recurrent-state-|test-save-load-state"

 jobs:
  ubuntu-24-openvino:
@@ -106,83 +106,3 @@ jobs:
          wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories260K/stories260K.bin
          ./bin/llama-convert-llama2c-to-ggml --copy-vocab-from-model ./tok512.bin --llama2c-model stories260K.bin --llama2c-output-model stories260K.gguf
          ./bin/llama-completion -m stories260K.gguf -p "One day, Lily met a Shoggoth" -n 500 -c 256
-
-  ubuntu-riscv64-native-sanitizer:
-    runs-on: ubuntu-24.04-riscv
-
-    continue-on-error: true
-
-    strategy:
-      matrix:
-        sanitizer: [ADDRESS, THREAD, UNDEFINED]
-        build_type: [Debug]
-
-    steps:
-      - name: Install dependencies
-        run: |
-          # Set gcc-14 and g++-14 as the default compilers
-          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-14 100
-          sudo update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-14 100
-
-          git lfs install
-
-      - name: GCC version check
-        run: |
-          gcc --version
-          g++ --version
-
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      # note: sparing some ccache since these jobs run on dedicated runners that are not part of the organitzation
-      #- name: ccache
-      #  uses: ggml-org/ccache-action@v1.2.24
-      #  with:
-      #    key: riscv-ubuntu-native-sanitizer-${{ matrix.sanitizer }}-${{ matrix.build_type }}
-      #    evict-old-files: 1d
-      #    save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-
-      - name: Build
-        id: cmake_build
-        if: ${{ matrix.sanitizer != 'THREAD' }}
-        run: |
-          cmake -B build \
-            -DLLAMA_OPENSSL=OFF \
-            -DCMAKE_BUILD_TYPE=${{ matrix.build_type }} \
-            -DGGML_OPENMP=ON \
-            -DLLAMA_BUILD_EXAMPLES=ON \
-            -DLLAMA_BUILD_TOOLS=ON \
-            -DLLAMA_BUILD_TESTS=OFF \
-            -DCMAKE_C_COMPILER_LAUNCHER=ccache \
-            -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
-            -DLLAMA_SANITIZE_${{ matrix.sanitizer }}=ON \
-            -DCMAKE_C_COMPILER=riscv64-linux-gnu-gcc-14 \
-            -DCMAKE_CXX_COMPILER=riscv64-linux-gnu-g++-14
-
-          cmake --build build --config ${{ matrix.build_type }} -j $(nproc)
-
-      - name: Build (no OpenMP)
-        id: cmake_build_no_openmp
-        if: ${{ matrix.sanitizer == 'THREAD' }}
-        run: |
-          cmake -B build \
-            -DLLAMA_OPENSSL=OFF \
-            -DCMAKE_BUILD_TYPE=${{ matrix.build_type }} \
-            -DGGML_OPENMP=OFF \
-            -DLLAMA_BUILD_EXAMPLES=ON \
-            -DLLAMA_BUILD_TOOLS=ON \
-            -DLLAMA_BUILD_TESTS=OFF \
-            -DCMAKE_C_COMPILER_LAUNCHER=ccache \
-            -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
-            -DLLAMA_SANITIZE_${{ matrix.sanitizer }}=ON \
-            -DCMAKE_C_COMPILER=riscv64-linux-gnu-gcc-14 \
-            -DCMAKE_CXX_COMPILER=riscv64-linux-gnu-g++-14
-
-          cmake --build build --config ${{ matrix.build_type }} -j $(nproc)
-
-      - name: Test
-        id: cmake_test
-        run: |
-          cd build
-          ctest -L main --verbose --timeout 900
@@ -101,8 +101,6 @@ jobs:

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

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

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

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

  gpu-rocm:
    runs-on: [self-hosted, Linux, AMD]
@@ -130,10 +160,10 @@ jobs:
      - name: Dawn Dependency
        id: dawn-depends
        run: |
-          DAWN_VERSION="v20260317.182325"
+          DAWN_VERSION="v20260908.214631"
          DAWN_OWNER="google"
          DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-18eb229ef5f707c1464cc581252e7603c73a3ef0-ubuntu-latest-Release"
+          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-ubuntu-latest-Release"
          echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
          curl -L -o artifact.tar.gz \
            "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
@@ -216,10 +246,10 @@ jobs:
      - name: Dawn Dependency
        id: dawn-depends
        run: |
-          DAWN_VERSION="v20260317.182325"
+          DAWN_VERSION="v20260908.214631"
          DAWN_OWNER="google"
          DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-18eb229ef5f707c1464cc581252e7603c73a3ef0-macos-latest-Release"
+          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-macos-latest-Release"
          echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
          curl -L -o artifact.tar.gz \
            "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
@@ -365,7 +395,11 @@ jobs:
      - name: Test
        id: ggml-ci
        run: |
-          LLAMA_ARG_THREADS=$(nproc) GG_BUILD_HIGH_PERF=1 GG_BUILD_NO_BF16=1 GG_BUILD_EXTRA_TESTS_0=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
+          LLAMA_ARG_THREADS=$(nproc) \
+          GG_BUILD_HIGH_PERF=1 \
+          GG_BUILD_NO_BF16=1 \
+          GG_BUILD_EXTRA_TESTS_0=1 \
+          bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp

  cpu-arm64-graviton4-kleidiai:
    runs-on: ah-ubuntu_22_04-c8g_8x
@@ -404,6 +438,8 @@ jobs:
      - name: Test
        id: ggml-ci
        run: |
+          LLAMA_ARG_THREADS=$(nproc) \
          GG_BUILD_KLEIDIAI=1 \
          GG_BUILD_EXTRA_TESTS_0=1 \
-          bash ./ci/run.sh ./tmp/results ./tmp/mnt
+          GG_BUILD_HIGH_PERF=1 \
+          bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
@@ -164,9 +164,7 @@ jobs:
          export GGML_VK_VISIBLE_DEVICES=0
          export GGML_VK_DISABLE_F16=1
          export GGML_VK_DISABLE_COOPMAT=1
-          # 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
+          ctest -L main --verbose --timeout 900

  windows:
    runs-on: windows-2025
@@ -68,7 +68,7 @@ jobs:

      - name: Fetch emdawnwebgpu
        run: |
-          DAWN_TAG="v20260317.182325"
+          DAWN_TAG="v20260908.214631"
          EMDAWN_PKG="emdawnwebgpu_pkg-${DAWN_TAG}.zip"
          echo "Downloading ${EMDAWN_PKG}"
          curl -L -o emdawn.zip \
@@ -77,10 +77,10 @@ jobs:
      - name: Dawn Dependency
        id: dawn-depends
        run: |
-          DAWN_VERSION="v20260317.182325"
+          DAWN_VERSION="v20260908.214631"
          DAWN_OWNER="google"
          DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-18eb229ef5f707c1464cc581252e7603c73a3ef0-macos-latest-Release"
+          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-macos-latest-Release"
          echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
          curl -L -o artifact.tar.gz \
            "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
@@ -147,10 +147,10 @@ jobs:
        id: dawn-depends
        run: |
          sudo apt-get install -y libxrandr-dev libxinerama-dev libxcursor-dev mesa-common-dev libx11-xcb-dev libxi-dev
-          DAWN_VERSION="v20260317.182325"
+          DAWN_VERSION="v20260908.214631"
          DAWN_OWNER="google"
          DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-18eb229ef5f707c1464cc581252e7603c73a3ef0-ubuntu-latest-Release"
+          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-ubuntu-latest-Release"
          echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
          curl -L -o artifact.tar.gz \
            "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
@@ -190,6 +190,4 @@ jobs:
        id: cmake_test
        run: |
          cd build
-          # 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
+          ctest -L main --verbose --timeout 900
@@ -0,0 +1,67 @@
+name: Fusion
+
+on:
+  workflow_dispatch: # allows manual triggering
+  push:
+    branches:
+      - master
+    paths: [
+      '.github/workflows/fusion.yml',
+      'ggml/**',
+      'tests/fusion/**',
+      'tests/test-fusion.cpp'
+    ]
+
+  pull_request:
+    types: [opened, synchronize, reopened]
+    paths: [
+      '.github/workflows/fusion.yml',
+      'ggml/**',
+      'tests/fusion/**',
+      'tests/test-fusion.cpp'
+    ]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
+  cancel-in-progress: true
+
+env:
+  GGML_NLOOP: 3
+  GGML_N_THREADS: 1
+  LLAMA_ARG_LOG_COLORS: 1
+  LLAMA_ARG_LOG_PREFIX: 1
+  LLAMA_ARG_LOG_TIMESTAMPS: 1
+
+jobs:
+  # TODO: add jobs for other backends as they adopt the fusion debug API
+  metal:
+    runs-on: [self-hosted, macOS, ARM64]
+
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+
+      - name: Build
+        id: cmake_build
+        run: |
+          cmake -B build \
+            -DCMAKE_BUILD_TYPE=Release \
+            -DLLAMA_FATAL_WARNINGS=ON \
+            -DLLAMA_OPENSSL=OFF \
+            -DGGML_SCHED_NO_REALLOC=ON \
+            -DGGML_BLAS=OFF \
+            -DGGML_METAL=ON
+          time cmake --build build --config Release --target test-llama-archs -j $(sysctl -n hw.logicalcpu)
+          time cmake --build build --config Release --target test-fusion -j $(sysctl -n hw.logicalcpu)
+
+      - name: Generate models
+        id: generate_models
+        run: |
+          rm -rf build-ci-models && mkdir -p build-ci-models
+          ./build/bin/test-llama-archs -o build-ci-models
+
+      - name: Test fusion
+        id: test_fusion
+        run: |
+          ./build/bin/test-fusion --models build-ci-models --device MTL0 --check tests/fusion/MTL.csv
@@ -1717,7 +1717,7 @@ jobs:
            - [Windows arm64 (OpenCL Adreno)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-opencl-adreno-arm64.zip)
            - [Windows x64 (CUDA 12)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-cuda-12.4-x64.zip) - [CUDA 12.4 DLLs](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-bin-win-cuda-12.4-x64.zip)
            - [Windows x64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-cuda-13.3-x64.zip) - [CUDA 13.3 DLLs](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-bin-win-cuda-13.3-x64.zip)
-            - [Windows arm64 (CUDA 13) (preview)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-cuda-13.4-arm64.zip) - [CUDA 13.4 DLLs](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-bin-win-cuda-13.4-arm64.zip)
+            - [Windows arm64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-cuda-13.4-arm64.zip) - [CUDA 13.4 DLLs](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-bin-win-cuda-13.4-arm64.zip)
            - [Windows x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-vulkan-x64.zip)
            - [Windows x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-openvino-${{ needs.windows-openvino.outputs.openvino_version }}-x64.zip)
            - [Windows x64 (SYCL)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-sycl-x64.zip)
@@ -32,6 +32,8 @@ on:
    ]

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

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

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

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

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

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

      - name: Install Python dependencies
        run: |
@@ -190,7 +190,7 @@ if [ ! -z ${GG_BUILD_OPENVINO} ]; then
    CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_OPENVINO=ON"

    # TODO: fix failing tests on OpenVINO backend
-    CTEST_EXTRA="-E test-llama-archs|^test-recurrent-state-|test-backend-ops|test-save-load-state"
+    CTEST_EXTRA="-E test-llama-archs|^test-recurrent-state-|test-save-load-state"
 fi

 ## helpers
@@ -250,7 +250,7 @@ function gg_run_ctest_debug {
    (cmake -G "${CMAKE_GENERATOR}" -DCMAKE_BUILD_TYPE=Debug ${CMAKE_EXTRA} .. ) 2>&1 | tee -a $OUT/${ci}-cmake.log
    (time cmake --build . --config Debug -j$(nproc)) 2>&1 | tee -a $OUT/${ci}-make.log

-    (time ctest -C Debug --output-on-failure -L main -E "test-opt|test-backend-ops|test-llama-archs" ${CTEST_EXTRA}) 2>&1 | tee -a $OUT/${ci}-ctest.log
+    (time ctest -C Debug --output-on-failure -L main -E "test-opt|test-llama-archs" ${CTEST_EXTRA}) 2>&1 | tee -a $OUT/${ci}-ctest.log

    set +e
 }
@@ -334,6 +334,35 @@ function gg_sum_test_llama_archs_tensor_split {
    gg_printf '```\n'
 }

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

 function gg_run_test_scripts {
@@ -739,25 +768,43 @@ function gg_check_build_requirements {
    fi
 }

-function gg_run_test_backend_ops_cpu {
+function gg_run_test_backend_ops {
    cd ${SRC}

    cd build-ci-release

    set -e

-    (time ./bin/test-backend-ops -b CPU ) 2>&1 | tee -a $OUT/${ci}-test-backend-ops-cpu.log
+    local args_extra="-j $(nproc)"
+
+    # TODO: fix multi-threaded for ROCm
+    #       https://github.com/ggml-org/llama.cpp/actions/runs/34576278519/job/103297889044?pr=28740#step:3:4865
+    if [ ! -z ${GG_BUILD_ROCM} ]; then
+        args_extra=""
+    fi
+
+    # TODO: MoltenVK bug?
+    #       https://github.com/ggml-org/llama.cpp/actions/runs/34611260059/job/103302413736?pr=28740#step:3:5897
+    if [ ! -z "${GG_BUILD_VULKAN}" ] && [ "$(uname -s)" = "Darwin" ]; then
+        args_extra=""
+    fi
+
+    if [ ! -z ${GG_BUILD_HIGH_PERF} ]; then
+        (time ./bin/test-backend-ops ${args_extra} -b CPU) 2>&1 | tee -a $OUT/${ci}-test-backend-ops.log
+    else
+        (time ./bin/test-backend-ops ${args_extra}       ) 2>&1 | tee -a $OUT/${ci}-test-backend-ops.log
+    fi

    set +e
 }

-function gg_sum_test_backend_ops_cpu {
+function gg_sum_test_backend_ops {
    gg_printf '### %s\n\n' "${ci}"

-    gg_printf 'Runs test-backend-ops for CPU backend\n'
+    gg_printf 'Runs test-backend-ops\n'
    gg_printf '- status: %s\n' "$(cat $OUT/${ci}.exit)"
    gg_printf '```\n'
-    gg_printf '%s\n' "$(cat $OUT/${ci}-test-backend-ops-cpu.log)"
+    gg_printf '%s\n' "$(cat $OUT/${ci}-test-backend-ops.log)"
    gg_printf '```\n'
    gg_printf '\n'
 }
@@ -790,11 +837,10 @@ ret=0
 test $ret -eq 0 && gg_run ctest_debug
 test $ret -eq 0 && gg_run ctest_release

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

-if [ ! -z ${GG_BUILD_HIGH_PERF} ]; then
-    test $ret -eq 0 && gg_run test_backend_ops_cpu
-fi
+test $ret -eq 0 && gg_run test_llama_archs_models
+test $ret -eq 0 && gg_run test_llama_archs_tensor_split

 if [ -z ${GG_BUILD_LOW_PERF} ]; then
    test $ret -eq 0 && gg_run embd_bge_small
@@ -134,6 +134,8 @@ set_target_properties(${TARGET} PROPERTIES
 target_include_directories(${TARGET} PUBLIC .)
 target_link_libraries     (${TARGET} PUBLIC vendor::nlohmann vendor::sheredom)
 target_compile_features   (${TARGET} PUBLIC cxx_std_17)
+target_precompile_headers (${TARGET} PRIVATE common.h)
+target_precompile_headers (${TARGET} PRIVATE chat.h)

 if (LLAMA_SUBPROCESS)
    target_compile_definitions(${TARGET} PUBLIC LLAMA_SUBPROCESS)
@@ -872,17 +872,6 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
                    arg.c_str(), e.what(), opt.to_string().c_str()));
            }
        }
-
-        // TODO: remove this check after deprecating --mmap|mlock|dio
-        auto has_arg = [&](std::initializer_list<const char *> names) {
-            return std::any_of(names.begin(), names.end(), [&](const char * name) {
-                return seen_args.count(name);
-            });
-        };
-        if (has_arg({"-lm", "--load-mode"}) &&
-            has_arg({"--mlock", "--mmap", "--no-mmap", "-dio", "--direct-io", "-ndio", "--no-direct-io"})) {
-            LOG_WRN("DEPRECATED: `--load-mode` and `--mlock`/`--mmap`/`--direct-io` should not be combined; only the last flag on the command line will take effect\n");
-        }
    };

    // parse all CLI args now, so that -hf is available below for remote preset resolution
@@ -2694,32 +2683,6 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            }
        ).set_env("LLAMA_ARG_RPC"));
    }
-    add_opt(common_arg(
-        {"--mlock"},
-        "DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing",
-        [](common_params & params) {
-            LOG_WRN("DEPRECATED: --mlock is deprecated. use --load-mode mlock instead\n");
-            params.load_mode = LLAMA_LOAD_MODE_MLOCK;
-        }
-    ).set_env("LLAMA_ARG_MLOCK"));
-    add_opt(common_arg(
-        {"--mmap"},
-        {"--no-mmap"},
-        "DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)",
-        [](common_params & params, bool value) {
-            LOG_WRN("DEPRECATED: --mmap and --no-mmap are deprecated. use --load-mode mmap instead\n");
-            params.load_mode = value ? LLAMA_LOAD_MODE_MMAP : LLAMA_LOAD_MODE_NONE;
-        }
-    ).set_env("LLAMA_ARG_MMAP"));
-    add_opt(common_arg(
-        {"-dio", "--direct-io"},
-        {"-ndio", "--no-direct-io"},
-        "DEPRECATED in favor of `--load-mode`: use DirectIO if available",
-        [](common_params & params, bool value) {
-            LOG_WRN("DEPRECATED: --direct-io and --no-direct-io are deprecated. use --load-mode dio instead\n");
-            params.load_mode = value ? LLAMA_LOAD_MODE_DIRECT_IO : LLAMA_LOAD_MODE_NONE;
-        }
-    ).set_env("LLAMA_ARG_DIO"));
    add_opt(common_arg(
        {"-lm", "--load-mode"}, "MODE",
        "model loading mode (default: auto)\n"
@@ -296,7 +296,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
            drafting[seq_id] = true;
            common_sampler_reset(smpls[seq_id].get());

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

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

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

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

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

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

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

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

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

            const int32_t n_draft = params.n_max;

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

        const int32_t n_tokens = batch_in.n_tokens;

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

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

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

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

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

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

-    llama_pos   n_past;
+    llama_pos   pos0;
    llama_token id_last;

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

@@ -188,7 +188,7 @@ llama_memory_breakdown_print: |   - Host               |                  439 =
 Op test for MUL_MAT:

 ```
-~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --hex-hostbuf 0 --devices HTP0:0 -- test-backend-ops -b HTP0:0 -o MUL_MAT
+~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --devices HTP0:0 -- test-backend-ops -b HTP0:0 -o MUL_MAT
 ...
 Backend 2/3: HTP0:0
 Device description: Hexagon
@@ -213,14 +213,109 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v
 | llama 1B Q4_0  | 729.75 MiB | 1.24 B | HTP        |  99 |       4 |     128 |    0 |  tg64 |  51.54 ± 1.13 |
 ```

+## Multi-Device Execution Modes
+
+The Hexagon backend supports multiple execution and partitioning modes to accommodate different model sizes, memory
+constraints, and single- or multi-NPU hardware topologies:
+
+### 1. Single-Device Mode with Dynamic Buffer Mapping
+
+Runs the model on a single NPU session (e.g. `HTP0` or `HTP0:0`).
+
+A single NPU session provides ~3.5GB of available virtual address space. For models larger than 3.5GB, the backend
+automatically maps and unmaps weight buffers during graph execution. This allows large models to run on a single NPU
+without manual configuration:
+
+```bash
+./scripts/snapdragon/run.py --target adb --devices HTP0:0 -- \
+    llama-cli -m models/Llama-3.2-3B-Instruct-Q4_0.gguf -ngl 99 -p "Hello"
+```
+
+### 2. Layer-Split Mode across Virtual Sessions (`HTP0,HTP1,...` or `HTP0:0,HTP0:1,...`)
+
+Partitions model layers at load time across multiple virtual sessions hosted on a single physical NPU.
+
+Each virtual session acts as an independent backend device from llama.cpp's perspective (similar to multiple GPUs).
+Because layers are permanently distributed across sessions, each session's allocated weights remain within its private 3.5GB
+address space window, eliminating runtime buffer re-mapping overhead.
+
+Here is an example of running the GPT-OSS-20B model on a Snapdragon device using 4 virtual sessions on a single NPU:
+
+```bash
+./scripts/snapdragon/run.py --target adb \
+    --devices HTP0:0,HTP0:1,HTP0:2,HTP0:3 -- \
+    llama-cli --load-mode none -m /data/local/tmp/gguf/gpt-oss-20b-Q4_0.gguf -t 4 \
+    --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 -no-cnv -f surfing.txt
+```
+
+Log output snippet:
+
+```
+...
+llama_model_loader: - type  f32:  289 tensors
+llama_model_loader: - type q4_0:   96 tensors
+llama_model_loader: - type q8_0:    2 tensors
+llama_model_loader: - type mxfp4:  72 tensors
+...
+load_tensors: offloaded 25/25 layers to GPU
+load_tensors:          CPU model buffer size =  1182.09 MiB
+load_tensors:       HTP0:1 model buffer size =  2512.58 MiB
+load_tensors:       HTP0:3 model buffer size =  2093.83 MiB
+load_tensors:       HTP0:0 model buffer size =  2931.34 MiB
+load_tensors:       HTP0:2 model buffer size =  2512.58 MiB
+...
+llama_perf_context_print: prompt eval time =    3843.67 ms /   197 tokens ( 19.51 ms per token, 51.25 tokens per second)
+llama_perf_context_print:        eval time =    1686.13 ms /    31 runs   ( 54.39 ms per token, 18.39 tokens per second)
+llama_perf_context_print:       total time =    6266.30 ms /   228 tokens
+llama_memory_breakdown_print: | memory breakdown [MiB] | total   free    self   model   context   compute    unaccounted |
+llama_memory_breakdown_print: |   - HTP0:0 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:1 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:2 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - HTP0:3 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
+llama_memory_breakdown_print: |   - Host               |                 1476 =  1208 +     105 +     162                |
+```
+
+### 3. Tensor-Split Mode across Physical Devices (`HTP0:0,HTP1:0,...`)
+
+Distributes model tensors across distinct physical NPU hardware cores using llama.cpp's tensor parallelism
+(`--split-mode tensor`).
+
+Tensors are partitioned across physical NPUs for parallel execution (proportions are distributed equally by default without
+needing an explicit `--tensor-split` option):
+
+```bash
+./scripts/snapdragon/run.py --target adb \
+    --devices HTP0:0,HTP1:0 -- \
+    llama-cli -m models/Llama-3.2-3B-Instruct-Q4_0.gguf --split-mode tensor -ngl 99 -p "Hello"
+```
+
+### 4. Row-Split Multi-Device Mode via Device Grouping (`HTP0[0-1]`)
+
+Groups multiple physical NPU cores into a single logical device using bracket notation (`HTP0[0-1]` or `HTP0[0,1]`).
+
+Unlike host-level tensor-splitting, row-splitting is executed entirely inside the Hexagon backend:
+
+```bash
+./scripts/snapdragon/run.py --target adb \
+    --devices 'HTP0[0-1]' -- \
+    llama-cli -m models/Llama-3.2-3B-Instruct-Q4_0.gguf -ngl 99 -p "Hello"
+```
+
+You can also combine row-splitting with layer-splitting across multiple grouped devices (e.g. `--devices 'HTP0[0-1],HTP1[2-3]'`
+on 4 physical NPUs, or `--devices 'HTP0[0-1:0],HTP1[0-1:1]'` on 2 physical NPUs using virtual sessions 0 and 1).
+
 ## Environment variables

 - `GGML_HEXAGON_DEVICES` (default: not set, defaults to HTP0 session)
-  Controls which NPU devices and sessions to allocate. Can be configured as:
-  - A single integer `N`: Allocates `N` sessions named `HTP0`, `HTP1`, ..., `HTP<N-1>` (behaves identically to `GGML_HEXAGON_NDEV=N`).
-  - A comma-separated list of device names in `HTP<physical_idx>:<virtual_idx>` format (or legacy `HTP<idx>` format). For example, `HTP0:0,HTP0:1` creates two virtual
-    sessions on the first physical NPU (useful for memory limits). `HTP0:0,HTP1:0` allocates one session on each of the two physical NPUs
-    on a dual-NPU device.
+  Controls which NPU devices and sessions to allocate. Configurable via `--devices` in `run.py`:
+  - `N` (single integer): Allocates `N` virtual sessions named `HTP0`, `HTP1`, ..., `HTP<N-1>` on physical NPU 0.
+  - `HTP<phys>:<virt>,...`: Comma-separated list of individual devices specifying physical and virtual index:
+    - `HTP0:0,HTP0:1`: Two virtual sessions on physical NPU 0 (layer-split on single NPU).
+    - `HTP0:0,HTP1:0`: One session on physical NPU 0 and one on physical NPU 1 (tensor-split across physical cores).
+  - `HTP<name>[<phys_spec>]`: Device grouping syntax for row-split multi-device execution:
+    - `HTP0[0-1]`: A single logical device `HTP0` that groups physical cores 0 and 1.
+    - `HTP0[0-1],HTP1[2-3]`: Two layer-split devices across 4 physical NPUs (cores 0-1 and 2-3).
+    - `HTP0[0-1:0],HTP1[0-1:1]`: Two layer-split devices across 2 physical NPUs using virtual sessions 0 and 1.

 - `GGML_HEXAGON_NDEV` (deprecated)
  Replaced by `GGML_HEXAGON_DEVICES`. Controls the number of virtual sessions to allocate on physical NPU `0`.
@@ -229,9 +324,8 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v
 - `GGML_HEXAGON_NHVX=0`
  Controls the number of HVX hardware threads to use. The default is all (actual number varies depending on the hardware version).

- `GGML_HEXAGON_HOSTBUF=1`
-  Controls whether the Hexagon backend allocates host buffers. By default, all buffers except for REPACK are host buffers.
-  This option is required for testing Ops that require REPACK buffers (MUL_MAT and MUL_MAT_ID).
+- `GGML_HEXAGON_HOSTBUF=1` (default: 0, disabled)
+  Enables allocating host buffers for debugging. By default, host buffers are disabled.

 - `GGML_HEXAGON_VERBOSE=1`
  Enables verbose logging of Ops from the backend. Example output:
@@ -246,23 +340,26 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v
  ```

 - `GGML_HEXAGON_PROFILE=1`
-  Enables Op profiling:
+  Enables Op profiling (configurable via `--hex-profile` in `run.py`):

-  - `1` Basic profile with per-op `usecs` and `cycles` counters
-  - `2` Extended profile with per-op `usecs`, `cycles` and default PMU counter data
-  - `0x1,...,0x8` Extended profile with per-op `usecs`, `cycles` and custom PMU counter data
+  - `1`: Basic profile with per-op `usecs` and `cycles` counters
+  - `2`: Extended profile with per-op `usecs`, `cycles` and default PMU counter data
+  - `0x1,...,0x8`: Extended profile with per-op `usecs`, `cycles` and custom PMU counter data

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

-      `GGML_HEXAGON_PROFILE=1 ./scripts/snapdragon/run.py --target adb -- llama-cli ... |& ./scripts/snapdragon/ggml-hexagon-profile.py -`
+  ```bash
+  ./scripts/snapdragon/run.py --target adb --hex-profile 1 -- llama-cli ... |& \
+      ./scripts/snapdragon/ggml-hexagon-profile.py -
+  ```

 - `GGML_HEXAGON_OPFILTER=regex`
-  Allows filtering (disabling) Ops that match the regex pattern:
+  Filters (disables) Ops matching the regex pattern (configurable via `--hex-opfilter` in `run.py`):

-  Examples:
-
-      `GGML_HEXAGON_OPFILTER="FLASH_ATTN_EXT" ./scripts/snapdragon/run.py --target adb -- llama-cli ...` - Disable Flash Attention on Hexagon (falls back to CPU or GPU)
-      `GGML_HEXAGON_OPFILTER="ADD\|SUB" ./scripts/snapdragon/run.py --target adb -- llama-cli ...` - Disable ADD and SUB on Hexagon (fall back to CPU or GPU)
+  ```bash
+  # Disable Flash Attention on Hexagon (falls back to CPU or GPU)
+  ./scripts/snapdragon/run.py --target adb --hex-opfilter "FLASH_ATTN_EXT" -- llama-cli ...

+  # Disable ADD and SUB on Hexagon (fall back to CPU or GPU)
+  ./scripts/snapdragon/run.py --target adb --hex-opfilter "ADD|SUB" -- llama-cli ...
+  ```
@@ -2,16 +2,16 @@

 ## Backend libraries

-The Hexagon backend consist of two parts:
+The Hexagon backend consists of two parts:

  - `libggml-hexagon`
-    This is the regular CPU-side GGML backend library, either shared or statically linked
+    This is the regular CPU-side GGML backend library, either shared or statically linked.

  - `libggml-htp-vNN`
    This is the NPU-side (HTP stands for Hexagon Tensor Processor) shared library that contains the Op dispatcher and kernels.
    The correct library is selected automatically at runtime based on the HW version.

-Here is an example of the build artifacts
+Here is an example of the build artifacts:

 ```
 ~/src/llama.cpp$ ls -l pkg-adb/llama.cpp/lib/libggml*
@@ -26,75 +26,307 @@ pkg-adb/llama.cpp/lib/libggml-htp-v81.so

 ## Memory buffers

-Hexagon NPU backend takes advantage of the Snapdragon's unified memory model where all buffers are fully accessible by the CPU and GPU.
-The NPU does have a dedicated tightly-coupled memory called VTCM but that memory is used only for intermediate data (e.g. dynamically
-quantized tensors) or temporary data (chunks of the weight tensors fetched via DMA).
-
-Please note that currently the Hexagon backend does not implement SET/GET_ROWS Ops because there is no advantage in offloading those
-to the NPU at this point.
-
-The backend does allocates non-host buffers for the tensors with datatypes that require repacking: Q4_0, Q8_0, MXFP4.
-From the MMU perspective these buffers are still regular buffers (normal access by the CPU) they are marked as non-host simply to force
-the repacking.
+The Hexagon NPU backend takes advantage of Snapdragon unified memory where all DDR buffers are accessible by CPU, GPU, and NPU.
+The NPU has dedicated tightly-coupled memory called VTCM (Vector Tightly-Coupled Memory). VTCM is used for intermediate data (such as
+dynamically quantized activations) and streaming buffers (chunks of weight and activation tensors fetched via DMA).

 ## Large model handling

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

-To support running models larger than 3.5GB on a single device, the Hexagon backend dynamically maps and unmaps execution buffers
-during the graph execution cycle to stay within the Process Domain window. This enables large models to run successfully on a single
-NPU device.
+To support running models larger than 3.5GB on a single device, the Hexagon backend dynamically maps and unmaps buffers:
+- Buffers are allocated in shared DDR (RPCMEM) via file descriptors (`fastrpc_mmap` using `FASTRPC_MAP_FD_DELAYED`).
+- Pinned buffers (such as KV cache and active compute buffers) remain mapped throughout execution.
+- Inactive weight buffers are dynamically mapped into the NPU session via `HAP_mmap()` during batch buffer preparation
+  (`prep_op_bufs()` in `htp/main.c`) and unmapped via `htp_iface_munmap()` when no longer needed by the active batch.
+- This dynamic sliding window allows a single NPU session to execute models that exceed the 3.5GB window.

-Alternatively, users can choose to use standard llama.cpp/GGML layer-splitting mode to partition and split the model across
-multiple Hexagon devices or virtual sessions (which behave like multiple GPUs from the offload and splitting perspective).
+Alternatively, users can partition and split the model across multiple virtual sessions or physical NPUs using layer-splitting,
+tensor-splitting, or row-splitting modes. For user-facing execution modes and examples, see the
+[Snapdragon user guide](README.md#multi-device-execution-modes).

-Here is an example of running GPT-OSS-20B model on a Snapdragon device using 4 virtual sessions on a single NPU (physical index 0).
+## Op and Kernel Development Guidelines
+
+Writing high-performance operators for Hexagon requires following specific guidelines.
+
+### DDR -> DMA -> VTCM Execution Pipeline
+
+- Strongly prefer the `DDR -> DMA -> VTCM -> compute (HVX/HMX) -> VTCM -> DMA -> DDR` data flow.
+- Direct HVX reads/writes from/to DDR are less efficient and should only be used as a fallback.
+- The DMA queue is a strict FIFO where operations must be pushed and popped in strict order.
+- Follow the pipelined multi-buffering sequence properly (typically 2x to 16x buffering) so every push has a corresponding pop:
+
+  1. In the prologue, push initial DDR -> VTCM transfers to prime the pipeline.
+  2. In the loop body, wait for buffer N via DMA pop, launch HVX/HMX compute on buffer N, push VTCM -> DDR writeback of result N,
+     and push DDR -> VTCM prefetch of buffer N+2.
+  3. In the epilogue, pop all remaining in-flight transfers to drain the pipeline.
+
+- Because every push must be matched by a pop, `dma_queue_flush()` is not required when the pipeline sequence is followed
+  properly. Flushing is only used in rare exceptions where a batch of operations is pushed without individual pops.
+- Use the DMA queue interface from [`dma-queue.h`](../../../ggml/src/ggml-hexagon/htp/dma-queue.h)
+  (`dma_queue_push_ddr_to_vtcm`, `dma_queue_pop`, `dma_queue_push_vtcm_to_ddr`).
+  See [`cumsum-ops.c`](../../../ggml/src/ggml-hexagon/htp/cumsum-ops.c) and
+  [`act-ops.c`](../../../ggml/src/ggml-hexagon/htp/act-ops.c) for reference implementations.
+
+### Avoid Scalar Reads and Writes to VTCM
+
+- Access VTCM data using DMA transfers or HVX/HMX vector instructions rather than scalar reads and writes.
+
+### Avoid Scalar Division in Inner Loops
+
+- Hexagon cores do not have hardware division instructions.
+- For recurring divisions across iterations or threads, use `fastdiv` from
+  [`hex-fastdiv.h`](../../../ggml/src/ggml-hexagon/htp/hex-fastdiv.h) with precomputed divisors (such as
+  `octx->ctx->mdev.count_div` or `octx->n_threads_div`).
+- Do not call `init_fastdiv_values()` for single-use divisions; use standard compiler division (`/`) instead.
+
+### Host-Side Precomputation via `kernel_params`
+
+- Precompute tensor shapes, strides, scale conversions, tiling layouts, and validation checks on the host CPU during graph
+  preparation in [`ggml-hexagon.cpp`](../../../ggml/src/ggml-hexagon/ggml-hexagon.cpp).
+- Pack precomputed parameters into the operator's fixed `kernel_params` structure in `htp_op_node` (such as
+  `htp_mm_kernel_params`, `htp_unary_kernel_params`, `htp_fa_kernel_params`, `htp_get_rows_kernel_params`).
+- The NPU executes directly using `octx->kernel_params` without redundant runtime metadata extraction or validation.
+- **Strict Host-Kernel Alignment**:
+  - Verify that parameters calculated by the host CPU are strictly honored by the NPU kernel.
+  - Ensure the kernel does not ignore host-computed fields (for example, falling back to `octx->n_threads` instead of
+    using `kparams->n_threads`, or ignoring precomputed `tasks_per_thread` and chunk counts).
+  - Both human developers and coding agents must audit both sides of the interface: ensure fields populated in `kernel_params`
+    in [`ggml-hexagon.cpp`](../../../ggml/src/ggml-hexagon/ggml-hexagon.cpp) are actively and consistently utilized by the
+    corresponding operator entry point and worker threads in `htp/*-ops.c`.
+
+### Tracing Instrumentation
+
+- All kernels must include trace events for performance profiling and timeline visualization in Perfetto
+  ([`hex-profile.h`](../../../ggml/src/ggml-hexagon/htp/hex-profile.h)).
+- Surround compute sections with `htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) info)` and
+  `htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) info)`.
+- Use specific event types for major phases:
+  - `HTP_TRACE_EVT_HVX_COMP`: Vector compute execution.
+  - `HTP_TRACE_EVT_DMA`: DMA transfer wait or poll cycles.
+  - `HTP_TRACE_EVT_FENCE`: Multi-device fence barrier synchronization.
+  - `HTP_TRACE_EVT_L2FLUSH`: L2 cache cleaning operations.
+- Pass meaningful progress metrics (such as row index, chunk index, or token index) in the 16-bit `info` parameter.
+
+### Work Queue and Threading
+
+- Distribute parallel work across NPU worker threads using the thread pool work queue:
+
+  ```c
+  work_queue_run(ctx->work_queue, worker_func, &op_ctx, n_threads);
+  ```
+
+- Keep worker functions independent and re-entrant. Worker threads should only operate on their designated chunk of rows or elements.
+
+### Avoid Redundant Defensive NULL Checks
+
+- Do not add defensive NULL checks or assertions for internal framework pointers or required graph operands and outputs.
+  Internal pointers include `ctx`, `octx`, local context structs like `*ctx`, `kparams`, and worker callback `data`.
+- These pointers are architectural invariants during kernel execution and host-side graph preparation.
+  Graph compute receives allocated nodes with valid required `node->src[N]` and `node->data` pointers.
+- Do not turn an invariant violation into an unsupported operation or missed fusion.
+  Checks such as `if (!octx || !octx->ctx)` clutter the code, obscure intent, and hide upstream errors.
+- **Distinction**: `octx->src[N]` pointers *can* be NULL by design and must be checked when optional.
+  Examples include attention masks, optional bias or weights in fused kernels, and frequency factors.
+
+### Multiline Macro Formatting
+
+- Keep trailing backslashes in multiline `#define` macros cleanly aligned to a consistent column.
+- Avoid trailing whitespace after macro backslashes.
+- Use [`scripts/snapdragon/ggml-hexagon-align-macros.py`](../../../scripts/snapdragon/ggml-hexagon-align-macros.py) to inspect, diff,
+  or automatically align macro definitions across Hexagon kernel sources:
+
+  ```bash
+  # Check for misaligned macros
+  python3 scripts/snapdragon/ggml-hexagon-align-macros.py ggml/src/ggml-hexagon/htp/
+
+  # Fix misaligned macros in-place
+  python3 scripts/snapdragon/ggml-hexagon-align-macros.py --fix ggml/src/ggml-hexagon/htp/
+  ```
+
+## Multi-Device Partitioning (mdev)
+
+Multi-device (mdev) mode enables row-level tensor parallel execution across multiple physical NPU cores or virtual NPU
+sessions.
+
+### 128-Byte Cache Line Alignment
+
+- Shared tensor buffers reside in DDR (RPCMEM) with a 128-byte cache line granularity
+  (`HEX_L2_LINE_SIZE` = 128 bytes, `HTP_TENSOR_MDEV_LINE_SIZE`).
+- **Rule**: Multi-device work partitions must align destination write regions to 128-byte cache line boundaries so distinct
+  devices never share or overwrite the same cache line.
+
+### Partitioning Helpers in `htp-tensor.h`
+
+Common partitioning logic is factored into reusable inline helpers in
+[`htp-tensor.h`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h):
+
+1. [`htp_tensor_mdev_rows_per_chunk`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L67):
+   Determines the minimum number of rows per chunk so that the chunk byte size is a multiple of 128 bytes:
+
+   ```
+   rows_per_chunk = 128 / hex_gcd_u32(row_size, 128)
+   ```
+
+   If row stride `nb[1]` is already a multiple of 128 bytes, `rows_per_chunk = 1`.
+   Returns `false` if the tensor cannot be safely row-partitioned (such as unaligned base pointer, permuted layout,
+   or non-128-byte aligned outer strides).
+
+2. [`htp_tensor_mdev_partition`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L94):
+   Calculates the per-device work range `struct htp_tensor_mdev_range { uint32_t start; uint32_t count; }` given
+   `total_units`, `units_per_chunk`, `mdev_idx`, `mdev_count`, and the precomputed `mdev_count_div`.
+   Handles chunk distribution across devices, assigns remainder units to the last device, and automatically triggers
+   single-device fallback when partitioning is unsafe.
+
+### Row-Partitioned Operators
+
+For row-wise operators
+(such as activations in [`act-ops.c`](../../../ggml/src/ggml-hexagon/htp/act-ops.c),
+binary ops in [`binary-ops.c`](../../../ggml/src/ggml-hexagon/htp/binary-ops.c),
+unary ops in [`unary-ops.c`](../../../ggml/src/ggml-hexagon/htp/unary-ops.c), and
+sameshape copies in [`cpy-ops.c`](../../../ggml/src/ggml-hexagon/htp/cpy-ops.c)):
+
+```c
+const uint32_t total_rows   = ne01 * ne02 * ne03;
+const size_t   dst_row_size = dst->ne[0] * elem_size;
+
+uint32_t row_start = 0;
+uint32_t nrows     = total_rows;
+
+if (octx->ctx->mdev.count > 1) {
+    uint32_t rows_per_chunk = 0;
+    htp_tensor_mdev_rows_per_chunk(dst, elem_size, (uint32_t) dst_row_size, &rows_per_chunk);
+    const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+        total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+    row_start = range.start;
+    nrows     = range.count;
+}
+
+if (nrows == 0) {
+    return HTP_STATUS_OK;
+}
+```
+
+### Element-Partitioned Operators
+
+For flat element-wise operations (such as reshape copies in
+[`cpy-ops.c`](../../../ggml/src/ggml-hexagon/htp/cpy-ops.c)):
+- Partition total linear elements N = ne0 * ne1 * ne2 * ne3 in 128-byte cache line chunks (`elems_per_line = (elem_size == 4) ? 32 : 64`).
+- Requires strict 1D contiguity:
+  [`htp_tensor_is_contiguous(dst, elem_size)`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L28)
+  and 128-byte aligned destination pointer
+  [`htp_tensor_mdev_data_aligned(dst)`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L47).
+- If contiguous and aligned, pass `elems_per_line` to
+  [`htp_tensor_mdev_partition`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L94);
+  otherwise pass 0 to trigger Device 0 fallback.
+
+### Single-Device Fallback (Device 0)
+
+- Fallback to Device 0 (`mdev.idx == 0`) when partitioning would cause cache line tearing or when work cannot be evenly distributed.
+- Triggers:
+  1. Destination tensor cannot be safely partitioned (`rows_per_chunk == 0` or non-contiguous/unaligned buffer).
+  2. Total aligned chunks < `mdev_count`.
+- Device 0 processes the entire tensor `[0, total_units)`.
+- Devices 1 ... N-1 receive `count = 0` and return `HTP_STATUS_OK` immediately.
+
+### Flatten Outer Dimensions Globally
+
+- **Never partition solely on `ne01` (dimension 1).**
+- Partitioning only on `ne01` repeats the device boundary across every 2D slice (`ne02`, `ne03`). If each 2D slice is small,
+  false sharing occurs repeatedly throughout the tensor.
+- Always flatten outer dimensions globally: `total_rows = ne01 * ne02 * ne03` and partition once across the combined row space.
+
+### Stateless Starting Coordinates
+
+- Do not use incremental state variables across slices that assume the thread or device starts at index 0.
+- Precompute starting multidimensional coordinates at `r = row_start` (or `e = elem_start`) once using `fastdiv`.
+- In inner loops, step base pointers directly (`ptr += stride`) or reset/wrap coordinates explicitly (`if (++i01 == ne01) { ... }`).
+
+### Clean Range Encapsulation
+
+- Initialize single-device default ranges at declaration:
+
+  ```c
+  uint32_t row_start = 0;
+  uint32_t nrows     = total_rows;
+  ```
+
+- Encapsulate all multi-device logic inside `if (octx->ctx->mdev.count > 1)`. If the block is omitted or compiled out,
+  the operator runs standard single-device execution untouched.
+- Do not propagate `mdev_` prefixes to worker functions or context structs. Worker threads are device-agnostic and
+  should only receive standard range parameters (`ctx.row_start`, `ctx.nrows`).
+- In worker threads, calculate row intervals using standard arithmetic:
+
+  ```c
+  const uint32_t ir0 = ctx->row_start + dr * ith;
+  const uint32_t ir1 = MIN(ir0 + dr, ctx->row_start + ctx->nrows);
+  ```
+
+  In single-device mode (`row_start == 0`), this naturally simplifies to `dr * ith` and `MIN(ir0 + dr, ctx->nrows)` with zero overhead.
+
+## Multi-Device Synchronization
+
+Multi-device execution synchronizes worker sessions across devices using explicit barriers and tensor cache flushing.
+
+### Synchronization Fence Protocol
+
+Multi-device execution synchronizes worker sessions through atomic fence slots and barriers defined in
+[`htp-fence.h`](../../../ggml/src/ggml-hexagon/htp/htp-fence.h):

 ```
-~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --devices HTP0:0,HTP0:1,HTP0:2,HTP0:3 -- llama-cli --load-mode none -m /data/local/tmp/gguf/gpt-oss-20b-Q4_0.gguf -t 4 --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 -no-cnv -f surfing.txt
-...
-llama_model_loader: - type  f32:  289 tensors
-llama_model_loader: - type q4_0:   96 tensors
-llama_model_loader: - type q8_0:    2 tensors
-llama_model_loader: - type mxfp4:  72 tensors
-...
-load_tensors: offloaded 25/25 layers to GPU
-load_tensors:          CPU model buffer size =  1182.09 MiB
-load_tensors:       HTP0:1 model buffer size =  2512.58 MiB
-load_tensors:       HTP0:3 model buffer size =  2093.83 MiB
-load_tensors:       HTP0:0 model buffer size =  2931.34 MiB
-load_tensors:       HTP0:2 model buffer size =  2512.58 MiB
-...
-llama_context: n_ctx_per_seq (8192) < n_ctx_train (131072) -- the full capacity of the model will not be utilized
-llama_context:        CPU  output buffer size =     0.77 MiB
-llama_kv_cache_iswa: creating non-SWA KV cache, size = 8192 cells
-llama_kv_cache:     HTP0:1 KV buffer size =    25.50 MiB
-llama_kv_cache:     HTP0:3 KV buffer size =    25.50 MiB
-llama_kv_cache:     HTP0:0 KV buffer size =    25.50 MiB
-llama_kv_cache:     HTP0:2 KV buffer size =    25.50 MiB
-llama_kv_cache: size =  102.00 MiB (  8192 cells,  12 layers,  1/1 seqs), K (q8_0):   51.00 MiB, V (q8_0):   51.00 MiB
-llama_kv_cache_iswa: creating     SWA KV cache, size = 256 cells
-llama_kv_cache:     HTP0:1 KV buffer size =     0.80 MiB
-llama_kv_cache:     HTP0:3 KV buffer size =     0.53 MiB
-llama_kv_cache:     HTP0:0 KV buffer size =     1.06 MiB
-llama_kv_cache:     HTP0:2 KV buffer size =     0.80 MiB
-llama_kv_cache: size =    3.19 MiB (   256 cells,  12 layers,  1/1 seqs), K (q8_0):    1.59 MiB, V (q8_0):    1.59 MiB
-llama_context:     HTP0:0 compute buffer size =    16.06 MiB
-llama_context:     HTP0:1 compute buffer size =    16.06 MiB
-llama_context:     HTP0:2 compute buffer size =    16.06 MiB
-llama_context:     HTP0:3 compute buffer size =    16.06 MiB
-llama_context:        CPU compute buffer size =    98.19 MiB
-...
-llama_perf_context_print: prompt eval time =    3843.67 ms /   197 tokens ( 19.51 ms per token, 51.25 tokens per second)
-llama_perf_context_print:        eval time =    1686.13 ms /    31 runs   ( 54.39 ms per token, 18.39 tokens per second)
-llama_perf_context_print:       total time =    6266.30 ms /   228 tokens
-llama_perf_context_print:    graphs reused =         30
-llama_memory_breakdown_print: | memory breakdown [MiB] | total   free    self   model   context   compute    unaccounted |
-llama_memory_breakdown_print: |   - HTP0:0 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP0:1 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP0:2 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - HTP0:3 (Hexagon)   |  2048 = 2048 + (   0 =     0 +       0 +       0) +           0 |
-llama_memory_breakdown_print: |   - Host               |                 1476 =  1208 +     105 +     162                |
+[NPU Session 0]                              [NPU Session 1]
+       |                                            |
+  (Input Prep)                                 (Input Prep)
+       |                                            |
+  Pre-Op Barrier ----------------------------- Pre-Op Barrier
+  (mdev_sync_fence)                            (mdev_sync_fence)
+       |                                            |
+  Kernel Execution                             Kernel Execution
+  (Output Slice 0)                             (Output Slice 1)
+       |                                            |
+  Tensor Cache Flush                           Tensor Cache Flush
+  (htp_tensor_flush_all)                       (htp_tensor_flush_all)
+       |                                            |
+  Post-Op/Batch Barrier ---------------------- Post-Op/Batch Barrier
+  (htp_mdev_group_barrier)                     (htp_mdev_group_barrier)
+       |                                            |
+  Return Response to Host                      Return Response to Host
 ```
+
+### Atomic Fence Slots and Cache Invalidation
+
+- Fence synchronization operates on dedicated RPCMEM shared memory mapped across all participating sessions (`ctx->mdev.fence_base`).
+- Each device owns a dedicated 128-byte cache-line aligned fence slot:
+
+  ```c
+  atomic_uint * my_fence = htp_mdev_fence_slot(fence_base, mdev_idx);
+  ```
+
+- **Writing to fence ([`htp_fence_write`](../../../ggml/src/ggml-hexagon/htp/htp-fence.h#L18))**:
+  Stores `seq` and `status`, issues a `syncht` thread synchronization barrier, and flushes/invalidates the line
+  using `Q6_dccleaninva_A(fence)`.
+- **Reading from peer fence ([`htp_fence_read`](../../../ggml/src/ggml-hexagon/htp/htp-fence.h#L26))**:
+  Executes `Q6_dccleaninva_A(fence)` and `syncht` before reading atomic values to ensure fresh data from DDR.
+
+### Deterministic Monotonic Sequence Numbers
+
+- Barrier fences use monotonically increasing sequence numbers:
+
+  ```c
+  const uint32_t seq = ++ctx->mdev.fence_seq;
+  ```
+
+- Comparing sequence numbers with signed arithmetic `(int32_t)(peer_seq - seq) >= 0` prevents race conditions or
+  misaligned barrier arrivals across iterations.
+- If any peer reports an error status (`peer_status > HTP_STATUS_OK`), the barrier propagates the error and unblocks immediately.
+
+### Tensor Cache Flush and Pipeline Completion
+
+- In the kernel, ensure all pushed DMA operations have been popped in strict FIFO order to drain the queue.
+- Use [`htp_tensor_flush_all()`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h) to flush specific dirty tensors back to DDR:
+  - [`htp_tensor_flush_all()`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h) flushes only modified tensor address ranges,
+    ensuring peer devices and the host CPU observe consistent data in DDR.
+- Never signal completion before all DMA transfers are drained and dirty tensor flushes have completed.
+
@@ -0,0 +1,122 @@
+## Build profiling
+This page is a working document for analyzing the current build and try to
+identify ways to improve the build time.
+
+### Requirements
+The profiling script requires clang to be used as the compiler tool chain and
+also requires that ClangBuildAnalyzer is installed.
+
+Mac:
+```console
+brew install clang-build-analyzer
+```
+
+Linux:
+```console
+git clone https://github.com/aras-p/ClangBuildAnalyzer.git
+cd ClangBuildAnalyzer
+cmake -B build -DCMAKE_BUILD_TYPE=Release
+cmake --build build -j$(nproc)
+sudo cp build/ClangBuildAnalyzer /usr/local/bin/
+```
+
+Windows: install LLVM/clang and Ninja (e.g. via the
+[LLVM releases page](https://github.com/llvm/llvm-project/releases) and
+`winget install Ninja-build.Ninja`), then build ClangBuildAnalyzer the same
+way as on Linux:
+```console
+git clone https://github.com/aras-p/ClangBuildAnalyzer.git
+cd ClangBuildAnalyzer
+cmake -B build -G Ninja -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_BUILD_TYPE=Release
+cmake --build build --config Release
+```
+Then add `ClangBuildAnalyzer\build` to `PATH`.
+
+### Usage
+Mac/Linux:
+```console
+$ ./scripts/build-profile.sh
+```
+
+Windows:
+```console
+> .\scripts\build-profile.ps1
+```
+
+Both accept `--full`/`-Full` (include Server, Tools, and Tests) and a jobs
+override (`-jN` / `-Jobs N`).
+
+Note: on Windows, `cmake` defaults to the Visual Studio generator, which
+ignores `CMAKE_C_COMPILER`/`CMAKE_CXX_COMPILER` and silently falls back to
+MSVC. `build-profile.ps1` passes `-G Ninja` so clang is actually used, this
+is required on ARM64.
+
+### Linux (Ubuntu 24.04)
+
+Environment:
+- Clang:     18.1.3 (Ubuntu clang version 18.1.3 (1ubuntu1))
+- libstdc++: GCC 13.3.0 (Ubuntu 13.3.0-6ubuntu2~24.04.1)
+- Target:    x86_64-pc-linux-gnu
+
+```console
+------------------------+-----+------------+------------+------------+
+| Build                  | TUs | Frontend   | Backend    | Total      |
+------------------------+-----+------------+------------+------------+
+| Minimal, master        | 249 |   468.2 s  |   270.3 s  |   738.5 s  |
+| Minimal, with PCH      | 253 |   177.1 s  |   265.8 s  |   442.9 s  |
+| Full,    master        | 396 |   811.0 s  |   692.2 s  | 1,503.2 s  |
+| Full,    with PCH      | 405 |   380.0 s  |   664.7 s  | 1,044.7 s  |
+| Full,    with PCH + UB | 264 |   357.7 s  |   635.7 s  |   993.4 s  |
+------------------------+-----+------------+------------+------------+
+
+PCH  = precompiled header.
+Full = includes building Server, Tools, and Tests.
+UB   = unity build for models
+```
+Note that the number of translation units (TUs) increases when using precompiled
+headers — each PCH target adds one extra TU for the precompilation step itself.
+
+### Mac (Apple M3)
+
+Environment:
+- Clang:  Apple clang version 17.0.0 (clang-1700.3.19.1)
+- libc++: ships with Apple clang 17.0.0 (Xcode toolchain)
+- Target: arm64-apple-macosx15.6
+
+```console
+------------------------+-----+------------+------------+------------+
+| Build                  | TUs | Frontend   | Backend    | Total      |
+------------------------+-----+------------+------------+------------+
+| Minimal, master        | 256 |   154.5 s  |    94.8 s  |   249.3 s  |
+| Minimal, with PCH      | 261 |    65.9 s  |    90.0 s  |   155.9 s  |
+| Full,    master        | 407 |   265.7 s  |   209.7 s  |   475.4 s  |
+| Full,    with PCH      | 414 |   154.6 s  |   197.5 s  |   352.1 s  |
+| Full,    with PCH + UB | 274 |   143.0 s  |   192.2 s  |   335.2 s  |
+------------------------+-----+------------+------------+------------+
+
+PCH = precompiled header.
+Full = includes building Server, Tools, and Tests.
+UB   = unity build for models
+```
+
+### Windows (ARM64)
+
+Environment:
+- Clang:  clang version 22.1.8 (LLVM, `C:\Program Files\LLVM`)
+- STL:    MSVC STL (Visual Studio 2022 Build Tools 14.44.35207)
+- Target: aarch64-pc-windows-msvc
+
+```console
+------------------------+-----+------------+------------+------------+
+| Build                  | TUs | Frontend   | Backend    | Total      |
+------------------------+-----+------------+------------+------------+
+| Minimal, master        | 249 |   159.4 s  |    82.2 s  |   241.6 s  |
+| Full,    master        | 373 |   337.2 s  |   167.4 s  |   504.6 s  |
+| Minimal, with PCH + UB | 113 |    62.3 s  |    82.4 s  |   144.7 s  |
+| Full,    with PCH + UB | 240 |   233.0 s  |   185.1 s  |   418.1 s  |
+------------------------+-----+------------+------------+------------+
+
+PCH = precompiled header.
+Full = includes building Server, Tools, and Tests.
+UB   = unity build for models
+```
@@ -243,6 +243,7 @@ IBM VXE/VXE2 SIMD acceleration depends on the BLAS implementation. It is strongl
 | FP32       | ✅           | ✅    | ❓     |
 | FP16       | ✅           | ✅    | ❓     |
 | BF16       | ✅           | ✅    | ❓     |
+| Q1_0       | ✅           | ❓    | ❓     |
 | Q4_0       | ✅           | ❓    | ❓     |
 | Q4_1       | ✅           | ❓    | ❓     |
 | MXFP4      | ✅           | ❓    | ❓     |
@@ -272,4 +273,4 @@ IBM VXE/VXE2 SIMD acceleration depends on the BLAS implementation. It is strongl
 -   🚫 - acceleration unavailable, will still run using scalar implementation
 -   ❓ - acceleration unknown, please contribute if you can test it yourself

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

 ## MUSA

@@ -805,7 +806,7 @@ To read documentation for how to build on Android, [click here](./android.md)

 ## WebGPU

-The WebGPU backend relies on [Dawn](https://dawn.googlesource.com/dawn). Follow the instructions [here](https://dawn.googlesource.com/dawn/+/refs/heads/main/docs/quickstart-cmake.md) to install Dawn locally so that llama.cpp can find it using CMake. The current implementation is up-to-date with Dawn commit `18eb229`.
+The WebGPU backend relies on [Dawn](https://dawn.googlesource.com/dawn). Follow the instructions [here](https://dawn.googlesource.com/dawn/+/refs/heads/main/docs/quickstart-cmake.md) to install Dawn locally so that llama.cpp can find it using CMake. The current implementation is up-to-date with Dawn commit `94c3c9c`.

 In the llama.cpp directory, build with CMake:

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    {
@@ -0,0 +1,281 @@
+static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_gcn(ggml_type type, int J, bool fallback) {
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q1_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_0, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_1, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_1, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2,  64,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q8_0, 512, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q2_K, 512, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q3_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q4_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q4_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q4_K, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 3,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q5_K, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_Q6_K, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ1_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ2_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_XXS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ3_S, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_XS, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_IQ4_NL, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false);
+
+// ---------------------------------------------------------------------------------------------
+
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_MXFP4, 512, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false);
+
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2, 128,  40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+    CASE(GGML_TYPE_NVFP4, 256, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
+
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+}
@@ -171,7 +171,7 @@ void ggml_cuda_mul_mat_q(
            ne00, ne01, ne1, s01, ne11, s1,
            ne02, ne12, s02, s12, s2,
            ne03, ne13, s03, s13, s3,
-            ne1};
+            ne1, ne1};
        ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
        return;
    }
@@ -244,6 +244,13 @@ void ggml_cuda_mul_mat_q(
                                         ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
    const int64_t s13 = ne12*s12;

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

    ggml_cuda_mul_mat_q_switch_type(ctx, args, stream);
 }
@@ -218,6 +218,7 @@ struct ggml_cuda_mmq_config {
 #include "mmq-config-ampere.cuh"
 #include "mmq-config-blackwell.cuh"

+#include "mmq-config-gcn.cuh"
 #include "mmq-config-cdna.cuh"
 #include "mmq-config-rdna2.cuh"
 #include "mmq-config-rdna3.cuh"
@@ -228,6 +229,9 @@ struct ggml_cuda_mmq_config {

 static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type type, const int J, const bool fallback, const int cc) {
    if (GGML_CUDA_CC_IS_AMD(cc)) {
+        if (GGML_CUDA_CC_IS_GCN(cc)) {
+            return ggml_cuda_mmq_get_config_gcn(type, J, fallback);
+        }
        if (GGML_CUDA_CC_IS_CDNA(cc)) {
            return ggml_cuda_mmq_get_config_cdna(type, J, fallback);
        }
@@ -256,7 +260,9 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty

 static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback) {
 #ifdef GGML_USE_HIP
-#ifdef CDNA
+#ifdef GCN
+    return ggml_cuda_mmq_get_config_gcn(type, J, fallback);
+#elif defined(CDNA)
    return ggml_cuda_mmq_get_config_cdna(type, J, fallback);
 #elif defined(RDNA4)
    return ggml_cuda_mmq_get_config_rdna4(type, J, fallback);
@@ -1376,6 +1382,7 @@ struct mmq_args {
    int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst;
    int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
    int64_t ncols_max;
+    int64_t ncols_opt; // value to optimize the tile size against, launch grid still uses ncols_max
 };

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

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

        if (ntiles_x < ntiles_J_best) {
            J_best = J;
@@ -344,6 +344,12 @@ struct htp_opformat {
        } else if (htp_op_is_unary(node.opcode)) {
            const auto * kparams = (const struct htp_unary_kernel_params *) node.kernel_params;
            snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
+        } else if (node.opcode == HTP_OP_MDEV_GROUP && node.node) {
+            snprintf(str, max_size, "idx %d count %d", (int) node.node->op_params[0], (int) node.dst()->ne[1]);
+        } else if ((node.opcode == HTP_OP_FENCE || node.opcode == HTP_OP_CPY_FENCE) && node.node) {
+            snprintf(str, max_size, "seq 0x%x", (uint32_t) node.node->op_params[0]);
+        } else if (node.opcode == HTP_OP_ALLREDUCE && node.node) {
+            snprintf(str, max_size, "seq 0x%x -> 0x%x", (uint32_t) node.node->op_params[0], (uint32_t) node.node->op_params[1]);
        } else {
            snprintf(str, max_size, "----");
        }
@@ -3,7 +3,6 @@
 #pragma clang diagnostic ignored "-Wunused-but-set-variable"

 #include <HAP_farf.h>
-#include <HAP_perf.h>

 #include <math.h>
 #include <string.h>
@@ -15,7 +14,7 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
+#include "hex-common.h"
 #include "htp-tensor.h"
 #include "htp-vtcm.h"

@@ -80,6 +79,7 @@ struct htp_act_context {
    uint32_t                 block;
    uint32_t                 src0_nrows;
    uint32_t                 src0_nrows_per_thread;
+    uint32_t                 row_start;
    int                      nc;

    uint8_t *                vtcm_src0;
@@ -329,104 +329,104 @@ static void geglu_f32(const float * restrict src0,
    }
 }

-#define DEFINE_GLU_PER_THREAD(NAME, OP_STR, CORE_EXPR)                                                                 \
-    static void glu_##NAME##_f32_per_thread(unsigned int nth, unsigned int ith, void * data) {                         \
-        struct htp_act_context * actx = (struct htp_act_context *) data;                                               \
-        htp_act_preamble;                                                                                              \
-                                                                                                                       \
-        struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL;                          \
-                                                                                                                       \
-        size_t src0_row_size = actx->src0_row_size;                                                                    \
-        size_t src1_row_size = actx->src1_row_size;                                                                    \
-        size_t dst_row_size  = actx->dst_row_size;                                                                     \
-                                                                                                                       \
-        size_t src0_row_stride = actx->src0_row_stride;                                                                \
-        size_t src1_row_stride = actx->src1_row_stride;                                                                \
-                                                                                                                       \
-        const uint32_t src0_nrows            = actx->src0_nrows;                                                       \
-        const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;                                            \
-                                                                                                                       \
-        const uint32_t src0_start_row = src0_nrows_per_thread * ith;                                                   \
-        const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);                       \
-                                                                                                                       \
-        /* no work for this thread */                                                                                  \
-        if (src0_start_row >= src0_end_row) {                                                                          \
-            return;                                                                                                    \
-        }                                                                                                              \
-                                                                                                                       \
-        const uint8_t * restrict data_src0 = actx->data_src0;                                                          \
-        const uint8_t * restrict data_src1 = actx->data_src1;                                                          \
-        uint8_t * restrict data_dst        = actx->data_dst;                                                           \
-                                                                                                                       \
-        const size_t src0_row_size_aligned = actx->src0_row_size_aligned;                                              \
-        const size_t src1_row_size_aligned = actx->src1_row_size_aligned;                                              \
-        const size_t dst_row_size_aligned  = actx->dst_row_size_aligned;                                               \
-                                                                                                                       \
-        uint8_t * restrict src0_spad_data = actx->vtcm_src0 + (ith * actx->vtcm_src0_size_per_thread);                 \
-        uint8_t * restrict src1_spad_data = actx->vtcm_src1 + (ith * actx->vtcm_src1_size_per_thread);                 \
-        uint8_t * restrict dst_spad_data  = actx->vtcm_dst  + (ith * actx->vtcm_dst_size_per_thread);                  \
-                                                                                                                       \
-        size_t src0_spad_half_size = actx->src0_spad_half_size;                                                        \
-        size_t src1_spad_half_size = actx->src1_spad_half_size;                                                        \
-        size_t dst_spad_half_size  = actx->dst_spad_half_size;                                                         \
-                                                                                                                       \
-        const int BLOCK = actx->block;                                                                                 \
-        if (BLOCK == 0) {                                                                                              \
-            FARF(ERROR,                                                                                                \
-                 OP_STR                                                                                                \
-                 " : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n",      \
-                 actx->vtcm_src0_size_per_thread, src0_row_size_aligned);                                              \
-            return;                                                                                                    \
-        }                                                                                                              \
-                                                                                                                       \
-        dma_queue * dma_queue = actx->octx->ctx->dma[ith];                                                             \
-                                                                                                                       \
-        /* See discussion: https://github.com/ggml-org/llama.cpp/pull/18151#issuecomment-3678235379 */                 \
-        for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) { \
-            const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                 \
-                                                                                                                       \
-            /* Dummy DMA transation for sequencing (interleaving dst,src,dst,...) */                                   \
-            dma_queue_push_vtcm_to_ddr(dma_queue,                                                                      \
-                                       dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),        \
-                                       dst_row_size, dst_row_size_aligned, 0);                                         \
-                                                                                                                       \
-            dma_queue_push(                                                                                            \
-                dma_queue,                                                                                             \
-                dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_stride)),   \
-                src0_row_size_aligned, src0_row_stride, src0_row_size, block_size);                                    \
-            dma_queue_push(                                                                                            \
-                dma_queue,                                                                                             \
-                dma_make_ptr(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_stride)),   \
-                src1_row_size_aligned, src1_row_stride, src1_row_size, block_size);                                    \
-        }                                                                                                              \
-                                                                                                                       \
-        for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) {                                           \
-            const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                 \
-                                                                                                                       \
-            float * dst_spad  = (float *) dma_queue_pop(dma_queue).src;                                                \
-            float * src0_spad = (float *) dma_queue_pop(dma_queue).dst;                                                \
-            float * src1_spad = (float *) dma_queue_pop(dma_queue).dst;                                                \
-                                                                                                                       \
-            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                     \
-            CORE_EXPR;                                                                                                 \
-            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                      \
-                                                                                                                       \
-            dma_queue_push_vtcm_to_ddr(dma_queue, dma_make_ptr(data_dst + (ir * dst_row_size), dst_spad),              \
-                                       dst_row_size, dst_row_size_aligned, block_size);                                \
-                                                                                                                       \
-            /* prefetch N+2 loop iteration if any */                                                                   \
-            const uint32_t pref_block = (ir + BLOCK * 2);                                                              \
-            if (pref_block < src0_end_row) {                                                                           \
-                const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block);                                \
-                dma_queue_push(dma_queue, dma_make_ptr(src0_spad, data_src0 + (pref_block * src0_row_stride)),         \
-                               src0_row_size_aligned, src0_row_stride, src0_row_size, pref_block_size);                \
-                dma_queue_push(dma_queue, dma_make_ptr(src1_spad, data_src1 + (pref_block * src1_row_stride)),         \
-                               src1_row_size_aligned, src1_row_stride, src1_row_size, pref_block_size);                \
-            }                                                                                                          \
-        }                                                                                                              \
-                                                                                                                       \
-        dma_queue_flush(dma_queue);                                                                                    \
-                                                                                                                       \
+#define DEFINE_GLU_PER_THREAD(NAME, OP_STR, CORE_EXPR)                                                                   \
+    static void glu_##NAME##_f32_per_thread(unsigned int nth, unsigned int ith, void * data) {                           \
+        struct htp_act_context * actx = (struct htp_act_context *) data;                                                 \
+        htp_act_preamble;                                                                                                \
+                                                                                                                         \
+        struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL;                            \
+                                                                                                                         \
+        size_t src0_row_size = actx->src0_row_size;                                                                      \
+        size_t src1_row_size = actx->src1_row_size;                                                                      \
+        size_t dst_row_size  = actx->dst_row_size;                                                                       \
+                                                                                                                         \
+        size_t src0_row_stride = actx->src0_row_stride;                                                                  \
+        size_t src1_row_stride = actx->src1_row_stride;                                                                  \
+                                                                                                                         \
+        const uint32_t src0_nrows            = actx->src0_nrows;                                                         \
+        const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;                                              \
+                                                                                                                         \
+        const uint32_t src0_start_row = actx->row_start + src0_nrows_per_thread * ith;                                   \
+        const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, actx->row_start + src0_nrows);       \
+                                                                                                                         \
+        /* no work for this thread */                                                                                    \
+        if (src0_start_row >= src0_end_row) {                                                                            \
+            return;                                                                                                      \
+        }                                                                                                                \
+                                                                                                                         \
+        const uint8_t * restrict data_src0 = actx->data_src0;                                                            \
+        const uint8_t * restrict data_src1 = actx->data_src1;                                                            \
+        uint8_t * restrict data_dst        = actx->data_dst;                                                             \
+                                                                                                                         \
+        const size_t src0_row_size_aligned = actx->src0_row_size_aligned;                                                \
+        const size_t src1_row_size_aligned = actx->src1_row_size_aligned;                                                \
+        const size_t dst_row_size_aligned  = actx->dst_row_size_aligned;                                                 \
+                                                                                                                         \
+        uint8_t * restrict src0_spad_data = actx->vtcm_src0 + (ith * actx->vtcm_src0_size_per_thread);                   \
+        uint8_t * restrict src1_spad_data = actx->vtcm_src1 + (ith * actx->vtcm_src1_size_per_thread);                   \
+        uint8_t * restrict dst_spad_data  = actx->vtcm_dst  + (ith * actx->vtcm_dst_size_per_thread);                    \
+                                                                                                                         \
+        size_t src0_spad_half_size = actx->src0_spad_half_size;                                                          \
+        size_t src1_spad_half_size = actx->src1_spad_half_size;                                                          \
+        size_t dst_spad_half_size  = actx->dst_spad_half_size;                                                           \
+                                                                                                                         \
+        const int BLOCK = actx->block;                                                                                   \
+        if (BLOCK == 0) {                                                                                                \
+            FARF(ERROR,                                                                                                  \
+                 OP_STR                                                                                                  \
+                 " : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n",        \
+                 actx->vtcm_src0_size_per_thread, src0_row_size_aligned);                                                \
+            return;                                                                                                      \
+        }                                                                                                                \
+                                                                                                                         \
+        dma_queue * dma_queue = actx->octx->ctx->dma[ith];                                                               \
+                                                                                                                         \
+        /* See discussion: https://github.com/ggml-org/llama.cpp/pull/18151#issuecomment-3678235379 */                   \
+        for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) {   \
+            const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                   \
+                                                                                                                         \
+            /* Dummy DMA transation for sequencing (interleaving dst,src,dst,...) */                                     \
+            dma_queue_push_vtcm_to_ddr(dma_queue,                                                                        \
+                                       dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),          \
+                                       dst_row_size, dst_row_size_aligned, 0);                                           \
+                                                                                                                         \
+            dma_queue_push(                                                                                              \
+                dma_queue,                                                                                               \
+                dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_stride)),     \
+                src0_row_size_aligned, src0_row_stride, src0_row_size, block_size);                                      \
+            dma_queue_push(                                                                                              \
+                dma_queue,                                                                                               \
+                dma_make_ptr(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_stride)),     \
+                src1_row_size_aligned, src1_row_stride, src1_row_size, block_size);                                      \
+        }                                                                                                                \
+                                                                                                                         \
+        for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) {                                             \
+            const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                   \
+                                                                                                                         \
+            float * dst_spad  = (float *) dma_queue_pop(dma_queue).src;                                                  \
+            float * src0_spad = (float *) dma_queue_pop(dma_queue).dst;                                                  \
+            float * src1_spad = (float *) dma_queue_pop(dma_queue).dst;                                                  \
+                                                                                                                         \
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                       \
+            CORE_EXPR;                                                                                                   \
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                        \
+                                                                                                                         \
+            dma_queue_push_vtcm_to_ddr(dma_queue, dma_make_ptr(data_dst + (ir * dst_row_size), dst_spad),                \
+                                       dst_row_size, dst_row_size_aligned, block_size);                                  \
+                                                                                                                         \
+            /* prefetch N+2 loop iteration if any */                                                                     \
+            const uint32_t pref_block = (ir + BLOCK * 2);                                                                \
+            if (pref_block < src0_end_row) {                                                                             \
+                const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block);                                  \
+                dma_queue_push(dma_queue, dma_make_ptr(src0_spad, data_src0 + (pref_block * src0_row_stride)),           \
+                               src0_row_size_aligned, src0_row_stride, src0_row_size, pref_block_size);                  \
+                dma_queue_push(dma_queue, dma_make_ptr(src1_spad, data_src1 + (pref_block * src1_row_stride)),           \
+                               src1_row_size_aligned, src1_row_stride, src1_row_size, pref_block_size);                  \
+            }                                                                                                            \
+        }                                                                                                                \
+                                                                                                                         \
+        dma_queue_flush(dma_queue);                                                                                      \
+                                                                                                                         \
    }

 DEFINE_GLU_PER_THREAD(swiglu, "swiglu-f32", swiglu_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
@@ -473,14 +473,30 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
    }

    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
-    const uint32_t n_threads  = MIN(octx->n_threads, src0_nrows);
+    const size_t dst_row_size = dst->ne[0] * SIZEOF_FP32;
+
+    uint32_t row_start = 0;
+    uint32_t nrows     = src0_nrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;

    // row_size   = bytes of useful data per row (what the kernel touches / what DMA copies).
    // row_stride = bytes between successive rows in DDR (may exceed row_size for non-contig src).
-    const size_t nc_bytes    = dst->ne[0] * SIZEOF_FP32;
-    const size_t src0_row_size = nc_bytes;
-    const size_t src1_row_size = nc_bytes;
-    const size_t dst_row_size  = nc_bytes;
+    const size_t nc_bytes        = dst_row_size;
+    const size_t src0_row_size   = nc_bytes;
+    const size_t src1_row_size   = nc_bytes;
    const size_t src0_row_stride = src0->nb[1];
    const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];

@@ -518,7 +534,7 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
    struct htp_act_context actx;
    actx.octx = octx;

-    actx.src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads;
+    actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);

    actx.src0_row_size = src0_row_size;
    actx.src1_row_size = src1_row_size;
@@ -545,7 +561,8 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
    actx.dst_spad_half_size  = L.dst_bytes_per_thread / 2;

    actx.block = actx.src0_spad_half_size / actx.src0_row_size_aligned;
-    actx.src0_nrows = src0_nrows;
+    actx.src0_nrows = nrows;
+    actx.row_start  = row_start;

    actx.nc = dst->ne[0];

@@ -570,7 +587,7 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
    actx.data_src1 = data_src1;
    actx.data_dst  = (uint8_t *) dst->data;

-    worker_pool_run_func(octx->ctx->worker_pool, act_op_func, &actx, n_threads);
+    work_queue_run(octx->ctx->work_queue, act_op_func, &actx, n_threads);
    return HTP_STATUS_OK;
 }

@@ -17,6 +17,7 @@
 #include "hex-dma.h"
 #include "hex-profile.h"
 #include "allreduce-ops.h"
+#include "htp-fence.h"

 struct htp_allreduce_context {
    struct htp_ops_context * octx;
@@ -242,7 +243,42 @@ DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f32,       float,  hvx_add_f32_aaa, 1, 0)
 DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f16, __fp16, hvx_add_f16_aaa, 1, 1)
 DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f32, float,  hvx_add_f32_aaa, 1, 1)

+static int validate_allreduce(
+    struct htp_ops_context * octx,
+    const struct htp_allreduce_kernel_params * kparams,
+    uint32_t n_ranks
+) {
+    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    if (kparams->vtcm_size_per_thread <= 0 || kparams->vtcm_size <= 0) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
+    const size_t n_vtcm_buffers = htp_allreduce_vtcm_buffer_count(
+        n_ranks, octx->n_threads, has_add, kparams->is_row_bcast != 0);
+    const size_t vtcm_size = n_vtcm_buffers * (size_t) kparams->vtcm_size_per_thread;
+    if (vtcm_size != (size_t) kparams->vtcm_size) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+    if (vtcm_size > octx->ctx->vtcm_size) {
+        return HTP_STATUS_VTCM_TOO_SMALL;
+    }
+
+    if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    return HTP_STATUS_OK;
+}
+
 int op_allreduce(struct htp_ops_context * octx) {
+    if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx > 0) {
+        return HTP_STATUS_OK;
+    }
+
    const struct htp_allreduce_kernel_params * kparams = (const struct htp_allreduce_kernel_params *) octx->kernel_params;
    const struct htp_tensor * dst = octx->dst;

@@ -253,38 +289,53 @@ int op_allreduce(struct htp_ops_context * octx) {
        return HTP_STATUS_INVAL_PARAMS;
    }

-    if (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    const uint32_t nelem = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
    const uint32_t fence_seq_entry = (uint32_t) octx->op_params[0];
    const uint32_t fence_seq_exit  = (uint32_t) octx->op_params[1];

+    const struct htp_tensor * my_sync = octx->src[n_ranks + rank];
+    atomic_uint * my_fence = (atomic_uint *) (uintptr_t) my_sync->data;
+
+    const int status = validate_allreduce(octx, kparams, n_ranks);
+    if (status != HTP_STATUS_OK) {
+        if (status == HTP_STATUS_NO_SUPPORT) {
+            FARF(ERROR, "ggml-hex: allreduce unsupported type %d : rank %u\n", dst->type, rank);
+        }
+        htp_fence_write(my_fence, fence_seq_exit, status);
+        return status;
+    }
+
+    const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
+    const uint32_t nelem = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
+
    // 1. Entry Barrier: Synchronize all ranks before reading
    struct htp_thread_trace * tr0 = &octx->ctx->trace[0];
    htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);

-    const struct htp_tensor * my_sync = octx->src[n_ranks + rank];
-    atomic_uint * my_fence = (atomic_uint *) my_sync->data;
-
-    atomic_store(&my_fence[0], fence_seq_entry);
-    asm volatile ("syncht" : : : "memory");
-    Q6_dccleaninva_A((void *) my_fence);
+    htp_fence_write(my_fence, fence_seq_entry, octx->status);

    for (uint32_t j = 0; j < n_ranks; j++) {
        if (j == rank) continue;
        const struct htp_tensor * peer_sync = octx->src[n_ranks + j];
-        atomic_uint * peer_fence = (atomic_uint *) peer_sync->data;
+        atomic_uint * peer_fence = (atomic_uint *) (uintptr_t) peer_sync->data;
        uint64_t spins = 0;
        while (1) {
-            Q6_dccleaninva_A((void *) peer_fence);
-            uint32_t val = atomic_load(&peer_fence[0]);
-            if (val == fence_seq_entry || val == fence_seq_exit) {
+            uint32_t peer_seq;
+            uint32_t peer_status;
+            htp_fence_read(peer_fence, &peer_seq, &peer_status);
+            if ((int32_t)(peer_seq - fence_seq_entry) >= 0) {
+                if (peer_status > HTP_STATUS_OK) {
+                    FARF(ERROR, "ggml-hex: allreduce entry peer %u failed with status %u\n", j, peer_status);
+                    htp_fence_write(my_fence, fence_seq_exit, peer_status);
+                    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
+                    return peer_status;
+                }
                break;
            }
            if (++spins > HTP_FENCE_TIMEOUT) {
-                FARF(ERROR, "ggml-hex: allreduce entry fence-wait TIMEOUT: rank %u waiting on %u (fence %p seq %u)\n", rank, j, peer_fence, fence_seq_entry);
+                FARF(ERROR, "ggml-hex: allreduce entry fence-wait TIMEOUT : rank %u waiting on %u fence %p seq 0x%x peer-seq 0x%x\n",
+                     rank, j, peer_fence, fence_seq_entry, peer_seq);
+                htp_fence_write(my_fence, fence_seq_exit, HTP_STATUS_INTERNAL_ERR);
+                htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
                return HTP_STATUS_INTERNAL_ERR;
            }
            hex_pause();
@@ -301,8 +352,6 @@ int op_allreduce(struct htp_ops_context * octx) {
        const uint32_t elems_per_thread     = (uint32_t) kparams->elems_per_thread;
        const uint32_t vtcm_size_per_thread = (uint32_t) kparams->vtcm_size_per_thread;

-        const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
-
        struct htp_allreduce_context actx;
        actx.octx                 = octx;
        actx.n_ranks              = n_ranks;
@@ -339,6 +388,8 @@ int op_allreduce(struct htp_ops_context * octx) {
                }
                break;
            default:
+                FARF(ERROR, "ggml-hex: allreduce unsupported kernel %d : rank %u\n", kparams->kernel_type, rank);
+                htp_fence_write(my_fence, fence_seq_exit, HTP_STATUS_NO_SUPPORT);
                return HTP_STATUS_NO_SUPPORT;
        }

@@ -368,23 +419,31 @@ int op_allreduce(struct htp_ops_context * octx) {
    // 4. Exit Barrier: Synchronize all ranks after writing
    htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);

-    atomic_store(&my_fence[0], fence_seq_exit);
-    asm volatile ("syncht" : : : "memory");
-    Q6_dccleaninva_A((void *) my_fence);
+    htp_fence_write(my_fence, fence_seq_exit, octx->status);

    for (uint32_t j = 0; j < n_ranks; j++) {
        if (j == rank) continue;
        const struct htp_tensor * peer_sync = octx->src[n_ranks + j];
-        atomic_uint * peer_fence = (atomic_uint *) peer_sync->data;
+        atomic_uint * peer_fence = (atomic_uint *) (uintptr_t) peer_sync->data;
        uint64_t spins = 0;
        while (1) {
-            Q6_dccleaninva_A((void *) peer_fence);
-            uint32_t val = atomic_load(&peer_fence[0]);
-            if (val == fence_seq_exit) {
+            uint32_t peer_seq;
+            uint32_t peer_status;
+            htp_fence_read(peer_fence, &peer_seq, &peer_status);
+            if ((int32_t)(peer_seq - fence_seq_exit) >= 0) {
+                if (peer_status > HTP_STATUS_OK) {
+                    FARF(ERROR, "ggml-hex: allreduce exit peer %u failed with status %u\n", j, peer_status);
+                    htp_fence_write(my_fence, fence_seq_exit, peer_status);
+                    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
+                    return peer_status;
+                }
                break;
            }
            if (++spins > HTP_FENCE_TIMEOUT) {
-                FARF(ERROR, "ggml-hex: allreduce exit fence-wait TIMEOUT: rank %u waiting on %u (fence %p seq %u)\n", rank, j, peer_fence, fence_seq_exit);
+                FARF(ERROR, "ggml-hex: allreduce exit fence-wait TIMEOUT : rank %u waiting on %u fence %p seq 0x%x peer-seq 0x%x\n",
+                     rank, j, peer_fence, fence_seq_exit, peer_seq);
+                htp_fence_write(my_fence, fence_seq_exit, HTP_STATUS_INTERNAL_ERR);
+                htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
                return HTP_STATUS_INTERNAL_ERR;
            }
            hex_pause();
@@ -394,5 +453,5 @@ int op_allreduce(struct htp_ops_context * octx) {

    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);

-    return HTP_STATUS_OK;
+    return octx->status;
 }
@@ -2,6 +2,8 @@
 #define ALLREDUCE_OPS_H

 #include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>

 #define HTP_ALLREDUCE_MAX_RANKS 4

@@ -15,6 +17,15 @@ enum htp_allreduce_kernel_type {
    HTP_ALLREDUCE_KERNEL_DMA_2D,
 };

+static inline size_t htp_allreduce_vtcm_buffer_count(
+    uint32_t n_ranks,
+    uint32_t n_threads,
+    bool has_add,
+    bool is_row_bcast
+) {
+    return (size_t) (n_ranks + 1) * n_threads + (has_add ? (is_row_bcast ? 1 : n_threads) : 0);
+}
+
 struct htp_allreduce_kernel_params {
    int32_t rank;
    int32_t n_ranks;
@@ -11,9 +11,10 @@
 #include "hvx-utils.h"
 #include "hex-dma.h"

+#include "hex-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
+#include "htp-tensor.h"

 #ifndef MIN
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -22,6 +23,9 @@
 struct htp_argsort_context {
    struct htp_ops_context * octx;
    uint32_t                 nrows_per_thread;
+    uint32_t                 total_rows;
+    uint32_t                 row_start;
+    uint32_t                 row_end;
    uint8_t *                vtcm_base;
    size_t                   vtcm_per_thread;
 };
@@ -336,10 +340,9 @@ static void htp_argsort_f32_##ne00##_##order_name(unsigned int n, unsigned int i
    const struct htp_tensor * src0 = octx->src[0];                                                             \
    const struct htp_tensor * dst = octx->dst;                                                                 \
    uint8_t * spad = actx->vtcm_base + actx->vtcm_per_thread * i;                                              \
-    uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3];                                             \
    uint32_t rows_per_thread = actx->nrows_per_thread;                                                         \
-    uint32_t start_row = rows_per_thread * i;                                                                  \
-    uint32_t end_row = MIN(start_row + rows_per_thread, total_rows);                                           \
+    uint32_t start_row = actx->row_start + rows_per_thread * i;                                                \
+    uint32_t end_row = MIN(start_row + rows_per_thread, actx->row_end);                                        \
    size_t values_size = hex_round_up(ne00 * sizeof(float), 128);                                              \
    float * values_buf = (float *) spad;                                                                       \
    int32_t * indices_buf = (int32_t *) (spad + values_size);                                                  \
@@ -386,9 +389,6 @@ static void htp_argsort_f32_fallback(unsigned int n, unsigned int i, void * data

    // Dimensions
    uint32_t ne00 = src0->ne[0];
-    uint32_t ne01 = src0->ne[1];
-    uint32_t ne02 = src0->ne[2];
-    uint32_t ne03 = src0->ne[3];

    uint32_t nb01 = src0->nb[1];

@@ -398,10 +398,9 @@ static void htp_argsort_f32_fallback(unsigned int n, unsigned int i, void * data
    enum ggml_sort_order order = (enum ggml_sort_order) octx->op_params[0];

    // Rows to process
-    uint32_t total_rows = ne01 * ne02 * ne03;
    uint32_t rows_per_thread = actx->nrows_per_thread;
-    uint32_t start_row = rows_per_thread * i;
-    uint32_t end_row = MIN(start_row + rows_per_thread, total_rows);
+    uint32_t start_row = actx->row_start + rows_per_thread * i;
+    uint32_t end_row = MIN(start_row + rows_per_thread, actx->row_end);

    size_t values_size = hex_round_up(ne00 * sizeof(float), 128);
    uint32_t num_vec_ind_values = hmx_ceil_div(ne00, VLEN/(sizeof(int32_t)));
@@ -451,8 +450,28 @@ int op_argsort(struct htp_ops_context * octx) {
        return HTP_STATUS_NO_SUPPORT;
    }

-    const uint32_t total_rows = octx->src[0]->ne[1] * octx->src[0]->ne[2] * octx->src[0]->ne[3];
-    const uint32_t n_threads = MIN(total_rows, octx->n_threads);
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;
+
+    const uint32_t total_rows  = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    const size_t dst_row_size  = dst->ne[0]  * sizeof(int32_t);
+
+    uint32_t row_start = 0;
+    uint32_t row_end   = total_rows;
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, sizeof(int32_t), (uint32_t) dst_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        row_end   = range.start + range.count;
+    }
+
+    const uint32_t nrows = row_end - row_start;
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;

    // Allocate scratchpad
    // We need 1 row of float + 1 row of int32 per thread.
@@ -478,7 +497,10 @@ int op_argsort(struct htp_ops_context * octx) {

    struct htp_argsort_context actx;
    actx.octx = octx;
-    actx.nrows_per_thread = (total_rows + n_threads - 1) / n_threads;
+    actx.nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    actx.total_rows       = nrows;
+    actx.row_start        = row_start;
+    actx.row_end          = row_end;
    actx.vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
    actx.vtcm_per_thread = spad_per_thread;

@@ -508,7 +530,7 @@ int op_argsort(struct htp_ops_context * octx) {
    }

    // Run jobs
-    worker_pool_run_func(octx->ctx->worker_pool, job_func, &actx, n_threads);
+    work_queue_run(octx->ctx->work_queue, job_func, &actx, n_threads);

    return HTP_STATUS_OK;
 }
@@ -13,9 +13,10 @@

 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+#include "hex-common.h"
+#include "hex-profile.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
 #include "htp-tensor.h"

 #ifndef MIN
@@ -36,6 +37,8 @@ struct htp_binary_context {

    uint32_t block_max;
    uint32_t nrows_per_thread;
+    uint32_t total_rows;
+    uint32_t row_start;
    size_t   src0_row_size_aligned;
    size_t   src1_row_size_aligned;
    size_t   dst_row_size_aligned;
@@ -48,27 +51,27 @@ struct htp_binary_context {
    const struct htp_tensor * src0 = octx->src[0]; \
    const struct htp_tensor * src1 = octx->src[1]; \
    const struct htp_tensor * dst  = octx->dst;    \
-                                       \
-    const uint32_t ne00 = src0->ne[0]; \
-    const uint32_t ne01 = src0->ne[1]; \
-    const uint32_t ne02 = src0->ne[2]; \
-    const uint32_t ne03 = src0->ne[3]; \
-                                       \
-    const uint32_t ne10 = src1->ne[0]; \
-    const uint32_t ne11 = src1->ne[1]; \
-    const uint32_t ne12 = src1->ne[2]; \
-    const uint32_t ne13 = src1->ne[3]; \
-                                       \
-    const uint32_t nb01 = src0->nb[1]; \
-    const uint32_t nb02 = src0->nb[2]; \
-    const uint32_t nb03 = src0->nb[3]; \
-                                       \
-    const uint32_t nb11 = src1->nb[1]; \
-    const uint32_t nb12 = src1->nb[2]; \
-    const uint32_t nb13 = src1->nb[3]; \
-                                       \
-    const uint32_t nb1 = dst->nb[1];   \
-    const uint32_t nb2 = dst->nb[2];   \
+                                                   \
+    const uint32_t ne00 = src0->ne[0];             \
+    const uint32_t ne01 = src0->ne[1];             \
+    const uint32_t ne02 = src0->ne[2];             \
+    const uint32_t ne03 = src0->ne[3];             \
+                                                   \
+    const uint32_t ne10 = src1->ne[0];             \
+    const uint32_t ne11 = src1->ne[1];             \
+    const uint32_t ne12 = src1->ne[2];             \
+    const uint32_t ne13 = src1->ne[3];             \
+                                                   \
+    const uint32_t nb01 = src0->nb[1];             \
+    const uint32_t nb02 = src0->nb[2];             \
+    const uint32_t nb03 = src0->nb[3];             \
+                                                   \
+    const uint32_t nb11 = src1->nb[1];             \
+    const uint32_t nb12 = src1->nb[2];             \
+    const uint32_t nb13 = src1->nb[3];             \
+                                                   \
+    const uint32_t nb1 = dst->nb[1];               \
+    const uint32_t nb2 = dst->nb[2];               \
    const uint32_t nb3 = dst->nb[3];

 static inline uint32_t calc_block_size(struct htp_binary_context * bctx, uint32_t ir, uint32_t end_row, uint32_t ne01, uint32_t ne02) {
@@ -93,87 +96,87 @@ static inline uint32_t calc_block_size(struct htp_binary_context * bctx, uint32_
 }

 // Macro for scalar op switch
-#define COMPUTE_SCALAR_OP(DST, SRC, VAL, TYPE, N) \
-    if(TYPE == HTP_TYPE_F32) { \
-        switch (octx->op) { \
-            case HTP_OP_ADD: hvx_add_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break; \
-            case HTP_OP_SUB: hvx_sub_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break; \
-            case HTP_OP_MUL: hvx_mul_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break; \
+#define COMPUTE_SCALAR_OP(DST, SRC, VAL, TYPE, N)                                               \
+    if(TYPE == HTP_TYPE_F32) {                                                                  \
+        switch (octx->op) {                                                                     \
+            case HTP_OP_ADD: hvx_add_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break;          \
+            case HTP_OP_SUB: hvx_sub_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break;          \
+            case HTP_OP_MUL: hvx_mul_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break;          \
            case HTP_OP_DIV: hvx_mul_scalar_f32_aa(DST, SRC, 1.0f / (*(float *)VAL), N); break; \
-            default: break; \
-        } \
-    } \
-    else { \
-        switch (octx->op) { \
-            case HTP_OP_ADD: hvx_add_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break; \
-            case HTP_OP_SUB: hvx_sub_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break; \
-            case HTP_OP_MUL: hvx_mul_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break; \
-            case HTP_OP_DIV: hvx_div_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break; \
-            default: break; \
-        } \
+            default: break;                                                                     \
+        }                                                                                       \
+    }                                                                                           \
+    else {                                                                                      \
+        switch (octx->op) {                                                                     \
+            case HTP_OP_ADD: hvx_add_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break;       \
+            case HTP_OP_SUB: hvx_sub_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break;       \
+            case HTP_OP_MUL: hvx_mul_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break;       \
+            case HTP_OP_DIV: hvx_div_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break;       \
+            default: break;                                                                     \
+        }                                                                                       \
    }

 // Macro for vector op switch (All Aligned)
-#define COMPUTE_VECTOR_OP_AAA(DST, SRC0, SRC1, TYPE, N) \
-    if(TYPE == HTP_TYPE_F32) { \
-        switch (octx->op) { \
+#define COMPUTE_VECTOR_OP_AAA(DST, SRC0, SRC1, TYPE, N)                  \
+    if(TYPE == HTP_TYPE_F32) {                                           \
+        switch (octx->op) {                                              \
            case HTP_OP_ADD: hvx_add_f32_aaa(DST, SRC0, SRC1, N); break; \
            case HTP_OP_SUB: hvx_sub_f32_aaa(DST, SRC0, SRC1, N); break; \
            case HTP_OP_MUL: hvx_mul_f32_aaa(DST, SRC0, SRC1, N); break; \
            case HTP_OP_DIV: hvx_div_f32_aaa(DST, SRC0, SRC1, N); break; \
-            default: break; \
-        } \
-    } \
-    else { \
-        switch (octx->op) { \
+            default: break;                                              \
+        }                                                                \
+    }                                                                    \
+    else {                                                               \
+        switch (octx->op) {                                              \
            case HTP_OP_ADD: hvx_add_f16_aaa(DST, SRC0, SRC1, N); break; \
            case HTP_OP_SUB: hvx_sub_f16_aaa(DST, SRC0, SRC1, N); break; \
            case HTP_OP_MUL: hvx_mul_f16_aaa(DST, SRC0, SRC1, N); break; \
            case HTP_OP_DIV: hvx_div_f16_aaa(DST, SRC0, SRC1, N); break; \
-            default: break; \
-        } \
+            default: break;                                              \
+        }                                                                \
    }

 // Macro for vector op switch (Dst Aligned, Src0 Aligned, Src1 Unaligned)
-#define COMPUTE_VECTOR_OP_AAU(DST, SRC0, SRC1, TYPE, N) \
-    if(TYPE == HTP_TYPE_F32) { \
-        switch (octx->op) { \
+#define COMPUTE_VECTOR_OP_AAU(DST, SRC0, SRC1, TYPE, N)                  \
+    if(TYPE == HTP_TYPE_F32) {                                           \
+        switch (octx->op) {                                              \
            case HTP_OP_ADD: hvx_add_f32_aau(DST, SRC0, SRC1, N); break; \
            case HTP_OP_SUB: hvx_sub_f32_aau(DST, SRC0, SRC1, N); break; \
            case HTP_OP_MUL: hvx_mul_f32_aau(DST, SRC0, SRC1, N); break; \
            case HTP_OP_DIV: hvx_div_f32_aau(DST, SRC0, SRC1, N); break; \
-            default: break; \
-        } \
-    } \
-    else { \
-        switch (octx->op) { \
+            default: break;                                              \
+        }                                                                \
+    }                                                                    \
+    else {                                                               \
+        switch (octx->op) {                                              \
            case HTP_OP_ADD: hvx_add_f16_aau(DST, SRC0, SRC1, N); break; \
            case HTP_OP_SUB: hvx_sub_f16_aau(DST, SRC0, SRC1, N); break; \
            case HTP_OP_MUL: hvx_mul_f16_aau(DST, SRC0, SRC1, N); break; \
            case HTP_OP_DIV: hvx_div_f16_aau(DST, SRC0, SRC1, N); break; \
-            default: break; \
-        } \
+            default: break;                                              \
+        }                                                                \
    }

 // Macro for vector op switch (All Unaligned - generic loop used in element repeat)
-#define COMPUTE_VECTOR_OP_UUU(DST, SRC0, SRC1, TYPE, N) \
-    if(TYPE == HTP_TYPE_F32) { \
-        switch (octx->op) { \
+#define COMPUTE_VECTOR_OP_UUU(DST, SRC0, SRC1, TYPE, N)                  \
+    if(TYPE == HTP_TYPE_F32) {                                           \
+        switch (octx->op) {                                              \
            case HTP_OP_ADD: hvx_add_f32_uuu(DST, SRC0, SRC1, N); break; \
            case HTP_OP_SUB: hvx_sub_f32_uuu(DST, SRC0, SRC1, N); break; \
            case HTP_OP_MUL: hvx_mul_f32_uuu(DST, SRC0, SRC1, N); break; \
            case HTP_OP_DIV: hvx_div_f32_uuu(DST, SRC0, SRC1, N); break; \
-            default: break; \
-        } \
-    } \
-    else { \
-        switch (octx->op) { \
+            default: break;                                              \
+        }                                                                \
+    }                                                                    \
+    else {                                                               \
+        switch (octx->op) {                                              \
            case HTP_OP_ADD: hvx_add_f16_uuu(DST, SRC0, SRC1, N); break; \
            case HTP_OP_SUB: hvx_sub_f16_uuu(DST, SRC0, SRC1, N); break; \
            case HTP_OP_MUL: hvx_mul_f16_uuu(DST, SRC0, SRC1, N); break; \
            case HTP_OP_DIV: hvx_div_f16_uuu(DST, SRC0, SRC1, N); break; \
-            default: break; \
-        } \
+            default: break;                                              \
+        }                                                                \
    }

 // 1. Scalar src1 (ne10 == 1)
@@ -184,9 +187,8 @@ static void binary_job_scalar(unsigned int nth, unsigned int ith, void * data) {

    const uint32_t src0_type = octx->src[0]->type;
    const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16);
-    const uint32_t total_rows = ne01 * ne02 * ne03;
-    const uint32_t start_row = bctx->nrows_per_thread * ith;
-    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, total_rows);
+    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
+    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
    if (start_row >= end_row) return;

    FARF(HIGH, "binary-scalar: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
@@ -222,6 +224,8 @@ static void binary_job_scalar(unsigned int nth, unsigned int ith, void * data) {
    }

    // Main loop
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
    for (uint32_t ir = start_row; ir < end_row; ) {
        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);

@@ -242,12 +246,14 @@ static void binary_job_scalar(unsigned int nth, unsigned int ith, void * data) {
        uint8_t * src1_ptr = (uint8_t *)src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11;
        uint32_t s1_stride = (ne11 == 1) ? 0 : nb11;

+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
        for (uint32_t r = 0; r < current_block_size; r++) {
            uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned;
            uint8_t * r_dst  = d_spad + r * bctx->dst_row_size_aligned;
            COMPUTE_SCALAR_OP(r_dst, r_src0, src1_ptr, src0_type, ne00);
            src1_ptr += s1_stride;
        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);

        uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size);
@@ -266,6 +272,7 @@ static void binary_job_scalar(unsigned int nth, unsigned int ith, void * data) {
        }
        ir += current_block_size;
    }
+
    dma_queue_flush(q);
 }

@@ -277,9 +284,8 @@ static void binary_job_vector_same_shape(unsigned int nth, unsigned int ith, voi

    const uint32_t src0_type = octx->src[0]->type;
    const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16);
-    const uint32_t total_rows = ne01 * ne02 * ne03;
-    const uint32_t start_row = bctx->nrows_per_thread * ith;
-    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, total_rows);
+    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
+    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
    if (start_row >= end_row) return;

    FARF(HIGH, "binary-same-shape: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
@@ -323,18 +329,22 @@ static void binary_job_vector_same_shape(unsigned int nth, unsigned int ith, voi
        spad_idx ^= 1;
    }

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
    for (uint32_t ir = start_row; ir < end_row; ) {
        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
        uint8_t * d_spad  = (uint8_t *) dma_queue_pop(q).src;
        uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst;
        uint8_t * s1_spad = (uint8_t *) dma_queue_pop(q).dst;

+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
        for (uint32_t r = 0; r < current_block_size; r++) {
            uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned;
            uint8_t * r_src1 = s1_spad + r * bctx->src1_row_size_aligned;
            uint8_t * r_dst  = d_spad  + r * bctx->dst_row_size_aligned;
            COMPUTE_VECTOR_OP_AAA(r_dst, r_src0, r_src1, src0_type, ne00);
        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);

        uint32_t i03, i02, i01, rem;
        i03 = fastdiv(ir, &bctx->src0_dim12_div);
@@ -366,6 +376,7 @@ static void binary_job_vector_same_shape(unsigned int nth, unsigned int ith, voi
        }
        ir += current_block_size;
    }
+
    dma_queue_flush(q);
 }

@@ -377,9 +388,8 @@ static void binary_job_vector_row_broadcast(unsigned int nth, unsigned int ith,

    const uint32_t src0_type  = octx->src[0]->type;
    const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16);
-    const uint32_t total_rows = ne01 * ne02 * ne03;
-    const uint32_t start_row  = bctx->nrows_per_thread * ith;
-    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, total_rows);
+    const uint32_t start_row  = bctx->row_start + bctx->nrows_per_thread * ith;
+    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
    if (start_row >= end_row) return;

    FARF(HIGH, "binary-row-bcast: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
@@ -416,17 +426,21 @@ static void binary_job_vector_row_broadcast(unsigned int nth, unsigned int ith,
        spad_idx ^= 1;
    }

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
    for (uint32_t ir = start_row; ir < end_row; ) {
        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
        uint8_t * d_spad  = (uint8_t *) dma_queue_pop(q).src;
        uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst;

+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
        for (uint32_t r = 0; r < current_block_size; r++) {
            uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned;
            uint8_t * r_src1 = (uint8_t *)s1_ptr; // Constant
            uint8_t * r_dst  = d_spad + r * bctx->dst_row_size_aligned;
            COMPUTE_VECTOR_OP_AAA(r_dst, r_src0, r_src1, src0_type, ne00);
        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);

        uint32_t i03 = fastdiv(ir, &bctx->src0_dim12_div);
        uint32_t rem = ir - i03 * (ne02 * ne01);
@@ -447,6 +461,7 @@ static void binary_job_vector_row_broadcast(unsigned int nth, unsigned int ith,
        }
        ir += current_block_size;
    }
+
    dma_queue_flush(q);
 }

@@ -458,9 +473,8 @@ static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void *

    const uint32_t src0_type = octx->src[0]->type;
    const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16);
-    const uint32_t total_rows = ne01 * ne02 * ne03;
-    const uint32_t start_row  = bctx->nrows_per_thread * ith;
-    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, total_rows);
+    const uint32_t start_row  = bctx->row_start + bctx->nrows_per_thread * ith;
+    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
    if (start_row >= end_row) return;

    FARF(HIGH, "binary-complex: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
@@ -493,6 +507,8 @@ static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void *
        spad_idx ^= 1;
    }

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
    for (uint32_t ir = start_row; ir < end_row; ) {
        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
        uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src;
@@ -503,6 +519,7 @@ static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void *
        uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
        uint32_t i01 = rem - i02 * ne01;

+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
        for (uint32_t r = 0; r < current_block_size; r++) {
            uint32_t r_i01 = i01 + r;
            uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div);
@@ -516,6 +533,7 @@ static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void *
            // Read src1 from DDR (unaligned)
            COMPUTE_VECTOR_OP_AAU(r_dst, r_src0, r_src1, src0_type, ne00);
        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);

        uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size);
@@ -532,6 +550,7 @@ static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void *
        }
        ir += current_block_size;
    }
+
    dma_queue_flush(q);
 }

@@ -544,9 +563,8 @@ static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void *
    const uint32_t src0_type = octx->src[0]->type;
    const uint32_t elem_size_bytes = (src0_type == HTP_TYPE_F32) ? sizeof(float) : sizeof(_Float16);
    const uint32_t row_size_bytes = ne00 * elem_size_bytes;;
-    const uint32_t total_rows = ne01 * ne02 * ne03;
-    const uint32_t start_row  = bctx->nrows_per_thread * ith;
-    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, total_rows);
+    const uint32_t start_row  = bctx->row_start + bctx->nrows_per_thread * ith;
+    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
    if (start_row >= end_row) return;

    uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
@@ -579,6 +597,8 @@ static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void *
        spad_idx ^= 1;
    }

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
    for (uint32_t ir = start_row; ir < end_row; ) {
        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
        uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src;
@@ -589,6 +609,7 @@ static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void *
        uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
        uint32_t i01 = rem - i02 * ne01;

+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
        for (uint32_t r = 0; r < current_block_size; r++) {
            uint32_t r_i01 = i01 + r;
            uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div);
@@ -606,6 +627,7 @@ static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void *
                COMPUTE_VECTOR_OP_UUU(r_dst + c * elem_size_bytes, r_src0 + c * elem_size_bytes, r_src1_row, src0_type, len);
            }
        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);

        uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size);
@@ -622,6 +644,7 @@ static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void *
        }
        ir += current_block_size;
    }
+
    dma_queue_flush(q);
 }

@@ -650,9 +673,8 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) {
    const uint32_t nb2 = dst->nb[2];
    const uint32_t nb3 = dst->nb[3];

-    const uint32_t total_rows = ne01 * ne02 * ne03;
-    const uint32_t start_row = bctx->nrows_per_thread * ith;
-    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, total_rows);
+    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
+    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
    if (start_row >= end_row) return;

    uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
@@ -683,6 +705,8 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) {
        spad_idx ^= 1;
    }

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
    for (uint32_t ir = start_row; ir < end_row; ) {
        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
        uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src;
@@ -693,6 +717,7 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) {
        uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
        uint32_t i01 = rem - i02 * ne01;

+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
        for (uint32_t r = 0; r < current_block_size; r++) {
            uint32_t r_i01 = i01 + r; // linear within block since we split at ne01

@@ -704,6 +729,7 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) {

            hvx_add_f32_aau(r_dst, r_src0, r_src1, ne00);
        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);

        uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, ne00 * sizeof(float), current_block_size);
@@ -720,6 +746,7 @@ static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) {
        }
        ir += current_block_size;
    }
+
    dma_queue_flush(q);
 }

@@ -729,15 +756,31 @@ static int execute_op_binary(struct htp_ops_context * octx) {
    const struct htp_tensor * dst  = octx->dst;

    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
-    const uint32_t n_threads  = MIN(octx->n_threads, src0_nrows);

-    // Use packed row sizes for VTCM allocation
+    // Use packed row sizes for VTCM allocation and alignment
    const uint32_t src0_type = octx->src[0]->type;
    const size_t elem_size = (src0_type == HTP_TYPE_F32) ? sizeof(float) : sizeof(_Float16);
    const size_t src0_row_size = src0->ne[0] * elem_size;
    const size_t src1_row_size = src1->ne[0] * elem_size;
    const size_t dst_row_size  = dst->ne[0]  * elem_size;

+    uint32_t row_start = 0;
+    uint32_t nrows     = src0_nrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, (uint32_t) elem_size, (uint32_t) dst_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
    size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN);
    size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN);
    size_t dst_row_size_aligned  = hex_round_up(dst_row_size,  VLEN);
@@ -815,7 +858,9 @@ static int execute_op_binary(struct htp_ops_context * octx) {

    struct htp_binary_context bctx;
    bctx.octx                  = octx;
-    bctx.nrows_per_thread      = (src0_nrows + n_threads - 1) / n_threads;
+    bctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    bctx.total_rows            = nrows;
+    bctx.row_start             = row_start;
    bctx.block_max             = rows_per_buffer;
    bctx.src0_row_size_aligned = src0_row_size_aligned;
    bctx.src1_row_size_aligned = src1_row_size_aligned;
@@ -850,7 +895,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
        dma_queue_pop(q);
    }

-    worker_pool_run_func(octx->ctx->worker_pool, worker_func, &bctx, n_threads);
+    work_queue_run(octx->ctx->work_queue, worker_func, &bctx, n_threads);

    return HTP_STATUS_OK;
 }
@@ -870,4 +915,3 @@ int op_binary(struct htp_ops_context * octx) {

    return HTP_STATUS_NO_SUPPORT;
 }
-
@@ -1,5 +1,8 @@
+#include "hex-common.h"
+#include "hex-profile.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 #include "hexagon_types.h"
 #include "hexagon_protos.h"
 #include "hvx_hexagon_protos.h"
@@ -13,6 +16,10 @@ struct htp_concat_context {
    struct htp_ops_context * octx;
    uint32_t dim;
    uint32_t nrows_per_thread;
+    uint32_t row_start;
+    uint32_t nrows;
+    uint32_t elem_start;
+    uint32_t nelems;
    struct fastdiv_values div_ne0;
    struct fastdiv_values div_ne1;
    struct fastdiv_values div_ne2;
@@ -28,10 +35,10 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *

    const uint32_t src0_ne0 = src0->ne[0];
    const uint32_t src1_ne0 = src1->ne[0];
-    const uint32_t ne1      = dst->ne[1];

-    const uint32_t start_i = ith * cctx->nrows_per_thread;
-    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < ne1) ? (start_i + cctx->nrows_per_thread) : ne1;
+    const uint32_t row_end = cctx->row_start + cctx->nrows;
+    const uint32_t start_i = cctx->row_start + ith * cctx->nrows_per_thread;
+    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < row_end) ? (start_i + cctx->nrows_per_thread) : row_end;
    if (start_i >= end_i) return;

    dma_queue * q = octx->ctx->dma[ith];
@@ -51,6 +58,8 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
    const uint32_t spad0_row_bytes = hex_round_up((src0_ne0 + src1_ne0_padded) * sizeof(float), VLEN);
    uint32_t mu = src1_ne0_padded * spad1_stride;

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
    for (uint32_t i = start_i; i < end_i; i += block_i) {
        uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i;

@@ -66,6 +75,7 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *

        HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride);

+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
        for (uint32_t j = 0; j < src1_ne0_padded; j += 32) {
            #pragma unroll(4)
            for (uint32_t ii = 0; ii < current_block_i; ii++) {
@@ -75,6 +85,7 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
                hvx_vmemu(dst_ptr) = vtcm_tmp[ii];
            }
        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);

        dma_queue_pop(q); // src0

@@ -95,10 +106,10 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *

    const uint32_t src0_ne0 = src0->ne[0];
    const uint32_t src1_ne0 = src1->ne[0];
-    const uint32_t ne1      = dst->ne[1];

-    const uint32_t start_i = ith * cctx->nrows_per_thread;
-    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < ne1) ? (start_i + cctx->nrows_per_thread) : ne1;
+    const uint32_t row_end = cctx->row_start + cctx->nrows;
+    const uint32_t start_i = cctx->row_start + ith * cctx->nrows_per_thread;
+    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < row_end) ? (start_i + cctx->nrows_per_thread) : row_end;
    if (start_i >= end_i) return;

    dma_queue * q = octx->ctx->dma[ith];
@@ -118,6 +129,8 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
    const uint32_t spad0_row_bytes = hex_round_up((src0_ne0 + src1_ne0_padded) * sizeof(__fp16), VLEN);
    uint32_t mu = src1_ne0_padded * spad1_stride;

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
    for (uint32_t i = start_i; i < end_i; i += block_i) {
        uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i;

@@ -133,6 +146,7 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *

        HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride);

+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
        for (uint32_t j = 0; j < src1_ne0_padded; j += 64) {
            #pragma unroll(4)
            for (uint32_t ii = 0; ii < current_block_i; ii++) {
@@ -142,6 +156,7 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
                hvx_vmemu(dst_ptr) = vtcm_tmp[ii];
            }
        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);

        dma_queue_pop(q); // src0

@@ -164,11 +179,14 @@ static void concat_generic(unsigned int nth, unsigned int ith, void * data) {
    const uint32_t type_size = (dst->type == HTP_TYPE_F32 || dst->type == HTP_TYPE_I32) ? 4 : 2;

    const uint32_t ne[4] = {dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]};
-    const uint32_t total_elements = ne[0] * ne[1] * ne[2] * ne[3];
-    const uint32_t chunk_size = (total_elements + nth - 1) / nth;

-    const uint32_t start_idx = MIN(ith * chunk_size, total_elements);
-    const uint32_t end_idx   = MIN(start_idx + chunk_size, total_elements);
+    // Per-device element range aligned to prevent false sharing
+    const uint32_t elem_start = cctx->elem_start;
+    const uint32_t nelems     = cctx->nelems;
+    const uint32_t chunk_size = (nelems + nth - 1) / nth;
+
+    const uint32_t start_idx = MIN(elem_start + ith * chunk_size, elem_start + nelems);
+    const uint32_t end_idx   = MIN(start_idx + chunk_size, elem_start + nelems);

    // Naive scalar element-wise copy
    for (uint32_t idx = start_idx; idx < end_idx; idx++) {
@@ -236,13 +254,28 @@ int op_concat(struct htp_ops_context * octx) {
    void (*worker_func)(unsigned int, unsigned int, void *) = concat_generic;

    if (dim == 0 && is_2d && is_src1_transposed && !is_src0_transposed) {
-        n_threads = MIN(dst->ne[1], n_threads);
-        if (n_threads < 1) {
-            n_threads = 1;
+        const uint32_t total_rows = dst->ne[1];
+        const size_t dst_data_row_size = dst->ne[0] * type_size;
+        uint32_t row_start = 0;
+        uint32_t nrows     = total_rows;
+        if (octx->ctx->mdev.count > 1) {
+            uint32_t rows_per_chunk = 0;
+            htp_tensor_mdev_rows_per_chunk(dst, type_size, (uint32_t) dst_data_row_size, &rows_per_chunk);
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            row_start = range.start;
+            nrows     = range.count;
        }
+
+        if (nrows == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        cctx.row_start = row_start;
+        cctx.nrows     = nrows;
+
        uint32_t block_i = (type_size == 4) ? 32 : 64;

-        cctx.nrows_per_thread = hmx_ceil_div(dst->ne[1], n_threads);
+        cctx.nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);

        // Allocate VTCM
        uint32_t spad1_stride = block_i * type_size;
@@ -270,8 +303,26 @@ int op_concat(struct htp_ops_context * octx) {
        } else {
            worker_func = concat_2d_f16_transposed;
        }
+    } else {
+        const uint32_t total_elements = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
+        uint32_t elem_start = 0;
+        uint32_t nelems     = total_elements;
+        if (octx->ctx->mdev.count > 1) {
+            const uint32_t elems_per_chunk = HEX_L2_LINE_SIZE / type_size;
+            const bool can_split = htp_tensor_mdev_data_aligned(dst) && htp_tensor_is_contiguous(dst, type_size) && !htp_tensor_is_permuted(dst);
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_elements, can_split ? elems_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            elem_start = range.start;
+            nelems     = range.count;
+        }
+
+        if (nelems == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        cctx.elem_start = elem_start;
+        cctx.nelems     = nelems;
    }

-    worker_pool_run_func(octx->ctx->worker_pool, worker_func, &cctx, n_threads);
+    work_queue_run(octx->ctx->work_queue, worker_func, &cctx, n_threads);
    return HTP_STATUS_OK;
 }
@@ -16,6 +16,7 @@
 #include "htp-ops.h"
 #include "hvx-utils.h"
 #include "htp-tensor.h"
+#include "htp-fence.h"

 struct htp_copy_context {
    struct htp_ops_context * octx;
@@ -29,7 +30,23 @@ struct htp_copy_context {
    uint32_t          src0_blocks_per_row;
    uint32_t          dst_blocks_per_row;

+    uint32_t          elem_start;
+    uint32_t          nelem;
+    uint32_t          elem_per_thread;
+
    uint32_t          src0_nrows_per_thread;
+    uint32_t          row_start;
+    uint32_t          nrows;
+
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne02_ne01;
+
+    struct fastdiv_values div_ne0;
+    struct fastdiv_values div_ne1_ne0;
+    struct fastdiv_values div_ne2_ne1_ne0;
+    struct fastdiv_values div_ne00;
+    struct fastdiv_values div_ne01_ne00;
+    struct fastdiv_values div_ne02_ne01_ne00;
 };

 #define cpy_preamble                              \
@@ -54,131 +71,113 @@ struct htp_copy_context {
    const uint32_t  nb0 = dst->nb[0];             \
    const uint32_t  nb1 = dst->nb[1];             \
    const uint32_t  nb2 = dst->nb[2];             \
-    const uint32_t  nb3 = dst->nb[3];             \
-                                                  \
-    const uint32_t   nr = ne01;
+    const uint32_t  nb3 = dst->nb[3];

-#define DEFINE_CPY_SAMESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                                       \
-static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, void * data) {                   \
-    struct htp_copy_context * ct = (struct htp_copy_context *) data;                                           \
-    struct htp_ops_context * octx = ct->octx;                                                                  \
-    cpy_preamble;                                                                                              \
-    const uint32_t dr  = ct->src0_nrows_per_thread;                                                            \
-    const uint32_t ir0 = dr * ith;                                                                             \
-    const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr;                                                    \
-    if (ir0 >= nr) return;                                                                                     \
-    for (uint32_t i03 = 0; i03 < ne03; i03++) {                                                                \
-        for (uint32_t i02 = 0; i02 < ne02; i02++) {                                                            \
-            _Pragma("unroll(4)")                                                                               \
-            for (uint32_t i01 = ir0; i01 < ir1; i01++) {                                                       \
-                uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;                     \
-                uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;                    \
-                hex_l2fetch(src0_ptr, ne00 * ELEM_SIZE, nb01, 2);                                              \
-                hvx_copy_uu(dst_ptr, src0_ptr, ne00, ELEM_SIZE);                                               \
-            }                                                                                                  \
-        }                                                                                                      \
-    }                                                                                                          \
+#define DEFINE_CPY_SAMESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                       \
+static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, void * data) {   \
+    struct htp_copy_context * ct = (struct htp_copy_context *) data;                           \
+    struct htp_ops_context * octx = ct->octx;                                                  \
+    cpy_preamble;                                                                              \
+    const uint32_t dr  = ct->src0_nrows_per_thread;                                            \
+    const uint32_t ir0 = ct->row_start + dr * ith;                                             \
+    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);                             \
+    if (ir0 >= ir1) return;                                                                    \
+    const bool contiguous = (nb01 == ne00 * ELEM_SIZE) && (nb1 == nb01) &&                     \
+                            (nb02 == ne01 * nb01)      && (nb2 == nb02) &&                     \
+                            (nb03 == ne02 * nb02)      && (nb3 == nb03);                       \
+    const uint32_t ne02_ne01 = ne02 * ne01;                                                    \
+    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);                                           \
+    uint32_t rem = ir0 - i03 * ne02_ne01;                                                      \
+    uint32_t i02 = fastdiv(rem, &ct->div_ne01);                                                \
+    uint32_t i01 = rem - i02 * ne01;                                                           \
+    uint8_t * dst_ptr  = (uint8_t *) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;               \
+    uint8_t * src0_ptr = (uint8_t *) src0->data + i01*nb01 + i02*nb02 + i03*nb03;              \
+    if (contiguous) {                                                                          \
+        hvx_copy_uu(dst_ptr, src0_ptr, (ir1 - ir0) * ne00, ELEM_SIZE);                         \
+        return;                                                                                \
+    }                                                                                          \
+    for (uint32_t r = ir0; r < ir1; r++) {                                                     \
+        hex_l2fetch(src0_ptr, ne00 * ELEM_SIZE, nb01, 2);                                      \
+        hvx_copy_uu(dst_ptr, src0_ptr, ne00, ELEM_SIZE);                                       \
+        dst_ptr  += nb1;                                                                       \
+        src0_ptr += nb01;                                                                      \
+        if (++i01 == ne01) {                                                                   \
+            i01 = 0;                                                                           \
+            if (++i02 == ne02) {                                                               \
+                i02 = 0;                                                                       \
+                i03++;                                                                         \
+            }                                                                                  \
+            dst_ptr  = (uint8_t *) dst->data  + i02*nb2  + i03*nb3;                            \
+            src0_ptr = (uint8_t *) src0->data + i02*nb02 + i03*nb03;                           \
+        }                                                                                      \
+    }                                                                                          \
 }

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

-#define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                                         \
-static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void * data) {                     \
-    struct htp_copy_context * ct = (struct htp_copy_context *) data;                                           \
-    struct htp_ops_context * octx = ct->octx;                                                                  \
-    cpy_preamble;                                                                                              \
-    const uint32_t dr  = ct->src0_nrows_per_thread;                                                            \
-    const uint32_t ir0 = dr * ith;                                                                             \
-    const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr;                                                    \
-    if (ir0 >= nr) return;                                                                                     \
-    const bool src0_contig = (nb00 == ELEM_SIZE)   &&                                                          \
-                             (nb01 == ne00 * nb00) &&                                                          \
-                             (nb02 == ne01 * nb01) &&                                                          \
-                             (nb03 == ne02 * nb02);                                                            \
-    const bool dst_contig  = (nb0  == ELEM_SIZE)   &&                                                          \
-                             (nb1  == ne0  * nb0)  &&                                                          \
-                             (nb2  == ne1  * nb1)  &&                                                          \
-                             (nb3  == ne2  * nb2);                                                             \
-    if (src0_contig && dst_contig) {                                                                           \
-        for (int64_t i03 = 0; i03 < ne03; i03++) {                                                             \
-            for (int64_t i02 = 0; i02 < ne02; i02++) {                                                         \
-                uint8_t * src_ptr = (uint8_t *) src0->data + i03*nb03 + i02*nb02 + ir0*nb01;                   \
-                uint32_t  flat    = ((i03*ne02 + i02)*ne01 + ir0) * ne00;                                      \
-                uint8_t * dst_ptr = (uint8_t *) dst->data  + flat * ELEM_SIZE;                                 \
-                hvx_copy_uu(dst_ptr, src_ptr, (ir1 - ir0) * ne00, ELEM_SIZE);                                  \
-            }                                                                                                  \
-        }                                                                                                      \
-        return;                                                                                                \
-    }                                                                                                          \
-    const bool reshape_flat_fast = (ne03 == 1 && ne2 == 1 && ne3 == 1) &&                                      \
-                                   (ne0 == ne00 * ne01) && (ne1 == ne02) &&                                    \
-                                   (nb00 == ELEM_SIZE) && (nb0 == ELEM_SIZE);                                  \
-    if (reshape_flat_fast) {                                                                                   \
-        for (uint32_t i02 = 0; i02 < ne02; i02++) {                                                            \
-            for (uint32_t i01 = ir0; i01 < ir1; i01++) {                                                       \
-                uint8_t * src0_ptr = (uint8_t *) src0->data + i01 * nb01 + i02 * nb02;                         \
-                uint8_t * dst_ptr  = (uint8_t *) dst->data  + i01 * ne00 * ELEM_SIZE + i02 * nb1;              \
-                hvx_copy_uu(dst_ptr, src0_ptr, ne00, ELEM_SIZE);                                               \
-            }                                                                                                  \
-        }                                                                                                      \
-        return;                                                                                                \
-    }                                                                                                          \
-    int64_t k10 = 0;                                                                                           \
-    int64_t i11 = 0;                                                                                           \
-    int64_t i12 = 0;                                                                                           \
-    int64_t i13 = 0;                                                                                           \
-    const int64_t nk00 = ct->src0_blocks_per_row;                                                              \
-    const int64_t nk0  = ct->dst_blocks_per_row;                                                               \
-    for (int64_t i03 = 0; i03 < ne03; i03++) {                                                                 \
-        for (int64_t i02 = 0; i02 < ne02; i02++) {                                                             \
-            k10 += nk00 * ir0;                                                                                 \
-            while (k10 >= nk0) {                                                                               \
-                k10 -= nk0;                                                                                    \
-                if (++i11 == ne1) {                                                                            \
-                    i11 = 0;                                                                                   \
-                    if (++i12 == ne2) {                                                                        \
-                        i12 = 0;                                                                               \
-                        if (++i13 == ne3) {                                                                    \
-                            i13 = 0;                                                                           \
-                        }                                                                                      \
-                    }                                                                                          \
-                }                                                                                              \
-            }                                                                                                  \
-            for (int64_t i01 = ir0; i01 < ir1; i01++) {                                                        \
-                for (int64_t k00 = 0; k00 < nk00; k00++) {                                                     \
-                    const char * src0_ptr = ((char *) src0->data + k00*nb00 + i01*nb01 + i02*nb02 + i03*nb03); \
-                          char * dst_ptr  = ((char *)  dst->data + k10*nb0  + i11*nb1  + i12*nb2  + i13*nb3);  \
-                    memcpy(dst_ptr, src0_ptr, ELEM_SIZE);                                                      \
-                    if (++k10 == nk0) {                                                                        \
-                        k10 = 0;                                                                               \
-                        if (++i11 == ne1) {                                                                    \
-                            i11 = 0;                                                                           \
-                            if (++i12 == ne2) {                                                                \
-                                i12 = 0;                                                                       \
-                                if (++i13 == ne3) {                                                            \
-                                    i13 = 0;                                                                   \
-                                }                                                                              \
-                            }                                                                                  \
-                        }                                                                                      \
-                    }                                                                                          \
-                }                                                                                              \
-            }                                                                                                  \
-            k10 += nk00 * (ne01 - ir1);                                                                        \
-            while (k10 >= nk0) {                                                                               \
-                k10 -= nk0;                                                                                    \
-                if (++i11 == ne1) {                                                                            \
-                    i11 = 0;                                                                                   \
-                    if (++i12 == ne2) {                                                                        \
-                        i12 = 0;                                                                               \
-                        if (++i13 == ne3) {                                                                    \
-                            i13 = 0;                                                                           \
-                        }                                                                                      \
-                    }                                                                                          \
-                }                                                                                              \
-            }                                                                                                  \
-        }                                                                                                      \
-    }                                                                                                          \
+#define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                               \
+static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void * data) {           \
+    struct htp_copy_context * ct = (struct htp_copy_context *) data;                                 \
+    struct htp_ops_context * octx = ct->octx;                                                        \
+    cpy_preamble;                                                                                    \
+    const uint32_t th_nelem = ct->elem_per_thread;                                                   \
+    const uint32_t th_start = ct->elem_start + ith * th_nelem;                                       \
+    const uint32_t th_end   = MIN(th_start + th_nelem, ct->elem_start + ct->nelem);                  \
+    if (th_start >= th_end) return;                                                                  \
+                                                                                                     \
+    const uint32_t ne01_ne00      = ne01 * ne00;                                                     \
+    const uint32_t ne02_ne01_ne00 = ne02 * ne01_ne00;                                                \
+    const uint32_t ne1_ne0        = ne1 * ne0;                                                       \
+    const uint32_t ne2_ne1_ne0    = ne2 * ne1_ne0;                                                   \
+                                                                                                     \
+    uint32_t e = th_start;                                                                           \
+    uint32_t i13 = fastdiv(e, &ct->div_ne2_ne1_ne0);                                                 \
+    uint32_t rem = e - i13 * ne2_ne1_ne0;                                                            \
+    uint32_t i12 = fastdiv(rem, &ct->div_ne1_ne0);                                                   \
+    uint32_t rem2 = rem - i12 * ne1_ne0;                                                             \
+    uint32_t i11 = fastdiv(rem2, &ct->div_ne0);                                                      \
+    uint32_t i10 = rem2 - i11 * ne0;                                                                 \
+                                                                                                     \
+    uint32_t i03 = fastdiv(e, &ct->div_ne02_ne01_ne00);                                              \
+    uint32_t rem_s = e - i03 * ne02_ne01_ne00;                                                       \
+    uint32_t i02 = fastdiv(rem_s, &ct->div_ne01_ne00);                                               \
+    uint32_t rem2_s = rem_s - i02 * ne01_ne00;                                                       \
+    uint32_t i01 = fastdiv(rem2_s, &ct->div_ne00);                                                   \
+    uint32_t i00 = rem2_s - i01 * ne00;                                                              \
+                                                                                                     \
+    char * dst_ptr        = (char *)       dst->data  + i10*nb0  + i11*nb1  + i12*nb2  + i13*nb3;    \
+    const char * src0_ptr = (const char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03;   \
+                                                                                                     \
+    for (; e < th_end; e++) {                                                                        \
+        *((ELEM_TYPE *) dst_ptr) = *((const ELEM_TYPE *) src0_ptr);                                  \
+                                                                                                     \
+        dst_ptr += nb0;                                                                              \
+        if (++i10 == ne0) {                                                                          \
+            i10 = 0;                                                                                 \
+            if (++i11 == ne1) {                                                                      \
+                i11 = 0;                                                                             \
+                if (++i12 == ne2) {                                                                  \
+                    i12 = 0;                                                                         \
+                    i13++;                                                                           \
+                }                                                                                    \
+            }                                                                                        \
+            dst_ptr = (char *) dst->data + i11*nb1 + i12*nb2 + i13*nb3;                              \
+        }                                                                                            \
+                                                                                                     \
+        src0_ptr += nb00;                                                                            \
+        if (++i00 == ne00) {                                                                         \
+            i00 = 0;                                                                                 \
+            if (++i01 == ne01) {                                                                     \
+                i01 = 0;                                                                             \
+                if (++i02 == ne02) {                                                                 \
+                    i02 = 0;                                                                         \
+                    i03++;                                                                           \
+                }                                                                                    \
+            }                                                                                        \
+            src0_ptr = (const char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03;                   \
+        }                                                                                            \
+    }                                                                                                \
 }

 DEFINE_CPY_RESHAPE(f32,  float, 4)
@@ -189,22 +188,33 @@ static void cpy_thread_f16_f32_sameshape(unsigned int nth, unsigned int ith, voi
    struct htp_ops_context * octx = ct->octx;
    cpy_preamble;

-    // parallelize by src0 rows
    const uint32_t dr  = ct->src0_nrows_per_thread;
-    const uint32_t ir0 = dr * ith;
-    const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr;
-    if (ir0 >= nr) return;
+    const uint32_t ir0 = ct->row_start + dr * ith;
+    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);
+    if (ir0 >= ir1) return;

-    // copy by rows
-    for (uint32_t i03 = 0; i03 < ne03; i03++) {
-        for (uint32_t i02 = 0; i02 < ne02; i02++) {
-            #pragma unroll(2)
-            for (uint32_t i01 = ir0; i01 < ir1; i01++) {
-                uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
-                uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
-                hex_l2fetch(src0_ptr, ne00 * sizeof(float), nb01, 2);
-                hvx_copy_f16_f32_uu(dst_ptr, src0_ptr, ne00);
+    const uint32_t ne02_ne01 = ne02 * ne01;
+    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);
+    uint32_t rem = ir0 - i03 * ne02_ne01;
+    uint32_t i02 = fastdiv(rem, &ct->div_ne01);
+    uint32_t i01 = rem - i02 * ne01;
+
+    uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
+    uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
+
+    for (uint32_t r = ir0; r < ir1; r++) {
+        hex_l2fetch(src0_ptr, ne00 * sizeof(float), nb01, 2);
+        hvx_copy_f16_f32_uu(dst_ptr, src0_ptr, ne00);
+        dst_ptr  += nb1;
+        src0_ptr += nb01;
+        if (++i01 == ne01) {
+            i01 = 0;
+            if (++i02 == ne02) {
+                i02 = 0;
+                i03++;
            }
+            dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
+            src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
        }
    }
 }
@@ -214,22 +224,33 @@ static void cpy_thread_f32_f16_sameshape(unsigned int nth, unsigned int ith, voi
    struct htp_ops_context * octx = ct->octx;
    cpy_preamble;

-    // parallelize by src0 rows
    const uint32_t dr  = ct->src0_nrows_per_thread;
-    const uint32_t ir0 = dr * ith;
-    const uint32_t ir1 = (ir0 + dr) < nr ? (ir0 + dr) : nr;
-    if (ir0 >= nr) return;
+    const uint32_t ir0 = ct->row_start + dr * ith;
+    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);
+    if (ir0 >= ir1) return;

-    // copy by rows
-    for (uint32_t i03 = 0; i03 < ne03; i03++) {
-        for (uint32_t i02 = 0; i02 < ne02; i02++) {
-            #pragma unroll(2)
-            for (uint32_t i01 = ir0; i01 < ir1; i01++) {
-                uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
-                uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
-                hex_l2fetch(src0_ptr, ne00 * sizeof(__fp16), nb01, 2);
-                hvx_copy_f32_f16_uu(dst_ptr, src0_ptr, ne00);
+    const uint32_t ne02_ne01 = ne02 * ne01;
+    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);
+    uint32_t rem = ir0 - i03 * ne02_ne01;
+    uint32_t i02 = fastdiv(rem, &ct->div_ne01);
+    uint32_t i01 = rem - i02 * ne01;
+
+    uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
+    uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
+
+    for (uint32_t r = ir0; r < ir1; r++) {
+        hex_l2fetch(src0_ptr, ne00 * sizeof(__fp16), nb01, 2);
+        hvx_copy_f32_f16_uu(dst_ptr, src0_ptr, ne00);
+        dst_ptr  += nb1;
+        src0_ptr += nb01;
+        if (++i01 == ne01) {
+            i01 = 0;
+            if (++i02 == ne02) {
+                i02 = 0;
+                i03++;
            }
+            dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
+            src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
        }
    }
 }
@@ -250,15 +271,19 @@ static inline void cpy_dma_sametype_sameshape(
    dma_queue * q = octx->ctx->dma[0];

    if (contiguous_outer) {
-        dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03);
-        dma_queue_pop(q);
+        if (!dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03)) {
+            dma_queue_flush(q);
+            dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03);
+        }
+        dma_queue_flush(q);
        return;
    }

    for (uint32_t i03 = 0; i03 < ne03; i03++) {
        for (uint32_t i02 = 0; i02 < ne02; i02++) {
-            uint8_t* dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
-            uint8_t* src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
+            uint8_t * dst_ptr  = (uint8_t *) dst->data  + i02 * nb2  + i03 * nb3;
+            uint8_t * src0_ptr = (uint8_t *) src0->data + i02 * nb02 + i03 * nb03;
+
            if (!dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01)) {
                dma_queue_flush(q);
                dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01);
@@ -269,10 +294,9 @@ static inline void cpy_dma_sametype_sameshape(
    dma_queue_flush(q);
 }

-int op_cpy(struct htp_ops_context * octx) {
+static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
    cpy_preamble;
-
-    const uint32_t n_threads = MIN(nr, octx->n_threads);
+    *use_dma = false;

    struct htp_copy_context ct;
    ct.octx = octx;
@@ -296,59 +320,117 @@ int op_cpy(struct htp_ops_context * octx) {
    }

    const bool sametype   = (src0->type == dst->type);
-    const bool transposed = (nb00 > nb01) || (nb0 > nb1);
+    const bool transposed = (nb00 > nb01) || (nb0 > nb1) ||
+                            (nb00 != ct.src0_type_size) || (nb0 != ct.dst_type_size) ||
+                            (nb01 < ne00 * ct.src0_type_size) || (nb1 < ne0 * ct.dst_type_size);
    const bool sameshape  = !transposed && (ne00 == ne0 && ne01 == ne1 && ne02 == ne2 && ne03 == ne3);

-    ct.src0_nrows_per_thread = (nr + n_threads - 1) / n_threads;
+    const uint32_t n_threads = octx->n_threads;

-    worker_callback_t copy_fun = NULL;
-    bool use_dma = false;
+    const bool dst_is_contiguous = htp_tensor_is_contiguous(dst, ct.dst_type_size);

-    if (sametype && sameshape) {
-        use_dma = true;
-    } else if (sameshape) {
-        /**/ if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32)
-            copy_fun = cpy_thread_f16_f32_sameshape;
-        else if (dst->type == HTP_TYPE_F32 && src0->type == HTP_TYPE_F16)
-            copy_fun = cpy_thread_f32_f16_sameshape;
-        else
-            return HTP_STATUS_NO_SUPPORT;
-    } else if (sametype) {
-        if (src0->type == HTP_TYPE_F32) {
-            copy_fun = cpy_thread_f32_reshape;
-        } else {
-            copy_fun = cpy_thread_f16_reshape;
+    if (sameshape) {
+        const uint32_t total_rows = ne01 * ne02 * ne03;
+        const uint32_t row_size   = ne00 * ct.dst_type_size;
+
+        ct.div_ne01      = init_fastdiv_values(ne01);
+        ct.div_ne02_ne01 = init_fastdiv_values(ne02 * ne01);
+
+        uint32_t row_start = 0;
+        uint32_t nrows     = total_rows;
+
+        if (octx->ctx->mdev.count > 1) {
+            const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
+            const bool can_split = htp_tensor_mdev_data_aligned(dst) && dst_is_contiguous;
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, can_split ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            row_start = range.start;
+            nrows     = range.count;
        }
+
+        if (nrows == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        ct.row_start = row_start;
+        ct.nrows     = nrows;
+        ct.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+
+        if (sametype && octx->ctx->mdev.count <= 1) {
+            *use_dma = true;
+            cpy_dma_sametype_sameshape(octx, dst, src0, ct.src0_type_size, ne00, ne01, ne02, ne03, nb01, nb02, nb03, nb1, nb2, nb3);
+        } else {
+            work_queue_func_t copy_fun = NULL;
+            if (sametype) {
+                copy_fun = (src0->type == HTP_TYPE_F32) ? cpy_thread_f32_sameshape : cpy_thread_f16_sameshape;
+            } else if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32) {
+                copy_fun = cpy_thread_f16_f32_sameshape;
+            } else if (dst->type == HTP_TYPE_F32 && src0->type == HTP_TYPE_F16) {
+                copy_fun = cpy_thread_f32_f16_sameshape;
+            } else {
+                return HTP_STATUS_NO_SUPPORT;
+            }
+            work_queue_run(octx->ctx->work_queue, copy_fun, &ct, n_threads);
+        }
+    } else if (sametype) {
+        const uint32_t total_elems = ne0 * ne1 * ne2 * ne3;
+        const uint32_t elems_per_line = (ct.dst_type_size == 4) ? 32 : 64;
+
+        ct.div_ne0            = init_fastdiv_values(ne0);
+        ct.div_ne1_ne0        = init_fastdiv_values(ne1 * ne0);
+        ct.div_ne2_ne1_ne0    = init_fastdiv_values(ne2 * ne1 * ne0);
+        ct.div_ne00           = init_fastdiv_values(ne00);
+        ct.div_ne01_ne00      = init_fastdiv_values(ne01 * ne00);
+        ct.div_ne02_ne01_ne00 = init_fastdiv_values(ne02 * ne01 * ne00);
+
+        uint32_t elem_start = 0;
+        uint32_t nelem      = total_elems;
+
+        if (octx->ctx->mdev.count > 1) {
+            const bool can_split = htp_tensor_mdev_data_aligned(dst) && dst_is_contiguous;
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_elems, can_split ? elems_per_line : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            elem_start = range.start;
+            nelem      = range.count;
+        }
+
+        if (nelem == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        ct.elem_start      = elem_start;
+        ct.nelem           = nelem;
+        ct.elem_per_thread = fastdiv(nelem + n_threads - 1, &octx->n_threads_div);
+
+        work_queue_func_t copy_fun = (src0->type == HTP_TYPE_F32) ? cpy_thread_f32_reshape : cpy_thread_f16_reshape;
+        work_queue_run(octx->ctx->work_queue, copy_fun, &ct, n_threads);
    } else {
        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 {
-        worker_pool_run_func(octx->ctx->worker_pool, copy_fun, &ct, n_threads);
-    }
-
-    const struct htp_tensor *sync = octx->src[1];
-    if (sync && (sync->flags & HTP_TENSOR_FENCE)) {
-        if (!use_dma) {
-            // htp_tensor_flush_all(octx->ctx, octx->dsts, 1);
-            qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
-        }
-
-        atomic_uint * sync_fence = (atomic_uint *) sync->data;
-        const uint32_t seq = (uint32_t) octx->op_params[0];
-
-        atomic_store(&sync_fence[0], seq);
-        asm volatile ("syncht" : : : "memory");
-        Q6_dccleaninva_A((void *) sync_fence);
-
-        FARF(HIGH, "ggml-hex: sync-release : fence %p seq %u\n", sync_fence, seq);
-    }
-
    return HTP_STATUS_OK;
 }
+
+int op_cpy(struct htp_ops_context * octx) {
+    bool use_dma = false;
+    int status = exec_cpy(octx, &use_dma);
+
+    htp_ops_context_set_status(octx, status);
+
+    if (octx->op == HTP_OP_CPY_FENCE) {
+        if (!use_dma) {
+            htp_flush_dirty_ranges(octx->ctx);
+        }
+
+        htp_mdev_group_barrier(octx);
+
+        if (octx->ctx->mdev.idx == 0) {
+            const struct htp_tensor * sync = octx->src[1];
+            const uint32_t seq = (uint32_t) octx->op_params[0];
+            atomic_uint * sync_fence = (atomic_uint *) (uintptr_t) sync->data;
+            htp_fence_write(sync_fence, seq, octx->status);
+
+            FARF(HIGH, "ggml-hex: sync-release : fence %p seq 0x%x status %d\n", sync_fence, seq, octx->status);
+        }
+    }
+
+    return octx->status;
+}
@@ -7,6 +7,8 @@

 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+#include "hex-common.h"
+#include "hex-profile.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
@@ -17,25 +19,25 @@
 #define htp_cumsum_tensors_preamble                         \
    const struct htp_tensor * restrict src0 = octx->src[0]; \
    const struct htp_tensor * restrict dst  = octx->dst;    \
-                                                     \
-    const uint32_t ne00 = src0->ne[0];               \
-    const uint32_t ne01 = src0->ne[1];               \
-    const uint32_t ne02 = src0->ne[2];               \
-    const uint32_t ne03 = src0->ne[3];               \
-                                                     \
-    const uint32_t ne0 = dst->ne[0];                 \
-    const uint32_t ne1 = dst->ne[1];                 \
-    const uint32_t ne2 = dst->ne[2];                 \
-    const uint32_t ne3 = dst->ne[3];                 \
-                                                     \
-    const uint32_t nb00 = src0->nb[0];               \
-    const uint32_t nb01 = src0->nb[1];               \
-    const uint32_t nb02 = src0->nb[2];               \
-    const uint32_t nb03 = src0->nb[3];               \
-                                                     \
-    const uint32_t nb0 = dst->nb[0];                 \
-    const uint32_t nb1 = dst->nb[1];                 \
-    const uint32_t nb2 = dst->nb[2];                 \
+                                                            \
+    const uint32_t ne00 = src0->ne[0];                      \
+    const uint32_t ne01 = src0->ne[1];                      \
+    const uint32_t ne02 = src0->ne[2];                      \
+    const uint32_t ne03 = src0->ne[3];                      \
+                                                            \
+    const uint32_t ne0 = dst->ne[0];                        \
+    const uint32_t ne1 = dst->ne[1];                        \
+    const uint32_t ne2 = dst->ne[2];                        \
+    const uint32_t ne3 = dst->ne[3];                        \
+                                                            \
+    const uint32_t nb00 = src0->nb[0];                      \
+    const uint32_t nb01 = src0->nb[1];                      \
+    const uint32_t nb02 = src0->nb[2];                      \
+    const uint32_t nb03 = src0->nb[3];                      \
+                                                            \
+    const uint32_t nb0 = dst->nb[0];                        \
+    const uint32_t nb1 = dst->nb[1];                        \
+    const uint32_t nb2 = dst->nb[2];                        \
    const uint32_t nb3 = dst->nb[3];

 struct htp_cumsum_context {
@@ -46,6 +48,7 @@ struct htp_cumsum_context {
    size_t          dst_row_size_aligned;
    uint32_t        rows_per_thread;
    uint32_t        total_rows;
+    uint32_t        row_start;
 };

 #define htp_cumsum_preamble                                                \
@@ -116,11 +119,8 @@ static inline void hvx_cumsum_row_f32(const float * restrict src, float * restri
 static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * data) {
    htp_cumsum_preamble;

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
-    const uint32_t ir0 = cctx->rows_per_thread * ith;
-    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_thread, cctx->total_rows);
+    const uint32_t ir0 = cctx->row_start + cctx->rows_per_thread * ith;
+    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_thread, cctx->row_start + cctx->total_rows);

    if (ir0 >= ir1) {
        return;
@@ -149,11 +149,15 @@ static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * dat
                                   src_row_size_aligned, src_row_size, 1);
    }

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
    for (uint32_t ir = ir0; ir < ir1; ir++) {
        float * dst_spad_row = (float *) dma_queue_pop(dma_queue).src;
        float * src_spad_row = (float *) dma_queue_pop(dma_queue).dst;

+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
        hvx_cumsum_row_f32(src_spad_row, dst_spad_row, ne00);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);

        dma_queue_push_vtcm_to_ddr(dma_queue,
                                   dma_make_ptr(dst_data + (ir * dst_row_size), (uint8_t *) dst_spad_row),
@@ -168,12 +172,10 @@ static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * dat
    }

    dma_queue_flush(dma_queue);
-    t2 = HAP_perf_get_qtimer_count();

-    FARF(HIGH, "cumsum-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n",
+    FARF(HIGH, "cumsum-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n",
         ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1,
-         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
 }

 // ---------------------------------------------------------------------------
@@ -183,14 +185,14 @@ static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * dat
 static void cumsum_thread_f32(unsigned int nth, unsigned int ith, void * data) {
    htp_cumsum_preamble;

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
    const uint8_t * src_data = (const uint8_t *) src0->data;
    uint8_t *       dst_data = (uint8_t *) dst->data;

-    const uint32_t ir0 = cctx->rows_per_thread * ith;
-    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_thread, cctx->total_rows);
+    const uint32_t ir0 = cctx->row_start + cctx->rows_per_thread * ith;
+    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_thread, cctx->row_start + cctx->total_rows);
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);

    for (uint32_t ir = ir0; ir < ir1; ir++) {
        const float * restrict src_row = (const float *) (src_data + ir * cctx->src_row_size);
@@ -198,12 +200,11 @@ static void cumsum_thread_f32(unsigned int nth, unsigned int ith, void * data) {
        hvx_cumsum_row_f32(src_row, dst_row, ne00);
    }

-    t2 = HAP_perf_get_qtimer_count();
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);

-    FARF(HIGH, "cumsum-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n",
+    FARF(HIGH, "cumsum-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n",
         ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1,
-         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
 }

 int op_cumsum_f32(struct htp_ops_context * octx) {
@@ -214,8 +215,25 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

-    const uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3];
-    const uint32_t n_threads  = MIN(octx->n_threads, total_rows);
+    const uint32_t total_rows      = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
+
+    uint32_t row_start = 0;
+    uint32_t nrows     = total_rows;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_data_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;

    const size_t src_row_size         = src0->nb[1];
    const size_t dst_row_size         = dst->nb[1];
@@ -240,14 +258,15 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
        .dst_row_size         = dst_row_size,
        .src_row_size_aligned = src_row_size_aligned,
        .dst_row_size_aligned = dst_row_size_aligned,
-        .rows_per_thread      = (total_rows + n_threads - 1) / n_threads,
-        .total_rows           = total_rows,
+        .rows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+        .total_rows           = nrows,
+        .row_start            = row_start,
    };

    if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
-        worker_pool_run_func(octx->ctx->worker_pool, cumsum_thread_f32, &cctx, n_threads);
+        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32, &cctx, n_threads);
    } else {
-        worker_pool_run_func(octx->ctx->worker_pool, cumsum_thread_f32_dma, &cctx, n_threads);
+        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32_dma, &cctx, n_threads);
    }

    return HTP_STATUS_OK;
@@ -5,8 +5,11 @@

 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+#include "hex-common.h"
+#include "hex-profile.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 #include "hvx-types.h"
 #include "hex-utils.h"
 #include "hvx-copy.h"
@@ -15,17 +18,17 @@
 #define htp_diag_tensors_preamble                           \
    const struct htp_tensor * restrict src0 = octx->src[0]; \
    const struct htp_tensor * restrict dst  = octx->dst;    \
-                                                     \
-    const uint32_t ne02 = src0->ne[2];               \
-                                                     \
-    const uint32_t ne0 = dst->ne[0];                 \
-    const uint32_t ne1 = dst->ne[1];                 \
-                                                     \
-    const uint32_t nb02 = src0->nb[2];               \
-    const uint32_t nb03 = src0->nb[3];               \
-                                                     \
-    const uint32_t nb1 = dst->nb[1];                 \
-    const uint32_t nb2 = dst->nb[2];                 \
+                                                            \
+    const uint32_t ne02 = src0->ne[2];                      \
+                                                            \
+    const uint32_t ne0 = dst->ne[0];                        \
+    const uint32_t ne1 = dst->ne[1];                        \
+                                                            \
+    const uint32_t nb02 = src0->nb[2];                      \
+    const uint32_t nb03 = src0->nb[3];                      \
+                                                            \
+    const uint32_t nb1 = dst->nb[1];                        \
+    const uint32_t nb2 = dst->nb[2];                        \
    const uint32_t nb3 = dst->nb[3];

 struct htp_diag_context {
@@ -36,6 +39,7 @@ struct htp_diag_context {
    size_t          dst_row_size_aligned;
    uint32_t        batches_per_thread;
    uint32_t        total_batches;
+    uint32_t        batch_start;
 };

 #define htp_diag_preamble                                              \
@@ -57,11 +61,8 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
    htp_diag_preamble;
    dma_queue * dma_queue = octx->ctx->dma[ith];

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
-    const uint32_t ib0 = dctx->batches_per_thread * ith;
-    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->total_batches);
+    const uint32_t ib0 = dctx->batch_start + dctx->batches_per_thread * ith;
+    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->batch_start + dctx->total_batches);

    if (ib0 >= ib1) {
        return;
@@ -79,6 +80,8 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
    uint8_t * src_spad = octx->src0_spad.data + (ith * src_batch_size_aligned);
    uint8_t * dst_spad = octx->dst_spad.data  + (ith * dst_row_size_aligned);

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
    for (uint32_t ib = ib0; ib < ib1; ib++) {
        const uint32_t i3 = ib / ne02;
        const uint32_t i2 = ib % ne02;
@@ -96,7 +99,9 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)

        for (uint32_t i1 = 0; i1 < ne1; i1++) {
            // Compute row in VTCM
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (ib * ne1 + i1));
            hvx_diag_row_f32(src_spad_f32, dst_spad_f32, i1, ne0);
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (ib * ne1 + i1));

            // Write completed row back to DDR
            uint8_t * dst_row = dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1;
@@ -107,12 +112,9 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
        }
    }

-    t2 = HAP_perf_get_qtimer_count();
-
-    FARF(HIGH, "diag-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n",
+    FARF(HIGH, "diag-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n",
         ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ib0, ib1,
-         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
 }

 // ---------------------------------------------------------------------------
@@ -122,14 +124,14 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
 static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
    htp_diag_preamble;

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
    const uint8_t * src_data = (const uint8_t *) src0->data;
    uint8_t *       dst_data = (uint8_t *) dst->data;

-    const uint32_t ib0 = dctx->batches_per_thread * ith;
-    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->total_batches);
+    const uint32_t ib0 = dctx->batch_start + dctx->batches_per_thread * ith;
+    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->batch_start + dctx->total_batches);
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ib0);

    for (uint32_t ib = ib0; ib < ib1; ib++) {
        const uint32_t i3 = ib / ne02;
@@ -143,12 +145,11 @@ static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
        }
    }

-    t2 = HAP_perf_get_qtimer_count();
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ib0);

-    FARF(HIGH, "diag-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u usec %u\n",
+    FARF(HIGH, "diag-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n",
         ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ib0, ib1,
-         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
 }

 int op_diag_f32(struct htp_ops_context * octx) {
@@ -160,7 +161,36 @@ int op_diag_f32(struct htp_ops_context * octx) {
    }

    const uint32_t total_batches = src0->ne[2] * src0->ne[3];
-    const uint32_t n_threads     = MIN(octx->n_threads, total_batches);
+    const size_t dst_batch_size  = dst->ne[1] * dst->nb[1];
+
+    uint32_t batch_start = 0;
+    uint32_t nbatches    = total_batches;
+
+    if (octx->ctx->mdev.count > 1) {
+        bool can_split = htp_tensor_mdev_data_aligned(dst) && (dst->ne[0] == 1 || dst->nb[0] == sizeof(float)) && !htp_tensor_is_permuted(dst);
+        uint32_t batches_per_chunk = 1;
+        if (can_split) {
+            if (dst->ne[2] > 1 && (dst->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0 &&
+                (dst->ne[3] <= 1 || (dst->nb[3] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0)) {
+                batches_per_chunk = 1;
+            } else if (dst->nb[2] == dst_batch_size &&
+                       (dst->ne[3] <= 1 || dst->nb[3] == dst->nb[2] * dst->ne[2])) {
+                batches_per_chunk = (dst_batch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(dst_batch_size, HEX_L2_LINE_SIZE)) : 1;
+            } else {
+                can_split = false;
+            }
+        }
+
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, can_split ? batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        batch_start = range.start;
+        nbatches    = range.count;
+    }
+
+    if (nbatches == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads     = octx->n_threads;

    const size_t src_batch_size         = src0->ne[0] * sizeof(float);
    const size_t dst_row_size           = dst->ne[0] * sizeof(float);
@@ -185,14 +215,15 @@ int op_diag_f32(struct htp_ops_context * octx) {
        .dst_row_size           = dst_row_size,
        .src_batch_size_aligned = src_batch_size_aligned,
        .dst_row_size_aligned   = dst_row_size_aligned,
-        .batches_per_thread     = (total_batches + n_threads - 1) / n_threads,
-        .total_batches          = total_batches,
+        .batches_per_thread     = fastdiv(nbatches + n_threads - 1, &octx->n_threads_div),
+        .total_batches          = nbatches,
+        .batch_start            = batch_start,
    };

    if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
-        worker_pool_run_func(octx->ctx->worker_pool, diag_thread_f32, &dctx, n_threads);
+        work_queue_run(octx->ctx->work_queue, diag_thread_f32, &dctx, n_threads);
    } else {
-        worker_pool_run_func(octx->ctx->worker_pool, diag_thread_f32_dma, &dctx, n_threads);
+        work_queue_run(octx->ctx->work_queue, diag_thread_f32_dma, &dctx, n_threads);
    }

    return HTP_STATUS_OK;
@@ -3,10 +3,11 @@
 #pragma clang diagnostic ignored "-Wunused-but-set-variable"

 #include <HAP_farf.h>
-#include <HAP_perf.h>
-
 #include <string.h>

+#include "hex-common.h"
+#include "hex-profile.h"
+
 #include "hvx-copy.h"
 #include "hvx-utils.h"

@@ -14,28 +15,30 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"

 // ggml op_params layout for FILL:
 //   op_params[0] (as float) - the scalar fill value

-#define fill_preamble \
+#define fill_preamble                          \
    const struct htp_tensor * dst = octx->dst; \
-    \
-    const uint32_t ne0 = dst->ne[0]; \
-    const uint32_t ne1 = dst->ne[1]; \
-    const uint32_t ne2 = dst->ne[2]; \
-    const uint32_t ne3 = dst->ne[3]; \
-    \
-    const uint32_t nb1 = dst->nb[1]; \
-    const uint32_t nb2 = dst->nb[2]; \
-    const uint32_t nb3 = dst->nb[3]; \
-    \
+                                               \
+    const uint32_t ne0 = dst->ne[0];           \
+    const uint32_t ne1 = dst->ne[1];           \
+    const uint32_t ne2 = dst->ne[2];           \
+    const uint32_t ne3 = dst->ne[3];           \
+                                               \
+    const uint32_t nb1 = dst->nb[1];           \
+    const uint32_t nb2 = dst->nb[2];           \
+    const uint32_t nb3 = dst->nb[3];           \
+                                               \
    const uint32_t nr = ne1 * ne2 * ne3;

 struct htp_fill_context {
    struct htp_ops_context * octx;
    uint32_t nrows_per_thread;
    uint32_t total_rows;  // ne1 * ne2 * ne3
+    uint32_t row_start;
    bool     opt_path;
    HVX_Vector splat_vec;
    uint32_t   elem_size;
@@ -47,10 +50,15 @@ static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
    fill_preamble;

    // Parallelise over the flat row index spanning ne1*ne2*ne3
-    const uint32_t ir0 = fctx->nrows_per_thread * ith;
-    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_thread, fctx->total_rows);
+    const uint32_t ir0 = fctx->row_start + fctx->nrows_per_thread * ith;
+    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_thread, fctx->row_start + fctx->total_rows);

-    uint64_t t1 = HAP_perf_get_qtimer_count();
+    if (ir0 >= ir1) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);

    if (fctx->opt_path) {
        // Opt path: tensor is fully contiguous, treat as flat array
@@ -69,9 +77,8 @@ static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
        }
    }

-    uint64_t t2 = HAP_perf_get_qtimer_count();
-    FARF(HIGH, "fill %u/%u: rows %u:%u usec %u\n",
-         ith, nth, ir0, ir1, (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1);
+    FARF(HIGH, "fill %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
 }

 int op_fill(struct htp_ops_context * octx) {
@@ -85,8 +92,23 @@ int op_fill(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

+    uint32_t row_start = 0;
+    uint32_t nrows     = nr;
+
+    if (octx->ctx->mdev.count > 1) {
+        const uint32_t row_size = nb1;
+        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(nr, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
    // nr = ne1*ne2*ne3 (flat row count across all outer dims); parallelise over it.
-    const uint32_t n_threads = MIN(nr, octx->n_threads);
+    const uint32_t n_threads = octx->n_threads;

    // Optimize if fully contiguous: skip stride arithmetic, treat as flat array
    const bool opt_path = (nb2 == nb1 * ne1) && (nb3 == nb2 * ne2);
@@ -99,8 +121,9 @@ int op_fill(struct htp_ops_context * octx) {

    struct htp_fill_context fctx = {
        .octx             = octx,
-        .nrows_per_thread = (nr + n_threads - 1) / n_threads,
-        .total_rows       = nr,
+        .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+        .total_rows       = nrows,
+        .row_start        = row_start,
        .opt_path         = opt_path,
    };

@@ -117,7 +140,7 @@ int op_fill(struct htp_ops_context * octx) {
        return HTP_STATUS_NO_SUPPORT;
    }

-    worker_pool_run_func(octx->ctx->worker_pool, fill_thread, &fctx, n_threads);
+    work_queue_run(octx->ctx->work_queue, fill_thread, &fctx, n_threads);

    return HTP_STATUS_OK;
 }
@@ -5,7 +5,6 @@
 #include <assert.h>
 #include <HAP_compute_res.h>
 #include <HAP_farf.h>
-#include <HAP_perf.h>
 #include <math.h>
 #include <stdbool.h>
 #include <stdatomic.h>
@@ -75,6 +74,7 @@ struct htp_fa_context {

    uint32_t qrows;
    uint32_t qrows_per_thread;
+    uint32_t qrow_start;

    bool is_q_fp32;

@@ -89,8 +89,6 @@ struct htp_fa_context {

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

 struct hmx_fa_context {
@@ -206,10 +204,9 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
    const uint32_t nb3 = dst->nb[3];

    // total rows in q
-    const uint32_t nr = factx->qrows;
-    const uint32_t dr = factx->qrows_per_thread;
-    const uint32_t ir0 = dr * ith;
-    const uint32_t ir1 = MIN(ir0 + dr, nr);
+    const uint32_t dr  = factx->qrows_per_thread;
+    const uint32_t ir0 = factx->qrow_start + dr * ith;
+    const uint32_t ir1 = MIN(ir0 + dr, factx->qrow_start + factx->qrows);

    if (ir0 >= ir1) return;

@@ -1888,6 +1885,24 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
    const uint32_t n_threads = factx.n_threads;
    const uint32_t G = factx.G;

+    // Multi-device: split Q blocks across devices
+    const uint32_t n_q_blocks = (neq1 + Br - 1) / Br;
+    uint32_t q_start_min = 0;
+    uint32_t q_start_max = neq1;
+
+    if (octx->ctx->mdev.count > 1) {
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(n_q_blocks, htp_tensor_mdev_data_aligned(dst) ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        const uint32_t block_start = range.start;
+        const uint32_t block_end   = range.start + range.count;
+
+        if (block_start >= block_end) {
+            return HTP_STATUS_OK;
+        }
+
+        q_start_min = block_start * Br;
+        q_start_max = MIN(block_end * Br, neq1);
+    }
+
    // ======== VTCM allocation (GQA-aware) ========
    // K/V row sizes drive the DMA descriptors (not the VTCM layout) and are used
    // throughout the KV loop below.
@@ -1977,7 +1992,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
    // ======== Main loop ========
    for (uint32_t ib3 = 0; ib3 < neq3; ++ib3) {
        const uint32_t im3 = mask ? fastmodulo(ib3, mask->ne[3], &factx.src3_div3) : 0;
-        for (uint32_t q_start = 0; q_start < neq1; q_start += Br) {
+        for (uint32_t q_start = q_start_min; q_start < q_start_max; q_start += Br) {
            const uint32_t n_rows_q    = hex_smin(Br, neq1 - q_start);
            const size_t   n_rows_g    = n_rows_q * G;
            const size_t   g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
@@ -1991,8 +2006,9 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {

                // 1. Push Q and KV DMAs for the very first iteration.
                // Subsequent iterations are enqueued early at the end of the previous iteration.
-                if (ib3 == 0 && q_start == 0 && kv_head == 0) {
-                    const uint8_t * q_ptr = (const uint8_t *) q->data;
+                if (ib3 == 0 && q_start == q_start_min && kv_head == 0) {
+                    const uint8_t * q_ptr = (const uint8_t *) q->data + q_start * q->nb[1] +
+                                            (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3];
                    const size_t q_row_bytes = q_transposed ? n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
                    const size_t n_rows      = q_transposed ? factx.G : n_rows_q;
                    dma_queue_push(dma, dma_make_ptr(factx.vtcm_q_dma, q_ptr), q_row_bytes, hex_smax(q_src_stride, q_row_bytes), q_row_bytes, n_rows);
@@ -2311,8 +2327,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                if (next_kv_head >= n_kv_heads) {
                    next_kv_head = 0;
                    next_q_start = q_start + Br;
-                    if (next_q_start >= neq1) {
-                        next_q_start = 0;
+                    if (next_q_start >= q_start_max) {
+                        next_q_start = q_start_min;
                        next_ib3     = ib3 + 1;
                    }
                }
@@ -2398,6 +2414,10 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
        return HTP_STATUS_NO_SUPPORT;
    }

+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
    if (kparams->kernel_type == HTP_FA_KERNEL_HMX) {
        return hmx_flash_attn_ext(octx);
    }
@@ -2407,8 +2427,6 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
    factx.k = k;
    factx.v = v;

-    factx.t_start = HAP_perf_get_qtimer_count();
-
    factx.src0_div21 = kparams->u.hvx.src0_div21;
    factx.src0_div1  = kparams->u.hvx.src0_div1;

@@ -2451,8 +2469,30 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
    }

    // total rows in q
-    factx.qrows = kparams->qrows;
-    factx.qrows_per_thread = kparams->qrows_per_thread;
+    const uint32_t neq1 = q->ne[1];
+    const uint32_t neq2 = q->ne[2];
+    const uint32_t neq3 = q->ne[3];
+    const uint32_t total_qrows = neq1 * neq2 * neq3;
+
+    uint32_t qrow_start = 0;
+    uint32_t qrows      = total_qrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_mdev_data_aligned(dst) && ((dst->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_qrows, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        qrow_start = range.start;
+        qrows      = range.count;
+    }
+
+    if (qrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
+    factx.qrows            = qrows;
+    factx.qrow_start       = qrow_start;
+    factx.qrows_per_thread = fastdiv(qrows + n_threads - 1, &octx->n_threads_div);

    size_t size_vkq_acc = hex_round_up(v->ne[0] * sizeof(float), 128); // VKQ32

@@ -2461,18 +2501,18 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {

    uint8_t * vtcm_cur = octx->ctx->vtcm_base;

-    factx.spad_q = vtcm_seq_alloc(&vtcm_cur, size_q_block * octx->n_threads);
-    factx.spad_k = vtcm_seq_alloc(&vtcm_cur, factx.size_k_block * 2 * octx->n_threads);
-    factx.spad_v = vtcm_seq_alloc(&vtcm_cur, factx.size_v_block * 2 * octx->n_threads);
-    factx.spad_m = vtcm_seq_alloc(&vtcm_cur, (mask ? factx.size_m_block * HVX_FA_DMA_CACHE_SIZE : 0) * octx->n_threads);
-    factx.spad_a = vtcm_seq_alloc(&vtcm_cur, size_vkq_acc * octx->n_threads);
+    factx.spad_q = vtcm_seq_alloc(&vtcm_cur, size_q_block * n_threads);
+    factx.spad_k = vtcm_seq_alloc(&vtcm_cur, factx.size_k_block * 2 * n_threads);
+    factx.spad_v = vtcm_seq_alloc(&vtcm_cur, factx.size_v_block * 2 * n_threads);
+    factx.spad_m = vtcm_seq_alloc(&vtcm_cur, (mask ? factx.size_m_block * HVX_FA_DMA_CACHE_SIZE : 0) * n_threads);
+    factx.spad_a = vtcm_seq_alloc(&vtcm_cur, size_vkq_acc * n_threads);

    if ((size_t) (vtcm_cur - octx->ctx->vtcm_base) > octx->ctx->vtcm_size) {
        return HTP_STATUS_VTCM_TOO_SMALL;
    }

    if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
-        work_queue_run(octx->ctx->work_queue, flash_attn_ext_f16_thread, &factx, octx->n_threads);
+        work_queue_run(octx->ctx->work_queue, flash_attn_ext_f16_thread, &factx, n_threads);
    }

    return HTP_STATUS_OK;
@@ -51,6 +51,7 @@ struct htp_fa_kernel_params {

    uint32_t qrows;
    uint32_t qrows_per_thread;
+    uint32_t qrow_start;
    float    m0;
    float    m1;
    uint32_t n_head_log2;
@@ -4,10 +4,13 @@

 #include "hvx-utils.h"
 #include "hex-fastdiv.h"
+#include "hex-common.h"
+#include "hex-profile.h"

 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
 #include "htp-ctx.h"
+#include "htp-tensor.h"

 #ifndef MIN
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -22,6 +25,8 @@ struct htp_gdn_context {
    size_t   state_bytes;
    uint8_t * vtcm_base;
    size_t   vtcm_per_thread;
+    uint32_t row_start;
+    uint32_t nrows;
 };

 static inline HVX_Vector gdn_mul_dot_f32(float * restrict dst, const float * restrict mul, const float * restrict dot, uint32_t n) {
@@ -586,8 +591,9 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
    const uint32_t n_seqs   = v->ne[3];
    const uint32_t K        = octx->op_params[0];

-    const uint32_t total_rows = H * n_seqs;
-    if (ith >= total_rows) {
+    const uint32_t row_end = gctx->row_start + gctx->nrows;
+
+    if (ith >= gctx->nrows) {
        return;
    }

@@ -621,11 +627,11 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
    const uint64_t state_seq_stride = state->nb[3] / sizeof(float);
    const uint64_t state_size_per_snap = (uint64_t) S_v * S_v * H * n_seqs;

-    uint32_t ir_prefetch = ith;
+    uint32_t ir_prefetch = gctx->row_start + ith;
    int spad_idx = 0;

    // Prefetch preamble (up to 2 steps)
-    for (int k = 0; k < 2 && ir_prefetch < total_rows; k++) {
+    for (int k = 0; k < 2 && ir_prefetch < row_end; k++) {
        const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H);
        const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H);
        const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v;
@@ -646,8 +652,11 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
        spad_idx ^= 1;
    }

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (gctx->row_start + ith));
+
    int curr_spad_idx = 0;
-    for (uint32_t ir = ith; ir < total_rows; ir += nth) {
+    for (uint32_t ir = gctx->row_start + ith; ir < row_end; ir += nth) {
        dma_queue_pop(dma);
        dma_queue_pop(dma);

@@ -812,7 +821,7 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
                       S_v * sizeof(float), S_v);

        // Prefetch next block (if any)
-        if (ir_prefetch < total_rows) {
+        if (ir_prefetch < row_end) {
            const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H);
            const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H);
            const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v;
@@ -828,6 +837,7 @@ static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, vo
        curr_spad_idx ^= 1;
    }
    dma_queue_flush(dma);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row_end);
 }


@@ -847,8 +857,9 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo
    const uint32_t H        = v->ne[1];
    const uint32_t n_seqs   = v->ne[3];

-    const uint32_t total_rows = H * n_seqs;
-    if (ith >= total_rows) {
+    const uint32_t row_end = gctx->row_start + gctx->nrows;
+
+    if (ith >= gctx->nrows) {
        return;
    }

@@ -881,11 +892,11 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo

    const uint64_t state_seq_stride = state->nb[3] / sizeof(float);

-    uint32_t ir_prefetch = ith;
+    uint32_t ir_prefetch = gctx->row_start + ith;
    int spad_idx = 0;

    // Prefetch preamble (up to 2 steps)
-    for (int k = 0; k < 2 && ir_prefetch < total_rows; k++) {
+    for (int k = 0; k < 2 && ir_prefetch < row_end; k++) {
        const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H);
        const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H);
        const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v;
@@ -906,8 +917,11 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo
        spad_idx ^= 1;
    }

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (gctx->row_start + ith));
+
    int curr_spad_idx = 0;
-    for (uint32_t ir = ith; ir < total_rows; ir += nth) {
+    for (uint32_t ir = gctx->row_start + ith; ir < row_end; ir += nth) {
        dma_queue_pop(dma);
        dma_queue_pop(dma);

@@ -1057,7 +1071,7 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo
                       S_v * sizeof(float), S_v);

        // Prefetch next block (if any)
-        if (ir_prefetch < total_rows) {
+        if (ir_prefetch < row_end) {
            const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H);
            const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H);
            const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v;
@@ -1073,6 +1087,7 @@ static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, vo
        curr_spad_idx ^= 1;
    }
    dma_queue_flush(dma);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row_end);
 }


@@ -1085,10 +1100,6 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
    const struct htp_tensor * state = octx->src[5];
    const struct htp_tensor * dst   = octx->dst;

-    if (!q || !k || !v || !g || !beta || !state || !dst) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-
    if (q->type != HTP_TYPE_F32 || k->type != HTP_TYPE_F32 || v->type != HTP_TYPE_F32 ||
        g->type != HTP_TYPE_F32 || beta->type != HTP_TYPE_F32 || state->type != HTP_TYPE_F32 ||
        dst->type != HTP_TYPE_F32) {
@@ -1124,16 +1135,37 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

+    const uint32_t total_rows = H * n_seqs;
+
+    uint32_t row_start = 0;
+    uint32_t nrows     = total_rows;
+
+    if (octx->ctx->mdev.count > 1) {
+        const uint32_t head_bytes = S_v * sizeof(float);
+        const uint32_t rows_per_chunk = (head_bytes > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(head_bytes, HEX_L2_LINE_SIZE)) : 1;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0,
+                                                                             octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
    struct htp_gdn_context gctx;
    gctx.octx = octx;
-    gctx.rows_per_thread = (H * n_seqs + octx->n_threads - 1) / octx->n_threads;
+    gctx.row_start       = row_start;
+    gctx.nrows           = nrows;
+    gctx.rows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
    gctx.state_bytes = (size_t) S_v * S_v * sizeof(float);

    size_t state_aligned = (size_t) S_v * S_v * sizeof(float);
    state_aligned = (state_aligned + 127) & ~(size_t)127;

-    assert(octx->ctx->vtcm_base != NULL);
-    assert(octx->ctx->vtcm_size >= 2 * state_aligned * octx->n_threads);
+    assert(octx->ctx->vtcm_size >= 2 * state_aligned * n_threads);

    gctx.vtcm_base = octx->ctx->vtcm_base;
    gctx.vtcm_per_thread = 2 * state_aligned;
@@ -1148,9 +1180,9 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
         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);
+        work_queue_run(octx->ctx->work_queue, gated_delta_net_f32_tg_thread, &gctx, n_threads);
    } else {
-        worker_pool_run_func(octx->ctx->worker_pool, gated_delta_net_f32_pp_thread, &gctx, octx->n_threads);
+        work_queue_run(octx->ctx->work_queue, gated_delta_net_f32_pp_thread, &gctx, n_threads);
    }

    return HTP_STATUS_OK;
@@ -10,6 +10,7 @@

 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+#include "hex-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
@@ -23,9 +24,12 @@ struct get_rows_context {
    const struct htp_get_rows_kernel_params * kparams;
    struct htp_get_rows_vtcm_layout vtcm_layout;
    uint8_t * vtcm_base;
+    uint32_t task_start;
+    uint32_t tasks;
+    uint32_t tasks_per_thread;
 };

-#define get_rows_preamble \
+#define get_rows_preamble                      \
    const uint32_t ne00 = octx->src[0]->ne[0]; \
    const uint32_t ne01 = octx->src[0]->ne[1]; \
    const uint32_t ne02 = octx->src[0]->ne[2]; \
@@ -61,12 +65,12 @@ static void get_rows_thread_st_##IDX_TYPE(unsigned int nth, unsigned int ith, vo
    struct htp_ops_context * octx = grctx->octx;                                                                       \
    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;                                                \
    get_rows_preamble;                                                                                                 \
-    const uint32_t dr  = kparams->tasks_per_thread;                                                                    \
-    const uint32_t ir0 = dr * ith;                                                                                     \
-    if (ir0 >= kparams->total_tasks) {                                                                                 \
+    const uint32_t dr  = grctx->tasks_per_thread;                                                                      \
+    const uint32_t ir0 = grctx->task_start + dr * ith;                                                                 \
+    if (ir0 >= grctx->task_start + grctx->tasks) {                                                                     \
        return;                                                                                                        \
    }                                                                                                                  \
-    const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks);                                                          \
+    const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);                                              \
    const uint32_t row_size_bytes = htp_tensor_get_row_size(octx->src[0]->type, ne00);                                 \
    dma_queue * dma_queue = octx->ctx->dma[ith];                                                                       \
    for (uint32_t i = ir0; i < ir1; ++i) {                                                                             \
@@ -101,12 +105,12 @@ static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned
    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;                                                \
    get_rows_preamble;                                                                                                 \
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                             \
-    const uint32_t dr  = kparams->tasks_per_thread;                                                                    \
-    const uint32_t ir0 = dr * ith;                                                                                     \
-    if (ir0 >= kparams->total_tasks) {                                                                                 \
+    const uint32_t dr  = grctx->tasks_per_thread;                                                                      \
+    const uint32_t ir0 = grctx->task_start + dr * ith;                                                                 \
+    if (ir0 >= grctx->task_start + grctx->tasks) {                                                                     \
        return;                                                                                                        \
    }                                                                                                                  \
-    const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks);                                                          \
+    const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);                                              \
    const uint32_t chunks_per_row = kparams->chunks_per_row;                                                           \
    const uint32_t chunk_size     = kparams->chunk_size;                                                               \
    dma_queue * dma_queue = octx->ctx->dma[ith];                                                                       \
@@ -225,13 +229,41 @@ int op_get_rows(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

+    const struct htp_tensor * dst = octx->dst;
+    const uint32_t total_tasks    = kparams->total_tasks;
+    const size_t dst_row_size     = htp_tensor_get_row_size(dst->type, dst->ne[0]);
+
+    uint32_t task_start = 0;
+    uint32_t tasks      = total_tasks;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t tasks_per_chunk = 1;
+        htp_tensor_mdev_rows_per_chunk(dst, dst_row_size / dst->ne[0], (uint32_t) dst_row_size, &tasks_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_tasks, tasks_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        task_start = range.start;
+        tasks      = range.count;
+    }
+
+    if (tasks == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
    struct get_rows_context grctx;
    grctx.octx = octx;
    grctx.kparams = kparams;
    grctx.vtcm_base = (uint8_t *)octx->ctx->vtcm_base;
+    grctx.task_start = task_start;
+    grctx.tasks = tasks;
+    grctx.tasks_per_thread = fastdiv(tasks + n_threads - 1, &octx->n_threads_div);

    const uint32_t ne00 = octx->src[0]->ne[0];
-    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, octx->src[0]->type, ne00, kparams->n_threads);
+    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, octx->src[0]->type, ne00, n_threads);

    const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);

@@ -247,14 +279,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",
+    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);
+         grctx.vtcm_layout.src0_bytes_per_thread * n_threads,
+         grctx.vtcm_layout.dst_bytes_per_thread  * n_threads,
+         kparams->use_dma, n_threads);

-    work_queue_run(octx->ctx->work_queue, q_func, &grctx, kparams->n_threads);
+    work_queue_run(octx->ctx->work_queue, q_func, &grctx, n_threads);
    return HTP_STATUS_OK;
 }
@@ -77,4 +77,13 @@ static inline bool hex_add_overflow(size_t a, size_t b, size_t *out) {
    return false;
 }

+static inline uint32_t hex_gcd_u32(uint32_t a, uint32_t b) {
+    while (b != 0) {
+        uint32_t t = b;
+        b = a % b;
+        a = t;
+    }
+    return a;
+}
+
 #endif // HEX_COMMON_H
@@ -39,7 +39,6 @@ static inline void hex_l2fetch_block(const void * addr, size_t size) {

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

@@ -27,7 +27,7 @@ static inline void hmx_init_column_scales(void *out_scales, HVX_Vector v_scale)
 // vscatter offsets for fused dequant+transpose: write K-values directly to [K][N] tile.
 // word[i] = i*128 maps K-row-pair i to byte offset i*128.
 // Column offset (n*4) is added at runtime.  Entries 0..15 cover one tile (region 2047);
-// entries 16..31 cover the next adjacent tile (region 4095) — pick region size at the
+// entries 16..31 cover the next adjacent tile (region 4095) - pick region size at the
 // call site to scatter into one tile (masked) or two contiguous tiles (unmasked).
 static const int32_t hmx_transpose_scatter_offsets[32] __attribute__((aligned(VLEN))) = {
    0 * 128,  1 * 128,  2 * 128,  3 * 128,  4 * 128,  5 * 128,  6 * 128,  7 * 128,  8 * 128,  9 * 128,  10 * 128,
@@ -198,16 +198,16 @@ static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out,
 }

 // --- HMX inline asm macros for load-store packetization ---
-#define HMX_LOAD_MPY_F16(act, wt, range) \
-    "{\n" \
+#define HMX_LOAD_MPY_F16(act, wt, range)              \
+    "{\n"                                             \
    "    activation.hf = mxmem(" act ", " range ")\n" \
-    "    weight.hf = mxmem(" wt ", " range ")\n" \
+    "    weight.hf = mxmem(" wt ", " range ")\n"      \
    "}\n"

-#define HMX_LOAD_MPY_DEEP_F16(act, wt, range) \
-    "{\n" \
+#define HMX_LOAD_MPY_DEEP_F16(act, wt, range)              \
+    "{\n"                                                  \
    "    activation.hf = mxmem(" act ", " range "):deep\n" \
-    "    weight.hf = mxmem(" wt ", " range ")\n" \
+    "    weight.hf = mxmem(" wt ", " range ")\n"           \
    "}\n"

 #define HMX_STORE_AFTER_F16(out, scale_reg) \
@@ -19,7 +19,7 @@
 #endif
 #define HTP_MAX_MMAPS    16

-#define HTP_MAX_DIRTY_RANGES 16
+#define HTP_MAX_DIRTY_RANGES 32

 // Memory mapping
 struct htp_mmap {
@@ -29,6 +29,11 @@ struct htp_mmap {
    uint32_t reserved;
 };

+struct htp_dirty_range {
+    uint32_t start;
+    uint32_t end;
+};
+
 // Scratchpad state
 struct htp_spad {
    const struct htp_tensor * src;             // original src of the data (for reuse)
@@ -38,6 +43,14 @@ struct htp_spad {
    uint32_t                  size_per_thread; // size per thread
 };

+struct htp_mdev_group {
+    uint16_t              idx;
+    uint16_t              count;
+    struct fastdiv_values count_div;
+    uint8_t *             fence_base;
+    uint32_t              fence_seq;
+};
+
 struct htp_context;

 // Context while processing an Op
@@ -65,8 +78,10 @@ struct htp_ops_context {
    struct htp_spad src3_spad;
    struct htp_spad dst_spad;

-    uint32_t n_threads;
-    uint32_t flags;
+    uint32_t              flags;
+    uint32_t              n_threads;
+    struct fastdiv_values n_threads_div;
+    int                   status;
 };

 // Main context for htp DSP backend
@@ -76,6 +91,7 @@ struct htp_context {
    struct htp_mmap        mmap[HTP_MAX_MMAPS];
    dma_queue_t            dma[HTP_MAX_NTHREADS];
    dma_queue_t            dma_cached[HTP_MAX_NTHREADS];
+    struct htp_thread_trace trace[HTP_MAX_NTHREADS + 1];
    work_queue_t           work_queue;
    hmx_queue_t            hmx_queue;

@@ -88,7 +104,6 @@ struct htp_context {
    bool                   hmx_enabled;
    bool                   etm;
    uint32_t               profiler;
-    struct htp_thread_trace trace[HTP_MAX_NTHREADS + 1];

    uint8_t *              vtcm_base;
    size_t                 vtcm_size;
@@ -97,16 +112,13 @@ struct htp_context {
    atomic_bool            vtcm_needs_release;

    uint64_t               max_vmem;
-    struct htp_dirty_range {
-        uint32_t start;
-        uint32_t end;
-        uint32_t bi;
-    } dirty_ranges[HTP_MAX_DIRTY_RANGES];
+    struct htp_dirty_range dirty_ranges[HTP_MAX_DIRTY_RANGES];

    // Persistent DDR scratchpad for MUL_MAT_ID mappings
    void *                 ddr_spad_base;
    size_t                 ddr_spad_size;

+    struct htp_mdev_group  mdev;
    struct htp_ops_context octx;

    qurt_thread_t          main_thread;
@@ -115,6 +127,27 @@ struct htp_context {
    size_t                 footprint;
 };

+static inline bool htp_ops_context_set_n_threads(struct htp_ops_context * octx, uint32_t n_threads) {
+    if (n_threads == 0 || n_threads > octx->ctx->n_threads) {
+        return false;
+    }
+
+    if (n_threads != octx->n_threads) {
+        octx->n_threads = n_threads;
+        octx->n_threads_div = n_threads == octx->ctx->n_threads
+            ? octx->ctx->n_threads_div
+            : init_fastdiv_values(n_threads);
+    }
+
+    return true;
+}
+
+static inline void htp_ops_context_set_status(struct htp_ops_context * octx, int status) {
+    if (status > HTP_STATUS_OK && octx->status == HTP_STATUS_OK) {
+        octx->status = status;
+    }
+}
+
 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);
@@ -0,0 +1,89 @@
+#ifndef HTP_FENCE_H
+#define HTP_FENCE_H
+
+#include <stdatomic.h>
+#include <stdint.h>
+
+#include <HAP_farf.h>
+
+#include "hex-utils.h"
+#include "htp-ops.h"
+#include "htp-ctx.h"
+
+static inline atomic_uint * htp_mdev_fence_slot(const void * fence_base, uint32_t idx) {
+    return (atomic_uint *) ((const uint8_t *) fence_base + (size_t) idx * HTP_FENCE_SLOT_SIZE);
+}
+
+static inline void htp_fence_write(void * fence_ptr, uint32_t seq, uint32_t status) {
+    atomic_uint * fence = (atomic_uint *) fence_ptr;
+    atomic_store(&fence[1], status);
+    atomic_store(&fence[0], seq);
+    asm volatile ("syncht" : : : "memory");
+    Q6_dccleaninva_A((void *) fence);
+}
+
+static inline void htp_fence_read(const void * fence_ptr, uint32_t * seq, uint32_t * status) {
+    const atomic_uint * fence = (const atomic_uint *) fence_ptr;
+    Q6_dccleaninva_A((void *) fence);
+    asm volatile ("syncht" : : : "memory");
+    *seq = atomic_load(&fence[0]);
+    *status = atomic_load(&fence[1]);
+}
+
+static inline void htp_mdev_group_barrier(struct htp_ops_context * octx) {
+    struct htp_context * ctx = octx->ctx;
+    if (ctx->mdev.count <= 1) {
+        return;
+    }
+
+    const uint32_t seq = ++ctx->mdev.fence_seq;
+
+    struct htp_thread_trace * tr = &ctx->trace[0];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);
+
+    const uint32_t mdev_idx   = ctx->mdev.idx;
+    const uint32_t mdev_count = ctx->mdev.count;
+
+    uint8_t * fence_base   = ctx->mdev.fence_base;
+    atomic_uint * my_fence = htp_mdev_fence_slot(fence_base, mdev_idx);
+    htp_fence_write(my_fence, seq, octx->status);
+
+    for (uint32_t d = 0; d < mdev_count; d++) {
+        if (d == mdev_idx) continue;
+        atomic_uint * peer_fence = htp_mdev_fence_slot(fence_base, d);
+        uint64_t spins = 0;
+        while (1) {
+            uint32_t peer_seq;
+            uint32_t peer_status;
+            htp_fence_read(peer_fence, &peer_seq, &peer_status);
+            if ((int32_t)(peer_seq - seq) >= 0) {
+                if (peer_status > HTP_STATUS_OK) {
+                    FARF(ERROR, "ggml-hex: mdev %u peer %u failed with status %u : seq 0x%08x\n",
+                         mdev_idx, d, peer_status, seq);
+                    htp_ops_context_set_status(octx, peer_status);
+                }
+                break;
+            }
+            if (++spins == 10000) {
+                FARF(ALWAYS, "ggml-hex: mdev %u waiting for mdev %u : seq 0x%08x (b %u op %u) my-fence %p peer-fence %p peer-seq 0x%08x (diff %d)\n",
+                     mdev_idx, d, seq, seq >> 12, seq & 0xfff, my_fence, peer_fence, peer_seq, (int32_t)(peer_seq - seq));
+            }
+            if (spins > HTP_FENCE_TIMEOUT) {
+                FARF(ERROR, "ggml-hex: mdev %u timeout waiting for mdev %u : seq 0x%08x (b %u op %u) peer-fence %p peer-seq 0x%08x\n",
+                     mdev_idx, d, seq, seq >> 12, seq & 0xfff, peer_fence, peer_seq);
+                htp_ops_context_set_status(octx, HTP_STATUS_INTERNAL_ERR);
+                break;
+            }
+            hex_pause();
+        }
+    }
+    asm volatile ("syncht" : : : "memory");
+
+    if (octx->status > HTP_STATUS_OK) {
+        htp_fence_write(my_fence, seq, octx->status);
+    }
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);
+}
+
+#endif // HTP_FENCE_H
@@ -77,6 +77,7 @@ enum htp_op_code {
    HTP_OP_GET_ROWS,
    HTP_OP_SCALE,
    HTP_OP_CPY,
+    HTP_OP_CPY_FENCE,
    HTP_OP_ARGSORT,
    HTP_OP_SQR,
    HTP_OP_SQRT,
@@ -100,6 +101,7 @@ enum htp_op_code {
    HTP_OP_ALLREDUCE,
    HTP_OP_ALLREDUCE_ADD,
    HTP_OP_GLU_SWIGLU_CLAMP,
+    HTP_OP_MDEV_GROUP,

    HTP_OP_INVALID
 };
@@ -114,6 +116,7 @@ enum htp_op_code {
 #define HTP_OP_MAX_TENSORS 8192 // must stay under 64K (uint16)

 #define HTP_FENCE_TIMEOUT  (1000000000ULL)
+#define HTP_FENCE_SLOT_SIZE 128

 #define HTP_OP_MAX_VMEM_DEFAULT (3355443200u)

@@ -214,30 +217,26 @@ struct htp_prof_desc {
 };

 struct htp_opbatch_req {
-    uint32_t id;          // Batch id
+    uint64_t seq;         // Sequence number
    uint32_t n_bufs;      // Number of buffers
    uint32_t n_tensors;   // Number of tensors
    uint32_t n_ops;       // Number of ops
    uint32_t n_traces;    // Number of trace descriptors per thread
-    uint32_t pad;         // unused
-    uint64_t seq;         // Sequence number
    // struct htp_buf_desc  bufs[];    -- dspqueue buf 0
    // struct htp_tensor    tensors[]; -- dspqueue buf 0
    // struct htp_op_desc   ops[];     -- dspqueue buf 0
 };

 struct htp_opbatch_rsp {
-    uint32_t id;         // Batch id
-    uint32_t status;     // HTP_STATUS_...
-    uint32_t n_bufs;     // Number of buffers
-    uint32_t n_tensors;  // Number of tensors
-    uint32_t n_ops;      // Number of op profile descriptors
-    uint32_t n_traces[HTP_MAX_NTHREADS + 1];
-    uint32_t usecs;          // Number of usec
-    uint32_t pad;            // align to 8 bytes
+    uint64_t seq;            // Sequence number
    uint64_t cycles_start;   // Start cycle counter
    uint64_t cycles_stop;    // Stop cycle counter
-    uint64_t seq;            // Sequence number
+    uint32_t status;         // HTP_STATUS_...
+    uint32_t n_bufs;         // Number of buffers
+    uint32_t n_tensors;      // Number of tensors
+    uint32_t n_ops;          // Number of op profile descriptors
+    uint32_t usecs;          // Number of usec
+    uint32_t n_traces[HTP_MAX_NTHREADS + 1];
    // struct htp_prof_desc profs[];  -- dspqueue buf 0
 };

@@ -20,7 +20,7 @@ struct l2flush_range {

 struct l2flush_multi_task {
    struct htp_thread_trace * trace;
-    struct l2flush_range      ranges[HTP_OP_MAX_INPUTS];
+    struct l2flush_range      ranges[HTP_MAX_DIRTY_RANGES];
    uint32_t                  n_ranges;
    uint32_t                  total_blocks;
    uint32_t                  blocks_per_thread;
@@ -73,6 +73,27 @@ static void l2flush_multi_worker(unsigned int n, unsigned int i, void * data) {
    htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, gb_first);
 }

+static void merge_dirty_ranges(struct htp_context * ctx) {
+    for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
+        struct htp_dirty_range * r = &ctx->dirty_ranges[i];
+        if (!r->start) continue;
+
+        for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES;) {
+            struct htp_dirty_range * s = &ctx->dirty_ranges[j];
+            if (i == j || !s->start || r->end < s->start || s->end < r->start) {
+                j++;
+                continue;
+            }
+
+            r->start = MIN(r->start, s->start);
+            r->end   = MAX(r->end, s->end);
+            s->start = 0;
+            s->end   = 0;
+            j = 0;
+        }
+    }
+}
+
 void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
    const struct htp_tensor * pending[HTP_OP_MAX_OUTPUTS];
    uint32_t n_pending = 0;
@@ -83,11 +104,6 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
            continue;
        }

-        if (t->size <= HEX_L2_FLUSH_IL_THRESHOLD) {
-            hex_l2flush((void *) (uintptr_t) t->data, t->size);
-            continue;
-        }
-
        uint32_t t_start = t->data;
        uint32_t t_end   = t_start + t->size;

@@ -110,6 +126,8 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
        }
    }

+    merge_dirty_ranges(ctx);
+
    if (n_pending == 0) {
        return;
    }
@@ -132,8 +150,8 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
            struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
            r->start = pending[i]->data;
            r->end   = pending[i]->data + pending[i]->size;
-            r->bi    = pending[i]->bi;
        }
+        merge_dirty_ranges(ctx);
        return;
    }

@@ -151,12 +169,12 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
            struct htp_dirty_range * r = &ctx->dirty_ranges[i];
            r->start = pending[i]->data;
            r->end   = pending[i]->data + pending[i]->size;
-            r->bi    = pending[i]->bi;
        }
+        merge_dirty_ranges(ctx);
        return;
    }

-    if (total_evict_size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1 && n_evict <= HTP_OP_MAX_INPUTS) {
+    if (total_evict_size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1 && n_evict <= HTP_MAX_DIRTY_RANGES) {
        struct l2flush_multi_task task;
        task.trace    = ctx->trace;
        task.n_ranges = n_evict;
@@ -195,7 +213,6 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
        struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
        r->start = pending[i]->data;
        r->end   = pending[i]->data + pending[i]->size;
-        r->bi    = pending[i]->bi;
    }

    for (uint32_t i = 0; i < n_empty; i++) {
@@ -203,8 +220,9 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co
        struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
        r->start = pending[n_evict + i]->data;
        r->end   = pending[n_evict + i]->data + pending[n_evict + i]->size;
-        r->bi    = pending[n_evict + i]->bi;
    }
+
+    merge_dirty_ranges(ctx);
 }

 static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor * t) {
@@ -242,15 +260,77 @@ static inline bool is_tensor_dirty(struct htp_context * ctx, const struct htp_te
    return false;
 }

+static void flush_dirty_ranges(struct htp_context * ctx, const struct htp_dirty_range * ranges, uint32_t n_ranges, uint64_t total_dirty) {
+    if (total_dirty >= HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
+        struct l2flush_multi_task task;
+        task.trace    = ctx->trace;
+        task.n_ranges = n_ranges;
+
+        uint32_t block_acc = 0;
+        for (uint32_t i = 0; i < n_ranges; i++) {
+            const struct htp_dirty_range * r = &ranges[i];
+            struct l2flush_range * rg = &task.ranges[i];
+            rg->start = hex_align_down((size_t) r->start, HEX_L2_LINE_SIZE);
+            rg->end   = hex_align_up((size_t) r->end, HEX_L2_LINE_SIZE);
+            rg->block_first = block_acc;
+            rg->n_blocks = (rg->end - rg->start + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
+            block_acc += rg->n_blocks;
+        }
+
+        task.total_blocks      = block_acc;
+        task.blocks_per_thread = fastdiv(block_acc + ctx->n_threads - 1, &ctx->n_threads_div);
+
+        work_queue_run(ctx->work_queue, l2flush_multi_worker, &task, ctx->n_threads);
+    } else {
+        struct htp_thread_trace * tr = &ctx->trace[0];
+        htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+        for (uint32_t i = 0; i < n_ranges; i++) {
+            const struct htp_dirty_range * r = &ranges[i];
+            hex_l2flush((void *) (uintptr_t) r->start, r->end - r->start);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+    }
+}
+
+void htp_flush_dirty_ranges(struct htp_context * ctx) {
+    struct htp_dirty_range ranges[HTP_MAX_DIRTY_RANGES];
+    uint32_t n_ranges = 0;
+    uint64_t total_dirty = 0;
+
+    for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
+        const struct htp_dirty_range * r = &ctx->dirty_ranges[i];
+        if (!r->start) {
+            continue;
+        }
+        ranges[n_ranges++] = *r;
+        total_dirty += r->end - r->start;
+    }
+
+    if (total_dirty == 0) {
+        return;
+    }
+
+    if (total_dirty > HEX_L2_FLUSH_ALL_THRESHOLD) {
+        flush_all_dcache(ctx);
+        return;
+    }
+
+    flush_dirty_ranges(ctx, ranges, n_ranges, total_dirty);
+    memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
+}
+
 void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
    const struct htp_tensor * dirty_tensors[HTP_OP_MAX_INPUTS];
+    struct htp_dirty_range ranges[HTP_OP_MAX_INPUTS];
    uint32_t n_dirty = 0;
    uint64_t total_dirty = 0;

    for (uint32_t i = 0; i < n; i++) {
        const struct htp_tensor * t = tensors[i];
-        if (t && !(t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE)) && is_tensor_dirty(ctx, t)) {
+        if (t && is_tensor_dirty(ctx, t)) {
            dirty_tensors[n_dirty++] = t;
+            ranges[n_dirty - 1].start = t->data;
+            ranges[n_dirty - 1].end   = t->data + t->size;
            total_dirty += t->size;
        }
    }
@@ -264,37 +344,8 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co
        return;
    }

-    if (total_dirty >= HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
-        struct l2flush_multi_task task;
-        task.trace    = ctx->trace;
-        task.n_ranges = 0;
-
-        uint32_t block_acc = 0;
-        for (uint32_t i = 0; i < n_dirty; i++) {
-            const struct htp_tensor * t = dirty_tensors[i];
-            make_tensor_clean(ctx, t);
-
-            struct l2flush_range * rg = &task.ranges[task.n_ranges++];
-            rg->start = hex_align_down((size_t) t->data, HEX_L2_LINE_SIZE);
-            rg->end   = hex_align_up((size_t) t->data + t->size, HEX_L2_LINE_SIZE);
-            rg->block_first = block_acc;
-            rg->n_blocks = (rg->end - rg->start + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
-            block_acc += rg->n_blocks;
-        }
-
-        task.total_blocks      = block_acc;
-        task.blocks_per_thread = fastdiv(block_acc + ctx->n_threads - 1, &ctx->n_threads_div);
-
-        work_queue_run(ctx->work_queue, l2flush_multi_worker, &task, ctx->n_threads);
-        return;
-    }
-
-    struct htp_thread_trace * tr = &ctx->trace[0];
+    flush_dirty_ranges(ctx, ranges, n_dirty, total_dirty);
    for (uint32_t i = 0; i < n_dirty; i++) {
-        const struct htp_tensor * t = dirty_tensors[i];
-        htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
-        hex_l2flush((void *) (uintptr_t) t->data, t->size);
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
-        make_tensor_clean(ctx, t);
+        make_tensor_clean(ctx, dirty_tensors[i]);
    }
 }
@@ -2,8 +2,20 @@
 #define HTP_TENSOR_H

 #include <stdint.h>
+#include <stdbool.h>
 #include "htp-ops.h"
 #include "hex-bitmap.h"
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+
+enum {
+    HTP_TENSOR_MDEV_LINE_SIZE = 128,
+};
+
+struct htp_tensor_mdev_range {
+    uint32_t start;
+    uint32_t count;
+};

 static inline void * htp_tensor_data(const struct htp_tensor * t) {
    return (void *) (uintptr_t) t->data;
@@ -13,6 +25,102 @@ static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) {
    return (uint32_t *) &t->flags;
 }

+static inline bool htp_tensor_is_contiguous(const struct htp_tensor * t, uint32_t type_size) {
+    uint32_t next_nb = type_size;
+    if (t->ne[0] != 1 && t->nb[0] != next_nb) {
+        return false;
+    }
+    next_nb *= t->ne[0];
+    for (int i = 1; i < HTP_OP_MAX_DIMS; i++) {
+        if (t->ne[i] != 1 && t->nb[i] != next_nb) {
+            return false;
+        }
+        next_nb *= t->ne[i];
+    }
+    return true;
+}
+
+static inline bool htp_tensor_is_permuted(const struct htp_tensor * t) {
+    return t->nb[0] > t->nb[1] || t->nb[1] > t->nb[2] || t->nb[2] > t->nb[3];
+}
+
+static inline bool htp_tensor_mdev_data_aligned(const struct htp_tensor * t) {
+    return ((uintptr_t) t->data & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0;
+}
+
+static inline bool htp_tensor_can_row_partition(const struct htp_tensor * t, uint32_t elem_size) {
+    if (!htp_tensor_mdev_data_aligned(t)) {
+        return false;
+    }
+    if (t->ne[0] != 1 && t->nb[0] != elem_size) {
+        return false;
+    }
+    if (htp_tensor_is_permuted(t)) {
+        return false;
+    }
+    if (t->ne[1] > 1 && (t->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) != 0) return false;
+    if (t->ne[2] > 1 && (t->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) != 0) return false;
+    if (t->ne[3] > 1 && (t->nb[3] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) != 0) return false;
+    return true;
+}
+
+static inline bool htp_tensor_mdev_rows_per_chunk(const struct htp_tensor * t, uint32_t elem_size, uint32_t row_size, uint32_t * rows_per_chunk) {
+    *rows_per_chunk = 0;
+
+    if (!htp_tensor_mdev_data_aligned(t)) {
+        return false;
+    }
+    if (t->ne[0] != 1 && t->nb[0] != elem_size) {
+        return false;
+    }
+    if (htp_tensor_is_permuted(t)) {
+        return false;
+    }
+    if (t->ne[1] > 1 && (t->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0 &&
+        (t->ne[2] <= 1 || (t->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0) &&
+        (t->ne[3] <= 1 || (t->nb[3] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0)) {
+        *rows_per_chunk = 1;
+        return true;
+    }
+    if (t->nb[1] == row_size &&
+        (t->ne[2] <= 1 || t->nb[2] == t->nb[1] * t->ne[1]) &&
+        (t->ne[3] <= 1 || t->nb[3] == t->nb[2] * t->ne[2])) {
+        *rows_per_chunk = (row_size > 0) ? (HTP_TENSOR_MDEV_LINE_SIZE / hex_gcd_u32(row_size, HTP_TENSOR_MDEV_LINE_SIZE)) : 1;
+        return true;
+    }
+    return false;
+}
+
+static inline struct htp_tensor_mdev_range htp_tensor_mdev_partition(uint32_t total_units, uint32_t units_per_chunk, uint32_t mdev_idx, uint32_t mdev_count, const struct fastdiv_values * mdev_count_div) {
+    struct htp_tensor_mdev_range range = { 0, total_units };
+
+    if (mdev_count <= 1) {
+        return range;
+    }
+
+    if (units_per_chunk == 0) {
+        range.start = (mdev_idx == 0) ? 0 : total_units;
+        range.count = (mdev_idx == 0) ? total_units : 0;
+        return range;
+    }
+
+    const uint32_t total_chunks = total_units / units_per_chunk;
+    if (total_chunks < mdev_count) {
+        range.start = (mdev_idx == 0) ? 0 : total_units;
+        range.count = (mdev_idx == 0) ? total_units : 0;
+        return range;
+    }
+
+    const uint32_t chunks_per_mdev = fastdiv(total_chunks + mdev_count - 1, mdev_count_div);
+    range.start = MIN(mdev_idx * chunks_per_mdev * units_per_chunk, total_units);
+    if (mdev_idx == mdev_count - 1) {
+        range.count = total_units - range.start;
+    } else {
+        range.count = MIN(chunks_per_mdev * units_per_chunk, total_units - range.start);
+    }
+    return range;
+}
+
 static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) {
    switch (type) {
        case HTP_TYPE_F32:  return ne00 * 4;
@@ -23,6 +131,7 @@ static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) {
 }

 struct htp_context;
+void htp_flush_dirty_ranges(struct htp_context * ctx);
 void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
 void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);

@@ -16,25 +16,25 @@
 #define UNUSED(x) (void)(x)

 #define hvx_arith_loop_body(dst_type, src0_type, src1_type, elem_size, vec_store, vec_op) \
-    do {                                                                       \
-        dst_type * vdst  = (dst_type *) dst;                                   \
-        src0_type * vsrc0 = (src0_type *) src0;                                \
-        src1_type * vsrc1 = (src1_type *) src1;                                \
-                                                                               \
-        const uint32_t epv  = 128 / (elem_size);                               \
-        const uint32_t nvec = n / epv;                                         \
-        const uint32_t nloe = n % epv;                                         \
-                                                                               \
-        uint32_t i = 0;                                                        \
-                                                                               \
-        _Pragma("unroll(4)")                                                   \
-        for (; i < nvec; i++) {                                                \
-            vdst[i] = vec_op(vsrc0[i], vsrc1[i]);                              \
-        }                                                                      \
-        if (nloe) {                                                            \
-            HVX_Vector v = vec_op(vsrc0[i], vsrc1[i]);                         \
-            vec_store((void *) &vdst[i], nloe * (elem_size), v);               \
-        }                                                                      \
+    do {                                                                                  \
+        dst_type * vdst  = (dst_type *) dst;                                              \
+        src0_type * vsrc0 = (src0_type *) src0;                                           \
+        src1_type * vsrc1 = (src1_type *) src1;                                           \
+                                                                                          \
+        const uint32_t epv  = 128 / (elem_size);                                          \
+        const uint32_t nvec = n / epv;                                                    \
+        const uint32_t nloe = n % epv;                                                    \
+                                                                                          \
+        uint32_t i = 0;                                                                   \
+                                                                                          \
+        _Pragma("unroll(4)")                                                              \
+        for (; i < nvec; i++) {                                                           \
+            vdst[i] = vec_op(vsrc0[i], vsrc1[i]);                                         \
+        }                                                                                 \
+        if (nloe) {                                                                       \
+            HVX_Vector v = vec_op(vsrc0[i], vsrc1[i]);                                    \
+            vec_store((void *) &vdst[i], nloe * (elem_size), v);                          \
+        }                                                                                 \
    } while(0)

 #if __HVX_ARCH__ < 79
@@ -56,43 +56,43 @@
 #define HVX_OP_MUL_F16(a, b) hvx_vec_mul_f16_f16(a, b)

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

 DEFINE_HVX_BINARY_OP_VARIANTS(hvx_add_f32, HVX_OP_ADD_F32, float)
 DEFINE_HVX_BINARY_OP_VARIANTS(hvx_sub_f32, HVX_OP_SUB_F32, float)
@@ -103,25 +103,25 @@ DEFINE_HVX_BINARY_OP_VARIANTS(hvx_sub_f16, HVX_OP_SUB_F16, _Float16)
 DEFINE_HVX_BINARY_OP_VARIANTS(hvx_mul_f16, HVX_OP_MUL_F16, _Float16)

 // Dispatcher logic
-#define HVX_BINARY_DISPATCHER(OP_NAME) \
+#define HVX_BINARY_DISPATCHER(OP_NAME)                                                                                                       \
 static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, const uint32_t num_elems) { \
-    if (hex_is_aligned((void *) dst, 128)) { \
-        if (hex_is_aligned((void *) src0, 128)) { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_aau(dst, src0, src1, num_elems); \
-        } else { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_auu(dst, src0, src1, num_elems); \
-        } \
-    } else { \
-        if (hex_is_aligned((void *) src0, 128)) { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_uau(dst, src0, src1, num_elems); \
-        } else { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_uuu(dst, src0, src1, num_elems); \
-        } \
-    } \
+    if (hex_is_aligned((void *) dst, 128)) {                                                                                                 \
+        if (hex_is_aligned((void *) src0, 128)) {                                                                                            \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_aau(dst, src0, src1, num_elems);                                               \
+        } else {                                                                                                                             \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_auu(dst, src0, src1, num_elems);                                               \
+        }                                                                                                                                    \
+    } else {                                                                                                                                 \
+        if (hex_is_aligned((void *) src0, 128)) {                                                                                            \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_uau(dst, src0, src1, num_elems);                                               \
+        } else {                                                                                                                             \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_uuu(dst, src0, src1, num_elems);                                               \
+        }                                                                                                                                    \
+    }                                                                                                                                        \
 }

 HVX_BINARY_DISPATCHER(hvx_add_f32)
@@ -166,44 +166,44 @@ static inline void hvx_mul_mul_f32_aa(uint8_t * restrict dst, const uint8_t * re

 // Scalar Operations

-#define hvx_scalar_loop_body(dst_type, src_type, elem_size, vec_store, scalar_op_macro)   \
-    do {                                                                       \
-        dst_type * restrict vdst = (dst_type *) dst;                           \
-        src_type * restrict vsrc = (src_type *) src;                           \
-                                                                               \
-        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++) {                                                \
-            HVX_Vector v = vsrc[i];                                            \
-            vdst[i] = scalar_op_macro(v);                                      \
-        }                                                                      \
-        if (nloe) {                                                            \
-            HVX_Vector v = vsrc[i];                                            \
-            v = scalar_op_macro(v);                                            \
-            vec_store((void *) &vdst[i], nloe * (elem_size), v);               \
-        }                                                                      \
+#define hvx_scalar_loop_body(dst_type, src_type, elem_size, vec_store, scalar_op_macro) \
+    do {                                                                                \
+        dst_type * restrict vdst = (dst_type *) dst;                                    \
+        src_type * restrict vsrc = (src_type *) src;                                    \
+                                                                                        \
+        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++) {                                                         \
+            HVX_Vector v = vsrc[i];                                                     \
+            vdst[i] = scalar_op_macro(v);                                               \
+        }                                                                               \
+        if (nloe) {                                                                     \
+            HVX_Vector v = vsrc[i];                                                     \
+            v = scalar_op_macro(v);                                                     \
+            vec_store((void *) &vdst[i], nloe * (elem_size), v);                        \
+        }                                                                               \
    } while(0)

-#define HVX_OP_ADD_SCALAR_F32(v) \
-    ({ \
+#define HVX_OP_ADD_SCALAR_F32(v)                                   \
+    ({                                                             \
        const HVX_VectorPred pred_inf = Q6_Q_vcmp_eq_VwVw(inf, v); \
-        HVX_Vector out = HVX_OP_ADD_F32(v, val_vec); \
-        Q6_V_vmux_QVV(pred_inf, inf, out); \
+        HVX_Vector out = HVX_OP_ADD_F32(v, val_vec);               \
+        Q6_V_vmux_QVV(pred_inf, inf, out);                         \
    })

 #define HVX_OP_MUL_SCALAR_F32(v) HVX_OP_MUL_F32(v, val_vec)
 #define HVX_OP_SUB_SCALAR_F32(v) HVX_OP_SUB_F32(v, val_vec)

-#define HVX_OP_ADD_SCALAR_F16(v) \
-    ({ \
+#define HVX_OP_ADD_SCALAR_F16(v)                                   \
+    ({                                                             \
        const HVX_VectorPred pred_inf = Q6_Q_vcmp_eq_VhVh(inf, v); \
-        HVX_Vector out = HVX_OP_ADD_F16(v, val_vec); \
-        Q6_V_vmux_QVV(pred_inf, inf, out); \
+        HVX_Vector out = HVX_OP_ADD_F16(v, val_vec);               \
+        Q6_V_vmux_QVV(pred_inf, inf, out);                         \
    })

 #define HVX_OP_MUL_SCALAR_F16(v) HVX_OP_MUL_F16(v, val_vec)
@@ -212,31 +212,31 @@ static inline void hvx_mul_mul_f32_aa(uint8_t * restrict dst, const uint8_t * re
 // Scalar Variants

 // Generic macro to define alignment permutations for an op
-#define DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(OP_NAME, OP_MACRO, SPLAT_MACRO, ELEM_TYPE) \
+#define DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(OP_NAME, OP_MACRO, SPLAT_MACRO, ELEM_TYPE)                                  \
 static inline void OP_NAME##_aa(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \
-    const HVX_Vector val_vec = SPLAT_MACRO(val); \
-    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \
-    assert((uintptr_t) dst % 128 == 0); \
-    assert((uintptr_t) src % 128 == 0); \
-    hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \
-} \
+    const HVX_Vector val_vec = SPLAT_MACRO(val);                                                                         \
+    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf);                                                \
+    assert((uintptr_t) dst % 128 == 0);                                                                                  \
+    assert((uintptr_t) src % 128 == 0);                                                                                  \
+    hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);                          \
+}                                                                                                                        \
 static inline void OP_NAME##_au(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \
-    const HVX_Vector val_vec = SPLAT_MACRO(val); \
-    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \
-    assert((uintptr_t) dst % 128 == 0); \
-    hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \
-} \
+    const HVX_Vector val_vec = SPLAT_MACRO(val);                                                                         \
+    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf);                                                \
+    assert((uintptr_t) dst % 128 == 0);                                                                                  \
+    hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);                         \
+}                                                                                                                        \
 static inline void OP_NAME##_ua(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \
-    const HVX_Vector val_vec = SPLAT_MACRO(val); \
-    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \
-    assert((uintptr_t) src % 128 == 0); \
-    hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
-} \
+    const HVX_Vector val_vec = SPLAT_MACRO(val);                                                                         \
+    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf);                                                \
+    assert((uintptr_t) src % 128 == 0);                                                                                  \
+    hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO);                         \
+}                                                                                                                        \
 static inline void OP_NAME##_uu(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \
-    const HVX_Vector val_vec = SPLAT_MACRO(val); \
-    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf); \
-    hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
-} \
+    const HVX_Vector val_vec = SPLAT_MACRO(val);                                                                         \
+    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf);                                                \
+    hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO);                        \
+}                                                                                                                        \

 DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_add_scalar_f32, HVX_OP_ADD_SCALAR_F32, hvx_vec_splat_f32, float)
 DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_sub_scalar_f32, HVX_OP_SUB_SCALAR_F32, hvx_vec_splat_f32, float)
@@ -247,17 +247,17 @@ DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_sub_scalar_f16, HVX_OP_SUB_SCALAR_F16,
 DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_mul_scalar_f16, HVX_OP_MUL_SCALAR_F16, hvx_vec_splat_f16, _Float16)

 // Dispatcher logic
-#define HVX_BINARY_SCALAR_DISPATCHER(OP_NAME, ELEM_TYPE) \
+#define HVX_BINARY_SCALAR_DISPATCHER(OP_NAME, ELEM_TYPE)                                                                          \
 static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, const uint32_t num_elems) { \
-    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { \
-        OP_NAME##_aa(dst, src, val, num_elems); \
-    } else if (hex_is_aligned((void *) dst, 128)) { \
-        OP_NAME##_au(dst, src, val, num_elems); \
-    } else if (hex_is_aligned((void *) src, 128)) { \
-        OP_NAME##_ua(dst, src, val, num_elems); \
-    } else { \
-        OP_NAME##_uu(dst, src, val, num_elems); \
-    } \
+    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) {                                                 \
+        OP_NAME##_aa(dst, src, val, num_elems);                                                                                   \
+    } else if (hex_is_aligned((void *) dst, 128)) {                                                                               \
+        OP_NAME##_au(dst, src, val, num_elems);                                                                                   \
+    } else if (hex_is_aligned((void *) src, 128)) {                                                                               \
+        OP_NAME##_ua(dst, src, val, num_elems);                                                                                   \
+    } else {                                                                                                                      \
+        OP_NAME##_uu(dst, src, val, num_elems);                                                                                   \
+    }                                                                                                                             \
 }

 HVX_BINARY_SCALAR_DISPATCHER(hvx_add_scalar_f32, float)
@@ -350,12 +350,12 @@ static inline void hvx_max_scalar_f32(uint8_t * restrict dst, const uint8_t * re

 // CLAMP Scalar variants

-#define HVX_OP_CLAMP_SCALAR(v) \
-    ({ \
+#define HVX_OP_CLAMP_SCALAR(v)                                           \
+    ({                                                                   \
        HVX_VectorPred pred_cap_right = Q6_Q_vcmp_gt_VsfVsf(v, max_vec); \
        HVX_VectorPred pred_cap_left  = Q6_Q_vcmp_gt_VsfVsf(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); \
+        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_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) {
@@ -219,64 +219,64 @@ static inline HVX_Vector hvx_vec_hybrid_div_f16(HVX_Vector vec1, HVX_Vector vec2
    } while(0)

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

 // Dispatcher logic
-#define HVX_DIV_DISPATCHER(OP_NAME) \
+#define HVX_DIV_DISPATCHER(OP_NAME)                                                                                                          \
 static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, const uint32_t num_elems) { \
-    if (hex_is_aligned((void *) dst, 128)) { \
-        if (hex_is_aligned((void *) src0, 128)) { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_aau(dst, src0, src1, num_elems); \
-        } else { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_auu(dst, src0, src1, num_elems); \
-        } \
-    } else { \
-        if (hex_is_aligned((void *) src0, 128)) { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_uau(dst, src0, src1, num_elems); \
-        } else { \
-            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems); \
-            else                                    OP_NAME##_uuu(dst, src0, src1, num_elems); \
-        } \
-    } \
+    if (hex_is_aligned((void *) dst, 128)) {                                                                                                 \
+        if (hex_is_aligned((void *) src0, 128)) {                                                                                            \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_aau(dst, src0, src1, num_elems);                                               \
+        } else {                                                                                                                             \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_auu(dst, src0, src1, num_elems);                                               \
+        }                                                                                                                                    \
+    } else {                                                                                                                                 \
+        if (hex_is_aligned((void *) src0, 128)) {                                                                                            \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_uau(dst, src0, src1, num_elems);                                               \
+        } else {                                                                                                                             \
+            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems);                                               \
+            else                                    OP_NAME##_uuu(dst, src0, src1, num_elems);                                               \
+        }                                                                                                                                    \
+    }                                                                                                                                        \
 }

 DEFINE_HVX_DIV_OP_VARIANTS(hvx_div_f32, hvx_div_f32_loop_body)
@@ -169,36 +169,36 @@ static inline HVX_Vector hvx_vec_inverse_f16_guard(HVX_Vector v_sf, HVX_Vector n
    } while(0)

 // Generic macro to define alignment permutations for an op
-#define DEFINE_HVX_INV_OP_VARIANTS(OP_NAME, OP_LOOP_BODY) \
+#define DEFINE_HVX_INV_OP_VARIANTS(OP_NAME, OP_LOOP_BODY)                                           \
 static inline void OP_NAME##_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \
-    assert((uintptr_t) dst % 128 == 0); \
-    assert((uintptr_t) src % 128 == 0); \
-    OP_LOOP_BODY(HVX_Vector, HVX_Vector, hvx_vec_store_a); \
-} \
+    assert((uintptr_t) dst % 128 == 0);                                                             \
+    assert((uintptr_t) src % 128 == 0);                                                             \
+    OP_LOOP_BODY(HVX_Vector, HVX_Vector, hvx_vec_store_a);                                          \
+}                                                                                                   \
 static inline void OP_NAME##_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \
-    assert((uintptr_t) dst % 128 == 0); \
-    OP_LOOP_BODY(HVX_Vector, HVX_UVector, hvx_vec_store_a); \
-} \
+    assert((uintptr_t) dst % 128 == 0);                                                             \
+    OP_LOOP_BODY(HVX_Vector, HVX_UVector, hvx_vec_store_a);                                         \
+}                                                                                                   \
 static inline void OP_NAME##_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \
-    assert((uintptr_t) src % 128 == 0); \
-    OP_LOOP_BODY(HVX_UVector, HVX_Vector, hvx_vec_store_u); \
-} \
+    assert((uintptr_t) src % 128 == 0);                                                             \
+    OP_LOOP_BODY(HVX_UVector, HVX_Vector, hvx_vec_store_u);                                         \
+}                                                                                                   \
 static inline void OP_NAME##_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { \
-    OP_LOOP_BODY(HVX_UVector, HVX_UVector, hvx_vec_store_u); \
-} \
+    OP_LOOP_BODY(HVX_UVector, HVX_UVector, hvx_vec_store_u);                                        \
+}                                                                                                   \

 // Dispatcher logic
-#define HVX_INV_DISPATCHER(OP_NAME) \
+#define HVX_INV_DISPATCHER(OP_NAME)                                                                          \
 static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t num_elems) { \
-    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) { \
-        OP_NAME##_aa(dst, src, num_elems); \
-    } else if (hex_is_aligned((void *) dst, 128)) { \
-        OP_NAME##_au(dst, src, num_elems); \
-    } else if (hex_is_aligned((void *) src, 128)) { \
-        OP_NAME##_ua(dst, src, num_elems); \
-    } else { \
-        OP_NAME##_uu(dst, src, num_elems); \
-    } \
+    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) {                            \
+        OP_NAME##_aa(dst, src, num_elems);                                                                   \
+    } else if (hex_is_aligned((void *) dst, 128)) {                                                          \
+        OP_NAME##_au(dst, src, num_elems);                                                                   \
+    } else if (hex_is_aligned((void *) src, 128)) {                                                          \
+        OP_NAME##_ua(dst, src, num_elems);                                                                   \
+    } else {                                                                                                 \
+        OP_NAME##_uu(dst, src, num_elems);                                                                   \
+    }                                                                                                        \
 }

 DEFINE_HVX_INV_OP_VARIANTS(hvx_inverse_f32, hvx_inverse_f32_loop_body)
@@ -68,30 +68,30 @@ static inline void hvx_scale_f32(uint8_t * restrict dst, const uint8_t * restric
    }
 }

-#define hvx_scale_offset_f32_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 elem_size = sizeof(float);                                    \
-        const uint32_t epv = 128 / elem_size;                                        \
-        const uint32_t nvec = n / epv;                                               \
-        const uint32_t nloe = n % epv;                                               \
-                                                                                     \
-        uint32_t i = 0;                                                              \
-                                                                                     \
-        _Pragma("unroll(4)")                                                         \
-        for (; i < nvec; ++i) {                                                      \
+#define hvx_scale_offset_f32_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 elem_size = sizeof(float);                                         \
+        const uint32_t epv = 128 / elem_size;                                             \
+        const uint32_t nvec = n / epv;                                                    \
+        const uint32_t nloe = n % epv;                                                    \
+                                                                                          \
+        uint32_t i = 0;                                                                   \
+                                                                                          \
+        _Pragma("unroll(4)")                                                              \
+        for (; i < nvec; ++i) {                                                           \
            HVX_Vector v = Q6_Vqf32_vadd_Vqf32Vsf(Q6_Vqf32_vmpy_VsfVsf(vsrc[i], vs), vo); \
-            vdst[i] = Q6_Vsf_equals_Vqf32(v);                                        \
-        }                                                                            \
-        if (nloe) {                                                                  \
+            vdst[i] = Q6_Vsf_equals_Vqf32(v);                                             \
+        }                                                                                 \
+        if (nloe) {                                                                       \
            HVX_Vector v = Q6_Vqf32_vadd_Vqf32Vsf(Q6_Vqf32_vmpy_VsfVsf(vsrc[i], vs), vo); \
-            vec_store((void *) &vdst[i], nloe * elem_size, Q6_Vsf_equals_Vqf32(v));  \
-        }                                                                            \
+            vec_store((void *) &vdst[i], nloe * elem_size, Q6_Vsf_equals_Vqf32(v));       \
+        }                                                                                 \
    } while(0)

 static inline void hvx_scale_offset_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const int n, const float scale, const float offset) {
@@ -68,50 +68,50 @@ static inline HVX_Vector hvx_vec_tanh_f32(HVX_Vector x) {
    return Q6_Vsf_equals_Vqf32(res);
 }

-#define hvx_sigmoid_loop_body(dst_type, src_type, vec_store)    \
-    do {                                                        \
-        dst_type * restrict vdst = (dst_type *) dst;            \
-        src_type * restrict vsrc = (src_type *) src;            \
-                                                                \
-        const HVX_Vector one     = hvx_vec_splat_f32(1.f);      \
-        const HVX_Vector max_exp = hvx_vec_splat_f32(87.f);     \
-        const HVX_Vector min_exp = hvx_vec_splat_f32(-87.f);    \
-                                                                \
-        const uint32_t epv  = 128 / sizeof(float);              \
-        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_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp); \
-        }                                                       \
-        if (nloe) {                                             \
+#define hvx_sigmoid_loop_body(dst_type, src_type, vec_store)                                  \
+    do {                                                                                      \
+        dst_type * restrict vdst = (dst_type *) dst;                                          \
+        src_type * restrict vsrc = (src_type *) src;                                          \
+                                                                                              \
+        const HVX_Vector one     = hvx_vec_splat_f32(1.f);                                    \
+        const HVX_Vector max_exp = hvx_vec_splat_f32(87.f);                                   \
+        const HVX_Vector min_exp = hvx_vec_splat_f32(-87.f);                                  \
+                                                                                              \
+        const uint32_t epv  = 128 / sizeof(float);                                            \
+        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_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp);        \
+        }                                                                                     \
+        if (nloe) {                                                                           \
             HVX_Vector tmp = hvx_vec_fast_sigmoid_f32_guard(vsrc[i], one, max_exp, min_exp); \
-             vec_store((void *) &vdst[i], nloe * sizeof(float), tmp); \
-        }                                                       \
+             vec_store((void *) &vdst[i], nloe * sizeof(float), tmp);                         \
+        }                                                                                     \
    } while(0)

-#define hvx_tanh_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 epv  = 128 / sizeof(float);              \
-        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_tanh_f32(vsrc[i]);               \
-        }                                                       \
-        if (nloe) {                                             \
-             HVX_Vector tmp = hvx_vec_tanh_f32(vsrc[i]);        \
+#define hvx_tanh_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 epv  = 128 / sizeof(float);                    \
+        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_tanh_f32(vsrc[i]);                     \
+        }                                                             \
+        if (nloe) {                                                   \
+             HVX_Vector tmp = hvx_vec_tanh_f32(vsrc[i]);              \
             vec_store((void *) &vdst[i], nloe * sizeof(float), tmp); \
-        }                                                       \
+        }                                                             \
    } while(0)

 static inline void hvx_sigmoid_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
@@ -4,87 +4,75 @@
 #include "hvx-base.h"
 #include "hvx-floor.h"

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

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

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

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

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

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

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

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

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

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

 #endif /* HVX_SIN_COS_H */
@@ -3,11 +3,12 @@
 #pragma clang diagnostic ignored "-Wunused-but-set-variable"

 #include <HAP_farf.h>
-#include <HAP_perf.h>
 #include <hexagon_protos.h>
 #include <hexagon_types.h>
 #include <string.h>

+#include "hex-common.h"
+
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
 #include "htp-ctx.h"
@@ -16,14 +17,19 @@
 #include "hex-dma.h"
 #include "hex-profile.h"
 #include "htp-vtcm.h"
+#include "htp-tensor.h"

 struct htp_im2col_context {
    struct htp_ops_context * octx;
+    uint32_t                 patch_base;           // first patch index assigned to this dev
+    uint32_t                 npatches;             // number of patches assigned to this dev
    uint32_t                 npatches_per_thread;  // patches = N*OH*OW (pure-DDR kernel)

-    uint32_t pe_rows_per_thread;                   // N*OH rows per worker
-    uint32_t pe_src_row_bytes;                     // one output row's source: IC*KH*IW*4, rounded 256
-    uint32_t pe_dst_row_bytes;                     // one output row's dst: OW*patch_stride*2, rounded 256
+    uint32_t pe_row_base;              // first N*OH row index assigned to this dev (DMA path)
+    uint32_t pe_nrows;                 // number of N*OH rows assigned to this dev (DMA path)
+    uint32_t pe_rows_per_thread;       // N*OH rows per worker
+    uint32_t pe_src_row_bytes;         // one output row's source: IC*KH*IW*4, rounded 256
+    uint32_t pe_dst_row_bytes;         // one output row's dst: OW*patch_stride*2, rounded 256

    // Patch-embed DMA path VTCM ping-pong.
    uint8_t * pe_vtcm_src;                         // base of the 2x src buffers region
@@ -58,33 +64,27 @@ static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout *
        struct htp_im2col_context * ictx        = (struct htp_im2col_context *) data;                     \
        struct htp_ops_context *    octx        = ictx->octx;                                             \
        struct htp_thread_trace * restrict tr   = &octx->ctx->trace[ith];                                 \
+        const struct htp_tensor * restrict src0 = octx->src[0];                                           \
        const struct htp_tensor * restrict src1 = octx->src[1];                                           \
        const struct htp_tensor * restrict dst  = octx->dst;                                              \
-        const int32_t  s0                       = octx->op_params[0];                                     \
-        const int32_t  s1                       = octx->op_params[1];                                     \
-        const int32_t  p0                       = octx->op_params[2];                                     \
-        const int32_t  p1                       = octx->op_params[3];                                     \
-        const int32_t  d0                       = octx->op_params[4];                                     \
-        const int32_t  d1                       = octx->op_params[5];                                     \
-        const uint32_t N                        = src1->ne[3];                                            \
-        const uint32_t IC                       = src1->ne[2];                                            \
-        const uint32_t IH                       = src1->ne[1];                                            \
-        const uint32_t IW                       = src1->ne[0];                                            \
-        const uint32_t KH                       = octx->src[0]->ne[1];                                    \
-        const uint32_t KW                       = octx->src[0]->ne[0];                                    \
+        const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1];                                   \
+        const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3];                                   \
+        const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5];                                   \
+        const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0];             \
+        const uint32_t KH                       = src0->ne[1], KW = src0->ne[0];                          \
        const uint32_t OH                       = dst->ne[2];                                             \
        const uint32_t OW                       = dst->ne[1];                                             \
        const uint32_t patch_stride             = IC * KH * KW;                                           \
        const float * restrict src_data         = (const float *) src1->data;                             \
        DST_CTYPE * restrict dst_data           = (DST_CTYPE *) dst->data;                                \
-        const uint32_t npatches                 = N * OH * OW;                                            \
-        const uint32_t patch_start              = ictx->npatches_per_thread * ith;                        \
-        const uint32_t patch_end                = MIN(patch_start + ictx->npatches_per_thread, npatches); \
-        if (patch_start >= patch_end) {                                                                   \
+        const uint32_t patch_end                = ictx->patch_base + ictx->npatches;                      \
+        const uint32_t patch_start              = ictx->patch_base + ictx->npatches_per_thread * ith;     \
+        const uint32_t patch_stop               = MIN(patch_start + ictx->npatches_per_thread, patch_end);\
+        if (patch_start >= patch_stop) {                                                                  \
            return;                                                                                       \
        }                                                                                                 \
        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, patch_start);                                   \
-        for (uint32_t p = patch_start; p < patch_end; p++) {                                              \
+        for (uint32_t p = patch_start; p < patch_stop; p++) {                                             \
            const uint32_t iow             = p % OW;                                                      \
            const uint32_t ioh             = (p / OW) % OH;                                               \
            const uint32_t in              = p / (OW * OH);                                               \
@@ -154,10 +154,10 @@ IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx
        uint8_t *      dst_base         = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread;                    \
        float *        srcb             = (float *) src_base;                                                        \
        DST_CTYPE *    dstb             = (DST_CTYPE *) dst_base;                                                    \
-        const uint32_t nrows            = N * OH;                                                                    \
+        const uint32_t row_end_max      = ictx->pe_row_base + ictx->pe_nrows;                                        \
        const uint32_t per_thread       = ictx->pe_rows_per_thread;                                                  \
-        const uint32_t row_start        = per_thread * ith;                                                          \
-        const uint32_t row_end          = MIN(row_start + per_thread, nrows);                                        \
+        const uint32_t row_start        = ictx->pe_row_base + per_thread * ith;                                      \
+        const uint32_t row_end          = MIN(row_start + per_thread, row_end_max);                                  \
        if (row_start >= row_end)                                                                                    \
            return;                                                                                                  \
        for (uint32_t r = row_start; r < row_end; r++) {                                                             \
@@ -266,26 +266,55 @@ int op_im2col(struct htp_ops_context * octx) {
        return HTP_STATUS_NO_SUPPORT;
    }

-    const uint32_t N         = src1->ne[3];
-    const uint32_t OH        = dst->ne[2];
-    const uint32_t OW        = dst->ne[1];
-    const uint32_t npatches  = N * OH * OW;
-    const uint32_t n_threads = MIN(octx->n_threads, npatches);
-
-    if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) || n_threads == 0) {
+    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
        return HTP_STATUS_OK;
    }

+    const uint32_t N             = src1->ne[3];
+    const uint32_t OH            = dst->ne[2];
+    const uint32_t OW            = dst->ne[1];
+    const uint32_t total_patches = N * OH * OW;
+    const uint32_t total_rows    = N * OH;
+
+    uint32_t patch_base = 0;
+    uint32_t npatches   = total_patches;
+    if (octx->ctx->mdev.count > 1) {
+        const uint32_t patch_size = dst->nb[1];
+        const uint32_t patches_per_chunk = (patch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(patch_size, HEX_L2_LINE_SIZE)) : 1;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_patches, htp_tensor_mdev_data_aligned(dst) ? patches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        patch_base = range.start;
+        npatches   = range.count;
+    }
+
+    uint32_t row_base = 0;
+    uint32_t nrows    = total_rows;
+    if (octx->ctx->mdev.count > 1) {
+        const uint32_t row_size = dst->nb[2];
+        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_base = range.start;
+        nrows    = range.count;
+    }
+
+    if (npatches == 0 && nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = MIN(octx->n_threads, MAX(npatches, 1));
+
    struct htp_im2col_context ictx = { 0 };
-    ictx.octx                      = octx;
-    ictx.npatches_per_thread       = (npatches + n_threads - 1) / n_threads;
+    ictx.octx                = octx;
+    ictx.patch_base          = patch_base;
+    ictx.npatches            = npatches;
+    ictx.npatches_per_thread = (npatches + n_threads - 1) / n_threads;

    // Clean non-overlapping patch-embed -> DMA kernel (if it fits VTCM);
    // everything else (padding/dilation/stride edges) -> pure-DDR kernel.
-    if (im2col_use_patchembed_dma(octx)) {
-        const uint32_t nrows = N * OH;
-        const uint32_t pth   = MIN(octx->n_threads, nrows);
+    if (im2col_use_patchembed_dma(octx) && nrows > 0) {
+        const uint32_t pth = MIN(octx->n_threads, nrows);
        if (pth > 0 && im2col_patchembed_dma_fits(octx, &ictx, pth)) {
+            ictx.pe_row_base        = row_base;
+            ictx.pe_nrows           = nrows;
            ictx.pe_rows_per_thread = (nrows + pth - 1) / pth;
            if (dst->type == HTP_TYPE_F16) {
                work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth);
@@ -297,6 +326,10 @@ int op_im2col(struct htp_ops_context * octx) {
        // else: doesn't fit -> fall through to the pure-DDR kernel below.
    }

+    if (npatches == 0) {
+        return HTP_STATUS_OK;
+    }
+
    if (dst->type == HTP_TYPE_F16) {
        work_queue_run(octx->ctx->work_queue, im2col_patchembed_thread, &ictx, n_threads);
    } else {
@@ -34,6 +34,7 @@
 #include "work-queue.h"
 #include "hex-profile.h"
 #include "allreduce-ops.h"
+#include "htp-fence.h"

 #define HMX_QUEUE_CAPACITY     16
 #define HMX_QUEUE_STACK_SIZE   16384
@@ -710,22 +711,43 @@ static inline void profile_stop(uint32_t mode, struct profile_data * d) {
 static int op_fence(struct htp_ops_context * octx) {
    struct htp_context *ctx = octx->ctx;
    struct htp_thread_trace * tr = &ctx->trace[0];
-    const uint32_t seq = (uint32_t) octx->op_params[0];
+    const uint32_t seq  = (uint32_t) octx->op_params[0];
+    const uint32_t mode = (uint32_t) octx->op_params[1];

    htp_trace_event_start(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);

    const struct htp_tensor * sync = octx->src[0];
-    atomic_uint * sync_fence = (atomic_uint *) sync->data;
+    atomic_uint * sync_fence = (atomic_uint *) (uintptr_t) sync->data;
+
+    if (mode == 1) {
+        htp_flush_dirty_ranges(ctx);
+
+        htp_mdev_group_barrier(octx);
+
+        if (ctx->mdev.idx == 0) {
+            htp_fence_write(sync_fence, seq, octx->status);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);
+        FARF(HIGH, "ggml-hex: sync-signal : fence %p seq 0x%x status %d\n", sync_fence, seq, octx->status);
+        return octx->status;
+    }
+
+    int status = HTP_STATUS_OK;
    uint64_t spins = 0;
    while (1) {
-        Q6_dccleaninva_A((void *) sync_fence);
-        asm volatile ("syncht" : : : "memory");
-        uint32_t val = atomic_load(&sync_fence[0]);
-        if ((int32_t)(val - seq) >= 0) {
+        uint32_t sync_seq;
+        uint32_t sync_status;
+        htp_fence_read(sync_fence, &sync_seq, &sync_status);
+        if ((int32_t)(sync_seq - seq) >= 0) {
+            if (sync_status > HTP_STATUS_OK) {
+                FARF(ERROR, "ggml-hex: sync-wait peer failed with status %u : fence %p seq 0x%x\n", sync_status, sync_fence, seq);
+                status = sync_status;
+            }
            break;
        }
        if (++spins > HTP_FENCE_TIMEOUT) {
-            FARF(ERROR, "ggml-hex: sync-wait TIMEOUT : fence %p spins %llu seq %u\n", sync_fence, spins, seq);
+            FARF(ERROR, "ggml-hex: sync-wait TIMEOUT : fence %p spins %llu seq 0x%x\n", sync_fence, spins, seq);
+            status = HTP_STATUS_INTERNAL_ERR;
            break;
        }
        hex_pause();
@@ -733,12 +755,27 @@ static int op_fence(struct htp_ops_context * octx) {

    htp_trace_event_stop(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq);

-    FARF(HIGH, "ggml-hex: sync-done : fence %p spins %llu seq %u\n", sync_fence, spins, seq);
+    FARF(HIGH, "ggml-hex: sync-done : fence %p spins %llu seq 0x%x\n", sync_fence, spins, seq);
+    return status;
+}
+
+static int op_mdev_group(struct htp_ops_context * octx) {
+    struct htp_context * ctx = octx->ctx;
+    const struct htp_tensor * sync = octx->src[0];
+    ctx->mdev.idx   = (uint16_t) octx->op_params[0];
+    ctx->mdev.count = (uint16_t) sync->ne[1];
+    if (ctx->mdev.count > 1) {
+        ctx->mdev.count_div = init_fastdiv_values(ctx->mdev.count);
+        ctx->mdev.fence_base = (uint8_t *) sync->data;
+    }
    return HTP_STATUS_OK;
 }

 static int execute_op(struct htp_ops_context * octx) {
    switch (octx->op) {
+        case HTP_OP_MDEV_GROUP:
+            return op_mdev_group(octx);
+
        case HTP_OP_FENCE:
            return op_fence(octx);

@@ -812,6 +849,7 @@ static int execute_op(struct htp_ops_context * octx) {
            return op_sum_rows(octx);

        case HTP_OP_CPY:
+        case HTP_OP_CPY_FENCE:
            return op_cpy(octx);

        case HTP_OP_REPEAT:
@@ -855,7 +893,7 @@ static int execute_op(struct htp_ops_context * octx) {
    }

    FARF(ERROR, "Unknown Op %u", octx->op);
-    return -1;
+    return HTP_STATUS_NO_SUPPORT;
 }

 static inline bool reuse_buf(struct htp_context *ctx, uint32_t *m_reuse, struct htp_buf_desc *b) {
@@ -984,11 +1022,19 @@ static void prep_tensors(struct htp_context *ctx, struct htp_buf_desc *bufs, str
    }
 }

-static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, uint32_t idx, struct htp_op_desc * op) {
-    memcpy(octx->op_params, op->params, sizeof(octx->op_params));
+static void mdev_group_init(struct htp_context * ctx, const struct htp_opbatch_req * req) {
+    memset(&ctx->mdev, 0, sizeof(ctx->mdev));
+    ctx->mdev.fence_seq = (uint32_t)((req->seq & 0xfffff) << 12);
+}
+
+static int proc_op_req(struct htp_ops_context * octx, struct htp_buf_desc * bufs, uint32_t n_bufs,
+                       struct htp_tensor * tens, uint32_t idx, struct htp_op_desc * op) {
+    memcpy(octx->op_params,     op->params, sizeof(octx->op_params));
    memcpy(octx->kernel_params, op->kernel_params, sizeof(octx->kernel_params));
-    octx->flags = op->flags;
-    octx->op    = op->opcode;
+    octx->flags         = op->flags;
+    octx->op            = op->opcode;
+    octx->n_threads     = octx->ctx->n_threads;
+    octx->n_threads_div = octx->ctx->n_threads_div;

    FARF(HIGH, "proc-op #%u: opcode %u flags 0x%x", idx, octx->op, octx->flags);

@@ -1027,9 +1073,13 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, u
            dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
    }

+    htp_tensor_dirty_all(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS);
+
+    htp_mdev_group_barrier(octx);
+
    int status = execute_op(octx);

-    htp_tensor_dirty_all(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS);
+    htp_ops_context_set_status(octx, status);

    octx->src0_spad.src = NULL;
    octx->src1_spad.src = NULL;
@@ -1037,7 +1087,7 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, u
    octx->src3_spad.src = NULL;
    octx->dst_spad.src  = NULL;

-    return status;
+    return octx->status;
 }

 static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_req * req, const struct dspqueue_buffer * dbuf) {
@@ -1059,7 +1109,7 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
        return;
    }

-    FARF(HIGH, "processing opbatch #%u: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", req->id,
+    FARF(HIGH, "processing opbatch #%llu: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", (unsigned long long) req->seq,
            n_bufs, n_tens, n_ops, req->n_traces, dbuf->size, b_size, t_size, o_size);

    // Setup descriptor pointers
@@ -1096,8 +1146,11 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r

    struct htp_ops_context *octx = &ctx->octx;
    memset(octx, 0, sizeof(*octx));
-    octx->n_threads = ctx->n_threads;
-    octx->ctx       = ctx;
+    octx->n_threads     = ctx->n_threads;
+    octx->n_threads_div = ctx->n_threads_div;
+    octx->ctx           = ctx;
+
+    mdev_group_init(ctx, req);

    work_queue_wakeup(ctx->work_queue);
    if (ctx->hmx_queue) {
@@ -1105,15 +1158,18 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
    }

    int op_status = HTP_STATUS_OK;
-    for (uint32_t i = 0; i < n_ops && op_status == HTP_STATUS_OK; i++) {
+    octx->status  = HTP_STATUS_OK;
+    for (uint32_t i = 0; i < n_ops; i++) {
        struct profile_data prof;

        profile_start(ctx->profiler, &prof);

-        op_status = proc_op_req(octx, tens, i, &ops[i]);
+        op_status = proc_op_req(octx, bufs, n_bufs, tens, i, &ops[i]);

        profile_stop(ctx->profiler, &prof);

+        htp_ops_context_set_status(octx, op_status);
+
        if (ctx->profiler) {
            pds[i].opcode = ops[i].opcode;
            pds[i].usecs  = prof.usecs;
@@ -1136,19 +1192,20 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r
    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
    htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);

+    htp_mdev_group_barrier(octx);
+
    profile_stop(HTP_PROF_BASIC, &batch_prof);

    struct htp_opbatch_rsp rsp;
    memset(&rsp, 0, sizeof(rsp));
-    rsp.id           = req->id;
-    rsp.status       = op_status;
+    rsp.seq          = req->seq;
+    rsp.status       = octx->status;
    rsp.n_bufs       = n_bufs;
    rsp.n_tensors    = n_tens;
    rsp.n_ops        = n_ops;
    rsp.usecs        = batch_prof.usecs;
    rsp.cycles_start = batch_prof.cycles_start;
    rsp.cycles_stop  = batch_prof.cycles_stop;
-    rsp.seq          = req->seq;

    if (ctx->profiler == HTP_PROF_TRACE) {
        for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
@@ -21,6 +21,7 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 #include "matmul-ops.h"
 #include "htp-vtcm.h"

@@ -89,6 +90,8 @@ struct htp_mm_context {

    // Precomputed values
    uint32_t src0_nrows_per_thread;
+    uint32_t src0_row_start;
+    uint32_t src0_row_end;
    uint32_t src0_row_size_padded;
    uint32_t src1_nrows;

@@ -135,6 +138,23 @@ struct htp_mm_context {
    uint32_t vtcm_dst_size_per_thread;
 };

+static int htp_mm_init_context(
+    struct htp_ops_context * octx,
+    const struct htp_mm_kernel_params * kparams
+) {
+    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    if (kparams->n_hmx) {
+        if (kparams->n_act_threads <= 0 || kparams->n_act_threads > (int32_t) octx->n_threads) {
+            return HTP_STATUS_INVAL_PARAMS;
+        }
+    }
+
+    return HTP_STATUS_OK;
+}
+
 // vdelta control to expand first 32 e8m0 values into 32 uint32 elements
 static const uint8_t __attribute__((aligned(128))) expand_x32_e8m0[128] = {
    0x00, 0x00, 0x00, 0x00, 0x01, 0x04, 0x00, 0x00, 0x02, 0x00, 0x08, 0x08, 0x01, 0x02, 0x00, 0x04, 0x04, 0x00, 0x00,
@@ -238,22 +258,24 @@ static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) {
    // This is the size of the rest of the dimensions of the result
    const uint32_t nr1 = ne1 * ne2 * ne3;

+    const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start;
+
    // distribute the thread work across the inner or outer loop based on which one is larger
    uint32_t dr0, dr1, ith0, ith1;
    if (nr0 > nr1) {
-        dr0  = fastdiv(nr0 + nth - 1, &octx->ctx->n_threads_div);
+        dr0  = fastdiv(src0_nrows + nth - 1, &octx->n_threads_div);
        dr1  = nr1;
        ith0 = ith;
        ith1 = 0;
    } else {
-        dr0  = nr0;
-        dr1  = fastdiv(nr1 + nth - 1, &octx->ctx->n_threads_div);
+        dr0  = src0_nrows;
+        dr1  = fastdiv(nr1 + nth - 1, &octx->n_threads_div);
        ith0 = 0;
        ith1 = ith;
    }

-    const uint32_t ir0_start = dr0 * ith0;
-    const uint32_t ir0_end   = MIN(ir0_start + dr0, nr0);
+    const uint32_t ir0_start = mmctx->src0_row_start + dr0 * ith0;
+    const uint32_t ir0_end   = MIN(ir0_start + dr0, mmctx->src0_row_end);

    const uint32_t ir1_start = dr1 * ith1;
    const uint32_t ir1_end   = MIN(ir1_start + dr1, nr1);
@@ -312,11 +334,11 @@ static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) {
 static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                                        \
    htp_matmul_preamble;                                                                                                          \
                                                                                                                                  \
-    const uint32_t src0_nrows = ne01 * ne02 * ne03;                                                                               \
+    const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start;                                                      \
    const uint32_t src1_nrows = ne11 * ne12 * ne13;                                                                               \
                                                                                                                                  \
-    const uint32_t src0_start_row  = src0_nrows_per_thread * ith;                                                                 \
-    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);                                     \
+    const uint32_t src0_start_row  = mmctx->src0_row_start + src0_nrows_per_thread * ith;                                         \
+    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end);                            \
                                                                                                                                  \
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                        \
                                                                                                                                  \
@@ -414,10 +436,10 @@ static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
 static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                                        \
    htp_matmul_preamble;                                                                                                          \
                                                                                                                                  \
-    const uint32_t src0_nrows = ne01;                                                                                             \
+    const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start;                                                      \
                                                                                                                                  \
-    const uint32_t src0_start_row  = src0_nrows_per_thread * ith;                                                                 \
-    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);                                     \
+    const uint32_t src0_start_row  = mmctx->src0_row_start + src0_nrows_per_thread * ith;                                         \
+    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end);                            \
                                                                                                                                  \
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                        \
                                                                                                                                  \
@@ -549,12 +571,22 @@ static void hvx_mm_nx_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, v
        uint32_t n_k_tiles_w = ne00 / 32;                                                                                         \
        uint32_t tile_row_stride = n_k_tiles_w * tile_size;                                                                       \
                                                                                                                                  \
-        const uint32_t src0_nrows = ne01 * src_w->ne[2] * src_w->ne[3];                                                           \
-        uint32_t src0_nrows_per_thread = fastdiv(src0_nrows + nth - 1, &octx->ctx->n_threads_div);                                \
+        uint32_t src0_start_row = 0;                                                                                              \
+        uint32_t src0_end_row   = ne01;                                                                                           \
+        if (octx->ctx->mdev.count > 1) {                                                                                          \
+            const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));                                              \
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0,                        \
+                                                         octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); \
+            src0_start_row = range.start;                                                                                         \
+            src0_end_row   = range.start + range.count;                                                                           \
+        }                                                                                                                         \
+                                                                                                                                  \
+        const uint32_t nrows = src0_end_row - src0_start_row;                                                                     \
+        uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div);                                          \
        src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32);                                                          \
                                                                                                                                  \
-        const uint32_t start_row = src0_nrows_per_thread * ith;                                                                   \
-        const uint32_t end_row   = MIN(start_row + src0_nrows_per_thread, src0_nrows);                                            \
+        const uint32_t start_row = src0_start_row + src0_nrows_per_thread * ith;                                                  \
+        const uint32_t end_row   = MIN(start_row + src0_nrows_per_thread, src0_end_row);                                          \
        if (start_row >= end_row) continue;                                                                                       \
                                                                                                                                  \
        uint32_t ct_start = start_row / 32;                                                                                       \
@@ -735,11 +767,11 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) {
    assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0);
    const uint32_t prefetch_mask = n_prefetch - 1;

-    const uint32_t src0_nrows = ne01 * ne02 * ne03;  // src0 rows
-    const uint32_t src1_nrows = ne11 * ne12 * ne13;  // src1 rows
+    const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start;  // src0 rows
+    const uint32_t src1_nrows = ne11 * ne12 * ne13;                          // src1 rows

-    const uint32_t src0_start_row  = src0_nrows_per_thread * ith;
-    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+    const uint32_t src0_start_row  = mmctx->src0_row_start + src0_nrows_per_thread * ith;
+    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end);
    const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U);

    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
@@ -781,7 +813,7 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) {
        const uint8_t * ss0 = dma_queue_pop(dma_queue).dst;

        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0);
-        // Process src1 columns in pairs (2×2 tiling)
+        // Process src1 columns in pairs (2x2 tiling)
        uint32_t ir1 = 0;
        for (; ir1 + 1 < src1_nrows; ir1 += 2) {
            const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride);
@@ -791,7 +823,7 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) {
            mmctx->vec_dot_2x2(ne00, &dst_row0[ir0], &dst_row1[ir0], ss0, ss0 + src0_stride, src1_col0, src1_col1);
        }

-        // Handle remaining src1 rows (fallback to 2×1)
+        // Handle remaining src1 rows (fallback to 2x1)
        for (; ir1 < src1_nrows; ++ir1) {
            const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride);
            float * restrict dst_row          = (float *) (dst->data + (ir1 * dst_row_size));
@@ -833,10 +865,10 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) {
 static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) {
    htp_matmul_preamble;

-    const uint32_t src0_nrows = ne01;
+    const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start;

-    const uint32_t src0_start_row  = src0_nrows_per_thread * ith;
-    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+    const uint32_t src0_start_row  = mmctx->src0_row_start + src0_nrows_per_thread * ith;
+    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end);

    struct htp_thread_trace * tr = &octx->ctx->trace[ith];

@@ -943,13 +975,10 @@ static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) {

    const struct htp_tensor * restrict ids = octx->src[2];

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
-    const uint32_t src0_nrows      = ne01;  // src0 rows per expert
+    const uint32_t src0_nrows      = mmctx->src0_row_end - mmctx->src0_row_start;  // src0 rows per expert
    const uint32_t src1_nrows      = ne11;
-    const uint32_t src0_start_row  = src0_nrows_per_thread * ith;
-    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+    const uint32_t src0_start_row  = mmctx->src0_row_start + src0_nrows_per_thread * ith;
+    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end);

    hvx_mm_run_quant_task(mmctx, ith);

@@ -1036,9 +1065,9 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) {

    const struct htp_tensor * restrict ids = octx->src[2];

-    const uint32_t src0_nrows      = ne01;  // src0 rows per expert
-    const uint32_t src0_start_row  = src0_nrows_per_thread * ith;
-    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+    const uint32_t src0_nrows      = mmctx->src0_row_end - mmctx->src0_row_start;  // src0 rows per expert
+    const uint32_t src0_start_row  = mmctx->src0_row_start + src0_nrows_per_thread * ith;
+    const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end);

    hvx_mm_run_quant_task(mmctx, ith);

@@ -1143,12 +1172,22 @@ static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) {
            const struct htp_tensor * restrict dst   = octx->dsts[p];
            if (!src_w || !dst) continue;

-            const uint32_t src0_nrows = src_w->ne[1];
-            uint32_t src0_nrows_per_thread = fastdiv(src0_nrows + nth - 1, &octx->ctx->n_threads_div);
+            const uint32_t ne01 = src_w->ne[1];
+            uint32_t start_row = 0;
+            uint32_t end_row   = ne01;
+            if (octx->ctx->mdev.count > 1) {
+                const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));
+                const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+                start_row = range.start;
+                end_row   = range.start + range.count;
+            }
+
+            const uint32_t nrows = end_row - start_row;
+            uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div);
            src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32);

-            const uint32_t src0_start_row = src0_nrows_per_thread * ith;
-            const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+            const uint32_t src0_start_row = start_row + src0_nrows_per_thread * ith;
+            const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, end_row);
            if (src0_start_row >= src0_end_row) continue;

            const uint8_t * restrict src0_row = (const uint8_t *) src_w->data + eid * src_w->nb[2];
@@ -1227,12 +1266,22 @@ static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) {
            const struct htp_tensor * restrict dst   = octx->dsts[p];
            if (!src_w || !dst) continue;

-            const uint32_t src0_nrows = src_w->ne[1];
-            uint32_t src0_nrows_per_thread = fastdiv(src0_nrows + nth - 1, &octx->ctx->n_threads_div);
+            const uint32_t ne01 = src_w->ne[1];
+            uint32_t start_row = 0;
+            uint32_t end_row   = ne01;
+            if (octx->ctx->mdev.count > 1) {
+                const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));
+                const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+                start_row = range.start;
+                end_row   = range.start + range.count;
+            }
+
+            const uint32_t nrows = end_row - start_row;
+            uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div);
            src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32);

-            const uint32_t src0_start_row = src0_nrows_per_thread * ith;
-            const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+            const uint32_t src0_start_row = start_row + src0_nrows_per_thread * ith;
+            const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, end_row);
            if (src0_start_row >= src0_end_row) continue;

            const uint8_t * src0_row = (const uint8_t *) src_w->data + cur_a * src_w->nb[2];
@@ -1323,15 +1372,33 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {

    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;

-    const uint32_t src0_nrows = ne01 * ne02 * ne03;
+    const uint32_t src0_nrows = ne01;
    const uint32_t src1_nrows = ne11 * ne12 * ne13;

+    uint32_t src0_row_start = 0;
+    uint32_t src0_row_end   = src0_nrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        src0_row_start = range.start;
+        src0_row_end   = range.start + range.count;
+    }
+
+    if (src0_row_start >= src0_row_end) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t nrows = src0_row_end - src0_row_start;
+    mmctx->src0_row_start = src0_row_start;
+    mmctx->src0_row_end   = src0_row_end;
+
    bool is_repacked = (src0->type == HTP_TYPE_Q4_0 || src0->type == HTP_TYPE_Q4_1 ||
                        src0->type == HTP_TYPE_Q8_0 || src0->type == HTP_TYPE_IQ4_NL ||
                        src0->type == HTP_TYPE_MXFP4);

    // Compute src0_nrows_per_thread
-    mmctx->src0_nrows_per_thread  = fastdiv(src0_nrows + octx->n_threads - 1, &octx->ctx->n_threads_div);
+    mmctx->src0_nrows_per_thread  = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div);
    if (is_repacked) {
        mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
    } else {
@@ -1503,13 +1570,13 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
        kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
        mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
    } else {
-        mmctx->vtcm_src1_size_per_thread = fastdiv(L.src1_bytes, &octx->ctx->n_threads_div);
+        mmctx->vtcm_src1_size_per_thread = fastdiv(L.src1_bytes, &octx->n_threads_div);
    }

-    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div);
-    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
+    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->n_threads_div);

-    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
+    const size_t vtcm_size = L.total_bytes;

    FARF(HIGH, "matmul-%s : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu (%zu)\n", mmctx->type,
         L.src0_bytes, L.src1_bytes, L.dst_bytes, vtcm_size);
@@ -1583,13 +1650,21 @@ static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) {

        const uint32_t ne00 = src_w->ne[0];
        const uint32_t ne01 = src_w->ne[1];
-        const uint32_t src0_nrows = ne01 * src_w->ne[2] * src_w->ne[3];
+        uint32_t start_row = 0;
+        uint32_t end_row   = ne01;
+        if (octx->ctx->mdev.count > 1) {
+            const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            start_row = range.start;
+            end_row   = range.start + range.count;
+        }

-        uint32_t src0_nrows_per_thread = fastdiv(src0_nrows + nth - 1, &octx->ctx->n_threads_div);
+        const uint32_t nrows = end_row - start_row;
+        uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div);
        src0_nrows_per_thread += (src0_nrows_per_thread & 1);

-        const uint32_t src0_start_row  = src0_nrows_per_thread * ith;
-        const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+        const uint32_t src0_start_row  = start_row + src0_nrows_per_thread * ith;
+        const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, end_row);
        const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U);
        if (src0_start_row >= src0_end_row) continue;

@@ -2638,10 +2713,6 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k
    const struct htp_tensor * restrict src0 = octx->src[0];
    const struct htp_tensor * restrict act  = octx->src[n_weights];

-    if (!src0 || !act) {
-        return HTP_STATUS_INVAL_PARAMS;
-    }
-
    const int weight_type = (int) src0->type;
    const int k           = (int) act->ne[0];
    const int k_valid     = (int) act->ne[0];
@@ -2714,16 +2785,31 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k

    hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00));  // scale: 1.0, bias: 0.0 in FP16

-    FARF(HIGH, "hmx-mm-nx-2d: n_weights %u m %d k %d wtype %d mc %d nc %d vtcm %zu/%zu",
-         n_weights, m, k, weight_type, m_chunk_n_rows, n_chunk_n_cols, L.total_bytes, vtcm_budget);
+    int m_start = 0;
+    int m_rows  = m;
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_can_row_partition(octx->dsts[0], sizeof(float));
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition((uint32_t) m, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        m_start = (int) range.start;
+        m_rows  = (int) range.count;
+    }
+
+    if (m_rows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    FARF(HIGH, "hmx-mm-nx-2d: n_weights %u m %d (%d..%d) k %d wtype %d mc %d nc %d vtcm %zu/%zu",
+         n_weights, m, m_start, m_start + m_rows, k, weight_type, m_chunk_n_rows, n_chunk_n_cols, L.total_bytes, vtcm_budget);

    htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);

+    const size_t mr_end = (size_t)(m_start + m_rows);
+
    if (pipeline) {
        hmx_matmul_job_t job_slots[2];

-        for (size_t mr = 0; mr < (size_t) m; mr += m_chunk_n_rows) {
-            const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows);
+        for (size_t mr = (size_t) m_start; mr < mr_end; mr += m_chunk_n_rows) {
+            const size_t n_rows = hex_smin(mr_end - mr, m_chunk_n_rows);

            void *vtcm_weight_bufs[2] = { vtcm_scratch0, vtcm_scratch1 };
            void *vtcm_output_bufs[2] = { vtcm_output,   vtcm_scratch2 };
@@ -2822,8 +2908,8 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k
        }
    } else {
        hmx_matmul_job_t job;
-        for (size_t mr = 0; mr < (size_t) m; mr += m_chunk_n_rows) {
-            const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows);
+        for (size_t mr = (size_t) m_start; mr < mr_end; mr += m_chunk_n_rows) {
+            const size_t n_rows = hex_smin(mr_end - mr, m_chunk_n_rows);

            struct activation_transfer_params act_params = {
                .ctx = ctx,
@@ -3095,7 +3181,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_
                            int chunk_dst_cols = params->n - (int)nc;
                            if (chunk_dst_cols > 0) {
                                transfer_output_chunk_threaded(ctx, output, src2_chunk, vtcm_output, (int) n_rows, (int) n_cols,
-                                                               params->dst_stride, params->src2_stride, chunk_dst_cols, ctx->n_threads);
+                                                               params->dst_stride, params->src2_stride, chunk_dst_cols, n_threads);
                            }
                        }
                    }
@@ -3216,7 +3302,10 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx,
                                         int weight_type,
                                         const struct mmid_row_mapping *matrix_rows,
                                         int cur_a,
-                                         int mapping_stride) {
+                                         int mapping_stride,
+                                         int m_start,
+                                         int m_end,
+                                         int n_threads) {
    struct htp_thread_trace * tr = &ctx->trace[0];
    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);

@@ -3247,7 +3336,6 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx,
    const int n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS;
    const struct fastdiv_values n_k_tiles_div = init_fastdiv_values(n_k_tiles);

-    const int n_threads = ctx->n_threads;
    const bool is_quant   = (weight_type != HTP_TYPE_F16 && weight_type != HTP_TYPE_F32);

    const size_t vec_dot_size = k * sizeof(__fp16);
@@ -3303,8 +3391,8 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx,

    hmx_matmul_job_t job;

-    for (size_t mr = 0; mr < (size_t) m_padded; mr += m_chunk_n_rows) {
-        const size_t n_rows = hex_smin(m_padded - mr, m_chunk_n_rows);
+    for (size_t mr = (size_t) m_start; mr < (size_t) m_end; mr += m_chunk_n_rows) {
+        const size_t n_rows = hex_smin((size_t) m_end - mr, m_chunk_n_rows);
        const size_t n_row_tiles = hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS);

        transfer_activation_chunk_gathered_threaded(
@@ -3368,31 +3456,48 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k
    const int act_stride = (int)(src1->nb[1] / sizeof(float));
    const int wgt_stride = (int)(src0->nb[1] / sizeof(__fp16));

+    int m_start = 0;
+    int m_rows  = m_total;
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition((uint32_t) m_total, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        m_start = (int) range.start;
+        m_rows  = (int) range.count;
+    }
+
+    if (m_rows == 0) {
+        return HTP_STATUS_OK;
+    }
+
    const float * src2_ptr = NULL;
    uint32_t src2_stride = 0;
    size_t src2_nb2 = 0;
    size_t src2_nb3 = 0;
    if (src2) {
-        src2_ptr = (const float *) src2->data;
        src2_stride = (src2->ne[1] == 1) ? 0 : (uint32_t) (src2->nb[1] / sizeof(float));
+        src2_ptr = (const float *) src2->data + m_start * src2_stride;
        src2_nb2 = (src2->ne[2] == 1) ? 0 : src2->nb[2];
        src2_nb3 = (src2->ne[3] == 1) ? 0 : src2->nb[3];
    }

+    const int dst_stride = (int)(dst->nb[1] / sizeof(float));
+    float       * dst_ptr = (float *)       dst->data  + m_start * dst_stride;
+    const float * act_ptr = (const float *) src1->data + m_start * act_stride;
+
    int ret = -1;
-    const int n_threads = MIN(kparams->n_threads, (int) octx->n_threads);
+    const int n_threads = kparams->n_threads;
    if (kparams->kernel_type == HTP_MM_KERNEL_HMX_F16_BATCHED) {
        hmx_mm_f16_f32_batched_params_t batch_params = {
-            .dst             = (float *) dst->data,
+            .dst             = dst_ptr,
            .src2            = src2_ptr,
-            .activation      = (float *) src1->data,
+            .activation      = act_ptr,
            .weight          = (const __fp16 *) src0->data,
-            .m               = m_total,
+            .m               = m_rows,
            .k               = k,
            .n               = n,
            .act_stride      = act_stride,
            .weight_stride   = wgt_stride,
-            .dst_stride      = (int) (dst->nb[1] / sizeof(float)),
+            .dst_stride      = dst_stride,
            .src2_stride     = src2_stride,
            .ne02            = ne02,
            .ne03            = ne03,
@@ -3420,9 +3525,9 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k
                                     kparams->vtcm_size);
    } else {
        ret = hmx_mm_2d_f32(
-            octx->ctx, (float*) dst->data, src2_ptr, (float*) src1->data, (const uint8_t *) src0->data,
-            m_total, k, n, act_stride, (int) src0->nb[1], (int) src0->type, (int) src1->ne[0],
-            (int)(dst->nb[1] / sizeof(float)), src2_stride, (int)dst->ne[0],
+            octx->ctx, dst_ptr, src2_ptr, act_ptr, (const uint8_t *) src0->data,
+            m_rows, k, n, act_stride, (int) src0->nb[1], (int) src0->type, (int) src1->ne[0],
+            dst_stride, src2_stride, (int)dst->ne[0],
            kparams->m_chunk, kparams->n_chunk, kparams->pipeline, n_threads,
            kparams->n_act_threads,
            &kparams->div_n_act_threads,
@@ -3441,6 +3546,11 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k
 int op_matmul(struct htp_ops_context * octx) {
    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;

+    const int status = htp_mm_init_context(octx, kparams);
+    if (status != HTP_STATUS_OK) {
+        return status;
+    }
+
    if (kparams->n_hmx) {
        return hmx_mm_op_matmul(octx, kparams);
    }
@@ -3463,6 +3573,16 @@ static int hmx_mm_op_matmul_id(
        const int32_t cne1 = matrix_row_counts[cur_a];
        if (cne1 == 0) continue;

+        const int m_padded = hex_align_up(cne1, 32);
+        int m_start = 0, m_end = m_padded;
+        if (octx->ctx->mdev.count > 1) {
+            const bool can_split = htp_tensor_mdev_data_aligned(dst) && (uint32_t) cne1 >= octx->ctx->mdev.count;
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition((uint32_t) m_padded, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            m_start = (int) range.start;
+            m_end   = (int) (range.start + range.count);
+        }
+        if (m_start >= m_end) continue;
+
        int ret = hmx_mm_id_2d_f32(octx->ctx, (float*) dst->data, (float*) src1->data,
                                   (const uint8_t *) src0->data + cur_a * nb02,
                                   cne1, ne00, ne01,
@@ -3471,7 +3591,8 @@ static int hmx_mm_op_matmul_id(
                                   nb11, nb12,
                                   nb1, nb2,
                                   (int) src0->nb[1], (int) src0->type,
-                                   matrix_rows, cur_a, mmctx->mapping_stride);
+                                   matrix_rows, cur_a, mmctx->mapping_stride,
+                                   m_start, m_end, (int) octx->n_threads);
        if (ret != 0) {
            FARF(ERROR, "HMX matmul failed for expert %u, error %d\n", cur_a, ret);
            return HTP_STATUS_NO_SUPPORT;
@@ -3524,7 +3645,7 @@ static int hvx_mm_matmul_id(
    htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
                                 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false);

-    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
+    const size_t vtcm_size = L.total_bytes;

    FARF(HIGH, "matmul-id-%s : src0-spad-size %zu src1-spad-size %zu src2-spad-size 0 dst-spad-size %zu (%zu)\n", mmctx->type,
         L.src0_bytes, L.src1_bytes, L.dst_bytes, vtcm_size);
@@ -3554,10 +3675,10 @@ static int hvx_mm_matmul_id(
    mmctx->vtcm_src0_stride = src0_row_size_padded;
    mmctx->vtcm_src1_stride = src1_row_size;

-    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
    mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
    mmctx->vtcm_src2_size_per_thread = 0;
-    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->n_threads_div);

    mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
    mmctx->quant_task_func = quant_task_func;
@@ -3587,6 +3708,20 @@ static int hmx_mm_op_matmul_id_nx(
        const int32_t cne1 = matrix_row_counts[cur_a];
        if (cne1 == 0) continue;

+        const int m_padded = hex_align_up(cne1, 32);
+        int m_start = 0, m_end = m_padded;
+        if (octx->ctx->mdev.count > 1) {
+            bool can_split = (uint32_t) cne1 >= octx->ctx->mdev.count;
+            for (uint32_t p = 0; p < n_weights && can_split; ++p) {
+                const struct htp_tensor * restrict dst = octx->dsts[p];
+                can_split = !dst || htp_tensor_mdev_data_aligned(dst);
+            }
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition((uint32_t) m_padded, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            m_start = (int) range.start;
+            m_end   = (int) (range.start + range.count);
+        }
+        if (m_start >= m_end) continue;
+
        for (uint32_t p = 0; p < n_weights; ++p) {
            const struct htp_tensor * restrict src_w = octx->src[p];
            const struct htp_tensor * restrict dst   = octx->dsts[p];
@@ -3600,7 +3735,8 @@ static int hmx_mm_op_matmul_id_nx(
                                       act->nb[1], act->nb[2],
                                       dst->nb[1], dst->nb[2],
                                       (int) src_w->nb[1], (int) src_w->type,
-                                       matrix_rows, cur_a, mmctx->mapping_stride);
+                                       matrix_rows, cur_a, mmctx->mapping_stride,
+                                       m_start, m_end, (int) octx->n_threads);
            if (ret != 0) {
                FARF(ERROR, "HMX matmul ID NX failed for expert %u weight %u, error %d\n", cur_a, p, ret);
                return HTP_STATUS_NO_SUPPORT;
@@ -3656,7 +3792,7 @@ static int hvx_mm_matmul_id_nx(
    htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], src1_nrows, octx->n_threads,
                                 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false);

-    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
+    const size_t vtcm_size = L.total_bytes;

    if (octx->ctx->vtcm_size < vtcm_size) {
        FARF(ERROR, "matmul-id-nx: current VTCM reservation %zu is too small, needed %zu\n",
@@ -3678,9 +3814,9 @@ static int hvx_mm_matmul_id_nx(
    mmctx->vtcm_src0_stride = 0;
    mmctx->vtcm_src1_stride = src1_row_size;

-    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
    mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
-    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->n_threads_div);

    mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
    mmctx->quant_task_func         = quant_task_func;
@@ -3769,16 +3905,21 @@ static inline void scan_expert_ids(
 int op_matmul_id(struct htp_ops_context * octx) {
    htp_matmul_tensors_preamble;

+    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
+    struct htp_mm_context mmctx_struct = {0};
+    struct htp_mm_context * mmctx = &mmctx_struct;
+
+    const int status = htp_mm_init_context(octx, kparams);
+    if (status != HTP_STATUS_OK) {
+        return status;
+    }
+
    struct htp_thread_trace * tr = &octx->ctx->trace[0];
    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);

-    struct htp_mm_context mmctx_struct = {0};
-    struct htp_mm_context * mmctx = &mmctx_struct;
    mmctx->octx = octx;
    mmctx->act = src1;

-    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
-
    const struct htp_tensor * restrict ids = octx->src[2];

    const size_t src0_row_size = nb01;
@@ -3789,9 +3930,6 @@ int op_matmul_id(struct htp_ops_context * octx) {
    const uint32_t src0_nrows = ne01;  // per expert
    const uint32_t src1_nrows = ne11 * ne12 * ne13;

-    mmctx->src0_nrows_per_thread = fastdiv(src0_nrows + octx->n_threads - 1, &octx->ctx->n_threads_div);
-    mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
-
    // row groups
    const int n_ids = ids->ne[0];  // n_expert_used
    const int n_as  = ne02;        // n_expert
@@ -3843,6 +3981,29 @@ int op_matmul_id(struct htp_ops_context * octx) {
    if (kparams->n_hmx) {
        s = hmx_mm_op_matmul_id(octx, mmctx);
    } else {
+        uint32_t src0_row_start = 0;
+        uint32_t src0_row_end   = src0_nrows;
+        if (octx->ctx->mdev.count > 1) {
+            const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float));
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            src0_row_start = range.start;
+            src0_row_end   = range.start + range.count;
+        }
+
+        if (src0_row_start >= src0_row_end) {
+            if (mapping_buf != octx->ctx->ddr_spad_base) {
+                free(mapping_buf);
+            }
+            return HTP_STATUS_OK;
+        }
+
+        const uint32_t nrows = src0_row_end - src0_row_start;
+        mmctx->src0_row_start = src0_row_start;
+        mmctx->src0_row_end   = src0_row_end;
+
+        mmctx->src0_nrows_per_thread = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div);
+        mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
+
        if (hvx_mm_init_vec_dot(mmctx, src0->type) == 0) {
            s = hvx_mm_matmul_id(octx, mmctx, src1_nrows > 1 ? hvx_mm_id : hvx_mv_id);
        } else {
@@ -3858,29 +4019,31 @@ int op_matmul_id(struct htp_ops_context * octx) {
 }

 int op_matmul_id_nx(struct htp_ops_context * octx) {
+    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
+    struct htp_mm_context mmctx_struct = {0};
+    struct htp_mm_context * mmctx = &mmctx_struct;
+
+    const int status = htp_mm_init_context(octx, kparams);
+    if (status != HTP_STATUS_OK) {
+        return status;
+    }
+
    struct htp_thread_trace * tr = &octx->ctx->trace[0];
    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);

-    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
+    mmctx->octx = octx;
    const uint32_t n_weights = kparams->n_weights;
    const struct htp_tensor * restrict src0 = octx->src[0];
    const struct htp_tensor * restrict act  = octx->src[n_weights];
    const struct htp_tensor * restrict ids  = octx->src[n_weights + 1];

-    struct htp_mm_context mmctx_struct = {0};
-    struct htp_mm_context * mmctx = &mmctx_struct;
-    mmctx->octx = octx;
    mmctx->act = act;

    const size_t src0_row_size = src0->nb[1];
    const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);

-    const uint32_t src0_nrows = src0->ne[1];
    const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3];

-    mmctx->src0_nrows_per_thread = fastdiv(src0_nrows + octx->n_threads - 1, &octx->ctx->n_threads_div);
-    mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32);
-
    const int n_ids = ids->ne[0];
    const int n_as  = src0->ne[2];

@@ -3946,6 +4109,12 @@ int op_matmul_id_nx(struct htp_ops_context * octx) {
 }
 int op_matmul_nx(struct htp_ops_context * octx) {
    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
+
+    const int status = htp_mm_init_context(octx, kparams);
+    if (status != HTP_STATUS_OK) {
+        return status;
+    }
+
    if (kparams->n_hmx) {
        return hmx_mm_nx_2d_f32(octx, kparams);
    }
@@ -4012,7 +4181,7 @@ int op_matmul_nx(struct htp_ops_context * octx) {
    htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], src1_nrows, octx->n_threads,
                                 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, true);

-    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
+    const size_t vtcm_size = L.total_bytes;

    if (octx->ctx->vtcm_size < vtcm_size) {
        FARF(ERROR, "matmul-nx: current VTCM reservation %zu is too small, needed %zu\n",
@@ -4034,9 +4203,9 @@ int op_matmul_nx(struct htp_ops_context * octx) {
    mmctx->vtcm_src0_stride = is_repacked ? 0 : src0_row_size_padded;
    mmctx->vtcm_src1_stride = src1_row_size;

-    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_src0_size_per_thread = fastdiv(L.src0_bytes, &octx->n_threads_div);
    mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
-    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->ctx->n_threads_div);
+    mmctx->vtcm_dst_size_per_thread  = fastdiv(L.dst_bytes, &octx->n_threads_div);

    mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
    mmctx->quant_task_func = quant_task_func;
@@ -12,8 +12,11 @@

 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+#include "hex-common.h"
+#include "hex-profile.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"

 /* Circular wrap: maps any integer x into [0, n) */
 static inline uint32_t wrap_around(int32_t x, uint32_t n) {
@@ -68,6 +71,7 @@ struct htp_pad_context {

    uint32_t nrows_per_thread;
    uint32_t total_dst_rows;
+    uint32_t row_start;

    size_t   type_size;

@@ -78,39 +82,39 @@ struct htp_pad_context {
    size_t   dst_row_size_aligned;
 };

-#define htp_pad_preamble                            \
-    const struct htp_tensor * src = octx->src[0];   \
-    const struct htp_tensor * dst = octx->dst;      \
-                                                    \
-    const uint32_t ne00 = src->ne[0];               \
-    const uint32_t nb00 = src->nb[0];               \
-                                                    \
-    const uint32_t ne0 = dst->ne[0];                \
-    const uint32_t ne1 = dst->ne[1];                \
-    const uint32_t ne2 = dst->ne[2];                \
-    const uint32_t ne3 = dst->ne[3];                \
-                                                    \
-    const uint32_t nb1 = dst->nb[1];                \
-    const uint32_t nb2 = dst->nb[2];                \
-    const uint32_t nb3 = dst->nb[3];                \
-                                                    \
-    const int32_t lp0 = pctx->lp0, rp0 = pctx->rp0; \
-    const int32_t lp1 = pctx->lp1, rp1 = pctx->rp1; \
-    const int32_t lp2 = pctx->lp2, rp2 = pctx->rp2; \
-    const int32_t lp3 = pctx->lp3, rp3 = pctx->rp3; \
-                                                    \
-    const size_t type_size = pctx->type_size;       \
-                                                    \
-    const uint32_t row_start = pctx->nrows_per_thread * ith;                                 \
-    const uint32_t row_end   = MIN(row_start + pctx->nrows_per_thread, pctx->total_dst_rows);
+#define htp_pad_preamble                                                             \
+    const struct htp_tensor * src = octx->src[0];                                    \
+    const struct htp_tensor * dst = octx->dst;                                       \
+                                                                                     \
+    const uint32_t ne00 = src->ne[0];                                                \
+    const uint32_t nb00 = src->nb[0];                                                \
+                                                                                     \
+    const uint32_t ne0 = dst->ne[0];                                                 \
+    const uint32_t ne1 = dst->ne[1];                                                 \
+    const uint32_t ne2 = dst->ne[2];                                                 \
+    const uint32_t ne3 = dst->ne[3];                                                 \
+                                                                                     \
+    const uint32_t nb1 = dst->nb[1];                                                 \
+    const uint32_t nb2 = dst->nb[2];                                                 \
+    const uint32_t nb3 = dst->nb[3];                                                 \
+                                                                                     \
+    const int32_t lp0 = pctx->lp0, rp0 = pctx->rp0;                                  \
+    const int32_t lp1 = pctx->lp1, rp1 = pctx->rp1;                                  \
+    const int32_t lp2 = pctx->lp2, rp2 = pctx->rp2;                                  \
+    const int32_t lp3 = pctx->lp3, rp3 = pctx->rp3;                                  \
+                                                                                     \
+    const size_t type_size = pctx->type_size;                                        \
+                                                                                     \
+    const uint32_t row_start = pctx->row_start + pctx->nrows_per_thread * ith;       \
+    const uint32_t row_end   = MIN(row_start + pctx->nrows_per_thread, pctx->row_start + pctx->total_dst_rows);


-#define htp_pad_dma_preamble                                        \
-    const size_t src_row_size         = pctx->src_row_size;         \
-    const size_t src_row_size_aligned = pctx->src_row_size_aligned; \
-    const size_t dst_row_size         = pctx->dst_row_size;         \
-    const size_t dst_row_size_aligned = pctx->dst_row_size_aligned; \
-                                                                    \
+#define htp_pad_dma_preamble                                                                \
+    const size_t src_row_size         = pctx->src_row_size;                                 \
+    const size_t src_row_size_aligned = pctx->src_row_size_aligned;                         \
+    const size_t dst_row_size         = pctx->dst_row_size;                                 \
+    const size_t dst_row_size_aligned = pctx->dst_row_size_aligned;                         \
+                                                                                            \
    uint8_t * src_spad_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread; \
    uint8_t * dst_spad_base = octx->dst_spad.data  + ith * octx->dst_spad.size_per_thread;  \
                                                                                            \
@@ -125,8 +129,8 @@ static void pad_job_per_thread_hvx(unsigned int nth, unsigned int ith, void * da
    struct htp_ops_context * octx = pctx->octx;
    htp_pad_preamble;

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, row_start);

    for (uint32_t dst_row = row_start; dst_row < row_end; dst_row++) {
        uint32_t i1, i2, i3;
@@ -165,18 +169,17 @@ static void pad_job_per_thread_hvx(unsigned int nth, unsigned int ith, void * da
        }
    }

-    t2 = HAP_perf_get_qtimer_count();
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, row_start);

-    FARF(HIGH, "pad-hvx %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u usec %u\n",
+    FARF(HIGH, "pad-hvx %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n",
         ith, nth,
         src->ne[0], src->ne[1], src->ne[2], src->ne[3],
         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         row_start, row_end,
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         row_start, row_end);
 }

 // ---------------------------------------------------------------------------
-// HVX + DMA PAD kernel — aligned, double-buffered
+// HVX + DMA PAD kernel - aligned, double-buffered
 // ---------------------------------------------------------------------------

 static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void * data) {
@@ -185,9 +188,6 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
    htp_pad_preamble;
    htp_pad_dma_preamble;

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
    // -----------------------------------------------------------------------
    // Priming phase: push 2 pairs of (dummy_dst_DMA, src_DMA) to seed the
    // double-buffer pipeline before the main loop begins.
@@ -222,6 +222,8 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
    // Main loop: pop completed DMAs, compute in VTCM with aligned HVX ops,
    // push dst DMA and prefetch src for the next+1 row.
    // -----------------------------------------------------------------------
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
    for (uint32_t ir = row_start; ir < row_end; ir++) {
        uint8_t * dst_spad_cur = (uint8_t *) dma_queue_pop(dma).src;
        uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma).dst;
@@ -236,6 +238,7 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
                                             lp2, rp2, ne2,
                                             lp3, rp3, ne3);

+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
        if (!interior) {
            hvx_splat_f32_a(dst_spad_cur, 0.0f, ne0);
        } else {
@@ -249,6 +252,7 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void
                hvx_copy_f32_ua(dst_interior, src_spad_cur, ne00);
            }
        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);

        dma_queue_push_vtcm_to_ddr(dma,
            dma_make_ptr(dst_ptr, dst_spad_cur),
@@ -274,14 +278,11 @@ static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void

    dma_queue_flush(dma);

-    t2 = HAP_perf_get_qtimer_count();
-
-    FARF(HIGH, "pad-hvx-dma %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u usec %u\n",
+    FARF(HIGH, "pad-hvx-dma %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n",
         ith, nth,
         src->ne[0], src->ne[1], src->ne[2], src->ne[3],
         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         row_start, row_end,
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         row_start, row_end);
 }

 // ---------------------------------------------------------------------------
@@ -293,8 +294,8 @@ static void pad_job_per_thread_hvx_circular(unsigned int nth, unsigned int ith,
    struct htp_ops_context * octx = pctx->octx;
    htp_pad_preamble;

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, row_start);

    for (uint32_t dst_row = row_start; dst_row < row_end; dst_row++) {
        uint32_t i1, i2, i3;
@@ -344,18 +345,17 @@ static void pad_job_per_thread_hvx_circular(unsigned int nth, unsigned int ith,
        }
    }

-    t2 = HAP_perf_get_qtimer_count();
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, row_start);

-    FARF(HIGH, "pad-hvx-circ %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u usec %u\n",
+    FARF(HIGH, "pad-hvx-circ %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n",
         ith, nth,
         src->ne[0], src->ne[1], src->ne[2], src->ne[3],
         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         row_start, row_end,
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         row_start, row_end);
 }

 // ---------------------------------------------------------------------------
-// HVX + DMA circular PAD kernel — aligned, double-buffered
+// HVX + DMA circular PAD kernel - aligned, double-buffered
 // ---------------------------------------------------------------------------

 static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int ith, void * data) {
@@ -364,9 +364,6 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
    htp_pad_preamble;
    htp_pad_dma_preamble;

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
    // -----------------------------------------------------------------------
    // Priming phase: push 2 pairs of (dummy_dst_DMA, src_DMA) to seed the
    // double-buffer pipeline.  Every row is a real src DMA (no null DMAs).
@@ -390,6 +387,8 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
    // Main loop: pop completed DMAs, assemble circular row in VTCM with
    // aligned HVX ops, push dst DMA and prefetch src for the next+1 row.
    // -----------------------------------------------------------------------
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
    for (uint32_t ir = row_start; ir < row_end; ir++) {
        uint8_t * dst_spad_cur = (uint8_t *) dma_queue_pop(dma).src;
        uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma).dst;
@@ -398,7 +397,7 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
        pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3);
        uint8_t * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;

-
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
        if (lp0 > 0) {
            uint8_t * dst_left       = dst_spad_cur;
            const uint8_t * src_left = src_spad_cur + (size_t)(ne00 - (uint32_t)lp0) * type_size;
@@ -430,6 +429,7 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i
                }
            }
        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);

        dma_queue_push_vtcm_to_ddr(dma,
            dma_make_ptr(dst_ptr, dst_spad_cur),
@@ -448,14 +448,11 @@ static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int i

    dma_queue_flush(dma);

-    t2 = HAP_perf_get_qtimer_count();
-
-    FARF(HIGH, "pad-hvx-circ-dma %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u usec %u\n",
+    FARF(HIGH, "pad-hvx-circ-dma %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n",
         ith, nth,
         src->ne[0], src->ne[1], src->ne[2], src->ne[3],
         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         row_start, row_end,
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         row_start, row_end);
 }

 int op_pad(struct htp_ops_context * octx) {
@@ -489,19 +486,33 @@ int op_pad(struct htp_ops_context * octx) {
    const uint32_t ne00 = src0->ne[0];

    const uint32_t total_dst_rows = dst->ne[1] * dst->ne[2] * dst->ne[3];
-    const uint32_t n_threads = MIN(octx->n_threads, total_dst_rows > 0 ? total_dst_rows : 1);
+    const size_t dst_row_size     = (size_t)ne0 * type_size;
+
+    uint32_t row_start = 0;
+    uint32_t nrows     = total_dst_rows;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, type_size, (uint32_t) dst_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_dst_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;

    const size_t src_row_size         = (size_t)ne00 * type_size;
-    const size_t dst_row_size         = (size_t)ne0  * type_size;
    const size_t src_row_size_aligned = hex_round_up(src_row_size, VLEN);
    const size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN);

    // Total VTCM needed: 2 buffers (ping+pong) for src and dst, per thread
    const size_t vtcm_needed = (size_t)n_threads * 2 * (src_row_size_aligned + dst_row_size_aligned);

-    const int use_dma = (src0->nb[0] == (uint32_t)type_size) &&
-                        (ne00 >= 512) &&
-                        (octx->ctx->vtcm_base != NULL) &&
+    const int use_dma = (src0->nb[0] == (uint32_t)type_size) && (ne00 >= 512) &&
                        (octx->ctx->vtcm_size >= vtcm_needed);

    if (use_dma) {
@@ -521,8 +532,9 @@ int op_pad(struct htp_ops_context * octx) {
        .lp1 = lp1, .rp1 = rp1,
        .lp2 = lp2, .rp2 = rp2,
        .lp3 = lp3, .rp3 = rp3,
-        .nrows_per_thread = (total_dst_rows + n_threads - 1) / n_threads,
-        .total_dst_rows   = total_dst_rows,
+        .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+        .total_dst_rows   = nrows,
+        .row_start        = row_start,
        .type_size        = type_size,
        .src_row_size         = src_row_size,
        .src_row_size_aligned = src_row_size_aligned,
@@ -537,11 +549,10 @@ int op_pad(struct htp_ops_context * octx) {
         dst->ne[0],  dst->ne[1],  dst->ne[2],  dst->ne[3],
         lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);

-    if      (circular && use_dma) { worker_pool_run_func(octx->ctx->worker_pool, pad_job_per_thread_hvx_circular_dma, &pctx, n_threads); }
-    else if (circular)            { worker_pool_run_func(octx->ctx->worker_pool, pad_job_per_thread_hvx_circular,     &pctx, n_threads); }
-    else if (use_dma)             { worker_pool_run_func(octx->ctx->worker_pool, pad_job_per_thread_hvx_dma,          &pctx, n_threads); }
-    else                          { worker_pool_run_func(octx->ctx->worker_pool, pad_job_per_thread_hvx,              &pctx, n_threads); }
+    if      (circular && use_dma) { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_circular_dma, &pctx, n_threads); }
+    else if (circular)            { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_circular,     &pctx, n_threads); }
+    else if (use_dma)             { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_dma,          &pctx, n_threads); }
+    else                          { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx,              &pctx, n_threads); }

    return HTP_STATUS_OK;
 }
-
@@ -12,8 +12,10 @@
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
 #include "htp-ctx.h"
+#include "hex-common.h"
+#include "hex-profile.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
+#include "htp-tensor.h"

 struct htp_repeat_context {
    struct htp_ops_context * octx;
@@ -25,6 +27,7 @@ struct htp_repeat_context {

    uint32_t nrows_per_thread;
    uint32_t total_dst_rows;  // ne1 * ne2 * ne3
+    uint32_t row_start;

    size_t   type_size;
 };
@@ -62,11 +65,11 @@ static void repeat_job_per_thread(unsigned int nth, unsigned int ith, void * dat

    const size_t row_bytes = ne00 * rctx->type_size;

-    const uint32_t row_start = rctx->nrows_per_thread * ith;
-    const uint32_t row_end   = MIN(row_start + rctx->nrows_per_thread, rctx->total_dst_rows);
+    const uint32_t row_start = rctx->row_start + rctx->nrows_per_thread * ith;
+    const uint32_t row_end   = MIN(row_start + rctx->nrows_per_thread, rctx->row_start + rctx->total_dst_rows);

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, row_start);

    for (uint32_t dst_row = row_start; dst_row < row_end; dst_row++) {
        // Decompose flat dst row index into (i1, i2, i3)
@@ -89,12 +92,12 @@ static void repeat_job_per_thread(unsigned int nth, unsigned int ith, void * dat
        }
    }

-    t2 = HAP_perf_get_qtimer_count();
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, row_start);

-    FARF(HIGH, "repeat %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u usec %u\n",
+    FARF(HIGH, "repeat %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n",
         ith, nth, src->ne[0], src->ne[1], src->ne[2], src->ne[3],
         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-         row_start, row_end, (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         row_start, row_end);
 }

 int op_repeat(struct htp_ops_context * octx) {
@@ -119,21 +122,39 @@ int op_repeat(struct htp_ops_context * octx) {
            return HTP_STATUS_NO_SUPPORT;
    }

-    const uint32_t total_dst_rows = dst->ne[1] * dst->ne[2] * dst->ne[3];
-    const uint32_t n_threads = MIN(octx->n_threads, total_dst_rows);
-
    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
        return HTP_STATUS_OK;
    }

+    const uint32_t total_dst_rows  = dst->ne[1] * dst->ne[2] * dst->ne[3];
+    const size_t dst_row_size = dst->ne[0] * type_size;
+
+    uint32_t row_start = 0;
+    uint32_t nrows     = total_dst_rows;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, type_size, (uint32_t) dst_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_dst_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
    struct htp_repeat_context rctx = {
        .octx             = octx,
        .nr0              = dst->ne[0] / src0->ne[0],
        .nr1              = dst->ne[1] / src0->ne[1],
        .nr2              = dst->ne[2] / src0->ne[2],
        .nr3              = dst->ne[3] / src0->ne[3],
-        .nrows_per_thread = (total_dst_rows + n_threads - 1) / n_threads,
-        .total_dst_rows   = total_dst_rows,
+        .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+        .total_dst_rows   = nrows,
+        .row_start        = row_start,
        .type_size        = type_size,
    };

@@ -142,7 +163,7 @@ int op_repeat(struct htp_ops_context * octx) {
         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
         rctx.nr0, rctx.nr1, rctx.nr2, rctx.nr3);

-    worker_pool_run_func(octx->ctx->worker_pool, repeat_job_per_thread, &rctx, n_threads);
+    work_queue_run(octx->ctx->work_queue, repeat_job_per_thread, &rctx, n_threads);

    return HTP_STATUS_OK;
 }
@@ -17,8 +17,8 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
 #include "htp-tensor.h"
+#include "rope-ops.h"

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

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

    uint32_t src0_nrows_per_thread;
-    size_t spad_stride;

    struct htp_ops_context * octx;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -466,31 +541,29 @@ static void rope_job_f32(unsigned int nth, unsigned int ith, void * data) {

    htp_rope_preamble;

-    const uint32_t src0_nrows = rctx->src0_nrows;
+    const uint32_t src0_nrows = rctx->nrows;
    const uint32_t src0_nrows_per_thread = rctx->src0_nrows_per_thread;

-    const uint32_t src0_start_row = src0_nrows_per_thread * ith;
-    const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+    const uint32_t src0_start_row = rctx->row_start + src0_nrows_per_thread * ith;
+    const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, rctx->row_start + src0_nrows);

    // no work for this thread
    if (src0_start_row >= src0_end_row) {
        return;
    }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

+    const uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
+
+    uint32_t row_start = 0;
+    uint32_t nrows     = total_rows;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_data_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+            total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
+    const uint32_t ne0 = dst->ne[0];
    const size_t src0_row_size   = src0->ne[0] * sizeof(float);
    const size_t src0_row_stride = src0->nb[1];
    const size_t dst_row_size    = dst->ne[0] * sizeof(float);
    const size_t dst_row_stride  = dst->nb[1];

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

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

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

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

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

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

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

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

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

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

+#include "hex-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
@@ -58,6 +59,9 @@ struct set_rows_context {
    const struct htp_set_rows_kernel_params * kparams;
    struct htp_set_rows_vtcm_layout vtcm_layout;
    uint8_t * vtcm_base;
+    uint32_t task_start;
+    uint32_t tasks;
+    uint32_t tasks_per_thread;
 };

 #define SET_ROWS_THREAD_DMA_FN(TYPE_NAME, IDX_TYPE, COMPUTE_EXPR)                                                \
@@ -67,12 +71,12 @@ static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsig
    const struct htp_set_rows_kernel_params * kparams = srctx->kparams;                                          \
    set_rows_preamble;                                                                                           \
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                       \
-    const uint32_t dr  = kparams->tasks_per_thread;                                                              \
-    const uint32_t ir0 = dr * ith;                                                                               \
-    if (ir0 >= kparams->total_tasks) {                                                                           \
+    const uint32_t dr  = srctx->tasks_per_thread;                                                                \
+    const uint32_t ir0 = srctx->task_start + dr * ith;                                                           \
+    if (ir0 >= srctx->task_start + srctx->tasks) {                                                               \
        return;                                                                                                  \
    }                                                                                                            \
-    const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks);                                                    \
+    const uint32_t ir1 = MIN(ir0 + dr, srctx->task_start + srctx->tasks);                                        \
    dma_queue * dma_queue = octx->ctx->dma[ith];                                                                 \
    const struct htp_set_rows_vtcm_layout * vtcm_layout = &srctx->vtcm_layout;                                   \
    uint8_t * vtcm_src0 = srctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread;   \
@@ -192,18 +196,44 @@ int op_set_rows(struct htp_ops_context * octx) {
        return HTP_STATUS_NO_SUPPORT;
    }

-    if (octx->src[1]->type != HTP_TYPE_I32 && octx->src[1]->type != HTP_TYPE_I64) {
-        return HTP_STATUS_NO_SUPPORT;
+    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
+        return HTP_STATUS_OK;
    }

+    const struct htp_tensor * dst = octx->dst;
+    const uint32_t total_tasks    = kparams->total_tasks;
+
+    uint32_t task_start = 0;
+    uint32_t tasks      = total_tasks;
+
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_mdev_data_aligned(dst) && (dst->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0 && !htp_tensor_is_permuted(dst);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_tasks, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        task_start = range.start;
+        tasks      = range.count;
+    }
+
+    if (tasks == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
    // l2fetch the src1 (indices) tensor in the main thread
    hex_l2fetch_block((const void *)octx->src[1]->data, octx->src[1]->ne[3] * octx->src[1]->nb[3]);

    struct set_rows_context srctx;
    srctx.octx = octx;
    srctx.kparams = kparams;
+    srctx.task_start = task_start;
+    srctx.tasks = tasks;
+    srctx.tasks_per_thread = fastdiv(tasks + n_threads - 1, &octx->n_threads_div);

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

    work_queue_func_t q_func = NULL;
@@ -216,15 +246,15 @@ 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",
+    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);
+         srctx.vtcm_layout.src0_bytes_per_thread * n_threads,
+         srctx.vtcm_layout.dst_bytes_per_thread  * n_threads,
+         n_threads);

-    work_queue_run(octx->ctx->work_queue, q_func, &srctx, kparams->n_threads);
+    work_queue_run(octx->ctx->work_queue, q_func, &srctx, n_threads);

    return HTP_STATUS_OK;
 }
@@ -14,9 +14,11 @@

 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+#include "hex-common.h"
+#include "hex-profile.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
+#include "htp-tensor.h"

 #define htp_softmax_preamble3                     \
    const uint32_t ne00 = src0->ne[0];            \
@@ -69,6 +71,8 @@ struct htp_softmax_context {
    struct fastdiv_values fastdiv_ne13; // For mask broadcasting

    uint32_t src0_nrows_per_thread;
+    uint32_t row_start;
+    uint32_t nrows;
 };

 static void apply_mask(float * restrict wp0,
@@ -223,19 +227,17 @@ static void softmax_job_f32(unsigned int nth, unsigned int ith, void * data) {

    htp_softmax_preamble3;

-    const uint32_t src0_nrows            = ne01 * ne02 * ne03;  // src0 rows
+    const uint32_t src0_nrows            = smctx->nrows;
    const uint32_t src0_nrows_per_thread = smctx->src0_nrows_per_thread;

-    const uint32_t src0_start_row = src0_nrows_per_thread * ith;
-    const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
+    const uint32_t src0_start_row = smctx->row_start + src0_nrows_per_thread * ith;
+    const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, smctx->row_start + src0_nrows);

    // no work for this thread
    if (src0_start_row >= src0_end_row) {
        return;
    }

-    uint64_t qt = HAP_perf_get_qtimer_count();
-
    int is_aligned = 1;
    int opt_path   = 0;

@@ -262,6 +264,9 @@ static void softmax_job_f32(unsigned int nth, unsigned int ith, void * data) {
    uint32_t prev_i2 = (uint32_t)-1;
    float slope = 1.0f;

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, src0_start_row);
+
    for (uint32_t r = src0_start_row; r < src0_end_row; ++r) {
        uint32_t i1 = fastmodulo(r, ne01, &smctx->fastdiv_ne01);
        uint32_t r_div_ne01 = fastdiv(r, &smctx->fastdiv_ne01);
@@ -323,10 +328,11 @@ static void softmax_job_f32(unsigned int nth, unsigned int ith, void * data) {
        }
    }

-    qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt);
-    FARF(HIGH, "softmax-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u : opt %u f16 %u usec %u\n", ith, nth,
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, src0_start_row);
+
+    FARF(HIGH, "softmax-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u : opt %u f16 %u\n", ith, nth,
         ne00, ne01, ne02, ne03, src0_start_row, src0_end_row, ne10, ne11, ne12, ne13,
-         ne0, ne1, ne2, ne3, opt_path, smctx->use_f16, (unsigned) qt);
+         ne0, ne1, ne2, ne3, opt_path, smctx->use_f16);
 }

 static int execute_op_softmax_f32(struct htp_ops_context * octx) {
@@ -342,13 +348,32 @@ static int execute_op_softmax_f32(struct htp_ops_context * octx) {
    init_softmax_ctx(&smctx, octx);

    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
-    const uint32_t n_threads  = MIN(octx->n_threads, src0_nrows);
+    const size_t elem_size = sizeof(float);
+    const size_t dst_row_size = dst->nb[1];

-    smctx.src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads;
+    uint32_t row_start = 0;
+    uint32_t nrows     = src0_nrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, (uint32_t) elem_size, (uint32_t) dst_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
+    smctx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    smctx.row_start             = row_start;
+    smctx.nrows                 = nrows;

    const size_t src0_row_size = src0->nb[1];
    const size_t src1_row_size = src0_row_size;
-    const size_t dst_row_size  = dst->nb[1];

    // VTCM scratchpads for all tensors
    // 4 rows per thread, padded to HVX vector size
@@ -383,9 +408,7 @@ static int execute_op_softmax_f32(struct htp_ops_context * octx) {
    octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size; octx->src1_spad.src = NULL;
    octx->dst_spad.data  = octx->src1_spad.data + octx->src1_spad.size; octx->dst_spad.src  = NULL;

-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) return err;
-
-    worker_pool_run_func(octx->ctx->worker_pool, softmax_job_f32, &smctx, n_threads);
+    work_queue_run(octx->ctx->work_queue, softmax_job_f32, &smctx, n_threads);

    return err;
 }
@@ -1,13 +1,16 @@
 #pragma clang diagnostic ignored "-Wunused-but-set-variable"

 #include <HAP_farf.h>
-#include <HAP_perf.h>
 #include <string.h>

+#include "hex-common.h"
+#include "hex-profile.h"
+
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 #include "hvx-types.h"
 #include "hvx-utils.h"

@@ -15,6 +18,7 @@ struct htp_solve_tri_context {
    struct htp_ops_context * octx;
    uint32_t                 jobs_per_thread;
    uint32_t                 total_jobs;
+    uint32_t                 job_start;
    uint32_t                 k_chunks;
    uint32_t                 col_block;
 };
@@ -89,11 +93,11 @@ static void solve_tri_batch_thread_f32(unsigned int nth, unsigned int ith, void
    const uint32_t col_block = VLEN_FP32;
    const uint32_t k_full    = (k / col_block) * col_block;

-    const uint32_t start_batch = sctx->jobs_per_thread * ith;
-    const uint32_t end_batch   = MIN(start_batch + sctx->jobs_per_thread, sctx->total_jobs);
+    const uint32_t start_batch = sctx->job_start + sctx->jobs_per_thread * ith;
+    const uint32_t end_batch   = MIN(start_batch + sctx->jobs_per_thread, sctx->job_start + sctx->total_jobs);

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_batch);

    for (uint32_t batch = start_batch; batch < end_batch; ++batch) {
        const uint32_t i03 = batch / ne02;
@@ -127,11 +131,10 @@ static void solve_tri_batch_thread_f32(unsigned int nth, unsigned int ith, void
        }
    }

-    t2 = HAP_perf_get_qtimer_count();
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) end_batch);

-    FARF(HIGH, "solve-tri-batch %d/%d: A=(%ux%u) B=(%ux%u) batch %u:%u usec %u\n",
-         ith, nth, n, n, k, n, start_batch, end_batch,
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+    FARF(HIGH, "solve-tri-batch %d/%d: A=(%ux%u) B=(%ux%u) batch %u:%u\n",
+         ith, nth, n, n, k, n, start_batch, end_batch);
 }

 // Chunk-level thread: each job is one (batch, col_chunk) pair.
@@ -148,11 +151,11 @@ static void solve_tri_chunk_thread_f32(unsigned int nth, unsigned int ith, void

    const uint32_t ne02 = src0->ne[2];

-    const uint32_t start_job = sctx->jobs_per_thread * ith;
-    const uint32_t end_job   = MIN(start_job + sctx->jobs_per_thread, sctx->total_jobs);
+    const uint32_t start_job = sctx->job_start + sctx->jobs_per_thread * ith;
+    const uint32_t end_job   = MIN(start_job + sctx->jobs_per_thread, sctx->job_start + sctx->total_jobs);

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_job);

    for (uint32_t job = start_job; job < end_job; ++job) {
        const uint32_t batch = job / sctx->k_chunks;
@@ -161,16 +164,14 @@ static void solve_tri_chunk_thread_f32(unsigned int nth, unsigned int ith, void
        const uint32_t i03 = batch / ne02;
        const uint32_t i02 = batch - i03 * ne02;

-        const uint32_t col0 = chunk * sctx->col_block;
-        const uint32_t coln = MIN(sctx->col_block, k - col0);
-
        const float * A_batch =
            (const float *) ((const uint8_t *) (uintptr_t) src0->data + i02 * src0->nb[2] + i03 * src0->nb[3]);
        const float * B_batch =
            (const float *) ((const uint8_t *) (uintptr_t) src1->data + i02 * src1->nb[2] + i03 * src1->nb[3]);
        float * X_batch = (float *) ((uint8_t *) (uintptr_t) dst->data + i02 * dst->nb[2] + i03 * dst->nb[3]);

-        const bool use_hvx = (coln >= 8);
+        const uint32_t col0 = chunk * sctx->col_block;
+        const uint32_t coln = MIN(sctx->col_block, k - col0);

        for (uint32_t row = 0; row < n; ++row) {
            const float diag     = A_batch[row * n + row];
@@ -179,7 +180,7 @@ static void solve_tri_chunk_thread_f32(unsigned int nth, unsigned int ith, void
            const float * A_row = A_batch + row * n;
            const float * B_row = B_batch + row * k;

-            if (use_hvx) {
+            if (coln >= 8) {
                solve_tri_row_hvx(A_row, B_row, X_batch, row, k, col0, coln, inv_diag);
            } else {
                solve_tri_row_scalar(A_row, B_row, X_batch, row, k, col0, coln, inv_diag);
@@ -187,11 +188,10 @@ static void solve_tri_chunk_thread_f32(unsigned int nth, unsigned int ith, void
        }
    }

-    t2 = HAP_perf_get_qtimer_count();
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) end_job);

-    FARF(HIGH, "solve-tri-chunk %d/%d: A=(%ux%u) B=(%ux%u) job %u:%u usec %u\n",
-         ith, nth, n, n, k, n, start_job, end_job,
-         (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+    FARF(HIGH, "solve-tri-chunk %d/%d: A=(%ux%u) B=(%ux%u) jobs %u:%u\n",
+         ith, nth, n, n, k, n, start_job, end_job);
 }

 int op_solve_tri(struct htp_ops_context * octx) {
@@ -235,32 +235,64 @@ int op_solve_tri(struct htp_ops_context * octx) {
         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], batched);

    if (batched) {
+        uint32_t job_start = 0;
+        uint32_t njobs     = total_batches;
+
+        if (octx->ctx->mdev.count > 1) {
+            const uint32_t batch_size = dst->nb[2];
+            const uint32_t batches_per_chunk = (batch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(batch_size, HEX_L2_LINE_SIZE)) : 1;
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, htp_tensor_mdev_data_aligned(dst) ? batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            job_start = range.start;
+            njobs     = range.count;
+        }
+
+        if (njobs == 0) {
+            return HTP_STATUS_OK;
+        }
+
        // Batch-level parallelism
-        const uint32_t n_threads = MIN((uint32_t) octx->n_threads, total_batches);
+        const uint32_t n_threads = octx->n_threads;

        struct htp_solve_tri_context sctx = {
            .octx            = octx,
-            .jobs_per_thread = (total_batches + n_threads - 1) / n_threads,
-            .total_jobs      = total_batches,
+            .jobs_per_thread = fastdiv(njobs + n_threads - 1, &octx->n_threads_div),
+            .total_jobs      = njobs,
+            .job_start       = job_start,
            .k_chunks        = k_chunks,
            .col_block       = col_block,
        };

-        worker_pool_run_func(octx->ctx->worker_pool, solve_tri_batch_thread_f32, &sctx, n_threads);
+        work_queue_run(octx->ctx->work_queue, solve_tri_batch_thread_f32, &sctx, n_threads);
    } else {
        // Chunk-level parallelism
        const uint32_t total_jobs = total_batches * k_chunks;
-        const uint32_t n_threads  = MIN((uint32_t) octx->n_threads, MAX(total_jobs, 1));
+
+        uint32_t job_start = 0;
+        uint32_t njobs     = total_jobs;
+
+        if (octx->ctx->mdev.count > 1) {
+            const bool can_split = htp_tensor_mdev_data_aligned(dst) && ((dst->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0);
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_jobs, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            job_start = range.start;
+            njobs     = range.count;
+        }
+
+        if (njobs == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        const uint32_t n_threads = octx->n_threads;

        struct htp_solve_tri_context sctx = {
            .octx            = octx,
-            .jobs_per_thread = (total_jobs + n_threads - 1) / n_threads,
-            .total_jobs      = total_jobs,
+            .jobs_per_thread = fastdiv(njobs + n_threads - 1, &octx->n_threads_div),
+            .total_jobs      = njobs,
+            .job_start       = job_start,
            .k_chunks        = k_chunks,
            .col_block       = col_block,
        };

-        worker_pool_run_func(octx->ctx->worker_pool, solve_tri_chunk_thread_f32, &sctx, n_threads);
+        work_queue_run(octx->ctx->work_queue, solve_tri_chunk_thread_f32, &sctx, n_threads);
    }

    return HTP_STATUS_OK;
@@ -4,7 +4,6 @@

 #include <HAP_farf.h>
 #include <HAP_mem.h>
-#include <HAP_perf.h>
 #include <HAP_ps.h>
 #include <hexagon_protos.h>
 #include <hexagon_types.h>
@@ -16,8 +15,9 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "hex-dma.h"
+#include "hex-profile.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
+#include "htp-tensor.h"
 #include "hvx-utils.h"

 #define htp_ssm_conv_tensors_preamble                           \
@@ -63,6 +63,8 @@ struct htp_ssm_conv_context {
    uint32_t nrows_per_thread;
    uint32_t d_inner_tile;
    uint64_t t_start;
+    uint32_t row_start;
+    uint32_t nrows;
 };

 #define htp_ssm_conv_preamble                                                   \
@@ -75,9 +77,6 @@ struct htp_ssm_conv_context {
 static void ssm_conv_thread_f32_f32(unsigned int nth, unsigned int ith, void *data) {
    htp_ssm_conv_preamble;

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
    const uint32_t d_conv  = src1->ne[0];
    const uint32_t d_inner = src0->ne[1];
    const uint32_t n_t     = dst->ne[1];
@@ -95,14 +94,17 @@ static void ssm_conv_thread_f32_f32(unsigned int nth, unsigned int ith, void *da

    // Calculate row range for this thread
    const uint32_t d_inner_per_thread = scctx->nrows_per_thread;
-    const uint32_t d_inner_start = d_inner_per_thread * ith;
-    const uint32_t d_inner_end   = MIN(d_inner_start + d_inner_per_thread, d_inner);
+    const uint32_t d_inner_start = scctx->row_start + d_inner_per_thread * ith;
+    const uint32_t d_inner_end   = MIN(d_inner_start + d_inner_per_thread, scctx->row_start + scctx->nrows);

    // No work for this thread
    if (d_inner_start >= d_inner_end) {
        return;
    }

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) d_inner_start);
+
    for (uint32_t i3 = 0; i3 < n_s; ++i3) {
        for (uint32_t i2 = 0; i2 < n_t; ++i2) {
            for (uint32_t i1 = d_inner_start; i1 < d_inner_end; ++i1) {
@@ -121,12 +123,12 @@ static void ssm_conv_thread_f32_f32(unsigned int nth, unsigned int ith, void *da
        }
    }

-    t2 = HAP_perf_get_qtimer_count();
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) d_inner_end);

-    FARF(HIGH, "ssm-conv-f32 %d/%d: %ux%ux%ux%u (%u:%u) * %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n",
+    FARF(HIGH, "ssm-conv-f32 %d/%d: %ux%ux%ux%u (%u:%u) * %ux%ux%ux%u -> %ux%ux%ux%u\n",
         ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], d_inner_start, d_inner_end,
         src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0], dst->ne[1],
-         dst->ne[2], dst->ne[3], (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         dst->ne[2], dst->ne[3]);
 }


@@ -257,9 +259,6 @@ static inline void transpose_src0_block(const float * src0_block,
 static void ssm_conv_thread_f32_f32_hvx(unsigned int nth, unsigned int ith, void *data) {
    htp_ssm_conv_preamble;

-    uint64_t t1, t2;
-    t1 = HAP_perf_get_qtimer_count();
-
    const uint32_t d_conv  = src1->ne[0];
    const uint32_t d_inner = src0->ne[1];
    const uint32_t n_t     = dst->ne[1];
@@ -273,13 +272,16 @@ static void ssm_conv_thread_f32_f32_hvx(unsigned int nth, unsigned int ith, void
    const uint32_t dst_stride_seq    = dst->nb[2]  / sizeof(float);

    const uint32_t dr  = scctx->nrows_per_thread;
-    const uint32_t ir0 = dr * ith;
-    const uint32_t ir1 = MIN(ir0 + dr, d_inner);
+    const uint32_t ir0 = scctx->row_start + dr * ith;
+    const uint32_t ir1 = MIN(ir0 + dr, scctx->row_start + scctx->nrows);

    if (ir0 >= ir1) {
        return;
    }

+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
+
    const uint32_t d_inner_per_thread = ir1 - ir0;
    const uint32_t d_inner_stride     = scctx->nrows_per_thread;
    const uint32_t d_inner_tile       = scctx->d_inner_tile;
@@ -319,97 +321,118 @@ static void ssm_conv_thread_f32_f32_hvx(unsigned int nth, unsigned int ith, void
                        HVX_Vector w = *(const HVX_Vector *) (src1_T + j * d_inner_stride + tile_off + cb);
                        acc          = Q6_Vqf32_vadd_Vqf32Vqf32(acc, Q6_Vqf32_vmpy_VsfVsf(x, w));
                    }
-                    HVX_Vector res = Q6_Vsf_equals_Vqf32(acc);

-                    float * dst_ptr = dst_data + i3 * dst_stride_seq + t * dst_stride_token + (ir0 + tile_off + cb);
+                    HVX_Vector y = Q6_Vsf_equals_Vqf32(acc);
+
+                    float * dst_ptr = dst_data + (ir0 + tile_off + cb) + t * dst_stride_token + i3 * dst_stride_seq;
                    if (cb_n == C_TILE) {
-                        *(HVX_UVector *) dst_ptr = res;
+                        *(HVX_UVector *) dst_ptr = y;
                    } else {
-                        hvx_vec_store_u(dst_ptr, cb_n * sizeof(float), res);
+                        hvx_vec_store_u(dst_ptr, cb_n * sizeof(float), y);
                    }
                }
            }
        }
    }

-    t2 = HAP_perf_get_qtimer_count();
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1);

-    FARF(HIGH, "ssm-conv-f32-hvx %d/%d: %ux%ux%ux%u (%u:%u) tile=%u * %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n",
-         ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1, d_inner_tile,
+    FARF(HIGH, "ssm-conv-f32-hvx %d/%d: %ux%ux%ux%u (%u:%u) * %ux%ux%ux%u -> %ux%ux%ux%u\n",
+         ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1,
         src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0], dst->ne[1],
-         dst->ne[2], dst->ne[3], (unsigned) HAP_perf_qtimer_count_to_us(t2 - t1));
+         dst->ne[2], dst->ne[3]);
 }

 int op_ssm_conv_f32(struct htp_ops_context * octx) {
-    htp_ssm_conv_tensors_preamble;
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * src1 = octx->src[1];
+    const struct htp_tensor * dst  = octx->dst;

    if (src0->type != HTP_TYPE_F32 || src1->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_F32) {
-        FARF(ERROR, "ssm_conv: only (F32 x F32 -> F32) OPs supported");
        return HTP_STATUS_NO_SUPPORT;
    }

-    struct htp_ssm_conv_context scctx = { 0 };
-    scctx.octx = octx;
-
    const uint32_t d_conv  = src1->ne[0];
    const uint32_t d_inner = src0->ne[1];
    const uint32_t n_t     = dst->ne[1];  // tokens per sequence
    const uint32_t n_s     = dst->ne[2];  // number of sequences in the batch

-    const uint32_t n_threads = MIN(octx->n_threads, d_inner);
+    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
+        return HTP_STATUS_OK;
+    }

-    if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
-        uint32_t use_hvx = 0;
-        if (d_inner >= VLEN_FP32 && n_t >= VLEN_FP32) {
-            use_hvx = 1;
-        }
+    uint32_t row_start = 0;
+    uint32_t nrows     = d_inner;

-        scctx.nrows_per_thread = hex_round_up((d_inner + n_threads - 1) / n_threads, VLEN_FP32);
+    if (octx->ctx->mdev.count > 1) {
+        const uint32_t elems_per_chunk = VLEN_FP32;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(d_inner, htp_tensor_mdev_data_aligned(dst) ? elems_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }

-        const uint32_t d_inner_per_thread = scctx.nrows_per_thread;
-        const uint32_t ncs                = src0->ne[0];
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }

-        const uint32_t src1_T_size = hex_round_up(d_conv * d_inner_per_thread * sizeof(float), 256);
-        const uint32_t src0_T_max = HTP_SSM_CONV_VTCM_BUDGET > src1_T_size ? HTP_SSM_CONV_VTCM_BUDGET - src1_T_size : 0;
+    const uint32_t n_threads = octx->n_threads;

-        uint32_t d_inner_tile = (src0_T_max / sizeof(float)) / ncs;
-        d_inner_tile -= (d_inner_tile % VLEN_FP32);
-        if (d_inner_tile == 0) {
-            FARF(HIGH, "ssm_conv-f32: inner tile rounds to 0 (ncs=%u), falling back to scalar\n", ncs);
+    struct htp_ssm_conv_context scctx = { 0 };
+    scctx.octx      = octx;
+    scctx.row_start = row_start;
+    scctx.nrows     = nrows;
+
+    uint32_t use_hvx = 0;
+    if (nrows >= VLEN_FP32 && n_t >= VLEN_FP32) {
+        use_hvx = 1;
+    }
+
+    const uint32_t raw_rpt = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    scctx.nrows_per_thread = hex_round_up(raw_rpt, VLEN_FP32);
+
+    const uint32_t d_inner_per_thread = scctx.nrows_per_thread;
+    const uint32_t ncs                = src0->ne[0];
+
+    const uint32_t src1_T_size = hex_round_up(d_conv * d_inner_per_thread * sizeof(float), 256);
+    const uint32_t src0_T_max = HTP_SSM_CONV_VTCM_BUDGET > src1_T_size ? HTP_SSM_CONV_VTCM_BUDGET - src1_T_size : 0;
+
+    uint32_t d_inner_tile = (src0_T_max / sizeof(float)) / ncs;
+    d_inner_tile -= (d_inner_tile % VLEN_FP32);
+    if (d_inner_tile == 0) {
+        FARF(HIGH, "ssm_conv-f32: inner tile rounds to 0 (ncs=%u), falling back to scalar\n", ncs);
+        use_hvx = 0;
+    } else {
+        scctx.d_inner_tile = d_inner_tile;
+
+        octx->src0_spad.size_per_thread = hex_round_up(d_inner_tile * ncs * sizeof(float), 256);
+        octx->src1_spad.size_per_thread = src1_T_size;
+        octx->dst_spad.size_per_thread  = 0;
+
+        octx->src0_spad.size = octx->src0_spad.size_per_thread * n_threads;
+        octx->src1_spad.size = octx->src1_spad.size_per_thread * n_threads;
+        octx->dst_spad.size  = 0;
+
+        octx->src0_spad.data = octx->ctx->vtcm_base;
+        octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size;
+        octx->src0_spad.src  = NULL;
+        octx->src1_spad.src  = NULL;
+
+        const size_t total_spad = octx->src0_spad.size + octx->src1_spad.size;
+        if (total_spad > octx->ctx->vtcm_size) {
+            FARF(HIGH, "ssm_conv-f32: scratchpad %zu exceeds VTCM %zu, falling back to scalar\n",
+                 total_spad, octx->ctx->vtcm_size);
            use_hvx = 0;
-        } else {
-            scctx.d_inner_tile = d_inner_tile;
-
-            octx->src0_spad.size_per_thread = hex_round_up(d_inner_tile * ncs * sizeof(float), 256);
-            octx->src1_spad.size_per_thread = src1_T_size;
-            octx->dst_spad.size_per_thread  = 0;
-
-            octx->src0_spad.size = octx->src0_spad.size_per_thread * n_threads;
-            octx->src1_spad.size = octx->src1_spad.size_per_thread * n_threads;
-            octx->dst_spad.size  = 0;
-
-            octx->src0_spad.data = octx->ctx->vtcm_base;
-            octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size;
-            octx->src0_spad.src  = NULL;
-            octx->src1_spad.src  = NULL;
-
-            const size_t total_spad = octx->src0_spad.size + octx->src1_spad.size;
-            if (total_spad > octx->ctx->vtcm_size) {
-                FARF(HIGH, "ssm_conv-f32: scratchpad %zu exceeds VTCM %zu, falling back to scalar\n",
-                     total_spad, octx->ctx->vtcm_size);
-                use_hvx = 0;
-            }
        }
+    }

-        FARF(HIGH, "ssm-conv-f32: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : use_hvx %d\n", src0->ne[0],
-             src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0],
-             dst->ne[1], dst->ne[2], dst->ne[3], use_hvx);
+    FARF(HIGH, "ssm-conv-f32: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : use_hvx %d\n", src0->ne[0],
+         src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0],
+         dst->ne[1], dst->ne[2], dst->ne[3], use_hvx);

-        if (use_hvx) {
-            worker_pool_run_func(octx->ctx->worker_pool, ssm_conv_thread_f32_f32_hvx, &scctx, n_threads);
-        } else {
-            worker_pool_run_func(octx->ctx->worker_pool, ssm_conv_thread_f32_f32, &scctx, n_threads);
-        }
+    if (use_hvx) {
+        work_queue_run(octx->ctx->work_queue, ssm_conv_thread_f32_f32_hvx, &scctx, n_threads);
+    } else {
+        work_queue_run(octx->ctx->work_queue, ssm_conv_thread_f32_f32, &scctx, n_threads);
    }

    return HTP_STATUS_OK;
@@ -13,35 +13,38 @@

 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+#include "hex-common.h"
+#include "hex-profile.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
+#include "htp-tensor.h"

 #define sum_rows_preamble                         \
    const struct htp_tensor *src0 = octx->src[0]; \
    const struct htp_tensor *dst  = octx->dst;    \
                                                  \
-    const uint32_t ne00 = src0->ne[0];     \
-    const uint32_t ne01 = src0->ne[1];     \
-    const uint32_t ne02 = src0->ne[2];     \
-    const uint32_t ne03 = src0->ne[3];     \
-                                           \
-    const uint32_t nb00 = src0->nb[0];     \
-    const uint32_t nb01 = src0->nb[1];     \
-    const uint32_t nb02 = src0->nb[2];     \
-    const uint32_t nb03 = src0->nb[3];     \
-                                           \
-    const uint32_t  ne0 = dst->ne[0];      \
-    const uint32_t  ne1 = dst->ne[1];      \
-    const uint32_t  ne2 = dst->ne[2];      \
-    const uint32_t  ne3 = dst->ne[3];      \
-                                           \
-    const uint32_t  nb0 = dst->nb[0];      \
-    const uint32_t  nb1 = dst->nb[1];      \
-    const uint32_t  nb2 = dst->nb[2];      \
-    const uint32_t  nb3 = dst->nb[3];      \
+    const uint32_t ne00 = src0->ne[0];            \
+    const uint32_t ne01 = src0->ne[1];            \
+    const uint32_t ne02 = src0->ne[2];            \
+    const uint32_t ne03 = src0->ne[3];            \
+                                                  \
+    const uint32_t nb00 = src0->nb[0];            \
+    const uint32_t nb01 = src0->nb[1];            \
+    const uint32_t nb02 = src0->nb[2];            \
+    const uint32_t nb03 = src0->nb[3];            \
+                                                  \
+    const uint32_t  ne0 = dst->ne[0];             \
+    const uint32_t  ne1 = dst->ne[1];             \
+    const uint32_t  ne2 = dst->ne[2];             \
+    const uint32_t  ne3 = dst->ne[3];             \
+                                                  \
+    const uint32_t  nb0 = dst->nb[0];             \
+    const uint32_t  nb1 = dst->nb[1];             \
+    const uint32_t  nb2 = dst->nb[2];             \
+    const uint32_t  nb3 = dst->nb[3];             \

 struct sum_rows_context {
+    struct htp_ops_context * octx;
    const uint8_t * src_data;
    uint8_t       * dst_data;
    uint32_t        ne00;
@@ -76,6 +79,9 @@ static void sum_rows_thread_f32(unsigned int nth, unsigned int ith, void *data)
    // Calculate actual number of rows for this thread
    const uint32_t n_rows = end_row - start_row;

+    struct htp_thread_trace * tr = &smctx->octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_row);
+
    for (uint32_t ir = 0; ir < n_rows; ir++) {
        const float * restrict src_local = src_th + (ir * (src_stride / sizeof(float)));

@@ -89,6 +95,8 @@ static void sum_rows_thread_f32(unsigned int nth, unsigned int ith, void *data)
            dst_th[ir] = hvx_reduce_sum_f32((const uint8_t *) src_local, ne00);
        }
    }
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_row);
 }

 int op_sum_rows(struct htp_ops_context * octx) {
@@ -102,9 +110,26 @@ int op_sum_rows(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

-    const uint32_t src0_nrows = ne01 * ne02 * ne03;
-    const uint32_t n_threads = MIN(octx->n_threads, src0_nrows);
-    const uint32_t rows_per_thread = (src0_nrows + n_threads - 1) / n_threads;
+    const uint32_t src0_nrows      = ne01 * ne02 * ne03;
+    const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
+
+    uint32_t row_start = 0;
+    uint32_t nrows     = src0_nrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_data_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+    const uint32_t rows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);

    bool opt_path = false;
    if ((0 == hex_is_aligned((void *) src0->data, VLEN)) && !(nb01 & (VLEN - 1))) {
@@ -112,17 +137,18 @@ int op_sum_rows(struct htp_ops_context * octx) {
    }

    struct sum_rows_context smctx = {
-        .src_data        = (const uint8_t *) src0->data,
-        .dst_data        = (uint8_t *) dst->data,
+        .octx            = octx,
+        .src_data        = (const uint8_t *) src0->data + row_start * nb01,
+        .dst_data        = (uint8_t *) dst->data + row_start * nb1,
        .ne00            = ne00,
        .src_stride      = nb01,
        .dst_stride      = nb1,
        .rows_per_thread = rows_per_thread,
-        .total_rows      = src0_nrows,
+        .total_rows      = nrows,
        .opt_path        = opt_path,
    };

-    worker_pool_run_func(octx->ctx->worker_pool, sum_rows_thread_f32, &smctx, n_threads);
+    work_queue_run(octx->ctx->work_queue, sum_rows_thread_f32, &smctx, n_threads);

    return HTP_STATUS_OK;
 }
@@ -46,6 +46,7 @@ struct htp_unary_context {
    uint32_t                  block;
    uint32_t                  src0_nrows;
    uint32_t                  src0_nrows_per_thread;
+    uint32_t                  row_start;
    uint32_t                  nc;
    uint32_t                  col_tile;             // tiled mode
    bool                      broadcast_weight;
@@ -496,7 +497,7 @@ static void tri_f32(const float * restrict src,
        }
        if (boundary > ne0) boundary = ne0;

-        // Full HVX vectors — each starts at a 128-byte aligned offset
+        // Full HVX vectors - each starts at a 128-byte aligned offset
        for (uint32_t i = 0; i < nvec; i++) {
            const uint32_t vec_start = i * VLEN_FP32;
            const uint32_t vec_end   = vec_start + VLEN_FP32;
@@ -563,7 +564,7 @@ static void softplus_f32(const float * restrict src,

        for (uint32_t i = 0; i < ne0; i++) {
            float x = src_f[i];
-            // For x > 20: softplus(x) ≈ x (avoids exp overflow)
+            // For x > 20: softplus(x) ~ x (avoids exp overflow)
            dst_f[i] = (x > 20.0f) ? x : logf(1.0f + expf(x));
        }
    }
@@ -661,8 +662,8 @@ static void unary_task_##SUFFIX##_##NAME(unsigned int nth, unsigned int ith, voi
    const size_t dst_row_size_aligned  = uctx->dst_row_size_aligned;                                                \
                                                                                                                    \
    const uint32_t src0_nrows = uctx->src0_nrows;                                                                   \
-    const uint32_t src0_start_row = src0_nrows_per_thread * ith;                                                    \
-    const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);                        \
+    const uint32_t src0_start_row = uctx->row_start + src0_nrows_per_thread * ith;                                  \
+    const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, uctx->row_start + src0_nrows);      \
                                                                                                                    \
    if (src0_start_row >= src0_end_row) {                                                                           \
        return;                                                                                                     \
@@ -833,124 +834,126 @@ DEFINE_UNARY_TASK_IMPL(unary_abs, _Float16, f16, false, false, abs_f16(src0_vtcm
 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) {                          \
-    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];                                                                   \
-    const struct htp_tensor * dst = octx->dst;                                                                      \
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                          \
-                                                                                                                    \
-    htp_unary_preamble;                                                                                             \
-                                                                                                                    \
-    int32_t *      op_params = octx->op_params;                                                                     \
-    const uint32_t col_tile  = uctx->col_tile;                                                                      \
-                                                                                                                    \
-    const uint32_t src0_nrows     = uctx->src0_nrows;                                                               \
-    const uint32_t src0_start_row = uctx->src0_nrows_per_thread * ith;                                              \
-    const uint32_t src0_end_row   = MIN(src0_start_row + uctx->src0_nrows_per_thread, src0_nrows);                  \
-                                                                                                                    \
-    if (src0_start_row >= src0_end_row) {                                                                           \
-        return;                                                                                                     \
-    }                                                                                                               \
-                                                                                                                    \
-    const uint8_t * restrict data_src = uctx->data_src0;                                                            \
-    uint8_t * restrict       data_dst = uctx->data_dst;                                                             \
-                                                                                                                    \
-    uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread);                           \
-    uint8_t * dst_vtcm_data  = uctx->vtcm_dst + (ith * uctx->vtcm_dst_size_per_thread);                             \
-                                                                                                                    \
-    const size_t src0_half = uctx->src0_vtcm_half_size;                                                             \
-    const size_t dst_half  = uctx->dst_vtcm_half_size;                                                              \
-                                                                                                                    \
-    dma_queue * dmaq = octx->ctx->dma[ith];                                                                         \
-                                                                                                                    \
-    const struct fastdiv_values * div_ne01  = &uctx->kparams->div_ne01;                                             \
-    const struct fastdiv_values * div_ne02  = &uctx->kparams->div_ne02;                                             \
-    const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012;                                            \
-    const struct fastdiv_values * div_tpr   = &uctx->kparams->div_tpr;                                              \
-                                                                                                                    \
-    const uint32_t tiles_per_row = (ne0 + col_tile - 1) / col_tile;                                                 \
-    const int32_t  tri_ttype     = (IS_TRI) ? op_params[0] : 0;                                                     \
-                                                                                                                    \
-    const bool src0_contig = (nb02 == (size_t)ne01 * nb01) &&                                                       \
-                             (nb03 == (size_t)ne02 * nb02);                                                         \
-    const bool dst_contig  = (nb2  == (size_t)ne1  * nb1)  &&                                                       \
-                             (nb3  == (size_t)ne2  * nb2);                                                          \
-                                                                                                                    \
-    const uint32_t total_tiles = (src0_end_row - src0_start_row) * tiles_per_row;                                   \
-                                                                                                                    \
-    for (uint32_t t = 0, vtcm_idx = 0; t < total_tiles && vtcm_idx < 2; t++, vtcm_idx++) {                          \
-        const uint32_t row  = src0_start_row + t / tiles_per_row;                                                   \
-        const uint32_t col  = (t % tiles_per_row) * col_tile;                                                       \
-        const uint32_t tw   = MIN(col_tile, ne0 - col);                                                             \
-        const size_t   tb   = (size_t) tw * sizeof(float);                                                          \
-        const size_t   soff = (src0_contig ? (row * nb01) :                                                         \
-                               unary_row_offset(row, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03)) +\
-                               (size_t) col * sizeof(float);                                                        \
-                                                                                                                    \
-        dma_queue_push(dmaq, dma_make_ptr(data_dst, dst_vtcm_data + (vtcm_idx * dst_half)), 0, 0, 0, 0);            \
-        dma_queue_push(dmaq, dma_make_ptr(src0_vtcm_data + (vtcm_idx * src0_half), data_src + soff), tb, tb, tb, 1);\
-    }                                                                                                               \
-                                                                                                                    \
-    uint32_t row = src0_start_row;                                                                                  \
-    uint32_t col = 0;                                                                                               \
-    uint32_t tile_in_row = 0;                                                                                       \
-    uint32_t i01 = fastmodulo(row, ne01, div_ne01);                                                                 \
-                                                                                                                    \
-    uint32_t prow = src0_start_row + fastdiv(2, div_tpr);                                                           \
-    uint32_t pcol = fastmodulo(2, tiles_per_row, div_tpr) * col_tile;                                               \
-    uint32_t ptile_in_row = fastmodulo(2, tiles_per_row, div_tpr);                                                  \
-                                                                                                                    \
-    for (uint32_t t = 0; t < total_tiles; t++) {                                                                    \
-        uint8_t * dst_vtcm = (uint8_t *) dma_queue_pop(dmaq).src;                                                   \
-        uint8_t * src_vtcm = (uint8_t *) dma_queue_pop(dmaq).dst;                                                   \
-                                                                                                                    \
-        const uint32_t tw  = MIN(col_tile, ne0 - col);                                                              \
-                                                                                                                    \
-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, t);                                                       \
-        CORE_TILE_EXPR;                                                                                             \
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, t);                                                        \
-                                                                                                                    \
-        const size_t doff = (dst_contig ? (row * nb1) :                                                             \
-                             unary_row_offset(row, ne1, ne2, div_ne01, div_ne02, div_ne012, nb1, nb2, nb3)) +       \
-                             (size_t) col * sizeof(float);                                                          \
-        const size_t tb   = (size_t) tw * sizeof(float);                                                            \
-        dma_queue_push(dmaq, dma_make_ptr(data_dst + doff, dst_vtcm), tb, tb, tb, 1);                               \
-                                                                                                                    \
-        const uint32_t pt = t + 2;                                                                                  \
-        if (pt < total_tiles) {                                                                                     \
-            const uint32_t ptw  = MIN(col_tile, ne0 - pcol);                                                        \
-            const size_t   ptb  = (size_t) ptw * sizeof(float);                                                     \
-            const size_t   psoff = (src0_contig ? (prow * nb01) :                                                   \
-                                    unary_row_offset(prow, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02,   \
-                                                     nb03)) +                                                       \
-                                   (size_t) pcol * sizeof(float);                                                   \
-            dma_queue_push(dmaq, dma_make_ptr(src_vtcm, data_src + psoff), ptb, ptb, ptb, 1);                       \
-        }                                                                                                           \
-                                                                                                                    \
-        tile_in_row++;                                                                                              \
-        col += col_tile;                                                                                            \
-        if (tile_in_row == tiles_per_row) {                                                                         \
-            tile_in_row = 0;                                                                                        \
-            col = 0;                                                                                                \
-            row++;                                                                                                  \
-            i01++;                                                                                                  \
-            if (i01 == ne01) {                                                                                      \
-                i01 = 0;                                                                                            \
-            }                                                                                                       \
-        }                                                                                                           \
-                                                                                                                    \
-        ptile_in_row++;                                                                                             \
-        pcol += col_tile;                                                                                           \
-        if (ptile_in_row == tiles_per_row) {                                                                        \
-            ptile_in_row = 0;                                                                                       \
-            pcol = 0;                                                                                               \
-            prow++;                                                                                                 \
-        }                                                                                                           \
-    }                                                                                                               \
-                                                                                                                    \
-    dma_queue_flush(dmaq);                                                                                          \
+#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) {                            \
+    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];                                                                     \
+    const struct htp_tensor * dst = octx->dst;                                                                        \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                            \
+                                                                                                                      \
+    htp_unary_preamble;                                                                                               \
+                                                                                                                      \
+    uint32_t     src0_nrows_per_thread = uctx->src0_nrows_per_thread;                                                 \
+                                                                                                                      \
+    int32_t *      op_params = octx->op_params;                                                                       \
+    const uint32_t col_tile  = uctx->col_tile;                                                                        \
+                                                                                                                      \
+    const uint32_t src0_nrows = uctx->src0_nrows;                                                                     \
+    const uint32_t src0_start_row = uctx->row_start + src0_nrows_per_thread * ith;                                    \
+    const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, uctx->row_start + src0_nrows);        \
+                                                                                                                      \
+    if (src0_start_row >= src0_end_row) {                                                                             \
+        return;                                                                                                       \
+    }                                                                                                                 \
+                                                                                                                      \
+    const uint8_t * restrict data_src = uctx->data_src0;                                                              \
+    uint8_t * restrict       data_dst = uctx->data_dst;                                                               \
+                                                                                                                      \
+    uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread);                             \
+    uint8_t * dst_vtcm_data  = uctx->vtcm_dst + (ith * uctx->vtcm_dst_size_per_thread);                               \
+                                                                                                                      \
+    const size_t src0_half = uctx->src0_vtcm_half_size;                                                               \
+    const size_t dst_half  = uctx->dst_vtcm_half_size;                                                                \
+                                                                                                                      \
+    dma_queue * dmaq = octx->ctx->dma[ith];                                                                           \
+                                                                                                                      \
+    const struct fastdiv_values * div_ne01  = &uctx->kparams->div_ne01;                                               \
+    const struct fastdiv_values * div_ne02  = &uctx->kparams->div_ne02;                                               \
+    const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012;                                              \
+    const struct fastdiv_values * div_tpr   = &uctx->kparams->div_tpr;                                                \
+                                                                                                                      \
+    const uint32_t tiles_per_row = (ne0 + col_tile - 1) / col_tile;                                                   \
+    const int32_t  tri_ttype     = (IS_TRI) ? op_params[0] : 0;                                                       \
+                                                                                                                      \
+    const bool src0_contig = (nb02 == (size_t)ne01 * nb01) &&                                                         \
+                             (nb03 == (size_t)ne02 * nb02);                                                           \
+    const bool dst_contig  = (nb2  == (size_t)ne1  * nb1)  &&                                                         \
+                             (nb3  == (size_t)ne2  * nb2);                                                            \
+                                                                                                                      \
+    const uint32_t total_tiles = (src0_end_row - src0_start_row) * tiles_per_row;                                     \
+                                                                                                                      \
+    for (uint32_t t = 0, vtcm_idx = 0; t < total_tiles && vtcm_idx < 2; t++, vtcm_idx++) {                            \
+        const uint32_t row  = src0_start_row + t / tiles_per_row;                                                     \
+        const uint32_t col  = (t % tiles_per_row) * col_tile;                                                         \
+        const uint32_t tw   = MIN(col_tile, ne0 - col);                                                               \
+        const size_t   tb   = (size_t) tw * sizeof(float);                                                            \
+        const size_t   soff = (src0_contig ? (row * nb01) :                                                           \
+                               unary_row_offset(row, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03)) +  \
+                               (size_t) col * sizeof(float);                                                          \
+                                                                                                                      \
+        dma_queue_push(dmaq, dma_make_ptr(data_dst, dst_vtcm_data + (vtcm_idx * dst_half)), 0, 0, 0, 0);              \
+        dma_queue_push(dmaq, dma_make_ptr(src0_vtcm_data + (vtcm_idx * src0_half), data_src + soff), tb, tb, tb, 1);  \
+    }                                                                                                                 \
+                                                                                                                      \
+    uint32_t row = src0_start_row;                                                                                    \
+    uint32_t col = 0;                                                                                                 \
+    uint32_t tile_in_row = 0;                                                                                         \
+    uint32_t i01 = fastmodulo(row, ne01, div_ne01);                                                                   \
+                                                                                                                      \
+    uint32_t prow = src0_start_row + fastdiv(2, div_tpr);                                                             \
+    uint32_t pcol = fastmodulo(2, tiles_per_row, div_tpr) * col_tile;                                                 \
+    uint32_t ptile_in_row = fastmodulo(2, tiles_per_row, div_tpr);                                                    \
+                                                                                                                      \
+    for (uint32_t t = 0; t < total_tiles; t++) {                                                                      \
+        uint8_t * dst_vtcm = (uint8_t *) dma_queue_pop(dmaq).src;                                                     \
+        uint8_t * src_vtcm = (uint8_t *) dma_queue_pop(dmaq).dst;                                                     \
+                                                                                                                      \
+        const uint32_t tw  = MIN(col_tile, ne0 - col);                                                                \
+                                                                                                                      \
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, t);                                                         \
+        CORE_TILE_EXPR;                                                                                               \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, t);                                                          \
+                                                                                                                      \
+        const size_t doff = (dst_contig ? (row * nb1) :                                                               \
+                             unary_row_offset(row, ne1, ne2, div_ne01, div_ne02, div_ne012, nb1, nb2, nb3)) +         \
+                             (size_t) col * sizeof(float);                                                            \
+        const size_t tb   = (size_t) tw * sizeof(float);                                                              \
+        dma_queue_push(dmaq, dma_make_ptr(data_dst + doff, dst_vtcm), tb, tb, tb, 1);                                 \
+                                                                                                                      \
+        const uint32_t pt = t + 2;                                                                                    \
+        if (pt < total_tiles) {                                                                                       \
+            const uint32_t ptw  = MIN(col_tile, ne0 - pcol);                                                          \
+            const size_t   ptb  = (size_t) ptw * sizeof(float);                                                       \
+            const size_t   psoff = (src0_contig ? (prow * nb01) :                                                     \
+                                    unary_row_offset(prow, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02,     \
+                                                     nb03)) +                                                         \
+                                   (size_t) pcol * sizeof(float);                                                     \
+            dma_queue_push(dmaq, dma_make_ptr(src_vtcm, data_src + psoff), ptb, ptb, ptb, 1);                         \
+        }                                                                                                             \
+                                                                                                                      \
+        tile_in_row++;                                                                                                \
+        col += col_tile;                                                                                              \
+        if (tile_in_row == tiles_per_row) {                                                                           \
+            tile_in_row = 0;                                                                                          \
+            col = 0;                                                                                                  \
+            row++;                                                                                                    \
+            i01++;                                                                                                    \
+            if (i01 == ne01) {                                                                                        \
+                i01 = 0;                                                                                              \
+            }                                                                                                         \
+        }                                                                                                             \
+                                                                                                                      \
+        ptile_in_row++;                                                                                               \
+        pcol += col_tile;                                                                                             \
+        if (ptile_in_row == tiles_per_row) {                                                                          \
+            ptile_in_row = 0;                                                                                         \
+            pcol = 0;                                                                                                 \
+            prow++;                                                                                                   \
+        }                                                                                                             \
+    }                                                                                                                 \
+                                                                                                                      \
+    dma_queue_flush(dmaq);                                                                                            \
 }

 static inline void tile_scale_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw, const int32_t * op_params) {
@@ -1146,14 +1149,32 @@ static int execute_op_unary(struct htp_ops_context * octx) {

    const struct htp_unary_kernel_params * kparams = (const struct htp_unary_kernel_params *) octx->kernel_params;

+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_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 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;

+    uint32_t row_start = 0;
+    uint32_t nrows     = src0_nrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        uint32_t rows_per_chunk = 0;
+        htp_tensor_mdev_rows_per_chunk(dst, (uint32_t) elem_size, (uint32_t) dst_data_row_size, &rows_per_chunk);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
    const size_t src0_row_size_aligned = kparams->src0_row_size_aligned;
    const size_t dst_row_size_aligned  = kparams->dst_row_size_aligned;

@@ -1191,8 +1212,9 @@ static int execute_op_unary(struct htp_ops_context * octx) {
        struct htp_unary_context uctx = {
            .octx                  = octx,
            .kparams               = kparams,
-            .src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads,
-            .src0_nrows            = src0_nrows,
+            .src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
+            .src0_nrows            = nrows,
+            .row_start             = row_start,

            .data_src0             = (const uint8_t *)src0->data,
            .data_src1             = (octx->op == HTP_OP_RMS_NORM_MUL) ? (const uint8_t *)src1->data : NULL,
@@ -1287,7 +1309,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
        }

        if (task_func) {
-            worker_pool_run_func(octx->ctx->worker_pool, task_func, &uctx, n_threads);
+            work_queue_run(octx->ctx->work_queue, task_func, &uctx, n_threads);
        } else {
            FARF(ERROR, "execute_op_unary: task function is NULL for op %d\n", octx->op);
            err = HTP_STATUS_NO_SUPPORT;
@@ -70,17 +70,8 @@ list(APPEND GGML_SOURCES_ROCM ${SRCS})
 file(GLOB   SRCS "../ggml-cuda/template-instances/mmf*.cu")
 list(APPEND GGML_SOURCES_ROCM ${SRCS})

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

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

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

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

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

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

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

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

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

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

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

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

 #import <Foundation/Foundation.h>

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

-    bool use_fusion;
    bool use_concurrency;
    bool use_graph_optimize;

    int debug_graph;
-    int debug_fusion;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

 const struct ggml_metal_device_props * ggml_metal_device_get_props(ggml_metal_device_t dev);

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

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

    struct ggml_metal_device_props props;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    int debug_graph;
-    int debug_fusion;

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

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

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

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

    auto pipeline = ggml_metal_library_get_pipeline_gated_delta_net(lib, op);

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

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

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

    const int nsg = pipeline.nsg;

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

-    return 1;
+    return n_fuse;
 }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

 typedef struct ggml_metal_op * ggml_metal_op_t;

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

 void ggml_metal_op_free(ggml_metal_op_t ctx);

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

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

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

    return NULL;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    iq1m_scale_t scale;

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

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

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

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

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

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

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

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