cont : clean-up

cont: more review comments
address review comments
2026-11-10 21:07:41 -06:00 · 2026-11-09 11:32:25 +02:00 · 2026-11-09 17:08:27 +08:00 · 2026-11-09 16:55:27 +08:00 · 2026-11-09 13:44:00 +08:00 · 2026-11-09 13:43:45 +08:00
526 changed files with 19212 additions and 12996 deletions
@@ -1,4 +1,4 @@
-ARG UBUNTU_VERSION=24.04
+ARG UBUNTU_VERSION=26.04
 ARG BUILD_DATE=N/A
 ARG APP_VERSION=N/A
 ARG APP_REVISION=N/A
@@ -22,9 +22,7 @@ FROM docker.io/ubuntu:$UBUNTU_VERSION AS build
 ARG TARGETARCH

 RUN apt-get update && \
-    apt-get install -y gcc-14 g++-14 build-essential git cmake libssl-dev
-
-ENV CC=gcc-14 CXX=g++-14
+    apt-get install -y build-essential git cmake libssl-dev

 WORKDIR /app

@@ -84,6 +84,9 @@ RUN apt-get update \
    && find /var/cache/apt/archives /var/lib/apt/lists -not -name lock -type f -delete \
    && find /var/cache -type f -delete

+# Some hosts (e.g. GKE) only mount the driver here, and CUDA 12.8.1-12.9.1 base images don't set these paths
+ENV LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/nvidia/lib:${LD_LIBRARY_PATH}
+
 COPY --from=build /app/lib/ /app

 ### Full
@@ -82,20 +82,14 @@ COPY --from=build /app/full /app

 WORKDIR /app

-ENV PATH="/root/.venv/bin:/root/.local/bin:${PATH}"
-
-# Flag for compatibility with pip
-ARG UV_INDEX_STRATEGY="unsafe-best-match"
 RUN apt-get update \
    && apt-get install -y \
-    build-essential \
-    curl \
    git \
-    ca-certificates \
-    && curl -LsSf https://astral.sh/uv/install.sh | sh \
-    && uv python install 3.13 \
-    && uv venv --python 3.13 /root/.venv \
-    && uv pip install --python /root/.venv/bin/python -r requirements.txt \
+    python3 \
+    python3-pip \
+    python3-wheel \
+    && pip install --break-system-packages --upgrade setuptools \
+    && pip install --break-system-packages -r requirements.txt \
    && apt autoremove -y \
    && apt clean -y \
    && rm -rf /tmp/* /var/tmp/* \
@@ -100,56 +100,56 @@ runs:
          echo "CUDA_PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.1" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
          echo "CUDA_PATH_V13_1=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.1" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8

-    - name: Install Cuda Toolkit 13.4 for x64
-      if: ${{ inputs.cuda_version == '13.4' && inputs.cuda_arch == 'x64' }}
+    - name: Install Cuda Toolkit 13.5 for x64
+      if: ${{ inputs.cuda_version == '13.5' && inputs.cuda_arch == 'x64' }}
      shell: pwsh
      run: |
-          mkdir -p "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4"
+          mkdir -p "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5"
          choco install unzip -y
-          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_cudart/windows-x86_64/cuda_cudart-windows-x86_64-13.4.49-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/cuda_nvrtc/windows-x86_64/cuda_nvrtc-windows-x86_64-13.4.59-archive.zip"
-          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/libcublas/windows-x86_64/libcublas-windows-x86_64-13.7.0.27-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_nvtx/windows-x86_64/cuda_nvtx-windows-x86_64-13.4.49-archive.zip"
-          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_profiler_api/windows-x86_64/cuda_profiler_api-windows-x86_64-13.4.49-archive.zip"
-          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/visual_studio_integration/windows-x86_64/visual_studio_integration-windows-x86_64-13.4.49-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"
-          unzip '*.zip' -d "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4"
-          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_cudart-windows-x86_64-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\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\cuda_nvrtc-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\libcublas-windows-x86_64-13.7.0.27-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_nvtx-windows-x86_64-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\cuda_profiler_api-windows-x86_64-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\visual_studio_integration-windows-x86_64-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\cccl-windows-x86_64-13.3.4.2.1-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
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_crt/windows-x86_64/cuda_crt-windows-x86_64-13.5.37-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/windows-x86_64/cuda_cudart-windows-x86_64-13.5.37-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/windows-x86_64/cuda_nvcc-windows-x86_64-13.5.37-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvrtc/windows-x86_64/cuda_nvrtc-windows-x86_64-13.5.37-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/libcublas/windows-x86_64/libcublas-windows-x86_64-13.10.0.20-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/libnvvm/windows-x86_64/libnvvm-windows-x86_64-13.5.37-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvtx/windows-x86_64/cuda_nvtx-windows-x86_64-13.5.25-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_profiler_api/windows-x86_64/cuda_profiler_api-windows-x86_64-13.5.25-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/visual_studio_integration/windows-x86_64/visual_studio_integration-windows-x86_64-13.5.25-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cccl/windows-x86_64/cccl-windows-x86_64-13.3.5.0.6-archive.zip"
+          unzip '*.zip' -d "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5"
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\cuda_crt-windows-x86_64-13.5.37-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\cuda_cudart-windows-x86_64-13.5.37-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\cuda_nvcc-windows-x86_64-13.5.37-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\cuda_nvrtc-windows-x86_64-13.5.37-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\libcublas-windows-x86_64-13.10.0.20-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\libnvvm-windows-x86_64-13.5.37-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\cuda_nvtx-windows-x86_64-13.5.25-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\cuda_profiler_api-windows-x86_64-13.5.25-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\visual_studio_integration-windows-x86_64-13.5.25-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\cccl-windows-x86_64-13.3.5.0.6-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          echo "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\bin" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append
+          echo "CUDA_PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
+          echo "CUDA_PATH_V13_5=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8

-    - name: Install Cuda Toolkit 13.4 for ARM64
-      if: ${{ inputs.cuda_version == '13.4' && inputs.cuda_arch == 'arm64' }}
+    - name: Install Cuda Toolkit 13.5 for ARM64
+      if: ${{ inputs.cuda_version == '13.5' && inputs.cuda_arch == 'arm64' }}
      shell: pwsh
      run: |
-          mkdir -p "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4"
+          mkdir -p "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5"
          choco install unzip -y
-          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.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
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cccl/windows-x86_64/cccl-windows-x86_64-13.3.5.0.6-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_crt/windows-x86_64/cuda_crt-windows-x86_64-13.5.37-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/windows-x86_64/cuda_nvcc-windows-x86_64-13.5.37-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/libnvvm/windows-x86_64/libnvvm-windows-x86_64-13.5.37-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/windows-arm64/cuda_cudart-windows-arm64-13.5.37-archive.zip"
+          curl -O "https://developer.download.nvidia.com/compute/cuda/redist/libcublas/windows-arm64/libcublas-windows-arm64-13.10.0.20-archive.zip"
+          unzip '*.zip' -d "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5"
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\cccl-windows-x86_64-13.3.5.0.6-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\cuda_crt-windows-x86_64-13.5.37-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\cuda_nvcc-windows-x86_64-13.5.37-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\libnvvm-windows-x86_64-13.5.37-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\cuda_cudart-windows-arm64-13.5.37-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\libcublas-windows-arm64-13.10.0.20-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" /E /I /H /Y
+          echo "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5\bin" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append
+          echo "CUDA_PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
+          echo "CUDA_PATH_V13_5=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.5" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
@@ -31,8 +31,8 @@ env:
  LLAMA_ARG_LOG_TIMESTAMPS: 1

 jobs:
-  ubuntu-24-llguidance:
-    runs-on: ${{ 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
+  ubuntu-26-llguidance:
+    runs-on: ${{ 'ubuntu-26.04-arm' || 'ubuntu-26.04' }}

    steps:
      - name: Clone
@@ -33,7 +33,7 @@ concurrency:

 jobs:
  android-ndk-snapdragon:
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04
    container:
      image: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.8'
    defaults:
@@ -67,7 +67,7 @@ jobs:
          path: pkg-snapdragon/llama.cpp

  linux-iot-snapdragon:
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04
    container:
      image: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.8'
    defaults:
@@ -103,7 +103,7 @@ jobs:
  test-snapdragon-qdc:
    name: Test on QDC Device (${{ matrix.device }})
    needs: [android-ndk-snapdragon, linux-iot-snapdragon]
-    runs-on: ubuntu-24.04-arm
+    runs-on: ubuntu-26.04-arm
    timeout-minutes: 180
    strategy:
      fail-fast: false
@@ -37,7 +37,7 @@ env:

 jobs:
  default:
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04

    steps:
      - name: Clone
@@ -56,6 +56,7 @@ jobs:
        uses: android-actions/setup-android@be39fa834029ff78f1a44aa3bb0819b8fc2bd8fd # v4.0.4
        with:
          log-accepted-android-sdk-licenses: false
+          packages: platform-tools cmake;3.31.6

      - name: Build
        run: |
@@ -63,7 +64,7 @@ jobs:
          ./gradlew build --no-daemon

  ndk:
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04
    container:
      image: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.3'
    defaults:
@@ -101,7 +102,7 @@ jobs:
          path: pkg-adb/llama.cpp

  arm64:
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04

    env:
      NDK_VERSION: "29.0.14206865"
@@ -40,8 +40,8 @@ jobs:
  #        path: ./spacemit_toolchain
  #        version: ${{ env.SPACEMIT_IME_TOOLCHAIN_VERSION }}

-  ubuntu-24-openvino-cache:
-    runs-on: ubuntu-24.04
+  ubuntu-26-openvino-cache:
+    runs-on: ubuntu-26.04

    env:
      # Sync versions in build-openvino.yml, ci-self-hosted-openvino.yml, release.yml, release-publish.yml, build-cache.yml, .devops/openvino.Dockerfile, docs/backend/OPENVINO.md
@@ -1,85 +0,0 @@
-name: Build relocatable cmake package
-on:
-  workflow_dispatch:
-  workflow_call:
-
-cache-mode: none
-permissions:
-  contents: read
-
-jobs:
-  linux:
-    runs-on: ubuntu-24.04
-    steps:
-      - name: Clone
-        uses: actions/checkout@v6
-        with:
-          fetch-depth: 0
-
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.24
-        with:
-          restore: false
-          save: false
-
-      - name: ccache-buckets-restore
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
-        with:
-          key: build-cmake-pkg
-          folder: llama.cpp
-          hf_bucket: ggml-org/cache
-
-      - name: Build
-        run: |
-          PREFIX="$(pwd)"/inst
-          cmake -S . -B build \
-                -DCMAKE_PREFIX_PATH="$PREFIX" \
-                -DLLAMA_OPENSSL=OFF \
-                -DLLAMA_BUILD_TESTS=OFF \
-                -DLLAMA_BUILD_TOOLS=OFF \
-                -DLLAMA_BUILD_EXAMPLES=OFF \
-                -DLLAMA_BUILD_APP=OFF \
-                -DLLAMA_BUILD_IS_DEV=OFF \
-                -DCMAKE_BUILD_TYPE=Release
-          cmake --build build --config Release -j $(nproc)
-          cmake --install build --prefix "$PREFIX" --config Release
-
-          export LLAMA_CONFIG="$PREFIX"/lib/cmake/llama/llama-config.cmake
-          build_commit=$(git rev-parse --short HEAD | xargs)
-          build_number=$(git rev-list --count HEAD | xargs)
-
-          major=$(grep -oE "set\(LLAMA_VERSION_MAJOR[[:space:]]+[0-9]+" CMakeLists.txt | grep -oE "[0-9]+$")
-          minor=$(grep -oE "set\(LLAMA_VERSION_MINOR[[:space:]]+[0-9]+" CMakeLists.txt | grep -oE "[0-9]+$")
-          patch=$(grep -oE "set\(LLAMA_VERSION_PATCH[[:space:]]+[0-9]+" CMakeLists.txt | grep -oE "[0-9]+$")
-          build_version="$major.$minor.$patch"
-
-          checks=("set\(LLAMA_VERSION[[:space:]]+$build_version\)"
-                  "set\(LLAMA_BUILD_COMMIT[[:space:]]+$build_commit\)"
-                  "set\(LLAMA_BUILD_NUMBER[[:space:]]+$build_number\)")
-
-          for check in "${checks[@]}"; do
-              if ! grep -qE "$check" "$LLAMA_CONFIG"; then
-                  echo "Checking llama-config.cmake version... \"$check\" failed!"
-                  exit 1
-              fi
-          done
-
-          echo "Checking llama-config.cmake version... success."
-
-          cd examples/simple-cmake-pkg
-          cmake -S . -B build -DCMAKE_PREFIX_PATH="$PREFIX"/lib/cmake
-          cmake --build build -j $(nproc)
-
-      - name: ccache-buckets-save
-        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
-        with:
-          key: build-cmake-pkg
-          folder: llama.cpp
-          evict-old-files: 1d
-          hf_bucket: ggml-org/cache
-          save: true
@@ -7,7 +7,6 @@ on:
      - master
    paths: [
      '.github/workflows/build-cpu.yml',
-      '.github/workflows/build-cmake-pkg.yml',
      '**/CMakeLists.txt',
      '**/.cmake',
      '**/*.h',
@@ -20,7 +19,6 @@ on:
    types: [opened, synchronize, reopened]
    paths: [
      '.github/workflows/build-cpu.yml',
-      '.github/workflows/build-cmake-pkg.yml',
      'ggml/src/ggml-rpc/**',
      '**/CMakeLists.txt',
      '**/.cmake',
@@ -47,17 +45,14 @@ env:
  LLAMA_ARG_LOG_TIMESTAMPS: 1

 jobs:
-  build-cmake-pkg:
-    uses: ./.github/workflows/build-cmake-pkg.yml
-
  ubuntu:
    strategy:
      matrix:
        include:
          - build: 'x64'
-            os: ubuntu-22.04
+            os: ubuntu-26.04
          - build: 'arm64'
-            os: ubuntu-24.04-arm
+            os: ubuntu-26.04-arm

    runs-on: ${{ matrix.os }}

@@ -81,13 +76,6 @@ jobs:
            libjpeg-dev build-essential libssl-dev \
            git-lfs

-      - name: Toolchain workaround (GCC 14)
-        if: ${{ contains(matrix.os, 'ubuntu-24.04') }}
-        run: |
-          sudo apt-get install -y gcc-14 g++-14
-          echo "CC=gcc-14" >> "$GITHUB_ENV"
-          echo "CXX=g++-14" >> "$GITHUB_ENV"
-
      - name: Python Dependencies
        id: python_depends
        run: |
@@ -242,3 +230,79 @@ jobs:
          older: 5m
          min: 1
          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}
+
+  build-cmake-pkg:
+    runs-on: ubuntu-24.04
+    steps:
+      - name: Clone
+        uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.24
+        with:
+          restore: false
+          save: false
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
+        with:
+          key: build-cmake-pkg
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
+
+      - name: Build
+        run: |
+          PREFIX="$(pwd)"/inst
+          cmake -S . -B build \
+                -DCMAKE_PREFIX_PATH="$PREFIX" \
+                -DLLAMA_OPENSSL=OFF \
+                -DLLAMA_BUILD_TESTS=OFF \
+                -DLLAMA_BUILD_TOOLS=OFF \
+                -DLLAMA_BUILD_EXAMPLES=OFF \
+                -DLLAMA_BUILD_APP=OFF \
+                -DLLAMA_BUILD_IS_DEV=OFF \
+                -DCMAKE_BUILD_TYPE=Release
+          cmake --build build --config Release -j $(nproc)
+          cmake --install build --prefix "$PREFIX" --config Release
+
+          export LLAMA_CONFIG="$PREFIX"/lib/cmake/llama/llama-config.cmake
+          build_commit=$(git rev-parse --short HEAD | xargs)
+          build_number=$(git rev-list --count HEAD | xargs)
+
+          major=$(grep -oE "set\(LLAMA_VERSION_MAJOR[[:space:]]+[0-9]+" CMakeLists.txt | grep -oE "[0-9]+$")
+          minor=$(grep -oE "set\(LLAMA_VERSION_MINOR[[:space:]]+[0-9]+" CMakeLists.txt | grep -oE "[0-9]+$")
+          patch=$(grep -oE "set\(LLAMA_VERSION_PATCH[[:space:]]+[0-9]+" CMakeLists.txt | grep -oE "[0-9]+$")
+          build_version="$major.$minor.$patch"
+
+          checks=("set\(LLAMA_VERSION[[:space:]]+$build_version\)"
+                  "set\(LLAMA_BUILD_COMMIT[[:space:]]+$build_commit\)"
+                  "set\(LLAMA_BUILD_NUMBER[[:space:]]+$build_number\)")
+
+          for check in "${checks[@]}"; do
+              if ! grep -qE "$check" "$LLAMA_CONFIG"; then
+                  echo "Checking llama-config.cmake version... \"$check\" failed!"
+                  exit 1
+              fi
+          done
+
+          echo "Checking llama-config.cmake version... success."
+
+          cd examples/simple-cmake-pkg
+          cmake -S . -B build -DCMAKE_PREFIX_PATH="$PREFIX"/lib/cmake
+          cmake --build build -j $(nproc)
+
+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: build-cmake-pkg
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
@@ -163,7 +163,7 @@ jobs:
  #         cmake --build build --config Release -j $(nproc)

  debian-13-loongarch64-cpu-cross:
-    runs-on: ${{ 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
+    runs-on: ${{ 'ubuntu-26.04-arm' || 'ubuntu-26.04' }}
    container: debian@sha256:653dfb9f86c3782e8369d5f7d29bb8faba1f4bff9025db46e807fa4c22903671

    steps:
@@ -218,7 +218,7 @@ jobs:
          cmake --build build --config Release -j $(nproc)

  debian-13-loongarch64-vulkan-cross:
-    runs-on: ${{ 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
+    runs-on: ${{ 'ubuntu-26.04-arm' || 'ubuntu-26.04' }}
    container: debian@sha256:653dfb9f86c3782e8369d5f7d29bb8faba1f4bff9025db46e807fa4c22903671

    steps:
@@ -276,8 +276,8 @@ jobs:

          cmake --build build --config Release -j $(nproc)

-  ubuntu-24-riscv64-cpu-spacemit-ime-cross:
-    runs-on: ubuntu-24.04
+  ubuntu-26-riscv64-cpu-spacemit-ime-cross:
+    runs-on: ubuntu-26.04

    env:
      # Make sure this is in sync with build-cache.yml
@@ -40,12 +40,12 @@ jobs:
          - cuda: '12.4'
            arch: x64
            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
-          - cuda: '13.4'
+          - cuda: '13.5'
            arch: x64
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
-          - cuda: '13.4'
+            defines: ''
+          - cuda: '13.5'
            arch: arm64
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3 -DCMAKE_TOOLCHAIN_FILE=cmake/arm64-windows-msvc-cuda.cmake'
+            defines: '-DCMAKE_TOOLCHAIN_FILE=cmake/arm64-windows-msvc-cuda.cmake'

    steps:
      - name: Clone
@@ -38,7 +38,7 @@ env:
  LLAMA_ARG_LOG_TIMESTAMPS: 1

 jobs:
-  ubuntu-24-sycl:
+  ubuntu-26-sycl:
    strategy:
      matrix:
        build: [fp32, fp16]
@@ -48,7 +48,7 @@ jobs:
          - build: fp16
            fp16: ON

-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-26.04

    env:
      ONEAPI_ROOT: /opt/intel/oneapi/
@@ -103,7 +103,7 @@ jobs:
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
        with:
-          key: sycl-ubuntu-24-${{ matrix.build }}
+          key: sycl-ubuntu-26-${{ matrix.build }}
          folder: llama.cpp
          hf_bucket: ggml-org/cache

@@ -131,7 +131,7 @@ jobs:
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
        with:
-          key: sycl-ubuntu-24-${{ matrix.build }}
+          key: sycl-ubuntu-26-${{ matrix.build }}
          folder: llama.cpp
          evict-old-files: 1d
          hf_bucket: ggml-org/cache
@@ -31,8 +31,8 @@ concurrency:
  cancel-in-progress: true

 jobs:
-  ubuntu-24-virtgpu:
-    runs-on: ${{ 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
+  ubuntu-26-virtgpu:
+    runs-on: ${{ 'ubuntu-26.04-arm' || 'ubuntu-26.04' }}

    steps:
      - name: Clone
@@ -42,7 +42,7 @@ env:

 jobs:
  ubuntu-arm64:
-    runs-on: ubuntu-24.04-arm
+    runs-on: ubuntu-26.04-arm

    steps:
      - name: Clone
@@ -53,9 +53,7 @@ jobs:
        id: depends
        run: |
          sudo apt-get update
-          sudo apt-get install -y gcc-14 g++-14 build-essential glslc libvulkan-dev spirv-headers libssl-dev ninja-build
-          echo "CC=gcc-14" >> "$GITHUB_ENV"
-          echo "CXX=g++-14" >> "$GITHUB_ENV"
+          sudo apt-get install -y build-essential glslc libvulkan-dev spirv-headers libssl-dev ninja-build

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.24
@@ -69,7 +67,7 @@ jobs:
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
        with:
-          key: vulkan-ubuntu-24.04-arm
+          key: vulkan-ubuntu-26.04-arm
          folder: llama.cpp
          hf_bucket: ggml-org/cache

@@ -92,14 +90,14 @@ jobs:
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
        with:
-          key: vulkan-ubuntu-24.04-arm
+          key: vulkan-ubuntu-26.04-arm
          folder: llama.cpp
          evict-old-files: 1d
          hf_bucket: ggml-org/cache
          save: true

  ubuntu-llvmpipe:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-26.04

    steps:
      - name: Clone
@@ -137,7 +135,7 @@ jobs:
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
        with:
-          key: vulkan-ubuntu-24.04-llvmpipe
+          key: vulkan-ubuntu-26.04-llvmpipe
          folder: llama.cpp
          hf_bucket: ggml-org/cache

@@ -156,7 +154,7 @@ jobs:
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
        with:
-          key: vulkan-ubuntu-24.04-llvmpipe
+          key: vulkan-ubuntu-26.04-llvmpipe
          folder: llama.cpp
          evict-old-files: 1d
          hf_bucket: ggml-org/cache
@@ -49,7 +49,7 @@ env:

 jobs:
  ubuntu-webgpu:
-    runs-on: ubuntu-24.04-arm
+    runs-on: ubuntu-26.04-arm

    steps:
      - name: Clone
@@ -97,7 +97,7 @@ jobs:
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
        with:
-          key: webgpu-ubuntu-24.04-arm-wasi
+          key: webgpu-ubuntu-26.04-arm-wasi
          folder: llama.cpp
          hf_bucket: ggml-org/cache

@@ -135,7 +135,7 @@ jobs:
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
        with:
-          key: webgpu-ubuntu-24.04-arm-wasi
+          key: webgpu-ubuntu-26.04-arm-wasi
          folder: llama.cpp
          evict-old-files: 1d
          hf_bucket: ggml-org/cache
@@ -50,7 +50,7 @@ env:

 jobs:
  ubuntu-webgpu:
-    runs-on: ubuntu-24.04-arm
+    runs-on: ubuntu-26.04-arm

    steps:
      - name: Clone
@@ -84,7 +84,7 @@ jobs:
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
        with:
-          key: webgpu-ubuntu-24.04-arm-wasm
+          key: webgpu-ubuntu-26.04-arm-wasm
          folder: llama.cpp
          hf_bucket: ggml-org/cache

@@ -107,7 +107,7 @@ jobs:
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
        with:
-          key: webgpu-ubuntu-24.04-arm-wasm
+          key: webgpu-ubuntu-26.04-arm-wasm
          folder: llama.cpp
          evict-old-files: 1d
          hf_bucket: ggml-org/cache
@@ -34,7 +34,7 @@ env:

 jobs:
  format:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-26.04

    steps:
      - name: Clone
@@ -45,7 +45,7 @@ jobs:
          wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key |
            sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc > /dev/null
          sudo add-apt-repository -y \
-            "deb http://apt.llvm.org/noble/ llvm-toolchain-noble-22 main"
+            "deb http://apt.llvm.org/resolute/ llvm-toolchain-resolute-22 main"
          sudo apt-get update
          sudo apt-get install -y clang-format-22

@@ -48,6 +48,12 @@ jobs:
  cpu-x64-high-perf:
    runs-on: hf-jobs-cpu-upgrade:ubuntu26_04

+    # note: this runner has 8 CPUs, but std::thread::hardware_concurrency reports the cores of the
+    #   host (https://github.com/ggml-org/llama.cpp/pull/31439#issuecomment-6683386059)
+    env:
+      TEST_ARG_HARDWARE_CONCURRENCY: 8
+      TEST_ARG_PARALLEL_WORKERS: 4
+
    steps:
      - name: Clone
        id: checkout
@@ -16,7 +16,7 @@ cache-mode: none
 jobs:
  # The job MUST be called `copilot-setup-steps` or it will not be picked up by Copilot.
  copilot-setup-steps:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-26.04

    # Set the permissions to the lowest permissions possible needed for your steps.
    # Copilot will be given its own token for its operations.
@@ -70,7 +70,7 @@ jobs:

  prepare_matrices:
    name: Prepare Docker matrices
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-26.04
    outputs:
      build_matrix: ${{ steps.matrices.outputs.build_matrix }}
      merge_matrix: ${{ steps.matrices.outputs.merge_matrix }}
@@ -87,19 +87,19 @@ jobs:
          # Keep all build targets in one place and derive merge targets from it.
          cat > build-matrix.json <<'JSON'
          [
-            { "tag": "cpu", "dockerfile": ".devops/cpu.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04" },
-            { "tag": "cpu", "dockerfile": ".devops/cpu.Dockerfile", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04-arm" },
-            { "tag": "cpu", "dockerfile": ".devops/s390x.Dockerfile", "platforms": "linux/s390x", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04-s390x", "prebuilt_ui": true },
-            { "tag": "cuda cuda12", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "12.8.1", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" },
-            { "tag": "cuda cuda12", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "12.8.1", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04-arm" },
-            { "tag": "cuda13", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "13.4.1", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" },
-            { "tag": "cuda13", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "13.4.1", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04-arm" },
-            { "tag": "musa", "dockerfile": ".devops/musa.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" },
-            { "tag": "intel", "dockerfile": ".devops/intel.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" },
-            { "tag": "vulkan", "dockerfile": ".devops/vulkan.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04" },
-            { "tag": "vulkan", "dockerfile": ".devops/vulkan.Dockerfile", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04-arm" },
-            { "tag": "rocm", "dockerfile": ".devops/rocm.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-24.04" },
-            { "tag": "openvino", "dockerfile": ".devops/openvino.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-24.04" }
+            { "tag": "cpu", "dockerfile": ".devops/cpu.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-26.04" },
+            { "tag": "cpu", "dockerfile": ".devops/cpu.Dockerfile", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-26.04-arm" },
+            { "tag": "cpu", "dockerfile": ".devops/s390x.Dockerfile", "platforms": "linux/s390x", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-26.04-s390x", "prebuilt_ui": true },
+            { "tag": "cuda cuda12", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "12.8.1", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04" },
+            { "tag": "cuda cuda12", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "12.8.1", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04-arm" },
+            { "tag": "cuda13", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "13.5.0", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04" },
+            { "tag": "cuda13", "dockerfile": ".devops/cuda.Dockerfile", "cuda_version": "13.5.0", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04-arm" },
+            { "tag": "musa", "dockerfile": ".devops/musa.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04" },
+            { "tag": "intel", "dockerfile": ".devops/intel.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04" },
+            { "tag": "vulkan", "dockerfile": ".devops/vulkan.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-26.04" },
+            { "tag": "vulkan", "dockerfile": ".devops/vulkan.Dockerfile", "platforms": "linux/arm64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-26.04-arm" },
+            { "tag": "rocm", "dockerfile": ".devops/rocm.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": true, "runs_on": "ubuntu-26.04" },
+            { "tag": "openvino", "dockerfile": ".devops/openvino.Dockerfile", "platforms": "linux/amd64", "full": true, "light": true, "server": true, "free_disk_space": false, "runs_on": "ubuntu-26.04" }
          ]
          JSON

@@ -400,7 +400,7 @@ jobs:
  merge_arch_tags:
    name: Create shared tags from digests
    needs: [prepare_matrices, push_to_registry, create_tag]
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-26.04
    permissions:
      contents: read
      packages: write
@@ -24,7 +24,7 @@ permissions:
 jobs:
  deploy:

-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04

    steps:
    - uses: actions/checkout@v6
@@ -39,7 +39,7 @@ permissions:

 jobs:
  make-release:
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04

    steps:
      - name: Checkout
@@ -11,7 +11,8 @@ on:
      'tests/fusion/**',
      'tests/test-fusion.cpp',
      'tests/test-llama-archs.cpp',
-      'src/llama-graph.cpp',
+      'src/llama-graph*',
+      'src/llama-kv-cache*',
      'src/llama-model*',
      'src/models/**'
    ]
@@ -24,7 +25,8 @@ on:
      'tests/fusion/**',
      'tests/test-fusion.cpp',
      'tests/test-llama-archs.cpp',
-      'src/llama-graph.cpp',
+      'src/llama-graph*',
+      'src/llama-kv-cache*',
      'src/llama-model*',
      'src/models/**'
    ]
@@ -145,6 +147,11 @@ jobs:
          rm -rf build-ci-models && mkdir -p build-ci-models
          ./build/bin/test-llama-archs -o build-ci-models

+      - name: Test KV shift
+        id: test_kv_shift
+        run: |
+          ./build/bin/test-llama-archs --kv-shift
+
      - name: Test fusion
        id: test_fusion
        run: |
@@ -293,6 +300,11 @@ jobs:
      #   run: |
      #     ./build/bin/test-fusion --models build-ci-models --device Vulkan0 --check tests/fusion/Vulkan.csv

+      - name: Test KV shift
+        id: test_kv_shift
+        run: |
+          ./build/bin/test-llama-archs --kv-shift
+
      - name: Test archs
        id: test_archs
        run: |
@@ -24,7 +24,7 @@ concurrency:

 jobs:
  python-check-requirements:
-    runs-on: ${{ 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
+    runs-on: ${{ 'ubuntu-26.04-arm' || 'ubuntu-26.04' }}
    name: check-requirements
    steps:
      - name: Check out source repository
@@ -171,14 +171,14 @@ jobs:
            **Linux:**
            - [Ubuntu x64 (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-x64.tar.gz)
            - [Ubuntu arm64 (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-arm64.tar.gz)
-            - [Ubuntu s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz)
+            - [Ubuntu 24.04 s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz)
            - [Ubuntu 26.04 s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-26.04-s390x.tar.gz)
            - [Ubuntu 26.04 s390x (zDNN)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-zdnn-s390x.tar.gz)
            - [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz)
            - [Ubuntu arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-arm64.tar.gz)
            - [Ubuntu x64 (CUDA 12)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-12.8-x64.tar.gz) - [CUDA 12.8 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-12.8-x64.tar.gz)
-            - [Ubuntu x64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.4-x64.tar.gz) - [CUDA 13.4 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.4-x64.tar.gz)
-            - [Ubuntu arm64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.4-arm64.tar.gz) - [CUDA 13.4 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.4-arm64.tar.gz)
+            - [Ubuntu x64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.5-x64.tar.gz) - [CUDA 13.5 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.5-x64.tar.gz)
+            - [Ubuntu arm64 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.5-arm64.tar.gz) - [CUDA 13.5 libraries](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-${{ steps.tag.outputs.name }}-bin-ubuntu-cuda-13.5-arm64.tar.gz)
            - [Ubuntu x64 (ROCm 10.0)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-10.0-x64.tar.gz)
            - [Ubuntu x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-openvino-2026.4.1-x64.tar.gz)
            - [Ubuntu x64 (SYCL FP32)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp32-x64.tar.gz)
@@ -194,8 +194,8 @@ jobs:
            - [Windows arm64 (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-cpu-arm64.zip)
            - [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.4-x64.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-x64.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 (CUDA 13)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-cuda-13.5-x64.zip) - [CUDA 13.5 DLLs](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-bin-win-cuda-13.5-x64.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.5-arm64.zip) - [CUDA 13.5 DLLs](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/cudart-llama-bin-win-cuda-13.5-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 arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-vulkan-arm64.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-2026.4.1-x64.zip)
@@ -345,14 +345,14 @@ jobs:
            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
          - build: 'x64'
            os: ubuntu-24.04
-            cuda: '13.4.1'
-            label: '13.4'
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
+            cuda: '13.5.0'
+            label: '13.5'
+            defines: ''
          - build: 'arm64'
            os: ubuntu-24.04-arm
-            cuda: '13.4.1'
-            label: '13.4'
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
+            cuda: '13.5.0'
+            label: '13.5'
+            defines: ''

    runs-on: ${{ matrix.os }}
    container: nvidia/cuda:${{ matrix.cuda }}-devel-ubuntu24.04
@@ -463,7 +463,7 @@ jobs:
    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-24.04

    env:
      NDK_VERSION: "29.0.14206865"
@@ -548,7 +548,7 @@ jobs:
    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
    container: 'ghcr.io/snapdragon-toolchain/arm64-android:v0.8'

    defaults:
@@ -601,7 +601,7 @@ jobs:
    needs: [check-release, ui-build]
    if: ${{ needs.check-release.outputs.should_release == 'true' }}

-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
    container: 'ghcr.io/snapdragon-toolchain/arm64-linux:v0.8'

    defaults:
@@ -1211,12 +1211,12 @@ jobs:
          - cuda: '12.4'
            arch: x64
            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
-          - cuda: '13.4'
+          - cuda: '13.5'
            arch: x64
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
-          - cuda: '13.4'
+            defines: ''
+          - cuda: '13.5'
            arch: arm64
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3 -DCMAKE_TOOLCHAIN_FILE=cmake/arm64-windows-msvc-cuda.cmake'
+            defines: '-DCMAKE_TOOLCHAIN_FILE=cmake/arm64-windows-msvc-cuda.cmake'

    steps:
      - name: Clone
@@ -59,7 +59,7 @@ concurrency:

 jobs:
  ubuntu:
-    runs-on: ubuntu-24.04-arm
+    runs-on: ubuntu-26.04-arm

    steps:
      - name: Dependencies
@@ -94,7 +94,7 @@ jobs:
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
        with:
-          key: server-ubuntu-24.04-arm
+          key: server-ubuntu-26.04-arm
          folder: llama.cpp
          hf_bucket: ggml-org/cache

@@ -111,7 +111,7 @@ jobs:
        env:
          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
        with:
-          key: server-ubuntu-24.04-arm
+          key: server-ubuntu-26.04-arm
          folder: llama.cpp
          evict-old-files: 1d
          hf_bucket: ggml-org/cache
@@ -12,7 +12,7 @@ permissions:
 jobs:
  update:
    name: Update Winget Package
-    runs-on: ubuntu-24.04  # previously ubuntu-latest
+    runs-on: ubuntu-26.04
    if: github.repository_owner == 'ggml-org'

    steps:
@@ -87,6 +87,7 @@
 /src/llama-graph.*                      @CISC
 /src/llama-model.*                      @CISC
 /src/llama-vocab.*                      @CISC
+/src/llama-moe-cache.*                  @am17an
 /src/models/                            @ggml-org/llama-models
 /tests/                                 @ggerganov
 /tests/test-chat.*                      @pwilkin
@@ -437,12 +437,8 @@ function gg_run_qwen3_0_6b {

    (time ./bin/llama-imatrix --model ${model_f16} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-imatrix.log

-    # TODO: these tests are extremely slow on WebGPU - re-enable when fixed
-    #       https://github.com/ggml-org/llama.cpp/actions/runs/39639985732
-    if [ -z ${GG_BUILD_WEBGPU} ]; then
-        (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
-        (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa off                ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
-    fi
+    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa off                ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa on  --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa on                 ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log

@@ -13,7 +13,7 @@ if(WIN32)
        endif()
        enable_language(RC)
        set(LLAMA_WIN_APP_MANIFEST_RC "${CMAKE_CURRENT_BINARY_DIR}/win-app-manifest.rc")
-        configure_file("${CMAKE_CURRENT_LIST_DIR}/win-app.rc.in" "${LLAMA_WIN_APP_MANIFEST_RC}" @ONLY)
+        file(CONFIGURE OUTPUT "${LLAMA_WIN_APP_MANIFEST_RC}" CONTENT "1 24 \"@LLAMA_WIN_APP_MANIFEST@\"\n" @ONLY)
        set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${LLAMA_WIN_APP_MANIFEST}")
    endif()
 endif()
@@ -1,3 +0,0 @@
-// generated from cmake/win-app.rc.in - do not edit
-// RT_MANIFEST (24), ID 1: the application manifest shared by all llama.cpp executables
-1 24 "@LLAMA_WIN_APP_MANIFEST@"
@@ -4,6 +4,7 @@
 #include "chat.h"
 #include "common.h"
 #include "download.h"
+#include "http.h"
 #include "json-schema-to-grammar.h"
 #include "json.h"
 #include "llama.h"
@@ -399,7 +400,26 @@ bool common_models_handler_is_preset_repo(const common_models_handler & handler)
    return !handler.plan.preset.url.empty();
 }

+static common_download_opts download_opts_for_url(const common_download_opts & opts, const std::string & url) {
+    common_download_opts result = opts;
+    try {
+        const auto hf    = common_http_parse_url(common_get_model_endpoint());
+        const auto parts = common_http_parse_url(url);
+
+        if (parts.scheme == hf.scheme &&
+            parts.host   == hf.host &&
+            parts.port   == hf.port &&
+            string_starts_with(result.bearer_token, "hf_")) {
+            return result;
+        }
+    } catch (const std::exception &) {
+    }
+    result.bearer_token.clear();
+    return result;
+}
+
 static std::vector<common_download_task> build_url_tasks(const common_params_model & model, common_download_opts opts) {
+    opts = download_opts_for_url(opts, model.url);
    auto parts = common_download_get_all_parts(model.url);
    std::vector<common_download_task> tasks;

@@ -488,7 +508,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params
        common_download_task task;
        task.url        = params.mmproj.url;
        task.local_path = params.mmproj.path;
-        task.opts       = opts;
+        task.opts       = download_opts_for_url(opts, task.url);
        tasks.push_back(task);
    }
    bool had_spec_url = false;
@@ -496,7 +516,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params
        common_download_task task;
        task.url        = params.speculative.draft.mparams.url;
        task.local_path = params.speculative.draft.mparams.path;
-        task.opts       = opts;
+        task.opts       = download_opts_for_url(opts, task.url);
        tasks.push_back(task);
        had_spec_url = true;
    }
@@ -2734,6 +2754,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            params.moe_cache_size = (size_t) value*1024*1024;
        }
    ).set_env("LLAMA_ARG_MOE_CACHE_MIB"));
+    add_opt(common_arg(
+        {"--moe-cache-cpu-share"}, "F",
+        string_format("share of the MoE cache misses that the CPU computes from the host experts instead of uploading them to the GPU once the cache is full, from 0 to 1 (default: %.2f)", (double) params.moe_cache_cpu_share),
+        [](common_params & params, const std::string & value) {
+            const float share = std::stof(value);
+            if (share < 0.0f || share > 1.0f) {
+                throw std::invalid_argument("invalid value");
+            }
+            params.moe_cache_cpu_share = share;
+        }
+    ).set_env("LLAMA_ARG_MOE_CACHE_CPU_SHARE"));
    add_opt(common_arg(
        {"-ncffn", "--n-cpu-ffn"}, "N",
        "keep the dense FFN weights of the first N layers in the CPU\n"
@@ -348,7 +348,8 @@ void common_chat_peg_mapper::map(const common_peg_ast_node & node) {
    }

    if (is_tool_id && current_tool) {
-        auto text = trim_trailing_space(node.text);
+        auto sanitized = node.sanitized_text();
+        auto text      = trim_trailing_space(sanitized);
        if (text.size() >= 2 && text.front() == '"' && text.back() == '"') {
            text = text.substr(1, text.size() - 2);
        }
@@ -401,7 +402,7 @@ void common_chat_peg_mapper::map(const common_peg_ast_node & node) {
    }

    if ((is_arg_value || is_arg_string_value) && current_tool) {
-        std::string value_content = std::string(node.text);
+        std::string value_content = node.sanitized_text();

        std::string value_to_add;
        if (value_content.empty() && is_arg_string_value) {
@@ -978,7 +979,7 @@ static std::string gemma4_to_json(const common_peg_ast_arena & arena, common_peg
    }

    if (node.rule == "gemma4-string-content") {
-        return escape_json_string_inner(std::string(node.text));
+        return escape_json_string_inner(node.sanitized_text());
    }

    if (node.rule == "gemma4-string") {
@@ -1195,7 +1196,7 @@ static std::string minimax_m3_value_to_json(const common_peg_ast_arena & arena,
    }

    if (node.tag == common_chat_peg_builder::TOOL_ARG_STRING_VALUE) {
-        return "\"" + escape_json_string_inner(std::string(node.text)) + (closed ? "\"" : "");
+        return "\"" + escape_json_string_inner(node.sanitized_text()) + (closed ? "\"" : "");
    }

    // Numbers and booleans are written verbatim by the template
@@ -895,27 +895,7 @@ bool fs_validate_filename(const std::string & filename, bool allow_subdirs) {
    return true;
 }

-#include <iostream>
-
-
 #ifdef _WIN32
-std::wstring utf8_to_wstring(const std::string & str) {
-    if (str.empty()) {
-        return std::wstring();
-    }
-
-    int size = MultiByteToWideChar(CP_UTF8, 0, str.c_str(), (int)str.size(), NULL, 0);
-
-    if (size <= 0) {
-        return std::wstring();
-    }
-
-    std::wstring wstr(size, 0);
-    MultiByteToWideChar(CP_UTF8, 0, str.c_str(), (int)str.size(), &wstr[0], size);
-
-    return wstr;
-}
-
 std::string wstring_to_utf8(const std::wstring & str) {
    if (str.empty()) {
        return std::string();
@@ -1333,6 +1313,12 @@ common_init_result::common_init_result(common_params & params, bool model_only)
        LOG_INF("%s", "decision model reads the embeddings output, enabling embedding mode\n");
    }

+    if (llama_model_classifier_free_guidance(model)) {
+        params.kv_unified   = true;
+        cparams.kv_unified  = true;
+        cparams.n_seq_max  *= 2;
+    }
+
    // embeddings need the whole batch in one ubatch, so n_batch must not be larger than n_ubatch
    // (server.cpp does this check for --embedding, but before the model is loaded)
    if (cparams.embeddings && cparams.n_batch > cparams.n_ubatch) {
@@ -1772,7 +1758,8 @@ struct llama_context_params common_context_params_to_llama(const common_params &
    cparams.type_k = params.cache_type_k;
    cparams.type_v = params.cache_type_v;

-    cparams.moe_cache_size = params.moe_cache_size;
+    cparams.moe_cache_size      = params.moe_cache_size;
+    cparams.moe_cache_cpu_share = params.moe_cache_cpu_share;

    return cparams;
 }
@@ -597,7 +597,8 @@ struct common_params {
    ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K
    ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V

-    size_t moe_cache_size = 0; // GPU cache size in bytes for the MoE experts kept in the CPU, split among the GPUs like the layers
+    size_t moe_cache_size      = 0;    // GPU cache size in bytes for the MoE experts kept in the CPU, split among the GPUs like the layers
+    float  moe_cache_cpu_share = 0.0f; // share of the MoE cache misses computed on the CPU instead of uploaded

    common_conversation_mode conversation_mode = COMMON_CONVERSATION_MODE_AUTO;

@@ -906,8 +907,7 @@ bool glob_match(const std::string & pattern, const std::string & str);
 //

 #ifdef _WIN32
-std::wstring utf8_to_wstring(const std::string & str);
-std::string  wstring_to_utf8(const std::wstring & str);
+std::string wstring_to_utf8(const std::wstring & str);
 #endif

 //
@@ -20,6 +20,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "NEOChatModel": "sensenova_u1",
    "AfmoeForCausalLM": "afmoe",
    "LagunaForCausalLM": "laguna",
+    "RicursiveForCausalLM": "ricursive",
    "FlintForCausalLM": "flint",
    "ApertusForCausalLM": "llama",
    "ArceeForCausalLM": "llama",
@@ -62,6 +63,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "DeepseekV2ForCausalLM": "deepseek",
    "DeepseekV3ForCausalLM": "deepseek",
    "DeepseekV32ForCausalLM": "deepseek",
+    "Xing4_0ForCausalLM": "xing",
    "DFlashDraftModel": "qwen",
    "DFlash2DraftModel": "qwen",
    "LiLiCorrDraftModel": "nemotron",
@@ -156,6 +158,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "KORMoForCausalLM": "qwen",
    "KimiK25ForConditionalGeneration": "deepseek",
    "KimiK3ForConditionalGeneration": "kimi_k3",
+    "KimiK31ForConditionalGeneration": "kimi_k3",
    "KimiLinearForCausalLM": "kimi_linear",
    "KimiLinearModel": "kimi_linear",
    "KimiVLForConditionalGeneration": "deepseek",
@@ -190,6 +193,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "Eagle3LlamaForCausalLM": "llama",
    "LlamaForCausalLMEagle3": "llama",
    "LlavaForConditionalGeneration": "llama",
+    "LlavaOnevision2ForConditionalGeneration": "qwen",
    "LlavaStableLMEpochForCausalLM": "stablelm",
    "MPTForCausalLM": "mpt",
    "MT5ForConditionalGeneration": "t5",
@@ -377,6 +381,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "LightOnOCRForConditionalGeneration": "lighton_ocr",
    "Llama4ForConditionalGeneration": "llama4",
    "LlavaForConditionalGeneration": "llava",
+    "LlavaOnevision2ForConditionalGeneration": "magevl",
    "MERaLiON2ForConditionalGeneration": "ultravox",
    "MiMoV2ForCausalLM": "mimo",
    "MiniMaxM3SparseForConditionalGeneration": "minimax",
@@ -3404,7 +3404,7 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st
    # Kimi-K3's text_config reports "KimiLinearForCausalLM", which is the older
    # Kimi-Linear-48B architecture and cannot load K3 (no attention residuals,
    # latent MoE, situ, ...). Route on the top-level architecture instead.
-    if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration", "Nemotron3_5TranscribeForConditionalGeneration"):
+    if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration", "KimiK31ForConditionalGeneration", "Nemotron3_5TranscribeForConditionalGeneration"):
        return arch

    # if "architectures" is found in the sub-config, use that instead
@@ -3412,6 +3412,9 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st
        arch = text_config["architectures"][0]
    elif model_type == ModelType.MMPROJ and vision_config.get("architectures") is not None:
        arch = vision_config["architectures"][0]
+    # some models (llava_onevision2) only declare the architecture inside text_config
+    if arch is None and text_config.get("architectures") is not None:
+        arch = text_config["architectures"][0]
    if arch is None:
        raise ValueError("Failed to detect model architecture")
    return arch
@@ -212,6 +212,8 @@ class ChatterboxTalkerModel(TextModel):
        self.gguf_writer.add_sampling_temp(sampling["temp"])
        self.gguf_writer.add_sampling_penalty_repeat(sampling["penalty_repeat"])
        self.gguf_writer.add_sampling_penalty_last_n(PENALTY_LAST_N)
+        if not self.is_turbo:
+            self.gguf_writer.add_classifier_free_guidance(True)

    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
        if self.is_turbo:
@@ -353,7 +355,7 @@ class ChatterboxMmprojModel(MmprojModel):
    def set_gguf_parameters(self):
        self.gguf_writer.add_file_type(self.ftype)

-        # speaker encoder, the front-end hparams are hard-coded in clip.cpp
+        # speaker encoder, the front-end hparams are hard-coded in tools/mtmd
        self.gguf_writer.add_clip_has_audio_encoder(True)
        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.CHATTERBOX_SPKENC)
        self.gguf_writer.add_audio_projection_dim(80)
@@ -1078,7 +1078,7 @@ class DeepseekV4FlashVisionModel(MmprojModel):
        self.gguf_writer.add_vision_use_silu(True) # SwiGLU MLP
        self.gguf_writer.add_vision_projector_scale_factor(self.hparams_vision["downsample_ratio"])
        self.gguf_writer.add_vision_min_pixels(self.hparams_vision["min_pixels"])
-        # hardcoded on the C++ side (see PROJECTOR_TYPE_DEEPSEEK4V in clip.cpp)
+        # hardcoded on the C++ side (see PROJECTOR_TYPE_DEEPSEEK4V in tools/mtmd)
        # if future models use different values, add GGUF keys for those
        assert self.global_config["vision_max_n_token"] == 384
        assert self.global_config["vision_max_wh_ratio"] == 8
@@ -458,7 +458,7 @@ class DeepseekV41VisionModel(DeepseekV4FlashVisionModel):
        self.gguf_writer.add_vision_use_silu(True)
        self.gguf_writer.add_vision_projector_scale_factor(self.hparams_vision["downsample_ratio"])
        self.gguf_writer.add_vision_min_pixels(self.hparams_vision["min_pixels"])
-        # hardcoded in clip.cpp (PROJECTOR_TYPE_DEEPSEEK41V); add GGUF keys if a later model differs
+        # hardcoded in tools/mtmd (PROJECTOR_TYPE_DEEPSEEK41V); add GGUF keys if a later model differs
        assert self.hparams_vision["max_image_tokens"] == 1024, self.hparams_vision["max_image_tokens"]
        assert self.hparams_vision["max_wh_ratio"] is None, self.hparams_vision["max_wh_ratio"]
        assert self.hparams_vision["rope_theta"] == 10000, self.hparams_vision["rope_theta"]
@@ -77,7 +77,7 @@ class Drama3MmprojModel(MmprojModel):
        self.gguf_writer.add_gen_audio_head_count(cfg["n_head"])
        self.gguf_writer.add_gen_audio_head_count_kv(cfg["n_local_heads"])
        self.gguf_writer.add_gen_audio_attention_layernorm_eps(cfg["norm_eps"])
-        # other tower and codec hparams are hardcoded in clip.cpp
+        # other tower and codec hparams are hardcoded in tools/mtmd

    @classmethod
    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
@@ -343,3 +343,33 @@ class KimiK3Model(TextModel):
            return

        yield from super().modify_tensors(data_torch, name, bid)
+
+
+@ModelBase.register("KimiK31ForConditionalGeneration")
+@ModelBase.example("moonshotai/Kimi-K3.1-Pro")
+class KimiK31Model(KimiK3Model):
+    """
+    Kimi-K3.1: K3 with the full-attention MLA layers replaced by gated DSA.
+
+    The tensor set is K3's plus a lightning indexer per full-attention layer,
+    and `dsa_config` adds the attention.indexer.* metadata.
+    """
+
+    model_arch = gguf.MODEL_ARCH.KIMI_K31
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+
+        # --- DSA indexer ---
+        dsa = self.hparams["dsa_config"]
+        dsa_layers = dsa["dsa_layers"]
+        # the layer lists are 1-indexed, like linear_attn_config's
+        assert sorted(dsa_layers) == sorted(self.hparams["linear_attn_config"]["full_attn_layers"]), \
+            "every full-attention layer is expected to own a DSA indexer"
+        self.gguf_writer.add_indexer_head_count(dsa["index_num_heads"])
+        self.gguf_writer.add_indexer_key_length(dsa["index_head_dim"])
+        self.gguf_writer.add_indexer_top_k(dsa["index_topk"])
+        self.gguf_writer.add_indexer_types([(il + 1) in dsa_layers
+                                            for il in range(self.hparams["num_hidden_layers"])])
+        # eps of the indexer key LayerNorm; the config declares no separate value
+        self.gguf_writer.add_layer_norm_eps(dsa.get("index_norm_eps", self.hparams["rms_norm_eps"]))
@@ -181,11 +181,11 @@ class LFM3VLVisionMixin(MmprojModel):

    def _set_lfm3vl_gguf_parameters(self):
        assert self.hparams_vision is not None
-        # clip.cpp hardcodes the GELU projector and its input norm eps
+        # mtmd hardcodes the GELU projector and its input norm eps
        if (self.global_config["projector_hidden_act"] != "gelu" or not self.global_config["projector_bias"]
                or self.global_config["projector_layer_norm_eps"] != 1e-6):
            raise ValueError("unsupported projector_hidden_act / projector_bias / projector_layer_norm_eps")
-        # clip.cpp hardcodes the vision rope base, the RGB patchifier and the 3x3 spatial conv
+        # mtmd hardcodes the vision rope base, the RGB patchifier and the 3x3 spatial conv
        if (self.hparams_vision["rope_parameters"]["rope_theta"] != 1e6 or self.hparams_vision["in_channels"] != 3
                or self.hparams_vision["conv_kernel_size"] != 3):
            raise ValueError("unsupported vision rope_theta / in_channels / conv_kernel_size")
@@ -347,7 +347,7 @@ class LFM3AudioMmprojModel(MmprojModel):
        "feed_forward.w3":       gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP,
        "feed_forward.w2":       gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN,
    }
-    # the codec graph in clip.cpp hardcodes the liquid-codec FSQ front end and decoder shape, conversion fails on a mismatch
+    # the codec graph in mtmd hardcodes the liquid-codec FSQ front end and decoder shape, conversion fails on a mismatch
    _codec = {
        "model_type": "liquid_codec", "fsq_levels": [8, 7, 6, 6], "n_q": 12, "quantizer_dim": 256,
        "encoder_dim": 512, "stft_n_fft": 1280, "stft_hop_length": 320, "stft_win_length": 1280, "stft_frame_stack": 6,
@@ -405,7 +405,7 @@ class LFM3AudioMmprojModel(MmprojModel):
            if fe[k] != v:
                raise ValueError(f"unsupported feature_extractor.{k}={fe[k]!r}, expected {v!r}")

-        # clip.cpp hardcodes the encoder rope base
+        # mtmd hardcodes the encoder rope base
        if self.hparams_audio["rope_parameters"]["rope_theta"] != 1e6:
            raise ValueError("unsupported audio encoder rope_theta")

@@ -420,7 +420,7 @@ class LFM3AudioMmprojModel(MmprojModel):
        if not self._audio_output:
            return

-        # audio generation: depthformer hparams, the codec ones are constant and live in clip.cpp
+        # audio generation: depthformer hparams, the codec ones are constant and live in mtmd
        gen = self.get_generation_config()
        for k, v in self._generation.items():
            if gen[k] != v:
@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from typing import Callable, Iterable, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import MmprojModel, ModelBase, gguf
+
+
+@ModelBase.register("LlavaOnevision2ForConditionalGeneration")
+@ModelBase.example("microsoft/Mage-VL-Cowork")
+class MageVLVisionModel(MmprojModel):
+    # a rope pair uses 2 different angles: the graph runs 2 ropes on heads permuted to [evens | odds]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.hparams_vision is not None
+
+        self.n_embd = self.find_vparam(["hidden_size"])
+        self.n_head = self.find_vparam(["num_attention_heads"])
+        self.d_head = self.n_embd // self.n_head
+
+        d_half = self.d_head // 2
+        assert d_half % 16 == 0, "the 4:6:6 split of the rope requires d_head to be a multiple of 32"
+
+        # 4:6:6 split for t:h:w
+        n_t = 4 * (d_half // 16)
+        n_s = 6 * (d_half // 16)
+        theta = float(self.find_vparam(["rope_theta"]))
+        freq_t = 1.0 / theta ** (torch.arange(n_t, dtype=torch.float64) / n_t)
+        freq_s = 1.0 / theta ** (torch.arange(n_s, dtype=torch.float64) / n_s)
+        freq = torch.cat([freq_t, freq_s, freq_s])
+        self.rope_freq = torch.cat([freq, freq])
+
+        def pair_axis(i: int) -> int:
+            dim = (2 * i) % d_half
+            return 0 if dim < n_t else 1 if dim < n_t + n_s else 2
+
+        self.rope_pairs = sorted(range(d_half), key=pair_axis)
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MAGEVL)
+        self.gguf_writer.add_vision_attention_layernorm_eps(self.find_vparam(["layer_norm_eps"]))
+        self.gguf_writer.add_vision_spatial_merge_size(self.find_vparam(["spatial_merge_size"]))
+        self.gguf_writer.add_vision_min_pixels(int(self.preprocessor_config["min_pixels"]))
+        self.gguf_writer.add_vision_max_pixels(int(self.preprocessor_config["max_pixels"]))
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, _ = item
+
+        if not name.startswith("visual."):
+            return None
+
+        return super().filter_tensors(item)
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # ggml uses theta = pos / rope_freqs
+        pairs = torch.tensor(self.rope_pairs)
+        freq_a = self.rope_freq[2 * pairs]
+        freq_b = self.rope_freq[2 * pairs + 1]
+        yield ("visual.rope_freqs.weight", (1.0 / torch.cat([freq_a, freq_b])).to(torch.float32))
+
+    def permute_qk(self, data_torch: Tensor) -> Tensor:
+        rows = [2 * i for i in self.rope_pairs] + [2 * i + 1 for i in self.rope_pairs]
+        return data_torch.reshape(self.n_head, self.d_head, -1)[:, rows].reshape(data_torch.shape)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if ".self_attn.qkv." in name:
+            q, k, v = data_torch.split(self.n_embd, dim=0)
+            data_torch = torch.cat([self.permute_qk(q), self.permute_qk(k), v], dim=0)
+
+        yield from super().modify_tensors(data_torch, name, bid)
@@ -161,7 +161,7 @@ class MiniCPMV4_6VisionModel(MmprojModel):
            # In MiniCPM-V 4.6 `vision_config.image_size` (980) describes the SigLIP
            # positional embedding bucket grid (70 x 70), while the per-slice processing
            # resolution is the preprocessor's `scale_resolution` (typically 448).
-            # The CLIP loader in tools/mtmd/clip.cpp consumes `clip.vision.image_size`
+            # The mmproj loader in tools/mtmd/mtmd-model.cpp consumes `clip.vision.image_size`
            # as the slice size and warmup resolution, so report `scale_resolution` there
            # to match the upstream MiniCPMV4_6ImageProcessorPil slicing rules.
            scale_resolution = self.preprocessor_config.get(
@@ -173,7 +173,7 @@ class MiniCPMV4_6VisionModel(MmprojModel):
        super().set_gguf_parameters()
        assert self.hparams_vision is not None

-        # projector type string is consumed by clip_projector_type_from_string() in clip.cpp
+        # projector type string is consumed by projector_type_from_string() in tools/mtmd/mtmd-impl.h
        self.gguf_writer.add_clip_projector_type(self.projector_type)

        self.gguf_writer.add_vision_projector_scale_factor(
@@ -243,7 +243,7 @@ class MiniCPMV4_7VisionModel(MiniCPMV4_6VisionModel):
    projector_type = gguf.VisionProjectorType.MINICPMV4_7
    # MiniCPMV4_7ImageProcessorPil default
    default_scale_resolution = 448
-    # rows of v.tok_embd_sep, the order must match clip_suffix_rows() in clip-impl.h
+    # rows of v.tok_embd_sep, the order must match mtmd_suffix_rows() in tools/mtmd/mtmd-impl.h
    tok_embd_sep = ["</image>", "<slice>", "</slice>", "\n"]

    def get_downsample_mode(self) -> str:
@@ -208,7 +208,7 @@ class NemotronNanoV2VLModel(MmprojModel):
        return name, gen

    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # RADIO's pos_embed doesn't have .weight suffix, but clip.cpp expects it
+        # RADIO's pos_embed doesn't have .weight suffix, but mtmd expects it
        if "patch_generator.pos_embed" in name:
            # Downsample position embeddings for fixed 512x512 image size
            import torch.nn.functional as F
@@ -302,7 +302,7 @@ class PocketTTSMmprojModel(MmprojModel):
        # voice-prompt encoder: mimi encoder + speaker_proj
        if self.has_audio_encoder:
            self.gguf_writer.add_clip_has_audio_encoder(True)
-            # note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models
+            # note: the 24kHz sample rate is hardcoded on the mtmd side, like the other audio models
            self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC)
            self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
            self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"])
@@ -314,7 +314,7 @@ class PocketTTSMmprojModel(MmprojModel):
            self.gguf_writer.add_audio_num_mel_bins(1)

        # generation: flow-matching decoder + mimi decoder
-        # the SEANet and flow net hparams are constant across the family, clip.cpp holds them
+        # the SEANet and flow net hparams are constant across the family, mtmd holds them
        self.gguf_writer.add_clip_has_gen_audio_encoder(True)
        self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_GEN)
        self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text)
@@ -170,7 +170,7 @@ class Qwen2MoeModel(TextModel):
                raise ValueError(f"Unprocessed experts: {experts}")


-@ModelBase.register("Qwen3ForCausalLM", "Qwen3Model")
+@ModelBase.register("Qwen3ForCausalLM", "Qwen3Model", "LlavaOnevision2ForConditionalGeneration")
@ModelBase.example("Qwen/Qwen3-8B")
 class Qwen3Model(Qwen2Model):
    model_arch = gguf.MODEL_ARCH.QWEN3
@@ -235,11 +235,11 @@ class Qwen3TTSSpeakerEncoderModel(MmprojModel):

        # handle speaker encoder config
        self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
-        # mel_spectrogram() front-end: sr=24000, n_fft=1024, hop=256, n_mels=128, fmin=0, fmax=12000 (=sr/2, the clip.cpp default)
+        # mel_spectrogram() front-end: sr=24000, n_fft=1024, hop=256, n_mels=128, fmin=0, fmax=12000 (=sr/2, the mtmd default)
        self.gguf_writer.add_audio_num_mel_bins(128)
        # 3 SE-Res2Net stages; the stem conv, mfa, asp and fc are not counted here
        self.gguf_writer.add_audio_block_count(3)
-        # ECAPA-TDNN has no attention/FFN, these are dummy to allow clip.cpp to load it
+        # ECAPA-TDNN has no attention/FFN, these are dummy to allow mtmd to load it
        self.gguf_writer.add_audio_embedding_length(1536)
        self.gguf_writer.add_audio_head_count(1)
        self.gguf_writer.add_audio_feed_forward_length(1536)
@@ -256,7 +256,7 @@ class Qwen3TTSSpeakerEncoderModel(MmprojModel):
        self.gguf_writer.add_gen_audio_head_count(code_predictor_config["num_attention_heads"])
        self.gguf_writer.add_gen_audio_head_count_kv(code_predictor_config["num_key_value_heads"])
        self.gguf_writer.add_gen_audio_attention_layernorm_eps(code_predictor_config["rms_norm_eps"])
-        # note: code2wav hparams are hardcoded on the mtmd/clip.cpp side for now, not written here
+        # note: code2wav hparams are hardcoded on the mtmd side for now, not written here

    def _wav_decoder_config(self) -> dict[str, Any] | None:
        # code2wav has its own config.json, inside the speech_tokenizer dir
@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import math
+
+from typing import Iterable, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import ModelBase, TextModel, gguf, logger
+
+
+@ModelBase.register("RicursiveForCausalLM")
+@ModelBase.example("ricursiveAI/Ripple-65B-Preview")
+class RicursiveModel(TextModel):
+    model_arch = gguf.MODEL_ARCH.RICURSIVE
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        hparams = self.hparams
+
+        if (act := hparams.get("hidden_act", "silu")) != "silu":
+            raise ValueError(f"Ricursive: unsupported hidden_act {act!r}")
+
+        # RicursiveConfig derives the layer pattern from layer_types alone, by default full attention iff i % 5 == 4
+        if (layer_types := hparams.get("layer_types")) is None:
+            layer_types = ["full_attention" if i % 5 == 4 else "sliding_attention" for i in range(self.block_count)]
+        if len(layer_types) != self.block_count:
+            raise ValueError(f"Ricursive: layer_types has {len(layer_types)} entries, expected {self.block_count}")
+        if unknown := set(layer_types) - {"sliding_attention", "full_attention"}:
+            raise ValueError(f"Ricursive: unsupported layer_types {sorted(unknown)}")
+        self.is_swa = [t == "sliding_attention" for t in layer_types]
+        # only sliding layers use RoPE, an all-full-attention model would have no positional encoding
+        if not any(self.is_swa):
+            raise ValueError("Ricursive: layer_types has no sliding_attention layer, all-full-attention models are not supported")
+        # RicursiveConfig pops a stored hybrid_layer_pattern and recomputes it from layer_types
+        if (pattern := hparams.get("hybrid_layer_pattern")) is not None and list(pattern) != [int(swa) for swa in self.is_swa]:
+            logger.warning("Ricursive: hybrid_layer_pattern does not match layer_types, ignoring it")
+
+        if not hparams.get("gate_full_precision", True):
+            logger.warning("Ricursive: gate_full_precision is false, llama.cpp computes the attention gate in f32")
+
+        # RicursiveConfig defaults; the reference ignores the tensors of a disabled feature
+        self.use_valres = hparams.get("use_valres", False)
+        self.use_gate = hparams.get("use_gated_attention", True)
+        names = list(self.model_tensors)
+        if bias := [name for name in names if name.endswith("bias")]:
+            raise ValueError(f"Ricursive: bias tensors are not supported, found {bias[0]!r}")
+        for enabled, suffix in ((self.use_valres, ".self_attn.value_res_lambda"), (self.use_gate, ".self_attn.g_proj.weight")):
+            if enabled and sum(name.endswith(suffix) for name in names) != self.block_count:
+                raise ValueError(f"Ricursive: expected one {suffix!r} tensor per layer")
+
+    def set_vocab(self):
+        self._set_vocab_gpt2()
+        # eos_token_id is [<|im_end|>, <|endoftext|>]: register the turn end as EOT so generation stops on it
+        eos = self.hparams.get("eos_token_id")
+        if isinstance(eos, list):
+            extra = [e for e in eos if e != self.hparams.get("bos_token_id")]
+            if extra:
+                self.gguf_writer.add_eot_token_id(extra[0])
+                logger.info(f"gguf: eot_token_id = {extra[0]}")
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        hparams = self.hparams
+        self.gguf_writer.add_vocab_size(hparams["vocab_size"])
+        self.gguf_writer.add_rope_dimension_count(hparams["head_dim"])
+
+        self.gguf_writer.add_sliding_window(hparams["sliding_window"])
+        self.gguf_writer.add_sliding_window_pattern(self.is_swa)
+        # sliding layers use RoPE, full attention layers use no positional encoding
+        self.gguf_writer.add_rope_pattern(self.is_swa)
+
+        if (embed_scale := hparams.get("embed_scale")) is not None:
+            if embed_scale == 0:
+                raise ValueError(f"Ricursive: unsupported embed_scale {embed_scale!r}")
+            self.gguf_writer.add_embedding_scale(float(embed_scale))
+        # absent and 0 both mean no clamp (the reference tests `if clamp_threshold:`)
+        if (clamp := hparams.get("ffn_clamp_threshold")) is not None:
+            if clamp < 0 or not math.isfinite(clamp):
+                raise ValueError(f"Ricursive: unsupported ffn_clamp_threshold {clamp!r}")
+            self.gguf_writer.add_swiglu_clamp([float(clamp)] * self.block_count)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # value_res_lambda has no .weight suffix in the checkpoint
+        if name.endswith(".value_res_lambda"):
+            if not self.use_valres:
+                return
+            name += ".weight"
+        elif name.endswith(".g_proj.weight") and not self.use_gate:
+            return
+        yield from super().modify_tensors(data_torch, name, bid)
@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import re
+from typing import Callable, Iterable, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import ModelBase, gguf
+from .deepseek import DeepseekV2Model
+
+
+@ModelBase.register("Xing4_0ForCausalLM")
+class Xing4_0Model(DeepseekV2Model):
+    """Xing4_0 : DeepSeek-V2/V3 backbone (MLA + MoE) + mHC residual mixing."""
+
+    model_arch = gguf.MODEL_ARCH.XING4_0
+
+    skip_mtp = False
+    supports_mtp_export = True
+    _n_main_layers: int | None = None
+
+    # map (prefix, kind) -> MODEL_TENSOR enum
+    _hc_tensor_map = {
+        ("hc_attn", "fn"):    gguf.MODEL_TENSOR.HC_ATTN_FN,
+        ("hc_attn", "base"):  gguf.MODEL_TENSOR.HC_ATTN_BASE,
+        ("hc_attn", "scale"): gguf.MODEL_TENSOR.HC_ATTN_SCALE,
+        ("hc_ffn",  "fn"):    gguf.MODEL_TENSOR.HC_FFN_FN,
+        ("hc_ffn",  "base"):  gguf.MODEL_TENSOR.HC_FFN_BASE,
+        ("hc_ffn",  "scale"): gguf.MODEL_TENSOR.HC_FFN_SCALE,
+    }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.block_count = self.hparams["num_hidden_layers"]
+        if not self.no_mtp:
+            self.block_count += self.hparams.get("num_nextn_predict_layers", 0)
+        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+    def index_tensors(self, remote_hf_model_id: str | None = None):
+        # needed by filter_tensors() before any tensor is filtered
+        type(self)._n_main_layers = self.hparams["num_hidden_layers"]
+        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        if (titem := super().filter_tensors(item)) is None:
+            return None
+        name, gen = titem
+
+        # the NextN/MTP block lives past num_hidden_layers (model.layers.40 -> blk.40)
+        assert cls._n_main_layers is not None
+        is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers
+
+        # --no-mtp: drop the appended NextN block entirely.
+        if is_mtp and cls.no_mtp:
+            return None
+        # --mtp: keep ONLY NextN-block tensors plus the shared embeddings/
+        # norm/lm_head (so the resulting GGUF carries just the draft head).
+        if cls.mtp_only and not is_mtp and name not in (
+            "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
+        ):
+            return None
+
+        return name, gen
+
+    def set_gguf_parameters(self):
+        rope_type = self.rope_parameters.get("rope_type") or self.rope_parameters.get("type")
+        if rope_type == "rope":
+            self.rope_parameters["rope_type"] = "yarn"
+            if "type" in self.rope_parameters:
+                self.rope_parameters["type"] = "yarn"
+            # base.py's YARN branch accesses this with a direct key lookup
+            self.rope_parameters.setdefault("original_max_position_embeddings",
+                                            self.hparams.get("original_max_position_embeddings", 4096))
+
+        super().set_gguf_parameters()
+        hparams = self.hparams
+
+        # NextN/MTP prediction layers
+        if not self.no_mtp and (num_nextn_predict_layers := hparams.get("num_nextn_predict_layers")) is not None:
+            self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
+
+        self.gguf_writer.add_hyper_connection_count(
+            hparams.get("hc_mult",  1))
+        self.gguf_writer.add_hyper_connection_sinkhorn_iterations(
+            hparams.get("hc_sinkhorn_iters",  20))
+        self.gguf_writer.add_hyper_connection_epsilon(
+            hparams.get("hc_eps",  1e-6))
+
+    def set_vocab(self):
+        # Xing4_0 uses a SentencePiece tokenizer (tokenizer.model + Xing4_0Tokenizer).
+        # V2's set_vocab tries GPT2/BPE first, which fails for SPM. Use the SPM path directly.
+        # Note: no add_tokenizer_pre override — the C++ SPM load path ignores
+        # tokenizer.ggml.pre entirely (pre-tokenizers only apply to BPE vocabs).
+        self._set_vocab_sentencepiece()
+
+    def prepare_metadata(self, vocab_only: bool):
+        from_dir = self.fname_out.is_dir()
+        super().prepare_metadata(vocab_only=vocab_only)
+
+        if not self.mtp_only or not from_dir:
+            return
+
+        output_type: str = self.ftype.name.partition("_")[2]
+        fname_default: str = gguf.naming_convention(
+            self.metadata.name, self.metadata.basename, self.metadata.finetune,
+            self.metadata.version, size_label=None, output_type=output_type, model_type=None)
+        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # handle mHC tensors: model.layers.{N}.{attn_hc|ffn_hc}.{hc_fn|hc_base|hc_scale}
+        match = re.match(r"model\.layers\.(\d+)\.(attn_hc|ffn_hc)\.(.+)$", name)
+        if match:
+            layer_idx = int(match.group(1))
+            hc_type = match.group(2)  # "attn_hc" or "ffn_hc"
+            param = match.group(3)    # "hc_fn", "hc_base", "hc_scale"
+
+            if bid is None:
+                bid = layer_idx
+
+            prefix = "hc_attn" if hc_type == "attn_hc" else "hc_ffn"
+
+            kind = {"hc_fn": "fn", "hc_base": "base", "hc_scale": "scale"}.get(param)
+            if kind is None:
+                raise ValueError(f"unexpected mHC tensor: {name}")
+            tensor_enum = self._hc_tensor_map[(prefix, kind)]
+            gguf_name = self.format_tensor_name(tensor_enum, bid)
+            yield (gguf_name, data_torch)
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)
@@ -817,7 +817,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_ENABLE_VMM | 0 or 1 (default) | Enable the virtual-memory device pool. |
 | GGML_SYCL_ENABLE_MKL_FA | 1 (default) or 0 | Enable oneMKL GEMM flash attention for XMX-accelerated prompt processing with quantized KV cache. Automatically activates during prefill (prompt processing) when all conditions are met: (1) flash-attn enabled (`-fa` or `--flash-attn on`), (2) KV cache quantized (`--cache-type-k q8_0 --cache-type-v q8_0` or other `*_0/*_1` types), (3) batch size ≥ 1024 (`--batch-size 1024`), (4) prompt length ≥ 1024 tokens. Set to 0 to force the TILE kernel for A/B testing. Example minimum command: `llama-cli -m model.gguf -fa -ngl 99 --cache-type-k q8_0 --cache-type-v q8_0 --batch-size 1024 -p "your prompt"` |
 | GGML_SYCL_MKL_FA_DEBUG | 0 (default) or 1 | Enable per-call diagnostic logging for MKL flash attention: GEMM/softmax timings, interleaved-head detection, and buffer memory usage. |
-| GGML_SYCL_MEMTRACE | 0 (default), 1, 2 | Enable record and output memory allocation diagnostics. Requires `-lv 4`. <br>0 -  Disable<br>1 - Basic memory info, including current and peak allocations, as well allocations from other sources, around 50 lines per model load.<br>2 - More verbose, logging around 900 specific allocations and deallocations. |
+| GGML_SYCL_MEMTRACE | 0 (default), 1, 2 | Enable record and output memory allocation diagnostics. Requires `-lv 4`. <br>0 -  Disable<br>1 - Basic memory info, including current and peak allocations, as well allocations from other sources, and the device free memory before each model buffer allocation, around 50 lines per model load.<br>2 - More verbose, logging around 900 specific allocations and deallocations. |
 | GGML_SYCL_MEMTRACE_STEP | 64 (default) or positive integer | With GGML_SYCL_MEMTRACE=1, the minimum growth in memory usage to trigger another log record. |
 | GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. |
 | GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. Unsupported types and layouts fall back to the standalone op kernels. See `ggml_sycl_can_fuse()`. |
@@ -836,6 +836,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | ZES_ENABLE_SYSMAN | 0 (default) or 1 | Support to get free memory of GPU by sycl::aspect::ext_intel_free_memory.<br>Recommended to use when --split-mode = layer |
 | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. |
 | GGML_SYCL_USM_SYSTEM | 0 (default) or 1 | Enable experimental support for [USM system allocations](https://github.khronos.org/SYCL_Reference/iface/usm_basic_concept.html#system-allocations) for large GPU buffers. This requires enough host memory for model weights and caches, an Intel Xe2+ GPU such as BMG or newer and supported on Linux only, with CONFIG_DRM_XE_GPUSVM enabled. |
+| GGML_SYCL_BMG_VRAM_CHECK | 1 (default) or 0 | Workaround for BMG (Arc B-series) GPUs, where the Level Zero driver accepts device allocations beyond VRAM and moves the overflow to system RAM, which can exhaust host memory when a model does not fit (e.g. `-ngl 999`, see https://github.com/intel/compute-runtime/issues/1047). When enabled, a model or split buffer that does not fit in the device's free memory is refused and the model load fails with an error. Set to 0 to disable. |

 1. <a id="sycl-graph-footnote"></a>SYCL Graph support (`GGML_SYCL_ENABLE_GRAPH=1`) is limited to single-GPU workloads and may also depend on the model architecture and kernels in use.
   Performance depends on how often a recorded graph can be reused, which will vary by workload.
@@ -325,10 +325,14 @@ on 4 physical NPUs, or `--devices 'HTP0[0-1:0],HTP1[0-1:1]'` on 2 physical NPUs
  Controls the number of HVX hardware threads to use. The default is all (actual number varies depending on the hardware version).

 - `GGML_HEXAGON_VTCM_SIZE=0`
-  Overrides the VTCM size used by each session. Accepts `K`, `M` and `G` suffixes; a plain number is in MB. The default (0) uses the size the hardware reports. Values above the hardware size will fail to start the session.
+  Overrides the VTCM size used by each session. Accepts `K`, `M` and `G` suffixes; a plain number is in MB.
+  The default (0) uses the size the hardware reports. Values above the hardware size will fail to start the session.

- `GGML_HEXAGON_HOSTBUF=1` (default: 0, disabled)
-  Enables allocating host buffers for debugging. By default, host buffers are disabled.
+- `GGML_HEXAGON_HOSTBUF=0` (default: 1, enabled)
+  Controls host buffer (`HTP*-HOST`) allocation in shared FastRPC memory to avoid host-device tensor copies.
+  CPU-assigned weights (input embeddings like `token_embd` and `per_layer_token_embd`, and tensors overridden with `-ot ...=CPU`)
+  are placed in host buffers as well. On older drivers/FW, large models can run out of IOVA space when these are mapped.
+  Pass `--no-host` to keep CPU-assigned weights in regular CPU memory (`-ot ...=CPU` alone is not sufficient).

 - `GGML_HEXAGON_DMA64=0` (default: enabled on v81+)
  Disables 64-bit DMA for model weights. Set to `1` to enable it explicitly on a supported architecture.
@@ -255,7 +255,7 @@ Writing high-performance operators for Hexagon requires following specific guide
 ### 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)).
+  ([`htp-trace.h`](../../../ggml/src/ggml-hexagon/htp/htp-trace.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:
@@ -128,7 +128,7 @@ Note: to debug the inference graph: you can use [llama-eval-callback](/examples/
 If the new model supports multimodal inputs, you will need to add a new encoder definition in `libmtmd`. You can find more information about llama.cpp's multimodal support in [the docs](../multimodal.md) and in the `tools/mtmd` source directory.

 1. In the conversion script, make sure you add a subclass that extends `MmprojModel` or another class that inherits from the same base class.
-2. Add the encoder definition in `clip.cpp`.
+2. Add the encoder definition: hparams and tensor loading in `mtmd-model.cpp`, graph inputs in `mtmd-subcontext.cpp`.
 3. Implement the preprocessor in `mtmd.cpp`. In most cases, you can reuse an existing preprocessor.
 4. Implement the encoder GGML graph, either in a dedicated file if the model is truly different from existing ones, or by reusing an existing implementation (for example: siglip, pixtral, or qwen) and adding a model-specific projector.

@@ -0,0 +1,26 @@
+# MoE cache
+
+When the experts of MoE layers stay in host memory (`-cmoe`, `-ncmoe`, `-ot`), the MoE cache keeps the most used ones in device memory for small batches (decode, speculative decoding). Large batches still upload the experts they use.
+
+| Argument | |
+|---|---|
+| `--moe-cache-mib N` | cache size in MiB, split among the GPUs like the layers (default: 0, disabled) |
+| `--moe-cache-cpu-share F` | share of the cache misses computed on the CPU from the host experts instead of uploaded, once the cache is full (default: 0) |
+
+The cache hits run on the device. With `--moe-cache-cpu-share F` the CPU computes a share `F` of the misses while the device computes the hits and uploads the other misses into the cache. This is faster when the CPU reads the host experts faster than the device can upload them, for example with fast RAM, a PCIe 4 link or experts that are not in pinned memory. The host experts are pinned with `-lm none`, `-lm mlock` and `-lm dio`, unless `--no-host`, but not when they are mmap'd (the default). The device uploads pinned experts at the full speed of the link, so the CPU share often does not help then. Until the cache is full every miss is uploaded, so the cache fills as fast as without the CPU share. The best share depends on the machine, the model and the cache size.
+
+## Picking the arguments
+
+Put the script in front of your server command:
+
+```bash
+python3 scripts/moe-cache-tune.py -- build/bin/llama-server -m model.gguf -cmoe -t 8 -c 32768
+```
+
+It starts the server 6 times, at two cache sizes × shares 0, 0.5 and 0.95, and times the generation of a few thousand tokens. That takes a few minutes, mostly model loading.
+
+From these runs it fits the upload and CPU speeds of the machine, then replays the routing of all the runs through a copy of the cache logic for other sizes and shares. It prints the measured speeds, the predicted t/s and the recommended arguments: the lowest share within 5% of the fastest.
+
+- **Server command:** keep your own arguments (`-t`, `-c`, expert placement, load mode); the speeds and the free device memory depend on them.
+- **`--max-mib N`:** tunes a smaller cache than what fits on the device. The runs in `--out` are reused while the command stays the same, so trying another value is instant.
+- **`--markdown`:** prints Markdown tables to paste in a PR.
@@ -34,7 +34,7 @@ Legend:
 |                CONV_TRANSPOSE_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                CONV_TRANSPOSE_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              COS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                      COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                      COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              CPY | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
 |               CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -62,7 +62,7 @@ Legend:
 |                       GELU_QUICK | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         GET_ROWS | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ❌ | ❌ |
 |                    GET_ROWS_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ |
-|                       GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                       GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -91,7 +91,7 @@ Legend:
 |                          POOL_2D | ❌ | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             RELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
+|                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                      REPEAT_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                         RMS_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    RMS_NORM_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -7682,19 +7682,19 @@
 "OpenCL","REPEAT","type=f32,ne=[10,5,4,1],nr=[1,2,1,1]","support","1","yes","OpenCL"
 "OpenCL","REPEAT","type=f32,ne=[10,5,4,1],nr=[1,1,2,1]","support","1","yes","OpenCL"
 "OpenCL","REPEAT","type=f32,ne=[10,5,4,1],nr=[1,1,1,2]","support","1","yes","OpenCL"
-"OpenCL","REPEAT","type=f16,ne=[10,5,4,1],nr=[2,1,1,1]","support","0","no","OpenCL"
-"OpenCL","REPEAT","type=i32,ne=[10,5,4,1],nr=[2,1,1,1]","support","0","no","OpenCL"
-"OpenCL","REPEAT","type=i16,ne=[10,5,4,1],nr=[1,1,1,2]","support","0","no","OpenCL"
-"OpenCL","REPEAT","type=bf16,ne=[10,5,4,1],nr=[2,1,1,1]","support","0","no","OpenCL"
+"OpenCL","REPEAT","type=f16,ne=[10,5,4,1],nr=[2,1,1,1]","support","1","yes","OpenCL"
+"OpenCL","REPEAT","type=i32,ne=[10,5,4,1],nr=[2,1,1,1]","support","1","yes","OpenCL"
+"OpenCL","REPEAT","type=i16,ne=[10,5,4,1],nr=[1,1,1,2]","support","1","yes","OpenCL"
+"OpenCL","REPEAT","type=bf16,ne=[10,5,4,1],nr=[2,1,1,1]","support","1","yes","OpenCL"
 "OpenCL","REPEAT","type=f32,ne=[10,5,4,3],nr=[1,1,1,1]","support","1","yes","OpenCL"
 "OpenCL","REPEAT","type=f32,ne=[10,5,4,3],nr=[2,1,1,1]","support","1","yes","OpenCL"
 "OpenCL","REPEAT","type=f32,ne=[10,5,4,3],nr=[1,2,1,1]","support","1","yes","OpenCL"
 "OpenCL","REPEAT","type=f32,ne=[10,5,4,3],nr=[1,1,2,1]","support","1","yes","OpenCL"
 "OpenCL","REPEAT","type=f32,ne=[10,5,4,3],nr=[1,1,1,2]","support","1","yes","OpenCL"
-"OpenCL","REPEAT","type=f16,ne=[10,5,4,3],nr=[2,1,1,1]","support","0","no","OpenCL"
-"OpenCL","REPEAT","type=i32,ne=[10,5,4,3],nr=[2,1,1,1]","support","0","no","OpenCL"
-"OpenCL","REPEAT","type=i16,ne=[10,5,4,3],nr=[1,1,1,2]","support","0","no","OpenCL"
-"OpenCL","REPEAT","type=bf16,ne=[10,5,4,3],nr=[2,1,1,1]","support","0","no","OpenCL"
+"OpenCL","REPEAT","type=f16,ne=[10,5,4,3],nr=[2,1,1,1]","support","1","yes","OpenCL"
+"OpenCL","REPEAT","type=i32,ne=[10,5,4,3],nr=[2,1,1,1]","support","1","yes","OpenCL"
+"OpenCL","REPEAT","type=i16,ne=[10,5,4,3],nr=[1,1,1,2]","support","1","yes","OpenCL"
+"OpenCL","REPEAT","type=bf16,ne=[10,5,4,3],nr=[2,1,1,1]","support","1","yes","OpenCL"
 "OpenCL","REPEAT_BACK","type=f32,ne=[8,6,4,2],nr=[1,1,1,1],v=0","support","0","no","OpenCL"
 "OpenCL","REPEAT_BACK","type=f32,ne=[8,6,4,2],nr=[2,1,1,1],v=0","support","0","no","OpenCL"
 "OpenCL","REPEAT_BACK","type=f32,ne=[8,6,4,2],nr=[1,2,1,1],v=0","support","0","no","OpenCL"
@@ -5465,8 +5465,10 @@
 "WebGPU: WebGPU","CONV_TRANSPOSE_2D","kernel_type=f16,ne_input=[3,2,3,1],ne_kernel=[2,2,1,3],stride=1","support","0","no","WebGPU"
 "WebGPU: WebGPU","CONV_TRANSPOSE_2D","kernel_type=f16,ne_input=[10,10,9,1],ne_kernel=[3,3,1,9],stride=2","support","0","no","WebGPU"
 "WebGPU: WebGPU","CONV_TRANSPOSE_2D","kernel_type=f16,ne_input=[129,63,35,1],ne_kernel=[3,3,48,35],stride=1","support","0","no","WebGPU"
-"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,500,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,5000,1,1]","support","0","no","WebGPU"
+"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,500,1,1],same_src=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,5000,1,1],same_src=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,500,1,1],same_src=1","support","1","yes","WebGPU"
+"WebGPU: WebGPU","COUNT_EQUAL","type=f32,ne=[4,5000,1,1],same_src=1","support","1","yes","WebGPU"
 "WebGPU: WebGPU","ARGMAX","type=f32,ne=[32,1,1,1]","support","1","yes","WebGPU"
 "WebGPU: WebGPU","ARGMAX","type=f32,ne=[32,513,1,1]","support","1","yes","WebGPU"
 "WebGPU: WebGPU","ARGMAX","type=f32,ne=[100,10,1,1]","support","1","yes","WebGPU"
@@ -12309,8 +12311,9 @@
 "WebGPU: WebGPU","SUM_ROWS","type=f32,ne=[33,1,1,1],permute=0,slice=0","support","1","yes","WebGPU"
 "WebGPU: WebGPU","SUM_ROWS","type=f32,ne=[33,1024,1,1],permute=0,slice=0","support","1","yes","WebGPU"
 "WebGPU: WebGPU","SUM_ROWS","type=f32,ne=[33,256,1,1],permute=0,slice=0","support","1","yes","WebGPU"
-"WebGPU: WebGPU","GROUP_NORM","type=f32,ne=[64,64,320,1],num_groups=32,eps=0.000001","support","0","no","WebGPU"
-"WebGPU: WebGPU","GROUP_NORM","type=f32,ne=[9,9,1280,1],num_groups=32,eps=0.000001","support","0","no","WebGPU"
+"WebGPU: WebGPU","GROUP_NORM","type=f32,ne=[64,64,320,1],num_groups=32,eps=0.000001","support","1","yes","WebGPU"
+"WebGPU: WebGPU","GROUP_NORM","type=f32,ne=[9,9,1280,1],num_groups=32,eps=0.000001","support","1","yes","WebGPU"
+"WebGPU: WebGPU","GROUP_NORM","type=f32,ne=[3,1,1,65537],num_groups=1,eps=0.000001","support","1","yes","WebGPU"
 "WebGPU: WebGPU","ACC","type=f32,ne_a=[256,17,1,1],ne_b=[256,16,1,1],stride_dim=-1","support","0","no","WebGPU"
 "WebGPU: WebGPU","ACC","type=f32,ne_a=[256,17,2,3],ne_b=[256,16,2,3],stride_dim=-1","support","0","no","WebGPU"
 "WebGPU: WebGPU","ACC","type=f32,ne_a=[256,17,2,3],ne_b=[128,16,2,3],stride_dim=-1","support","0","no","WebGPU"
@@ -103,6 +103,9 @@ int main(int argc, char ** argv) {
    params.cpuparams_batch.n_threads = params.speculative.draft.cpuparams_batch.n_threads;
    params.tensor_buft_overrides     = params.speculative.draft.tensor_buft_overrides;

+    // the MoE cache is only used by the target context
+    params.moe_cache_size = 0;
+
    auto llama_init_dft = common_init_from_params(params);

    model_dft = llama_init_dft->model();
@@ -6,9 +6,9 @@
 extern "C" {
 #endif

-#define RPC_PROTO_MAJOR_VERSION    8
+#define RPC_PROTO_MAJOR_VERSION    9
 #define RPC_PROTO_MINOR_VERSION    0
-#define RPC_PROTO_PATCH_VERSION    1
+#define RPC_PROTO_PATCH_VERSION    0

 #ifdef  __cplusplus
 static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
@@ -1519,6 +1519,7 @@ extern "C" {
    // note: this op has undefined behaviour if a given `a` is referenced more than once for a given `b`
    //       (i.e. duplicate experts per token are not guaranteed to be supported in all backends)
    // TODO: add backend support: https://github.com/ggml-org/llama.cpp/pull/29998#issuecomment-6153742533
+    // the CPU backend skips negative ids, their rows of the result are zero
    GGML_API struct ggml_tensor * ggml_mul_mat_id(
            struct ggml_context * ctx,
            struct ggml_tensor  * as,
@@ -216,6 +216,12 @@ set_target_properties(ggml-base PROPERTIES
 configure_file(ggml-version.h.in ${CMAKE_CURRENT_BINARY_DIR}/ggml-version.h @ONLY)

 target_include_directories(ggml-base PRIVATE . ${CMAKE_CURRENT_BINARY_DIR})
+
+# icx/icpx default to -fp-model=fast, which contracts and reassociates the reference quantizers, so their output
+# depends on the compiler; keep ggml-base IEEE-exact so the CPU reference matches gcc/clang builds
+if (CMAKE_C_COMPILER_ID STREQUAL "IntelLLVM")
+    target_compile_options(ggml-base PRIVATE -fp-model=precise)
+endif()
 if (GGML_BACKEND_DL)
    target_compile_definitions(ggml-base PUBLIC GGML_BACKEND_DL)
 endif()
@@ -1808,13 +1808,14 @@ static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer_n(ggml_b
    const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft);

    constexpr size_t compute_headroom = 16; // Maximum number of views per statically allocated tensor that can be created between evals.
+    constexpr size_t compute_headroom_min = 256; // A buffer with few tensors can still get many views, e.g. a KV cache shared by many layers.
    const ggml_init_params params_static = {
        /*.mem_size   =*/ n_tensors * ggml_tensor_overhead(),
        /*.mem_buffer =*/ nullptr,
        /*.no_alloc   =*/ true,
    };
    const ggml_init_params params_compute = {
-        /*.mem_size   =*/ compute_headroom * n_tensors * ggml_tensor_overhead(),
+        /*.mem_size   =*/ std::max(compute_headroom * n_tensors, compute_headroom_min) * ggml_tensor_overhead(),
        /*.mem_buffer =*/ nullptr,
        /*.no_alloc   =*/ true,
    };
@@ -1246,6 +1246,14 @@ struct ggml_backend_sched {
    int prev_graph_size;

    int debug_realloc;
+
+    // inputs of the next split read from the backend of the current split before it computes, see ggml_backend_sched_prefetch_inputs
+    ggml_backend_event_t prefetch_events[GGML_SCHED_MAX_BACKENDS]; // created on first use, NULL if the backend has no events
+    bool                 prefetch_tried [GGML_SCHED_MAX_BACKENDS];
+    struct ggml_tensor * prefetched[GGML_SCHED_MAX_SPLIT_INPUTS];
+    int n_prefetched;
+    int prefetch_split;     // the split that reads the prefetched inputs
+    int prefetch_backend_id; // the backend that copies them
 };

 #define hash_id(tensor) ggml_hash_find_or_insert(&sched->hash_set, tensor)
@@ -2086,6 +2094,16 @@ static void ggml_backend_sched_copy_input(ggml_backend_sched_t sched, struct ggm
    ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, input);
    struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy);

+    // read before the previous split, wait for the copy only
+    if (split == &sched->splits[sched->prefetch_split]) {
+        for (int i = 0; i < sched->n_prefetched; i++) {
+            if (sched->prefetched[i] == input) {
+                ggml_backend_event_synchronize(sched->prefetch_events[sched->prefetch_backend_id]);
+                return;
+            }
+        }
+    }
+
    if (input->flags & GGML_TENSOR_FLAG_INPUT) {
        // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done
        if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
@@ -2122,6 +2140,62 @@ static void ggml_backend_sched_copy_input(ggml_backend_sched_t sched, struct ggm
    }
 }

+// the inputs of the next split that are computed on the backend of this split before it, and copied to host memory,
+// are copied before this split computes, so that the next split waits for these copies instead of this split
+// e.g. a CPU split runs while a GPU split computes, if it reads only the results of the earlier splits
+static void ggml_backend_sched_prefetch_inputs(ggml_backend_sched_t sched, int split_id) {
+    sched->n_prefetched = 0;
+    if (split_id + 1 >= sched->n_splits) {
+        return;
+    }
+    struct ggml_backend_sched_split * split = &sched->splits[split_id];
+    struct ggml_backend_sched_split * next  = &sched->splits[split_id + 1];
+    const int backend_id = split->backend_id;
+    ggml_backend_t backend = sched->backends[backend_id];
+    if (backend->iface.get_tensor_async == NULL) {
+        return;
+    }
+
+    for (int i = 0; i < next->n_inputs && sched->n_prefetched < GGML_SCHED_MAX_SPLIT_INPUTS; i++) {
+        struct ggml_tensor * input = next->inputs[i];
+        if ((input->flags & GGML_TENSOR_FLAG_INPUT) || tensor_backend_id(input) != backend_id) {
+            continue;
+        }
+        struct ggml_tensor * input_cpy = tensor_copy(input, next->backend_id, sched->cur_copy);
+        if (input_cpy->buffer == NULL || !ggml_backend_buffer_is_host(input_cpy->buffer)) {
+            continue;
+        }
+
+        // not if this split computes or writes the input
+        const struct ggml_tensor * base = input->view_src ? input->view_src : input;
+        bool written = false;
+        for (int j = 0; j < split->graph.n_nodes && !written; j++) {
+            const struct ggml_tensor * node = split->graph.nodes[j];
+            written = !ggml_is_view_op(node->op) && (node == base || node->view_src == base);
+        }
+        if (written) {
+            continue;
+        }
+
+        if (!sched->prefetch_tried[backend_id]) {
+            sched->prefetch_tried[backend_id]  = true;
+            sched->prefetch_events[backend_id] = ggml_backend_event_new(backend->device);
+        }
+        if (sched->prefetch_events[backend_id] == NULL) {
+            return;
+        }
+
+        ggml_backend_tensor_get_async(backend, input, input_cpy->data, 0, ggml_nbytes(input));
+        sched->prefetched[sched->n_prefetched++] = input;
+    }
+
+    if (sched->n_prefetched > 0) {
+        sched->prefetch_split      = split_id + 1;
+        sched->prefetch_backend_id = backend_id;
+        ggml_backend_event_record(sched->prefetch_events[backend_id], backend);
+    }
+}
+
 static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) {
    GGML_ASSERT(sched);
    struct ggml_backend_sched_split * splits = sched->splits;
@@ -2150,6 +2224,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                ggml_backend_sched_copy_input(sched, split, split->inputs[input_id]);
            }
        }
+        // after the copies above, which may read host memory reused by the copies of the next split
+        // before the weights, whose copy callback may queue large uploads
+        ggml_backend_sched_prefetch_inputs(sched, split_id);
        for (int input_id = 0; input_id < split->n_inputs; input_id++) {
            if (ggml_backend_sched_is_host_weight(split->inputs[input_id])) {
                ggml_backend_sched_copy_input(sched, split, split->inputs[input_id]);
@@ -2286,6 +2363,7 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
        for (int c = 0; c < sched->n_copies; c++) {
            ggml_backend_event_free(sched->events[b][c]);
        }
+        ggml_backend_event_free(sched->prefetch_events[b]);
    }
    ggml_gallocr_free(sched->galloc);
    ggml_free(sched->ctx);
@@ -224,8 +224,8 @@ static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = {
    },
    [GGML_TYPE_F16] = {
        .from_float               = (ggml_from_float_t) ggml_cpu_fp32_to_fp16,
-        .vec_dot                  = (ggml_vec_dot_t) ggml_vec_dot_f16,
-        .vec_dot_type             = GGML_TYPE_F16,
+        .vec_dot                  = (ggml_vec_dot_t) ggml_vec_dot_f16_f32,
+        .vec_dot_type             = GGML_TYPE_F32,
        .nrows                    = 1,
    },
    [GGML_TYPE_Q1_0] = {
@@ -1284,9 +1284,12 @@ void ggml_compute_forward_mul_mat(
    const int ith = params->ith;
    const int nth = params->nth;

-    enum ggml_type           const vec_dot_type         = type_traits_cpu[src0->type].vec_dot_type;
-    ggml_from_float_t        const from_float           = type_traits_cpu[vec_dot_type].from_float;
-    int64_t                  const vec_dot_num_rows     = type_traits_cpu[src0->type].nrows;
+    enum ggml_type    const vec_dot_type     = type_traits_cpu[src0->type].vec_dot_type;
+    ggml_from_float_t const from_float       = type_traits_cpu[vec_dot_type].from_float;
+    int64_t           const vec_dot_num_rows = type_traits_cpu[src0->type].nrows;
+
+    // TODO: a more efficient F16 x F16 path can be used here if ggml_prec_get_src(dst, 1) allows converting src1 to F16
+    //       the code needs to be refactored to enable this (src1 is currently converted based on vec_dot_type alone)

    GGML_ASSERT(ne0 == ne01);
    GGML_ASSERT(ne1 == ne11);
@@ -1509,8 +1512,8 @@ static void ggml_compute_forward_mul_mat_id_one_chunk(

    const enum ggml_type type = src0->type;

-    ggml_vec_dot_t    const vec_dot      = type_traits_cpu[type].vec_dot;
-    enum ggml_type    const vec_dot_type = type_traits_cpu[type].vec_dot_type;
+    ggml_vec_dot_t const vec_dot      = type_traits_cpu[type].vec_dot;
+    enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type;

    const int64_t blck_0 = 16;
    const int64_t blck_1 = 16;
@@ -1577,8 +1580,11 @@ static void ggml_compute_forward_mul_mat_id(

    const bool src1_cont = ggml_is_contiguous(src1);

-    enum ggml_type    const vec_dot_type    = type_traits_cpu[type].vec_dot_type;
-    ggml_from_float_t const from_float      = type_traits_cpu[vec_dot_type].from_float;
+    enum ggml_type    const vec_dot_type = type_traits_cpu[type].vec_dot_type;
+    ggml_from_float_t const from_float   = type_traits_cpu[vec_dot_type].from_float;
+
+    // TODO: a more efficient F16 x F16 path can be used here if ggml_prec_get_src(dst, 1) allows converting src1 to F16
+    //       the code needs to be refactored to enable this (src1 is currently converted based on vec_dot_type alone)

    // we don't support permuted src0 or src1
    GGML_ASSERT(nb00 == ggml_type_size(type));
@@ -2965,8 +2971,8 @@ struct ggml_cplan ggml_graph_plan(
                        if ((node->src[0]->type == GGML_TYPE_F16 ||
                             node->src[0]->type == GGML_TYPE_BF16) &&
                            node->src[1]->type == GGML_TYPE_F32) {
+                            cur += sizeof(float)*ne10*ne11;
                            cur += sizeof(ggml_fp16_t)*ne00*ne01*ne02;
-                            cur += sizeof(ggml_fp16_t)*ne10*ne11;
                        } else if (node->src[0]->type == GGML_TYPE_F32 &&
                                   node->src[1]->type == GGML_TYPE_F32) {
                            cur += sizeof(float)*ne00*ne01*ne02;
@@ -2995,8 +3001,8 @@ struct ggml_cplan ggml_graph_plan(
                        GGML_ASSERT(node->src[0]->type == GGML_TYPE_F16 || node->src[0]->type == GGML_TYPE_F32);
                        GGML_ASSERT(node->src[1]->type == GGML_TYPE_F32);

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

                    } break;
                case GGML_OP_TOP_K:
@@ -455,10 +455,11 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
                src0->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32) {
                return src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16;
            }
-            // BF16 in src1 is widened into the F32 work buffer
+            // F16 or BF16 in src1 is widened into the F32 work buffer
            return src1->type == GGML_TYPE_F32 ||
                   src1->type == ggml_get_type_traits_cpu(src0->type)->vec_dot_type ||
-                   (src1->type == GGML_TYPE_BF16 && ggml_get_type_traits_cpu(src0->type)->vec_dot_type == GGML_TYPE_F32);
+                   ((src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_BF16) &&
+                    ggml_get_type_traits_cpu(src0->type)->vec_dot_type == GGML_TYPE_F32);
        case GGML_OP_SOFT_MAX_BACK: {
            if (op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32) {
                return false;
@@ -326,6 +326,12 @@ template <> inline float32x4_t load(const ggml_fp16_t * p) {

    return vec_xl(0, (const float *)(tmp));
 }
+template <> inline float32x4_t load(const ggml_bf16_t * p) {
+    uint64_t bits;
+    memcpy(&bits, p, sizeof(bits));
+    const uint64x2_t v_bits = { bits, bits };
+    return GGML_BF16_TO_F32_HI((GGML_BF16_VEC) v_bits);
+}
 template <> inline float32x4_t load(const float * p) {
    return vec_xl(0, p);
 }
@@ -1302,6 +1308,88 @@ class tinyBLAS_RVV {
 //////////////////////////////////////////////////////////////////////////////////////////
 // QUANT ZERO MATRIX MULTIPLICATION

+#if defined(__VXE__) || defined(__VXE2__)
+template <typename TA>
+class tinyBLAS_Q0_VXE {
+    static_assert(sizeof(TA::qs) == QK8_0, "VXE weight loads require 32 quantized bytes per block");
+
+  public:
+    tinyBLAS_Q0_VXE(int64_t k,
+                    const TA * A, int64_t lda,
+                    const block_q8_0 * B, int64_t ldb,
+                    float * C, int64_t ldc,
+                    int32_t ith, int32_t nth)
+        : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) {
+    }
+
+    void matmul(int64_t m, int64_t n) {
+        constexpr int32_t tile_size = 2;
+        const int64_t row_tiles = m / tile_size;
+        const int64_t tiles     = row_tiles * (n / tile_size);
+        const int64_t duty      = (tiles + nth - 1) / nth;
+        const int64_t start     = duty * ith;
+        const int64_t end       = MIN(start + duty, tiles);
+
+        for (int64_t job = start; job < end; ++job) {
+            const int64_t ii = tile_size * (job % row_tiles);
+            const int64_t jj = tile_size * (job / row_tiles);
+            float32x4_t v_sum[tile_size][tile_size] = {};
+
+            for (int64_t l = 0; l < k; ++l) {
+                int8x16_t v_xl[tile_size], v_xh[tile_size], v_yl[tile_size], v_yh[tile_size];
+                float x_scale[tile_size], y_scale[tile_size];
+
+                for (int32_t i = 0; i < tile_size; ++i) {
+                    const TA & x = A[lda * (ii + i) + l];
+                    const block_q8_0 & y = B[ldb * (jj + i) + l];
+
+                    const int8_t * x_qs = x.qs;
+                    const int8_t * y_qs = y.qs;
+                    v_xl[i] = vec_xl(0, x_qs);
+                    v_xh[i] = vec_xl(0, x_qs + QK8_0 / 2);
+                    v_yl[i] = vec_xl(0, y_qs);
+                    v_yh[i] = vec_xl(0, y_qs + QK8_0 / 2);
+                    x_scale[i] = unhalf(x.d);
+                    y_scale[i] = unhalf(y.d);
+                }
+
+                for (int32_t j = 0; j < tile_size; ++j) {
+                    for (int32_t i = 0; i < tile_size; ++i) {
+                        const int32x4_t v_dot = dot_s8_wide(v_xl[i], v_yl[j]) + dot_s8_wide(v_xh[i], v_yh[j]);
+                        const float32x4_t v_scale = vec_splats(x_scale[i] * y_scale[j]);
+                        v_sum[j][i] = vec_madd(vec_float(v_dot), v_scale, v_sum[j][i]);
+                    }
+                }
+            }
+
+            for (int32_t j = 0; j < tile_size; ++j) {
+                for (int32_t i = 0; i < tile_size; ++i) {
+                    C[ldc * (jj + j) + ii + i] = hsum(v_sum[j][i]);
+                }
+            }
+        }
+    }
+
+  private:
+    static int32x4_t dot_s8_wide(int8x16_t v_x, int8x16_t v_y) {
+        const int16x8_t v_even = vec_mule(v_x, v_y);
+        const int16x8_t v_odd  = vec_mulo(v_x, v_y);
+        // Two (-128) * (-128) products can overflow a signed 16-bit sum.
+        return vec_unpackh(v_even) + vec_unpackl(v_even) + vec_unpackh(v_odd) + vec_unpackl(v_odd);
+    }
+
+    const TA * const A;
+    const block_q8_0 * const B;
+    float * const C;
+    const int64_t k;
+    const int64_t lda;
+    const int64_t ldb;
+    const int64_t ldc;
+    const int32_t ith;
+    const int32_t nth;
+};
+#endif
+
 #if defined(__ARM_FEATURE_DOTPROD)
 template <typename TA>
 class tinyBLAS_Q0_ARM {
@@ -4006,6 +4094,14 @@ bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64
                (float *)C, ldc};
            return tb.matmul(m, n);
        }
+#elif defined(__VXE__) || defined(__VXE2__)
+        if (Btype == GGML_TYPE_BF16) {
+            tinyBLAS<4, float32x4_t, float32x4_t, ggml_bf16_t, ggml_bf16_t, float> tb{ params, k,
+                (const ggml_bf16_t *)A, lda,
+                (const ggml_bf16_t *)B, ldb,
+                (float *)C, ldc};
+            return tb.matmul(m, n);
+        }
 #elif defined(__MMA__)
        if (k % 8) {
            return false;
@@ -4054,6 +4150,13 @@ bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64
                (float *)C, ldc};
            return tb.matmul(m, n);
        }
+        if (Btype == GGML_TYPE_F32) {
+            tinyBLAS<16, __m512, __m512, ggml_fp16_t, float, float> tb{ params, k,
+                (const ggml_fp16_t *)A, lda,
+                (const float *)B, ldb,
+                (float *)C, ldc};
+            return tb.matmul(m, n);
+        }
 #elif (defined(__AVX__) || defined(__AVX2__)) && defined(__F16C__)
        if (Btype == GGML_TYPE_F16) {
            tinyBLAS<8, __m256, __m256, ggml_fp16_t, ggml_fp16_t, float> tb{ params, k,
@@ -4062,6 +4165,13 @@ bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64
                (float *)C, ldc};
            return tb.matmul(m, n);
        }
+        if (Btype == GGML_TYPE_F32) {
+            tinyBLAS<8, __m256, __m256, ggml_fp16_t, float, float> tb{ params, k,
+                (const ggml_fp16_t *)A, lda,
+                (const float *)B, ldb,
+                (float *)C, ldc};
+            return tb.matmul(m, n);
+        }
 #elif defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER)
        if (n < 8)
            return false;
@@ -4148,6 +4258,17 @@ bool llamafile_sgemm(const struct ggml_compute_params * params, int64_t m, int64
            params->ith, params->nth};
        tb.matmul(m, n);
        return true;
+#elif defined(__VXE__) || defined(__VXE2__)
+        if (m % 2 || n % 2) {
+            return false;
+        }
+        tinyBLAS_Q0_VXE<block_q8_0> tb{
+            k, (const block_q8_0 *)A, lda,
+            (const block_q8_0 *)B, ldb,
+            (float *)C, ldc,
+            params->ith, params->nth};
+        tb.matmul(m, n);
+        return true;
 #elif defined(__MMA__)
    //TO-DO: Remove this condition once gemv forwarding is enabled.
        if (n < 8 && n != 4)
@@ -6514,17 +6514,28 @@ static void ggml_compute_forward_conv_transpose_1d_f16_f32(
    const int ith = params->ith;
    const int nth = params->nth;

-    const int nk = ne00*ne01*ne02;
-
    GGML_ASSERT(nb00 == sizeof(ggml_fp16_t));
    GGML_ASSERT(nb10 == sizeof(float));

    if (ith == 0) {
        memset(params->wdata, 0, params->wsize);

+        // permute source data (src1) from (L x Cin) to (Cin x L)
+        {
+            float * const wdata = (float *) params->wdata + 0;
+            float * dst_data = wdata;
+
+            for (int64_t i11 = 0; i11 < ne11; i11++) {
+                const float * const src = (float *)((char *) src1->data + i11*nb11);
+                for (int64_t i10 = 0; i10 < ne10; i10++) {
+                    dst_data[i10*ne11 + i11] = src[i10];
+                }
+            }
+        }
+
        // permute kernel data (src0) from (K x Cout x Cin) to (Cin x K x Cout)
        {
-            ggml_fp16_t * const wdata = (ggml_fp16_t *) params->wdata + 0;
+            ggml_fp16_t * const wdata = (ggml_fp16_t *)((char *) params->wdata + ne10*ne11*sizeof(float));

            for (int64_t i02 = 0; i02 < ne02; i02++) {
                for (int64_t i01 = 0; i01 < ne01; i01++) {
@@ -6537,19 +6548,6 @@ static void ggml_compute_forward_conv_transpose_1d_f16_f32(
            }
        }

-        // permute source data (src1) from (L x Cin) to (Cin x L)
-        {
-            ggml_fp16_t * const wdata = (ggml_fp16_t *) params->wdata + nk;
-            ggml_fp16_t * dst_data = wdata;
-
-            for (int64_t i11 = 0; i11 < ne11; i11++) {
-                const float * const src = (float *)((char *) src1->data + i11*nb11);
-                for (int64_t i10 = 0; i10 < ne10; i10++) {
-                    dst_data[i10*ne11 + i11] = GGML_CPU_FP32_TO_FP16(src[i10]);
-                }
-            }
-        }
-
        // need to zero dst since we are accumulating into it
        memset(dst->data, 0, ggml_nbytes(dst));
    }
@@ -6567,19 +6565,19 @@ static void ggml_compute_forward_conv_transpose_1d_f16_f32(
    const int ir0 = dr*ith;
    const int ir1 = MIN(ir0 + dr, nr);

-    ggml_fp16_t * const wdata     = (ggml_fp16_t *) params->wdata + 0;
-    ggml_fp16_t * const wdata_src = wdata + nk;
+    float * const wdata_src = (float *) params->wdata + 0;
+    ggml_fp16_t * const wdata_kernel = (ggml_fp16_t *)((char *) params->wdata + ne10*ne11*sizeof(float));

    for (int i1 = ir0; i1 < ir1; i1++) {
        float * dst_data = (float *)((char *) dst->data + i1*nb1);
-        ggml_fp16_t * wdata_kernel = wdata + i1*ne02*ne00;
+        ggml_fp16_t * wdata_kernel_row = wdata_kernel + i1*ne02*ne00;
        for (int i10 = 0; i10 < ne10; i10++) {
            const int i1n = i10*ne11;
            for (int i00 = 0; i00 < ne00; i00++) {
                float v = 0;
-                ggml_vec_dot_f16(ne02, &v, 0,
-                        (ggml_fp16_t *)    wdata_src + i1n, 0,
-                        (ggml_fp16_t *) wdata_kernel + i00*ne02, 0, 1);
+                ggml_vec_dot_f16_f32(ne02, &v, 0,
+                        wdata_kernel_row + i00*ne02, 0,
+                        wdata_src + i1n, 0, 1);
                dst_data[i10*s0 + i00] += v;
            }
        }
@@ -7172,29 +7170,30 @@ void ggml_compute_forward_im2col_3d(
    }
 }

-static void ggml_call_mul_mat(ggml_type type, const ggml_compute_params * params, int64_t m, int64_t n, int64_t k,
+static void ggml_call_mul_mat(ggml_type src0_type, ggml_type src1_type, const ggml_compute_params * params, int64_t m, int64_t n, int64_t k,
                              void * a, void * b, float * c) {
-    const ggml_type_traits * traits = ggml_get_type_traits(type);
+    const ggml_type_traits * traits0 = ggml_get_type_traits(src0_type);
+    const ggml_type_traits * traits1 = ggml_get_type_traits(src1_type);
    struct ggml_tensor src1 = {};
-    src1.type  = type;
+    src1.type  = src1_type;
    src1.ne[0] = k;
    src1.ne[1] = m;
    src1.ne[2] = 1;
    src1.ne[3] = 1;
-    src1.nb[0] = traits->type_size;
-    src1.nb[1] = k * traits->type_size;
+    src1.nb[0] = traits1->type_size;
+    src1.nb[1] = k * traits1->type_size;
    src1.nb[2] = src1.nb[1];
    src1.nb[3] = src1.nb[2];
    src1.data  = a;

    struct ggml_tensor src0 = {};
-    src0.type  = type;
+    src0.type  = src0_type;
    src0.ne[0] = k;
    src0.ne[1] = n;
    src0.ne[2] = 1;
    src0.ne[3] = 1;
-    src0.nb[0] = traits->type_size;
-    src0.nb[1] = k * traits->type_size;
+    src0.nb[0] = traits0->type_size;
+    src0.nb[1] = k * traits0->type_size;
    src0.nb[2] = src0.nb[1];
    src0.nb[3] = src0.nb[2];
    src0.data  = b;
@@ -7304,8 +7303,6 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
    GGML_ASSERT(kernel_type == GGML_TYPE_F16 || kernel_type == GGML_TYPE_F32);
    GGML_ASSERT(kernel->type == kernel_type);

-    const ggml_type_traits * traits = ggml_get_type_traits(kernel_type);
-
    const int32_t stride_x   = dst->op_params[0];
    const int32_t stride_y   = dst->op_params[1];
    const int32_t pad_x      = dst->op_params[2];
@@ -7331,7 +7328,7 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
    const int64_t knl_n           = knl_w * knl_h * c_in;
    const int64_t patch_total     = dst->ne[3] * dst_w * dst_h;

-    const int64_t space_per_patch   = knl_n * traits->type_size + c_out * sizeof(float);
+    const int64_t space_per_patch   = knl_n * sizeof(float) + c_out * sizeof(float);
    const int64_t batch_size        = params->wsize / space_per_patch;
    const int64_t patches_per_batch = batch_size > 8 ? (batch_size / 8) * 8 : batch_size;
    const int64_t batch_n           = (patch_total + patches_per_batch - 1) / patches_per_batch;
@@ -7358,7 +7355,7 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
            const int64_t  src_y       =  p % dst_w;

            const float * src_base = (const float *)((const char *)src_data + batch_n * src->nb[3]);
-            char *        dst_row  = (char *) tmp + (p % patches_per_batch) * knl_n * traits->type_size;
+            float *       dst_row  = (float *) tmp + (p % patches_per_batch) * knl_n;

            for (int64_t ic = 0; ic < c_in; ++ic) {
                for (int64_t ky = 0; ky < knl_h; ++ky) {
@@ -7376,12 +7373,7 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params
                            src_val               = *src_ptr;
                        }

-                        char * element_ptr = dst_row + dst_idx * traits->type_size;
-                        if (kernel_type == GGML_TYPE_F32) {
-                            *(float *) element_ptr = src_val;
-                        } else if (kernel_type == GGML_TYPE_F16) {
-                            *(ggml_fp16_t *) element_ptr = GGML_CPU_FP32_TO_FP16(src_val);
-                        }
+                        dst_row[dst_idx] = src_val;
                    }
                }
            }
@@ -7389,12 +7381,12 @@ static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params

        ggml_barrier(params->threadpool);

-        float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n * traits->type_size);
+        float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n * sizeof(float));

        GGML_ASSERT(gemm_output + patch_n * c_out <= (float*)tmp + params->wsize);

        // GEMM: patches[patch_n, knl_n] × kernel[knl_n, c_out] = output[patch_n, c_out]
-        ggml_call_mul_mat(kernel_type, params, patch_n, c_out, knl_n, tmp, knl_data, gemm_output);
+        ggml_call_mul_mat(kernel_type, GGML_TYPE_F32, params, patch_n, c_out, knl_n, tmp, knl_data, gemm_output);

        ggml_barrier(params->threadpool);

@@ -7441,8 +7433,6 @@ static void ggml_compute_forward_conv_3d_impl(const ggml_compute_params * params
    GGML_ASSERT(kernel_type == GGML_TYPE_F16 || kernel_type == GGML_TYPE_F32);
    GGML_ASSERT(kernel->type == kernel_type);

-    const ggml_type_traits * traits = ggml_get_type_traits(kernel_type);
-
    const int32_t s0 = dst->op_params[0];
    const int32_t s1 = dst->op_params[1];
    const int32_t s2 = dst->op_params[2];
@@ -7474,7 +7464,7 @@ static void ggml_compute_forward_conv_3d_impl(const ggml_compute_params * params
    const int64_t knl_n_total       = knl_n_per_channel * c;
    const int64_t patch_total       = n * dst_w * dst_h * dst_d;

-    const int64_t space_per_patch   = knl_n_total * traits->type_size + oc * sizeof(float);
+    const int64_t space_per_patch   = knl_n_total * sizeof(float) + oc * sizeof(float);
    const int64_t batch_size        = params->wsize / space_per_patch;
    const int64_t patches_per_batch = batch_size > 8 ? (batch_size / 8) * 8 : batch_size;
    const int64_t batch_n           = (patch_total + patches_per_batch - 1) / patches_per_batch;
@@ -7500,7 +7490,7 @@ static void ggml_compute_forward_conv_3d_impl(const ggml_compute_params * params
            const int64_t dst_y      = p_in_depth / dst_w;
            const int64_t dst_x      = p_in_depth % dst_w;

-            char * dst_row = (char *) tmp + (p % patches_per_batch) * knl_n_total * traits->type_size;
+            float * dst_row = (float *) tmp + (p % patches_per_batch) * knl_n_total;

            for (int64_t ic = 0; ic < c; ++ic) {
                for (int64_t kz = 0; kz < knl_d; ++kz) {
@@ -7521,12 +7511,7 @@ static void ggml_compute_forward_conv_3d_impl(const ggml_compute_params * params
                                src_val = *src_ptr;
                            }

-                            char * element_ptr = dst_row + dst_idx * traits->type_size;
-                            if (kernel_type == GGML_TYPE_F32) {
-                                *(float *)element_ptr = src_val;
-                            } else if (kernel_type == GGML_TYPE_F16) {
-                                *(ggml_fp16_t *)element_ptr = GGML_CPU_FP32_TO_FP16(src_val);
-                            }
+                            dst_row[dst_idx] = src_val;
                        }
                    }
                }
@@ -7535,8 +7520,8 @@ static void ggml_compute_forward_conv_3d_impl(const ggml_compute_params * params

        ggml_barrier(params->threadpool);

-        float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n_total * traits->type_size);
-        ggml_call_mul_mat(kernel_type, params, patch_n_in_batch, oc, knl_n_total, tmp, knl_data, gemm_output);
+        float * gemm_output = (float *) ((char *) tmp + patches_per_batch * knl_n_total * sizeof(float));
+        ggml_call_mul_mat(kernel_type, GGML_TYPE_F32, params, patch_n_in_batch, oc, knl_n_total, tmp, knl_data, gemm_output);

        ggml_barrier(params->threadpool);

@@ -7588,7 +7573,7 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
    const int ith = params->ith;
    const int nth = params->nth;

-    const int nk = ne00*ne01*ne02*ne03;
+    const int nsrc = ne10*ne11*ne12*ne13;

    GGML_ASSERT(nb00 == ggml_type_size(src0->type));
    GGML_ASSERT(nb10 == sizeof(float));
@@ -7596,9 +7581,26 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
    if (ith == 0) {
        memset(params->wdata, 0, params->wsize);

+        // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh), for all batches
+        {
+            float * const wdata = (float *) params->wdata + 0;
+            for (int i13 = 0; i13 < ne13; i13++) {
+                float * const wdata_b = wdata + i13*ne10*ne11*ne12;
+                for (int i12 = 0; i12 < ne12; i12++) {
+                    for (int i11 = 0; i11 < ne11; i11++) {
+                        const float * const src = (float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11);
+                        float * dst_data = wdata_b + i11*ne10*ne12;
+                        for (int i10 = 0; i10 < ne10; i10++) {
+                            dst_data[i10*ne12 + i12] = src[i10];
+                        }
+                    }
+                }
+            }
+        }
+
        // permute kernel data (src0) from (Kw x Kh x Cout x Cin) to (Cin x Kw x Kh x Cout)
        {
-            kernel_t * const wdata = (kernel_t *) params->wdata + 0;
+            kernel_t * const wdata = (kernel_t *)((char *) params->wdata + nsrc*sizeof(float));

            for (int64_t i03 = 0; i03 < ne03; i03++) {
                for (int64_t i02 = 0; i02 < ne02; i02++) {
@@ -7613,27 +7615,6 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
            }
        }

-        // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh), for all batches
-        {
-            kernel_t * const wdata = (kernel_t *) params->wdata + nk;
-            for (int i13 = 0; i13 < ne13; i13++) {
-                kernel_t * const wdata_b = wdata + i13*ne10*ne11*ne12;
-                for (int i12 = 0; i12 < ne12; i12++) {
-                    for (int i11 = 0; i11 < ne11; i11++) {
-                        const float * const src = (float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11);
-                        kernel_t * dst_data = wdata_b + i11*ne10*ne12;
-                        for (int i10 = 0; i10 < ne10; i10++) {
-                            if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
-                                dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]);
-                            } else {
-                                dst_data[i10*ne12 + i12] = src[i10];
-                            }
-                        }
-                    }
-                }
-            }
-        }
-
        memset(dst->data, 0, ggml_nbytes(dst));
    }
    ggml_barrier(params->threadpool);
@@ -7650,14 +7631,14 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
    const int ip0 = dp*ith;
    const int ip1 = MIN(ip0 + dp, np);

-    kernel_t * const wdata = (kernel_t *) params->wdata + 0;
-    kernel_t * const wdata_src = wdata + nk;
+    float * const wdata_src = (float *) params->wdata + 0;
+    kernel_t * const wdata_kernel = (kernel_t *)((char *) params->wdata + nsrc*sizeof(float));

    for (int i2 = ip0; i2 < ip1; i2++) { // Cout
-        kernel_t * wdata_kernel = wdata + i2*ne01*ne00*ne03;
+        kernel_t * wdata_kernel_c = wdata_kernel + i2*ne01*ne00*ne03;
        for (int i3 = 0; i3 < ne3; i3++) { // batch
            float * dst_data = (float *)((char *) dst->data + i3*nb3 + i2*nb2);
-            kernel_t * wdata_src_b = wdata_src + i3*ne10*ne11*ne12;
+            float * wdata_src_b = wdata_src + i3*ne10*ne11*ne12;
            for (int i11 = 0; i11 < ne11; i11++) {
                for (int i10 = 0; i10 < ne10; i10++) {
                    const int i1n = i11*ne10*ne12 + i10*ne12;
@@ -7665,13 +7646,13 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
                        for (int i00 = 0; i00 < ne00; i00++) {
                            float v = 0;
                            if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
-                                ggml_vec_dot_f16(ne03, &v, 0,
-                                        wdata_src_b + i1n, 0,
-                                        wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
+                                ggml_vec_dot_f16_f32(ne03, &v, 0,
+                                        wdata_kernel_c + i01*ne00*ne03 + i00*ne03, 0,
+                                        wdata_src_b + i1n, 0, 1);
                            } else {
                                ggml_vec_dot_f32(ne03, &v, 0,
-                                        wdata_src_b + i1n, 0,
-                                        wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
+                                        wdata_kernel_c + i01*ne00*ne03 + i00*ne03, 0,
+                                        wdata_src_b + i1n, 0, 1);
                            }
                            dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v;
                        }
@@ -9591,6 +9572,7 @@ void ggml_compute_forward_flash_attn_ext(
    switch (dst->op_params[3]) {
        case GGML_PREC_DEFAULT:
        case GGML_PREC_F32:
+        case GGML_PREC_BF16:
            {
                // uses F32 accumulators
                ggml_compute_forward_flash_attn_ext_f16(params, dst);
@@ -198,7 +198,7 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
                    case GGML_TYPE_Q5_0:
                    case GGML_TYPE_Q5_1:
                    case GGML_TYPE_Q5_K:
-                        //case GGML_TYPE_MXFP4:
+                    case GGML_TYPE_MXFP4:
                        forward_mul_mat(params, op);
                        return true;
                    default:
@@ -1180,7 +1180,7 @@ class tensor_traits_common : public tensor_traits_base {
        const int64_t DK = nek0;
        const int64_t DV = nev0;

-        const bool supported_prec  = (dst->op_params[3] == GGML_PREC_F32 || dst->op_params[3] == GGML_PREC_DEFAULT);
+        const bool supported_prec  = (dst->op_params[3] == GGML_PREC_F32 || dst->op_params[3] == GGML_PREC_BF16 || dst->op_params[3] == GGML_PREC_DEFAULT);
        const bool supported_types = (q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16);
        const bool supported_shape = (DK > 0 && DK <= 128 && DV > 0 && DV <= 128);
        const bool supported_vlen  = (__riscv_vlenb() == 128);
@@ -1375,10 +1375,10 @@ static const ggml::cpu::tensor_traits * ggml_riscv64_spacemit_get_optimal_repack
        case GGML_TYPE_MXFP4:
            {
 #if defined(RISCV64_SPACEMIT_IME2)
-                // TODO
-                // if (cur->ne[1] % 32 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
-                //     return &ggml::cpu::riscv64_spacemit::mxfp4_32x32_q8_0;
-                // }
+                if (ggml_n_dims(cur) == 2 && cur->ne[0] % 32 == 0 && cur->ne[1] % 32 == 0 &&
+                    (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime2)) {
+                    return &ggml::cpu::riscv64_spacemit::mxfp4_32x32_q8_0;
+                }
 #endif
            }
            break;
@@ -1613,6 +1613,9 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
                    if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) {
                        return false;
                    }
+                    if (op->src[0]->type == GGML_TYPE_MXFP4 && !ggml_is_contiguous(op->src[1])) {
+                        return false;
+                    }
                    if (op->src[1]->type == GGML_TYPE_F32) {
                        return true;
                    }
@@ -3754,7 +3754,7 @@ void gemm_kernel_i8mxfp4_m1(size_t          blk_len,
    using blk_type           = nrow_block_mxfp4<NB_COLS>;

    GGML_ASSERT(blk_len == K_TILE);
-    GGML_ASSERT(count_m == 1);
+    GGML_ASSERT(count_m >= 1);
    GGML_UNUSED(quant_b_zp);

    const size_t a_blk_stride  = q8_blk_size(blk_len, true);
@@ -3923,7 +3923,7 @@ void gemm_kernel_i8mxfp4_m1(size_t          blk_len,
                :
                : [BCK] "r"(cnt), [NBLKS] "r"(nblks), [pA] "r"(quant_a_ptr), [pB] "r"(QuantBDataPtr), [pC] "r"(CPtr)
                : "cc", "memory", "t0", "t1", "t2", "t3", "t4", "f0", "s2", "s3", "s4", "s5", "s6", "s7", "v0", "v1",
-                  "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v16", "v17", "v18", "v19",
+                  "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19",
                  "v20", "v21", "v22", "v23", "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31");
        }
    }
@@ -3943,7 +3943,7 @@ void gemm_kernel_i8mxfp4_m4(size_t          blk_len,
    using blk_type           = nrow_block_mxfp4<NB_COLS>;

    GGML_ASSERT(blk_len == K_TILE);
-    GGML_ASSERT(count_m == 4);
+    GGML_ASSERT(count_m >= 4);
    GGML_UNUSED(quant_b_zp);

    const size_t a_blk_stride  = q8_blk_size(blk_len, true);
@@ -5687,20 +5687,10 @@ size_t gemm_kernel_i8mxfp4(size_t          blk_len,
                           size_t          k_blks,
                           size_t          ldc) {
    if (count_m >= 4) {
-#if 1
-        gemm_kernel_i8mxfp4_mrow_ref<4, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
-                                            k_blks, ldc);
-#else
        gemm_kernel_i8mxfp4_m4(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
-#endif
        return 4;
    } else {
-#if 1
-        gemm_kernel_i8mxfp4_mrow_ref<1, 32>(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n,
-                                            k_blks, ldc);
-#else
        gemm_kernel_i8mxfp4_m1(blk_len, quant_a_ptr, quant_b_data, quant_b_zp, c_ptr, count_m, count_n, k_blks, ldc);
-#endif
        return 1;
    }
 }
@@ -377,6 +377,85 @@ void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * G
    *s = sumf;
 }

+void ggml_vec_dot_f16_f32(int n, float * GGML_RESTRICT s, size_t bs, const ggml_fp16_t * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc) {
+    assert(nrc == 1);
+    GGML_UNUSED(nrc);
+    GGML_UNUSED(bx);
+    GGML_UNUSED(by);
+    GGML_UNUSED(bs);
+
+    ggml_float sumf = 0.0;
+    int i = 0;
+
+    // on AVX512 and AVX with F16C, x is widened to F32 in-register and accumulated in F32 over 4 independent sums
+#if defined(__AVX512F__)
+    __m512 acc0 = _mm512_setzero_ps();
+    __m512 acc1 = _mm512_setzero_ps();
+    __m512 acc2 = _mm512_setzero_ps();
+    __m512 acc3 = _mm512_setzero_ps();
+    for (; i + 64 <= n; i += 64) {
+        acc0 = _mm512_fmadd_ps(_mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x + i +  0))), _mm512_loadu_ps(y + i +  0), acc0);
+        acc1 = _mm512_fmadd_ps(_mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x + i + 16))), _mm512_loadu_ps(y + i + 16), acc1);
+        acc2 = _mm512_fmadd_ps(_mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x + i + 32))), _mm512_loadu_ps(y + i + 32), acc2);
+        acc3 = _mm512_fmadd_ps(_mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x + i + 48))), _mm512_loadu_ps(y + i + 48), acc3);
+    }
+    for (; i + 16 <= n; i += 16) {
+        acc0 = _mm512_fmadd_ps(_mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x + i))), _mm512_loadu_ps(y + i), acc0);
+    }
+    sumf = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3)));
+#elif defined(__AVX__) && defined(__F16C__)
+    __m256 acc0 = _mm256_setzero_ps();
+    __m256 acc1 = _mm256_setzero_ps();
+    __m256 acc2 = _mm256_setzero_ps();
+    __m256 acc3 = _mm256_setzero_ps();
+    for (; i + 32 <= n; i += 32) {
+        acc0 = GGML_F32x8_FMA(acc0, _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i +  0))), _mm256_loadu_ps(y + i +  0));
+        acc1 = GGML_F32x8_FMA(acc1, _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i +  8))), _mm256_loadu_ps(y + i +  8));
+        acc2 = GGML_F32x8_FMA(acc2, _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i + 16))), _mm256_loadu_ps(y + i + 16));
+        acc3 = GGML_F32x8_FMA(acc3, _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i + 24))), _mm256_loadu_ps(y + i + 24));
+    }
+    for (; i + 8 <= n; i += 8) {
+        acc0 = GGML_F32x8_FMA(acc0, _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i))), _mm256_loadu_ps(y + i));
+    }
+    const __m256 acc = _mm256_add_ps(_mm256_add_ps(acc0, acc1), _mm256_add_ps(acc2, acc3));
+    __m128 r = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps(acc, 1));
+    r = _mm_add_ps(r, _mm_movehl_ps(r, r));
+    r = _mm_add_ss(r, _mm_movehdup_ps(r));
+    sumf = _mm_cvtss_f32(r);
+#elif defined(GGML_SIMD)
+    // convert chunks of x to F32 and reuse the F32 SIMD dot
+    // TODO: convert F16 to F32 in-register and fuse the dot (needs a new SIMD macro); also avoid the per-chunk call overhead for tiny k
+    float tx[64];
+    for (; i + 64 <= n; i += 64) {
+        const ggml_fp16_t * GGML_RESTRICT xr = x + i;
+        float * GGML_RESTRICT tr = tx;
+        for (int j = 0; j < 64; ++j) {
+            tr[j] = GGML_CPU_FP16_TO_FP32(xr[j]);
+        }
+        float v;
+        ggml_vec_dot_f32(64, &v, 0, tx, 0, y + i, 0, 1);
+        sumf += v;
+    }
+    if (i < n) {
+        const int nr = n - i;
+        const ggml_fp16_t * GGML_RESTRICT xr = x + i;
+        float * GGML_RESTRICT tr = tx;
+        for (int j = 0; j < nr; ++j) {
+            tr[j] = GGML_CPU_FP16_TO_FP32(xr[j]);
+        }
+        float v;
+        ggml_vec_dot_f32(nr, &v, 0, tx, 0, y + i, 0, 1);
+        sumf += v;
+        i = n;
+    }
+#endif
+    for (; i < n; ++i) {
+        sumf += (ggml_float)(GGML_CPU_FP16_TO_FP32(x[i])*y[i]);
+    }
+
+    *s = sumf;
+}
+
 void ggml_vec_silu_f32(const int n, float * y, const float * x) {
    int i = 0;
 #if defined(__AVX512F__) && defined(__AVX512DQ__)
@@ -70,6 +70,7 @@ extern ggml_fp16_t ggml_table_gelu_quick_f16[1 << 16];
 void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc);
 void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t * GGML_RESTRICT x, size_t bx, ggml_bf16_t * GGML_RESTRICT y, size_t by, int nrc);
 void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * GGML_RESTRICT x, size_t bx, ggml_fp16_t * GGML_RESTRICT y, size_t by, int nrc);
+void ggml_vec_dot_f16_f32(int n, float * GGML_RESTRICT s, size_t bs, const ggml_fp16_t * GGML_RESTRICT x, size_t bx, const float * GGML_RESTRICT y, size_t by, int nrc);

 void ggml_vec_silu_f32(const int n, float * y, const float * x);
 ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean); //it will also center y ( y = y - mean )
@@ -428,17 +428,18 @@ static __device__ __forceinline__ void dequantize_V_q4_0(const void * __restrict

    const int8_t * q8 = (const int8_t *) &q;

-#ifdef FP16_AVAILABLE
    if constexpr (std::is_same_v<T, half>) {
+#ifdef FP16_AVAILABLE
        const half2 d = __half2half2(x[ib].d);

 #pragma unroll
        for (int l0 = 0; l0 < ne; l0 += 2) {
            ((half2 *) dst)[l0/2] = d * make_half2(q8[l0 + 0], q8[l0 + 1]);
        }
-    } else
+#else
+        NO_DEVICE_CODE;
 #endif // FP16_AVAILABLE
-    if constexpr (std::is_same_v<T, float>) {
+    } else if constexpr (std::is_same_v<T, float>) {
        const float d = x[ib].d;

 #pragma unroll
@@ -466,8 +467,8 @@ static __device__ __forceinline__ void dequantize_V_q4_1(const void * __restrict

    const int8_t * q8 = (const int8_t *) &q;

-#ifdef FP16_AVAILABLE
    if constexpr (std::is_same_v<T, half>) {
+#ifdef FP16_AVAILABLE
        const half2 dm = x[ib].dm;
        const half2 d  = __half2half2( __low2half(dm));
        const half2 m  = __half2half2(__high2half(dm));
@@ -476,9 +477,10 @@ static __device__ __forceinline__ void dequantize_V_q4_1(const void * __restrict
        for (int l0 = 0; l0 < ne; l0 += 2) {
            ((half2 *) dst)[l0/2] = d * make_half2(q8[l0 + 0], q8[l0 + 1]) + m;
        }
-    } else
+#else
+        NO_DEVICE_CODE;
 #endif // FP16_AVAILABLE
-    if constexpr (std::is_same_v<T, float>) {
+    } else if constexpr (std::is_same_v<T, float>) {
        const float2 dm = __half22float2(x[ib].dm);

 #pragma unroll
@@ -518,17 +520,18 @@ static __device__ __forceinline__ void dequantize_V_q5_0(const void * __restrict

    const int8_t * q8 = (const int8_t *) &q;

-#ifdef FP16_AVAILABLE
    if constexpr (std::is_same_v<T, half>) {
+#ifdef FP16_AVAILABLE
        const half2 d = __half2half2(x[ib].d);

 #pragma unroll
        for (int l0 = 0; l0 < ne; l0 += 2) {
            ((half2 *) dst)[l0/2] = d * make_half2(q8[l0 + 0], q8[l0 + 1]);
        }
-    } else
+#else
+        NO_DEVICE_CODE;
 #endif // FP16_AVAILABLE
-    if constexpr (std::is_same_v<T, float>) {
+    } else if constexpr (std::is_same_v<T, float>) {
        const float d = x[ib].d;

 #pragma unroll
@@ -566,8 +569,8 @@ static __device__ __forceinline__ void dequantize_V_q5_1(const void * __restrict

    const int8_t * q8 = (const int8_t *) &q;

-#ifdef FP16_AVAILABLE
    if constexpr (std::is_same_v<T, half>) {
+#ifdef FP16_AVAILABLE
        const half2 dm = x[ib].dm;
        const half2 d  = __half2half2( __low2half(dm));
        const half2 m  = __half2half2(__high2half(dm));
@@ -576,9 +579,10 @@ static __device__ __forceinline__ void dequantize_V_q5_1(const void * __restrict
        for (int l0 = 0; l0 < ne; l0 += 2) {
            ((half2 *) dst)[l0/2] = d * make_half2(q8[l0 + 0], q8[l0 + 1]) + m;
        }
-    } else
+#else
+        NO_DEVICE_CODE;
 #endif // FP16_AVAILABLE
-    if constexpr (std::is_same_v<T, float>) {
+    } else if constexpr (std::is_same_v<T, float>) {
        const float2 dm = __half22float2(x[ib].dm);

 #pragma unroll
@@ -601,17 +605,18 @@ static __device__ __forceinline__ void dequantize_V_q8_0(const void * __restrict
    int8_t qs[ne];
    ggml_cuda_memcpy_1<ne, 2>(qs, x[ib].qs + iqs);

-#ifdef FP16_AVAILABLE
    if constexpr (std::is_same<T, half>::value) {
+#ifdef FP16_AVAILABLE
        const half2 d = __half2half2(x[ib].d);

 #pragma unroll
        for (int l0 = 0; l0 < ne; l0 += 2) {
            ((half2 *) dst)[l0/2] = d * make_half2(qs[l0 + 0], qs[l0 + 1]);
        }
-    } else
+#else
+        NO_DEVICE_CODE;
 #endif // FP16_AVAILABLE
-    if constexpr (std::is_same<T, float>::value) {
+    } else if constexpr (std::is_same<T, float>::value) {
        const float d = x[ib].d;

 #pragma unroll
@@ -5978,7 +5978,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16 && !ggml_cuda_op_mul_mat_use_fwht(op)) {
                    return false;
                }
-                if (op->op == GGML_OP_MUL_MAT_ID && ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) {
+                if (op->op == GGML_OP_MUL_MAT_ID && (ggml_get_op_params_i32(op, 3) == GGML_PREC_BF16 || ggml_get_op_params_i32(op, 3) == GGML_PREC_F32)) {
                    return false;
                }
 #ifdef GGML_USE_MUSA
@@ -6252,9 +6252,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g

            if (op->src[3]->ne[0] == 1) {
                // Mamba2
-                // (kernel only supports (d_state == 96 || d_state == 128 || d_state == 256) && d_head % 16 == 0)
+                // (kernel only supports (d_state == 64 || d_state == 96 || d_state == 128 || d_state == 256) && d_head % 16 == 0)
                const int64_t d_state = op->src[0]->ne[0];
-                return (d_state == 96 || d_state == 128 || d_state == 256) && op->src[0]->ne[1] % 16 == 0;
+                return (d_state == 64 || d_state == 96 || d_state == 128 || d_state == 256) && op->src[0]->ne[1] % 16 == 0;
            } else {
                if (K > 1) {
                    return false;
@@ -2,239 +2,7 @@
 #include "lightning-indexer.cuh"
 #include "fattn-common.cuh"
 #include "convert.cuh"
-
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
-#if defined(TURING_MMA_AVAILABLE)
-
-typedef union {
-    int2 i2;
-    half2 h2[2];
-} half4;
-
-// TODO add support for AMD cards via rocWMMA
-#include <mma.h>
-namespace wmma = nvcuda::wmma;
-
-template <int WARPS_PER_BLOCK, int K_VECS_PER_BLOCK, int64_t N_EMBD, int64_t N_HEAD, ggml_type TYPE_K>
-static __global__ void lightning_indexer_kernel_wmma(
-        const float * Q, const char * K, const float * W, const half * M, float * dst,
-        int64_t n_stream, int64_t n_batch, int64_t n_kv,
-        size_t nb1, size_t nb2, size_t nb3,
-        size_t nbq1, size_t nbq2, size_t nbq3,
-        size_t nbk1, size_t nbk2, size_t nbk3,
-        size_t nbw1, size_t nbw2, size_t nbw3,
-        size_t nbm1, size_t nbm2, size_t nbm3,
-        int64_t nem3
-    ) {
-
-    constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WARP_SIZE;
-    constexpr int HEADS_PER_INNER_LOOP = 8;
-    constexpr int K_EMBD_PER_INNER_LOOP = 16;
-    constexpr int N_EMBD_PADDED = N_EMBD + 8;
-
-    const int i_batch  = blockIdx.y;
-    const int i_stream = blockIdx.z;
-    const int i_warp   = threadIdx.y;
-    const int i_lane   = threadIdx.x;
-    const int tid      = i_warp * WARP_SIZE + i_lane;
-
-    // each block processes K_VECS_PER_BLOCK K vectors
-    const int start_kv = blockIdx.x * K_VECS_PER_BLOCK;
-
-    const char  * q_base = (const char  *)                 Q + i_batch*nbq2 + i_stream*nbq3;
-    const float * w_base = (const float *) ((const char *) W + i_batch*nbw1 + i_stream*nbw3);
-
-    // phase 1 - load weights and first Q tile to shared memory
-
-    __shared__ float w_shared[N_HEAD];
-    __shared__ int2  q_shared_h[HEADS_PER_INNER_LOOP][N_EMBD_PADDED / 4];
-
-    if (tid < N_HEAD) {
-        w_shared[tid] = w_base[tid];
-    }
-
-    // total number of half4 elements in HEADS_PER_INNER_LOOP x N_EMBD Q tile
-    constexpr int N_Q_TILE = HEADS_PER_INNER_LOOP * (N_EMBD / 4);
-    // number of registers needed in each thread to store Q tile in thread block
-    constexpr int N_Q_NEXT = (N_Q_TILE + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
-
-#pragma unroll
-    for (int i_q = tid; i_q < N_Q_TILE; i_q += THREADS_PER_BLOCK) {
-        const int i_head = i_q / (N_EMBD / 4);
-        const int i_embd = i_q % (N_EMBD / 4);
-        const float4 q = *(const float4 *) (q_base + i_head*nbq1 + i_embd*sizeof(float4));
-        half4 q_packed;
-        q_packed.h2[0] = __float22half2_rn(make_float2(q.x, q.y));
-        q_packed.h2[1] = __float22half2_rn(make_float2(q.z, q.w));
-        q_shared_h[i_head][i_embd] = q_packed.i2;
-    }
-
-    // phase 2 - load (and dequantize if needed) K to shared mem
-
-    __shared__ half2 k_shared_h[K_VECS_PER_BLOCK][N_EMBD_PADDED / 4][2];
-
-    constexpr int n_k = K_VECS_PER_BLOCK * (N_EMBD / 4);
-
-    if constexpr (TYPE_K == GGML_TYPE_F16) {
-#pragma unroll
-        for (int i_k = tid; i_k < n_k; i_k += THREADS_PER_BLOCK) {
-            const int i_k_vec = i_k / (N_EMBD / 4);
-            const int i_embd = i_k % (N_EMBD / 4);
-            const int i_kv = start_kv + i_k_vec;
-            if (i_kv < n_kv) {
-                const int2 * k_base = (const int2 *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3);
-                *(int2*) &k_shared_h[i_k_vec][i_embd] = k_base[i_embd];
-            } else {
-                *(int2*) &k_shared_h[i_k_vec][i_embd] = make_int2(0, 0);
-            }
-        }
-    } else {
-        constexpr dequantize_V_t dequantize_k = get_dequantize_V<TYPE_K, half, 4>();
-#pragma unroll
-        for (int i_k = tid; i_k < n_k; i_k += THREADS_PER_BLOCK) {
-            const int i_k_vec = i_k / (N_EMBD / 4);
-            const int i_embd = i_k % (N_EMBD / 4);
-            const int i_kv = start_kv + i_k_vec;
-            if (i_kv < n_kv) {
-                const void * k_base = (const void *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3);
-                dequantize_k(k_base, &k_shared_h[i_k_vec][i_embd][0], i_embd * 4);
-            } else {
-                *(int2*) &k_shared_h[i_k_vec][i_embd] = make_int2(0, 0);
-            }
-        }
-    }
-
-    __syncthreads();
-
-    // phase 3 - calculate lightning indexer scores
-
-    __shared__ float qk_shared[WARPS_PER_BLOCK][HEADS_PER_INNER_LOOP][K_VECS_PER_BLOCK];
-
-    // load K fragment
-    wmma::fragment<wmma::matrix_b, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, half, wmma::col_major> frag_k;
-    wmma::load_matrix_sync(frag_k, (half*) &k_shared_h[0][i_warp * K_EMBD_PER_INNER_LOOP / 4], N_EMBD_PADDED);
-
-    float score_k = 0.0f;
-
-    for (int i_head_0 = 0; i_head_0 < N_HEAD; i_head_0 += HEADS_PER_INNER_LOOP) {
-        const int i_head_next = i_head_0 + HEADS_PER_INNER_LOOP;
-
-        // we don't use accumulator for anything, fill it with zeros
-        wmma::fragment<wmma::accumulator, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, float> frag_acc;
-        wmma::fill_fragment(frag_acc, 0.0f);
-
-        // load Q fragment
-        wmma::fragment<wmma::matrix_a, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, half, wmma::row_major> frag_q;
-        wmma::load_matrix_sync(frag_q, (half*) &q_shared_h[0][i_warp * K_EMBD_PER_INNER_LOOP / 4], N_EMBD_PADDED);
-
-        // preload next Q tile to registers during matrix multiplication
-        float4 q_next[N_Q_NEXT];
-
-        if (i_head_next < N_HEAD) {
-#pragma unroll
-            for (int i_q = tid, i_q_next = 0; i_q < N_Q_TILE; i_q += THREADS_PER_BLOCK) {
-                const int i_head = i_head_next + i_q / (N_EMBD / 4);
-                const int i_embd =               i_q % (N_EMBD / 4);
-                q_next[i_q_next++] = *(const float4 *) (q_base + i_head*nbq1 + i_embd*sizeof(float4));
-            }
-        }
-
-        // perform matrix multiplication
-        wmma::mma_sync(frag_acc, frag_q, frag_k, frag_acc);
-        wmma::store_matrix_sync((float*) &qk_shared[i_warp][0][0], frag_acc, K_VECS_PER_BLOCK, wmma::mem_row_major);
-
-        // make sure all threads finished using q_shared_h so we can store next tile
-        __syncthreads();
-
-        // write preloaded Q tile to shared memory
-        if (i_head_next < N_HEAD) {
-#pragma unroll
-            for (int i_q = tid, i_q_next = 0; i_q < N_Q_TILE; i_q += THREADS_PER_BLOCK) {
-                const int i_head = i_q / (N_EMBD / 4);
-                const int i_embd = i_q % (N_EMBD / 4);
-                half4 q_packed;
-                q_packed.h2[0] = __float22half2_rn(make_float2(q_next[i_q_next].x, q_next[i_q_next].y));
-                q_packed.h2[1] = __float22half2_rn(make_float2(q_next[i_q_next].z, q_next[i_q_next].w));
-                q_shared_h[i_head][i_embd] = q_packed.i2;
-                ++i_q_next;
-            }
-        }
-
-        // accumulate QK multiplication results from all block warps
-        // (there are 256 threads in block and 256 matmul outputs)
-        // TODO it will break if WARP_SIZE is not 32
-        const int h = tid / K_VECS_PER_BLOCK;
-        const int k = tid % K_VECS_PER_BLOCK;
-        const float w_val = w_shared[i_head_0 + h];
-
-        float sum = 0.0f;
-#pragma unroll
-        for (int w = 0; w < WARPS_PER_BLOCK; ++w) {
-            sum += qk_shared[w][h][k];
-        }
-
-        // ReLU, weight
-        sum = sum > 0.0f ? sum : 0.0f;
-        sum *= w_val;
-
-        // wait until qk_shared[0] is no longer used
-        __syncthreads();
-
-        // reuse qk_shared[0] for storing partial results
-        qk_shared[0][h][k] = sum;
-
-        // wait until all threads write their results
-        __syncthreads();
-
-        // accumulate result over heads
-        if (tid < K_VECS_PER_BLOCK) {
-#pragma unroll
-            for (int i_head = 0; i_head < HEADS_PER_INNER_LOOP; ++i_head) {
-                score_k += qk_shared[0][i_head][tid];
-            }
-        }
-
-        // make sure all threads finished using qk_shared
-        __syncthreads();
-    }
-
-    // phase 4 - store output to VRAM
-
-    if (tid < K_VECS_PER_BLOCK) {
-        const int i_kv = start_kv + tid;
-        if (i_kv < n_kv) {
-            const half * m_base = (const half *) ((const char *) M + i_batch*nbm1 + (i_stream%nem3)*nbm3);
-            float * dst_base = (float *) ((char *) dst + i_batch*nb1 + i_stream*nb3);
-            dst_base[i_kv] = score_k + __half2float(m_base[i_kv]);
-        }
-    }
-}
-
-#else // defined(TURING_MMA_AVAILABLE)
-
-template <int WARPS_PER_BLOCK, int K_VECS_PER_BLOCK, int64_t N_EMBD, int64_t N_HEAD, ggml_type TYPE_K>
-static __global__ void lightning_indexer_kernel_wmma(
-        const float * Q, const char * K, const float * W, const half * M, float * dst,
-        int64_t n_stream, int64_t n_batch, int64_t n_kv,
-        size_t nb1, size_t nb2, size_t nb3,
-        size_t nbq1, size_t nbq2, size_t nbq3,
-        size_t nbk1, size_t nbk2, size_t nbk3,
-        size_t nbw1, size_t nbw2, size_t nbw3,
-        size_t nbm1, size_t nbm2, size_t nbm3,
-        int64_t nem3
-    ) {
-    GGML_UNUSED_VARS(Q, K, W, M, dst,
-        n_stream, n_batch, n_kv,
-        nb1, nb2, nb3,
-        nbq1, nbq2, nbq3,
-        nbk1, nbk2, nbk3,
-        nbw1, nbw2, nbw3,
-        nem3);
-    NO_DEVICE_CODE;
-}
-
-#endif // defined(TURING_MMA_AVAILABLE)
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#include "mma.cuh"

 // tokens scored per block by the tile kernel
 #define LIGHTNING_INDEXER_TILE_TOKENS 8
@@ -247,6 +15,173 @@ static __global__ void lightning_indexer_kernel_wmma(
 #define LIGHTNING_INDEXER_TILE_HEADS_PER_PASS 4
 #endif

+
+template <int WARPS_PER_BLOCK, int K_VECS_PER_BLOCK, int64_t N_EMBD, int64_t N_HEAD, ggml_type TYPE_K>
+static __global__ void lightning_indexer_kernel_mma(
+        const float * Q, const char * K, const float * W, const half * M, float * dst,
+        int64_t n_stream, int64_t n_batch, int64_t n_kv,
+        size_t nb1, size_t nb2, size_t nb3,
+        size_t nbq1, size_t nbq2, size_t nbq3,
+        size_t nbk1, size_t nbk2, size_t nbk3,
+        size_t nbw1, size_t nbw2, size_t nbw3,
+        size_t nbm1, size_t nbm2, size_t nbm3,
+        int64_t nem3
+    ) {
+    constexpr int MMA_DIM = 16;
+    constexpr int WAVE_SIZE = ggml_cuda_get_physical_warp_size();
+    constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WAVE_SIZE;
+    constexpr int N_EMBD_PADDED = N_EMBD + 8;
+    constexpr int K_VECS_PADDED = K_VECS_PER_BLOCK + 8;
+    constexpr int HEAD_CHUNK    = MMA_DIM;
+    constexpr int N_HEAD_CHUNKS = N_HEAD / HEAD_CHUNK;
+    constexpr int N_KV_TILES    = K_VECS_PER_BLOCK / MMA_DIM;
+    constexpr int N_EMBD_TILES  = N_EMBD / MMA_DIM;
+
+    static_assert(N_HEAD           % MMA_DIM == 0, "N_HEAD must be a multiple of 16");
+    static_assert(N_EMBD           % MMA_DIM == 0, "N_EMBD must be a multiple of 16");
+    static_assert(K_VECS_PER_BLOCK % MMA_DIM == 0, "K_VECS_PER_BLOCK must be a multiple of 16");
+
+#if defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
+    using tile_qk = ggml_cuda_mma::tile<MMA_DIM, MMA_DIM, float, ggml_cuda_mma::DATA_LAYOUT_J_MAJOR>;
+#else
+    using tile_qk = ggml_cuda_mma::tile<MMA_DIM, MMA_DIM, float>;
+#endif
+
+    const int i_batch  = blockIdx.y;
+    const int i_stream = blockIdx.z;
+    const int i_warp   = threadIdx.y;
+    const int i_lane   = threadIdx.x;
+    const int tid      = i_warp * WAVE_SIZE + i_lane;
+
+    const int start_kv = blockIdx.x * K_VECS_PER_BLOCK;
+
+    const char  * q_base = (const char  *)                 Q + i_batch*nbq2 + i_stream*nbq3;
+    const float * w_base = (const float *) ((const char *) W + i_batch*nbw1 + i_stream*nbw3);
+
+    __shared__ float w_shared[N_HEAD];
+    __shared__ half  q_shared[2][HEAD_CHUNK][N_EMBD_PADDED];
+    __shared__ half  k_shared[K_VECS_PER_BLOCK][N_EMBD_PADDED];
+    __shared__ float qk_chunk[HEAD_CHUNK][K_VECS_PADDED];
+    __shared__ float score_shared[K_VECS_PER_BLOCK];
+
+    for (int i = tid; i < N_HEAD; i += THREADS_PER_BLOCK) {
+        w_shared[i] = w_base[i];
+    }
+    for (int i = tid; i < K_VECS_PER_BLOCK; i += THREADS_PER_BLOCK) {
+        score_shared[i] = 0.0f;
+    }
+
+    constexpr int N_K_VEC = K_VECS_PER_BLOCK * (N_EMBD / 4);
+    if constexpr (TYPE_K == GGML_TYPE_F16) {
+        for (int i = tid; i < N_K_VEC; i += THREADS_PER_BLOCK) {
+            const int i_k_vec = i / (N_EMBD / 4);
+            const int i_embd4 = i % (N_EMBD / 4);
+            const int i_kv    = start_kv + i_k_vec;
+            if (i_kv < n_kv) {
+                const int2 * k_ptr = (const int2 *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3);
+                *(int2 *) &k_shared[i_k_vec][i_embd4*4] = k_ptr[i_embd4];
+            } else {
+                *(int2 *) &k_shared[i_k_vec][i_embd4*4] = make_int2(0, 0);
+            }
+        }
+    } else {
+        constexpr dequantize_V_t dequantize_k = get_dequantize_V<TYPE_K, half, 4>();
+        for (int i = tid; i < N_K_VEC; i += THREADS_PER_BLOCK) {
+            const int i_k_vec = i / (N_EMBD / 4);
+            const int i_embd4 = i % (N_EMBD / 4);
+            const int i_kv    = start_kv + i_k_vec;
+            if (i_kv < n_kv) {
+                const void * k_ptr = (const void *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3);
+                dequantize_k(k_ptr, &k_shared[i_k_vec][i_embd4*4], i_embd4*4);
+            } else {
+                *(int2 *) &k_shared[i_k_vec][i_embd4*4] = make_int2(0, 0);
+            }
+        }
+    }
+
+    constexpr int N_Q_CHUNK_VEC = HEAD_CHUNK * (N_EMBD / 4);
+    constexpr int N_Q_NEXT = (N_Q_CHUNK_VEC + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+
+    for (int i = tid; i < N_Q_CHUNK_VEC; i += THREADS_PER_BLOCK) {
+        const int i_head  = i / (N_EMBD / 4);
+        const int i_embd4 = i % (N_EMBD / 4);
+        const float4 q = *(const float4 *) (q_base + i_head*nbq1 + i_embd4*sizeof(float4));
+        const half2 qh[2] = {make_half2(q.x, q.y), make_half2(q.z, q.w)};
+        ggml_cuda_memcpy_1<sizeof(qh)>(&q_shared[0][i_head][i_embd4*4], qh);
+    }
+
+    __syncthreads();
+
+    constexpr ggml_cuda_mma::data_layout input_layout = ggml_cuda_mma::get_input_data_layout();
+
+    for (int hc = 0; hc < N_HEAD_CHUNKS; ++hc) {
+        const int buf  = hc & 1;
+        const int nbuf = (hc + 1) & 1;
+        const bool has_next = (hc + 1) < N_HEAD_CHUNKS;
+
+        float4 q_next[N_Q_NEXT];
+        if (has_next) {
+#pragma unroll
+            for (int i = tid, j = 0; i < N_Q_CHUNK_VEC; i += THREADS_PER_BLOCK, ++j) {
+                const int i_head  = (hc + 1) * HEAD_CHUNK + i / (N_EMBD / 4);
+                const int i_embd4 =                         i % (N_EMBD / 4);
+                q_next[j] = *(const float4 *) (q_base + i_head*nbq1 + i_embd4*sizeof(float4));
+            }
+        }
+
+        for (int kt = i_warp; kt < N_KV_TILES; kt += WARPS_PER_BLOCK) {
+            tile_qk Dqk;
+#pragma unroll
+            for (int et = 0; et < N_EMBD_TILES; ++et) {
+                ggml_cuda_mma::tile<MMA_DIM, MMA_DIM/2, half2, input_layout> Aq;
+                ggml_cuda_mma::tile<MMA_DIM, MMA_DIM/2, half2, input_layout> Bk;
+                ggml_cuda_mma::load_generic(Aq, (const half2 *) &q_shared[buf][0][et*MMA_DIM], N_EMBD_PADDED/2);
+                ggml_cuda_mma::load_generic(Bk, (const half2 *) &k_shared[kt*MMA_DIM][et*MMA_DIM], N_EMBD_PADDED/2);
+                ggml_cuda_mma::mma(Dqk, Aq, Bk);
+            }
+#pragma unroll
+            for (int l = 0; l < Dqk.ne; ++l) {
+                const int head_local = Dqk.get_i(l);
+                const int kv_local   = kt*MMA_DIM + Dqk.get_j(l);
+                qk_chunk[head_local][kv_local] = Dqk.x[l];
+            }
+        }
+
+        __syncthreads();
+
+        if (has_next) {
+#pragma unroll
+            for (int i = tid, j = 0; i < N_Q_CHUNK_VEC; i += THREADS_PER_BLOCK, ++j) {
+                const int i_head  = i / (N_EMBD / 4);
+                const int i_embd4 = i % (N_EMBD / 4);
+                const half2 qh[2] = {make_half2(q_next[j].x, q_next[j].y), make_half2(q_next[j].z, q_next[j].w)};
+                ggml_cuda_memcpy_1<sizeof(qh)>(&q_shared[nbuf][i_head][i_embd4*4], qh);
+            }
+        }
+
+        for (int kv = tid; kv < K_VECS_PER_BLOCK; kv += THREADS_PER_BLOCK) {
+            float s = 0.0f;
+#pragma unroll
+            for (int h = 0; h < HEAD_CHUNK; ++h) {
+                const float qk = qk_chunk[h][kv];
+                s += (qk > 0.0f ? qk : 0.0f) * w_shared[hc*HEAD_CHUNK + h];
+            }
+            score_shared[kv] += s;
+        }
+
+        __syncthreads();
+    }
+
+    for (int i_kv_local = tid; i_kv_local < K_VECS_PER_BLOCK; i_kv_local += THREADS_PER_BLOCK) {
+        const int i_kv = start_kv + i_kv_local;
+        if (i_kv < n_kv) {
+            const half  * m_base   = (const half *) ((const char *) M + i_batch*nbm1 + (i_stream%nem3)*nbm3);
+            float       * dst_base = (float *) ((char *) dst + i_batch*nb1 + i_stream*nb3);
+            dst_base[i_kv] = score_shared[i_kv_local] + __half2float(m_base[i_kv]);
+        }
+    }
+}
+
 // TODO there is one ugly assumption used in this kernel - that WARP_SIZE is equal to 32
 // thanks to that one warp operating on float4 processes whole indexer K/Q vectors
 // 32 * 4 = 128 (N_EMBD)
@@ -616,27 +551,22 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor *
    const int cc     = ggml_cuda_info().devices[device].cc;

    if (n_embd == 128 && n_head == 64) {
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
-        if (GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16) {
-            // use wmma kernel
+        if ((turing_mma_available(cc) || amd_mfma_available(cc) || amd_wmma_available(cc)) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16) {
            constexpr int K_VECS_PER_BLOCK = 32;
-            constexpr int WARPS_PER_BLOCK = 8;
+            constexpr int WARPS_PER_BLOCK  = 4;

-            dim3 block(32, WARPS_PER_BLOCK);
+            dim3 block(ggml_cuda_info().devices[device].warp_size, WARPS_PER_BLOCK);
            int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK);
            dim3 grid(num_kv_blocks, n_batch, n_stream);

-            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_F16)
-            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q4_0)
-            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q4_1)
-            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q5_0)
-            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q5_1)
-            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q8_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 64, k, GGML_TYPE_F16)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 64, k, GGML_TYPE_Q4_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 64, k, GGML_TYPE_Q4_1)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 64, k, GGML_TYPE_Q5_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 64, k, GGML_TYPE_Q5_1)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 64, k, GGML_TYPE_Q8_0)
            GGML_ABORT("fatal error");
        } else {
-#else // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
-        {
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
            // use vector kernel
            constexpr int K_VECS_PER_WARP = 8;
            constexpr int WARPS_PER_BLOCK = 8;
@@ -657,27 +587,22 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor *
            GGML_ABORT("fatal error");
        }
    } else if (n_embd == 128 && n_head == 32) {
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
-        if (GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16) {
-            // use wmma kernel
+        if ((turing_mma_available(cc) || amd_mfma_available(cc) || amd_wmma_available(cc)) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16) {
            constexpr int K_VECS_PER_BLOCK = 32;
-            constexpr int WARPS_PER_BLOCK = 8;
+            constexpr int WARPS_PER_BLOCK  = 4;

-            dim3 block(32, WARPS_PER_BLOCK);
+            dim3 block(ggml_cuda_info().devices[device].warp_size, WARPS_PER_BLOCK);
            int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK);
            dim3 grid(num_kv_blocks, n_batch, n_stream);

-            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_F16)
-            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q4_0)
-            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q4_1)
-            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q5_0)
-            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q5_1)
-            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q8_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 32, k, GGML_TYPE_F16)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 32, k, GGML_TYPE_Q4_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 32, k, GGML_TYPE_Q4_1)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 32, k, GGML_TYPE_Q5_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 32, k, GGML_TYPE_Q5_1)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_mma, 128, 32, k, GGML_TYPE_Q8_0)
            GGML_ABORT("fatal error");
        } else {
-#else // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
-        {
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
            // use vector kernel
            constexpr int K_VECS_PER_WARP = 8;
            constexpr int WARPS_PER_BLOCK = 8;
@@ -50,7 +50,8 @@ void ggml_cuda_mul_mat_f(ggml_backend_cuda_context & ctx, const ggml_tensor * sr
    mmf_ids_data * ids_info_ptr = nullptr;
    ggml_cuda_pool_alloc<int32_t> ids_src_compact_dev;
    ggml_cuda_pool_alloc<int32_t> ids_dst_compact_dev;
-    ggml_cuda_pool_alloc<int32_t> expert_bounds_dev;
+
+    ggml_cuda_pool_alloc<ggml_cuda_mm_id_tile_spec> tile_spec_dev;

    // For MUL_MAT_ID the memory layout is different than for MUL_MAT:
    const int64_t ncols_dst          = ids ? ne2  : ne1;
@@ -74,25 +75,28 @@ void ggml_cuda_mul_mat_f(ggml_backend_cuda_context & ctx, const ggml_tensor * sr
        const int64_t n_experts     = ne02;
        const int64_t n_tokens      = ne12;
        const int64_t ne_get_rows   = n_tokens * n_expert_used;
+        const int64_t max_tiles     = (ne_get_rows + 16 - 1) / 16 + ne02;

        ids_src_compact_dev.alloc(ctx.pool(), ne_get_rows);
        ids_dst_compact_dev.alloc(ctx.pool(), ne_get_rows);
-        expert_bounds_dev.alloc(ctx.pool(), n_experts + 1);
+        tile_spec_dev.alloc(ctx.pool(), max_tiles);

        const int si1  = static_cast<int>(ids_s1);
        const int sis1 = static_cast<int>(src1->nb[2] / src1->nb[1]);

        GGML_ASSERT(sis1 > 0);

-        ggml_cuda_launch_mm_ids_helper(ids_d, ids_src_compact_dev.get(), ids_dst_compact_dev.get(), expert_bounds_dev.get(),
-            static_cast<int>(n_experts), static_cast<int>(n_tokens), static_cast<int>(n_expert_used), static_cast<int>(ne11), si1, sis1, /*write_inverse =*/ false, ctx.stream());
+        ggml_cuda_launch_mm_ids_helper(ids_d, ids_src_compact_dev.get(), ids_dst_compact_dev.get(), tile_spec_dev.get(),
+            static_cast<int>(n_experts), static_cast<int>(n_tokens), static_cast<int>(n_expert_used), static_cast<int>(ne11), 16,
+            si1, sis1, /*write_inverse =*/ false, ctx.stream());
        CUDA_CHECK(cudaGetLastError());

-        ids_info.ids_src_compact   = ids_src_compact_dev.get();
-        ids_info.ids_dst_compact   = ids_dst_compact_dev.get();
-        ids_info.expert_bounds_dev = expert_bounds_dev.get();
-        ids_info.n_experts         = static_cast<int>(n_experts);
-        ids_info.sis1              = sis1;
+        ids_info.ids_src_compact = ids_src_compact_dev.get();
+        ids_info.ids_dst_compact = ids_dst_compact_dev.get();
+        ids_info.tile_spec_dev   = tile_spec_dev.get();
+        ids_info.n_experts       = static_cast<int>(n_experts);
+        ids_info.max_tiles       = static_cast<int>(max_tiles);
+        ids_info.sis1            = sis1;
        ids_info_ptr = &ids_info;
    }

@@ -1,7 +1,8 @@
 #pragma once

-#include "mma.cuh"
 #include "common.cuh"
+#include "mma.cuh"
+#include "mmid.cuh"
 #include "convert.cuh"

 using namespace ggml_cuda_mma;
@@ -36,8 +37,9 @@ static constexpr __device__ int mmf_get_padding() {
 struct mmf_ids_data {
    const int32_t * ids_src_compact = nullptr;
    const int32_t * ids_dst_compact = nullptr;
-    const int32_t * expert_bounds_dev = nullptr;
+    const ggml_cuda_mm_id_tile_spec * tile_spec_dev = nullptr;
    int n_experts = 0;
+    int max_tiles = 0;
    int sis1 = 0;
 };

@@ -308,7 +310,7 @@ __launch_bounds__(ggml_cuda_get_physical_warp_size()*nwarps, 1)
 static __global__ void mul_mat_f_ids(
        const T * __restrict__ x, const float * __restrict__ y,
        const int32_t * __restrict__ ids_src_compact, const int32_t * __restrict__ ids_dst_compact,
-        const int32_t * __restrict__ expert_bounds, float * __restrict__ dst,
+        const ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec, float * __restrict__ dst,
        const int ncols, const int ncols_dst_total, const int nchannels_dst, const int stride_row, const int stride_col_y, const int stride_col_dst,
        const int channel_ratio, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
        const int sample_ratio, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
@@ -349,21 +351,20 @@ static __global__ void mul_mat_f_ids(
    constexpr int ntA = rows_per_block / tile_A::I;
    constexpr int ntB = (cols_per_block + tile_B::I - 1) / tile_B::I;

-    const int row0        = blockIdx.x * rows_per_block;
+    const int row0     = blockIdx.x * rows_per_block;
+    const int tile_idx = blockIdx.y;

-    const int expert_idx = blockIdx.y;
-    const int expert_start = expert_bounds[expert_idx];
-    const int expert_end   = expert_bounds[expert_idx + 1];
+    const ggml_cuda_mm_id_tile_spec ts = tile_spec[tile_idx];
+
+    const int expert_idx   = ts.expert;
+    const int expert_start = ts.col_low;
+    const int expert_end   = ts.col_high;
    const int ncols_expert = expert_end - expert_start;

-    const int tiles_for_expert = (ncols_expert + cols_per_block - 1) / cols_per_block;
-    const int tile_idx = blockIdx.z;
-    if (tile_idx >= tiles_for_expert) {
+    if (ncols_expert == 0) {
        return;
    }

-    const int col_base = tile_idx * cols_per_block;
-
    GGML_UNUSED(channel_ratio);

    const int channel_x   = expert_idx;
@@ -409,10 +410,9 @@ static __global__ void mul_mat_f_ids(
 #pragma unroll
                for (int j0 = 0; j0 < tile_B::I; ++j0) {
                    const int j = j0 + tile_idx_local*tile_B::I;
-                    const int global_j = col_base + j;
                    float val = 0.0f;
-                    if (j < cols_per_block && global_j < ncols_expert) {
-                        const int src_entry = ids_src_expert[global_j];
+                    if (j < cols_per_block && j < ncols_expert) {
+                        const int src_entry = ids_src_expert[j];
                        const uint2 qrm = fast_div_modulo((uint32_t) src_entry, sis1_fd);
                        const int token   = (int) qrm.x;
                        const int channel = (int) qrm.y;
@@ -462,10 +462,9 @@ static __global__ void mul_mat_f_ids(
 #pragma unroll
                for (int j0 = 0; j0 < tile_B::I; ++j0) {
                    const int j = j0 + tile_idx_local*tile_B::I;
-                    const int global_j = col_base + j;
                    float2 tmp = make_float2(0.0f, 0.0f);
-                    if (j < cols_per_block && global_j < ncols_expert) {
-                        const int src_entry = ids_src_expert[global_j];
+                    if (j < cols_per_block && j < ncols_expert) {
+                        const int src_entry = ids_src_expert[j];
                        const uint2 qrm = fast_div_modulo((uint32_t) src_entry, sis1_fd);
                        const int token   = (int) qrm.x;
                        const int channel = (int) qrm.y;
@@ -564,9 +563,8 @@ static __global__ void mul_mat_f_ids(
            }
        }

-        const int global_j = col_base + j;
-        if (j < cols_per_block && global_j < ncols_expert && nchannels_dst > 0) {
-            const int dst_entry = ids_dst_expert[global_j];
+        if (j < cols_per_block && j < ncols_expert && nchannels_dst > 0) {
+            const int dst_entry = ids_dst_expert[j];
            const uint2 qrm = fast_div_modulo((uint32_t) dst_entry, nch_fd);
            const int token = (int) qrm.x;
            if (token < ncols_dst_total) {
@@ -580,7 +578,7 @@ static __global__ void mul_mat_f_ids(
    }
    }
 #else
-    GGML_UNUSED_VARS(x, y, ids_src_compact, ids_dst_compact, expert_bounds, dst,
+    GGML_UNUSED_VARS(x, y, ids_src_compact, ids_dst_compact, tile_spec, dst,
        ncols, ncols_dst_total, nchannels_dst, stride_row, stride_col_y, stride_col_dst,
        channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
        sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, sis1_fd, nch_fd);
@@ -603,17 +601,16 @@ static inline void mul_mat_f_switch_ids(
    // Use the compact-ids kernel only for larger tiles; for small ncols_dst (< 16)
    // we prefer the normal mul_mat_f path with has_ids=true.
    if (has_ids_data && ncols_dst > 16) {
-        const int max_tiles = (int) ((ncols_dst + cols_per_block - 1) / cols_per_block);
-        if (max_tiles == 0) {
+        if (ids_data->max_tiles == 0) {
            return;
        }
-        dim3 block_nums_ids(block_nums.x, ids_data->n_experts, max_tiles);
+        dim3 block_nums_ids(block_nums.x, ids_data->max_tiles, 1);

        const uint3 sis1_fd = ids_data->sis1 > 0 ? init_fastdiv_values((uint32_t) ids_data->sis1) : make_uint3(0, 0, 1);
        const uint3 nch_fd  = init_fastdiv_values((uint32_t) nchannels_dst);

        mul_mat_f_ids<T, rows_per_block, cols_per_block, nwarps><<<block_nums_ids, block_dims, nbytes_shared_total, stream>>>
-            (x, y, ids_data->ids_src_compact, ids_data->ids_dst_compact, ids_data->expert_bounds_dev, dst,
+            (x, y, ids_data->ids_src_compact, ids_data->ids_dst_compact, ids_data->tile_spec_dev, dst,
            ncols_x, ncols_dst, nchannels_dst, stride_row, stride_col_y, stride_col_dst,
            channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
            sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst,
@@ -27,12 +27,16 @@ template <>      struct mm_ids_pow2<0> { static constexpr int value = 1; };
 // Helper function for mul_mat_id, converts ids to a more convenient format.
 // ids_src1 describes how to permute the flattened column indices of src1 in order to get a compact src1 tensor sorted by expert.
 // ids_dst describes the same mapping but for the dst tensor.
-// The upper and lower bounds for the ith expert in the compact src1 tensor are stored in expert_bounds[i:i+1].
+// The first (n_tokens*n_expert_used + J - 1) / tile_width elements of tile_spec contain the specication for
+//   matrix multiplication tiles at token positions which are exactly divided by tile_width.
+// Additionally the following n_expert elements of tile_spec may contain additional tiles that fall between these seams.
 template <int n_expert_used_template>
 __launch_bounds__(ggml_cuda_get_physical_warp_size(), 1)
 static __global__ void mm_ids_helper(
-        const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
-        const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse) {
+        const int32_t * __restrict__ ids,
+        int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec,
+        const int n_tokens, const int n_expert_used_var, const int nchannels_y, const uint3 tile_width,
+        const int si1, const int sis1, const bool write_inverse) {
    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
    const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
    const int expert = blockIdx.x;
@@ -118,19 +122,29 @@ static __global__ void mm_ids_helper(
        return;
    }

-    expert_bounds[expert] = nex_prev;
+    int col_low  = nex_prev;
+    int col_high = nex_prev + min(int(tile_width.z), it_compact);

-    if (expert < static_cast<int>(gridDim.x) - 1) {
-        return;
+    // The while loop covers all tiles at the seams containing the column with % tile_width == 0.
+    while (col_high > fastdiv(col_low + tile_width.z - 1, tile_width) * tile_width.z) {
+        tile_spec[fastdiv(col_low + tile_width.z - 1, tile_width)]
+            = {expert, col_low, col_high, /*padding =*/ 0};
+
+        col_low  = col_high;
+        col_high = min(col_high + int(tile_width.z), nex_prev + it_compact);
    }

-    expert_bounds[gridDim.x] = nex_prev + it_compact;
+    // Finally, at most one extra tile_spec between the seams per expert.
+    tile_spec[fastdiv(n_tokens*n_expert_used + tile_width.z - 1, tile_width) + expert]
+        = {expert, col_low, col_high, /*padding =*/ 0};
 }

 template <int n_expert_used_template>
 static void launch_mm_ids_helper(
-        const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
-        const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
+        const int32_t * __restrict__ ids,
+        int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec,
+        const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int tile_width,
+        const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
    GGML_ASSERT(n_tokens          < (1 << 22) && "too few bits in mm_ids_helper_store");
    GGML_ASSERT(n_expert_used_var < (1 << 10) && "too few bits in mm_ids_helper_store");

@@ -139,41 +153,53 @@ static void launch_mm_ids_helper(
    const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
    CUDA_SET_SHARED_MEMORY_LIMIT(mm_ids_helper<n_expert_used_template>, smpbo);

+    const uint3 tile_width_fd = init_fastdiv_values(tile_width);
+
    const dim3 num_blocks(n_experts, 1, 1);
    const dim3 block_size(warp_size, 1, 1);
    const size_t nbytes_shared = n_tokens*sizeof(mm_ids_helper_store);
    GGML_ASSERT(nbytes_shared <= smpbo);
    mm_ids_helper<n_expert_used_template><<<num_blocks, block_size, nbytes_shared, stream>>>
-        (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1, write_inverse);
+        (ids, ids_src1, ids_dst, tile_spec, n_tokens, n_expert_used_var, nchannels_y, tile_width_fd, si1, sis1, write_inverse);
 }

 void ggml_cuda_launch_mm_ids_helper(
-        const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
-        const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
+        const int32_t * __restrict__ ids,
+        int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec,
+        const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int tile_width,
+        const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
    switch (n_expert_used) {
        case  2:
-            launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, tile_spec,
+                n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
            break;
        case  4:
-            launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, tile_spec,
+                n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
            break;
        case  6:
-            launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, tile_spec,
+                n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
            break;
        case  8:
-            launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, tile_spec,
+                n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
            break;
        case 10:
-            launch_mm_ids_helper<10>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            launch_mm_ids_helper<10>(ids, ids_src1, ids_dst, tile_spec,
+                n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
            break;
        case 16:
-            launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, tile_spec,
+                n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
            break;
        case 32:
-            launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, tile_spec,
+                n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
            break;
        default:
-            launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+            launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, tile_spec,
+                n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
            break;
    }
 }
@@ -1,5 +1,16 @@
 #pragma once

+struct ggml_cuda_mm_id_tile_spec {
+    int32_t expert;
+    int32_t col_low;
+    int32_t col_high;
+
+    int32_t padding;
+};
+static_assert(sizeof(ggml_cuda_mm_id_tile_spec) == 16, "bad ggml_cuda_mm_id_tile_spec size");
+
 void ggml_cuda_launch_mm_ids_helper(
-        const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds,
-        int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, bool write_inverse, cudaStream_t stream);
+        const int32_t * __restrict__ ids,
+        int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec,
+        const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int tile_width,
+        const int si1, const int sis1, const bool write_inverse, cudaStream_t stream);
@@ -247,10 +247,7 @@ void ggml_cuda_mul_mat_q(
            ncols_opt = ne12;

            // 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.
-            if (GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4_M(cc) || GGML_CUDA_CC_IS_RDNA4(cc)) {
-                ncols_opt = (ne12*n_expert_used + ne02 - 1) / ne02;
-            }
+            ncols_opt = (ne12*n_expert_used + ne02 - 1) / ne02;
        }

        J_best = ggml_cuda_mmq_get_J_best(src0->type, ncols_opt, fallback, cc, smpbo, prec_src1, &nthreads_best);
@@ -264,6 +261,8 @@ void ggml_cuda_mul_mat_q(
        * src1_load_chunk_size;

    if (!ids) {
+        const int ntj = (ne1 + J_best - 1) / J_best;
+
        const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block + src1_q8_1_padding;
        ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
        ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
@@ -301,7 +300,7 @@ void ggml_cuda_mul_mat_q(
            ne00, ne01, ne1, s01, ne11, s1,
            ne02, ne12, s02, s12, s2,
            ne03, ne13, s03, s13, s3,
-            ne1, J_best, fusion_device};
+            ntj, J_best, fusion_device};
        ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
        return;
    }
@@ -316,7 +315,9 @@ void ggml_cuda_mul_mat_q(

    ggml_cuda_pool_alloc<int32_t> ids_src1(ctx.pool(), ne_get_rows);
    ggml_cuda_pool_alloc<int32_t> ids_dst(ctx.pool(), ne_get_rows + J_best-1); // Needs to be padded for unconditional memory access.
-    ggml_cuda_pool_alloc<int32_t> expert_bounds(ctx.pool(), ne02 + 1);
+
+    const int ntj = (ne_get_rows + J_best - 1) / J_best + ne02;
+    ggml_cuda_pool_alloc<ggml_cuda_mm_id_tile_spec> tile_spec(ctx.pool(), ntj);

    // gate/up activations are broadcast across experts (ne11 == 1): quantize each token once and
    // scatter to its slots. ids_src1 then holds the inverse map (token slot -> compact row).
@@ -327,8 +328,8 @@ void ggml_cuda_mul_mat_q(
        const int si1  = ids->nb[1] / ggml_element_size(ids);
        const int sis1 = nb12 / nb11;

-        ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(),
-            ne02, ne12, n_expert_used, ne11, si1, sis1, /*write_inverse =*/ dedup_bcast, stream);
+        ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), tile_spec.get(),
+            ne02, ne12, n_expert_used, ne11, J_best, si1, sis1, /*write_inverse =*/ dedup_bcast, stream);
        CUDA_CHECK(cudaGetLastError());
    }

@@ -375,12 +376,12 @@ void ggml_cuda_mul_mat_q(

    // 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,
+        src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), tile_spec.get(), dst_d,
        src1_scale.ptr,
        ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
-        ne02, ne02, s02, s12, s2,
+        ne02,    1, s02, s12, s2,
        ne03, ne13, s03, s13, s3,
-        ne12, J_best, fusion_device};
+        ntj, J_best, fusion_device};

    ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
 }
@@ -1,6 +1,7 @@
 #pragma once

 #include "common.cuh"
+#include "mmid.cuh"
 #include "unary.cuh"

 #include <climits>
@@ -14,8 +15,8 @@
 typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, const char * __restrict__ x_gate, int * x_tile,
    const int kbx0, const int i_max, const int stride);
 typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00);
-typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted,
-    float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ ids_dst,
+    float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_sup,
    const ggml_glu_op glu_op, const float glu_limit, const float * __restrict__ x_scale, const float * __restrict__ gate_scale);

 enum mmq_q8_1_ds_layout {
@@ -462,22 +463,21 @@ static __device__ __forceinline__ const block_t * ggml_cuda_mmq_x_row(

 template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a(
        const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst,
-        const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+        const float * __restrict__ y_scale, const int stride, const int i_max, const int j_sup,
        const ggml_glu_op glu_op, const float glu_limit, const float * __restrict__ x_scale, const float * __restrict__ gate_scale) {
    GGML_UNUSED_VARS(glu_op, glu_limit, x_scale, gate_scale);
-
    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
    constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
    static_assert(I % warp_size == 0, "DP4A MMQ tile height must be a multiple of the physical warp size");

-    const bool y_scale_used = y_scale != nullptr;
+    const bool y_scale_used = type == GGML_TYPE_NVFP4 && y_scale != nullptr;

 #pragma unroll
    for (int j0 = 0; j0 < J; j0 += nwarps) {
        const int j = j0 + threadIdx.y;

-        if (j > j_max) {
+        if (j >= j_sup) {
            return;
        }

@@ -489,15 +489,10 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
                continue;
            }

-            if constexpr (type == GGML_TYPE_NVFP4) {
-                if (y_scale_used) {
-                    dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
-                } else {
-                    dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
-                }
+            if (y_scale_used) {
+                dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
            } else {
                dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
-                GGML_UNUSED(y_scale_used);
            }
        }
    }
@@ -508,7 +503,7 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
 template<ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
            const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst,
-            const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+            const float * __restrict__ y_scale, const int stride, const int i_max, const int j_sup,
            const ggml_glu_op glu_op, const float glu_limit, const float * __restrict__ x_scale, const float * __restrict__ gate_scale) {

 #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
@@ -523,7 +518,7 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(

    const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I);

-    const bool y_scale_used = y_scale != nullptr;
+    const bool y_scale_used = type == GGML_TYPE_NVFP4 && y_scale != nullptr;

    // Fused NVFP4: load the weight scales once, 1 without scales.
    constexpr bool fused_nvfp4 = has_fusion && type == GGML_TYPE_NVFP4;
@@ -538,7 +533,7 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
            for (int l = 0; l < (has_fusion ? tile_C::ne/2 : tile_C::ne); ++l) {
                const int j = j0 + (threadIdx.y % ntx) * tile_C::J + tile_C::get_j(l);

-                if (j > j_max) {
+                if (j >= j_sup) {
                    continue;
                }

@@ -548,7 +543,7 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
                    continue;
                }

-                if constexpr (has_fusion) {
+                if (has_fusion) {
                    const int sum_index = (j0/tile_C::J + n)*tile_C::ne + l;
                    float up   = sum[sum_index];
                    float gate = sum[sum_index + tile_C::ne/2];
@@ -561,15 +556,10 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
                        gate *= scale_gate;
                    }
                    dst[ids_dst[j]*stride + i] = ggml_cuda_op_glu_single(glu_op, gate, up, glu_limit);
-                } else if constexpr (type == GGML_TYPE_NVFP4) {
-                    if (y_scale_used) {
-                        dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l];
-                    } else {
-                        dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
-                    }
+                } else if (y_scale_used) {
+                    dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l];
                } else {
                    dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
-                    GGML_UNUSED(y_scale_used);
                }
            }
        }
@@ -937,7 +927,7 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
        const float * __restrict__ y_scale, const ggml_glu_op glu_op, const float glu_limit,
        const float * __restrict__ x_scale, const float * __restrict__ gate_scale,
        const int stride_row_x, const int ncols_y, const int stride_col_dst,
-        const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) {
+        const int tile_x_i_max, const int tile_y_j_sup, const int kb0_start, const int kb0_stop) {
    static_assert(!(fixup && has_fusion), "the fused gate/up path does not split K");
    static_assert(!(fallback && has_fusion), "the fused gate/up path needs nrows_x % 128 == 0");

@@ -970,7 +960,7 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
    constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int);

    for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) {
-        load_tiles(x, x_gate, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x);
+        load_tiles(x, x_gate, tile_x, offset_x + kb0, tile_x_i_max, stride_row_x);
        {
            const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz;
 #pragma unroll
@@ -1007,7 +997,7 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
    if (fixup) {
        write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J, glu_op, glu_limit, x_scale, gate_scale);
    } else {
-        write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j, glu_op, glu_limit, x_scale, gate_scale);
+        write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_i_max, tile_y_j_sup, glu_op, glu_limit, x_scale, gate_scale);
    }
 }

@@ -1018,13 +1008,13 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fu
 __launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1), ggml_cuda_mmq_get_occupancy(type, J, fallback, prec_src1))
 static __global__ void mul_mat_q(
        const char * __restrict__ x, const char * __restrict__ x_gate, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
-        const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup,
+        const ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec, float * __restrict__ dst, float * __restrict__ tmp_fixup,
        const float * __restrict__ y_scale, const ggml_glu_op glu_op, const float glu_limit,
        const float * __restrict__ x_scale, const float * __restrict__ gate_scale,
        const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst,
        const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
        const uint3 sample_ratio, const uint3 nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
-        const uint3 ntx) {
+        const uint3 ntj) {

 #ifdef TURING_MMA_AVAILABLE
    constexpr bool fusion_available = true;
@@ -1049,7 +1039,7 @@ static __global__ void mul_mat_q(
    constexpr int qk        = ggml_cuda_type_traits<type>::qk;
    constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);

-    const uint32_t nty = (nrows_x + I - 1) / I; // Number of tiles y
+    const uint32_t nti = (nrows_x + I - 1) / I; // Number of tiles in I direction

    // Initialize the ids for writing back data with just the index.
    // For regular matrix multiplications this is never changed.
@@ -1074,35 +1064,28 @@ static __global__ void mul_mat_q(

        const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y);
        const int wt = tmp2.x;
-        const int zt = tmp2.y;
+        int       zt = tmp2.y;
        const int jt = blockIdx.x;
        const int it = blockIdx.y;

        // Defaults for regular matrix multiplication:
-        int col_low    = 0;
-        int col_high   = ncols_dst;
-        int col_diff   = ncols_dst;
+        int col_low        = jt*J;
+        int col_high       = min(col_low + J, ncols_dst);
        int offset_y       = wt*stride_sample_y   + zt*stride_channel_y;
+        int offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
        int offset_dst     = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
-        int offset_y_scale;
-        if constexpr (type == GGML_TYPE_NVFP4) {
-            offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
-        } else {
-            GGML_UNUSED(offset_y_scale);
-        }

        if (ids_dst) {
-            col_low  = expert_bounds[zt + 0];
-            col_high = expert_bounds[zt + 1];
-            col_diff = col_high - col_low;
+            const ggml_cuda_mm_id_tile_spec tsj = tile_spec[jt];
+            zt       = tsj.expert;
+            col_low  = tsj.col_low;
+            col_high = tsj.col_high;

-            offset_y   = 0;
-            offset_dst = 0;
-            if constexpr (type == GGML_TYPE_NVFP4) {
-                offset_y_scale = 0;
-            }
+            offset_y       = 0;
+            offset_y_scale = 0;
+            offset_dst     = 0;

-            if (jt*J >= col_diff) {
+            if (col_low == col_high) {
                return;
            }

@@ -1115,21 +1098,18 @@ static __global__ void mul_mat_q(
                    break;
                }

-                ids_dst_shared[j] = ids_dst[col_low + jt*J + j];
+                ids_dst_shared[j] = ids_dst[col_low + j];
            }
            __syncthreads();
        }

-        offset_y   += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int));
-        offset_dst += it*I_dst;
-        const float * y_scale_tile = nullptr;
-        if constexpr (type == GGML_TYPE_NVFP4) {
-            offset_y_scale += col_low + jt*J;
-            y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
-        }
+        offset_y       += col_low * (sizeof(block_q8_1_mmq)/sizeof(int));
+        offset_y_scale += col_low;
+        offset_dst     += it*I_dst;

-        const int tile_x_max_i = nrows_x  - it*I_dst - 1;
-        const int tile_y_max_j = col_diff - jt*J - 1;
+        const float * y_scale_tile = type == GGML_TYPE_NVFP4 && y_scale ? y_scale + offset_y_scale : nullptr;
+
+        const int tile_x_i_max = nrows_x - it*I_dst - 1;

        const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I_dst*stride_row_x;

@@ -1141,7 +1121,7 @@ static __global__ void mul_mat_q(
        mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, has_fusion>
            (x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
             glu_op, glu_limit, x_scale_tile, gate_scale_tile, stride_row_x, ncols_y, stride_col_dst,
-             tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z);
+             tile_x_i_max, col_high - col_low, 0, blocks_per_ne00.z);
        return;
    }

@@ -1149,8 +1129,8 @@ static __global__ void mul_mat_q(
    constexpr int blocks_per_iter = ITER_K / qk;

    // kbc == k block continuous, current index in continuous ijk space.
-    int kbc      = int64_t(blockIdx.x)    *(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
-    int kbc_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
+    int kbc      = int64_t(blockIdx.x)    *(nsamples_y.z*nchannels_y.z*ntj.z*nti*blocks_per_ne00.z) / gridDim.x;
+    int kbc_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntj.z*nti*blocks_per_ne00.z) / gridDim.x;

    kbc      -= fastmodulo(kbc,      blocks_per_ne00) % blocks_per_iter;
    kbc_stop -= fastmodulo(kbc_stop, blocks_per_ne00) % blocks_per_iter;
@@ -1160,41 +1140,34 @@ static __global__ void mul_mat_q(
    int kb0_stop  = min(blocks_per_ne00.z, uint32_t(kb0_start + kbc_stop - kbc));
    while (kbc < kbc_stop && kb0_stop == int(blocks_per_ne00.z)) {
        int tmp = fastdiv(kbc, blocks_per_ne00);
-        uint2 tmp2 = fast_div_modulo(tmp, ntx);
+        uint2 tmp2 = fast_div_modulo(tmp, ntj);
        const int jt = tmp2.y;
        tmp = tmp2.x;
        tmp2 = fast_div_modulo(tmp, nchannels_y);
-        const int zt = tmp2.y;
+        int zt = tmp2.y;
        tmp = tmp2.x;
        tmp2 = fast_div_modulo(tmp, nsamples_y);
        const int wt = tmp2.y;
        const int it = tmp2.x;

        // Defaults for regular matrix multiplication:
-        int col_low    = 0;
-        int col_high   = ncols_dst;
-        int col_diff   = ncols_dst;
+        int col_low        = jt*J;
+        int col_high       = min(col_low + J, ncols_dst);
        int offset_y       = wt*stride_sample_y   + zt*stride_channel_y;
+        int offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
        int offset_dst     = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
-        int offset_y_scale;
-        if constexpr (type == GGML_TYPE_NVFP4) {
-            offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
-        } else {
-            GGML_UNUSED(offset_y_scale);
-        }

        if (ids_dst) {
-            col_low  = expert_bounds[zt + 0];
-            col_high = expert_bounds[zt + 1];
-            col_diff = col_high - col_low;
+            const ggml_cuda_mm_id_tile_spec tsj = tile_spec[jt];
+            zt       = tsj.expert;
+            col_low  = tsj.col_low;
+            col_high = tsj.col_high;

-            offset_y   = 0;
-            offset_dst = 0;
-            if constexpr (type == GGML_TYPE_NVFP4) {
-                offset_y_scale = 0;
-            }
+            offset_y       = 0;
+            offset_y_scale = 0;
+            offset_dst     = 0;

-            if (jt*J >= col_diff) {
+            if (col_low == col_high) {
                kbc += blocks_per_ne00.z;
                kbc -= fastmodulo(kbc, blocks_per_ne00);

@@ -1213,21 +1186,18 @@ static __global__ void mul_mat_q(
                    break;
                }

-                ids_dst_shared[j] = ids_dst[col_low + jt*J + j];
+                ids_dst_shared[j] = ids_dst[col_low + j];
            }
            __syncthreads();
        }

-        offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
-        offset_dst += it*I;
-        const float * y_scale_tile = nullptr;
-        if constexpr (type == GGML_TYPE_NVFP4) {
-            offset_y_scale += col_low + jt * J;
-            y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
-        }
+        offset_y       += col_low * (sizeof(block_q8_1_mmq) / sizeof(int));
+        offset_y_scale += col_low;
+        offset_dst     += it*I;

-        const int tile_x_max_i = nrows_x  - it*I - 1;
-        const int tile_y_max_j = col_diff - jt*J - 1;
+        const float * y_scale_tile = type == GGML_TYPE_NVFP4 && y_scale ? y_scale + offset_y_scale : nullptr;
+
+        const int tile_x_i_max = nrows_x - it*I - 1;

        const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;

@@ -1235,7 +1205,7 @@ static __global__ void mul_mat_q(
        mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, false>
            (x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
             glu_op, glu_limit, x_scale, gate_scale, stride_row_x, ncols_y, stride_col_dst,
-             tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
+             tile_x_i_max, col_high - col_low, kb0_start, kb0_stop);

        kbc += blocks_per_ne00.z;
        kbc -= fastmodulo(kbc, blocks_per_ne00);
@@ -1249,41 +1219,34 @@ static __global__ void mul_mat_q(
    }

    int tmp = fastdiv(kbc, blocks_per_ne00);
-    uint2 tmp2 = fast_div_modulo(tmp, ntx);
+    uint2 tmp2 = fast_div_modulo(tmp, ntj);
    const int jt = tmp2.y;
    tmp = tmp2.x;
    tmp2 = fast_div_modulo(tmp, nchannels_y);
-    const int zt = tmp2.y;
+    int zt = tmp2.y;
    tmp = tmp2.x;
    tmp2 = fast_div_modulo(tmp, nsamples_y);
    const int wt = tmp2.y;
    const int it = tmp2.x;

    // Defaults for regular matrix multiplication:
-    int col_low    = 0;
-    int col_high   = ncols_dst;
-    int col_diff   = ncols_dst;
+    int col_low        = jt*J;
+    int col_high       = min(col_low + J, ncols_dst);
    int offset_y       = wt*stride_sample_y   + zt*stride_channel_y;
+    int offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
    int offset_dst     = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
-    int offset_y_scale;
-    if constexpr (type == GGML_TYPE_NVFP4) {
-        offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
-    } else {
-        GGML_UNUSED(offset_y_scale);
-    }

    if (ids_dst) {
-        col_low  = expert_bounds[zt + 0];
-        col_high = expert_bounds[zt + 1];
-        col_diff = col_high - col_low;
+        const ggml_cuda_mm_id_tile_spec tsj = tile_spec[jt];
+        zt       = tsj.expert;
+        col_low  = tsj.col_low;
+        col_high = tsj.col_high;

-        offset_y   = 0;
-        offset_dst = 0;
-        if constexpr (type == GGML_TYPE_NVFP4) {
-            offset_y_scale = 0;
-        }
+        offset_y       = 0;
+        offset_y_scale = 0;
+        offset_dst     = 0;

-        if (jt*J >= col_diff) {
+        if (col_low == col_high) {
            return;
        }

@@ -1302,16 +1265,13 @@ static __global__ void mul_mat_q(
        __syncthreads();
    }

-    offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
-    offset_dst += it*I;
-    const float * y_scale_tile = nullptr;
-    if constexpr (type == GGML_TYPE_NVFP4) {
-        offset_y_scale += col_low + jt * J;
-        y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
-    }
+    offset_y       += col_low * (sizeof(block_q8_1_mmq) / sizeof(int));
+    offset_y_scale += col_low;
+    offset_dst     += it*I;

-    const int tile_x_max_i = nrows_x  - it*I - 1;
-    const int tile_y_max_j = col_diff - jt*J - 1;
+    const float * y_scale_tile = type == GGML_TYPE_NVFP4 && y_scale ? y_scale + offset_y_scale : nullptr;
+
+    const int tile_x_i_max = nrows_x - it*I - 1;

    const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;

@@ -1319,16 +1279,16 @@ static __global__ void mul_mat_q(
    mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, false>
        (x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
         glu_op, glu_limit, x_scale, gate_scale, stride_row_x, ncols_y, stride_col_dst,
-         tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
+         tile_x_i_max, col_high - col_low, kb0_start, kb0_stop);
 }

 template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
 __launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1)/2, 1)
 static __global__ void mul_mat_q_stream_k_fixup(
-        const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ expert_bounds, float * __restrict__ dst,
+        const int32_t * __restrict__ ids_dst, const ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec, float * __restrict__ dst,
        float * __restrict__ tmp_last_tile, const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst,
        const int stride_col_dst, const uint3 nchannels_y, const int stride_channel_dst, const uint3 nsamples_y,
-        const int stride_sample_dst, const uint3 ntx) {
+        const int stride_sample_dst, const uint3 ntj) {
    constexpr int warp_size       = ggml_cuda_get_physical_warp_size();
    constexpr int nwarps          = (ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / 2) / warp_size;
    constexpr int I               = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
@@ -1339,13 +1299,13 @@ static __global__ void mul_mat_q_stream_k_fixup(
    float sum[J / nwarps] = {0.0f};
    const int i = blockIdx.y*warp_size + threadIdx.x;

-    const int nty = (nrows_x + I - 1) / I;
+    const int nti = (nrows_x + I - 1) / I;

    const int bidx0 = blockIdx.x;

    // kbc == k block continuous, current index in continuous ijk space.
-    int kbc0      = int64_t(blockIdx.x)    *(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
-    int kbc0_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
+    int kbc0      = int64_t(blockIdx.x)    *(nsamples_y.z*nchannels_y.z*ntj.z*nti*blocks_per_ne00.z) / gridDim.x;
+    int kbc0_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntj.z*nti*blocks_per_ne00.z) / gridDim.x;

    kbc0      -= fastmodulo(kbc0,      blocks_per_ne00) % blocks_per_iter;
    kbc0_stop -= fastmodulo(kbc0_stop, blocks_per_ne00) % blocks_per_iter;
@@ -1364,7 +1324,7 @@ static __global__ void mul_mat_q_stream_k_fixup(
    int bidx = bidx0 - 1;
    int kbc_stop = kbc0;
    while(true) {
-        int kbc = int64_t(bidx)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
+        int kbc = int64_t(bidx)*(nsamples_y.z*nchannels_y.z*ntj.z*nti*blocks_per_ne00.z) / gridDim.x;
        kbc -= fastmodulo(kbc, blocks_per_ne00) % blocks_per_iter;

        if (kbc == kbc_stop) { // Did not have any data.
@@ -1396,11 +1356,11 @@ static __global__ void mul_mat_q_stream_k_fixup(
    }

    int tmp = fastdiv(kbc0, blocks_per_ne00);
-    uint2 tmp2 = fast_div_modulo(tmp, ntx);
+    uint2 tmp2 = fast_div_modulo(tmp, ntj);
    const int jt = tmp2.y;
    tmp = tmp2.x;
    tmp2 = fast_div_modulo(tmp, nchannels_y);
-    const int zt = tmp2.y;
+    int zt = tmp2.y;
    tmp = tmp2.x;
    tmp2 = fast_div_modulo(tmp, nsamples_y);
    const int wt = tmp2.y;
@@ -1411,7 +1371,7 @@ static __global__ void mul_mat_q_stream_k_fixup(
        dst += offset_dst;

        const int i_max = nrows_x   - it*I - 1;
-        const int j_max = ncols_dst - jt*J - 1;
+        const int j_sup = ncols_dst - jt*J;
        if (fallback && i > i_max) {
            return;
        }
@@ -1420,7 +1380,7 @@ static __global__ void mul_mat_q_stream_k_fixup(
        for (int j0 = 0; j0 < J; j0 += nwarps) {
            const int j = j0 + threadIdx.y;

-            if (j > j_max) {
+            if (j >= j_sup) {
                return;
            }

@@ -1430,12 +1390,14 @@ static __global__ void mul_mat_q_stream_k_fixup(
    }

    __shared__ int ids_dst_shared[J];
-    const int col_low  = expert_bounds[zt + 0];
-    const int col_high = expert_bounds[zt + 1];
-    const int col_diff = col_high - col_low;
+    const ggml_cuda_mm_id_tile_spec tsj = tile_spec[jt];
+    zt                 = tsj.expert;
+    const int col_low  = tsj.col_low;
+    const int col_high = tsj.col_high;
+    const int j_sup    = col_high - col_low;

    for (int j = threadIdx.y*warp_size + threadIdx.x; j < J; j += nwarps*warp_size) {
-        ids_dst_shared[j] = ids_dst[col_low + jt*J + j];
+        ids_dst_shared[j] = ids_dst[col_low + j];
    }
    __syncthreads();

@@ -1443,7 +1405,6 @@ static __global__ void mul_mat_q_stream_k_fixup(
    dst += offset_dst;

    const int i_max = nrows_x  - it*I - 1;
-    const int j_max = col_diff - jt*J - 1;
    if (fallback && i > i_max) {
        return;
    }
@@ -1452,7 +1413,7 @@ static __global__ void mul_mat_q_stream_k_fixup(
    for (int j0 = 0; j0 < J; j0 += nwarps) {
        const int j = j0 + threadIdx.y;

-        if (j > j_max) {
+        if (j >= j_sup) {
            return;
        }

@@ -1461,12 +1422,12 @@ static __global__ void mul_mat_q_stream_k_fixup(
 }

 struct mmq_args {
-    const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; float * dst;
-    const float * y_scale;
+    const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst;
+    const ggml_cuda_mm_id_tile_spec * tile_spec; float * dst; const float * y_scale;
    int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst;
    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 ntj;
    int J_best; // Tile width in ne11(dense)/ne12(MoE) direction to use for optimal performance.
    ggml_cuda_mm_fusion_args_device fusion = {}; // gate != nullptr selects the fused gate/up/GLU path
 };
@@ -1498,39 +1459,38 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
    }

    const int I_dst = has_fusion ? config.I/2 : config.I; // Rows of dst per tile, see mul_mat_q.
-    const int nty  = (args.nrows_x   + I_dst - 1) / I_dst;
-    const int ntx  = (args.ncols_max + config.J - 1) / config.J;
-    const int ntzw = args.nchannels_y * args.nsamples_y;
+    const int nti   = (args.nrows_x + I_dst - 1) / I_dst;
+    const int ntzw  = args.nchannels_y * args.nsamples_y;
    // Column tiles innermost so that the concurrently running blocks share the same tile of x and read it from L2.
-    const dim3 block_nums_xy_tiling(ntx, nty, ntzw);
+    const dim3 block_nums_xy_tiling(args.ntj, nti, ntzw);

-    GGML_ASSERT(args.nchannels_y % args.nchannels_x == 0);
-    GGML_ASSERT(args.nsamples_y  % args.nsamples_x  == 0);
-    const int channel_ratio = args.nchannels_y / args.nchannels_x;
-    const int sample_ratio  = args.nsamples_y  / args.nsamples_x;
+    GGML_ASSERT(args.ids_dst || args.nchannels_y % args.nchannels_x == 0);
+    GGML_ASSERT(                args.nsamples_y  % args.nsamples_x  == 0);
+    const int channel_ratio = args.ids_dst ? 1 : args.nchannels_y / args.nchannels_x;
+    const int sample_ratio  =                    args.nsamples_y  / args.nsamples_x;

    const uint3 blocks_per_ne00_fd = init_fastdiv_values(args.ncols_x / ggml_cuda_type_traits<type>::qk);
-    const uint3 ntx_fd             = init_fastdiv_values(ntx);
+    const uint3 ntj_fd             = init_fastdiv_values(args.ntj);
    const uint3 nchannels_y_fd     = init_fastdiv_values(args.nchannels_y);
    const uint3 nsamples_y_fd      = init_fastdiv_values(args.nsamples_y);
    const uint3 channel_ratio_fd   = init_fastdiv_values(channel_ratio);
    const uint3 sample_ratio_fd    = init_fastdiv_values(sample_ratio);

    if (has_fusion || !config.stream_k) {
-        GGML_ASSERT(nty <= 65535); // gridDim.y limit
+        GGML_ASSERT(nti <= 65535); // gridDim.y limit
        mul_mat_q<type, J, fallback, prec_src1, has_fusion><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
-            (args.x, (const char *) args.fusion.gate, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale, args.fusion.glu_op, args.fusion.glu_limit,
-             (const float *) args.fusion.x_scale, (const float *) args.fusion.gate_scale,
+            (args.x, (const char *) args.fusion.gate, args.y, args.ids_dst, args.tile_spec, args.dst, nullptr,
+             args.y_scale, args.fusion.glu_op, args.fusion.glu_limit, (const float *) args.fusion.x_scale, (const float *) args.fusion.gate_scale,
             blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
             channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
             sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
-             ntx_fd);
+             ntj_fd);
        return;
    }

    // For the stream-k kernel it is possible to run it with tiling by setting the number of CUDA blocks equal to the number of tiles.
    // This is worthwhile if the efficiency of tiling is high and skipping the fixup kernel is more important.
-    const int ntiles_dst = ntx * nty * ntzw;
+    const int ntiles_dst = args.ntj * nti * ntzw;
    const int tiles_nwaves = (ntiles_dst + nsm - 1) / nsm;
    const int tiles_efficiency_percent = 100 * ntiles_dst / (nsm*tiles_nwaves);

@@ -1551,12 +1511,12 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
    const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);

    mul_mat_q<type, J, fallback, prec_src1, has_fusion><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
-        (args.x, (const char *) args.fusion.gate, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale, args.fusion.glu_op, args.fusion.glu_limit,
-         (const float *) args.fusion.x_scale, (const float *) args.fusion.gate_scale,
+        (args.x, (const char *) args.fusion.gate, args.y, args.ids_dst, args.tile_spec, args.dst, tmp_fixup.ptr,
+         args.y_scale, args.fusion.glu_op, args.fusion.glu_limit, (const float *) args.fusion.x_scale, (const float *) args.fusion.gate_scale,
         blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
         channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
         sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
-         ntx_fd);
+         ntj_fd);

    if (!fixup_needed) {
        return;
@@ -1564,9 +1524,9 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a

    CUDA_CHECK(cudaGetLastError());
    mul_mat_q_stream_k_fixup<type, J, fallback, prec_src1><<<block_nums_fixup, block_dims_fixup, 0, stream>>>
-        (args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst,
+        (args.ids_dst, args.tile_spec, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst,
         args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst,
-         ntx_fd);
+         ntj_fd);
 }

 template <ggml_type type, bool fallback, ggml_prec prec_src1, bool has_fusion>
@@ -285,6 +285,7 @@ static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_rdna4(ggml_type
 int get_mmvq_mmid_max_batch(ggml_type type, int cc) {
    // NVIDIA: P100, Volta, Ada Lovelace, and Blackwell always use MMVQ for MUL_MAT_ID.
    if (GGML_CUDA_CC_IS_NVIDIA(cc)) {
+        cc = ggml_cuda_highest_compiled_arch(cc);
        if (cc == GGML_CUDA_CC_PASCAL || cc == GGML_CUDA_CC_VOLTA || cc >= GGML_CUDA_CC_ADA_LOVELACE) {
            return MMVQ_MAX_BATCH_SIZE;
        }
@@ -256,7 +256,17 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
    // NOTE: if you change conditions here, be sure to update the corresponding supports_op condition!
    if (src3_nb1 == sizeof(float)) {
        // Mamba-2
-        if (d_state == 96) {
+        if (d_state == 64) { // Zamba2
+            constexpr int threads   = 64;
+            constexpr int num_warps = threads/WARP_SIZE;
+
+            const dim3 blocks((n_head * head_dim + (num_warps - 1)) / num_warps, n_seq, 1);
+            const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, 0, stream);
+            ggml_cuda_kernel_launch(ssm_scan_f32_group<64/WARP_SIZE, 64>, launch_params,
+                    src0, src1, src2, src3, src4, src5, src6, dst,
+                    src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
+                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_base, s_slot_bytes, n_head, head_dim, n_group, n_tok, K);
+        } else if (d_state == 96) {
            constexpr int threads   = 96;
            constexpr int num_warps = threads/WARP_SIZE;

@@ -287,7 +297,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
                    src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_base, s_slot_bytes, n_head, head_dim, n_group, n_tok, K);
        } else {
-            GGML_ABORT("doesn't support d_state!=(96, 128 or 256).");
+            GGML_ABORT("doesn't support d_state!=(64, 96, 128 or 256).");
        }
    } else {
        // Mamba-1
@@ -88,18 +88,18 @@ static __device__ __forceinline__ uint32_t top_k_float_to_ordered(float value) {
    return bits ^ mask;
 }

-struct top_k_radix_state {
+// one state per cache line, so the gather atomics of different rows do not compete
+struct alignas(128) top_k_radix_state {
    uint32_t prefix;
    uint32_t prefix_mask;
    int rank;
    int greater_count;
-    int equal_count;
 };

 static __global__ void top_k_radix_init(top_k_radix_state * states, int nrows, int k) {
    const int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < nrows) {
-        states[row] = {0, 0, k, 0, 0};
+        states[row] = {0, 0, k, 0};
    }
 }

@@ -109,6 +109,7 @@ static __global__ void top_k_radix_histogram(
        const top_k_radix_state * __restrict__ states,
        int * __restrict__ block_histograms,
        int ncols,
+        int block_cols,
        int blocks_per_row,
        int shift) {
    constexpr int NBINS = 1 << RADIX_BITS;
@@ -117,15 +118,15 @@ static __global__ void top_k_radix_histogram(
    const int row_block = blockIdx.x % blocks_per_row;
    const int tid = threadIdx.x;
    const float * row_src = src + (size_t) row * ncols;
+    const int64_t col_begin = (int64_t) row_block * block_cols;
+    const int64_t col_end   = min(col_begin + block_cols, (int64_t) ncols);
    __shared__ int histogram[NBINS];

    histogram[tid] = 0;
    __syncthreads();

    const top_k_radix_state state = states[row];
-    for (int64_t col = row_block * BLOCK_SIZE + tid;
-         col < ncols;
-         col += blocks_per_row * BLOCK_SIZE) {
+    for (int64_t col = col_begin + tid; col < col_end; col += BLOCK_SIZE) {
        const uint32_t key = top_k_float_to_ordered(row_src[col]);
        if ((key & state.prefix_mask) == state.prefix) {
            atomicAdd(&histogram[(key >> shift) & (NBINS - 1)], 1);
@@ -170,43 +171,91 @@ static __global__ void top_k_radix_select(
    }
 }

-static __global__ void top_k_radix_reset_counters(top_k_radix_state * states, int nrows) {
-    const int row = blockIdx.x * blockDim.x + threadIdx.x;
-    if (row < nrows) {
-        states[row].greater_count = 0;
-        states[row].equal_count = 0;
-    }
-}
-
-template<int BLOCK_SIZE>
+template<int BLOCK_SIZE, int RADIX_BITS>
 static __global__ void top_k_radix_gather(
        const float * __restrict__ src,
        int * __restrict__ dst,
+        const int * __restrict__ block_histograms,
        top_k_radix_state * __restrict__ states,
        int ncols,
+        int block_cols,
        int k,
        int blocks_per_row) {
+    constexpr int NBINS     = 1 << RADIX_BITS;
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+    constexpr int nwarps    = BLOCK_SIZE / warp_size;
+    using lane_mask_t = std::conditional_t<warp_size == 64, uint64_t, uint32_t>;
+    // HIP requires a 64-bit mask also for a 32-wide warp
+    constexpr uint64_t ballot_mask = ~(lane_mask_t) 0;
+    const auto popc = [](const lane_mask_t mask) {
+        if constexpr (warp_size == 64) {
+            return __popcll(mask);
+        } else {
+            return __popc(mask);
+        }
+    };
+    __shared__ int warp_ties[nwarps];
+
    const int row = blockIdx.x / blocks_per_row;
    const int row_block = blockIdx.x % blocks_per_row;
-    const int tid = threadIdx.x;
+    const int lane = threadIdx.x % warp_size;
+    const int warp = threadIdx.x / warp_size;
    const float * row_src = src + (size_t) row * ncols;
    int * row_dst = dst + (size_t) row * k;
    top_k_radix_state * state = &states[row];
+    const uint32_t threshold = state->prefix;
+    const int rank = state->rank;

-    for (int64_t col = row_block * BLOCK_SIZE + tid;
-         col < ncols;
-         col += blocks_per_row * BLOCK_SIZE) {
+    // each warp reads a contiguous part of the contiguous range of the block
+    const int     warp_cols  = block_cols / nwarps;
+    const int64_t warp_begin = (int64_t) row_block * block_cols + warp * warp_cols;
+    const int64_t warp_end   = min(warp_begin + warp_cols, (int64_t) ncols);
+
+    int nties = 0;
+    for (int64_t col = warp_begin + lane; col < warp_end; col += warp_size) {
        const uint32_t key = top_k_float_to_ordered(row_src[col]);
-        if (key > state->prefix) {
+        if (key > threshold) {
            const int pos = atomicAdd(&state->greater_count, 1);
            row_dst[pos] = col;
-        } else if (key == state->prefix) {
-            const int pos = atomicAdd(&state->equal_count, 1);
-            if (pos < state->rank) {
-                row_dst[k - state->rank + pos] = col;
-            }
+        } else if (key == threshold) {
+            nties++;
        }
    }
+
+    // the ties fill the last rank slots in column order, the lowest columns win.
+    // the histograms of the last pass count the ties of each block in the bin of the threshold.
+    const int * block_ties = block_histograms + (size_t) row * blocks_per_row * NBINS + (threshold & (NBINS - 1));
+    int pos = 0;
+    for (int b = lane; b < row_block; b += warp_size) {
+        pos += block_ties[b * NBINS];
+    }
+    pos = warp_reduce_sum<warp_size>(pos);
+    if (pos >= rank || block_ties[row_block * NBINS] == 0) {
+        return; // the same result in all threads of the block
+    }
+
+    nties = warp_reduce_sum<warp_size>(nties);
+    if (lane == 0) {
+        warp_ties[warp] = nties;
+    }
+    __syncthreads();
+    for (int w = 0; w < warp; ++w) {
+        pos += warp_ties[w];
+    }
+
+    // pos is the same in all lanes of the warp, so the lanes leave the loop together
+    const int pos_end = min(pos + nties, rank);
+    for (int64_t tile = warp_begin; tile < warp_end && pos < pos_end; tile += warp_size) {
+        const int64_t col    = tile + lane;
+        const bool    is_tie = col < warp_end && top_k_float_to_ordered(row_src[col]) == threshold;
+
+        const lane_mask_t ties    = (lane_mask_t) __ballot_sync(ballot_mask, is_tie);
+        const int         tie_pos = pos + popc(ties & (((lane_mask_t) 1 << lane) - 1));
+        if (is_tie && tie_pos < pos_end) {
+            row_dst[k - rank + tie_pos] = col;
+        }
+        pos += popc(ties);
+    }
 }

 static void top_k_radix_cuda(
@@ -216,6 +265,8 @@ static void top_k_radix_cuda(
    constexpr int RADIX_BITS = 8;
    constexpr int NBINS = 1 << RADIX_BITS;
    const int blocks_per_row = (int) std::min<int64_t>(((int64_t) ncols + 1023) / 1024, 64);
+    // each block of a row reads a contiguous range of whole tiles
+    const int block_cols = (int) (((int64_t) ncols + blocks_per_row * BLOCK_SIZE - 1) / (blocks_per_row * BLOCK_SIZE) * BLOCK_SIZE);

    // chunk the rows to bound the histogram memory to 64 MB
    const int64_t chunk_nrows = ggml_cuda_chunk_nrows((size_t) blocks_per_row * NBINS * sizeof(int), nrows);
@@ -234,16 +285,14 @@ static void top_k_radix_cuda(
        for (int shift = 32 - RADIX_BITS; shift >= 0; shift -= RADIX_BITS) {
            top_k_radix_histogram<BLOCK_SIZE, RADIX_BITS>
                <<<row_grid, BLOCK_SIZE, 0, stream>>>(
-                    src, states, histograms, ncols, blocks_per_row, shift);
+                    src, states, histograms, ncols, block_cols, blocks_per_row, shift);
            top_k_radix_select<BLOCK_SIZE, RADIX_BITS>
                <<<iter_nrows, BLOCK_SIZE, 0, stream>>>(histograms, states, blocks_per_row, shift);
        }

-        top_k_radix_reset_counters
-            <<<(iter_nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, iter_nrows);
-        top_k_radix_gather<BLOCK_SIZE>
+        top_k_radix_gather<BLOCK_SIZE, RADIX_BITS>
            <<<row_grid, BLOCK_SIZE, 0, stream>>>(
-                src, dst, states, ncols, k, blocks_per_row);
+                src, dst, histograms, states, ncols, block_cols, k, blocks_per_row);

        src += (size_t) ncols * iter_nrows;
        dst += (size_t) k     * iter_nrows;
@@ -217,6 +217,9 @@ __global__ void topk_moe_cuda(const float *         logits,
            const int expert = i + threadIdx.x;
            selection_wt[i / WARP_SIZE] =
                (n_experts % WARP_SIZE == 0 || expert < n_experts) ? wt[i / WARP_SIZE] + bias[expert] : -INFINITY;
+            if (__isnanf(selection_wt[i / WARP_SIZE])) {
+                selection_wt[i / WARP_SIZE] = -FLT_MAX;
+            }
        }
    }

@@ -1523,7 +1523,7 @@ bool ggml_et_op_flash_attn_ext(ggml_backend_et_device_context * dev_ctx, const g
    }

    const ggml_prec prec = ggml_flash_attn_ext_get_prec(node);
-    if (prec != GGML_PREC_F32 && prec != GGML_PREC_DEFAULT) {
+    if (prec != GGML_PREC_F32 && prec != GGML_PREC_BF16 && prec != GGML_PREC_DEFAULT) {
        GGML_LOG_ERROR("ET: FLASH_ATTN_EXT baseline kernel only supports F32 precision\n");
        return false;
    }
@@ -1299,7 +1299,7 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm
                const bool me_eligible = op->src[1]->type == GGML_TYPE_F16 && op->src[2]->type == GGML_TYPE_F16 &&
                                         (op->src[0]->ne[0] % 32) == 0;

-                supported = me_eligible && mask_ok && (prec == GGML_PREC_F32 || prec == GGML_PREC_DEFAULT) &&
+                supported = me_eligible && mask_ok && (prec == GGML_PREC_F32 || prec == GGML_PREC_BF16 || prec == GGML_PREC_DEFAULT) &&
                            max_bias == 0.0f && logit_softcap == 0.0f && op->src[0]->nb[0] == sizeof(float) &&
                            op->src[1]->nb[0] == k_elem && op->src[2]->nb[0] == v_elem && op->nb[0] == sizeof(float) &&
                            op->src[0]->ne[0] == op->src[1]->ne[0] &&  // dk matches
@@ -90,7 +90,7 @@ size_t opt_vmem       = HTP_OP_MAX_VMEM_DEFAULT;  // max available va space for
 int    opt_etm        = 0;
 int    opt_verbose    = 0;
 int    opt_profile    = 0; // profiling mode (0-disabled, 1-basic, 2-pmu)
-bool   opt_hostbuf    = false;
+bool   opt_hostbuf    = true;
 bool   opt_dma64      = false;
 int    opt_power_mode = HTP_POWER_MODE_BURST; // DCVS/HMX power profile requested at session start

@@ -150,7 +150,7 @@ int    opt_opqueue   = 8;    // max number of pending batches
 int    opt_optrace   = 0;    // trace buffer size per thread (0 means default)
 int    opt_oppoll    = 0;    // polling for batch completions
 int    opt_opfusion  = 1;    // enable/disable op fusion
-int    opt_rt_repack = 1;    // runtime repacking: 0 = disabled, 1 = quantized only (default), 2 = always
+int    opt_rt_repack = 1;    // runtime repacking: 0 = disabled, 1 = all tileable weights (default)

 static std::regex* opt_opfilter = NULL; // regex of ops to not claim

@@ -242,10 +242,11 @@ static void ggml_hexagon_dump_batch_prof(const char * sess_name, const htp_opbat

    char evt_str[256] = "----";
    if (opt_profile == 3) {
-        snprintf(evt_str, sizeof(evt_str), "evt-cnt %u,%u,%u,%u,%u,%u,%u,%u,%u,%u,%u",
+        static_assert(HTP_OP_MAX_NTHREADS == 10, "current implementation assumes 10 trace counters");
+        snprintf(evt_str, sizeof(evt_str), "evt-cnt %u,%u,%u,%u,%u,%u,%u,%u,%u,%u",
                rsp.n_traces[0], rsp.n_traces[1], rsp.n_traces[2], rsp.n_traces[3],
                rsp.n_traces[4], rsp.n_traces[5], rsp.n_traces[6], rsp.n_traces[7],
-                rsp.n_traces[8], rsp.n_traces[9], rsp.n_traces[10]);
+                rsp.n_traces[8], rsp.n_traces[9]);
    }

    GGML_LOG_DEBUG("ggml-hex: %s profile-op OPBATCH|----|n-ops %u|%s|----|----|usec %u cycles %llu start %llu mhz %.1f\n",
@@ -253,16 +254,11 @@ static void ggml_hexagon_dump_batch_prof(const char * sess_name, const htp_opbat
 }

 static void ggml_hexagon_dump_trace_events(const char * sess_name, const htp_opbatch_rsp & rsp,
-                                           const htp_trace_desc * trace_events, uint32_t n_traces) {
+                                           const htp_trace_desc * trace_events, uint32_t n_threads, uint32_t n_traces) {
    if (opt_profile == 3 && trace_events) {
-        uint32_t valid_cnt[HTP_MAX_NTHREADS + 1] = {0};
-        for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
-            uint32_t count = rsp.n_traces[t];
-            valid_cnt[t] = count > n_traces ? n_traces : count;
-        }
-
-        for (uint32_t t = 0; t <= HTP_MAX_NTHREADS; t++) {
-            for (uint32_t idx = 0; idx < valid_cnt[t]; idx++) {
+        for (uint32_t t = 0; t < n_threads; t++) {
+            uint32_t count = std::min(rsp.n_traces[t], n_traces);
+            for (uint32_t idx = 0; idx < count; idx++) {
                const auto & e = trace_events[t * n_traces + idx];
                bool is_stop = (e.info & 0x8000) != 0;
                uint16_t info = e.info & 0x7FFF;
@@ -393,7 +389,7 @@ struct ggml_hexagon_session {

 const char * htp_session_name(const struct ggml_hexagon_session * sess) { return sess->c_name(); }
 size_t       htp_session_vtcm_size(const struct ggml_hexagon_session * sess) { return sess->vtcm_size; }
-uint32_t     htp_session_n_threads(const struct ggml_hexagon_session * sess) { return sess->n_threads; }
+uint32_t     htp_session_n_hvx(const struct ggml_hexagon_session * sess) { return sess->n_hvx; }
 uint32_t     htp_session_n_hmx(const struct ggml_hexagon_session * sess) { return sess->n_hmx; }
 uint32_t     htp_session_mdev_count(const struct ggml_hexagon_session * sess) { return sess->mdev.count; }

@@ -463,7 +459,6 @@ struct ggml_hexagon_rpcmem_block {
 struct ggml_hexagon_shared_buffer {
    ggml_hexagon_session *                     sess;
    std::shared_ptr<ggml_hexagon_rpcmem_block> mem;
-    std::vector<htp_tensor_extra *>            tensor_extra;
    bool     mapped;
    bool     pinned;
    bool     extended;
@@ -486,6 +481,9 @@ struct ggml_hexagon_shared_buffer {
        if (err != 0) {
            GGML_LOG_ERROR("ggml-hex: %s buffer mapping failed : domain_id %d size %zu fd %d error 0x%08x\n", sess->c_name(),
                    sess->domain_id, size(), fd(), (unsigned) err);
+            if (!this->pinned) {
+                GGML_LOG_ERROR("ggml-hex: %s possibly out of IOVA space, try --no-host\n", sess->c_name());
+            }
            throw std::runtime_error("ggml-hex: fastrpc_mmap failed (see log for details)");
        }

@@ -561,9 +559,6 @@ struct ggml_hexagon_shared_buffer {

    ~ggml_hexagon_shared_buffer() {
        free();
-        for (auto * extra : tensor_extra) {
-            delete extra;
-        }
    }
 };

@@ -622,11 +617,9 @@ inline void ggml_hexagon_session::free_fence(void * ptr, uint32_t n_slots) {
 }

 inline ggml_tensor ggml_hexagon_session::make_fence_tensor(const volatile void * slot, uint32_t n_slots) const {
-    static htp_tensor_extra fence_extra { {}, 0, HTP_TENSOR_FENCE };
-
    ggml_tensor t {};
    t.buffer = &this->fence_buf->backend_buffer;
-    t.extra  = &fence_extra;
+    t.extra  = (void *) (uintptr_t) HTP_TENSOR_FENCE;
    t.data   = const_cast<void *>(slot);
    t.type   = GGML_TYPE_I8;
    t.ne[0]  = HTP_FENCE_SLOT_SIZE;
@@ -659,10 +652,7 @@ static enum ggml_status ggml_backend_hexagon_buffer_init_tensor(ggml_backend_buf
    HEX_VERBOSE("ggml-hex: %s init-tensor %s : base %p data %p nbytes %zu\n", sess->c_name(),
                tensor->name, (void *) sbuf->base(), tensor->data, ggml_nbytes(tensor));

-    auto extra = new htp_tensor_extra();
-    sbuf->tensor_extra.push_back(extra);
-
-    tensor->extra = extra;
+    tensor->extra = (void *) (uintptr_t) 0;

    return GGML_STATUS_SUCCESS;
 }
@@ -672,42 +662,19 @@ static void ggml_backend_hexagon_buffer_set_tensor(ggml_backend_buffer_t buffer,
                                                   const void *          data,
                                                   size_t                offset,
                                                   size_t                size) {
-    auto extra = (htp_tensor_extra *)           tensor->extra;
-    auto sbuf  = (ggml_hexagon_shared_buffer *) buffer->context;
-    auto sess  = sbuf->sess;
+    auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context;
+    auto sess = sbuf->sess;
+
+    sess->flush_sync();

    if (ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
-        extra->flags |= HTP_TENSOR_WEIGHT;
-        if (htp_is_tiled_weight_tensor(tensor)) {
-            extra->flags |= HTP_TENSOR_TILED;
-        }
+        ggml_hexagon_tensor_set_flags(tensor, HTP_TENSOR_WEIGHT);
    }

    HEX_VERBOSE("ggml-hex: %s set-tensor %s : data %p offset %zu size %zu usage %d flags 0x%x\n",
-        sess->c_name(), tensor->name, data, offset, size, (int) buffer->usage, extra->flags);
+        sess->c_name(), tensor->name, data, offset, size, (int) buffer->usage, ggml_hexagon_tensor_get_flags(tensor));

-    if ((extra->flags & HTP_TENSOR_TILED) == 0) {
-        memcpy((char *) tensor->data + offset, data, size);
-        return;
-    }
-
-    if (offset == 0 && size == ggml_nbytes(tensor) && extra->shadow_buf.empty()) {
-        htp_repack_tensor_tiled(tensor, data, size);
-        return;
-    }
-
-    if (extra->shadow_buf.size() < ggml_nbytes(tensor)) {
-        extra->shadow_buf.resize(ggml_nbytes(tensor));
-    }
-    memcpy(extra->shadow_buf.data() + offset, data, size);
-    extra->shadow_size += size;
-
-    if (extra->shadow_size >= ggml_nbytes(tensor)) {
-        htp_repack_tensor_tiled(tensor, extra->shadow_buf.data(), extra->shadow_buf.size());
-        extra->shadow_buf.clear();
-        extra->shadow_buf.shrink_to_fit();
-        extra->shadow_size = 0;
-    }
+    memcpy((char *) tensor->data + offset, data, size);
 }

 static void ggml_backend_hexagon_buffer_get_tensor(ggml_backend_buffer_t buffer,
@@ -715,14 +682,14 @@ static void ggml_backend_hexagon_buffer_get_tensor(ggml_backend_buffer_t buffer,
                                                   void *                data,
                                                   size_t                offset,
                                                   size_t                size) {
-    auto extra = (htp_tensor_extra *)           tensor->extra;
-    auto sbuf  = (ggml_hexagon_shared_buffer *) buffer->context;
-    auto sess  = sbuf->sess;
+    auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context;
+    auto sess = sbuf->sess;
+    const uint32_t flags = ggml_hexagon_tensor_get_flags(tensor);

    HEX_VERBOSE("ggml-hex: %s get-tensor %s : data %p offset %zu size %zu usage %d flags 0x%x\n",
-            sess->c_name(), tensor->name, data, offset, size, (int) buffer->usage, extra->flags);
+            sess->c_name(), tensor->name, data, offset, size, (int) buffer->usage, flags);

-    if ((extra->flags & HTP_TENSOR_TILED) == 0) {
+    if ((flags & HTP_TENSOR_TILED) == 0) {
        memcpy(data, (const char *) tensor->data + offset, size);
        return;
    }
@@ -742,47 +709,27 @@ static bool ggml_backend_hexagon_buffer_cpy_tensor(ggml_backend_buffer_t      bu
 }

 static void ggml_backend_hexagon_buffer_set_tensor_2d(ggml_backend_buffer_t buffer,
-                                                          ggml_tensor *         tensor,
-                                                          const void *          data,
-                                                          size_t                offset,
-                                                          size_t                size,
-                                                          size_t                n_copies,
-                                                          size_t                stride_tensor,
-                                                          size_t                stride_data) {
-    auto extra = (htp_tensor_extra *)           tensor->extra;
-    auto sbuf  = (ggml_hexagon_shared_buffer *) buffer->context;
-    auto sess  = sbuf->sess;
+                                                      ggml_tensor *         tensor,
+                                                      const void *          data,
+                                                      size_t                offset,
+                                                      size_t                size,
+                                                      size_t                n_copies,
+                                                      size_t                stride_tensor,
+                                                      size_t                stride_data) {
+    auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context;
+    auto sess = sbuf->sess;
+
+    sess->flush_sync();

    if (ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
-        extra->flags |= HTP_TENSOR_WEIGHT;
-        if (htp_is_tiled_weight_tensor(tensor)) {
-            extra->flags |= HTP_TENSOR_TILED;
-        }
+        ggml_hexagon_tensor_set_flags(tensor, HTP_TENSOR_WEIGHT);
    }

    HEX_VERBOSE("ggml-hex: %s set-tensor-2d %s : data %p offset %zu size %zu n_copies %zu stride_tensor %zu stride_data %zu usage %d flags 0x%x\n",
-                sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, (int) buffer->usage, extra->flags);
+                sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, (int) buffer->usage, ggml_hexagon_tensor_get_flags(tensor));

-    if ((extra->flags & HTP_TENSOR_TILED) == 0) {
-        for (size_t i = 0; i < n_copies; i++) {
-            memcpy((uint8_t *) tensor->data + offset + i * stride_tensor, (const uint8_t *) data + i * stride_data, size);
-        }
-        return;
-    }
-
-    if (extra->shadow_buf.size() < ggml_nbytes(tensor)) {
-        extra->shadow_buf.resize(ggml_nbytes(tensor));
-    }
    for (size_t i = 0; i < n_copies; i++) {
-        memcpy(extra->shadow_buf.data() + offset + i * stride_tensor, (const uint8_t *) data + i * stride_data, size);
-    }
-    extra->shadow_size += n_copies * size;
-
-    if (extra->shadow_size >= ggml_nbytes(tensor)) {
-        htp_repack_tensor_tiled(tensor, extra->shadow_buf.data(), extra->shadow_buf.size());
-        extra->shadow_buf.clear();
-        extra->shadow_buf.shrink_to_fit();
-        extra->shadow_size = 0;
+        memcpy((uint8_t *) tensor->data + offset + i * stride_tensor, (const uint8_t *) data + i * stride_data, size);
    }
 }

@@ -794,14 +741,14 @@ static void ggml_backend_hexagon_buffer_get_tensor_2d(ggml_backend_buffer_t buff
                                                      size_t                n_copies,
                                                      size_t                stride_tensor,
                                                      size_t                stride_data) {
-    auto extra = (htp_tensor_extra *)           tensor->extra;
-    auto sbuf  = (ggml_hexagon_shared_buffer *) buffer->context;
-    auto sess  = sbuf->sess;
+    auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context;
+    auto sess = sbuf->sess;
+    const uint32_t flags = ggml_hexagon_tensor_get_flags(tensor);

    HEX_VERBOSE("ggml-hex: %s get-tensor-2d %s : data %p offset %zu size %zu n_copies %zu stride_tensor %zu stride_data %zu usage %d\n",
                sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, (int) buffer->usage);

-    if ((extra->flags & HTP_TENSOR_TILED) == 0) {
+    if ((flags & HTP_TENSOR_TILED) == 0) {
        for (size_t i = 0; i < n_copies; i++) {
            memcpy((uint8_t *)data + i * stride_data, (const uint8_t *)tensor->data + offset + i * stride_tensor, size);
        }
@@ -825,6 +772,29 @@ static void ggml_backend_hexagon_buffer_get_tensor_2d(ggml_backend_buffer_t buff
    }
 }

+static void ggml_backend_hexagon_buffer_set_tensor_n(ggml_backend_buffer_t buffer,
+                                                     struct ggml_tensor * const * tensors,
+                                                     const void * const * data,
+                                                     const size_t * offsets,
+                                                     const size_t * sizes,
+                                                     size_t n_tensors) {
+    const bool is_weights = ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS;
+    auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context;
+    auto sess = sbuf->sess;
+
+    sess->flush_sync();
+
+    for (size_t i = 0; i < n_tensors; i++) {
+        if (sizes[i] == 0) {
+            continue;
+        }
+        if (is_weights) {
+            ggml_hexagon_tensor_set_flags(tensors[i], HTP_TENSOR_WEIGHT);
+        }
+        memcpy((char *) tensors[i]->data + offsets[i], data[i], sizes[i]);
+    }
+}
+
 static void ggml_backend_hexagon_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
    auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context;
    auto sess = sbuf->sess;
@@ -844,7 +814,7 @@ static ggml_backend_buffer_i ggml_backend_hexagon_buffer_interface = {
    /* .cpy_tensor      = */ ggml_backend_hexagon_buffer_cpy_tensor,
    /* .clear           = */ ggml_backend_hexagon_buffer_clear,
    /* .reset           = */ NULL,
-    /* .set_tensor_n    = */ NULL,
+    /* .set_tensor_n    = */ ggml_backend_hexagon_buffer_set_tensor_n,
    /* .get_tensor_n    = */ NULL,
    /* .cpy_tensor_n    = */ NULL,
 };
@@ -856,8 +826,12 @@ static void ggml_backend_hexagon_host_buffer_set_tensor(ggml_backend_buffer_t bu
                                                        const void *          data,
                                                        size_t                offset,
                                                        size_t                size) {
+    auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context;
+    auto sess = sbuf->sess;
+
+    sess->flush_sync();
+
    memcpy((char *) tensor->data + offset, data, size);
-    GGML_UNUSED(buffer);
 }

 static void ggml_backend_hexagon_host_buffer_get_tensor(ggml_backend_buffer_t buffer,
@@ -869,6 +843,25 @@ static void ggml_backend_hexagon_host_buffer_get_tensor(ggml_backend_buffer_t bu
    GGML_UNUSED(buffer);
 }

+static void ggml_backend_hexagon_host_buffer_set_tensor_n(ggml_backend_buffer_t buffer,
+                                                          struct ggml_tensor * const * tensors,
+                                                          const void * const * data,
+                                                          const size_t * offsets,
+                                                          const size_t * sizes,
+                                                          size_t n_tensors) {
+    auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context;
+    auto sess = sbuf->sess;
+
+    sess->flush_sync();
+
+    for (size_t i = 0; i < n_tensors; i++) {
+        if (sizes[i] == 0) {
+            continue;
+        }
+        memcpy((char *) tensors[i]->data + offsets[i], data[i], sizes[i]);
+    }
+}
+
 static ggml_backend_buffer_i ggml_backend_hexagon_host_buffer_interface = {
    /* .free_buffer     = */ ggml_backend_hexagon_buffer_free_buffer,
    /* .get_base        = */ ggml_backend_hexagon_buffer_get_base,
@@ -881,7 +874,7 @@ static ggml_backend_buffer_i ggml_backend_hexagon_host_buffer_interface = {
    /* .cpy_tensor      = */ ggml_backend_hexagon_buffer_cpy_tensor,
    /* .clear           = */ ggml_backend_hexagon_buffer_clear,
    /* .reset           = */ NULL,
-    /* .set_tensor_n    = */ NULL,
+    /* .set_tensor_n    = */ ggml_backend_hexagon_host_buffer_set_tensor_n,
    /* .get_tensor_n    = */ NULL,
    /* .cpy_tensor_n    = */ NULL,
 };
@@ -896,7 +889,7 @@ static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer(
            ggml_backend_buffer_type_t buffer_type, size_t size) {
    auto dev_ctx = static_cast<ggml_backend_hexagon_buffer_type_context *>(buffer_type->context)->dev_ctx;
    auto sess    = dev_ctx->session();
-    if (sess && sess->max_vmem && size > sess->max_vmem) {
+    if (size > sess->max_vmem) {
        GGML_LOG_ERROR("ggml-hex: %s buffer size %zu exceeds max_vmem %zu\n",
                       dev_ctx->c_name(), size, sess->max_vmem);
        return nullptr;
@@ -914,7 +907,7 @@ static ggml_backend_buffer_t ggml_backend_hexagon_host_buffer_type_alloc_buffer(
            ggml_backend_buffer_type_t buffer_type, size_t size) {
    auto dev_ctx = static_cast<ggml_backend_hexagon_buffer_type_context *>(buffer_type->context)->dev_ctx;
    auto sess    = dev_ctx->session();
-    if (sess && sess->max_vmem && size > sess->max_vmem) {
+    if (size > sess->max_vmem) {
        GGML_LOG_ERROR("ggml-hex: %s host buffer size %zu exceeds max_vmem %zu\n",
                       dev_ctx->c_name(), size, sess->max_vmem);
        return nullptr;
@@ -934,7 +927,7 @@ static size_t ggml_backend_hexagon_buffer_type_get_alignment(ggml_backend_buffer
 }

 static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const struct ggml_tensor * t) {
-    if (!ggml_backend_buft_is_host(buft) && htp_is_tiled_weight_tensor(t)) {
+    if (opt_rt_repack != 0 && !ggml_backend_buft_is_host(buft) && ggml_op_is_empty(t->op) && htp_is_tiled_weight_tensor(t)) {
        int64_t ne0 = hex_round_up(t->ne[0], 32);
        int64_t ne1 = hex_round_up(t->ne[1], 32);
        int64_t ne2 = t->ne[2];
@@ -1212,8 +1205,7 @@ struct ggml_hexagon_opbatch {
    std::vector<htp_tensor>          h_tens;    // htp tensor descriptors
    std::vector<htp_op_desc>         h_ops;     // htp op descriptors

-    std::unordered_map<const ggml_tensor*, int> t_map; // tensor ptr  to index
-    std::unordered_multimap<void*, int>         d_map; // tensor data to index
+    std::unordered_map<uint64_t, int> t_map; // tensor hash to index

    unsigned int n_bufs;     // num buffers in the batch
    unsigned int n_tens;     // num tensors ...
@@ -1232,7 +1224,6 @@ struct ggml_hexagon_opbatch {
        b_vmem = 0;

        t_map.clear();
-        d_map.clear();
        ops.resize(n_ops_max);
    }

@@ -1252,7 +1243,6 @@ struct ggml_hexagon_opbatch {
        h_ops.resize(n_ops_max);

        t_map.reserve(n_tens_max);
-        d_map.reserve(n_tens_max);

        GGML_LOG_INFO("ggml-hex: %s op batching: n-bufs %u n-tensors %u n-ops %u vmem %zu\n",
                sess->c_name(), n_bufs_max, n_tens_max, n_ops_max, b_vmem_max);
@@ -1298,7 +1288,7 @@ struct ggml_hexagon_opbatch {
    }

    static void init_tensor_shape(htp_tensor & h, const ggml_tensor * t) {
-        auto extra = (htp_tensor_extra *) t->extra;
+        const uint32_t flags = ggml_hexagon_tensor_get_flags(t);

        h.type  = t->type;
        h.ne[0] = t->ne[0];
@@ -1306,7 +1296,7 @@ struct ggml_hexagon_opbatch {
        h.ne[2] = t->ne[2];
        h.ne[3] = t->ne[3];

-        if (extra->flags & HTP_TENSOR_TILED) {
+        if (flags & HTP_TENSOR_TILED) {
            if (t->type == GGML_TYPE_F32 || t->type == GGML_TYPE_BF16) {
                h.type  = GGML_TYPE_F16;
                h.nb[0] = sizeof(ggml_fp16_t);
@@ -1327,31 +1317,42 @@ struct ggml_hexagon_opbatch {
            h.size  = ggml_nbytes(t);
        }

-        h.flags = extra->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_TILED | HTP_TENSOR_FENCE);
+        h.flags = flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_TILED | HTP_TENSOR_FENCE);
    }

-    static bool same_shape(const htp_tensor & a, const htp_tensor & b) {
-        return (a.type == b.type) &&
-               (a.ne[0] == b.ne[0]) && (a.ne[1] == b.ne[1]) && (a.ne[2] == b.ne[2]) && (a.ne[3] == b.ne[3]) &&
-               (a.nb[0] == b.nb[0]) && (a.nb[1] == b.nb[1]) && (a.nb[2] == b.nb[2]) && (a.nb[3] == b.nb[3]);
+    static uint64_t hash_tensor(const ggml_tensor * t, const htp_tensor & shape) {
+        uint64_t hash = 1469598103934665603ull;
+        auto hash_value = [&hash](uint64_t value) {
+            hash ^= value;
+            hash *= 1099511628211ull;
+        };
+
+        hash_value((uintptr_t) t->data);
+        hash_value(shape.type);
+        hash_value(shape.size);
+        hash_value(shape.flags);
+        for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+            hash_value(shape.ne[i]);
+            hash_value(shape.nb[i]);
+        }
+
+        // avalanche high bits down to low bits for power-of-two table indexing
+        hash ^= hash >> 33;
+        hash *= 0xff51afd7ed558ccdull;
+        hash ^= hash >> 33;
+
+        return hash;
    }

    // add tensor and return its index
    int add_tensor(const ggml_tensor * t) {
-        // Lookup by tensor ptr
-        auto it = t_map.find(t);
-        if (it != t_map.end()) { return it->second; }
-
        htp_tensor shape {};
        init_tensor_shape(shape, t);

-        // Lookup by tensor data and shape
-        auto range = d_map.equal_range(t->data);
-        for (auto dit = range.first; dit != range.second; ++dit) {
-            if (same_shape(h_tens[dit->second], shape)) {
-                t_map.insert({t, dit->second});
-                return dit->second;
-            }
+        const uint64_t hash = hash_tensor(t, shape);
+        auto it = t_map.find(hash);
+        if (it != t_map.end()) {
+            return it->second;
        }

        auto sbuf = static_cast<ggml_hexagon_shared_buffer *>(t->buffer->context);
@@ -1360,8 +1361,7 @@ struct ggml_hexagon_opbatch {
        int ti = n_tens++;
        GGML_ASSERT(n_tens <= n_tens_max);

-        t_map.insert({t,       ti});
-        d_map.insert({t->data, ti});
+        t_map[hash] = ti;

        uint64_t t_offset = (uint8_t *) t->data - sbuf->base();

@@ -1391,8 +1391,11 @@ struct ggml_hexagon_opbatch {

        auto fit_tensor = [&](const ggml_tensor * t) {
            if (!t) return;
+
            if (check_tens) {
-                if (t_map.count(t)) return;
+                htp_tensor shape {};
+                init_tensor_shape(shape, t);
+                if (t_map.count(hash_tensor(t, shape))) return;
                extra_tens++;
            }

@@ -1540,7 +1543,7 @@ struct ggml_hexagon_opqueue {

        size_t tr_size = 0;
        if (opt_profile == 3) {
-            tr_size = (HTP_MAX_NTHREADS + 1) * opt_optrace * sizeof(htp_trace_desc);
+            tr_size = sess->n_threads * opt_optrace * sizeof(htp_trace_desc);
        }

        shm_blk_size = sizeof(htp_buf_desc)  * n_bufs    +
@@ -1595,7 +1598,7 @@ struct ggml_hexagon_opqueue {
        size_t tr_size = 0;
        if (opt_profile == 3) {
            req.n_traces = opt_optrace;
-            tr_size = (HTP_MAX_NTHREADS + 1) * req.n_traces * sizeof(htp_trace_desc);
+            tr_size = shm_buf->sess->n_threads * req.n_traces * sizeof(htp_trace_desc);
        } else {
            req.n_traces = 0;
        }
@@ -1650,7 +1653,7 @@ struct ggml_hexagon_opqueue {
        uint32_t n_traces = 0;
        if (opt_profile == 3) {
            n_traces = opt_optrace;
-            tr_size = (HTP_MAX_NTHREADS + 1) * n_traces * sizeof(htp_trace_desc);
+            tr_size = shm_buf->sess->n_threads * n_traces * sizeof(htp_trace_desc);
        }

        const size_t m_size = b_size + t_size + o_size + p_size + tr_size;
@@ -1679,7 +1682,7 @@ struct ggml_hexagon_opqueue {
                ggml_hexagon_dump_op_prof(shm_buf->sess->c_name(), ops[i], pd[i]);
            }

-            ggml_hexagon_dump_trace_events(shm_buf->sess->c_name(), rsp, trace_events, n_traces);
+            ggml_hexagon_dump_trace_events(shm_buf->sess->c_name(), rsp, trace_events, shm_buf->sess->n_threads, n_traces);
        }
    }
 };
@@ -2114,34 +2117,80 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
        }
    }

-    char session_uri[256];
-    {
-        char htp_uri[256];
-        snprintf(htp_uri, sizeof(htp_uri), "file:///libggml-htp-v%u.so?htp_iface_skel_handle_invoke&_modver=1.0", opt_arch);
+    auto open_session = [&](const std::string & prefix, std::string & uri_out) -> int {
+        const std::string htp_uri = prefix + "libggml-htp-v" + std::to_string(opt_arch) +
+                                    ".so?htp_iface_skel_handle_invoke&_modver=1.0";
+
+        std::vector<char> session_uri(htp_uri.size() + 256);

        struct remote_rpc_get_uri u = {};
        u.session_id      = this->session_id;
        u.domain_name     = const_cast<char *>(dom_name.c_str());
        u.domain_name_len = dom_name.size();
-        u.module_uri      = const_cast<char *>(htp_uri);
-        u.module_uri_len  = strlen(htp_uri);
-        u.uri             = session_uri;
-        u.uri_len         = sizeof(session_uri);
+        u.module_uri      = const_cast<char *>(htp_uri.c_str());
+        u.module_uri_len  = htp_uri.size();
+        u.uri             = session_uri.data();
+        u.uri_len         = session_uri.size();

        int err = remote_session_control(FASTRPC_GET_URI, (void *) &u, sizeof(u));
        if (err != AEE_SUCCESS) {
-            snprintf(session_uri, sizeof(session_uri), "%s&_dom=%s&_session=%u",
-                     htp_uri, dom_name.c_str(), this->session_id);
+            snprintf(session_uri.data(), session_uri.size(), "%s&_dom=%s&_session=%u",
+                     htp_uri.c_str(), dom_name.c_str(), this->session_id);

-            GGML_LOG_WARN("ggml-hex: %s failed to get URI (physical %d, virtual %d) : error 0x%x. Falling back to single session URI: %s\n",
-                          this->c_name(), phys_idx, virt_idx, err, session_uri);
+            HEX_VERBOSE("ggml-hex: %s failed to get URI (physical %d, virtual %d) : error 0x%x. Falling back to single session URI: %s\n",
+                        this->c_name(), phys_idx, virt_idx, err, session_uri.data());
+        }
+
+        uri_out = session_uri.data();
+        return htp_iface_open(session_uri.data(), &this->handle);
+    };
+
+    const std::string lib_dir = htp_lib_dir();
+
+    std::string session_uri;
+    int err = AEE_EFAILED;
+
+    // 1. Try opening session with absolute path in URI
+    if (!lib_dir.empty()) {
+        std::string abs_prefix = "file:///" + lib_dir;
+        if (abs_prefix.back() != '/' && abs_prefix.back() != '\\') {
+            abs_prefix += "/";
+        }
+
+        HEX_VERBOSE("ggml-hex: %s trying absolute path URI: %s\n", this->c_name(), abs_prefix.c_str());
+        err = open_session(abs_prefix, session_uri);
+        if (err != AEE_SUCCESS) {
+            HEX_VERBOSE("ggml-hex: %s failed to open session with absolute path : uri %s error 0x%x\n",
+                        this->c_name(), session_uri.c_str(), err);
        }
    }

-    // Open session
-    int err = htp_iface_open(session_uri, &this->handle);
+    // 2. Fallback to DSP search path via FASTRPC_MANAGE_DSP_LIBRARY_PATH
+    if (err != AEE_SUCCESS && !lib_dir.empty()) {
+        int sp_err = htpdrv_set_dsp_search_path(this->domain_id, lib_dir, this->c_name());
+        if (sp_err == AEE_SUCCESS) {
+            err = open_session("file:///", session_uri);
+            if (err != AEE_SUCCESS) {
+                HEX_VERBOSE("ggml-hex: %s failed to open session with search path : uri %s error 0x%x\n",
+                            this->c_name(), session_uri.c_str(), err);
+            }
+        }
+    }
+
+    // 3. Fallback to process environment variables
+    if (err != AEE_SUCCESS && !lib_dir.empty()) {
+        HEX_VERBOSE("ggml-hex: %s falling back to process env for DSP search path\n", this->c_name());
+        htpdrv_set_dsp_env_path(lib_dir, this->c_name());
+        err = open_session("file:///", session_uri);
+    }
+
+    // Bare name attempt if lib_dir is empty and nothing was tried yet
+    if (err != AEE_SUCCESS && lib_dir.empty()) {
+        err = open_session("file:///", session_uri);
+    }
+
    if (err != AEE_SUCCESS) {
-        GGML_LOG_ERROR("ggml-hex: %s failed to open session : uri %s error 0x%x\n", this->c_name(), session_uri, err);
+        GGML_LOG_ERROR("ggml-hex: %s failed to open session : uri %s error 0x%x\n", this->c_name(), session_uri.c_str(), err);
        throw std::runtime_error("ggml-hex: failed to open session (see log for details)");
    }

@@ -2155,11 +2204,11 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
        unsigned long long hw_vtcm_size = 0;
        int hw_err = htp_iface_hwinfo(this->handle, &hw_n_threads, &hw_n_hvx, &hw_n_hmx, &hw_vtcm_size);
        if (hw_err == 0) {
-            const uint32_t max_n_threads = (std::min)((uint32_t) HTP_MAX_NTHREADS, (uint32_t) hw_n_threads);
-            this->n_threads = opt_nhvx > 0 ? (uint32_t) (std::min)(opt_nhvx, (size_t) max_n_threads) : max_n_threads;
-            this->n_hvx     = this->n_threads;
-            this->n_hmx     = (opt_nhmx != 0) ? (uint32_t)hw_n_hmx : 0;
-            this->vtcm_size = (uint64_t)hw_vtcm_size;
+            this->n_threads          = (std::min)((uint32_t) HTP_OP_MAX_NTHREADS, (uint32_t) hw_n_threads);
+            const uint32_t max_n_hvx = (std::min)({this->n_threads - 1, (uint32_t) hw_n_hvx, (uint32_t) HTP_OP_MAX_NHVX});
+            this->n_hvx              = opt_nhvx > 0 ? (uint32_t) (std::min)(opt_nhvx, (size_t) max_n_hvx) : max_n_hvx;
+            this->n_hmx              = (opt_nhmx != 0) ? (uint32_t)hw_n_hmx : 0;
+            this->vtcm_size          = (uint64_t)hw_vtcm_size;
            if (opt_vtcm_size > 0) {
                if (opt_vtcm_size > hw_vtcm_size) {
                    GGML_LOG_WARN("ggml-hex: %s VTCM override %zu exceeds HW size %llu\n", this->c_name(), opt_vtcm_size, hw_vtcm_size);
@@ -2171,11 +2220,11 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
                          (unsigned long long)(this->vtcm_size / (1024 * 1024)));
        } else {
            GGML_LOG_WARN("ggml-hex: %s failed to query hwinfo (0x%x), using defaults\n", this->c_name(), hw_err);
-            const uint32_t default_n_threads = (std::min)(8u, (uint32_t) HTP_MAX_NTHREADS);
-            this->n_threads = opt_nhvx > 0 ? (uint32_t) (std::min)(opt_nhvx, (size_t) HTP_MAX_NTHREADS) : default_n_threads;
-            this->n_hvx     = this->n_threads;
-            this->n_hmx     = (opt_nhmx != 0) ? 1 : 0;
-            this->vtcm_size = opt_vtcm_size > 0 ? opt_vtcm_size : 8 * 1024 * 1024;
+            this->n_threads          = (std::min)(8u, (uint32_t) HTP_OP_MAX_NTHREADS);
+            const uint32_t max_n_hvx = (std::min)(this->n_threads - 1, (uint32_t) HTP_OP_MAX_NHVX);
+            this->n_hvx              = opt_nhvx > 0 ? (uint32_t) (std::min)(opt_nhvx, (size_t) max_n_hvx) : (std::min)(4u, max_n_hvx);
+            this->n_hmx              = (opt_nhmx != 0) ? 1 : 0;
+            this->vtcm_size          = opt_vtcm_size > 0 ? opt_vtcm_size : 8 * 1024 * 1024;
        }
    }

@@ -2191,7 +2240,7 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
    }

    GGML_LOG_INFO("ggml-hex: %s new session : session-id %d domain-id %d uri %s handle 0x%lx\n", this->c_name(),
-                  this->session_id, this->domain_id, session_uri, (unsigned long) this->handle);
+                  this->session_id, this->domain_id, session_uri.c_str(), (unsigned long) this->handle);

    const size_t req_q_size = (sizeof(htp_opbatch_req) * opt_opqueue * 2) + 1024;
    const size_t rsp_q_size = (sizeof(htp_opbatch_rsp) * opt_opqueue * 2) + 1024;
@@ -2244,7 +2293,7 @@ void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) n
    this->op_batch = new ggml_hexagon_opbatch(this, opt_opbatch, this->max_vmem);

    // Start dspqueue/opbatch processing
-    err = htp_iface_start(this->handle, this->session_id, this->queue_id, this->n_threads, opt_nhmx, opt_vtcm_size, this->max_vmem,
+    err = htp_iface_start(this->handle, this->session_id, this->queue_id, this->n_hvx, opt_nhmx, opt_vtcm_size, this->max_vmem,
                          (uint32_t) opt_power_mode);
    if (err != 0) {
        GGML_LOG_ERROR("ggml-hex: %s failed to start session: 0x%08x\n", this->c_name(), (unsigned) err);
@@ -2364,11 +2413,33 @@ static void ggml_backend_hexagon_free(ggml_backend_t backend) {
    delete backend;
 }

-static inline bool op_is_compute(ggml_tensor *node)
-{
+static inline bool ggml_hexagon_op_is_compute(const struct ggml_tensor * node) {
    return !ggml_op_is_empty(node->op) && !ggml_is_empty(node) && (node->flags & GGML_TENSOR_FLAG_COMPUTE);
 }

+static void ggml_hexagon_tensor_repack(ggml_tensor * tensor, ggml_hexagon_session * sess) {
+    if (tensor->data == nullptr || ggml_hexagon_tensor_is_tiled(tensor) || !ggml_hexagon_tensor_can_tile(tensor)) {
+        return;
+    }
+    HEX_VERBOSE("ggml-hex: %s rt-repack %s|size %zu\n", sess->c_name(), htp_format_tensor(tensor, HTP_OPFMT_TENSOR).c_str(), ggml_nbytes(tensor));
+    std::vector<uint8_t> raw(ggml_nbytes(tensor));
+    memcpy(raw.data(), tensor->data, raw.size());
+    htp_repack_tensor_tiled(tensor, raw.data(), raw.size());
+    ggml_hexagon_tensor_set_flags(tensor, HTP_TENSOR_TILED);
+}
+
+static void ggml_hexagon_graph_repack(const ggml_cgraph * graph, ggml_hexagon_session * sess) {
+    for (int i = 0; i < graph->n_nodes; ++i) {
+        ggml_tensor * n = graph->nodes[i];
+        if (!ggml_hexagon_op_is_compute(n)) {
+            continue;
+        }
+        if (n->op == GGML_OP_MUL_MAT || n->op == GGML_OP_MUL_MAT_ID) {
+            ggml_hexagon_tensor_repack(n->src[0], sess);
+        }
+    }
+}
+
 static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, ggml_cgraph * graph) {
    auto sess = static_cast<ggml_hexagon_session *>(backend->context);

@@ -2385,13 +2456,15 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg
    if (cache_hit) {
        nodes_ptr = &sess->cached_nodes;
    } else {
+        ggml_hexagon_graph_repack(graph, sess);
+
        htp_fusion_tag_graph(graph);

        computed_nodes.reserve(graph->n_nodes);

        for (int i = 0; i < graph->n_nodes; ++i) {
            ggml_tensor * n = graph->nodes[i];
-            if (!op_is_compute(n)) {
+            if (!ggml_hexagon_op_is_compute(n)) {
                continue;
            }

@@ -2401,22 +2474,7 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg
            }

            htp_opnode node(opcode, n);
-            if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID || node.opcode == HTP_OP_GET_ROWS) {
-                ggml_tensor * src0 = node.node->src[0];
-                if (htp_needs_repack(src0)) {
-                    HEX_VERBOSE("ggml-hex: %s rt-repack %s|size %zu\n", sess->c_name(), htp_format_tensor(src0, HTP_OPFMT_TENSOR).c_str(), ggml_nbytes(src0));
-                    std::vector<uint8_t> raw(ggml_nbytes(src0));
-                    memcpy(raw.data(), src0->data, raw.size());
-                    htp_repack_tensor_tiled(src0, raw.data(), raw.size());
-                    auto * extra = (htp_tensor_extra *) src0->extra;
-                    extra->flags |= HTP_TENSOR_TILED;
-                }
-            }
-
            GGML_ASSERT(htp_kparams_solve(sess, node));
-            if (node.opcode == HTP_OP_CPY && node.kparams<htp_copy_kernel_params>()->total_elems == 0) {
-                continue;
-            }

            computed_nodes.push_back(std::move(node));
        }
@@ -2626,7 +2684,11 @@ static std::vector<int> ggml_hexagon_graph_optimize_reorder(const std::vector<no
 }

 static void ggml_backend_hexagon_graph_optimize(ggml_backend_t backend, ggml_cgraph * gf, ggml_backend_graph_optimize_params * params) {
-    GGML_UNUSED(params);
+    auto sess = static_cast<ggml_hexagon_session *>(backend->context);
+
+    HEX_VERBOSE("ggml-hex: %s graph-optimize n_nodes %d\n", sess->c_name(), gf->n_nodes);
+
+    ggml_hexagon_graph_repack(gf, sess);

    const int n = gf->n_nodes;

@@ -2661,7 +2723,7 @@ static void ggml_backend_hexagon_graph_optimize(ggml_backend_t backend, ggml_cgr
        }
    }

-    GGML_UNUSED(backend);
+    GGML_UNUSED(params);
 }

 static uint64_t ggml_hexagon_session_key(const ggml_hexagon_session * sess) {
@@ -2723,14 +2785,7 @@ static bool ggml_backend_hexagon_cpy_tensor_async(ggml_backend_t backend_src, gg
        return false;
    }

-    // FIXME: ggml-meta needs to call init_tensor on auxiliary tensors
-    if (!dst->extra) {
-        ggml_backend_buffer_init_tensor(dst->buffer, dst);
-    }
-
-    auto * dst_extra = static_cast<htp_tensor_extra *>(dst->extra);
-    const auto * src_extra = static_cast<const htp_tensor_extra *>(src->extra);
-    dst_extra->flags = src_extra->flags & ~HTP_TENSOR_FUSEABLE;
+    dst->extra = (void *) (uintptr_t) (ggml_hexagon_tensor_get_flags(src) & ~HTP_TENSOR_FUSEABLE);

    auto sess_src = static_cast<ggml_hexagon_session *>(backend_src->context);
    auto sess_dst = static_cast<ggml_hexagon_session *>(backend_dst->context);
@@ -3003,18 +3058,18 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
    }

    bool supp = false;
-    switch (op->op) {
-        case GGML_OP_NONE:
-        case GGML_OP_RESHAPE:
-        case GGML_OP_VIEW:
-        case GGML_OP_PERMUTE:
-        case GGML_OP_TRANSPOSE:
-            supp = true;
-            break;
-
-        default:
+    if (ggml_op_is_empty(op->op) || ggml_is_empty(op)) {
+        supp = true;
+    } else {
+        supp = ggml_backend_hexagon_buffer_type_get_alloc_size(&dev_ctx->buffer_type, op) <= sess->max_vmem;
+        for (int i = 0; supp && i < GGML_MAX_SRC; i++) {
+            if (op->src[i] && ggml_backend_hexagon_buffer_type_get_alloc_size(&dev_ctx->buffer_type, op->src[i]) > sess->max_vmem) {
+                supp = false;
+            }
+        }
+        if (supp) {
            supp = htp_kparams_check(sess, op);
-            break;
+        }
    }

    ggml_hexagon_dump_op_supp(sess->c_name(), op, supp);
@@ -3340,8 +3395,8 @@ static void ggml_hexagon_discover_devices() {
            discovery_supported = true;
            const int n_domains = std::min(domain_info.sys.num_domains, (int) domains.size());
            for (int i = 0; i < n_domains; i++) {
-                GGML_LOG_INFO("ggml-hex: FASTRPC_GET_DOMAINS[%d]: type %d id %d name '%s' status %d instance-id %d\n",
-                              i, (int) domains[i].type, domains[i].id, domains[i].name, domains[i].status, domains[i].instance_id);
+                HEX_VERBOSE("ggml-hex: FASTRPC_GET_DOMAINS[%d]: type %d id %d name '%s' status %d instance-id %d\n",
+                            i, (int) domains[i].type, domains[i].id, domains[i].name, domains[i].status, domains[i].instance_id);
                if (domains[i].type != FASTRPC_NSP) {
                    GGML_LOG_DEBUG("ggml-hex:   skipping non-CDSP domain (type=%d)\n", (int) domains[i].type);
                    continue;
@@ -3355,10 +3410,10 @@ static void ggml_hexagon_discover_devices() {
                              domains[i].instance_id, domains[i].id, domains[i].name);
            }
        } else {
-            GGML_LOG_WARN("ggml-hex: FASTRPC_GET_DOMAINS fetch failed (0x%x), using static CDSP domains\n", (unsigned) err);
+            HEX_VERBOSE("ggml-hex: FASTRPC_GET_DOMAINS fetch failed (0x%x), using static CDSP domains\n", (unsigned) err);
        }
    } else if (err != AEE_SUCCESS) {
-        GGML_LOG_DEBUG("ggml-hex: FASTRPC_GET_DOMAINS query failed (0x%x), using static CDSP domains\n", (unsigned) err);
+        HEX_VERBOSE("ggml-hex: FASTRPC_GET_DOMAINS query failed (0x%x), using static CDSP domains\n", (unsigned) err);
    }

    // Populate domain IDs and names for all configured devices
@@ -4,12 +4,6 @@
 #include "htp-opnode.h"
 #include "htp-repack.h"

-struct htp_tensor_extra {
-    std::vector<uint8_t> shadow_buf;
-    size_t               shadow_size { 0 };
-    uint32_t             flags { 0 };
-};
-
 extern int    opt_arch;
 extern size_t opt_ndev;
 extern size_t opt_nhvx;
@@ -46,36 +40,41 @@ extern int    opt_rt_repack;
    if (opt_verbose) GGML_LOG_DEBUG(__VA_ARGS__)
 #endif

-static inline bool htp_tensor_is_fuseable(const struct ggml_tensor * t) {
-    if (!t || !t->extra) return false;
-    const auto * extra = (const struct htp_tensor_extra *) t->extra;
-    return (extra->flags & HTP_TENSOR_FUSEABLE) != 0;
+static inline uint32_t ggml_hexagon_tensor_get_flags(const struct ggml_tensor * t) {
+    return (uint32_t) (uintptr_t) t->extra;
 }

-static inline bool htp_can_rt_repack(const struct ggml_tensor * t) {
-    if (opt_rt_repack == 0 || !htp_is_tiled_weight_tensor(t)) {
-        return false;
-    }
-    if (t->buffer && ggml_backend_buft_is_host(t->buffer->buft)) {
-        return false;
-    }
-    return opt_rt_repack > 1 || ggml_is_quantized(t->type);
+static inline void ggml_hexagon_tensor_set_flags(struct ggml_tensor * t, uint32_t flags) {
+    t->extra = (void *) (uintptr_t) (ggml_hexagon_tensor_get_flags(t) | flags);
 }

-static inline bool htp_tensor_is_tiled(const struct ggml_tensor * t) {
-    const auto * extra = (const struct htp_tensor_extra *) t->extra;
-    if (extra && (extra->flags & HTP_TENSOR_TILED)) {
-        return true;
-    }
-    return htp_can_rt_repack(t);
+static inline void ggml_hexagon_tensor_clear_flags(struct ggml_tensor * t, uint32_t flags) {
+    t->extra = (void *) (uintptr_t) (ggml_hexagon_tensor_get_flags(t) & ~flags);
 }

-static inline bool htp_needs_repack(const struct ggml_tensor * t) {
-    const auto * extra = (const struct htp_tensor_extra *) t->extra;
-    if (extra && (extra->flags & HTP_TENSOR_TILED)) {
+static inline bool ggml_hexagon_tensor_is_fuseable(const struct ggml_tensor * t) {
+    return (ggml_hexagon_tensor_get_flags(t) & HTP_TENSOR_FUSEABLE) != 0;
+}
+
+static inline bool ggml_hexagon_tensor_is_tiled(const struct ggml_tensor * t) {
+    return (ggml_hexagon_tensor_get_flags(t) & HTP_TENSOR_TILED) != 0;
+}
+
+static inline bool ggml_hexagon_tensor_can_tile(const struct ggml_tensor * t) {
+    if (opt_rt_repack == 0) {
        return false;
    }
-    return htp_can_rt_repack(t);
+    if (t->buffer) {
+        if (ggml_backend_buft_is_host(t->buffer->buft) ||
+            ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+            return false;
+        }
+    }
+    return htp_is_tiled_weight_tensor(t);
+}
+
+static inline bool ggml_hexagon_tensor_will_tile(const struct ggml_tensor * t) {
+    return ggml_hexagon_tensor_is_tiled(t) || ggml_hexagon_tensor_can_tile(t);
 }

 static inline bool htp_tensors_overlap(const struct ggml_tensor * a, const struct ggml_tensor * b) {
@@ -98,7 +97,7 @@ struct ggml_hexagon_session;

 const char * htp_session_name(const struct ggml_hexagon_session * sess);
 size_t       htp_session_vtcm_size(const struct ggml_hexagon_session * sess);
-uint32_t     htp_session_n_threads(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_n_hvx(const struct ggml_hexagon_session * sess);
 uint32_t     htp_session_n_hmx(const struct ggml_hexagon_session * sess);
 uint32_t     htp_session_mdev_count(const struct ggml_hexagon_session * sess);

@@ -1,7 +1,9 @@
+#include <algorithm>
 #include <filesystem>
 #include <set>
 #include <sstream>
 #include <string>
+#include <vector>

 #ifdef _WIN32
 #   define WIN32_LEAN_AND_MEAN
@@ -28,6 +30,13 @@

 #include <domain.h>

+extern int opt_verbose;
+
+#ifndef HEX_VERBOSE
+#define HEX_VERBOSE(...) \
+    if (opt_verbose) GGML_LOG_DEBUG(__VA_ARGS__)
+#endif
+
 //
 // Driver API types
 //
@@ -423,3 +432,188 @@ int htpdrv_get_arch(int domain, int * arch) {
    *arch = (int) ((val >> 4) * 10 + (val & 0x0f));
    return 0;
 }
+
+// Directory of the HTP library as a path with a trailing slash, or empty if it cannot be resolved.
+// Checks GGML_HEXAGON_LIBHTP_DIR first, then the module directory and ../lib.
+// On Windows, non-ASCII directories or paths with spaces use the 8.3 short name.
+std::string htp_lib_dir() {
+    namespace fs = std::filesystem;
+
+    std::vector<fs::path> candidates;
+
+    const char * env_dir = getenv("GGML_HEXAGON_LIBHTP_DIR");
+    if (env_dir && *env_dir) {
+        std::error_code ec;
+        fs::path p = fs::absolute(env_dir, ec);
+        if (!ec) {
+            candidates.push_back(p);
+        }
+    }
+
+#ifdef _WIN32
+    HMODULE hm = nullptr;
+    if (GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+                           (LPCWSTR) &htp_lib_dir, &hm)) {
+        std::wstring buf(32768, L'\0');
+        DWORD n = GetModuleFileNameW(hm, buf.data(), (DWORD) buf.size());
+        if (n > 0 && n < buf.size()) {
+            buf.resize(n);
+            std::error_code ec;
+            fs::path mod_path = fs::canonical(buf, ec);
+            if (ec) {
+                mod_path = fs::absolute(buf, ec);
+            }
+            const fs::path mod_dir = mod_path.parent_path();
+            candidates.push_back(mod_dir);
+            candidates.push_back(mod_dir / "lib");
+            candidates.push_back(mod_dir / ".." / "lib");
+        }
+    }
+#elif defined(__linux__) || defined(__ANDROID__)
+    Dl_info info;
+    if (dladdr((void *) &htp_lib_dir, &info) && info.dli_fname) {
+        std::error_code ec;
+        fs::path mod_path = fs::canonical(info.dli_fname, ec);
+        if (ec) {
+            mod_path = fs::absolute(info.dli_fname, ec);
+        }
+        const fs::path mod_dir = mod_path.parent_path();
+        candidates.push_back(mod_dir);
+        candidates.push_back(mod_dir / "lib");
+        candidates.push_back(mod_dir / ".." / "lib");
+    }
+#endif
+
+    for (size_t i = 0; i < candidates.size(); ++i) {
+        const fs::path & d = candidates[i];
+        std::error_code ec;
+        if (!fs::is_directory(d, ec)) {
+            continue;
+        }
+
+        // For auto-discovered directories, verify it contains an HTP library
+        const bool is_env = (env_dir && i == 0);
+        if (!is_env) {
+            bool found = false;
+            for (const auto & entry : fs::directory_iterator(d, ec)) {
+                if (ec) {
+                    break;
+                }
+                const std::string fn = entry.path().filename().string();
+                if (fn.rfind("libggml-htp-", 0) == 0) {
+                    found = true;
+                    break;
+                }
+            }
+            if (!found) {
+                continue;
+            }
+        }
+
+        std::error_code ec_canon;
+        fs::path abs_d = fs::canonical(d, ec_canon);
+        if (ec_canon) {
+            abs_d = fs::absolute(d, ec_canon);
+        }
+
+#ifdef _WIN32
+        std::wstring dir = abs_d.wstring();
+        auto is_ascii = [](const std::wstring & s) {
+            return std::all_of(s.begin(), s.end(), [](wchar_t c) { return c < 0x80; });
+        };
+        auto has_spaces = [](const std::wstring & s) {
+            return std::any_of(s.begin(), s.end(), [](wchar_t c) { return c == L' '; });
+        };
+
+        if (!is_ascii(dir) || has_spaces(dir)) {
+            DWORD m = GetShortPathNameW(dir.c_str(), nullptr, 0);
+            if (m > 0) {
+                std::wstring sp(m, L'\0');
+                DWORD res = GetShortPathNameW(dir.c_str(), sp.data(), (DWORD) sp.size());
+                if (res > 0 && res < sp.size()) {
+                    sp.resize(res);
+                    if (is_ascii(sp)) {
+                        dir = sp;
+                    }
+                }
+            }
+            if (!is_ascii(dir)) {
+                continue;
+            }
+        }
+        if (dir.back() != L'\\' && dir.back() != L'/') {
+            dir += L'\\';
+        }
+        return std::string(dir.begin(), dir.end());
+#else
+        std::string dir = abs_d.string();
+        if (dir.back() != '/') {
+            dir += '/';
+        }
+        return dir;
+#endif
+    }
+
+    return {};
+}
+
+int htpdrv_set_dsp_search_path(int domain_id, const std::string & path, const char * dev_name) {
+    if (path.empty()) {
+        return AEE_SUCCESS;
+    }
+
+    std::string clean_path = path;
+    while (clean_path.size() > 1 && (clean_path.back() == '/' || clean_path.back() == '\\')) {
+        clean_path.pop_back();
+    }
+
+    fastrpc_manage_dsp_lib_path p {};
+    p.effec_domain_id = domain_id;
+    p.type            = FASTRPC_SET_PATH;
+    p.path            = const_cast<char *>(clean_path.c_str());
+    p.pathlen         = clean_path.size() + 1;
+
+    int err = remote_session_control(FASTRPC_MANAGE_DSP_LIBRARY_PATH, (void *) &p, sizeof(p));
+    if (err != AEE_SUCCESS && domain_id != -1) {
+        p.effec_domain_id = -1;
+        err = remote_session_control(FASTRPC_MANAGE_DSP_LIBRARY_PATH, (void *) &p, sizeof(p));
+    }
+
+    const char * name = dev_name ? dev_name : "driver";
+    if (err == AEE_SUCCESS) {
+        GGML_LOG_INFO("ggml-hex: %s set DSP search path (domain %d) to %s\n", name, p.effec_domain_id, clean_path.c_str());
+        return AEE_SUCCESS;
+    }
+
+    HEX_VERBOSE("ggml-hex: %s FASTRPC_MANAGE_DSP_LIBRARY_PATH failed (0x%x)\n", name, err);
+    return err;
+}
+
+void htpdrv_set_dsp_env_path(const std::string & path, const char * dev_name) {
+    if (path.empty()) {
+        return;
+    }
+
+    std::string clean_path = path;
+    while (clean_path.size() > 1 && (clean_path.back() == '/' || clean_path.back() == '\\')) {
+        clean_path.pop_back();
+    }
+
+    const char * name = dev_name ? dev_name : "driver";
+    HEX_VERBOSE("ggml-hex: %s setting DSP search path via process env to %s\n", name, clean_path.c_str());
+
+    static const char * const env_names[] = {
+        "ADSP_LIBRARY_PATH",
+        "CDSP_LIBRARY_PATH",
+        "CDSP1_LIBRARY_PATH",
+        "CDSP2_LIBRARY_PATH",
+        "CDSP3_LIBRARY_PATH",
+    };
+    for (const char * env_name : env_names) {
+#ifdef _WIN32
+        SetEnvironmentVariableA(env_name, clean_path.c_str());
+#else
+        setenv(env_name, clean_path.c_str(), 1);
+#endif
+    }
+}
@@ -120,4 +120,10 @@ HTPDRV_API int remote_system_request(system_req_payload * req);

 #ifdef __cplusplus
 }
+
+#include <string>
+
+std::string htp_lib_dir();
+int         htpdrv_set_dsp_search_path(int domain_id, const std::string & path, const char * dev_name = nullptr);
+void        htpdrv_set_dsp_env_path(const std::string & path, const char * dev_name = nullptr);
 #endif
@@ -35,7 +35,7 @@ static inline bool htp_match_consumer(
        return false;
    }
    const ggml_tensor * prev_out = prev.dst();
-    if (!htp_tensor_is_fuseable(prev_out)) {
+    if (!ggml_hexagon_tensor_is_fuseable(prev_out)) {
        return false;
    }
    if (htp_same_tensor_data(cur.src0(), prev_out)) {
@@ -53,14 +53,14 @@ static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struc
    if (kparams->n_hmx) {
        return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D;
    }
-    return htp_tensor_is_tiled(src0) && (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW);
+    return ggml_hexagon_tensor_is_tiled(src0) && (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW);
 }

 static bool is_supported_mul_mat_id_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams) {
    if (kparams->n_hmx) {
        return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D;
    }
-    return htp_tensor_is_tiled(src0) &&
+    return ggml_hexagon_tensor_is_tiled(src0) &&
           (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK);
 }

@@ -74,7 +74,7 @@ static bool fuse_allreduce_add(const ggml_hexagon_session * sess, htp_opnode & p
    const uint32_t n_ranks = (uint32_t) ar_kparams->n_ranks;

    const ggml_tensor * ar_local = prev.inputs[rank];
-    if (!htp_tensor_is_fuseable(ar_local)) return false;
+    if (!ggml_hexagon_tensor_is_fuseable(ar_local)) return false;

    const ggml_tensor * add_dst    = cur.dst();
    const ggml_tensor * res_tensor = nullptr;
@@ -157,6 +157,30 @@ static bool fuse_rms_norm_mul(const ggml_hexagon_session * sess, htp_opnode & pr
    return true;
 }

+// post-norm residual: dst = rms_norm(x) * w + r
+static bool fuse_rms_norm_mul_add(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_RMS_NORM_MUL) return false;
+
+    const ggml_tensor * r = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &r)) return false;
+
+    const ggml_tensor * src0   = prev.src0();
+    const ggml_tensor * weight = prev.inputs[1];
+    const ggml_tensor * dst    = cur.dst();
+    if (!ggml_are_same_layout(src0, r)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst, r)) return false;
+
+    struct htp_unary_kernel_params kparams;
+    if (!htp_kparams_solve_rms_norm_mul_add(sess, src0, weight, dst, &kparams)) return false;
+
+    prev.opcode = HTP_OP_RMS_NORM_MUL_ADD;
+    prev.inputs.push_back(r);
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
 static bool fuse_rms_norm_scale(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
    GGML_UNUSED(sess);
    if (prev.opcode != HTP_OP_RMS_NORM && prev.opcode != HTP_OP_RMS_NORM_MUL) return false;
@@ -215,11 +239,11 @@ static bool fuse_mul_mat_add(const ggml_hexagon_session * sess, htp_opnode & pre
    if (kparams.n_hmx > 0 && orig_kparams->n_hmx > 0) {
        if (kparams.m_chunk < orig_kparams->m_chunk ||
            kparams.n_chunk < orig_kparams->n_chunk ||
-            kparams.n_act_threads < orig_kparams->n_act_threads) {
-            HEX_VERBOSE("ggml-hex: %s skip MUL_MAT_ADD fusion: HMX efficiency reduced (m %d->%d, n %d->%d, th %d->%d)\n",
+            kparams.n_act_hvx < orig_kparams->n_act_hvx) {
+            HEX_VERBOSE("ggml-hex: %s skip MUL_MAT_ADD fusion: HMX efficiency reduced (m %d->%d, n %d->%d, hvx %d->%d)\n",
                        htp_session_name(sess), orig_kparams->m_chunk, kparams.m_chunk,
                        orig_kparams->n_chunk, kparams.n_chunk,
-                        orig_kparams->n_act_threads, kparams.n_act_threads);
+                        orig_kparams->n_act_hvx, kparams.n_act_hvx);
            return false;
        }
    }
@@ -250,7 +274,7 @@ static bool fuse_mul_mat_nx(const ggml_hexagon_session * sess, htp_opnode & prev

    if (x_in != x_orig || w_in->type != w0->type || w_in->ne[0] != w0->ne[0] ||
        w0->ne[2] != 1 || w0->ne[3] != 1 || w_in->ne[2] != 1 || w_in->ne[3] != 1 ||
-        htp_tensor_is_tiled(w_in) != htp_tensor_is_tiled(w0) ||
+        ggml_hexagon_tensor_is_tiled(w_in) != ggml_hexagon_tensor_is_tiled(w0) ||
        orig_kparams->n_hmx != cur_kparams->n_hmx) {
        return false;
    }
@@ -295,7 +319,7 @@ static bool fuse_mul_mat_id_nx(const ggml_hexagon_session * sess, htp_opnode & p

    if (x_in != x || ids_in != ids || w_in->type != w0->type ||
        w_in->ne[0] != w0->ne[0] || w_in->ne[2] != w0->ne[2] ||
-        htp_tensor_is_tiled(w_in) != htp_tensor_is_tiled(w0) ||
+        ggml_hexagon_tensor_is_tiled(w_in) != ggml_hexagon_tensor_is_tiled(w0) ||
        orig_kparams->n_hmx != cur_kparams->n_hmx) {
        return false;
    }
@@ -522,8 +546,8 @@ static bool fuse_binary_nx(const ggml_hexagon_session * sess, htp_opnode & prev,

        const uint32_t total_elems = (uint32_t) ggml_nelements(src0);
        const uint32_t elem_size   = (uint32_t) ggml_type_size(src0->type);
-        const uint32_t max_threads = htp_session_n_threads(sess);
-        if (!htp_binary_solve_chunked(&L, &kparams, total_elems, elem_size, max_threads, n_srcs, false, vtcm_size)) {
+        const uint32_t max_hvx     = htp_session_n_hvx(sess);
+        if (!htp_binary_solve_chunked(&L, &kparams, total_elems, elem_size, max_hvx, n_srcs, false, vtcm_size)) {
            return false;
        }
    }
@@ -598,19 +622,17 @@ int htp_fusion_max(const struct ggml_cgraph * graph, int idx) {

 void htp_fusion_tag_graph(struct ggml_cgraph * graph) {
    for (int i = 0; i < graph->n_nodes; i++) {
-        auto * extra = (htp_tensor_extra *) graph->nodes[i]->extra;
-        if (!extra) continue;
+        ggml_tensor * node = graph->nodes[i];
+        ggml_hexagon_tensor_clear_flags(node, HTP_TENSOR_FUSEABLE);

-        extra->flags &= ~HTP_TENSOR_FUSEABLE;
-
-        const ggml_op op = graph->nodes[i]->op;
+        const ggml_op op = node->op;
        if ((op == GGML_OP_RMS_NORM || op == GGML_OP_SCALE || op == GGML_OP_REPEAT ||
             op == GGML_OP_UNARY    || op == GGML_OP_ADD   || op == GGML_OP_SUB    ||
             op == GGML_OP_MUL      || op == GGML_OP_DIV   || op == GGML_OP_MUL_MAT_ID ||
             op == GGML_OP_SSM_CONV ||
-             (op == GGML_OP_MUL_MAT && !htp_is_hadamard(graph->nodes[i]))) &&
+             (op == GGML_OP_MUL_MAT && !htp_is_hadamard(node))) &&
            ggml_node_has_n_uses(graph, i, 1)) {
-            extra->flags |= HTP_TENSOR_FUSEABLE;
+            ggml_hexagon_tensor_set_flags(node, HTP_TENSOR_FUSEABLE);
        }
    }
 }
@@ -622,6 +644,7 @@ bool htp_fusion_try_fuse(const struct ggml_hexagon_session * sess, struct htp_op
        case HTP_OP_ADD:
            return (htp_fusion_enabled(HTP_FUSE_ALLREDUCE_ADD) && fuse_allreduce_add(sess, prev, cur)) ||
                   (htp_fusion_enabled(HTP_FUSE_MUL_MAT_ADD)   && fuse_mul_mat_add(sess, prev, cur))   ||
+                   (htp_fusion_enabled(HTP_FUSE_RMS_NORM_MUL_ADD) && fuse_rms_norm_mul_add(sess, prev, cur)) ||
                   (htp_fusion_enabled(HTP_FUSE_BINARY_NX)     && fuse_binary_nx(sess, prev, cur));
        case HTP_OP_MUL:
            return (htp_fusion_enabled(HTP_FUSE_RMS_NORM_MUL)   && fuse_rms_norm_mul(sess, prev, cur))   ||
@@ -16,6 +16,7 @@ enum htp_fusion_flags {
    HTP_FUSE_BINARY_NX       = (1 << 10), // 1024
    HTP_FUSE_RMS_NORM_SCALE  = (1 << 11), // 2048
    HTP_FUSE_SSM_CONV_SILU   = (1 << 12), // 4096
+    HTP_FUSE_RMS_NORM_MUL_ADD = (1 << 13), // 8192
 };

 static inline bool htp_fusion_enabled(int flag) {
@@ -95,6 +95,14 @@ bool htp_kparams_solve_rms_norm_mul(
    struct htp_unary_kernel_params * kparams
 );

+bool htp_kparams_solve_rms_norm_mul_add(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    struct htp_unary_kernel_params * kparams
+);
+
 bool htp_kparams_solve_matmul_add(
    const struct ggml_hexagon_session * sess,
    const struct ggml_tensor * src0,
@@ -157,7 +157,7 @@ static std::string format_kernel_params(const htp_opnode & node) {
        } else if (type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
            path = "hvx-tiled";
        }
-        snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", path, (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
        const auto * kparams = node.kparams<htp_fa_kernel_params>();
        const char * path = "unknown";
@@ -167,10 +167,11 @@ static std::string format_kernel_params(const htp_opnode & node) {
        } else if (type == HTP_FA_KERNEL_HVX) {
            path = "hvx";
        }
-        snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", path, (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (htp_op_is_unary(node.opcode)) {
        const auto * kparams = node.kparams<htp_unary_kernel_params>();
-        snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", kparams->col_tile ? "wide-row" : "row-block",
+                 (unsigned int) kparams->n_hvx, (unsigned 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) {
@@ -179,41 +180,43 @@ static std::string format_kernel_params(const htp_opnode & 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 if (node.opcode == HTP_OP_SSM_CONV) {
        const auto * kparams = node.kparams<htp_ssm_conv_kernel_params>();
-        snprintf(str, max_size, "%s vtcm %d", kparams->n_t == 1 ? "decode" : "prefill", (int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", kparams->n_t == 1 ? "decode" : "prefill",
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_SOFTMAX) {
        const auto * kparams = node.kparams<htp_softmax_kernel_params>();
-        snprintf(str, max_size, "k%d nth %d vtcm %d", (int) kparams->kernel_type, (int) kparams->n_threads, (int) kparams->vtcm_size);
+        snprintf(str, max_size, "k%d nhvx %u vtcm %u", (int) kparams->kernel_type,
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_GATED_DELTA_NET) {
        const auto * kparams = node.kparams<htp_gdn_kernel_params>();
        const char * path = (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) ? "hmx-chunked" : "hvx-recurrent";
-        snprintf(str, max_size, "%s-%s vtcm %u", path, kparams->kda ? "kda" : "scalar",
-                 (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s-%s nhvx %u vtcm %u", path, kparams->kda ? "kda" : "scalar",
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_DSV4_HC_PRE) {
        const auto * kparams = node.kparams<htp_dsv4_hc_pre_kernel_params>();
        const char * ktype_str = kparams->kernel_type == HTP_DSV4_HC_PRE_KERNEL_GATED ? "gated" : "plain";
-        snprintf(str, max_size, "%s nth %u vtcm %u", ktype_str,
-                 (unsigned int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", ktype_str,
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_DSV4_HC_POST) {
        const auto * kparams = node.kparams<htp_dsv4_hc_post_kernel_params>();
        const char * ktype_str = kparams->kernel_type == HTP_DSV4_HC_POST_KERNEL_COMB ? "comb" : "identity";
-        snprintf(str, max_size, "%s nth %u vtcm %u", ktype_str,
-                 (unsigned int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", ktype_str,
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_LIGHTNING_INDEXER) {
        const auto * kparams = node.kparams<htp_lid_kernel_params>();
-        snprintf(str, max_size, "%s tok %u bc %u nth %u vtcm %u",
+        snprintf(str, max_size, "%s nhvx %u tok %u bc %u vtcm %u",
                 kparams->kernel_type == HTP_LID_KERNEL_HMX ? "hmx" : "hvx",
-                 (unsigned int) kparams->n_tokens_tile, (unsigned int) kparams->kv_block,
-                 (unsigned int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->n_tokens_tile,
+                 (unsigned int) kparams->kv_block, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_MUL || node.opcode == HTP_OP_ADD || node.opcode == HTP_OP_ADD_ID ||
               node.opcode == HTP_OP_SUB || node.opcode == HTP_OP_DIV) {
        const auto * kparams = node.kparams<htp_binary_kernel_params>();
-        snprintf(str, max_size, "vtcm %u", (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "nhvx %u vtcm %u", (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
        const auto * kparams = node.kparams<htp_sort_kernel_params>();
-        snprintf(str, max_size, "%s nth %d nchk %d chk %d vtcm %d",
+        snprintf(str, max_size, "%s nhvx %u nchk %u chk %u vtcm %u",
                 node.opcode == HTP_OP_TOP_K ? "top_k" : "argsort",
-                 (int) kparams->n_threads, (int) kparams->n_chunks,
-                 (int) kparams->chunk_elems, (int) kparams->vtcm_size);
+                 (unsigned int) kparams->n_hvx, (unsigned int) kparams->n_chunks,
+                 (unsigned int) kparams->chunk_elems, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_GET_ROWS) {
        const auto * kparams = node.kparams<htp_get_rows_kernel_params>();
        const char * ktype_str = "unknown";
@@ -222,16 +225,16 @@ static std::string format_kernel_params(const htp_opnode & node) {
            case HTP_GET_ROWS_KERNEL_TILED:    ktype_str = "tiled"; break;
            case HTP_GET_ROWS_KERNEL_FLAT:     ktype_str = "flat"; break;
        }
-        snprintf(str, max_size, "%s%s vtcm %u", ktype_str, kparams->n_threads > 1 ? "-multi" : "", (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", ktype_str, (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (htp_op_is_glu(node.opcode)) {
        const auto * kparams = node.kparams<htp_glu_kernel_params>();
-        snprintf(str, max_size, "nth %d vtcm %u", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "nhvx %u vtcm %u", (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_PAD) {
        const auto * kparams = node.kparams<htp_pad_kernel_params>();
-        snprintf(str, max_size, "k%d nth %d vtcm %u", (int) kparams->kernel_type, (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "k%d nhvx %u vtcm %u", (int) kparams->kernel_type, (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_CUMSUM) {
        const auto * kparams = node.kparams<htp_cumsum_kernel_params>();
-        snprintf(str, max_size, "%s nth %d vtcm %u", kparams->use_dma ? "dma" : "hvx", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", kparams->use_dma ? "dma" : "hvx", (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_TOP_K_MASK) {
        const auto * kparams = node.kparams<htp_top_k_mask_kernel_params>();
        const char * ktype_str = "?";
@@ -239,19 +242,19 @@ static std::string format_kernel_params(const htp_opnode & node) {
            case HTP_TOP_K_MASK_KERNEL_ROWS:  ktype_str = "rows"; break;
            case HTP_TOP_K_MASK_KERNEL_PICKS: ktype_str = "picks"; break;
        }
-        snprintf(str, max_size, "%s nth %d tile %u x%u vtcm %u", ktype_str, (int) kparams->n_threads, (unsigned int) kparams->tile_cols, (unsigned int) kparams->n_col_tiles, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u tile %u x%u vtcm %u", ktype_str, (unsigned int) kparams->n_hvx, (unsigned int) kparams->tile_cols, (unsigned int) kparams->n_col_tiles, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_FILL) {
        const auto * kparams = node.kparams<htp_fill_kernel_params>();
-        snprintf(str, max_size, "%s nth %d", kparams->kernel_type == HTP_FILL_KERNEL_CONTIGUOUS ? "contig" : "strided", (int) kparams->n_threads);
+        snprintf(str, max_size, "%s nhvx %u", kparams->kernel_type == HTP_FILL_KERNEL_CONTIGUOUS ? "contig" : "strided", (unsigned int) kparams->n_hvx);
    } else if (node.opcode == HTP_OP_ROLL) {
        const auto * kparams = node.kparams<htp_roll_kernel_params>();
-        snprintf(str, max_size, "k%d nth %d", (int) kparams->kernel_type, (int) kparams->n_threads);
+        snprintf(str, max_size, "k%d nhvx %u", (int) kparams->kernel_type, (unsigned int) kparams->n_hvx);
    } else if (node.opcode == HTP_OP_DIAG) {
        const auto * kparams = node.kparams<htp_diag_kernel_params>();
-        snprintf(str, max_size, "%s nth %d vtcm %u", kparams->use_dma ? "dma" : "hvx", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u vtcm %u", kparams->use_dma ? "dma" : "hvx", (unsigned int) kparams->n_hvx, (unsigned int) kparams->vtcm_size);
    } else if (node.opcode == HTP_OP_SOLVE_TRI) {
        const auto * kparams = node.kparams<htp_solve_tri_kernel_params>();
-        snprintf(str, max_size, "%s nth %d", kparams->batched ? "batched" : "chunked", (int) kparams->n_threads);
+        snprintf(str, max_size, "%s nhvx %u", kparams->batched ? "batched" : "chunked", (unsigned int) kparams->n_hvx);
    } else if (node.opcode == HTP_OP_UPSCALE) {
        const auto * kparams = node.kparams<htp_upscale_kernel_params>();
        const char * ktype_str = "unknown";
@@ -264,7 +267,7 @@ static std::string format_kernel_params(const htp_opnode & node) {
            case HTP_UPSCALE_KERNEL_BILINEAR_AA: ktype_str = "bilinear-aa"; break;
            case HTP_UPSCALE_KERNEL_BICUBIC:     ktype_str = "bicubic";     break;
        }
-        snprintf(str, max_size, "%s nth %d chk %d vtcm %u", ktype_str, (int) kparams->n_threads, (int) kparams->dst_nrows_per_chunk, (unsigned int) kparams->vtcm_size);
+        snprintf(str, max_size, "%s nhvx %u chk %u vtcm %u", ktype_str, (unsigned int) kparams->n_hvx, (unsigned int) kparams->dst_nrows_per_chunk, (unsigned int) kparams->vtcm_size);
    } else {
        snprintf(str, max_size, "----");
    }
@@ -8,31 +8,30 @@
 #include "htp/htp-tensor.h"

 static inline bool htp_is_tiled_weight_type(enum ggml_type type) {
-    return type == GGML_TYPE_Q4_0  || type == GGML_TYPE_Q4_1   ||
-           type == GGML_TYPE_Q8_0  || type == GGML_TYPE_IQ4_NL ||
-           type == GGML_TYPE_MXFP4 || type == GGML_TYPE_Q1_0   ||
-           type == GGML_TYPE_Q6_K   ||
+    return type == GGML_TYPE_Q1_0  || type == GGML_TYPE_Q2_0   ||
+           type == GGML_TYPE_Q4_0  || type == GGML_TYPE_Q4_1   ||
+           type == GGML_TYPE_MXFP4 || type == GGML_TYPE_IQ4_NL ||
+           type == GGML_TYPE_Q8_0  || type == GGML_TYPE_Q6_K   ||
+           type == GGML_TYPE_Q2_K  || type == GGML_TYPE_Q3_K   ||
           type == GGML_TYPE_Q4_K  || type == GGML_TYPE_Q5_K   ||
-           type == GGML_TYPE_Q3_K  || type == GGML_TYPE_Q2_K   ||
-           type == GGML_TYPE_Q2_0  ||
           type == GGML_TYPE_F16   || type == GGML_TYPE_F32    ||
           type == GGML_TYPE_BF16;
 }

 static inline bool htp_is_tiled_weight_tensor(const struct ggml_tensor * t) {
-    if (!ggml_is_contiguous(t) || ggml_is_view(t)) {
-        return false;
-    }
    if (!htp_is_tiled_weight_type(t->type)) {
        return false;
    }
-    if (t->ne[0] < 32) {
+    if (!ggml_is_contiguous(t) || ggml_is_view(t)) {
+        return false;
+    }
+    if (t->ne[0] < 32 || (t->ne[0] % 32 != 0) || t->ne[3] != 1) {
        return false;
    }
    if (!ggml_is_quantized(t->type)) {
-        return (t->ne[0] % 32 == 0) && (t->ne[1] % 32 == 0) && (t->ne[2] == 1) && (t->ne[3] == 1);
+        return t->ne[1] > 1;
    }
-    return t->ne[3] == 1;
+    return true;
 }

 // Size of one tiled row in the DSP tiled layout. The Q3_K..Q6_K tiles store uncompressed scales/mins,
@@ -59,10 +58,6 @@ static inline size_t htp_tiled_row_size(enum ggml_type type, int64_t ne0) {
    return ggml_row_size(type, ne0);
 }

-static inline bool htp_is_hmx_weight_type(enum ggml_type type) {
-    return htp_is_tiled_weight_type(type);
-}
-
 void htp_repack_tensor_tiled(struct ggml_tensor * t, const void * data, size_t size);
 void htp_unpack_tensor_tiled(void * data, const struct ggml_tensor * t, size_t offset, size_t size);

--- a/Show More
+++ b/Show More