llama_moe_cache: run missed experts on CPU, add tuning script (#30966 )

* moe-cache : compute part of the misses on the CPU --moe-cache-cpu-share F computes a share F of the MoE cache misses on the CPU from the host experts instead of uploading them. The graph builds the expert FFN twice: on the device from the cache slots, where the misses left to the CPU read zero slots, and on the CPU from the host experts, where the experts computed on the device have id -1. The two results are added before the down bias. - ggml-cpu: MUL_MAT_ID skips negative ids and writes zero rows - sched: the copy callback is called for all the inputs of a split, and before the host sync when there are no events - the slot lookup callback reads the inputs of the CPU part together with the selected experts, so the CPU part does not wait for the device, and its result is copied back asynchronously - each token uses distinct zero slots, mul_mat_id does not support an expert selected twice by a token - LLAMA_MOE_CACHE_TRACE writes the selected experts of each lookup, for a simulator of the cache * scripts : add moe-cache-tune.py to pick the MoE cache arguments The script wraps a llama-server command: it starts the server at two cache sizes and three CPU shares, fits the upload and CPU speeds of the machine, and replays the routing of the runs through a copy of the cache logic to predict other sizes and shares. It prints the recommended --moe-cache-mib and --moe-cache-cpu-share. docs/moe-cache.md describes the MoE cache arguments and the script. * add am17an as codeowner for MoE cache * moe-cache : upload all the misses until the cache is full With a high --moe-cache-cpu-share, few misses are uploaded, so an empty cache fills slowly and more misses than needed go to the CPU. The CPU share now starts once all the slots of a cache group hold an expert. moe-cache-tune.py replays the same rule. * add backend sync for async transfer queue * moe-cache : let the scheduler copy the results of the CPU part * Fix MSVC compilation * remove sync change * llama-graph : extract the CPU share condition into can_cpu_share * llama-graph : make llm_graph_params::moe_cache const The graph construction no longer updates the MoE cache directly: the CPU parts of the layers are accumulated in the graph result and applied to the cache in llama_context::process_ubatch, after the graph inputs are set and before the graph is computed. Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (Mac Studio, M5 Ultra, 256GB) * scripts : replay the ClassLRU policy in moe-cache-tune.py * moe-cache : pick CPU misses by CRF instead of random (#31853) * moe-cache : remove the backend logic of the CPU part from the graph The CPU part reads the host experts, so the scheduler places it on the CPU without set_tensor_backend. The cache gives no CPU map for batches that a device would offload. Graphs with a CPU part look up their own slot map, so the graph does not register the part with the cache. ggml-backend : the inputs of the next split that are already computed are copied to host memory before the current split runs, so a CPU split does not wait for the device split before it. This replaces the stash of the cache, and the copy callback is again called for host weights only. * llama-graph : return the MoE cache slots and CPU ids in a struct The scheduler prefetch no longer checks that the next split runs on another backend: the inputs of a split are never on its own backend, so the input loop already skips that case. --------- Co-authored-by: Gaurav Garg <gaugarg@nvidia.com> Co-authored-by: Georgi Gerganov <ggerganov@gmail.com> Co-authored-by: Alexey Kopytko <alexey@kopytko.com>
sycl: add f16 support for SWIGLU_OAI (#30857 )
2026-11-07 11:27:38 -06:00 · 2026-11-08 00:20:58 +08:00 · 2026-11-07 09:02:58 -05:00 · 2026-11-07 13:04:13 +02:00 · 2026-11-07 13:03:38 +02:00 · 2026-11-07 10:01:27 +02:00
474 changed files with 28072 additions and 6485 deletions
@@ -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
@@ -64,7 +64,7 @@ RUN mkdir -p /app/full \
    && cp requirements.txt /app/full \
    && cp .devops/tools.sh /app/full/tools.sh

-FROM docker.io/intel/oneapi-toolkit:$ONEAPI_VERSION AS base
+FROM docker.io/intel/deep-learning-essentials:$ONEAPI_VERSION AS base

 ARG BUILD_DATE=N/A
 ARG APP_VERSION=N/A
@@ -79,6 +79,8 @@ LABEL org.opencontainers.image.created=$BUILD_DATE \
      org.opencontainers.image.url=$IMAGE_URL \
      org.opencontainers.image.source=$IMAGE_SOURCE

+COPY --from=build /opt/intel/oneapi/dnnl /opt/intel/oneapi/dnnl
+
 #Following versions are for multiple GPUs, since 26.x has known issue:
 #   https://github.com/ggml-org/llama.cpp/issues/21747,
 #   https://github.com/intel/compute-runtime/issues/921.
@@ -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
@@ -1,57 +0,0 @@
-name: Build relocatable cmake package
-on:
-  workflow_dispatch:
-  workflow_call:
-
-cache-mode: none
-permissions:
-  contents: read
-
-jobs:
-  linux:
-    runs-on: [self-hosted, Linux, fast]
-    steps:
-      - uses: actions/checkout@v6
-        with:
-          fetch-depth: 0
-
-      - 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)
@@ -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,9 +45,6 @@ env:
  LLAMA_ARG_LOG_TIMESTAMPS: 1

 jobs:
-  build-cmake-pkg:
-    uses: ./.github/workflows/build-cmake-pkg.yml
-
  ubuntu:
    strategy:
      matrix:
@@ -242,3 +237,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
@@ -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
@@ -43,13 +43,12 @@ jobs:
    strategy:
      matrix:
        include:
-        # thread and address doesn't run properly on some self hosted machines, so run it on Github instead
        - sanitizer: ADDRESS
-          machine: ubuntu-24.04
+          machine: ubuntu-26.04
        - sanitizer: THREAD
-          machine: ubuntu-24.04
+          machine: ubuntu-26.04
        - sanitizer: UNDEFINED
-          machine: [self-hosted, X64, Linux, CPU]
+          machine: ubuntu-26.04

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

@@ -58,31 +57,24 @@ jobs:
        id: checkout
        uses: actions/checkout@v6

-      # - name: ccache
-      #   uses: ggml-org/ccache-action@v1.2.24
-      #   if: ${{ matrix.sanitizer != 'UNDEFINED' }}
-      #   with:
-      #     key: ctest-${{ matrix.sanitizer }}-ubuntu-24.04
-      #     variant: ccache
-      #     evict-old-files: 1d
-      #     save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.24
+        with:
+          restore: false
+          save: false

-      # with UNDEFINED sanitizer, we have to build in Debug to avoid GCC 13 false-positive warnings
-      - name: Build (undefined)
-        id: cmake_build_undefined
-        if: ${{ matrix.sanitizer == 'UNDEFINED' }}
-        run: |
-          cmake -B build \
-            -DCMAKE_BUILD_TYPE=Debug \
-            -DLLAMA_FATAL_WARNINGS=ON \
-            -DLLAMA_SANITIZE_${{ matrix.sanitizer }}=ON \
-            -DGGML_SANITIZE_${{ matrix.sanitizer }}=ON
-
-          cmake --build build --config Debug -j $(nproc)
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
+        with:
+          key: ctest-sanitizer-${{ matrix.sanitizer }}
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache

      - name: Build
        id: cmake_build
-        if: ${{ matrix.sanitizer == 'ADDRESS' }}
+        if: ${{ matrix.sanitizer == 'ADDRESS' || matrix.sanitizer == 'UNDEFINED' }}
        run: |
          cmake -B build \
            -DCMAKE_BUILD_TYPE=RelWithDebInfo \
@@ -103,6 +95,18 @@ jobs:

          cmake --build build --config RelWithDebInfo -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: ctest-sanitizer-${{ matrix.sanitizer }}
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
      - name: Test
        id: cmake_test
        run: |
@@ -46,18 +46,55 @@ env:

 jobs:
  cpu-x64-high-perf:
-    runs-on: [self-hosted, Linux, X64, CPU]
+    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
        uses: actions/checkout@v6

+      - name: Dependencies
+        id: depends
+        run: |
+          sudo apt update
+          sudo apt install -y build-essential cmake libssl-dev time unzip wget python3 python3-venv python3-dev
+
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.24
+        with:
+          restore: false
+          save: false
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        with:
+          key: self-hosted-cpu-x64-high-perf
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
+
      - name: Test
        id: ggml-ci
        run: |
          LLAMA_ARG_THREADS=$(nproc) GG_BUILD_HIGH_PERF=1 GG_BUILD_EXTRA_TESTS_0=1 bash ./ci/run.sh ~/results/llama.cpp/$RUNNER_NAME ~/mnt/llama.cpp

+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: self-hosted-cpu-x64-high-perf
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
  cpu-arm64-high-perf-graviton4:
    runs-on: ah-ubuntu_24_04-c8g_8x

@@ -92,8 +92,8 @@ jobs:
            { "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": "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-24.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-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" },
@@ -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: |
@@ -0,0 +1,37 @@
+name: "PR video demo label"
+
+on:
+  pull_request_target:
+    types: [opened, edited, reopened, synchronize]
+
+cache-mode: none
+permissions:
+  pull-requests: write
+  issues: write
+
+jobs:
+  video-demo-label:
+    runs-on: ubuntu-slim
+    steps:
+      - name: Add or remove the video demo label
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          PR_URL: ${{ github.event.pull_request.html_url }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          # passed through env, so the body is only matched, never run, and is safe to read from a fork
+          BODY: ${{ github.event.pull_request.body }}
+          HAS_LABEL: ${{ contains(github.event.pull_request.labels.*.name, 'video demo') }}
+        run: |
+          # sections written by demo:pr (see tools/ui/demo/cli.ts): the managed comment
+          # markers, or a preview linking to a demo video uploaded for this PR
+          has_demo=false
+          if { grep -qF '<!-- demo:start -->' <<< "$BODY" && grep -qF '<!-- demo:end -->' <<< "$BODY"; } ||
+            grep -qE "/resolve/videos/pr/$PR_NUMBER/[0-9a-f]{10}(-[0-9]+)?\.mp4" <<< "$BODY"; then
+            has_demo=true
+          fi
+
+          if [ "$has_demo" = true ] && [ "$HAS_LABEL" = false ]; then
+            gh pr edit "$PR_URL" --add-label 'video demo'
+          elif [ "$has_demo" = false ] && [ "$HAS_LABEL" = true ]; then
+            gh pr edit "$PR_URL" --remove-label 'video demo'
+          fi
@@ -177,8 +177,8 @@ jobs:
            - [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
@@ -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
@@ -1,4 +1,4 @@
-# date: Wed Oct 28 12:38:26 EET 2026
+# date: Wed Nov  4 12:03:31 EET 2026
 # this file is auto-generated by scripts/gen-authors.sh

 Нияз Гарифзянов <112617865+garrnizon@users.noreply.github.com>
@@ -48,6 +48,7 @@ Abheek Gulati <abheekg@hotmail.com>
 abhijain1204fujitsu <139222713+abhijain1204fujitsu@users.noreply.github.com>
 Abhijit Ramesh <abhijitramesh2k@gmail.com>
 abhijitb11 <113058133+abhijitb11@users.noreply.github.com>
+Abhijith Jajam <129675815+jajamabhijith@users.noreply.github.com>
 Abhilash Majumder <30946547+abhilash1910@users.noreply.github.com>
 Abhinay Krishna <abhinaykrishna60@gmail.com>
 Abhiram <78226909+geckguy@users.noreply.github.com>
@@ -203,14 +204,17 @@ Ankit Khandelwal <ankk98@gmail.com>
 Ankur Verma <31362771+ankurvdev@users.noreply.github.com>
 anon998 <131767832+anon998@users.noreply.github.com>
 Anri Lombard <anri.m.lombard@gmail.com>
+Anthony Maltsev <77708730+anthonymaltsev@users.noreply.github.com>
 Anthony Umfer <aumfer@gmail.com>
 Anthony Van de Gejuchte <anthonyvdgent@gmail.com>
 Anthony Wang <anthonywang903@gmail.com>
 antichristHater <142441588+antichristHater@users.noreply.github.com>
 Antoine Viallon <antoine@lesviallon.fr>
+Anton I <inflyer@hotmail.com>
 Anton Mitkov <anton_b_mitkov@abv.bg>
 Anton Mitkov <anton.mitkov@codeplay.com>
 Antonis Makropoulos <benuix@gmail.com>
+anubhavfujitsu <Anubhav.Srivastava@fujitsu.com>
 Anudit Nagar <nagaranudit@gmail.com>
 Anuj Attri <anujattri01@gmail.com>
 anujj <ajalota@nvidia.com>
@@ -293,6 +297,7 @@ Bernat Vadell <hounter.caza@gmail.com>
 Bernhard M. Wiedemann <githubbmwprimary@lsmod.de>
 Bert Wagner <github@bertwagner.com>
 Bertay Eren <39909689+bertaye@users.noreply.github.com>
+Bhargav Krish <bkrishnamurt@nvidia.com>
 Bhavik Sharda <10757940+BLSharda@users.noreply.github.com>
 bhubbb <79117352+bhubbb@users.noreply.github.com>
 Bill Sideris <bill88t@feline.gr>
@@ -533,6 +538,7 @@ Don Mahurin <dmahurin@users.noreply.github.com>
 Dong Won Kim <63934649+ddwkim@users.noreply.github.com>
 Donghyeon Jeong <54725479+djeong20@users.noreply.github.com>
 Dongliang Wei <121270393+wdl339@users.noreply.github.com>
+Dongsheng Luo <luods1995@outlook.com>
 Doomsdayrs <38189170+Doomsdayrs@users.noreply.github.com>
 DooWoong Lee (David) <manics99@naver.com>
 DorianRudolph <dorianrudo97@googlemail.com>
@@ -612,6 +618,7 @@ Faez Shakil <faez.shakil@gmail.com>
 fairydreaming <166155368+fairydreaming@users.noreply.github.com>
 Faisal Zaghloul <faisal.zaghloul@gmail.com>
 Faisal Zaghloul <quic_fzaghlou@quicinc.com>
+FamousMonsterr <famousmonster@ya.ru>
 Fan Shupei <dymarkfan@outlook.com>
 FantasyGmm <16450052+FantasyGmm@users.noreply.github.com>
 fanyang <fanyang89@outlook.com>
@@ -654,6 +661,7 @@ Frederik Vogel <Schaltfehler@users.noreply.github.com>
 Fredrik Hultin <noname@nurd.se>
 fredzillman <fzillman@gmail.com>
 frob <rick+github@frob.com.au>
+frogscooper <frogscooper@gmail.com>
 Frost-54 <91168148+Frost-54@users.noreply.github.com>
 Frosty40 <newjordan@gmail.com>
 Funtowicz Morgan <mfuntowicz@users.noreply.github.com>
@@ -694,6 +702,7 @@ Gilad S. <7817232+giladgd@users.noreply.github.com>
 Giovanni Rivera <rivera.giovanni271@gmail.com>
 github-actions[bot] <github-actions[bot]@users.noreply.github.com>
 GittyBurstein <g0534163997@gmail.com>
+Giulianno <lernexbusiness@gmail.com>
 Giuseppe Scrivano <giuseppe@scrivano.org>
 Giuseppe Scrivano <gscrivan@redhat.com>
 GiviMAD <GiviMAD@users.noreply.github.com>
@@ -762,6 +771,7 @@ Henry147147 <44851451+Henry147147@users.noreply.github.com>
 Herman Semenoff <GermanAizek@yandex.ru>
 Herman Semenov <GermanAizek@yandex.ru>
 Hesen Peng <hesen.peng@gmail.com>
+hesiru <siruhe666@gmail.com>
 hey-gm <heitorgm@outlook.com>
 HighDoping <highdoping@gmail.com>
 HimariO <dsfhe49854@gmail.com>
@@ -919,6 +929,7 @@ Jesse Ikonen <jesse.ikonen@gmail.com>
 Jesse Jojo Johnson <williamsaintgeorge@gmail.com>
 Jesse LaRose <jesse@taey.ai>
 Jesse Posner <jesse.posner@gmail.com>
+Jesse Symons <155509533+Torchit1@users.noreply.github.com>
 Jesus Gulfo <jesdga95@gmail.com>
 Jesus Talavera <145992175+jesus-talavera-ibm@users.noreply.github.com>
 Jetson Tan <tanzongyouyi@outlook.com>
@@ -1071,6 +1082,7 @@ Kevin Liu <4396kevinliu@gmail.com>
 Kevin Lo <kevlo@kevlo.org>
 Kevin Pouget <kpouget@redhat.com>
 Kevin Wang <kevmo314@gmail.com>
+Keyu Zhu <inuebisu@outlook.com>
 Khashayar Ghafouri <43180261+khashayarghafouri@users.noreply.github.com>
 khimaros <me@khimaros.com>
 Kilian Hu <90606809+kilian-hu@users.noreply.github.com>
@@ -1078,6 +1090,7 @@ Kilian Krampf <kilian@krampf.de>
 kiltyj <kiltyj@gmail.com>
 Kim S. <polydecay@users.noreply.github.com>
 kimminsu <80271594+kimminsu38oo@users.noreply.github.com>
+Kirill A. Korinsky <kirill@korins.ky>
 KITAITI Makoto <KitaitiMakoto@gmail.com>
 kiwi <122582483+kiwi142857@users.noreply.github.com>
 klosax <131523366+klosax@users.noreply.github.com>
@@ -1136,6 +1149,7 @@ leejet <leejet714@gmail.com>
 Leng Yue <lengyue@lengyue.me>
 Lennart Austenfeld <53152202+l-austenfeld@users.noreply.github.com>
 leo-pony <nengjunma@outlook.com>
+Leon <118652405+LeON-Nie-code@users.noreply.github.com>
 Leon Knauer <git@leonknauer.com>
 Leonard Mosescu <tlemo@users.noreply.github.com>
 leonardHONG <2695316095@qq.com>
@@ -1146,8 +1160,10 @@ lexasub <lexakopp2212@gmail.com>
 lgai-exaone <exaonemodels@lgresearch.ai>
 lhez <lih@qti.qualcomm.com>
 lhez <quic_lih@quicinc.com>
+Li Chuyang <1418587903@qq.com>
 Li Pengzhan <151381994+Lpzhan931@users.noreply.github.com>
 Li Tan <tanliboy@gmail.com>
+li-lizhe <82501526+li-lizhe@users.noreply.github.com>
 liangliangchang <lichang@amd.com>
 liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>
 limitedAtonement <limitedAtonement@users.noreply.github.com>
@@ -1166,6 +1182,7 @@ LoganDark <github@logandark.mozmail.com>
 Loïc Carrère <loic.carrere@gmail.com>
 lon <114724657+longregen@users.noreply.github.com>
 loonerin <132926317+loonerin@users.noreply.github.com>
+Lorenzo Bona <lorenzobona@tuta.com>
 losewayy <1134285612@qq.com>
 LostRuins <39025047+LostRuins@users.noreply.github.com>
 LostRuins Concedo <39025047+LostRuins@users.noreply.github.com>
@@ -1253,6 +1270,7 @@ Masaya, Kato <62578291+msy-kato@users.noreply.github.com>
 mashdragon <122402293+mashdragon@users.noreply.github.com>
 Mason Milburn <masonmilby@gmail.com>
 MasterYi1024 <39848311+MasterYi1024@users.noreply.github.com>
+Maternion <98753158+maternion@users.noreply.github.com>
 Mateusz Charytoniuk <mateusz.charytoniuk@protonmail.com>
 Matheus C. França <matheus-catarino@hotmail.com>
 Matheus Gabriel Alves Silva <matheusgasource@gmail.com>
@@ -1350,6 +1368,7 @@ Miwa / Ensan <63481257+ensan-hcl@users.noreply.github.com>
 miyan <1138989048@qq.com>
 mj-shifu <77107165+mj-shifu@users.noreply.github.com>
 mkoker <132301062+mkoker@users.noreply.github.com>
+mmalpe-nv <mmalpe@nvidia.com>
 mmyjona <jonathan.gonse@gmail.com>
 mnehete32 <33429707+mnehete32@users.noreply.github.com>
 Mohamed Elashri <git@elashri.com>
@@ -1368,6 +1387,7 @@ mtmcp <141645996+mtmcp@users.noreply.github.com>
 muggle-stack <promuggle@qq.com>
 Muhammad Saad <saadydev@gmail.com>
 Muhammad Salem <salem.ebo@gmail.com>
+MuoDoo <36472129+MuoDoo@users.noreply.github.com>
 Murilo Santana <mvrilo@gmail.com>
 Musab Gultekin <musabgultekin@users.noreply.github.com>
 musoles <135031143+musoles@users.noreply.github.com>
@@ -1416,10 +1436,12 @@ Nicolas Budyn <16596345+nbud@users.noreply.github.com>
 Nicolas Mowen <nickmowen213@gmail.com>
 Nicolás Pérez <nicolas_perez@brown.edu>
 Nicolò Scipione <nicolo.scipione@codeplay.com>
+Nidhi <nidhisharma16635@gmail.com>
 Nigel Bosch <pnigelb@gmail.com>
 Nigel Jones <nigel.l.jones+git@gmail.com>
 Nik Bogatyrev <50734723+notforu777@users.noreply.github.com>
 Nikhil Jain <nikhil.jain0987@gmail.com>
+Nikhil Sharma <150424968+nikhilfujitsu@users.noreply.github.com>
 Nikita Gordeev <iamfobey@gmail.com>
 Nikita Sarychev <42014488+sARY77@users.noreply.github.com>
 Niklas Korz <niklas@niklaskorz.de>
@@ -1592,6 +1614,7 @@ Raman Shinde <raman.shinde15@gmail.com>
 Rand Xie <randxiexyy29@gmail.com>
 Randall Fitzgerald <randall@dasaku.net>
 Random Fly <renfei8@live.cn>
+randomized <286843573@qq.com>
 rankaiyx <rankaiyx@foxmail.com>
 rankaiyx <rankaiyx@rankaiyx.com>
 RapidMark <32768622+RapidMark@users.noreply.github.com>
@@ -2009,6 +2032,7 @@ Tungsten842 <quantmint@protonmail.com>
 Tushar <ditsuke@protonmail.com>
 tv1wnd <55383215+tv1wnd@users.noreply.github.com>
 tyronecai <tyronecai@163.com>
+Uber <20715993+YagUber@users.noreply.github.com>
 ubergarm <leimgrub@gmail.com>
 ubik2 <ubik2@users.noreply.github.com>
 UEXTM.com <84163508+uextm@users.noreply.github.com>
@@ -2242,4 +2266,5 @@ Zijun Yu <zijun.yu@intel.com>
 zql <37731799+zqlcode@users.noreply.github.com>
 zrm <trustiosity.zrm@gmail.com>
 Zsapi <martin1.zsapka@gmail.com>
+zyang-dev <267119621+zyang-dev@users.noreply.github.com>
 zzzzwc <tiddar@foxmail.com>
@@ -4,7 +4,7 @@ include(CheckIncludeFileCXX)

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

@@ -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
@@ -296,25 +296,6 @@ function gg_run_test_llama_archs_tensor_split {
    set +e
 }

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

 function gg_run_test_scripts {
@@ -456,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

@@ -701,7 +678,6 @@ test $ret -eq 0 && gg_run ctest_release

 test $ret -eq 0 && gg_run test_backend_ops

-test $ret -eq 0 && gg_run test_llama_archs_models
 test $ret -eq 0 && gg_run test_llama_archs_tensor_split

 if [ -z ${GG_BUILD_LOW_PERF} ]; then
@@ -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;
    }
@@ -710,7 +730,7 @@ static void common_params_apply_system_config(common_params & params, llama_exam
    std::vector<std::filesystem::path> paths;

 #if defined(_WIN32)
-    const std::filesystem::path program_data = common_get_path_from_env("PROGRAMDATA");
+    const std::filesystem::path program_data = common_get_env("PROGRAMDATA");
    if (!program_data.empty()) {
        paths.push_back(program_data / "llama.cpp" / "config.ini");
    }
@@ -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"
@@ -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();
@@ -983,20 +963,9 @@ void common_set_env(const std::string & name, const std::string & value) {
 #endif
 }

-std::filesystem::path common_get_path_from_env(const std::string & name) {
-#if defined(_WIN32)
-    const std::wstring wname = utf8_to_wstring(name);
-    const wchar_t * wvalue = _wgetenv(wname.c_str());
-    return wvalue ? std::filesystem::path(wvalue) : std::filesystem::path();
-#else
-    const char * value = std::getenv(name.c_str());
-    return value ? std::filesystem::path(value) : std::filesystem::path();
-#endif
-}
-
 #if !defined(_WIN32)
 static std::filesystem::path get_home_directory() {
-    std::filesystem::path home = common_get_path_from_env("HOME");
+    std::filesystem::path home = common_get_env("HOME");
    if (!home.empty()) {
        return home;
    }
@@ -1009,19 +978,19 @@ static std::filesystem::path get_home_directory() {
 #endif

 std::filesystem::path fs_get_cache_directory() {
-    std::filesystem::path cache_directory = common_get_path_from_env("LLAMA_CACHE");
+    std::filesystem::path cache_directory = common_get_env("LLAMA_CACHE");
    if (!cache_directory.empty()) {
        return cache_directory;
    }
 #if defined(_WIN32)
-    cache_directory = common_get_path_from_env("LOCALAPPDATA");
+    cache_directory = common_get_env("LOCALAPPDATA");
    if (cache_directory.empty()) {
        throw std::runtime_error("Failed to find %LOCALAPPDATA% directory");
    }
 #elif defined(__APPLE__)
    cache_directory = get_home_directory() / "Library/Caches";
 #else
-    cache_directory = common_get_path_from_env("XDG_CACHE_HOME");
+    cache_directory = common_get_env("XDG_CACHE_HOME");
    if (cache_directory.empty()) {
        cache_directory = get_home_directory() / ".cache";
    }
@@ -1032,12 +1001,12 @@ std::filesystem::path fs_get_cache_directory() {
 std::filesystem::path fs_get_config_directory() {
    std::filesystem::path config_directory;
 #if defined(_WIN32)
-    config_directory = common_get_path_from_env("APPDATA");
+    config_directory = common_get_env("APPDATA");
    if (config_directory.empty()) {
        throw std::runtime_error("Failed to find %APPDATA% directory");
    }
 #else
-    config_directory = common_get_path_from_env("XDG_CONFIG_HOME");
+    config_directory = common_get_env("XDG_CONFIG_HOME");
    if (config_directory.empty()) {
        config_directory = get_home_directory() / ".config";
    }
@@ -1053,7 +1022,7 @@ std::filesystem::path fs_get_cache_file(const std::string & filename) {
    if (ec) {
        throw std::runtime_error("failed to create cache directory: " + cache_directory.string());
    }
-    return cache_directory / std::filesystem::u8path(filename);
+    return cache_directory / filename;
 }

 //
@@ -1344,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) {
@@ -1783,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

 //
@@ -919,9 +919,6 @@ std::string  wstring_to_utf8(const std::wstring & str);
 std::string common_get_env(const std::string & name);
 void        common_set_env(const std::string & name, const std::string & value);

-// reads a path from the environment, an unset variable gives an empty path
-std::filesystem::path common_get_path_from_env(const std::string & name);
-
 //
 // Filesystem utils
 //
@@ -1293,6 +1290,20 @@ enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *);
 // prompt utils
 //

+// leave default-inserted elements uninitialized; callers must fill them before use
+template <typename T>
+struct common_no_init_allocator : std::allocator<T> {
+    template <typename U> struct rebind { using other = common_no_init_allocator<U>; };
+
+    common_no_init_allocator() = default;
+    template <typename U> common_no_init_allocator(const common_no_init_allocator<U> &) {}
+
+    template <typename U> void construct(U * p) { ::new ((void *) p) U; }
+};
+
+// sequence state data, filled by llama_state_seq_get_data_ext()
+using common_state_data = std::vector<uint8_t, common_no_init_allocator<uint8_t>>;
+
 struct common_prompt_checkpoint {
    int64_t n_tokens;

@@ -1302,8 +1313,8 @@ struct common_prompt_checkpoint {
    llama_pos pos_min;
    llama_pos pos_max;

-    std::vector<uint8_t> data_tgt;
-    std::vector<uint8_t> data_dft;
+    common_state_data data_tgt;
+    common_state_data data_dft;

    // (optional) speculative-decoding implementation state stashed with the checkpoint
    // (e.g. eagle3's deferred-boundary g_embd row)
@@ -45,7 +45,7 @@ static const std::string PARTIAL_SUFFIX = ".downloadInProgress";

 static void write_etag(const std::string & path, const std::string & etag) {
    const std::string etag_path = path + ".etag";
-    fs_write_atomic(std::filesystem::u8path(etag_path), etag);
+    fs_write_atomic(etag_path, etag);
    LOG_DBG("%s: file etag saved: %s\n", __func__, etag_path.c_str());
 }

@@ -293,8 +293,7 @@ static int common_download_file_single_online(const std::string & url,

    // alt_path is the HF snapshot: the blob may have been moved there when symlinks are unavailable
    auto cache_hit = [&] {
-        return std::filesystem::exists(std::filesystem::u8path(path)) ||
-               (!alt_path.empty() && std::filesystem::exists(std::filesystem::u8path(alt_path)));
+        return std::filesystem::exists(path) || (!alt_path.empty() && std::filesystem::exists(alt_path));
    };

    if (skip_etag && cache_hit()) {
@@ -511,7 +510,7 @@ int common_download_file_single(const std::string & url,
        return common_download_file_single_online(url, path, online_opts, skip_etag, lock_path, alt_path);
    }

-    if (!std::filesystem::exists(std::filesystem::u8path(path))) {
+    if (!std::filesystem::exists(path)) {
        LOG_ERR("%s: required file is not available in cache (offline mode): %s\n", __func__, path.c_str());
        return -1;
    }
@@ -748,7 +747,8 @@ static hf_cache::hf_file find_best_model(const hf_cache::hf_files & files,
    }

    for (const auto & t : tags) {
-        std::regex pattern(t + "[.-]", std::regex::icase);
+        // match the tag as a full token, so ":F16" does not match a BF16 file
+        std::regex pattern("(^|[-./])" + t + "[.-]", std::regex::icase);
        for (const auto & f : files) {
            if (gguf_filename_is_model(f.path) &&
                std::regex_search(f.path, pattern)) {
@@ -1211,7 +1211,7 @@ bool common_download_remove_sidecars(const std::string & hf_repo_with_tag, const
    for (const auto & s : sidecars) {
        const auto & f = plan.*s.file;
        if (!f.local_path.empty()) {
-            to_remove.emplace_back(std::filesystem::u8path(f.local_path));
+            to_remove.emplace_back(f.local_path);
        }
    }
    if (to_remove.empty()) {
@@ -61,6 +61,7 @@ static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(
        const llama_context_params * cparams,
        std::vector<ggml_backend_dev_t> & devs,
        uint32_t & hp_ngl,
+        uint32_t & hp_ngl_skip,
        uint32_t & hp_n_ctx_train,
        uint32_t & hp_n_expert,
        ggml_log_level log_level) {
@@ -152,6 +153,8 @@ static std::vector<llama_device_memory_data> common_get_device_memory_data_impl(
    if (mparams->load_mtp) {
        hp_ngl    += llama_model_n_layer_nextn(model);
    }
+    // without MTP the NextN layers hold no weights but the model loader still counts them for n_gpu_layers
+    hp_ngl_skip    = mparams->load_mtp ? 0 : llama_model_n_layer_nextn(model);
    hp_n_ctx_train = llama_model_n_ctx_train(model);
    hp_n_expert    = llama_model_n_expert(model);

@@ -240,11 +243,12 @@ common_device_memory_data_vec common_get_device_memory_data(
        const llama_context_params * cparams,
        std::vector<ggml_backend_dev_t> & devs,
        uint32_t & hp_ngl,
+        uint32_t & hp_ngl_skip,
        uint32_t & hp_n_ctx_train,
        uint32_t & hp_n_expert,
        ggml_log_level log_level) {
    std::vector<llama_device_memory_data> impl = common_get_device_memory_data_impl(
-            path_model, mparams, cparams, devs, hp_ngl, hp_n_ctx_train, hp_n_expert, log_level);
+            path_model, mparams, cparams, devs, hp_ngl, hp_ngl_skip, hp_n_ctx_train, hp_n_expert, log_level);

    common_device_memory_data_vec ret(impl.size());
    for (size_t i = 0; i < impl.size(); i++) {
@@ -270,6 +274,7 @@ static void common_params_fit_impl(

    std::vector<ggml_backend_dev_t> devs;
    uint32_t hp_ngl = 0; // hparams.n_gpu_layers
+    uint32_t hp_nsk = 0; // number of skipped NextN layers, counted by the model loader but not by hp_ngl
    uint32_t hp_nct = 0; // hparams.n_ctx_train
    uint32_t hp_nex = 0; // hparams.n_expert

@@ -291,6 +296,7 @@ static void common_params_fit_impl(
        if (dmds_extra.empty() || n_ctx_extra != cparams->n_ctx) {
            std::vector<ggml_backend_dev_t> devs_extra;
            uint32_t ngl_extra = 0;
+            uint32_t nsk_extra = 0;
            uint32_t nct_extra = 0;
            uint32_t nex_extra = 0;

@@ -302,7 +308,7 @@ static void common_params_fit_impl(
            dmds_t measured;
            try {
                measured = common_get_device_memory_data_impl(
-                    extra->path_model, extra->mparams, extra->cparams, devs_extra, ngl_extra, nct_extra, nex_extra, log_level);
+                    extra->path_model, extra->mparams, extra->cparams, devs_extra, ngl_extra, nsk_extra, nct_extra, nex_extra, log_level);
            } catch (const std::runtime_error & e) {
                // the extra model is optional, fit the main model alone rather than giving up
                LOG_WRN("%s: failed to measure the memory of the extra model, fitting without it: %s\n", __func__, e.what());
@@ -342,7 +348,7 @@ static void common_params_fit_impl(
    // step 1: get data for default parameters and check whether any changes are necessary in the first place

    LOG_TRC("%s: getting device memory data for initial parameters:\n", __func__);
-    dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+    dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nsk, hp_nct, hp_nex, log_level);

    // saturate instead of overflowing, this also preserves the UINT32_MAX sentinel of n_ctx_min:
    const uint32_t n_ctx_max       = (uint32_t) std::min<uint64_t>(uint64_t(hp_nct)    * n_streams, UINT32_MAX);
@@ -354,7 +360,7 @@ static void common_params_fit_impl(
        if (n_streams > 1) {
            LOG_TRC("%s: context size unset and KV cache not unified -> using %" PRIu32 " for %" PRIu32 " sequences:\n",
                __func__, n_ctx_max, n_streams);
-            dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+            dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nsk, hp_nct, hp_nex, log_level);
        }
    }
    add_extra_memory(dmds_full);
@@ -493,7 +499,7 @@ static void common_params_fit_impl(

                    int64_t sum_projected_used_min_ctx = 0;
                    cparams->n_ctx = n_ctx_min_total;
-                    dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+                    dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nsk, hp_nct, hp_nex, log_level);
                    add_extra_memory(dmds_min_ctx);
                    if (nd == 0) {
                        sum_projected_used_min_ctx = dmds_min_ctx.back().mb.total();
@@ -640,6 +646,20 @@ static void common_params_fit_impl(
        assert(uint32_t(mparams.n_gpu_layers) <= hp_ngl + 1);
        uint32_t il0 = hp_ngl + 1 - mparams.n_gpu_layers; // start index for tensor buft overrides

+        // the skipped NextN layers come right before the output layer, put them on the same device as the output layer
+        // so that the layers assigned by the model loader match the layers above
+        if (hp_nsk > 0 && mparams.n_gpu_layers > 0) {
+            mparams.n_gpu_layers += hp_nsk;
+            if (nd > 1) {
+                // the last devices have no layers if not even the output layer fits
+                size_t id_last = nd - 1;
+                while (ngl_per_device[id_last].n_layer == 0) {
+                    id_last--;
+                }
+                tensor_split[id_last] += hp_nsk;
+            }
+        }
+
        mparams.tensor_split = tensor_split;

        size_t itbo = 0;
@@ -675,7 +695,7 @@ static void common_params_fit_impl(
        set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, mparams_copy);

        dmds_t dmd_nl = common_get_device_memory_data_impl(
-            path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+            path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nsk, hp_nct, hp_nex, log_level);
        add_extra_memory(dmd_nl);

        LOG_TRC("%s: memory for test allocation by device:\n", func_name);
@@ -704,7 +724,7 @@ static void common_params_fit_impl(

        LOG_TRC("%s: getting device memory data with all MoE tensors moved to system memory:\n", __func__);
        dmds_t dmds_cpu_moe = common_get_device_memory_data_impl(
-            path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
+            path_model, mparams, cparams, devs, hp_ngl, hp_nsk, hp_nct, hp_nex, log_level);
        add_extra_memory(dmds_cpu_moe);

        for (size_t id = 0; id < nd; id++) {
@@ -1185,10 +1205,11 @@ void common_fit_print(
        llama_context_params * cparams) {
    std::vector<ggml_backend_dev_t> devs;
    uint32_t hp_ngl = 0; // hparams.n_gpu_layers
+    uint32_t hp_nsk = 0; // number of skipped NextN layers, counted by the model loader but not by hp_ngl
    uint32_t hp_nct = 0; // hparams.n_ctx_train
    uint32_t hp_nex = 0; // hparams.n_expert

-    auto dmd = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, GGML_LOG_LEVEL_ERROR);
+    auto dmd = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nsk, hp_nct, hp_nex, GGML_LOG_LEVEL_ERROR);
    GGML_ASSERT(dmd.size() == devs.size() + 1);

    for (size_t id = 0; id < devs.size(); id++) {
@@ -104,6 +104,7 @@ common_device_memory_data_vec common_get_device_memory_data(
         const llama_context_params * cparams,
    std::vector<ggml_backend_dev_t> & devs,
                           uint32_t & hp_ngl,
+                           uint32_t & hp_ngl_skip,
                           uint32_t & hp_n_ctx_train,
                           uint32_t & hp_n_expert,
                     ggml_log_level   log_level);
@@ -45,7 +45,7 @@ static fs::path get_cache_directory() {
            {HOME_DIR,                fs::path(".cache") / "huggingface" / "hub"}
        };
        for (const auto & entry : entries) {
-            if (fs::path base = common_get_path_from_env(entry.var); !base.empty()) {
+            if (fs::path base = common_get_env(entry.var); !base.empty()) {
                return entry.path.empty() ? base : base / entry.path;
            }
        }
@@ -242,7 +242,7 @@ static std::string get_repo_commit(const std::string & repo_id,
                LOG_WRN("%s: skip invalid commit: %s\n", __func__, _commit.c_str());
                continue;
            }
-            const fs::path candidate = fs::u8path(_name);
+            const fs::path candidate = _name;

            if (!is_valid_subpath(refs_path, candidate)) {
                LOG_WRN("%s: skip invalid branch: %s\n", __func__, _name.c_str());
@@ -317,7 +317,7 @@ hf_files get_repo_files(const std::string & repo_id,
            file.repo_id = repo_id;
            file.path = item["path"].get<std::string>();

-            const fs::path subpath = fs::u8path(file.path);
+            const fs::path subpath = file.path;

            if (!is_valid_subpath(commit_path, subpath)) {
                LOG_WRN("%s: skip invalid path: %s\n", __func__, file.path.c_str());
@@ -434,8 +434,7 @@ hf_files get_cached_files(const std::string & repo_id) {
            if (!path.empty()) {
                hf_file file;
                file.repo_id = _repo_id;
-                const auto generic_path = path.generic_u8string();
-                file.path = std::string(generic_path.begin(), generic_path.end());
+                file.path = path.generic_string();
                file.local_path = entry.path().string();
                file.final_path = file.local_path;
                files.push_back(std::move(file));
@@ -450,8 +449,8 @@ std::string finalize_file(const hf_file & file, const std::function<bool()> & ke
    static std::atomic<bool> symlinks_disabled{false};

    std::error_code ec;
-    fs::path local_path = fs::u8path(file.local_path);
-    fs::path final_path = fs::u8path(file.final_path);
+    fs::path local_path = file.local_path;
+    fs::path final_path = file.final_path;

    if (local_path == final_path || fs::exists(final_path, ec)) {
        return file.final_path;
@@ -385,7 +385,7 @@ static bool is_draft_file(const std::string & fname) {
 }

 common_presets common_preset_context::load_from_models_dir(const std::string & models_dir) const {
-    const std::filesystem::path dir = std::filesystem::u8path(models_dir);
+    const std::filesystem::path dir = models_dir;
    if (!std::filesystem::exists(dir) || !std::filesystem::is_directory(dir)) {
        throw std::runtime_error(string_format("error: '%s' does not exist or is not a directory\n", models_dir.c_str()));
    }
@@ -1613,12 +1613,12 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
            }
        }

-        llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ false);
-        llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
-
        is_mem_shared = llama_get_ctx_other(ctx_dft) == ctx_tgt;
        chain_heads   = n_mtp_layers > 1 && !is_mem_shared;

+        llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ is_mem_shared);
+        llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
+
        if (chain_heads) {
            this->params.n_max = std::min(this->params.n_max, n_mtp_layers);

@@ -1772,16 +1772,28 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
            }

            const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1;
-            verify_h_rows[seq_id] = n_rows;
-            verify_h[seq_id].resize((size_t) n_rows * n_embd);
+            if (is_mem_shared && !batch_in.tokens[i_batch_end[seq_id]].output) {
+                verify_h_rows[seq_id] = 0;
+                continue;
+            }

-            for (int32_t i = 0; i < n_rows; ++i) {
-                const float * h = llama_get_embeddings_nextn_ith(ctx_tgt, i_batch_beg[seq_id] + i);
+            const bool all_rows_output = std::all_of(
+                    batch_in.tokens.begin() + i_batch_beg[seq_id],
+                    batch_in.tokens.begin() + i_batch_end[seq_id] + 1,
+                    [](const common_batch::token & token) { return token.output; });
+            const int32_t h_rows = is_mem_shared && !all_rows_output ? 1 : n_rows;
+            const int32_t h_offset = is_mem_shared && !all_rows_output ? i_batch_end[seq_id] : i_batch_beg[seq_id];
+
+            verify_h_rows[seq_id] = h_rows;
+            verify_h[seq_id].resize((size_t) h_rows * n_embd);
+
+            for (int32_t i = 0; i < h_rows; ++i) {
+                const float * h = llama_get_embeddings_nextn_ith(ctx_tgt, h_offset + i);
                std::memcpy(verify_h[seq_id].data() + (size_t) i * n_embd, h, row_bytes);
            }

            std::memcpy(pending_h[seq_id].data(),
-                    verify_h[seq_id].data() + (size_t) (n_rows - 1) * n_embd, row_bytes);
+                    verify_h[seq_id].data() + (size_t) (h_rows - 1) * n_embd, row_bytes);
        }

        return true;
@@ -1967,6 +1979,29 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
        const size_t row_bytes = (size_t) n_embd * sizeof(float);
        std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) i_h * n_embd, row_bytes);
    }
+
+    bool get_state(llama_seq_id seq_id, std::vector<uint8_t> & data) const override {
+        if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
+            return false;
+        }
+
+        const size_t row_bytes = (size_t) n_embd * sizeof(float);
+        data.resize(row_bytes);
+        std::memcpy(data.data(), pending_h[seq_id].data(), row_bytes);
+        return true;
+    }
+
+    void set_state(llama_seq_id seq_id, const std::vector<uint8_t> & data) override {
+        if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
+            return;
+        }
+
+        const size_t row_bytes = (size_t) n_embd * sizeof(float);
+        if (data.size() != row_bytes) {
+            return;
+        }
+        std::memcpy(pending_h[seq_id].data(), data.data(), row_bytes);
+    }
 };

 // state of self-speculation (simple implementation, not ngram-map)
@@ -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",
@@ -45,6 +46,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "ChameleonForConditionalGeneration": "chameleon",
    "ChatGLMForConditionalGeneration": "chatglm",
    "ChatGLMModel": "chatglm",
+    "ChatterboxModel": "chatterbox",
    "ClefModel": "clef",
    "CloverMoeModelLM": "clover",
    "CodeShellForCausalLM": "codeshell",
@@ -61,6 +63,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "DeepseekV2ForCausalLM": "deepseek",
    "DeepseekV3ForCausalLM": "deepseek",
    "DeepseekV32ForCausalLM": "deepseek",
+    "Xing4_0ForCausalLM": "xing",
    "DFlashDraftModel": "qwen",
    "DFlash2DraftModel": "qwen",
    "LiLiCorrDraftModel": "nemotron",
@@ -94,6 +97,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "FalconForCausalLM": "falcon",
    "FalconH1ForCausalLM": "falcon_h1",
    "FalconMambaForCausalLM": "mamba",
+    "FishOmniModel": "drama3",
    "GPT2LMHeadModel": "gpt2",
    "GPTBigCodeForCausalLM": "starcoder",
    "GPTNeoXForCausalLM": "gptneox",
@@ -154,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",
@@ -197,6 +202,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "MambaLMHeadModel": "mamba",
    "MapleForCausalLM": "maple",
    "MellumForCausalLM": "mellum",
+    "Metis16ForCausalLM": "metis",
    "MiMoV2FlashForCausalLM": "mimo",
    "MiMoV2ForCausalLM": "mimo",
    "MiniCPM3ForCausalLM": "minicpm",
@@ -223,6 +229,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "ModernBertForSequenceClassification": "bert",
    "ModernBertModel": "bert",
    "NanbeigeForCausalLM": "nanbeige",
+    "Nemotron3_5TranscribeForConditionalGeneration": "nemotron",
    "NemotronForCausalLM": "nemotron",
    "NemotronHForCausalLM": "nemotron",
    "NemotronHPuzzleForCausalLM": "nemotron",
@@ -328,6 +335,7 @@ TEXT_MODEL_MAP: dict[str, str] = {

 MMPROJ_MODEL_MAP: dict[str, str] = {
    "AudioFlamingo3ForConditionalGeneration": "ultravox",
+    "ChatterboxModel": "chatterbox",
    "ClefModel": "clef",
    "CogVLMForCausalLM": "cogvlm",
    "Cohere2VisionForConditionalGeneration": "command_r",
@@ -342,6 +350,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "DotsOCRForCausalLM": "dotsocr",
    "EmbeddingGemma2Model": "gemma",
    "Exaone4_5_ForConditionalGeneration": "exaone",
+    "FishOmniModel": "drama3",
    "Gemma3ForConditionalGeneration": "gemma",
    "Gemma3nForConditionalGeneration": "gemma",
    "Gemma4ForConditionalGeneration": "gemma",
@@ -380,6 +389,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "MiniCPMOcrForConditionalGeneration": "minicpm",
    "Mistral3ForConditionalGeneration": "llava",
    "Mistral3Model": "llava",
+    "Nemotron3_5TranscribeForConditionalGeneration": "nemotron",
    "NemotronH_Omni_Reasoning_V3": "nemotron",
    "NemotronH_Nano_VL_V2": "nemotron",
    "MuseGlimmerForConditionalGeneration": "muse_glimmer",
@@ -167,6 +167,8 @@ class ModelBase:
        self._fusable_qkv_weight_layers: set[int] = set()
        self._fusable_qkv_bias_layers: set[int] = set()
        self.hparams = ModelBase.load_hparams(self.dir_model, self.is_mistral_format) if hparams is None else hparams
+        # source tensor name to filtered tensor name mapping.
+        self._source_tensor_names: dict[str, str | None] = {}
        self.model_tensors = self.index_tensors(remote_hf_model_id=remote_hf_model_id)
        self.metadata_override = metadata_override
        self.model_name = model_name
@@ -225,21 +227,25 @@ class ModelBase:
        tensors: dict[str, Callable[[], Tensor]] = {}

        if remote_hf_model_id is not None:
-            is_safetensors = True
-
            logger.info(f"Using remote model with HuggingFace id: {remote_hf_model_id}")
            remote_tensors = gguf.utility.SafetensorRemote.get_list_tensors_hf_model(remote_hf_model_id)
            for name, remote_tensor in remote_tensors.items():
                data_gen = lambda r=remote_tensor: LazyTorchTensor.from_remote_tensor(r)  # noqa: E731
-                if titem := self.filter_tensors((name, data_gen)):
+                titem = self.filter_tensors((name, data_gen))
+                self._source_tensor_names[name] = titem[0] if titem else None
+                if titem:
                    tname, tgen = titem
                    tensors[tname] = tgen

            return tensors

-        prefix = "model" if not self.is_mistral_format else "consolidated"
+        return self.index_local_tensors("model" if not self.is_mistral_format else "consolidated")
+
+    # index the local "<prefix>*.safetensors" files (or pytorch_model*.bin), through filter_tensors
+    def index_local_tensors(self, prefix: str) -> dict[str, Callable[[], Tensor]]:
+        tensors: dict[str, Callable[[], Tensor]] = {}
        part_names: list[str] = ModelBase.get_model_part_names(self.dir_model, prefix, ".safetensors")
-        is_safetensors: bool = len(part_names) > 0 or (not self.is_mistral_format and (self.dir_model / "model.safetensors.index.json").is_file())
+        is_safetensors: bool = len(part_names) > 0 or (not self.is_mistral_format and (self.dir_model / f"{prefix}.safetensors.index.json").is_file())
        if not is_safetensors:
            part_names = ModelBase.get_model_part_names(self.dir_model, "pytorch_model", ".bin")

@@ -247,7 +253,7 @@ class ModelBase:
        tensor_names_from_parts: set[str] = set()

        if not self.is_mistral_format:
-            index_name = "model.safetensors" if is_safetensors else "pytorch_model.bin"
+            index_name = f"{prefix}.safetensors" if is_safetensors else "pytorch_model.bin"
            index_name += ".index.json"
            index_file = self.dir_model / index_name

@@ -292,7 +298,9 @@ class ModelBase:
                            data_gen = lambda data=data_torch: LazyTorchTensor.from_eager(data)  # noqa: E731
                        else:
                            data_gen = lambda data=data_torch: data  # noqa: E731
-                    if titem := self.filter_tensors((name, data_gen)):
+                    titem = self.filter_tensors((name, data_gen))
+                    self._source_tensor_names[name] = titem[0] if titem else None
+                    if titem:
                        tname, tgen = titem
                        tensors[tname] = tgen

@@ -492,7 +500,7 @@ class ModelBase:
                        s = self.model_tensors[name]
                        self.model_tensors[weight_name] = lambda w=w, s=s: dequant_bitnet(w(), s())
                        tensors_to_remove.append(name)
-            elif quant_method == "fp8":
+            elif quant_method in ("fp8", "bsgemm-moe"):  # bsgemm-moe: drama-3, same block layout
                block_size = quant_config.get("weight_block_size")
                for name in self.model_tensors.keys():
                    if name.endswith("_scale_inv"):
@@ -1277,6 +1285,7 @@ class ModelBase:
        expert_scales: dict[tuple[int, str], list[tuple[int, float]]] = {}
        expert_input_scales: dict[tuple[int, str], list[tuple[int, float]]] = {}
        expert_shapes: dict[tuple[int, str], list[int]] = {}
+        expert_names: dict[tuple[int, str], str] = {}
        n_experts = self.find_hparam(["num_local_experts", "num_experts", "n_routed_experts"], optional=True) or 0
        consumed: list[str] = []

@@ -1325,6 +1334,7 @@ class ModelBase:
                    expert_scales[key] = []
                    expert_input_scales[key] = []
                    expert_shapes[key] = shape
+                    expert_names[key] = name
                expert_blocks[key].append((expert_id, raw.copy()))
                # Collect per-expert scale2 (scalar per expert)
                expert_scales[key].append((expert_id, float(scale2.float().sum())))
@@ -1333,13 +1343,13 @@ class ModelBase:

                # Flush when all experts for this (layer, proj) are collected
                if n_experts > 0 and len(expert_blocks[key]) >= n_experts:
-                    self._flush_nvfp4_experts(key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type)
+                    self._flush_nvfp4_experts(key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, expert_names, bid, proj_type)
            else:
                self._repack_nvfp4(name, weight, scale, scale2, input_scale)

        # Flush any remaining experts (fallback if n_experts was unknown)
        for bid, proj_type in list(expert_blocks.keys()):
-            self._flush_nvfp4_experts((bid, proj_type), expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type)
+            self._flush_nvfp4_experts((bid, proj_type), expert_blocks, expert_scales, expert_input_scales, expert_shapes, expert_names, bid, proj_type)

        # Remove consumed tensors so get_tensors/modify_tensors won't see them
        for name in consumed:
@@ -1350,7 +1360,7 @@ class ModelBase:
            if name.endswith((".k_scale", ".v_scale")):
                del self.model_tensors[name]

-    def _flush_nvfp4_experts(self, key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type):
+    def _flush_nvfp4_experts(self, key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, expert_names, bid, proj_type):
        experts = expert_blocks.pop(key)
        scales = expert_scales.pop(key)
        input_scales = expert_input_scales.pop(key)
@@ -1362,7 +1372,7 @@ class ModelBase:
        new_name = self.map_tensor_name(merged_name)
        logger.info(f"Repacked {new_name} with shape [{len(experts)}, {shape[0]}, {shape[1]}] and quantization NVFP4")
        self.gguf_writer.add_tensor(new_name, merged, raw_dtype=gguf.GGMLQuantizationType.NVFP4)
-        self._tag_prec_a4(merged_name, new_name)
+        self._tag_prec_a4(expert_names.pop(key), new_name)

        scales.sort(key=lambda x: x[0])
        self._write_scales_tensor(new_name.replace(".weight", ".scale"), [s[1] for s in scales])
@@ -1413,18 +1423,25 @@ class ModelBase:
        # Per-tensor NVFP4 precision.
        self._nvfp4_layer_algo = {}
        if quant_layers:
-            # store all possible module paths and assert if a quantized layer is not in the model
-            modules: set[str] = set()
-            for name in self.model_tensors:
-                while name := name.rpartition(".")[0]:
-                    modules.add(name)
-
-            for layer_name, entry in quant_layers.items():
-                if not isinstance(entry, dict):
+            # quantized_layers keys are module paths in the source checkpoint, not tensor names
+            # resolve each source tensor to the closest quantized_layers parent and record the algo under its filtered name for _tag_prec_a4
+            layer_algo = {k: v.get("quant_algo") for k, v in quant_layers.items() if isinstance(v, dict)}
+            matched: set[str] = set()
+            for src_name, filtered_name in self._source_tensor_names.items():
+                module = src_name
+                while module := module.rpartition(".")[0]:
+                    if module in layer_algo:
+                        break
+                if not module:
                    continue
-                if titem := self.filter_tensors((layer_name, lambda: torch.empty(0))):
-                    assert titem[0] in modules, f"quantized_layers entry {layer_name!r} is not in the model tensors"
-                    self._nvfp4_layer_algo[titem[0]] = entry.get("quant_algo")
+
+                matched.add(module)
+                # filter_tensors may drop a tensors so we need to check for that here.
+                if filtered_name is not None:
+                    self._nvfp4_layer_algo[filtered_name.removesuffix(".weight").removesuffix(".bias")] = layer_algo[module]
+
+            unmatched = layer_algo.keys() - matched
+            assert not unmatched, f"quantized_layers entries not in the model tensors: {sorted(unmatched)}"

        # NVFP4 weights are repacked and written directly to gguf_writer.
        # This must run before dequant_model so NVFP4 tensors are removed
@@ -2114,6 +2131,9 @@ class TextModel(ModelBase):
        if chkhsh == "0fe1cf6eda062318a1af7270f3331a85c539a01778ff948e24388e949c5282f4":
            # ref: https://huggingface.co/evilfreelancer/ruGPT3XL
            res = "gpt-2"
+        if chkhsh == "df1bc34a6366466bcb36a9a1b1cc0afdb08cfe21c9250602c930f159085390a9":
+            # ref: https://huggingface.co/Lernex/Metis-1.6-Logos
+            res = "gpt-2"
        if chkhsh == "9e454714343b69b99b71795c1d27a68c2a1d15dab111f4d353109f966af29da7":
            # ref: https://huggingface.co/LiquidAI/LFM2.5-8B-A1B
            res = "lfm2"
@@ -3384,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"):
+    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
@@ -0,0 +1,473 @@
+# Chatterbox (ResembleAI) turbo and multilingual v3, from the raw checkpoints of the official repos
+# the checkpoints ship no config.json, the variant is detected from the talker file and the hparams come from the reference code
+# talker: text_emb and speech_emb are fused into one [text | speech] vocab, speech token i is <|speech_i|>
+# mmproj: the s3gen sidecar, the voice and conditioning encoders, the position and speech tables, and the conds.pt voice
+
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Iterable
+
+import torch
+import torch.nn.functional as F
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import ModelBase, TextModel, MmprojModel, LazyTorchTensor, gguf, logger
+from .llama import LlamaModel
+
+# checkpoint file prefixes, the files are "<prefix>.safetensors"
+TURBO_TALKER = "t3_turbo_v1"
+MTL_TALKER   = "t3_mtl23ls_v3"
+TURBO_S3GEN  = "s3gen_meanflow"
+MTL_S3GEN    = "s3gen_v3"
+VE           = "ve"
+
+# relative speech ids shared by both variants (start/stop_speech_token in the reference)
+SPEECH_BOS = 6561
+SPEECH_EOS = 6562
+
+# sampling defaults of the reference, samplers it does not use are disabled so the common defaults do not apply
+TURBO_TOP_K = 1000
+MTL_TOP_K   = 0
+TURBO_SAMPLING = {"min_p": 0.0,  "top_p": 0.95, "temp": 0.8, "penalty_repeat": 1.2}
+MTL_SAMPLING   = {"min_p": 0.05, "top_p": 1.0,  "temp": 0.8, "penalty_repeat": 1.2}
+# the reference repetition penalty covers the whole generated history
+PENALTY_LAST_N = -1
+
+
+def _has_file(dir_model: Path, prefix: str) -> bool:
+    return (dir_model / f"{prefix}.safetensors").is_file()
+
+
+def _is_turbo(dir_model: Path) -> bool:
+    if _has_file(dir_model, TURBO_TALKER):
+        return True
+    if _has_file(dir_model, MTL_TALKER):
+        return False
+    raise FileNotFoundError(f"no chatterbox talker checkpoint in {dir_model}")
+
+
+# talker hparams of the reference (GPT2_MEDIUM_CONFIG / LLAMA_520M_CONFIG_DICT, T3Config)
+TURBO_TEXT_CONFIG = {
+    "n_layer": 24, "n_embd": 1024, "n_inner": 4096, "n_head": 16, "n_ctx": 8196, "layer_norm_epsilon": 1e-05,
+    "vocab_size": 50276, "speech_vocab_size": 6563,
+}
+MTL_TEXT_CONFIG = {
+    "num_hidden_layers": 30, "hidden_size": 1024, "intermediate_size": 4096, "num_attention_heads": 16,
+    "num_key_value_heads": 16, "head_dim": 64, "max_position_embeddings": 131072, "rms_norm_eps": 1e-05,
+    "rope_theta": 500000.0,
+    "rope_scaling": {"factor": 8.0, "high_freq_factor": 4.0, "low_freq_factor": 1.0, "original_max_position_embeddings": 8192},
+    "vocab_size": 2454, "speech_vocab_size": 8194,
+}
+
+
+def _is_chatterbox_checkpoint(dir_model: Path) -> bool:
+    return _has_file(dir_model, VE) and (_has_file(dir_model, TURBO_TALKER) or _has_file(dir_model, MTL_TALKER))
+
+
+@ModelBase.register_hparams_loader(_is_chatterbox_checkpoint)
+def _load_chatterbox_hparams(dir_model: Path) -> dict[str, Any]:
+    turbo = _is_turbo(dir_model)
+    logger.info(f"gguf: detected Chatterbox {'turbo' if turbo else 'multilingual v3'} checkpoint")
+    text_config = dict(TURBO_TEXT_CONFIG if turbo else MTL_TEXT_CONFIG)
+    # the mmproj tensors are named explicitly, they do not go through the block tensor map
+    return {"architectures": ["ChatterboxModel"], "text_config": text_config, "audio_config": {"num_hidden_layers": 0}}
+
+
+@ModelBase.register("ChatterboxModel")
+class ChatterboxTalkerModel(TextModel):
+    model_arch = gguf.MODEL_ARCH.LLAMA  # multilingual; the turbo constructor switches to GPT2
+
+    def __init__(self, dir_model: Path, *args, **kwargs):
+        self.is_turbo = _is_turbo(dir_model)
+        if self.is_turbo:
+            self.model_arch = gguf.MODEL_ARCH.GPT2
+        super().__init__(dir_model, *args, **kwargs)
+        self._text_embd: Tensor | None = None
+        self._speech_embd: Tensor | None = None
+        self._text_head: Tensor | None = None
+        self._speech_head: Tensor | None = None
+
+    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+        return self.index_local_tensors(TURBO_TALKER if self.is_turbo else MTL_TALKER)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, _ = item
+        # the reference feeds inputs_embeds, the transformer input embeddings are unused
+        if name in ("tfmr.wte.weight", "tfmr.embed_tokens.weight"):
+            return None
+        if name.startswith(("cond_enc.", "text_pos_emb.", "speech_pos_emb.")):
+            return None
+        return super().filter_tensors(item)
+
+    def set_vocab(self):
+        if self.is_turbo:
+            self._set_vocab_turbo()
+        else:
+            self._set_vocab_mtl()
+
+    def _speech_token_names(self, n_speech: int) -> list[str]:
+        return [f"<|speech_{i}|>" for i in range(n_speech)]
+
+    def _set_vocab_turbo(self):
+        # stock GPT-2 BPE from the checkpoint dir, extended with the speech tokens
+        tokens, toktypes, tokpre = self.get_vocab_base()
+        n_text = len(tokens)
+        n_speech = self.hparams["speech_vocab_size"]
+        speech = self._speech_token_names(n_speech)
+        tokens += speech
+        toktypes += [gguf.TokenType.CONTROL] * n_speech
+
+        with open(self.dir_model / "merges.txt", "r", encoding="utf-8") as f:
+            merges = [line.rstrip("\n") for line in f if line.strip() and not line.startswith("#version")]
+
+        self._write_vocab(tokens, toktypes, merges, tokpre, n_text)
+
+    def _write_vocab(self, tokens: list[str], toktypes: list[int], merges: list[str], tokpre: str, n_text: int):
+        self.gguf_writer.add_tokenizer_model("gpt2")
+        self.gguf_writer.add_tokenizer_pre(tokpre)
+        self.gguf_writer.add_token_list(tokens)
+        self.gguf_writer.add_token_types(toktypes)
+        self.gguf_writer.add_token_merges(merges)
+
+        self.gguf_writer.add_bos_token_id(n_text + SPEECH_BOS)
+        self.gguf_writer.add_eos_token_id(n_text + SPEECH_EOS)
+        self.gguf_writer.add_add_bos_token(False)
+        self.gguf_writer.add_add_eos_token(False)
+
+        # the reference samples the speech head only
+        self.gguf_writer.add_suppress_tokens(list(range(n_text)))
+
+    def _set_vocab_mtl(self):
+        # the reference BPE is char-level, it is re-encoded to byte-level with extra merges that rebuild each multi-byte char
+        # its lowercasing is folded into the embedding table, so the prompt needs no text preprocessing
+        with open(self.dir_model / "mtl_tokenizer.json", "r", encoding="utf-8") as f:
+            tok = json.load(f)
+
+        byte_map = gguf.vocab.bytes_to_unicode()
+
+        def enc(s: str) -> str:
+            return "".join(byte_map[b] for b in s.encode("utf-8"))
+
+        n_text = self.hparams["vocab_size"]
+        tokens: list[str] = [f"[unused_{i}]" for i in range(n_text)]
+        toktypes = [int(gguf.TokenType.UNUSED)] * n_text
+        char_merges: list[str] = []
+        for t, i in tok["model"]["vocab"].items():
+            if t == " ":
+                # the reference substitutes [SPACE] before encoding, the raw space token is never used
+                continue
+            tokens[i] = enc(t)
+            toktypes[i] = int(gguf.TokenType.NORMAL)
+            if len(t) == 1 and len(t.encode("utf-8")) > 1:
+                parts = [byte_map[b] for b in t.encode("utf-8")]
+                for k in range(1, len(parts)):
+                    char_merges.append("".join(parts[:k]) + " " + parts[k])
+        for entry in tok.get("added_tokens", []):
+            tokens[entry["id"]] = entry["content"]
+            toktypes[entry["id"]] = int(gguf.TokenType.CONTROL)
+
+        # raw spaces tokenize to [SPACE] directly
+        space_id = tok["model"]["vocab"]["[SPACE]"]
+        tokens[space_id] = enc(" ")
+        toktypes[space_id] = int(gguf.TokenType.NORMAL)
+
+        n_speech = self.hparams["speech_vocab_size"]
+        tokens += self._speech_token_names(n_speech)
+        toktypes += [int(gguf.TokenType.CONTROL)] * n_speech
+
+        # chars are atomic in the reference, they must form before any of its merges
+        merges = char_merges
+        for m in tok["model"].get("merges", []):
+            a, b = m if isinstance(m, list) else m.split(" ")
+            merges.append(enc(a) + " " + enc(b))
+
+        self._write_vocab(tokens, toktypes, merges, "default", n_text)
+
+    def _mtl_case_fold_pairs(self) -> list[tuple[int, int]]:
+        # (upper_id, lower_id) of the single-char vocab entries, for every script in the vocab
+        with open(self.dir_model / "mtl_tokenizer.json", "r", encoding="utf-8") as f:
+            vocab = json.load(f)["model"]["vocab"]
+        pairs: list[tuple[int, int]] = []
+        for t, i in vocab.items():
+            if len(t) != 1:
+                continue
+            low = t.lower()
+            if low != t and len(low) == 1 and low in vocab:
+                pairs.append((i, vocab[low]))
+        return pairs
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+
+        sampling = TURBO_SAMPLING if self.is_turbo else MTL_SAMPLING
+        self.gguf_writer.add_sampling_top_k(TURBO_TOP_K if self.is_turbo else MTL_TOP_K)
+        self.gguf_writer.add_sampling_min_p(sampling["min_p"])
+        self.gguf_writer.add_sampling_top_p(sampling["top_p"])
+        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:
+            return
+        # llama3 rope scaling baked into the rope_freqs factors tensor
+        rp = self.hparams["rope_scaling"]
+        dim = self.hparams["head_dim"]
+        base = self.hparams["rope_theta"]
+        freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+        factor = rp["factor"]
+        low_freq_wavelen = rp["original_max_position_embeddings"] / rp["low_freq_factor"]
+        high_freq_wavelen = rp["original_max_position_embeddings"] / rp["high_freq_factor"]
+        rope_factors = []
+        for freq in freqs:
+            wavelen = 2 * math.pi / freq
+            if wavelen < high_freq_wavelen:
+                rope_factors.append(1)
+            elif wavelen > low_freq_wavelen:
+                rope_factors.append(factor)
+            else:
+                smooth = (rp["original_max_position_embeddings"] / wavelen - rp["low_freq_factor"]) / (rp["high_freq_factor"] - rp["low_freq_factor"])
+                rope_factors.append(1 / ((1 - smooth) / factor + smooth))
+        yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))
+
+    def _maybe_emit_fused(self) -> Iterable[tuple[str, Tensor]]:
+        if self._text_embd is not None and self._speech_embd is not None:
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD),
+                   torch.cat([self._text_embd, self._speech_embd], dim=0))
+            self._text_embd = None
+            self._speech_embd = None
+        if self._text_head is not None and self._speech_head is not None:
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT),
+                   torch.cat([self._text_head, self._speech_head], dim=0))
+            self._text_head = None
+            self._speech_head = None
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # generate_extra_tensors output comes back through here with its final name
+        if name.startswith("rope_freqs"):
+            yield (name, data_torch)
+            return
+        # fused [text | speech] vocab: embeddings and output head
+        if name == "text_emb.weight":
+            if not self.is_turbo:
+                # the reference lowercases its input, the uppercase rows take the lowercase ones
+                fold = list(range(data_torch.shape[0]))
+                for upper, lower in self._mtl_case_fold_pairs():
+                    fold[upper] = lower
+                data_torch = data_torch[torch.tensor(fold)]
+            self._text_embd = data_torch
+            yield from self._maybe_emit_fused()
+            return
+        if name == "speech_emb.weight":
+            self._speech_embd = data_torch
+            yield from self._maybe_emit_fused()
+            return
+        if name == "text_head.weight":
+            self._text_head = data_torch
+            yield from self._maybe_emit_fused()
+            return
+        if name == "speech_head.weight":
+            self._speech_head = data_torch
+            yield from self._maybe_emit_fused()
+            return
+        if name == "speech_head.bias":
+            # the text head has no bias, its zone is suppressed at sampling anyway
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT, suffix=".bias"), F.pad(data_torch, (self.hparams["vocab_size"], 0)))
+            return
+
+        assert name.startswith("tfmr.")
+        name = name[len("tfmr."):]
+
+        if self.is_turbo:
+            # HF GPT-2 layout: Conv1D style weights are stored transposed
+            if name.endswith((".c_attn.weight", ".c_proj.weight", ".c_fc.weight")):
+                data_torch = data_torch.transpose(1, 0)
+            yield (self.map_tensor_name(name), data_torch)
+            return
+
+        if name.endswith("q_proj.weight"):
+            data_torch = LlamaModel.permute(data_torch, self.hparams["num_attention_heads"], None)
+        if name.endswith("k_proj.weight"):
+            data_torch = LlamaModel.permute(data_torch, self.hparams["num_attention_heads"], self.hparams["num_key_value_heads"])
+        yield (self.map_tensor_name("model." + name), data_torch)
+
+
+@ModelBase.register("ChatterboxModel")
+class ChatterboxMmprojModel(MmprojModel):
+    has_vision_encoder = False
+    has_audio_encoder = True
+
+    def __init__(self, dir_model: Path, *args, **kwargs):
+        self.is_turbo = _is_turbo(dir_model)
+        super().__init__(dir_model, *args, **kwargs)
+        self._wnorm_g: dict[str, Tensor] = {}
+        self._wnorm_v: dict[str, Tensor] = {}
+
+    def get_audio_config(self) -> dict[str, Any] | None:
+        return self.global_config.get("audio_config")
+
+    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+        tensors: dict[str, Callable[[], Tensor]] = {}
+        for prefix in (TURBO_S3GEN if self.is_turbo else MTL_S3GEN, VE, TURBO_TALKER if self.is_turbo else MTL_TALKER):
+            tensors.update(self.index_local_tensors(prefix))
+        return tensors
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        # the s3gen, ve and talker checkpoints have disjoint top-level names
+        name, gen = item
+        new_name: str | None = None
+        # batchnorm bookkeeping (over the gguf name length cap) and the tokenizer dsp buffers are unused
+        if name.endswith("num_batches_tracked") or name in ("tokenizer.window", "tokenizer._mel_filters"):
+            return None
+        for src, dst in (
+            ("flow.encoder.",           "a.gen.fenc."),
+            ("flow.decoder.estimator.", "a.gen.est."),
+            ("mel2wav.",                "a.gen.hift."),
+            ("speaker_encoder.",        "a.spk."),
+            ("tokenizer.",              "a.s3tok."),
+            ("flow.",                   "a.gen.flow."),  # input_embedding, encoder_proj, spk_embed_affine_layer
+            ("lstm.",                   "a.ve.lstm."),
+            ("proj.",                   "a.ve.proj."),
+            ("cond_enc.",               "a.cenc."),
+        ):
+            if name.startswith(src):
+                new_name = dst + name[len(src):]
+                break
+        if name == "text_pos_emb.emb.weight":
+            new_name = "a.gen.t3.text_pos_emb"
+        elif name == "speech_pos_emb.emb.weight":
+            new_name = "a.gen.t3.speech_pos_emb"
+        elif name == "speech_emb.weight":
+            new_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.A_GEN_CODE_OUT_EMBD] + ".weight"
+        if new_name is None:
+            return None
+        return super().filter_tensors((new_name, gen))
+
+    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
+        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)
+        self.gguf_writer.add_audio_num_mel_bins(80)
+        self.gguf_writer.add_audio_block_count(0)
+        self.gguf_writer.add_audio_embedding_length(192)
+        self.gguf_writer.add_audio_head_count(1)
+        self.gguf_writer.add_audio_feed_forward_length(192)
+        self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
+
+        # audio generator (s3gen, chatterbox projector); the flow encoder shape
+        self.gguf_writer.add_clip_has_gen_audio_encoder(True)
+        self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.CHATTERBOX)
+        self.gguf_writer.add_gen_audio_projection_dim(80)
+        self.gguf_writer.add_gen_audio_embedding_length(512)
+        self.gguf_writer.add_gen_audio_feed_forward_length(2048)
+        self.gguf_writer.add_gen_audio_block_count(6)
+        self.gguf_writer.add_gen_audio_head_count(8)
+        self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5)
+
+    def _fuse_weight_norm(self, name: str, data_torch: Tensor) -> tuple[str, Tensor] | None:
+        # torch weight_norm parametrization: weight = g * v / |v| over dims 1..n
+        base = name.split(".parametrizations.weight.original")[0]
+        if name.endswith("original0"):
+            self._wnorm_g[base] = data_torch
+        else:
+            self._wnorm_v[base] = data_torch
+        if base in self._wnorm_g and base in self._wnorm_v:
+            g = self._wnorm_g.pop(base)
+            v = self._wnorm_v.pop(base)
+            norm = v.float().norm(dim=tuple(range(1, v.dim())), keepdim=True)
+            return (base + ".weight", g.float() * v.float() / norm)
+        return None
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if ".parametrizations.weight.original" in name:
+            fused = self._fuse_weight_norm(name, data_torch)
+            if fused is not None:
+                yield fused
+            return
+        yield (name, data_torch)
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        del name, bid
+        # tensors read on the host stay F32
+        if new_name.startswith(("a.ve.", "a.cenc.", "voice.", "a.gen.t3.", "a.gen.hift.m_source.")):
+            return gguf.GGMLQuantizationType.F32
+        # the conv ops have no BF16 kernels
+        if n_dims >= 2 and new_name.endswith((".weight", ".weight_v")):
+            return gguf.GGMLQuantizationType.F16
+        return gguf.GGMLQuantizationType.F32
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        conds = torch.load(self.dir_model / "conds.pt", map_location="cpu", weights_only=True)
+        t3c = conds["t3"] if isinstance(conds, dict) else conds.t3
+        genc = conds["gen"] if isinstance(conds, dict) else conds.gen
+        t3c = vars(t3c) if not isinstance(t3c, dict) else t3c
+        genc = vars(genc) if not isinstance(genc, dict) else genc
+
+        def talker_tensor(name: str) -> Tensor:
+            return LazyTorchTensor.to_eager(self.model_tensors[name]())
+
+        # the conds.pt voice: voice.default holds the talker conditioning rows, voice.default.ref_* the flow decoder reference
+        # 1D tensors are written as F32, the speech ids stay exact
+        yield ("voice.default.ref_tokens", genc["prompt_token"][0].float())
+        yield ("voice.default.ref_feat", genc["prompt_feat"][0].float())
+        yield ("voice.default.ref_spk", genc["embedding"][0].float())
+
+        spkr_w = talker_tensor("a.cenc.spkr_enc.weight").float()
+        spkr_b = talker_tensor("a.cenc.spkr_enc.bias").float()
+        spkr_row = spkr_w @ t3c["speaker_emb"][0].float() + spkr_b
+        speech_emb = talker_tensor(self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_OUT_EMBD)).float()
+        cond_tokens = t3c["cond_prompt_speech_tokens"][0]
+
+        if self.is_turbo:
+            # turbo talker conditioning: [spkr, speech rows of the cond tokens]
+            yield ("voice.default", torch.cat([spkr_row[None], speech_emb[cond_tokens]], dim=0))
+            return
+
+        # multilingual talker conditioning [spkr, perceiver x32, emotion], from the reference perceiver over the cond speech tokens
+        with torch.no_grad():
+            speech_pos = talker_tensor("a.gen.t3.speech_pos_emb").float()
+            pse = speech_emb[cond_tokens] + speech_pos[: cond_tokens.shape[0]]
+
+            ln_w = talker_tensor("a.cenc.perceiver.attn.norm.weight").float()
+            ln_b = talker_tensor("a.cenc.perceiver.attn.norm.bias").float()
+            wq = talker_tensor("a.cenc.perceiver.attn.to_q.weight").float()
+            bq = talker_tensor("a.cenc.perceiver.attn.to_q.bias").float()
+            wk = talker_tensor("a.cenc.perceiver.attn.to_k.weight").float()
+            bk = talker_tensor("a.cenc.perceiver.attn.to_k.bias").float()
+            wv = talker_tensor("a.cenc.perceiver.attn.to_v.weight").float()
+            bv = talker_tensor("a.cenc.perceiver.attn.to_v.bias").float()
+            wo = talker_tensor("a.cenc.perceiver.attn.proj_out.weight").float()
+            bo = talker_tensor("a.cenc.perceiver.attn.proj_out.bias").float()
+            query = talker_tensor("a.cenc.perceiver.pre_attention_query")[0].float()
+
+            n_head = 4
+            n_e = query.shape[1]
+
+            def attn_block(x1: Tensor, x2: Tensor) -> Tensor:
+                nx1 = F.layer_norm(x1, (n_e,), ln_w, ln_b)
+                nx2 = F.layer_norm(x2, (n_e,), ln_w, ln_b)
+                q = (nx1 @ wq.T + bq).view(-1, n_head, n_e // n_head).transpose(0, 1)
+                k = (nx2 @ wk.T + bk).view(-1, n_head, n_e // n_head).transpose(0, 1)
+                v = (nx2 @ wv.T + bv).view(-1, n_head, n_e // n_head).transpose(0, 1)
+                ctx = F.scaled_dot_product_attention(q, k, v)
+                ctx = ctx.transpose(0, 1).reshape(-1, n_e)
+                return ctx @ wo.T + bo + x1
+
+            pre = attn_block(query, pse)
+            p32 = attn_block(pre, pre)
+
+            emo = talker_tensor("a.cenc.emotion_adv_fc.weight").float()
+            emo_row = emo[:, 0] * t3c["emotion_adv"].reshape(-1)[0]
+            yield ("voice.default", torch.cat([spkr_row[None], p32, emo_row[None]], dim=0))
@@ -0,0 +1,251 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Iterable, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import ModelBase, MmprojModel, gguf
+from .qwen import Qwen3_5MoeTextModel
+
+# Fish Audio Drama 3 (FishOmniModel)
+# text model: Qwen3.5-MoE backbone, samples the main token of each frame
+# mmproj: codebook tower and Fish DAC decoder (codec.pth)
+
+_TOWER_SHARED = "audio_decoder_mtp_shared.layers."
+_TOWER_RESIDUAL = "audio_decoder_mtp.layers."
+
+
+@ModelBase.register("FishOmniModel")
+@ModelBase.example("fishaudio/drama-3")
+class Drama3TextModel(Qwen3_5MoeTextModel):
+    model_arch = gguf.MODEL_ARCH.QWEN35MOE
+    no_mtp = True  # the checkpoint has no MTP head
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        if item[0].startswith("audio_decoder"):
+            return None
+        return super().filter_tensors(item)
+
+
+@ModelBase.register("FishOmniModel")
+@ModelBase.example("fishaudio/drama-3")
+class Drama3MmprojModel(MmprojModel):
+    has_vision_encoder = False
+    has_audio_encoder = False
+
+    _LAYER_MAP = {
+        "enorm.weight":                                 (gguf.MODEL_TENSOR.A_GEN_CODE_ENORM, True),
+        "hnorm.weight":                                 (gguf.MODEL_TENSOR.A_GEN_CODE_HNORM, True),
+        "eh_proj.weight":                               (gguf.MODEL_TENSOR.A_GEN_CODE_EH_PROJ, False),
+        "semantic_embedding.weight":                    (gguf.MODEL_TENSOR.A_GEN_CODE_BLK_EMBD, False),
+        "codebook_embedding.weight":                    (gguf.MODEL_TENSOR.A_GEN_CODE_BLK_EMBD, False),
+        "final_layernorm.weight":                       (gguf.MODEL_TENSOR.A_GEN_CODE_BLK_OUT_NORM, True),
+        "transformer_layer.attention_norm.weight":      (gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_NORM, True),
+        "transformer_layer.ffn_norm.weight":            (gguf.MODEL_TENSOR.A_GEN_CODE_FFN_NORM, True),
+        "transformer_layer.attention.q_norm.weight":    (gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_Q_NORM, True),
+        "transformer_layer.attention.k_norm.weight":    (gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_K_NORM, True),
+        "transformer_layer.attention.wo.weight":        (gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_OUT, False),
+        "transformer_layer.feed_forward.w1.weight":     (gguf.MODEL_TENSOR.A_GEN_CODE_FFN_GATE, False),
+        "transformer_layer.feed_forward.w3.weight":     (gguf.MODEL_TENSOR.A_GEN_CODE_FFN_UP, False),
+        "transformer_layer.feed_forward.w2.weight":     (gguf.MODEL_TENSOR.A_GEN_CODE_FFN_DOWN, False),
+    }
+
+    _heads: dict[int, Tensor]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._heads = {}
+
+    def get_audio_config(self) -> dict[str, Any] | None:
+        cfg = dict(self.global_config["audio_decoder_config"])
+        cfg["num_hidden_layers"] = cfg["shared_enrichment_layers"] + cfg["num_codebooks"] - 1
+        return cfg
+
+    def set_gguf_parameters(self):
+        cfg = self.hparams
+        self.gguf_writer.add_file_type(self.ftype)
+        self.gguf_writer.add_clip_has_gen_audio_encoder(True)
+        self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.DRAMA3_GEN)
+        self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text)
+        self.gguf_writer.add_gen_audio_embedding_length(cfg["dim"])
+        self.gguf_writer.add_gen_audio_feed_forward_length(cfg["intermediate_size"])
+        self.gguf_writer.add_gen_audio_block_count(cfg["num_hidden_layers"])
+        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
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name = item[0]
+        if name.startswith("audio_decoder") or name == "model.language_model.embed_tokens.weight":
+            return item
+        return None
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        # conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path
+        if new_name.endswith(".weight") and (
+            new_name in ("a.gen.wav.dac.entry.weight", "a.gen.wav.dac.post_conv.weight")
+            or (".up.blk." in new_name and new_name.endswith(".dwconv.weight"))
+            or (".dac.blk." in new_name and (new_name.endswith(".conv1.weight") or new_name.endswith(".conv2.weight")))
+        ):
+            return gguf.GGMLQuantizationType.F16
+        # keep ConvTranspose1d kernels and folded codebook tables in F32
+        if new_name.endswith(".conv.weight") and (".up.blk." in new_name or ".dac.blk." in new_name):
+            return gguf.GGMLQuantizationType.F32
+        if new_name.startswith("a.gen.wav.quant."):
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        T = gguf.MODEL_TENSOR
+        cfg = self.hparams
+
+        if name.startswith("a.gen.wav."):
+            yield (name, data_torch)
+            return
+
+        if name == "model.language_model.embed_tokens.weight":
+            # semantic band of the backbone embeddings
+            begin = self.global_config["semantic_start_token_id"]
+            yield (self.format_tensor_name(T.A_GEN_CODE_OUT_EMBD), data_torch[begin:begin + cfg["vocab_size"]])
+            return
+
+        if name == "audio_decoder.codebook_embeddings.weight":
+            n_cb = cfg["num_codebooks"]
+            yield (self.format_tensor_name(T.A_GEN_CODE_EMBD), data_torch.reshape(n_cb, -1, data_torch.shape[-1]))
+            return
+
+        if name.startswith("audio_decoder_mtp.heads."):
+            n_heads = cfg["num_codebooks"] - 1
+            self._heads[int(name.split(".")[2])] = data_torch
+            if len(self._heads) < n_heads:
+                return
+            yield (self.format_tensor_name(T.A_GEN_CODE_HEAD), torch.stack([self._heads.pop(i) for i in range(n_heads)]))
+            return
+
+        for prefix, offset in ((_TOWER_SHARED, 0), (_TOWER_RESIDUAL, cfg["shared_enrichment_layers"])):
+            if not name.startswith(prefix):
+                continue
+            idx, key = name[len(prefix):].split(".", 1)
+            il = int(idx) + offset
+            if key == "transformer_layer.attention.wqkv.weight":
+                # rows are [query | gate | key | value]
+                n_q = cfg["n_head"] * cfg["head_dim"]
+                n_kv = cfg["n_local_heads"] * cfg["head_dim"]
+                q, gate, k, v = data_torch.split([n_q, n_q, n_kv, n_kv], dim=0)
+                yield (self.format_tensor_name(T.A_GEN_CODE_ATTN_Q, il), q)
+                yield (self.format_tensor_name(T.A_GEN_CODE_ATTN_GATE, il), gate)
+                yield (self.format_tensor_name(T.A_GEN_CODE_ATTN_K, il), k)
+                yield (self.format_tensor_name(T.A_GEN_CODE_ATTN_V, il), v)
+                return
+            tensor, zero_centered = self._LAYER_MAP[key]
+            if zero_centered:
+                data_torch = data_torch.float() + 1
+            yield (self.format_tensor_name(tensor, il), data_torch)
+            return
+
+        raise ValueError(f"unexpected tensor {name}")
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # only the DAC decoder, the encoder is for voice cloning (not supported yet)
+        sd = torch.load(self.dir_model / "codec.pth", map_location="cpu", mmap=True, weights_only=True)
+        sd = sd.get("state_dict", sd)
+        T = gguf.MODEL_TENSOR
+
+        def get(name: str) -> Tensor:
+            return sd[name].float()
+
+        def conv(prefix: str) -> Tensor:
+            # fold weight norm
+            if prefix + ".conv.weight" in sd:
+                return get(prefix + ".conv.weight")
+            g = get(prefix + ".conv.parametrizations.weight.original0")
+            v = get(prefix + ".conv.parametrizations.weight.original1")
+            return g * v / v.norm(dim=(1, 2), keepdim=True)
+
+        def quant_table(prefix: str) -> Tensor:
+            # fold out_proj into the codebook, so decoding is a row lookup
+            g = get(prefix + ".out_proj.weight_g")
+            v = get(prefix + ".out_proj.weight_v")
+            w = (g * v / v.norm(dim=(1, 2), keepdim=True)).squeeze(-1)
+            return get(prefix + ".codebook.weight") @ w.T + get(prefix + ".out_proj.bias")
+
+        def snake(alpha: Tensor) -> tuple[Tensor, Tensor]:
+            alpha = alpha.flatten()
+            return alpha, 1.0 / (alpha + 1e-9)
+
+        q = "quantizer."
+        yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_FIRST_CB), quant_table(q + "semantic_quantizer.quantizers.0"))
+        n_rest = sum(1 for k in sd if k.startswith(q + "quantizer.quantizers.") and k.endswith(".codebook.weight"))
+        yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_REST_CB),
+               torch.stack([quant_table(f"{q}quantizer.quantizers.{i}") for i in range(n_rest)]))
+
+        # post_module uses interleaved RoPE, permute q/k rows so NEOX RoPE gives the same result
+        tfm = q + "post_module."
+        n_head, d_head = 16, 64
+        order = torch.cat([torch.arange(0, d_head, 2), torch.arange(1, d_head, 2)])
+
+        def neox_rows(w: Tensor) -> Tensor:
+            return w.reshape(n_head, d_head, -1)[:, order].reshape(w.shape)
+
+        n_tfm = sum(1 for k in sd if k.startswith(tfm + "layers.") and k.endswith(".ffn_norm.weight"))
+        for il in range(n_tfm):
+            p = f"{tfm}layers.{il}."
+            wq, wk, wv = get(p + "attention.wqkv.weight").chunk(3, dim=0)
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_ATTN_Q, il), neox_rows(wq))
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_ATTN_K, il), neox_rows(wk))
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_ATTN_V, il), wv)
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_ATTN_OUT, il), get(p + "attention.wo.weight"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_ATTN_NORM, il), get(p + "attention_norm.weight"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_ATTN_SCALE, il), get(p + "attention_layer_scale.gamma"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_FFN_NORM, il), get(p + "ffn_norm.weight"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_FFN_GATE, il), get(p + "feed_forward.w1.weight"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_FFN_UP, il), get(p + "feed_forward.w3.weight"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_FFN_DOWN, il), get(p + "feed_forward.w2.weight"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_TFM_FFN_SCALE, il), get(p + "ffn_layer_scale.gamma"))
+        yield (self.format_tensor_name(T.A_GEN_WAV_TFM_OUTPUT_NORM), get(tfm + "norm.weight"))
+
+        # upsample: 2x (ConvTranspose1d + ConvNeXt)
+        for il in range(2):
+            p = f"{q}upsample.{il}."
+            yield (self.format_tensor_name(T.A_GEN_WAV_UP_CONV, il, suffix=".weight"), conv(p + "0"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_UP_CONV, il, suffix=".bias"), get(p + "0.conv.bias"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_UP_DWCONV, il, suffix=".weight"), conv(p + "1.dwconv"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_UP_DWCONV, il, suffix=".bias"), get(p + "1.dwconv.conv.bias"))
+            for key, tensor in (("norm", T.A_GEN_WAV_UP_NORM), ("pwconv1", T.A_GEN_WAV_UP_PW1), ("pwconv2", T.A_GEN_WAV_UP_PW2)):
+                yield (self.format_tensor_name(tensor, il, suffix=".weight"), get(f"{p}1.{key}.weight"))
+                yield (self.format_tensor_name(tensor, il, suffix=".bias"), get(f"{p}1.{key}.bias"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_UP_GAMMA, il, suffix=""), get(p + "1.gamma"))
+
+        # DAC: entry conv, 4x (Snake, ConvTranspose1d, 3 residual units), Snake, post conv
+        d = "decoder.model."
+        yield (self.format_tensor_name(T.A_GEN_WAV_DAC_ENTRY, suffix=".weight"), conv(d + "0"))
+        yield (self.format_tensor_name(T.A_GEN_WAV_DAC_ENTRY, suffix=".bias"), get(d + "0.conv.bias"))
+        for bid in range(4):
+            p = f"{d}{bid + 1}.block."
+            a, b = snake(get(p + "0.alpha"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_SNAKE, bid, suffix=".alpha"), a)
+            yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_SNAKE, bid, suffix=".beta"), b)
+            yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_CONV, bid, suffix=".weight"), conv(p + "1"))
+            yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_CONV, bid, suffix=".bias"), get(p + "1.conv.bias"))
+            for xid in range(3):
+                r = f"{p}{xid + 2}.block."
+                for act, tensor in (("0", T.A_GEN_WAV_DAC_RES_ACT1), ("2", T.A_GEN_WAV_DAC_RES_ACT2)):
+                    a, b = snake(get(r + act + ".alpha"))
+                    name = gguf.TENSOR_NAMES[tensor].format(bid=bid, xid=xid)
+                    yield (name + ".alpha", a)
+                    yield (name + ".beta", b)
+                for cv, tensor in (("1", T.A_GEN_WAV_DAC_RES_CONV1), ("3", T.A_GEN_WAV_DAC_RES_CONV2)):
+                    name = gguf.TENSOR_NAMES[tensor].format(bid=bid, xid=xid)
+                    yield (name + ".weight", conv(r + cv))
+                    yield (name + ".bias", get(r + cv + ".conv.bias"))
+        a, b = snake(get(d + "5.alpha"))
+        yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_SNAKE, suffix=".alpha"), a)
+        yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_SNAKE, suffix=".beta"), b)
+        yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_CONV, suffix=".weight"), conv(d + "6"))
+        yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_CONV, suffix=".bias"), get(d + "6.conv.bias"))
@@ -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"]))
@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+import math
+from typing import Callable, Iterable, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import LazyTorchTensor, ModelBase, TextModel, gguf
+
+
+@ModelBase.register("Metis16ForCausalLM")
+class Metis16Model(TextModel):
+    model_arch = gguf.MODEL_ARCH.METIS16
+
+    # read by generate_extra_tensors() or not used at inference: dropped in modify_tensors()
+    _folded = {
+        "pass_embeddings", "lm_head.weight", "depth_budget.multiplier",
+        "depth_memory.pass_embeddings", "depth_memory.anchor_embeddings",
+        "depth_memory.state_write.bias", "depth_memory.metadata_write.bias",
+    }
+    _folded_layer = {
+        "mixer_connection.controller.weight", "mixer_connection.controller.bias", "mixer_connection.mix_logits",
+        "mixer_connection.read_logits", "mixer_connection.write_logits",
+        "moe_connection.controller.weight", "moe_connection.controller.bias", "moe_connection.mix_logits",
+        "moe_connection.read_logits", "moe_connection.write_logits",
+        "moe.k_router.weight", "moe.k_router.bias", "moe.k_budget.multiplier",
+    }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        hp = self.hparams
+        self.n_pass = hp["num_passes"]
+        self.attn_layers = list(hp["attention_layer_indices"])
+        self._generated: set[str] = set()
+        # one cache slot per pass and layer, while the tensors exist once per physical layer
+        self.block_count = self.n_pass * hp["num_hidden_layers"]
+
+    def set_vocab(self):
+        self._set_vocab_gpt2()
+        # turns end on the first eos id, so llama.cpp must stop on it too
+        eos = self.hparams.get("eos_token_id")
+        if isinstance(eos, list) and eos:
+            self.gguf_writer.add_eot_token_id(eos[0])
+
+    def set_gguf_parameters(self):
+        hp = self.hparams
+        super().set_gguf_parameters()
+
+        n_layer = hp["num_hidden_layers"]
+        self.gguf_writer.add_head_count_kv([hp["num_key_value_heads"] if il % n_layer in self.attn_layers else 0 for il in range(self.block_count)])
+
+        self.gguf_writer.add_expert_feed_forward_length(hp["expert_intermediate_size"])
+        self.gguf_writer.add_expert_shared_feed_forward_length(hp["shared_expert_intermediate_size"])
+        self.gguf_writer.add_expert_shared_count(1)
+        self.gguf_writer.add_moe_latent_size(hp["latent_dim"])
+
+        d_inner = hp["mamba_expand"] * hp["hidden_size"]
+        self.gguf_writer.add_ssm_conv_kernel(hp["mamba_d_conv"])
+        self.gguf_writer.add_ssm_inner_size(d_inner)
+        self.gguf_writer.add_ssm_state_size(hp["mamba_d_state"])
+        self.gguf_writer.add_ssm_time_step_rank(d_inner // hp["mamba_head_dim"])
+        self.gguf_writer.add_ssm_group_count(hp["mamba_n_groups"])
+        self.gguf_writer.add_ssm_norm_epsilon(hp["mamba_norm_eps"])
+
+        self.gguf_writer.add_hyper_connection_count(hp["n_streams"])
+        self.gguf_writer.add_metis_pass_count(self.n_pass)
+        self.gguf_writer.add_metis_memory_length(hp["memory_dim"])
+        self.gguf_writer.add_metis_route_length(hp["route_feature_dim"])
+        self.gguf_writer.add_metis_continuation_logit_limit(hp["continuation_logit_limit"])
+        self.gguf_writer.add_attention_pass_lora_rank(hp["attention_pass_lora_rank"])
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        # the checkpoint leaves these names suffix-less; add the suffix of the GGUF tensor so the mapper keeps it
+        if name.endswith((".mixer.impl.mixer.dt_bias", ".moe.selection_bias")):
+            name += ".bias"
+        elif (name.endswith((".moe.expert_embeddings", ".mixer.pass_gates"))
+              or name in ("depth_memory.stream_gate", "stream_embeddings")):
+            name += ".weight"
+        elif name == "depth_memory.stream_gate_bias":
+            name = "depth_memory.stream_gate.bias"
+        return super().filter_tensors((name, gen))
+
+    def _load(self, name: str) -> Tensor:
+        return LazyTorchTensor.to_eager(self.model_tensors[name]())
+
+    def _emit(self, name: str, data: Tensor) -> tuple[str, Tensor]:
+        self._generated.add(name)
+        return name, data
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        hp = self.hparams
+        r = hp["route_feature_dim"]
+        n_stream = hp["n_streams"]
+        T = gguf.MODEL_TENSOR
+        passes = self._load("pass_embeddings")[: self.n_pass]
+
+        # the mixes depend only on the weights and the pass, so Sinkhorn runs once here
+        # the controller runs in the checkpoint dtype, as the reference does
+        for il in range(hp["num_hidden_layers"]):
+            for src, (mix, read, write) in (
+                ("mixer_connection", (T.HC_ATTN_MIX, T.HC_ATTN_READ, T.HC_ATTN_WRITE)),
+                ("moe_connection", (T.HC_FFN_MIX, T.HC_FFN_READ, T.HC_FFN_WRITE)),
+            ):
+                prefix = f"layers.{il}.{src}."
+                ctrl_w, ctrl_b = self._load(prefix + "controller.weight"), self._load(prefix + "controller.bias")
+                mixes, reads, writes = [], [], []
+                for p in range(self.n_pass):
+                    delta = torch.nn.functional.linear(passes[p].to(ctrl_w.dtype), ctrl_w, ctrl_b).float()
+                    m_delta, r_delta, w_delta = torch.split(delta, [n_stream * n_stream, n_stream, n_stream])
+                    log_m = self._load(prefix + "mix_logits").float() + m_delta.view(n_stream, n_stream)
+                    for _ in range(hp["mhc_sinkhorn_iterations"]):
+                        log_m = log_m - torch.logsumexp(log_m, dim=-1, keepdim=True)
+                        log_m = log_m - torch.logsumexp(log_m, dim=-2, keepdim=True)
+                    mixes.append(log_m.exp())
+                    reads.append(torch.softmax(self._load(prefix + "read_logits").float() + r_delta, dim=-1))
+                    writes.append(torch.softmax(self._load(prefix + "write_logits").float() + w_delta, dim=-1))
+                act = ctrl_w.dtype
+                yield self._emit(self.format_tensor_name(mix, il), torch.stack(mixes).to(act).float())
+                yield self._emit(self.format_tensor_name(read, il), torch.stack(reads).to(act).float())
+                yield self._emit(self.format_tensor_name(write, il), torch.stack(writes).to(act).float())
+
+        # one bias per memory entry: write biases, pass and anchor embeddings, and the constant routed width
+        meta_w = self._load("depth_memory.metadata_write.weight").float()
+        mean_k = hp["num_experts_per_tok"] / hp["max_routed_k"]
+        base = (self._load("depth_memory.state_write.bias").float() + self._load("depth_memory.metadata_write.bias").float()
+                + mean_k * meta_w[:, r])
+        pass_e = self._load("depth_memory.pass_embeddings").float()
+        anchor_e = self._load("depth_memory.anchor_embeddings").float()
+        n_anchor = len(self.attn_layers) + 1
+        entry_bias = torch.stack([base + pass_e[p] + anchor_e[a] for p in range(self.n_pass) for a in range(n_anchor)])
+        yield self._emit(self.format_tensor_name(T.DEPTH_MEM_ENTRY_BIAS), entry_bias)
+
+    def _split_features(self, data_torch: Tensor, template: str, parts: tuple[str, ...], d: int, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # the graph reads the concatenated features one block at a time; each block is its own tensor
+        for i, part in enumerate(parts):
+            piece = data_torch[:, i * d:(i + 1) * d] if i < 3 else data_torch[:, 3 * d:]
+            yield from super().modify_tensors(piece, template.format(part), bid)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name in self._generated:
+            yield name, data_torch
+            return
+        if name in self._folded:
+            return
+        hp = self.hparams
+        d = hp["hidden_size"]
+
+        if bid is None:
+            if name == "stream_embeddings.weight":
+                # the graph keeps the streams innermost: [n_stream, n_embd, n_tokens]
+                yield from super().modify_tensors(data_torch.T.contiguous(), name, bid)
+                return
+            if name == "depth_memory.metadata_write.weight":
+                # metadata input becomes [summary, entropy, confidence, continuation]; entropy arrives unnormalized
+                r = hp["route_feature_dim"]
+                meta_w = data_torch.float()
+                entropy_col = meta_w[:, r + 1:r + 2] / math.log(hp["n_routed_experts"])
+                data_torch = torch.cat((meta_w[:, :r], entropy_col, meta_w[:, r + 2:]), dim=1)
+            elif name == "depth_memory.stream_gate.weight":
+                data_torch = data_torch.float() / math.sqrt(d)
+            elif name == "depth_memory.route_projection.weight":
+                yield from self._split_features(data_torch, "depth_memory.route_{}.weight", ("state", "retr", "diff", "hist"), d, bid)
+                return
+            elif name == "continuation.hidden.weight":
+                yield from self._split_features(data_torch, "continuation.hidden_{}.weight", ("state", "retr", "diff", "route"), d, bid)
+                return
+            yield from super().modify_tensors(data_torch, name, bid)
+            return
+
+        local = name.removeprefix(f"layers.{bid}.")
+        if local in self._folded_layer:
+            return
+
+        if local == "mixer.impl.mixer.A_log":
+            data_torch = -torch.exp(data_torch).unsqueeze(-1)
+        elif local == "mixer.impl.mixer.D":
+            data_torch = data_torch.unsqueeze(-1)
+        elif local == "mixer.impl.mixer.conv1d.weight":
+            data_torch = data_torch.squeeze(1)
+        elif local == "mixer.impl.mixer.norm.weight":
+            data_torch = data_torch.reshape(hp["mamba_n_groups"], -1)
+        elif local == "mixer.pass_gates.weight":
+            # the reference applies tanh(gate) in the checkpoint dtype
+            data_torch = torch.tanh(data_torch[: self.n_pass]).to(torch.bfloat16).float()
+        elif local == "moe.shared_expert.gate_up.weight":
+            gate, up = data_torch.chunk(2, dim=0)
+            yield from super().modify_tensors(gate, name.replace("gate_up", "gate"), bid)
+            yield from super().modify_tensors(up, name.replace("gate_up", "up"), bid)
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+    def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
+        # FP32 in the checkpoint, tiny and read as rows, or computed here at full precision: keep FP32
+        if (new_name.startswith(("cont_", "depth_mem_entry_bias", "depth_mem_stream_gate"))
+                or new_name.endswith(("ffn_exp_embd.weight", "ffn_gate_inp.weight", "hc_attn_mix.weight", "hc_attn_read.weight",
+                                      "hc_attn_write.weight", "hc_ffn_mix.weight", "hc_ffn_read.weight", "hc_ffn_write.weight"))):
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
@@ -428,6 +428,9 @@ class NemotronHModel(GraniteHybridModel):
    is_moe: bool = False
    supports_mtp_export = True
    _experts: list[dict[str, Tensor]] | None = None
+    # tokens past the embedding rows (left out of the vocab), and special tokens that are printed as text
+    _vocab_skip: tuple[str, ...] = ()
+    _user_defined_tokens: re.Pattern[str] | None = None

    _SSM_LAYER_TYPES = {"mamba", "linear_attention"}
    _ATTN_LAYER_TYPES = {"attention", "full_attention"}
@@ -616,7 +619,11 @@ class NemotronHModel(GraniteHybridModel):
        vocab_size = -(vocab_size // -pad_vocab) * pad_vocab
        self.hparams["vocab_size"] = vocab_size

-        assert max(tokenizer.vocab.values()) < vocab_size  # ty: ignore[unresolved-attribute]
+        # draft-model converters (DFlash, LiLiCorr) call this with their own instance, which has neither hook
+        vocab_skip = getattr(self, "_vocab_skip", ())
+        user_defined_tokens = getattr(self, "_user_defined_tokens", None)
+
+        assert max(id_ for tok, id_ in tokenizer.vocab.items() if tok not in vocab_skip) < vocab_size  # ty: ignore[unresolved-attribute]

        tokpre = self.get_vocab_base_pre(tokenizer)

@@ -638,7 +645,8 @@ class NemotronHModel(GraniteHybridModel):
                        if previous_token != token:
                            logger.info(f"{repr(previous_token)} is encoded and decoded back to {repr(token)} using AutoTokenizer")

-                    if added_tokens_decoder[i].special or self.does_token_look_special(token):
+                    printed = user_defined_tokens is not None and user_defined_tokens.fullmatch(token)
+                    if (added_tokens_decoder[i].special or self.does_token_look_special(token)) and not printed:
                        toktypes.append(gguf.TokenType.CONTROL)
                    else:
                        token = token.replace(b"\xe2\x96\x81".decode("utf-8"), " ")  # pre-normalize user-defined spaces
@@ -834,3 +842,140 @@ class NemotronHPuzzleModel(NemotronHModel):
        if item[0].startswith("mtp."):
            return None
        return super().filter_tensors(item)
+
+
+# NeMo speechlm2 SALM checkpoint: a Nemotron-H decoder under llm.*, the speech encoder under perception.*
+@ModelBase.register("Nemotron3_5TranscribeForConditionalGeneration")
+@ModelBase.example("nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B")
+class NemotronTranscribeModel(NemotronHModel):
+    model_arch = gguf.MODEL_ARCH.NEMOTRON_H
+    # the audio placeholder (id vocab_size) has no embedding row; speaker, event and style tags print with the transcript
+    _vocab_skip = ("<|audio|>",)
+    _user_defined_tokens = re.compile(r"<spk:\d+>|<ev:[a-z_]+>|<sty:(?:stt|end):[a-z_]+>")
+    _renames = {
+        "llm.model.embed_tokens.weight": "backbone.embeddings.weight",
+        "llm.model.norm.weight":         "backbone.norm_f.weight",
+        "llm.lm_head.weight":            "lm_head.weight",
+    }
+    # grouped experts: [n_expert, n_in, n_out] -> [n_expert, n_out, n_in]
+    _grouped_experts = {"gate_and_up_projs": "up_proj", "down_projs": "down_proj"}
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if name.startswith(("perception.", "ctc_timestamp.")) or "_extra_state" in name:
+            return None
+        if not name.startswith("llm."):
+            # the FP8 and NVFP4 releases store the decoder as backbone.*
+            raise ValueError(f"unexpected tensor {name}: only the BF16 checkpoint layout (llm.*) is supported")
+        if name in cls._renames:
+            name = cls._renames[name]
+        elif name.startswith("llm.model.layers."):
+            name = "backbone.layers." + name[len("llm.model.layers."):]
+        elif name.startswith("llm.mtp.layers."):
+            name = "mtp.layers." + name[len("llm.mtp.layers."):]
+        else:
+            raise ValueError(f"unexpected decoder tensor {name}")
+        return super().filter_tensors((name.replace("._fp32_params.", "."), gen))
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        leaf = name.rsplit(".", 1)[-1]
+        if ".mixer.experts." in name and leaf in self._grouped_experts:
+            if name.startswith("mtp."):
+                assert self._mtp_bid is not None, f"{name}: MTP head is not exported"
+                bid = self._mtp_bid
+            assert bid is not None
+            data_torch = data_torch.permute(0, 2, 1).contiguous()
+            name = f"model.layers.{bid}.mlp.experts.{self._grouped_experts[leaf]}.weight"
+            yield from ModelBase.modify_tensors(self, data_torch, name, bid)
+            return
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
+@ModelBase.register("Nemotron3_5TranscribeForConditionalGeneration")
+class NemotronTranscribeAudioModel(MmprojModel):
+    # NeMo ParallelExpertEncoder: ASR and speaker transformers, speaker activity fused into the ASR states
+    has_vision_encoder = False
+    has_audio_encoder = True
+
+    def get_audio_config(self) -> dict[str, Any] | None:
+        asr = self.global_config["pe_encoder_config"]["asr_encoder_cfg"]
+        return {
+            "hidden_size":         asr["d_model"],
+            "intermediate_size":   int(float(asr.get("ff_expansion", 4)) * asr["d_model"]),
+            "num_hidden_layers":   asr["n_layers"],
+            "num_attention_heads": asr["n_heads"],
+        }
+
+    @staticmethod
+    def _rope_dims(cfg: dict[str, Any]) -> int:
+        assert cfg.get("self_attention_model") == "rope" and cfg.get("subsampling", "feature_stacking") == "feature_stacking"
+        assert cfg.get("attn_mode", "full") == "full" and not cfg.get("xscaling", False) and cfg.get("pre_block_norm", True)
+        assert float(cfg.get("rope_base", 10000.0)) == 10000.0
+        n_rot = int(cfg["d_model"] // cfg["n_heads"] * float(cfg.get("rotary_fraction", 1.0)))
+        assert n_rot >= 2 and n_rot % 2 == 0
+        return n_rot
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        pe = self.global_config["pe_encoder_config"]
+        asr = pe["asr_encoder_cfg"]
+        diar = pe["diarization_model_cfg"]
+        spk = diar["encoder"]
+        assert pe.get("asr_encoder_type") == "transformer" and pe.get("speaker_feature_mode", "thresholded") == "thresholded"
+        assert pe.get("asr_normalize_type") == "per_feature" and pe.get("diar_normalize_type") in (None, "NA")
+        assert asr["subsampling_factor"] == spk["subsampling_factor"] == 8 and asr["feat_in"] == spk["feat_in"]
+        assert self._rope_dims(spk) == spk["d_model"] // spk["n_heads"] and not spk.get("qkv_bias", False)
+        assert diar.get("high_resolution", False)
+        # constants in clip: the ParallelExpertEncoder online settings and the Sortformer streaming config
+        assert (pe.get("online_inference_length", 500), pe.get("chunk_left_context", 50), pe.get("chunk_right_context", 50),
+                pe.get("diar_spkcache_len", 188), pe.get("diar_fifo_len", 40), pe.get("diar_spkcache_update_period", 300),
+                pe.get("speaker_activity_threshold", 0.5)) == (500, 50, 50, 188, 40, 300, 0.5), \
+            "streaming settings differ from the constants in tools/mtmd/models/nemotron-transcribe.cpp"
+        sm = diar["sortformer_modules"]
+        assert (sm["spkcache_sil_frames_per_spk"], sm["pred_score_threshold"], sm["scores_boost_latest"], sm["strong_boost_rate"],
+                sm["weak_boost_rate"], sm["min_pos_scores_rate"], sm["use_learnable_sil_emb"]) == (1, 0.25, 0.05, 0.75, 1.5, 0.5, True), \
+            "speaker cache settings differ from the constants in tools/mtmd/models/nemotron-transcribe.cpp"
+
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_TRANSCRIBE)
+        self.gguf_writer.add_audio_num_mel_bins(asr["feat_in"])
+        self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
+        self.gguf_writer.add_audio_subsampling_factor(asr["subsampling_factor"])
+        self.gguf_writer.add_audio_rope_dims(self._rope_dims(asr))
+        self.gguf_writer.add_audio_spk_block_count(spk["n_layers"])
+        self.gguf_writer.add_audio_spk_head_count(spk["n_heads"])
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # ParallelExpertEncoder._build_sinusoid_position_encoding
+        pe = self.global_config["pe_encoder_config"]
+        n_spk = int(pe["diarization_model_cfg"]["sortformer_modules"]["num_spks"])
+        d = int(pe["asr_encoder_cfg"]["d_model"])
+        scale = float(pe.get("spk_kernel_scale", 1.0))
+        pos = torch.arange(n_spk, dtype=torch.float32)[:, None]
+        div = torch.exp(torch.arange(0, d, 2, dtype=torch.float32) * -(np.log(10000.0) / d))
+        kernel = torch.zeros(n_spk, d, dtype=torch.float32)
+        kernel[:, 0::2] = torch.sin(pos * div)
+        kernel[:, 1::2] = torch.cos(pos * div)
+        # stored as [d, n_spk]: the graph multiplies it with the speaker activity directly
+        yield self.format_tensor_name(gguf.MODEL_TENSOR.A_FUSION_SPK_KERNEL, suffix=""), (kernel * scale).T.contiguous()
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.A_SPK_HEAD_UPSAMPLE, None):
+            return gguf.GGMLQuantizationType.F16  # ggml_conv_1d kernel
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if not name.startswith("perception.") or "_extra_state" in name:
+            return None
+        # training-only head; the speaker model's unused front end
+        if name.startswith(("perception.encoder.diarization_model.sortformer_modules.hidden_to_spks.",
+                            "perception.encoder.diarization_model.preprocessor.")):
+            return None
+        return name, gen
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name == "perception.preprocessor.featurizer.fb":
+            data_torch = data_torch.squeeze(0)
+        yield from super().modify_tensors(data_torch, name, bid)
@@ -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)
@@ -198,6 +198,7 @@ pre_computed_hashes = [
    # jina-v2-de variants
    {"name": "jina-v2-de", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/aari1995/German_Semantic_V3", "chkhsh": "b3d1dd861f1d4c5c0d2569ce36baf3f90fe8a102db3de50dd71ff860d91be3df"},
    {"name": "gpt-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/evilfreelancer/ruGPT3XL", "chkhsh": "0fe1cf6eda062318a1af7270f3331a85c539a01778ff948e24388e949c5282f4"},
+    {"name": "gpt-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/Lernex/Metis-1.6-Logos", "chkhsh": "df1bc34a6366466bcb36a9a1b1cc0afdb08cfe21c9250602c930f159085390a9"},
    # lfm2 variants
    {"name": "lfm2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/LiquidAI/LFM2.5-8B-A1B", "chkhsh": "9e454714343b69b99b71795c1d27a68c2a1d15dab111f4d353109f966af29da7"},
    # hrm-text (DFM Mimir) is SPM-style BPE: normalizer maps ' ' -> '▁', merges
@@ -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.
@@ -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.
@@ -62,7 +62,7 @@ Legend:
 |                       GELU_QUICK | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         GET_ROWS | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ❌ | ❌ |
 |                    GET_ROWS_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ |
-|                       GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                       GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -71,7 +71,7 @@ Legend:
 |                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                LIGHTNING_INDEXER | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ | ❌ |
+|                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ❌ |
 |                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
 |                 MUL_MAT_HADAMARD | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | 🟡 |
@@ -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"
@@ -12278,13 +12278,23 @@
 "WebGPU: WebGPU","SUM","type=f32,ne=[11,5,6,3],permute=[0,2,1,3]","support","1","yes","WebGPU"
 "WebGPU: WebGPU","SUM","type=f32,ne=[11,5,6,3],permute=[0,3,2,1]","support","1","yes","WebGPU"
 "WebGPU: WebGPU","SUM","type=f32,ne=[11,5,6,3],permute=[0,1,3,2]","support","1","yes","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[10,5,4,3]","support","0","no","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[33,1,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[33,256,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[32769,1,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[32,1,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[32,256,1,1]","support","0","no","WebGPU"
-"WebGPU: WebGPU","MEAN","type=f32,ne=[32768,1,1,1]","support","0","no","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[10,5,4,3],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[33,1,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[33,256,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[32769,1,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[32,1,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[32,256,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[32768,1,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[11,5,6,3],permute=1,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[11,5,6,3],permute=0,slice=1","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[11,5,6,3],permute=1,slice=1","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[1,7,3,2],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[1023,5,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[1024,5,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[1025,5,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[17,65537,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[512,1,1,1],permute=0,slice=0","support","1","yes","WebGPU"
+"WebGPU: WebGPU","MEAN","type=f32,ne=[281,1536,1,1],permute=0,slice=0","support","1","yes","WebGPU"
 "WebGPU: WebGPU","SUM","type=f32,ne=[33,1,1,1]","support","1","yes","WebGPU"
 "WebGPU: WebGPU","SUM","type=f32,ne=[33,1024,1,1]","support","1","yes","WebGPU"
 "WebGPU: WebGPU","SUM","type=f32,ne=[33,256,1,1]","support","1","yes","WebGPU"
@@ -12299,8 +12309,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"
@@ -4,7 +4,7 @@ project("ggml" C CXX ASM)

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

@@ -211,6 +211,18 @@ extern "C" {

    // Context management and operations for faster communication between backends, used for tensor parallelism (meta backend)
    typedef void * (*ggml_backend_comm_init_t)(ggml_backend_t * backends, size_t n_backends);
+
+    // Builds the communicator for the case where each rank is a separate process, so a process
+    // owns one backend of the communicator rather than all of them.
+    //
+    // The caller provides bootstrap_callback, which moves bytes between the ranks before the
+    // communicator exists. Each rank passes size bytes in send_buf, and recv_buf receives
+    // world_size * size bytes in rank order. The backend decides what to exchange and how many
+    // times to call it, and every rank must make the same calls in the same order.
+    // comm_init_rank blocks until all world_size ranks have joined.
+    typedef bool   (*ggml_backend_comm_bootstrap_callback)(void * user_data, const void * send_buf, void * recv_buf, size_t size);
+    typedef void * (*ggml_backend_comm_init_rank_t)(ggml_backend_t backend, int rank, int world_size, ggml_backend_comm_bootstrap_callback bootstrap_callback, void * bootstrap_callback_data);
+
    typedef void   (*ggml_backend_comm_free_t)(void * comm_ctx);
    typedef bool   (*ggml_backend_comm_allreduce_tensor_t)(void * comm_ctx, struct ggml_tensor ** tensors);

@@ -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,
@@ -2034,10 +2035,6 @@ extern "C" {
            float                 beta_slow),
        "use ggml_rope_ext_inplace instead");

-    // compute correction dims for YaRN RoPE scaling
-    GGML_API void ggml_rope_yarn_corr_dims(
-        int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]);
-
    // rotary position embedding backward, i.e compute dx from dy
    // a - dy
    GGML_API struct ggml_tensor * ggml_rope_ext_back(
@@ -2080,6 +2077,11 @@ extern "C" {
            struct ggml_tensor  * a,
            int                   n_offs);

+    // control rounding of YaRN correction bounds (default: true)
+    GGML_API struct ggml_tensor * ggml_rope_set_truncate(
+            struct ggml_tensor * a,
+            bool                 truncate);
+
    // im2col
    // converts data into a format that effectively results in a convolution when combined with matrix multiplication
    GGML_API struct ggml_tensor * ggml_im2col(
@@ -489,6 +489,32 @@ static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct
    return it->second[index];
 }

+// bytes spanned by all dims except skip_dim
+static size_t ggml_backend_meta_inner_extent(const struct ggml_tensor * tensor, int skip_dim) {
+    size_t extent = 0;
+    for (int i = 0; i < GGML_MAX_DIMS; i++) {
+        if (i != skip_dim && tensor->ne[i] > 1) {
+            extent += (size_t) (tensor->ne[i] - 1) * tensor->nb[i];
+        }
+    }
+    return extent;
+}
+
+// the view dim that holds the split axis of src, or -1
+// all other dims of src and view must stay inside one slab of that axis
+static int ggml_backend_meta_view_slab_dim(const struct ggml_tensor * view, const struct ggml_tensor * src, int axis) {
+    const size_t slab = src->nb[axis];
+    if (ggml_backend_meta_inner_extent(src, axis) >= slab) {
+        return -1;
+    }
+    for (int dim = 0; dim < GGML_MAX_DIMS; dim++) {
+        if (view->ne[dim] == src->ne[axis] && view->nb[dim] == slab) {
+            return ggml_backend_meta_inner_extent(view, dim) < slab ? dim : -1;
+        }
+    }
+    return -1;
+}
+
 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync);

 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
@@ -700,9 +726,39 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
            }
            GGML_ABORT("fatal error");
        }
+        // permuted view: split the view dim that spans the src split axis, other dims must be inside one element of it or outside the span
+        if (ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0]) && axis >= 0 && axis < GGML_MAX_DIMS-1) {
+            const size_t nb_axis = tensor->src[0]->nb[axis];
+            const size_t nb_next = tensor->src[0]->nb[axis+1];
+            for (int dim = 0; dim < GGML_MAX_DIMS; dim++) {
+                if (tensor->ne[dim] == 1 || tensor->nb[dim] < nb_axis || tensor->nb[dim]*tensor->ne[dim] != nb_next) {
+                    continue;
+                }
+                bool others_ok = true;
+                for (int i = 0; i < GGML_MAX_DIMS; i++) {
+                    if (i == dim || tensor->ne[i] == 1) {
+                        continue;
+                    }
+                    others_ok = others_ok && (tensor->nb[i]*tensor->ne[i] <= tensor->nb[dim] || tensor->nb[i] % nb_next == 0);
+                }
+                if (others_ok) {
+                    return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1};
+                }
+            }
+        }
        if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED || src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
            return src_ss[0];
        }
+        // a view that only regroups dims inside one slab of the split axis keeps the split,
+        // the per-device strides are only rescaled above nb[axis] (e.g. the KV streams of a permuted q)
+        if (axis >= 0 && axis < GGML_MAX_DIMS && tensor->view_offs == tensor->src[0]->view_offs) {
+            const int dim = ggml_backend_meta_view_slab_dim(tensor, tensor->src[0], axis);
+            if (dim >= 0) {
+                ggml_backend_meta_split_state ret = src_ss[0];
+                ret.axis = ggml_backend_meta_split_axis(dim);
+                return ret;
+            }
+        }
        GGML_ABORT("view of permuted tensor not implemented");
        //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1};
    };
@@ -1000,8 +1056,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
                split_state = handle_pad(src_ss);
            } break;
            case GGML_OP_PAD_REFLECT_1D:
-            case GGML_OP_ROLL:
-            case GGML_OP_ARANGE:
+            case GGML_OP_ROLL: {
+                split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+            } break;
+            case GGML_OP_ARANGE: {
+                // no sources, every device computes the full range
+                split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+            } break;
            case GGML_OP_TIMESTEP_EMBEDDING: {
                split_state = handle_generic(src_ss, /*scalar_only =*/ true);
            } break;
@@ -1747,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);
@@ -2934,7 +2934,7 @@ void ggml_cann_rope(ggml_backend_cann_context & ctx, ggml_tensor * dst) {
    const float theta_scale = powf(freq_base, -2.0f / n_dims);

    float corr_dims[2];
-    ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims);
+    ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, true, corr_dims);

    bool       is_neox    = mode & GGML_ROPE_TYPE_NEOX;
    const bool is_imrope  = mode == GGML_ROPE_TYPE_IMROPE;  // qwen3vl apply interleaved mrope
@@ -3289,7 +3289,7 @@ void ggml_cann_rope_cache_preload(ggml_backend_cann_context & ctx, ggml_tensor *
    const float theta_scale = powf(freq_base, -2.0f / n_dims);

    float corr_dims[2];
-    ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims);
+    ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, true, corr_dims);

    bool       is_neox    = mode & GGML_ROPE_TYPE_NEOX;
    const bool is_imrope  = mode == GGML_ROPE_TYPE_IMROPE;
@@ -2547,6 +2547,9 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
            }
        case GGML_OP_ROPE:
            {
+                if (!ggml_rope_get_truncate(op)) {
+                    return false;
+                }
                if (((const int32_t *) op->op_params)[15] != 0) {
                    return false; // FIXME: support ggml_rope_set_offset
                }
@@ -29,7 +29,7 @@ void ggml_quantize_mat_q8_0_4x8(const float * GGML_RESTRICT x, void * GGML_RESTR
    assert(k % QK8_0 == 0);
    const int nb = k / QK8_0;

-#if defined(__riscv_v_intrinsic)
+#if defined(__riscv_v)
    block_q8_0x4 * GGML_RESTRICT y = (block_q8_0x4 *) vy;
    const size_t vl_calc = __riscv_vsetvl_e32m8(QK8_0);
    const size_t vl_save = __riscv_vsetvl_e64m2(4);
@@ -345,7 +345,7 @@ static inline int32x4_t ggml_nvfp4_dot8(const int8x8_t q4_lo, const int8x8_t q8_
 #include <immintrin.h>
 #endif

-#ifdef __riscv_v_intrinsic
+#if defined(__riscv_v)
 #include <riscv_vector.h>
 #endif

@@ -746,7 +746,7 @@ static void ggml_init_arm_arch_features(void) {}
 #endif
 #endif // __ARM_ARCH

-#if defined(__riscv) && defined(__riscv_v_intrinsic)
+#if defined(__riscv) && defined(__riscv_v)
 #include <riscv_vector.h>
 static void ggml_init_riscv_arch_features(void) {
    ggml_riscv_arch_features.rvv_vlen = __riscv_vlenb();
@@ -3556,7 +3556,7 @@ void ggml_cpu_fp16_to_fp32(const ggml_fp16_t * x, float * y, int64_t n) {
        _mm_storeu_ps(y + i, y_vec);
    }

-#elif defined(__riscv_v_intrinsic) && defined(__riscv_zvfhmin)
+#elif defined(__riscv_v) && defined(__riscv_zvfhmin)
    // calculate step size
    const int epr = __riscv_vsetvlmax_e16m2();
    const int step = epr * 2;
@@ -3626,7 +3626,7 @@ void ggml_cpu_bf16_to_fp32(const ggml_bf16_t * x, float * y, int64_t n) {
                                        (const __m128i *)(x + i))),
                                16)));
    }
-#elif defined(__riscv_v_intrinsic) && defined(__riscv_zvfbfmin)
+#elif defined(__riscv_v) && defined(__riscv_zvfbfmin)
    // calculate step size
    const int epr = __riscv_vsetvlmax_e16m2();
    const int step = epr * 2;
@@ -3746,7 +3746,7 @@ int ggml_cpu_has_arm_fma(void) {
 }

 int ggml_cpu_has_riscv_v(void) {
-#if defined(__riscv_v_intrinsic)
+#if defined(__riscv_v)
    return 1;
 #else
    return 0;
@@ -3754,7 +3754,7 @@ int ggml_cpu_has_riscv_v(void) {
 }

 int ggml_cpu_get_rvv_vlen(void) {
-#if defined(__riscv) && defined(__riscv_v_intrinsic)
+#if defined(__riscv) && defined(__riscv_v)
    return ggml_riscv_arch_features.rvv_vlen;
 #else
    return 0;
@@ -69,7 +69,7 @@
 #define VECTOR_REGISTERS 16
 #endif

-#if defined(__riscv_v_intrinsic)
+#if defined(__riscv_v)
 #define LMUL 4
 #endif

@@ -180,7 +180,7 @@ inline float32x4_t madd(float32x4_t a, float32x4_t b, float32x4_t c) {
 }
 #endif

-#if defined(__riscv_v_intrinsic)
+#if defined(__riscv_v)
 template <> inline vfloat32m1_t madd(vfloat32m1_t a, vfloat32m1_t b, vfloat32m1_t c) {
    return __riscv_vfmacc_vv_f32m1(c, a, b, __riscv_vsetvlmax_e32m1());
 }
@@ -277,7 +277,7 @@ inline float hsum(__m512 x) {
 }
 #endif // __AVX512F__

-#if defined(__riscv_v_intrinsic)
+#if defined(__riscv_v)
 inline float hsum(vfloat32m1_t x) {
    return __riscv_vfmv_f_s_f32m1_f32(
        __riscv_vfredusum_vs_f32m1_f32m1(x, __riscv_vfmv_v_f_f32m1(0, 1), __riscv_vsetvlmax_e32m1()));
@@ -458,7 +458,7 @@ template <> inline __m512bh load_partial(const ggml_bf16_t *p, int n) {
 #endif
 #endif

-#if defined(__riscv_v_intrinsic)
+#if defined(__riscv_v)
 template <> inline vfloat32m1_t load(const float *p) {
    return __riscv_vle32_v_f32m1(p, __riscv_vsetvlmax_e32m1());
 }
@@ -503,7 +503,7 @@ template <> inline vbfloat16m4_t load(const ggml_bf16_t *p) {
 }
 #endif

-#if defined(__riscv_v_intrinsic)
+#if defined(__riscv_v)
 template <typename T> T set_zero();

 template <> inline vfloat32m1_t set_zero() {
@@ -520,7 +520,7 @@ template <> inline vfloat32m8_t set_zero() {
 }
 #endif

-#if defined(__riscv_v_intrinsic)
+#if defined(__riscv_v)
 template <typename T> size_t vlmax() {
    if constexpr (std::is_same_v<T, vfloat32m1_t>) { return  __riscv_vsetvlmax_e32m1(); }
    else if constexpr (std::is_same_v<T, vfloat32m2_t>) { return  __riscv_vsetvlmax_e32m2(); }
@@ -732,7 +732,7 @@ class tinyBLAS {
    const int64_t ldc;
 };

-#if defined(__riscv_v_intrinsic)
+#if defined(__riscv_v)
 template <typename D, typename V, typename TA, typename TB, typename TC>
 class tinyBLAS_RVV {
  public:
@@ -3958,7 +3958,7 @@ 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(__riscv_v_intrinsic)
+#elif defined(__riscv_v)
    #if LMUL == 1
        tinyBLAS_RVV<vfloat32m1_t, vfloat32m1_t, float, float, float> tb{ params,
            k, (const float *)A, lda,
@@ -6310,8 +6310,8 @@ static void ggml_compute_forward_rope_flt(
    //const int n_past     = ((int32_t *) dst->op_params)[0];
    const int n_dims     = ((int32_t *) dst->op_params)[1];
    const int mode       = ((int32_t *) dst->op_params)[2];
-    //const int n_ctx      = ((int32_t *) dst->op_params)[3];
    const int n_ctx_orig = ((int32_t *) dst->op_params)[4];
+    const bool truncate = ggml_rope_get_truncate(dst);

    memcpy(&freq_base,   (int32_t *) dst->op_params +  5, sizeof(float));
    memcpy(&freq_scale,  (int32_t *) dst->op_params +  6, sizeof(float));
@@ -6356,7 +6356,7 @@ static void ggml_compute_forward_rope_flt(
    const float theta_scale = powf(freq_base, -2.0f/n_dims);

    float corr_dims[2];
-    ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims);
+    ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, truncate, corr_dims);

    const bool is_imrope = mode == GGML_ROPE_TYPE_IMROPE; // qwen3vl apply interleaved mrope
    const bool mrope_used = mode & GGML_ROPE_TYPE_MROPE;  // ggml_rope_multi, note: also true for vision (24 & 8 == true) and for imrope
@@ -10114,7 +10114,7 @@ static void ggml_compute_forward_ssm_scan_f32(
                        }

                        sumf = GGML_F32xt_REDUCE_ONE(sum);
-    #elif defined(__riscv_v_intrinsic)
+    #elif defined(__riscv_v)
                        // todo: RVV implementation
                        const int np = 0;
    #else
@@ -11778,7 +11778,7 @@ static void ggml_compute_forward_rwkv_wkv7_f32(
    int64_t h_stride_2d = head_size * head_size;

    #if defined(GGML_SIMD)
-        #if defined(__ARM_FEATURE_SVE) || defined(__riscv_v_intrinsic)
+        #if defined(__ARM_FEATURE_SVE) || defined(__riscv_v)
            // scalar Route to scalar implementation       //TODO: Write SVE code and RVV code
            for (int64_t t = 0; t < T; t++) {
                int64_t t_offset = t * t_stride;
@@ -5,7 +5,7 @@
 #include "simd-mappings.h"

 // TODO: add support for sizeless vector types
-#if defined(GGML_SIMD) && !defined(__ARM_FEATURE_SVE) && !defined(__riscv_v_intrinsic)
+#if defined(GGML_SIMD) && !defined(__ARM_FEATURE_SVE) && !defined(__riscv_v)

 // TODO: untested on avx512
 // These are in units of GGML_F32_EPR
@@ -149,7 +149,7 @@ static void simd_gemm(
        C += N;
    }
 }
-#elif defined(GGML_SIMD) && defined(__riscv_v_intrinsic)
+#elif defined(GGML_SIMD) && defined(__riscv_v)
 // RM accumulators + 1 B vector = RM + 1 <= 8  =>  RM <= 7
 // Microkernel: C[RM x vl] += A[RM x K] * B[K x N]
 template <int RM>
@@ -14,7 +14,9 @@
 #include <arm_neon.h>
 #endif

-#if defined(__riscv_v_intrinsic)
+// __riscv_v_intrinsic reports compiler API support and can be defined without
+// enabling V; use the ISA feature macro to select RVV code.
+#if defined(__riscv_v)
 #include <riscv_vector.h>
 #endif

@@ -1306,7 +1308,7 @@ static inline void __lzs_f16cx4_store(ggml_fp16_t * x, float32x4_t v_y) {
 #define GGML_BF16_FMA_HI(acc, x, y) \
    (acc) = GGML_F32x4_FMA((acc), GGML_BF16_TO_F32_HI(x), GGML_BF16_TO_F32_HI(y))

-#elif defined(__riscv_v_intrinsic)
+#elif defined(__riscv_v)

 // compatible with vlen >= 128

@@ -84,7 +84,7 @@ void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * G
        }
        // reduce sum1,sum2 to sum1
        GGML_F32_VEC_REDUCE(sumf, sum1, sum2, sum3, sum4, sum5, sum6, sum7, sum8);
-    #elif defined(__riscv_v_intrinsic)
+    #elif defined(__riscv_v)
        int vl = __riscv_vsetvlmax_e32m8();
        vfloat32m1_t vs = __riscv_vfmv_v_f_f32m1(0.0f, 1);
        vfloat32m8_t vsum;
@@ -195,7 +195,7 @@ void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t *
    sumf += (ggml_float)_mm_cvtss_f32(g);

 #undef LOAD
-#elif defined(__riscv_v_intrinsic) && defined(__riscv_zvfbfwma)
+#elif defined(__riscv_v) && defined(__riscv_zvfbfwma)
    size_t vl = __riscv_vsetvlmax_e32m4();

    // initialize accumulators to all zeroes
@@ -318,7 +318,7 @@ void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * G
        sum1_hi = svadd_f32_m(DEFAULT_PG32, sum1_hi, sum3_hi);

        sumf = ggml_sve_sum_f32x2(sum1_lo, sum1_hi);
-    #elif defined(__riscv_v_intrinsic)
+    #elif defined(__riscv_v)
        #if defined(__riscv_zvfh)
            int vl = __riscv_vsetvlmax_e32m2();
            vfloat32m1_t vs = __riscv_vfmv_v_f_f32m1(0.0f, 1);
@@ -405,7 +405,7 @@ void ggml_vec_silu_f32(const int n, float * y, const float * x) {
    for (; i + 3 < n; i += 4) {
        vec_xst(ggml_v_silu(vec_xl(0, x + i)), 0, y + i);
    }
-#elif defined(__riscv_v_intrinsic)
+#elif defined(__riscv_v)
    for (int vl; i < n; i += vl) {
        vl = __riscv_vsetvl_e32m2(n - i);
        vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl);
@@ -446,7 +446,7 @@ void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float *
    for (; i + 3 < n; i += 4) {
        vec_xst(vec_mul(ggml_v_silu(vec_xl(0, x + i)), vec_xl(0, g + i)), 0, y + i);
    }
-#elif defined(__riscv_v_intrinsic)
+#elif defined(__riscv_v)
    for (int vl; i < n; i += vl) {
        vl = __riscv_vsetvl_e32m2(n - i);
        vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl);
@@ -602,7 +602,7 @@ ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const floa
        val = vec_mul(val, val);
        sum += (ggml_float)vec_hsum_f32x4(val);
    }
-#elif defined(__riscv_v_intrinsic)
+#elif defined(__riscv_v)
    vfloat64m1_t vsum = __riscv_vfmv_v_f_f64m1(0, 1);
    for (int vl; i < n; i += vl) {
        vl = __riscv_vsetvl_e32m2(n - i);
@@ -681,7 +681,7 @@ ggml_float ggml_vec_soft_max_f32(const int n, float * y, const float * x, float
        vec_xst(v_val, 0, y + i);
        sum += (ggml_float)vec_hsum_f32x4(v_val);
    }
-#elif defined(__riscv_v_intrinsic)
+#elif defined(__riscv_v)
    vfloat64m1_t vsum = __riscv_vfmv_v_f_f64m1(0, 1);
    for (int avl; i < n; i += avl) {
        avl = __riscv_vsetvl_e32m2(n - i);
@@ -206,7 +206,7 @@ inline static void ggml_vec_dot_f16_unroll(const int n, const int xs, float * GG
        sumf[0] = ggml_sve_sum_f32x2(sum_0_lo, sum_0_hi);
        sumf[1] = ggml_sve_sum_f32x2(sum_1_lo, sum_1_hi);
        np = n;
-    #elif defined(__riscv_v_intrinsic)
+    #elif defined(__riscv_v)
        #if defined(__riscv_zvfh)
            size_t vl = __riscv_vsetvlmax_e32m4();

@@ -396,7 +396,7 @@ inline static void ggml_vec_mad_f32(const int n, float * GGML_RESTRICT y, const

            svst1_f32(pg, y + np2, ay1);
        }
-    #elif defined(__riscv_v_intrinsic)
+    #elif defined(__riscv_v)
        for (int i = 0, avl; i < n; i += avl) {
            avl = __riscv_vsetvl_e32m8(n - i);
            vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[i], avl);
@@ -513,7 +513,7 @@ inline static void ggml_vec_mad_f16(const int n, ggml_fp16_t * GGML_RESTRICT y,
        svst1_f16(pg, (__fp16 *)(y + np2), hy);
    }
    np = n;
-#elif defined(__riscv_v_intrinsic) // implies __riscv_v_intrinsic
+#elif defined(__riscv_v)
    #if defined (__riscv_zvfh)
        const ggml_fp16_t s = GGML_CPU_FP32_TO_FP16(v);
        const _Float16 scale = *(const _Float16*)(&s);
@@ -599,7 +599,7 @@ inline static void ggml_vec_mad_f32_unroll(const int n, const int xs, const int
                y[i] += x[k][i]*v[k][0];
            }
        }
-    #elif defined(__riscv_v_intrinsic)
+    #elif defined(__riscv_v)
        for (int i = 0, avl; i < n; i += avl) {
            avl = __riscv_vsetvl_e32m8(n - i);
            vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl);
@@ -660,7 +660,7 @@ inline static void ggml_vec_mad1_f32(const int n, float * y, const float * x, co
        for (int i = 0; i < n; ++i) {
            y[i] = x[i]*s + b;
        }
-    #elif defined(__riscv_v_intrinsic)
+    #elif defined(__riscv_v)
        for (int i = 0, avl; i < n; i += avl) {
            avl = __riscv_vsetvl_e32m8(n - i);
            vfloat32m8_t ax = __riscv_vle32_v_f32m8(&x[i], avl);
@@ -729,7 +729,7 @@ inline static void ggml_vec_scale_f32(const int n, float * y, const float   v) {
            ay1 = svmul_f32_m(pg, ay1, vx);
            svst1_f32(pg, y + i, ay1);
        }
-    #elif defined(__riscv_v_intrinsic)
+    #elif defined(__riscv_v)
        for (int i = 0, avl; i < n; i += avl) {
            avl = __riscv_vsetvl_e32m8(n - i);
            vfloat32m8_t ay = __riscv_vle32_v_f32m8(&y[i], avl);
@@ -793,7 +793,7 @@ inline static void ggml_vec_scale_f16(const int n, ggml_fp16_t * y, const float
        svst1_f16(pg, (__fp16 *)(y + np), out);
    }
    np = n;
-#elif defined(__riscv_v_intrinsic)
+#elif defined(__riscv_v)
    #if defined(__riscv_zvfh)
        const ggml_fp16_t s = GGML_CPU_FP32_TO_FP16(v);
        const _Float16 scale = *(const _Float16*)(&s);
@@ -1411,7 +1411,7 @@ inline static float32x4_t ggml_v_gelu_quick(float32x4_t x) {
    return vec_div(x, vec_add(one, exp_cx));
 }

-#elif defined(__riscv_v_intrinsic)
+#elif defined(__riscv_v)

 // adapted from arm limited optimized routine
 // the maximum error is 1.45358 plus 0.5 ulps
@@ -1469,7 +1469,7 @@ inline static vfloat32m2_t ggml_v_gelu_quick_m2(vfloat32m2_t x, int vl) {
    return __riscv_vfdiv_vv_f32m2(x, __riscv_vfadd_vf_f32m2(exp_cx, 1.0f, vl), vl);
 }

-#endif // __ARM_NEON / __AVX2__ / __SSE2__ / __VXE__ / __riscv_v_intrinsic
+#endif // __ARM_NEON / __AVX2__ / __SSE2__ / __VXE__ / __riscv_v

 inline static void ggml_vec_silu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
    for (int i = 0; i < n; ++i) {
@@ -0,0 +1,198 @@
+#include "batch-get.cuh"
+#include "common.cuh"
+#include "ggml-cpp.h"
+
+#include <climits>
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+struct ggml_cuda_batch_get_descriptor {
+    const unsigned char * src;
+    unsigned char * dst;
+    size_t size;
+};
+
+static __global__ void ggml_cuda_batch_get_kernel(const ggml_cuda_batch_get_descriptor * copies, size_t count) {
+    // Blocks along y cooperate on each descriptor.
+    const size_t index = blockIdx.x;
+    if (index >= count) {
+        return;
+    }
+    const size_t first = (size_t) blockIdx.y * blockDim.x + threadIdx.x;
+    const size_t stride = (size_t) gridDim.y * blockDim.x;
+    const auto copy = copies[index];
+    // Use 16-, 8-, 4-, or 2-byte copies when both pointers meet that alignment; threads copy separate chunks.
+    const uintptr_t alignment = (uintptr_t) copy.src | (uintptr_t) copy.dst;
+    size_t tail = 0;
+    if (alignment % 16 == 0) {
+        for (size_t i = first; i < copy.size / 16; i += stride) {
+            ggml_cuda_memcpy_1<16>(copy.dst + 16*i, copy.src + 16*i);
+        }
+        tail = copy.size / 16 * 16;
+    } else if (alignment % 8 == 0) {
+        for (size_t i = first; i < copy.size / 8; i += stride) {
+            ggml_cuda_memcpy_1<8>(copy.dst + 8*i, copy.src + 8*i);
+        }
+        tail = copy.size / 8 * 8;
+    } else if (alignment % 4 == 0) {
+        for (size_t i = first; i < copy.size / 4; i += stride) {
+            ggml_cuda_memcpy_1<4>(copy.dst + 4*i, copy.src + 4*i);
+        }
+        tail = copy.size / 4 * 4;
+    } else if (alignment % 2 == 0) {
+        for (size_t i = first; i < copy.size / 2; i += stride) {
+            ggml_cuda_memcpy_1<2>(copy.dst + 2*i, copy.src + 2*i);
+        }
+        tail = copy.size / 2 * 2;
+    }
+    // Copy trailing bytes, or the entire range when no wider alignment is available.
+    for (size_t i = tail + first; i < copy.size; i += stride) {
+        ggml_cuda_memcpy_1<1>(copy.dst + i, copy.src + i);
+    }
+}
+
+struct ggml_cuda_batch_get_mapping {
+    void * host = nullptr;
+    unsigned char * device = nullptr;
+
+    ggml_cuda_batch_get_mapping(void * host, size_t size) {
+        const cudaError_t err = cudaHostRegister(host, size, cudaHostRegisterMapped | cudaHostRegisterPortable);
+        if (err != cudaSuccess) {
+            (void) cudaGetLastError();
+            GGML_LOG_DEBUG("%s: failed to register %.2f MiB of pinned memory: %s\n", __func__,
+                           size / 1024.0 / 1024.0, cudaGetErrorString(err));
+            return;
+        }
+        this->host = host;
+        if (cudaHostGetDevicePointer((void **) &device, host, 0) != cudaSuccess) {
+            (void) cudaGetLastError();
+            device = nullptr;
+            return;
+        }
+    }
+
+    ggml_cuda_batch_get_mapping(const ggml_cuda_batch_get_mapping &) = delete;
+    ggml_cuda_batch_get_mapping & operator=(const ggml_cuda_batch_get_mapping &) = delete;
+
+    ~ggml_cuda_batch_get_mapping() {
+        if (host) {
+            // Release only the registration acquired by this object.
+            CUDA_CHECK(cudaHostUnregister(host));
+        }
+    }
+};
+#endif
+
+bool ggml_cuda_get_tensor_n(int device,
+        const ggml_tensor * const * tensors, void * const * data, const size_t * offsets,
+        const size_t * sizes, size_t n_tensors) {
+#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+    GGML_UNUSED(device);
+    GGML_UNUSED(tensors);
+    GGML_UNUSED(data);
+    GGML_UNUSED(offsets);
+    GGML_UNUSED(sizes);
+    GGML_UNUSED(n_tensors);
+    return false;
+#else
+    // Provisional dispatch values based on experiments.
+    constexpr size_t min_batch_reads = 128;
+    constexpr size_t max_gather_read_bytes = 128 * 1024;
+    // Collect non-empty CPU destination ranges.
+    std::vector<std::pair<uintptr_t, uintptr_t>> ranges;
+    ranges.reserve(n_tensors);
+    size_t payload_bytes = 0;
+    // Number of non-empty reads whose size is at most max_gather_read_bytes.
+    size_t n_reads_small = 0;
+    for (size_t i = 0; i < n_tensors; ++i) {
+        if (sizes[i] == 0) {
+            continue;
+        }
+        const uintptr_t address = (uintptr_t) data[i];
+        GGML_ASSERT(address <= UINTPTR_MAX - sizes[i]);
+        GGML_ASSERT(payload_bytes <= SIZE_MAX - sizes[i]);
+        payload_bytes += sizes[i];
+        n_reads_small += sizes[i] <= max_gather_read_bytes;
+        ranges.emplace_back(address, address + sizes[i]);
+    }
+    if (ranges.empty()) {
+        return true;
+    }
+    // Total number of non-empty reads.
+    const size_t n_reads = ranges.size();
+    if (n_reads > SIZE_MAX / sizeof(ggml_cuda_batch_get_descriptor) || n_reads > INT_MAX) {
+        return false;
+    }
+    if (n_reads < 2 || (n_reads < min_batch_reads && n_reads_small == n_reads)) {
+        return false;
+    }
+    std::sort(ranges.begin(), ranges.end());
+    for (size_t i = 1; i < ranges.size(); ++i) {
+        // Preserve the ordered copy path when CPU destinations overlap.
+        if (ranges[i].first < ranges[i - 1].second) {
+            return false;
+        }
+    }
+    // Find the CPU span to register, including gaps between destinations.
+    const uintptr_t begin = ranges.front().first;
+    const uintptr_t end = ranges.back().second;
+    if (end - begin - payload_bytes > payload_bytes) {
+        return false;
+    }
+    // Register and map the existing CPU destination so the kernel can write to it directly.
+    ggml_cuda_batch_get_mapping mapping((void *) begin, end - begin);
+    if (!mapping.host) {
+        return false;
+    }
+    // Reads selected for gather; zero means all reads use DMA.
+    size_t n_reads_gather = mapping.device && n_reads_small >= min_batch_reads ? n_reads_small : 0;
+    // Keep descriptor storage alive until the final stream synchronization.
+    constexpr size_t max_batch_descriptors = 128 * 1024;
+    const size_t bytes = std::min(n_reads_gather, max_batch_descriptors) * sizeof(ggml_cuda_batch_get_descriptor);
+    ggml_backend_buffer_ptr descriptors(n_reads_gather ? ggml_backend_buft_alloc_buffer(ggml_backend_cuda_buffer_type(device), bytes) : nullptr);
+    if (!descriptors) {
+        n_reads_gather = 0;
+    }
+    const auto host_type = ggml_backend_cuda_host_buffer_type();
+    ggml_backend_buffer_ptr host_descriptors(n_reads_gather ? ggml_backend_buft_alloc_buffer(host_type, bytes) : nullptr);
+    if (!host_descriptors || ggml_backend_buffer_get_type(host_descriptors.get()) != host_type) {
+        n_reads_gather = 0;
+    }
+    if (n_reads_gather) {
+        auto * copies = static_cast<ggml_cuda_batch_get_descriptor *>(ggml_backend_buffer_get_base(host_descriptors.get()));
+        auto * device_copies = static_cast<ggml_cuda_batch_get_descriptor *>(ggml_backend_buffer_get_base(descriptors.get()));
+        const int threads = 256;
+        const size_t nsm = ggml_cuda_info().devices[device].nsm;
+        size_t i = 0;
+        for (size_t first = 0; first < n_reads_gather; first += max_batch_descriptors) {
+            const size_t n_batch = std::min(n_reads_gather - first, max_batch_descriptors);
+            size_t count = 0;
+            for (; i < n_tensors && count < n_batch; ++i) {
+                if (sizes[i] == 0 || sizes[i] > max_gather_read_bytes) {
+                    continue;
+                }
+                copies[count++] = {(const unsigned char *) tensors[i]->data + offsets[i],
+                        mapping.device + ((uintptr_t) data[i] - begin), sizes[i]};
+            }
+            const size_t blocks_per_copy = std::min<size_t>(UINT16_MAX, std::max<size_t>(1, (nsm + n_batch - 1) / n_batch));
+            const dim3 blocks((unsigned int) n_batch, (unsigned int) blocks_per_copy);
+            // Stream ordering keeps each kernel ahead of the next descriptor overwrite.
+            CUDA_CHECK(cudaMemcpyAsync(device_copies, copies, n_batch * sizeof(copies[0]), cudaMemcpyHostToDevice, cudaStreamPerThread));
+            ggml_cuda_batch_get_kernel<<<blocks, threads, 0, cudaStreamPerThread>>>(device_copies, n_batch);
+            CUDA_CHECK(cudaGetLastError());
+            if (first + n_batch < n_reads_gather) {
+                // Wait before refilling the pinned host descriptor buffer.
+                CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+            }
+        }
+    }
+    // Keep the destination pinned for DMA too; overlapping destinations already declined above.
+    for (size_t i = 0; i < n_tensors; ++i) {
+        if (sizes[i] && (!n_reads_gather || sizes[i] > max_gather_read_bytes)) {
+            CUDA_CHECK(cudaMemcpyAsync(data[i], (const char *) tensors[i]->data + offsets[i], sizes[i], cudaMemcpyDeviceToHost, cudaStreamPerThread));
+        }
+    }
+    // Wait for GPU writes before releasing the CPU mapping and descriptor buffers.
+    CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+    return true;
+#endif
+}
@@ -0,0 +1,7 @@
+#pragma once
+
+#include "ggml-backend.h"
+
+bool ggml_cuda_get_tensor_n(int device,
+        const ggml_tensor * const * tensors, void * const * data, const size_t * offsets,
+        const size_t * sizes, size_t n_tensors);
@@ -1469,16 +1469,21 @@ struct ggml_backend_cuda_context {
    // when the computation is split across CPU/GPU (e.g., with --n-cpu-moe)
    std::unordered_map<const void *, std::unique_ptr<ggml_cuda_graph>> cuda_graphs;

-    int64_t last_graph_eviction_sweep = 0;
+    int64_t next_graph_eviction_time = 0;
+    int64_t last_graph_use_time = 0;

-    ggml_cuda_graph * cuda_graph(const void * first_node_ptr) {
-        const int64_t time_now = ggml_time_us();
+    ggml_cuda_graph * cuda_graph(const void * first_node_ptr, int64_t time_now = ggml_time_us()) {
+        // Give all cached graphs 10s to be reused after an idle period, without updating each graph.
+        if (!cuda_graphs.empty() && time_now - last_graph_use_time >= 10'000'000) {
+            next_graph_eviction_time = time_now + 10'000'000;
+        }
+        last_graph_use_time = time_now;

-        // sweep every 5s, evicting cuda graphs unused for >=10s
-        if (time_now - last_graph_eviction_sweep >= 5'000'000) {
-            last_graph_eviction_sweep = time_now;
+        // Sweep every 5s during active use, evicting graphs unused for at least 10s.
+        if (time_now >= next_graph_eviction_time) {
+            next_graph_eviction_time = time_now + 5'000'000;
            for (auto it = cuda_graphs.begin(); it != cuda_graphs.end(); ) {
-                if (time_now - it->second->last_used_time >= 10'000'000) {
+                if (it->first != first_node_ptr && time_now - it->second->last_used_time >= 10'000'000) {
                    it = cuda_graphs.erase(it);
                } else {
                    ++it;
@@ -1580,6 +1585,7 @@ struct ggml_cuda_mm_fusion_args_host {
    const ggml_tensor * gate_scale = nullptr;
    ggml_glu_op glu_op;
    float glu_limit = 0.0f;
+    ggml_prec prec_src1 = GGML_PREC_UNDEFINED; // src1 precision requested by the fused MUL_MATs (MMQ)
    const ggml_tensor * shared_up = nullptr;
    const ggml_tensor * shared_gate = nullptr;
    ggml_tensor * shared_dst = nullptr;
@@ -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
@@ -9,6 +9,7 @@
 #include "ggml-cuda/arange.cuh"
 #include "ggml-cuda/argmax.cuh"
 #include "ggml-cuda/argsort.cuh"
+#include "ggml-cuda/batch-get.cuh"
 #include "ggml-cuda/binbcast.cuh"
 #include "ggml-cuda/clamp.cuh"
 #include "ggml-cuda/col2im-1d.cuh"
@@ -886,6 +887,11 @@ static void ggml_backend_cuda_buffer_get_tensor_n(ggml_backend_buffer_t buffer,

    ggml_cuda_set_device(ctx->device);

+    if (!ggml_cuda_pinned_memory_disabled() &&
+            ggml_cuda_get_tensor_n(ctx->device, tensors, data, offsets, sizes, n_tensors)) {
+        return;
+    }
+
    for (size_t i = 0; i < n_tensors; i++) {
        if (sizes[i] == 0) {
            continue;
@@ -1117,6 +1123,18 @@ struct ggml_backend_cuda_comm_context {

    ggml_cuda_ar_pipeline *     ar_pipeline = nullptr;

+    // Number of ranks in the communicator when each rank is a separate process. Zero unless
+    // comm_init_rank created this context.
+    int                         world_size = 0;
+
+    // Number of ranks taking part in the reduction. For a communicator from comm_init_rank this is
+    // world_size. Otherwise it is the number of backends this process holds. The NCCL reduction
+    // uses it to choose the tensor size below which it reduces in F32 rather than converting to
+    // BF16 first.
+    size_t n_ranks() const {
+        return world_size > 0 ? (size_t) world_size : backends.size();
+    }
+
 #ifdef GGML_USE_NCCL
    std::vector<ncclComm_t>     comms;
 #endif // GGML_USE_NCCL
@@ -1144,6 +1162,7 @@ static bool ggml_backend_cuda_comm_allreduce_nccl(
    }

    const size_t n_backends = comm_ctx->backends.size();
+    const size_t n_ranks    = comm_ctx->n_ranks();

    for (size_t i = 0; i < n_backends; ++i) {
        GGML_ASSERT(tensors[i] != nullptr);
@@ -1153,7 +1172,7 @@ static bool ggml_backend_cuda_comm_allreduce_nccl(

    // For small tensors, simply reduce them as FP32.
    // The following heuristic for how "small" a tensor should be is based on RTX 4090s connected via 16x PCIe 4.0.
-    if ((n_backends <= 2 && ne < 32768) || (n_backends == 3 && ne < 131072) || (n_backends >= 4 && ne < 262144)) {
+    if ((n_ranks <= 2 && ne < 32768) || (n_ranks == 3 && ne < 131072) || (n_ranks >= 4 && ne < 262144)) {
        for (size_t i = 0; i < n_backends; ++i) {
            if ((tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
                ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context;
@@ -1384,6 +1403,90 @@ static void * ggml_backend_cuda_comm_init(ggml_backend_t * backends, size_t n_ba
    return ret;
 }

+#ifdef GGML_USE_NCCL
+// Builds the communicator with ncclCommInitRank, for the case where each rank is its own
+// process. Blocks until all world_size ranks have joined.
+static void * ggml_backend_cuda_comm_init_rank(ggml_backend_t backend, int rank, int world_size,
+                                               ggml_backend_comm_bootstrap_callback bootstrap_callback, void * bootstrap_callback_data) {
+    if (world_size < 1 || rank < 0 || rank >= world_size) {
+        GGML_LOG_ERROR("%s: invalid rank %d for world size %d\n", __func__, rank, world_size);
+        return nullptr;
+    }
+
+    // Every rank sends whether its local checks passed, and rank 0 also sends the NCCL unique id.
+    // So all ranks either join together or fail together, and no rank waits in ncclCommInitRank alone.
+    struct bootstrap_msg {
+        uint8_t      ok;
+        ncclUniqueId id;
+    };
+    bootstrap_msg local;
+    memset(&local, 0, sizeof(local));
+    local.ok = 1;
+
+    if (!ggml_backend_is_cuda(backend)) {
+        GGML_LOG_ERROR("%s: backend is not a CUDA backend\n", __func__);
+        local.ok = 0;
+    } else {
+        // NCCL requires one distinct physical GPU per rank, see ggml_backend_cuda_comm_init_nccl.
+        const ggml_cuda_device_info & info = ggml_cuda_info();
+        if (info.device_count > info.physical_device_count) {
+            GGML_LOG_ERROR("%s: CUDA virtual devices are in use, which NCCL does not support\n", __func__);
+            local.ok = 0;
+        }
+    }
+    if (local.ok && rank == 0) {
+        const ncclResult_t rc = ncclGetUniqueId(&local.id);
+        if (rc != ncclSuccess) {
+            GGML_LOG_ERROR("%s: ncclGetUniqueId failed (%s)\n", __func__, ncclGetErrorString(rc));
+            local.ok = 0;
+        }
+    }
+
+    std::vector<bootstrap_msg> all(world_size);
+    if (!bootstrap_callback(bootstrap_callback_data, &local, all.data(), sizeof(local))) {
+        GGML_LOG_ERROR("%s: bootstrap exchange failed on rank %d\n", __func__, rank);
+        return nullptr;
+    }
+    for (int i = 0; i < world_size; i++) {
+        if (!all[i].ok) {
+            GGML_LOG_ERROR("%s: rank %d failed its setup checks\n", __func__, i);
+            return nullptr;
+        }
+    }
+
+    ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
+
+    auto * ret = new ggml_backend_cuda_comm_context;
+    ret->backends.push_back(backend);
+    ret->dev_ids.push_back(cuda_ctx->device);
+    ret->world_size = world_size;
+    ret->comms.resize(1);
+    ret->try_allreduce = ggml_backend_cuda_comm_try_allreduce_nccl;
+
+    // On DGX Spark, NCCL picks 64 channels by default, which is slower and uses more memory than 8.
+    // NCCL_MAX_CTAS still overrides this.
+    ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
+    const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc;
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_DGX_SPARK) {
+        config.maxCTAs = 8;
+    }
+
+    ggml_cuda_set_device(cuda_ctx->device);
+    const ncclResult_t rc = ncclCommInitRankConfig(&ret->comms[0], world_size, all[0].id, rank, &config);
+    if (rc != ncclSuccess) {
+        GGML_LOG_ERROR("%s: ncclCommInitRankConfig failed for rank %d/%d (%s)\n",
+                       __func__, rank, world_size, ncclGetErrorString(rc));
+        ret->comms.clear();
+        delete ret;
+        return nullptr;
+    }
+
+    GGML_LOG_INFO("%s: joined NCCL communicator as rank %d/%d on device %d\n",
+                  __func__, rank, world_size, cuda_ctx->device);
+    return ret;
+}
+#endif // GGML_USE_NCCL
+
 // Top-level dispatch -- calls the function pointer chosen by comm_init.
 // Returns false to let the meta-backend's butterfly run.
 static bool ggml_backend_cuda_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tensor ** tensors) {
@@ -2048,8 +2151,50 @@ static bool ggml_cuda_match_shared_expert(const ggml_cgraph * graph, int routed_
    return true;
 }

+// The NVFP4 weight scale that scale_node (MUL) applies to the output of mm_node (MUL_MAT), or nullptr.
+// Shared by ggml_cuda_try_fuse and graph_optimize so that both accept the same patterns.
+static const ggml_tensor * ggml_cuda_get_mul_mat_scale(const ggml_tensor * scale_node, const ggml_tensor * mm_node) {
+    const bool scale_lhs_mm = scale_node->src[0] == mm_node;
+    const bool scale_rhs_mm = scale_node->src[1] == mm_node;
+    if (!scale_lhs_mm && !scale_rhs_mm) {
+        return nullptr;
+    }
+
+    const ggml_tensor * scale = scale_lhs_mm ? scale_node->src[1] : scale_node->src[0];
+    if (mm_node->src[0]->type != GGML_TYPE_NVFP4 || scale_node->type != GGML_TYPE_F32 ||
+            scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != 1 ||
+            !ggml_are_same_shape(scale_node, mm_node)) {
+        return nullptr;
+    }
+
+    return scale;
+}
+
+// The same for MUL_MAT_ID with one scale per expert: RESHAPE, REPEAT and GET_ROWS (by the ids) of the scales, then the MUL.
+static const ggml_tensor * ggml_cuda_get_mul_mat_id_scale(const ggml_tensor * reshape, const ggml_tensor * repeat, const ggml_tensor * getrows,
+        const ggml_tensor * scale_node, const ggml_tensor * mm_node) {
+    if (repeat->src[0] != reshape || getrows->src[0] != repeat || getrows->src[1] != mm_node->src[2]) {
+        return nullptr;
+    }
+    if (!((scale_node->src[0] == mm_node && scale_node->src[1] == getrows) ||
+            (scale_node->src[0] == getrows && scale_node->src[1] == mm_node))) {
+        return nullptr;
+    }
+
+    const ggml_tensor * scale = reshape->src[0];
+    if (mm_node->src[0]->type != GGML_TYPE_NVFP4 || scale_node->type != GGML_TYPE_F32 ||
+            scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != mm_node->src[0]->ne[2] ||
+            !ggml_are_same_shape(scale_node, mm_node)) {
+        return nullptr;
+    }
+
+    return scale;
+}
+
 // Buffer- and batch-size independent, so graph_optimize produces the same graph topology for every ubatch size (see #29986).
-static bool ggml_cuda_mul_mat_q_fusion_matches(const ggml_tensor * up, const ggml_tensor * gate, const ggml_tensor * glu, const int device) {
+// up_scale/gate_scale: the MUL nodes that apply the NVFP4 weight scales (per tensor or per expert) between the MUL_MATs and the GLU.
+static bool ggml_cuda_mul_mat_q_fusion_matches(const ggml_tensor * up, const ggml_tensor * gate, const ggml_tensor * glu, const int device,
+        const ggml_tensor * up_scale = nullptr, const ggml_tensor * gate_scale = nullptr) {
    const ggml_tensor * x_up = up->src[0];
    const ggml_tensor * y    = up->src[1];

@@ -2062,7 +2207,12 @@ static bool ggml_cuda_mul_mat_q_fusion_matches(const ggml_tensor * up, const ggm
        return false;
    }

-    if (!ggml_cuda_should_fuse_mul_mat(up, gate, glu)) {
+    if (!ggml_cuda_should_fuse_mul_mat(up, gate, glu, nullptr, nullptr, up_scale, gate_scale)) {
+        return false;
+    }
+
+    // The fused kernel runs both MUL_MATs at one src1 precision, see ggml_cuda_mul_mat_q.
+    if (ggml_prec_get_src(up, 1) != ggml_prec_get_src(gate, 1)) {
        return false;
    }

@@ -2075,7 +2225,8 @@ static bool ggml_cuda_mul_mat_q_fusion_matches(const ggml_tensor * up, const ggm
    return turing_mma_available(ggml_cuda_info().devices[device].cc);
 }

-static bool ggml_cuda_should_fuse_mul_mat_q(const ggml_tensor * up, const ggml_tensor * gate, const ggml_tensor * glu) {
+static bool ggml_cuda_should_fuse_mul_mat_q(const ggml_tensor * up, const ggml_tensor * gate, const ggml_tensor * glu,
+        const ggml_tensor * up_scale = nullptr, const ggml_tensor * gate_scale = nullptr) {
    // A hinted matmul can be fused, but its unfused path may be faster (e.g. FWHT for GGML_HINT_SRC0_IS_HADAMARD).
    if (ggml_get_op_params_i32(up, 1) != GGML_HINT_NONE || ggml_get_op_params_i32(gate, 1) != GGML_HINT_NONE) {
        return false;
@@ -2087,7 +2238,7 @@ static bool ggml_cuda_should_fuse_mul_mat_q(const ggml_tensor * up, const ggml_t
            ggml_nbytes(x) != ggml_backend_buffer_get_alloc_size(x->buffer, x) && x->view_src;
    };
    const int device = ggml_cuda_get_device();
-    if (!ggml_cuda_mul_mat_q_fusion_matches(up, gate, glu, device) || bad_padding_clear(up->src[0]) || bad_padding_clear(gate->src[0])) {
+    if (!ggml_cuda_mul_mat_q_fusion_matches(up, gate, glu, device, up_scale, gate_scale) || bad_padding_clear(up->src[0]) || bad_padding_clear(gate->src[0])) {
        return false;
    }

@@ -2106,7 +2257,7 @@ static bool ggml_cuda_should_fuse_mul_mat_q(const ggml_tensor * up, const ggml_t
    }

    // The fused kernel needs one block per tile (no stream-k split), so only fuse if the tiling is efficient.
-    return ggml_cuda_mmq_fusion_is_efficient(x_up->type, x_up->ne[1], ncols, nchannels, device);
+    return ggml_cuda_mmq_fusion_is_efficient(x_up->type, ggml_prec_get_src(up, 1), x_up->ne[1], ncols, nchannels, device);
 }

 static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
@@ -4181,43 +4332,6 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
    bool fused_mul_mat_vec = false;
    int  fused_node_count  = 0;

-    auto get_mul_mat_scale = [](const ggml_tensor * scale_node, const ggml_tensor * mm_node) -> const ggml_tensor * {
-        const bool scale_lhs_mm = scale_node->src[0] == mm_node;
-        const bool scale_rhs_mm = scale_node->src[1] == mm_node;
-        if (!scale_lhs_mm && !scale_rhs_mm) {
-            return nullptr;
-        }
-
-        const ggml_tensor * scale = scale_lhs_mm ? scale_node->src[1] : scale_node->src[0];
-        if (mm_node->src[0]->type != GGML_TYPE_NVFP4 || scale_node->type != GGML_TYPE_F32 ||
-                scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != 1 ||
-                !ggml_are_same_shape(scale_node, mm_node)) {
-            return nullptr;
-        }
-
-        return scale;
-    };
-
-    auto get_mul_mat_id_scale = [](const ggml_tensor * reshape, const ggml_tensor * repeat, const ggml_tensor * getrows,
-            const ggml_tensor * scale_node, const ggml_tensor * mm_node) -> const ggml_tensor * {
-        if (repeat->src[0] != reshape || getrows->src[0] != repeat || getrows->src[1] != mm_node->src[2]) {
-            return nullptr;
-        }
-        if (!((scale_node->src[0] == mm_node && scale_node->src[1] == getrows) ||
-                (scale_node->src[0] == getrows && scale_node->src[1] == mm_node))) {
-            return nullptr;
-        }
-
-        const ggml_tensor * scale = reshape->src[0];
-        if (mm_node->src[0]->type != GGML_TYPE_NVFP4 || scale_node->type != GGML_TYPE_F32 ||
-                scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != mm_node->src[0]->ne[2] ||
-                !ggml_are_same_shape(scale_node, mm_node)) {
-            return nullptr;
-        }
-
-        return scale;
-    };
-
    auto get_bias_tensor = [](const ggml_tensor * bias_node, const ggml_tensor * mul_node, ggml_op op_bias) -> const ggml_tensor * {
        if (op_bias == GGML_OP_ADD) {
            if (bias_node->src[0] == mul_node) {
@@ -4284,8 +4398,8 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
                    continue;
                }

-                const ggml_tensor * gate_scale = get_mul_mat_scale(gate_scale_n, gate_n);
-                const ggml_tensor * up_scale   = get_mul_mat_scale(up_scale_n, up_n);
+                const ggml_tensor * gate_scale = ggml_cuda_get_mul_mat_scale(gate_scale_n, gate_n);
+                const ggml_tensor * up_scale   = ggml_cuda_get_mul_mat_scale(up_scale_n, up_n);
                if (!gate_scale || !up_scale) {
                    continue;
                }
@@ -4316,6 +4430,14 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
                    fused_node_count  = n_ops;
                    break;
                }
+
+                if (!with_bias && ggml_cuda_should_fuse_mul_mat_q(up_n, gate_n, glu, up_scale_n, gate_scale_n)) {
+                    fusion_data.prec_src1 = ggml_prec_get_src(up_n, 1);
+                    ggml_cuda_mul_mat_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data);
+                    fused_mul_mat_vec = true;
+                    fused_node_count  = n_ops;
+                    break;
+                }
            }

            if (fused_mul_mat_vec) {
@@ -4380,9 +4502,9 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
                    continue;
                }

-                const ggml_tensor * gate_scale = get_mul_mat_id_scale(cgraph->nodes[gate_idx + 1], cgraph->nodes[gate_idx + 2],
+                const ggml_tensor * gate_scale = ggml_cuda_get_mul_mat_id_scale(cgraph->nodes[gate_idx + 1], cgraph->nodes[gate_idx + 2],
                        cgraph->nodes[gate_idx + 3], gate_scale_n, gate_n);
-                const ggml_tensor * up_scale = get_mul_mat_id_scale(cgraph->nodes[up_idx + 1], cgraph->nodes[up_idx + 2],
+                const ggml_tensor * up_scale = ggml_cuda_get_mul_mat_id_scale(cgraph->nodes[up_idx + 1], cgraph->nodes[up_idx + 2],
                        cgraph->nodes[up_idx + 3], up_scale_n, up_n);
                if (!gate_scale || !up_scale) {
                    continue;
@@ -4410,6 +4532,14 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
                    fused_node_count  = n_ops;
                    break;
                }
+
+                if (!with_bias && ggml_cuda_should_fuse_mul_mat_q(up_n, gate_n, glu, up_scale_n, gate_scale_n)) {
+                    fusion_data.prec_src1 = ggml_prec_get_src(up_n, 1);
+                    ggml_cuda_mul_mat_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data);
+                    fused_mul_mat_vec = true;
+                    fused_node_count  = n_ops;
+                    break;
+                }
            }

            if (fused_mul_mat_vec) {
@@ -4525,6 +4655,7 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
                fusion_data.gate      = gate->src[0];
                fusion_data.glu_op    = ggml_get_glu_op(glu);
                fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3);
+                fusion_data.prec_src1 = ggml_prec_get_src(up, 1);

                ggml_cuda_mul_mat_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data);
                fused_mul_mat_vec = true;
@@ -4586,9 +4717,9 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph

            const ggml_tensor * scale = nullptr;
            if (op == GGML_OP_MUL_MAT) {
-                scale = get_mul_mat_scale(scale_node, mm_node);
+                scale = ggml_cuda_get_mul_mat_scale(scale_node, mm_node);
            } else {
-                scale = get_mul_mat_id_scale(cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 3], scale_node, mm_node);
+                scale = ggml_cuda_get_mul_mat_id_scale(cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 3], scale_node, mm_node);
            }
            if (!scale) {
                continue;
@@ -5139,6 +5270,33 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph
                    continue;
                }
            }
+            // Same with the NVFP4 scales: MUL_MAT(gate), MUL(gate scale), MUL_MAT(up), MUL(up scale), GLU.
+            if (i + 4 < cgraph->n_nodes && cgraph->nodes[i + 4]->op == GGML_OP_GLU) {
+                ggml_tensor * glu        = cgraph->nodes[i + 4];
+                ggml_tensor * gate_scale = cgraph->nodes[i + 1];
+                ggml_tensor * up_scale   = cgraph->nodes[i + 3];
+                if (glu->src[0] == gate_scale && glu->src[1] == up_scale &&
+                        ggml_cuda_mul_mat_q_fusion_matches(cgraph->nodes[i + 2], cgraph->nodes[i], glu, cuda_ctx->device, up_scale, gate_scale) &&
+                        ggml_cuda_get_mul_mat_scale(gate_scale, cgraph->nodes[i]) && ggml_cuda_get_mul_mat_scale(up_scale, cgraph->nodes[i + 2])) {
+                    add_alloc_deps(i, i + 4);
+                    i += 4;
+                    continue;
+                }
+            }
+            // Same with per-expert NVFP4 scales: MUL_MAT_ID, RESHAPE, REPEAT, GET_ROWS, MUL for gate and for up, GLU.
+            if (i + 10 < cgraph->n_nodes && cgraph->nodes[i + 10]->op == GGML_OP_GLU) {
+                ggml_tensor * glu        = cgraph->nodes[i + 10];
+                ggml_tensor * gate_scale = cgraph->nodes[i + 4];
+                ggml_tensor * up_scale   = cgraph->nodes[i + 9];
+                if (glu->src[0] == gate_scale && glu->src[1] == up_scale &&
+                        ggml_cuda_mul_mat_q_fusion_matches(cgraph->nodes[i + 5], cgraph->nodes[i], glu, cuda_ctx->device, up_scale, gate_scale) &&
+                        ggml_cuda_get_mul_mat_id_scale(cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 3], gate_scale, cgraph->nodes[i]) &&
+                        ggml_cuda_get_mul_mat_id_scale(cgraph->nodes[i + 6], cgraph->nodes[i + 7], cgraph->nodes[i + 8], up_scale, cgraph->nodes[i + 5])) {
+                    add_alloc_deps(i, i + 10);
+                    i += 10;
+                    continue;
+                }
+            }

            ggml_cuda_moe_weighted_reduction_match match;
            if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) {
@@ -6384,6 +6542,11 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con
    if (strcmp(name, "ggml_backend_comm_allreduce_tensor") == 0) {
        return (void *)ggml_backend_cuda_comm_allreduce_tensor;
    }
+#ifdef GGML_USE_NCCL
+    if (strcmp(name, "ggml_backend_comm_init_rank") == 0) {
+        return (void *)ggml_backend_cuda_comm_init_rank;
+    }
+#endif // GGML_USE_NCCL
    if (strcmp(name, "ggml_backend_register_host_buffer") == 0) {
        return (void *)ggml_backend_cuda_register_host_buffer;
    }
@@ -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);
@@ -1724,10 +1724,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    }
 }

-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4(
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4(
        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kb0, const int i_max, const int stride) {
-    GGML_UNUSED(x_gate);
-
    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);
@@ -1755,7 +1753,7 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
            i = min(i, i_max);
        }

-        const block_nvfp4 * bxi = (const block_nvfp4 *) x + kb0 + i * stride + kbx;
+        const block_nvfp4 * bxi = ggml_cuda_mmq_x_row<block_nvfp4, has_fusion>(x, x_gate, kb0, i, stride) + kbx;
        const uint32_t * __restrict__ src_qs = reinterpret_cast<const uint32_t *>(bxi->qs);
        const int kqs = 16 * kbx;
        const int ksc = 4 * kbx;
@@ -1782,10 +1780,8 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
    }
 }

-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4_nvfp4(
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4_nvfp4(
        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
-    GGML_UNUSED(x_gate);
-
    constexpr int warp_size       = ggml_cuda_get_physical_warp_size();
    constexpr int nwarps          = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size;
    constexpr int I               = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4);
@@ -1800,7 +1796,6 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    const int kbx = txi % threads_per_row;
    const int row_in_warp = txi / threads_per_row;

-    const block_nvfp4 * bxi_base = (const block_nvfp4 *) x + kbx0 + kbx;
    uint32_t * x_u32_scale = x_u32 + 64 + kbx;

 #pragma unroll
@@ -1811,7 +1806,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

-        const block_nvfp4 * bxi = bxi_base + i * stride;
+        const block_nvfp4 * bxi = ggml_cuda_mmq_x_row<block_nvfp4, has_fusion>(x, x_gate, kbx0, i, stride) + kbx;

        const uint32_t * src_qs = reinterpret_cast<const uint32_t *>(bxi->qs);

@@ -114,17 +114,17 @@ static ggml_prec ggml_cuda_mmq_get_prec_env() {

 // src1 is quantized to Q8_1 unless the FP4 types can use 4-bit activations, in which case they
 // default to the native W4A4 instructions on Blackwell.
-static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggml_tensor * dst, const int cc) {
+static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_type type, const ggml_prec prec_requested, const int cc) {
    static const ggml_prec prec_env = ggml_cuda_mmq_get_prec_env();

    ggml_prec prec = prec_env;
    if (prec == GGML_PREC_UNDEFINED) {
-        prec = (ggml_prec) ggml_get_op_params_i32(dst, 3);
+        prec = prec_requested;
    }

    // Q4 only for the FP4 types on Blackwell
    GGML_ASSERT(prec == GGML_PREC_UNDEFINED || prec == GGML_PREC_Q8 || prec == GGML_PREC_Q4);
-    const bool can_use_q4 = (src0->type == GGML_TYPE_NVFP4 || src0->type == GGML_TYPE_MXFP4) && blackwell_mma_available(cc);
+    const bool can_use_q4 = (type == GGML_TYPE_NVFP4 || type == GGML_TYPE_MXFP4) && blackwell_mma_available(cc);
    if (prec == GGML_PREC_Q8 || !can_use_q4) {
        return GGML_PREC_Q8;
    }
@@ -172,11 +172,18 @@ void ggml_cuda_mul_mat_q(
    const ggml_tensor * gate = fusion ? fusion->gate : nullptr;
    ggml_cuda_mm_fusion_args_device fusion_device = {};
    if (fusion) {
-        GGML_ASSERT(gate && !fusion->x_bias && !fusion->gate_bias && !fusion->x_scale && !fusion->gate_scale);
+        GGML_ASSERT(gate && !fusion->x_bias && !fusion->gate_bias && !fusion->x_scale == !fusion->gate_scale);
        GGML_ASSERT(gate->type == src0->type && ggml_are_same_shape(gate, src0) && ggml_are_same_stride(gate, src0));
-        fusion_device.gate      = gate->data;
-        fusion_device.glu_op    = fusion->glu_op;
-        fusion_device.glu_limit = fusion->glu_limit;
+        // NVFP4 weight scales (only applied for NVFP4): one value per tensor, or one per expert with ids.
+        for (const ggml_tensor * scale : { fusion->x_scale, fusion->gate_scale }) {
+            GGML_ASSERT(!scale || (src0->type == GGML_TYPE_NVFP4 && scale->type == GGML_TYPE_F32 && ggml_is_contiguous(scale) &&
+                ggml_nelements(scale) == (ids ? src0->ne[2] : 1)));
+        }
+        fusion_device.gate       = gate->data;
+        fusion_device.x_scale    = fusion->x_scale    ? fusion->x_scale->data    : nullptr;
+        fusion_device.gate_scale = fusion->gate_scale ? fusion->gate_scale->data : nullptr;
+        fusion_device.glu_op     = fusion->glu_op;
+        fusion_device.glu_limit  = fusion->glu_limit;
    }

    GGML_TENSOR_BINARY_OP_LOCALS;
@@ -224,9 +231,8 @@ void ggml_cuda_mul_mat_q(

    const bool fallback = ggml_cuda_mmq_needs_fallback(ne01);

-    // With fusion dst is the GLU output, not a MUL_MAT; the fused types always quantize src1 to Q8_1.
-    // TODO support for Q4 precision if both tensors allow it
-    const ggml_prec prec_src1 = gate ? GGML_PREC_Q8 : ggml_cuda_mmq_get_prec_src1(src0, dst, cc);
+    // With fusion dst is the GLU output, not a MUL_MAT; the src1 precision requested by both fused MUL_MATs is in the fusion args.
+    const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0->type, gate ? fusion->prec_src1 : ggml_prec_get_src(dst, 1), cc);

    const bool use_native_fp4 = prec_src1 == GGML_PREC_Q4;
    const size_t y_block_size       = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
@@ -241,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);
@@ -258,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());
@@ -295,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;
    }
@@ -310,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).
@@ -321,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());
    }

@@ -369,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);
 }
@@ -513,18 +520,20 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
    return (!GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
 }

-bool ggml_cuda_mmq_fusion_is_efficient(const ggml_type type, const int64_t nrows_x, const int64_t ncols, const int64_t nchannels, const int device) {
+bool ggml_cuda_mmq_fusion_is_efficient(
+        const ggml_type type, const ggml_prec prec_requested, const int64_t nrows_x, const int64_t ncols, const int64_t nchannels, const int device) {
    const int    cc    = ggml_cuda_info().devices[device].cc;
    const int    nsm   = ggml_cuda_info().devices[device].nsm;
    const size_t smpbo = ggml_cuda_info().devices[device].smpbo;

-    // Same tile width as ggml_cuda_mul_mat_q; the fused kernel has no fallback, quantizes src1 to Q8_1 and writes I/2 output rows per tile.
-    const bool fallback = false;
-    const int  J_best   = ggml_cuda_mmq_get_J_best(type, ncols, fallback, cc, smpbo, GGML_PREC_Q8);
+    // Same tile width as ggml_cuda_mul_mat_q; the fused kernel has no fallback and writes I/2 output rows per tile.
+    const bool      fallback  = false;
+    const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(type, prec_requested, cc);
+    const int       J_best    = ggml_cuda_mmq_get_J_best(type, ncols, fallback, cc, smpbo, prec_src1);
    if (J_best == 0) {
        return false;
    }
-    const int64_t nrows_tile = ggml_cuda_mmq_get_config(type, J_best, fallback, cc, GGML_PREC_Q8).I / 2;
+    const int64_t nrows_tile = ggml_cuda_mmq_get_config(type, J_best, fallback, cc, prec_src1).I / 2;
    const int64_t ntiles     = (ncols + J_best - 1)/J_best * ((nrows_x + nrows_tile - 1)/nrows_tile) * nchannels;
    // Same condition as launch_mul_mat_q uses to run the stream-k kernel with one block per tile.
    const int64_t tiles_nwaves             = (ntiles + nsm - 1) / nsm;
@@ -1,6 +1,7 @@
 #pragma once

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

 #include <climits>
@@ -14,9 +15,9 @@
 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,
-    const ggml_glu_op glu_op, const float glu_limit);
+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 {
    MMQ_Q8_1_DS_LAYOUT_D4,
@@ -436,7 +437,8 @@ static __host__ int ggml_cuda_mmq_get_nbytes_shared_x(const ggml_cuda_mmq_config

 // Weight types with a fused gate/up/GLU path (see ggml_cuda_mmq_get_util_funcs).
 static constexpr __host__ __device__ bool ggml_cuda_mmq_fusion_supported(const ggml_type type) {
-    return type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K || type == GGML_TYPE_Q6_K || type == GGML_TYPE_Q8_0 || type == GGML_TYPE_IQ4_XS;
+    return type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K || type == GGML_TYPE_Q6_K || type == GGML_TYPE_Q8_0 || type == GGML_TYPE_IQ4_XS ||
+        type == GGML_TYPE_NVFP4;
 }

 // Fused gate/up/GLU: the weight tile alternates chunks of this many up and gate rows, so rows r and r + 8 of an MMA tile (same thread) pair up.
@@ -461,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 ggml_glu_op glu_op, const float glu_limit) {
-    GGML_UNUSED_VARS(glu_op, glu_limit);
-
+        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;
        }

@@ -488,26 +489,22 @@ 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);
            }
        }
    }
 }

 // has_fusion: accumulators l and l + ne/2 hold the up and gate values of one output element, the GLU of the two is written.
+// x_scale/gate_scale: optional NVFP4 scales of the up and gate weights (per tensor, or of the expert of this tile).
 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 ggml_glu_op glu_op, const float glu_limit) {
+            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)
    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
@@ -521,7 +518,12 @@ 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;
+    const float scale_up   = fused_nvfp4 && x_scale ? *x_scale    : 1.0f;
+    const float scale_gate = fused_nvfp4 && gate_scale ? *gate_scale : 1.0f;

 #pragma unroll
    for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
@@ -531,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;
                }

@@ -541,24 +543,29 @@ 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;
-                    dst[ids_dst[j]*stride + i] = ggml_cuda_op_glu_single(glu_op, sum[sum_index + tile_C::ne/2], sum[sum_index], 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];
+                    float up   = sum[sum_index];
+                    float gate = sum[sum_index + tile_C::ne/2];
+                    if constexpr (type == GGML_TYPE_NVFP4) {
+                        if (y_scale_used) {
+                            up   *= y_scale[j];
+                            gate *= y_scale[j];
+                        }
+                        up   *= scale_up;
+                        gate *= scale_gate;
                    }
+                    dst[ids_dst[j]*stride + i] = ggml_cuda_op_glu_single(glu_op, gate, up, glu_limit);
+                } 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);
                }
            }
        }
    }

-    GGML_UNUSED_VARS(y_scale_used, glu_op, glu_limit);
+    GGML_UNUSED_VARS(y_scale_used, glu_op, glu_limit, x_scale, gate_scale, scale_up, scale_gate);
 }

 // -------------------------------------------------------------------------------------------------------------------------------------
@@ -712,7 +719,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
            case GGML_TYPE_NVFP4:
                return ggml_cuda_mmq_util_funcs(
                    VDR_NVFP4_Q8_1_MMQ,
-                    ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback, prec_src1>,
+                    ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback, prec_src1, false>,
                    ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_dp4a<type, J, fallback>,
                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
            default:
@@ -737,9 +744,9 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
            if (prec_src1 == GGML_PREC_Q4) {
                return ggml_cuda_mmq_util_funcs(
                    -1,
-                    ggml_cuda_mmq_load_tiles_nvfp4_nvfp4<type, J, fallback>,
+                    ggml_cuda_mmq_load_tiles_nvfp4_nvfp4<type, J, fallback, has_fusion>,
                    ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
-                    ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
+                    ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
            }
            break;
        default:
@@ -882,9 +889,9 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
        case GGML_TYPE_NVFP4:
            return ggml_cuda_mmq_util_funcs(
                -1,
-                ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback, prec_src1>,
+                ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback, prec_src1, has_fusion>,
                ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
        default:
            return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr);
    }
@@ -918,8 +925,9 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
        const char * __restrict__ x, const char * __restrict__ x_gate, const int offset_x, const int * __restrict__ y,
        const int * __restrict__ ids_dst, 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 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");

@@ -952,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
@@ -987,9 +995,9 @@ 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);
+        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);
+        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);
    }
 }

@@ -1000,21 +1008,28 @@ 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;
 #else
    constexpr bool fusion_available = false;
 #endif // TURING_MMA_AVAILABLE
+#ifdef BLACKWELL_MMA_AVAILABLE
+    constexpr bool native_fp4_available = true;
+#else
+    constexpr bool native_fp4_available = false;
+#endif // BLACKWELL_MMA_AVAILABLE

-    // Skip unused template specializations for faster compilation:
-    if (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type == GGML_TYPE_COUNT || (has_fusion && !fusion_available)) {
+    // Skip unused template specializations for faster compilation (the host only uses 4-bit src1 with native FP4):
+    if (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type == GGML_TYPE_COUNT ||
+            (has_fusion && (!fusion_available || (prec_src1 == GGML_PREC_Q4 && !native_fp4_available)))) {
        NO_DEVICE_CODE;
        return;
    }
@@ -1024,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.
@@ -1049,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;
            }

@@ -1090,29 +1098,30 @@ 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;

+        // With ids the NVFP4 weight scales are per expert.
+        const float * x_scale_tile    = x_scale    && ids_dst ? x_scale    + zt : x_scale;
+        const float * gate_scale_tile = gate_scale && ids_dst ? gate_scale + zt : gate_scale;
+
        constexpr bool fixup = false;
        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, stride_row_x, ncols_y, stride_col_dst,
-             tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z);
+             glu_op, glu_limit, x_scale_tile, gate_scale_tile, stride_row_x, ncols_y, stride_col_dst,
+             tile_x_i_max, col_high - col_low, 0, blocks_per_ne00.z);
        return;
    }

@@ -1120,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;
@@ -1131,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);

@@ -1184,29 +1186,26 @@ 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;

        constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
        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, stride_row_x, ncols_y, stride_col_dst,
-             tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
+             glu_op, glu_limit, x_scale, gate_scale, stride_row_x, ncols_y, stride_col_dst,
+             tile_x_i_max, col_high - col_low, kb0_start, kb0_stop);

        kbc += blocks_per_ne00.z;
        kbc -= fastmodulo(kbc, blocks_per_ne00);
@@ -1220,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;
        }

@@ -1273,33 +1265,30 @@ 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;

    constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
    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, stride_row_x, ncols_y, stride_col_dst,
-         tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
+         glu_op, glu_limit, x_scale, gate_scale, stride_row_x, ncols_y, stride_col_dst,
+         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);
@@ -1310,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;
@@ -1335,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.
@@ -1367,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;
@@ -1382,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;
        }
@@ -1391,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;
            }

@@ -1401,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();

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

@@ -1432,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
 };
@@ -1469,38 +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,
+            (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);

@@ -1521,11 +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,
+        (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;
@@ -1533,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>
@@ -1661,4 +1652,5 @@ void ggml_cuda_mul_mat_q(

 bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts);

-bool ggml_cuda_mmq_fusion_is_efficient(ggml_type type, int64_t nrows_x, int64_t ncols, int64_t nchannels, int device);
+// prec_requested: the src1 precision requested by the fused MUL_MATs.
+bool ggml_cuda_mmq_fusion_is_efficient(ggml_type type, ggml_prec prec_requested, int64_t nrows_x, int64_t ncols, int64_t nchannels, int device);
@@ -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;
        }
@@ -83,21 +83,21 @@ static __global__ void quantize_q8_1(
    ggml_cuda_pdl_sync();
    const float xi = i0 < ne00 ? x[i03*s03 + i02*s02 + i01*s01 + i00] : 0.0f;
    float amax = fabsf(xi);
-    float sum = xi;

    amax = warp_reduce_max<QK8_1>(amax);
-    sum  = warp_reduce_sum<QK8_1>(sum);

    const float  d = amax / 127.0f;
    const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);

+    const int sum = warp_reduce_sum<QK8_1>((int) q);
+
    y[ib].qs[iqs] = q;

    if (iqs > 0) {
        return;
    }

-    y[ib].ds = make_half2(d, sum);
+    y[ib].ds = make_half2(d, d * sum);
 }

 __device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) {
@@ -500,17 +500,6 @@ static __global__ void quantize_mmq_q8_1(
        amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE));
    }

-    float sum;
-    if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {
-        sum = xi.x + xi.y + xi.z + xi.w;
-
-        // Calculate sums across vals_per_sum/4 threads.
-#pragma unroll
-        for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {
-            sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset, WARP_SIZE);
-        }
-    }
-
    const float d_inv = 127.0f / amax;
    char4 q;
    q.x = roundf(xi.x*d_inv);
@@ -519,6 +508,19 @@ static __global__ void quantize_mmq_q8_1(
    q.w = roundf(xi.w*d_inv);
    const float d = 1.0f / d_inv;

+    float sum = 0.0f;
+    if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {
+        int sum_q = int(q.x) + int(q.y) + int(q.z) + int(q.w);
+
+        // Calculate sums across vals_per_sum/4 threads.
+#pragma unroll
+        for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {
+            sum_q += __shfl_xor_sync(0xFFFFFFFF, sum_q, offset, WARP_SIZE);
+        }
+
+        sum = d * float(sum_q);
+    }
+
    // write the block once (normal) or to each of the token's compact rows (scatter)
    const int nwrite = scatter ? n_expert_used : 1;
 #pragma unroll
@@ -591,7 +591,6 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
    //const int n_past     = ((int32_t *) dst->op_params)[0];
    const int n_dims     = ((int32_t *) dst->op_params)[1];
    const int mode       = ((int32_t *) dst->op_params)[2];
-    //const int n_ctx      = ((int32_t *) dst->op_params)[3];
    const int n_ctx_orig = ((int32_t *) dst->op_params)[4];
    const int n_offs     = ((int32_t *) dst->op_params)[15];
    mrope_sections sections;
@@ -637,7 +636,7 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
    }

    rope_corr_dims corr_dims;
-    ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims.v);
+    ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, ggml_rope_get_truncate(dst), corr_dims.v);

    // compute
    if (is_neox) {
@@ -929,7 +928,7 @@ void ggml_cuda_op_rms_norm_mul_rope_fused(ggml_backend_cuda_context & ctx,
    const float * freq_factors = rope->src[2] != nullptr ? (const float *) rope->src[2]->data : nullptr;

    rope_corr_dims corr_dims;
-    ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims.v);
+    ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, ggml_rope_get_truncate(rope), corr_dims.v);

    const size_t ts0 = ggml_type_size(x->type);
    GGML_ASSERT(x->nb[0] == ts0);
@@ -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;
+            }
        }
    }

@@ -1069,7 +1069,7 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm
                const bool zero_rot_offset  = ggml_get_op_params_i32(op, 15) == 0;

                supported =
-                    zero_view_offset && zero_rot_offset && ndims <= 512 &&
+                    ggml_rope_get_truncate(op) && zero_view_offset && zero_rot_offset && ndims <= 512 &&
                    (is_normal || (is_neox && ndims % 16 == 0) || (is_imrope && ndims % 16 == 0 && has_sections));
            } else {
                supported = false;
@@ -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

@@ -135,6 +135,7 @@ size_t opt_mbuf_total  = 0;                        // total buffer space limit (
 int    opt_mm_select     = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
 int    opt_fa_select     = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
 int    opt_fa_split      = 1; // 1 = head-split + team-split when n_kv_heads < N_cores (default), 0 = off
+int    opt_li_select     = 2; // 2 = HMX -> HVX, 1 = HVX, 0 = CPU (unsupported)
 int    opt_gdn_select    = 2; // 2 = HMX -> HVX, 1 = HVX, 0 = CPU (unsupported)
 int    opt_ar_select     = 2; // 2 = fused ALLREDUCE+ADD (default), 1 = unfused ALLREDUCE, 0 = fallback to CPY+FENCE
 int    opt_ar_scatter    = 1; // 1 = reduce-scatter the fused ALLREDUCE+ADD (default), 0 = full reduction
@@ -149,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

@@ -462,7 +463,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;
@@ -485,6 +485,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)");
        }

@@ -560,9 +563,6 @@ struct ggml_hexagon_shared_buffer {

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

@@ -621,11 +621,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;
@@ -658,10 +656,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;
 }
@@ -671,42 +666,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,
@@ -714,14 +686,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;
    }
@@ -741,47 +713,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);
    }
 }

@@ -793,14 +745,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);
        }
@@ -824,6 +776,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;
@@ -843,7 +818,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,
 };
@@ -855,8 +830,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,
@@ -868,6 +847,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,
@@ -880,7 +878,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,
 };
@@ -895,7 +893,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;
@@ -913,7 +911,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;
@@ -933,7 +931,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];
@@ -1211,8 +1209,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 ...
@@ -1231,7 +1228,6 @@ struct ggml_hexagon_opbatch {
        b_vmem = 0;

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

@@ -1251,7 +1247,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);
@@ -1297,7 +1292,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];
@@ -1305,7 +1300,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);
@@ -1326,31 +1321,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);
@@ -1359,8 +1365,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();

@@ -1390,8 +1395,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++;
            }

@@ -2351,6 +2359,7 @@ ggml_hexagon_session::~ggml_hexagon_session() noexcept(true) {

 // ** backend interface

+
 static const char * ggml_backend_hexagon_name(ggml_backend_t backend) {
    auto sess = static_cast<ggml_hexagon_session *>(backend->context);
    return sess->c_name();
@@ -2362,11 +2371,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);

@@ -2383,13 +2414,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;
            }

@@ -2399,22 +2432,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));
        }
@@ -2624,7 +2642,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;

@@ -2659,7 +2681,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) {
@@ -2721,14 +2743,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);
@@ -3001,18 +3016,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);
@@ -3409,6 +3424,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
    const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT");
    const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT");
    const char * str_fa_split  = getenv("GGML_HEXAGON_FA_SPLIT");
+    const char * str_li_select = getenv("GGML_HEXAGON_LI_SELECT");
    const char * str_gdn_select = getenv("GGML_HEXAGON_GDN_SELECT");
    const char * str_ar_select = getenv("GGML_HEXAGON_AR_SELECT");
    const char * str_ar_scatter = getenv("GGML_HEXAGON_AR_SCATTER");
@@ -3475,6 +3491,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
    opt_mm_select = str_mm_select ? atoi(str_mm_select)                   : opt_mm_select;
    opt_fa_select = str_fa_select ? atoi(str_fa_select)                   : opt_fa_select;
    opt_fa_split  = str_fa_split  ? atoi(str_fa_split)                   : opt_fa_split;
+    opt_li_select = str_li_select ? atoi(str_li_select)                   : opt_li_select;
    opt_gdn_select    = str_gdn_select    ? atoi(str_gdn_select)          : opt_gdn_select;
    opt_ar_select     = str_ar_select     ? atoi(str_ar_select)           : opt_ar_select;
    opt_ar_scatter    = str_ar_scatter    ? atoi(str_ar_scatter)          : opt_ar_scatter;
@@ -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;
@@ -29,6 +23,7 @@ extern size_t opt_mbuf_total;
 extern int    opt_mm_select;
 extern int    opt_fa_select;
 extern int    opt_fa_split;
+extern int    opt_li_select;
 extern int    opt_gdn_select;
 extern int    opt_ar_select;
 extern int    opt_ar_scatter;
@@ -45,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) {
@@ -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;
@@ -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;
    }
@@ -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) {
@@ -11,12 +11,14 @@
 #include "htp/cpy-ops.h"
 #include "htp/cumsum-ops.h"
 #include "htp/diag-ops.h"
+#include "htp/dsv4-hc-ops.h"
 #include "htp/fill-ops.h"
 #include "htp/flash-attn-ops.h"
 #include "htp/gated-delta-net-ops.h"
 #include "htp/get-rows-ops.h"
 #include "htp/glu-ops.h"
 #include "htp/im2col-ops.h"
+#include "htp/lightning-indexer-ops.h"
 #include "htp/matmul-ops.h"
 #include "htp/pad-ops.h"
 #include "htp/pool-ops.h"
@@ -28,6 +30,7 @@
 #include "htp/solve-tri-ops.h"
 #include "htp/ssm-conv.h"
 #include "htp/sum-rows-ops.h"
+#include "htp/top-k-mask-ops.h"
 #include "htp/unary-ops.h"
 #include "htp/upscale-ops.h"

@@ -92,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,
@@ -5,11 +5,13 @@
 #include "htp/binary-ops.h"
 #include "htp/cumsum-ops.h"
 #include "htp/diag-ops.h"
+#include "htp/dsv4-hc-ops.h"
 #include "htp/fill-ops.h"
 #include "htp/flash-attn-ops.h"
 #include "htp/gated-delta-net-ops.h"
 #include "htp/get-rows-ops.h"
 #include "htp/glu-ops.h"
+#include "htp/lightning-indexer-ops.h"
 #include "htp/matmul-ops.h"
 #include "htp/pad-ops.h"
 #include "htp/repeat-ops.h"
@@ -18,6 +20,7 @@
 #include "htp/solve-tri-ops.h"
 #include "htp/ssm-conv.h"
 #include "htp/sum-rows-ops.h"
+#include "htp/top-k-mask-ops.h"
 #include "htp/unary-ops.h"
 #include "htp/upscale-ops.h"

@@ -185,6 +188,22 @@ static std::string format_kernel_params(const htp_opnode & node) {
        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);
+    } 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);
+    } 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);
+    } 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",
+                 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);
    } 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>();
@@ -213,6 +232,14 @@ static std::string format_kernel_params(const htp_opnode & node) {
    } 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);
+    } else if (node.opcode == HTP_OP_TOP_K_MASK) {
+        const auto * kparams = node.kparams<htp_top_k_mask_kernel_params>();
+        const char * ktype_str = "?";
+        switch (kparams->kernel_type) {
+            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);
    } 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);
@@ -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);

@@ -37,6 +37,7 @@ add_library(${HTP_LIB} SHARED
    repeat-ops.c
    ssm-conv.c
    cumsum-ops.c
+    top-k-mask-ops.c
    fill-ops.c
    concat-ops.c
    diag-ops.c
@@ -49,6 +50,8 @@ add_library(${HTP_LIB} SHARED
    pool-ops.c
    upscale-ops.c
    allreduce-ops.c
+    dsv4-hc-ops.c
+    lightning-indexer-ops.c
 )

 target_compile_definitions(${HTP_LIB} PRIVATE
@@ -315,55 +315,61 @@ static void bitonic_sort_vtcm_desc(uint8_t * values, uint8_t * indices, uint32_t
    }
 }

-static void top_k_select_tiled(dma_queue * dma_q, struct htp_thread_trace * tr,
-                               dma_addr_t src_addr, uint32_t n, uint32_t k,
-                               float * values_buf, int32_t * indices_buf,
-                               float * out_values, int32_t * out_indices,
-                               uint32_t chunk_base) {
-    const uint32_t tile_elems = 1024;
-    uint32_t n_tiles            = (n + tile_elems - 1) / tile_elems;
-    uint32_t candidate_count    = n_tiles * k;
-    uint32_t merge_n_vec        = hmx_ceil_div(candidate_count, 32);
-    uint32_t merge_n_vec_pow2   = 1;
-    while (merge_n_vec_pow2 < merge_n_vec) merge_n_vec_pow2 <<= 1;
-    uint32_t merge_elems        = merge_n_vec_pow2 * 32;
-    float * candidate_values    = values_buf + tile_elems;
-    int32_t * candidate_indices = indices_buf + tile_elems;
-    uint32_t candidate_pos      = 0;
+// top-k of one row held in VTCM, k <= 64. Keeps a sorted list of the k best, rejects 256-float blocks with one max-reduce.
+// Writes k values and k indices (ties: first seen wins). Slots with no finite value get -INF and index 0.
+static void top_k_scan(const float * vals, uint32_t n, uint32_t k, float * out_v, int32_t * out_i) {
+    float   best_v[64];
+    int32_t best_i[64];
+    uint32_t cnt = 0;
+    float    thr = -INFINITY;

-    for (uint32_t offset = 0; offset < n; offset += tile_elems) {
-        uint32_t tile_count = MIN(tile_elems, n - offset);
-        dma_sync_read(dma_q, values_buf, src_addr + (size_t) offset * sizeof(float), tile_count * sizeof(float));
-
-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) offset);
-        if (tile_count < tile_elems) {
-            hvx_splat_f32_u((uint8_t *) (values_buf + tile_count), -INFINITY, tile_elems - tile_count);
+    const uint32_t blk = 256;
+    for (uint32_t off = 0; off < n; off += blk) {
+        const uint32_t bn = MIN(blk, n - off);
+        if (hvx_reduce_max_f32((const uint8_t *) (vals + off), bn) <= thr) {
+            continue;
        }
-
-        bitonic_sort_vtcm_desc((uint8_t *) values_buf, (uint8_t *) indices_buf, tile_elems / 32, true);
-        uint32_t tile_k = MIN(k, tile_count);
-        for (uint32_t j = 0; j < tile_k; j++) {
-            candidate_values[candidate_pos] = values_buf[j];
-            candidate_indices[candidate_pos] = indices_buf[j] + (int32_t) offset;
-            candidate_pos++;
-        }
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) offset);
-    }
-
-    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
-    if (merge_elems > candidate_pos) {
-        hvx_splat_f32_u((uint8_t *) (candidate_values + candidate_pos), -INFINITY, merge_elems - candidate_pos);
-        for (uint32_t j = candidate_pos; j < merge_elems; j++) {
-            candidate_indices[j] = 0;
+        for (uint32_t vi = 0; vi < bn; vi += 32) {
+            const uint32_t vn = MIN(32, bn - vi);
+            HVX_VectorAlias v_local;
+            v_local.v = hvx_vmem(vals + off + vi);
+            for (uint32_t i = 0; i < vn; i++) {
+                const float v = v_local.fp32[i];
+                if (!(v > thr)) continue;
+                uint32_t pos = cnt < k ? cnt++ : k - 1;
+                while (pos > 0 && best_v[pos - 1] < v) {
+                    best_v[pos] = best_v[pos - 1];
+                    best_i[pos] = best_i[pos - 1];
+                    pos--;
+                }
+                best_v[pos] = v;
+                best_i[pos] = (int32_t) (off + vi + i);
+                if (cnt == k) thr = best_v[k - 1];
+            }
        }
    }

-    bitonic_sort_vtcm_desc((uint8_t *) candidate_values, (uint8_t *) candidate_indices, merge_n_vec_pow2, false);
-    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, 0);
+    // fewer than k values above -INF: take -INF or NaN elements, then pad with index 0
+    for (uint32_t off = 0; off < n && cnt < k; off += 32) {
+        const uint32_t vn = MIN(32, n - off);
+        HVX_VectorAlias v_local;
+        v_local.v = hvx_vmem(vals + off);
+        for (uint32_t j = 0; j < vn && cnt < k; j++) {
+            if (!(v_local.fp32[j] > -INFINITY)) {
+                best_v[cnt] = -INFINITY;
+                best_i[cnt] = (int32_t) (off + j);
+                cnt++;
+            }
+        }
+    }
+    for (; cnt < k; cnt++) {
+        best_v[cnt] = -INFINITY;
+        best_i[cnt] = 0;
+    }

    for (uint32_t j = 0; j < k; j++) {
-        out_values[j]  = candidate_values[j];
-        out_indices[j] = candidate_indices[j] + (int32_t) chunk_base;
+        out_v[j] = best_v[j];
+        out_i[j] = best_i[j];
    }
 }

@@ -509,10 +515,14 @@ static void htp_sort_chunk_job(unsigned int n, unsigned int i, void * data) {
        }

        if (local_k > 1 && local_k <= 64 && chunk_elems > 1024) {
-            top_k_select_tiled(octx->ctx->dma[i], tr, src_addr, real_count, local_k,
-                               values_buf, indices_buf,
-                               &merge_values[i * local_k], &merge_indices[i * local_k],
-                               chunk_base);
+            dma_sync_read(octx->ctx->dma[i], values_buf, src_addr, real_count * sizeof(float));
+
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
+            top_k_scan(values_buf, real_count, local_k, &merge_values[i * local_k], &merge_indices[i * local_k]);
+            for (uint32_t j = 0; j < local_k; j++) {
+                merge_indices[i * local_k + j] += (int32_t) chunk_base;
+            }
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
            return;
        }

@@ -668,10 +678,14 @@ static inline void compute_row_sort(
    uint32_t ne00,
    uint32_t chunk_elems,
    bool is_top_k,
-    enum ggml_sort_order order
+    enum ggml_sort_order order,
+    uint32_t k
 ) {
    if (is_top_k) {
-        if (ne00 == 32 || ne00 == 64 || ne00 == 128 || ne00 == 256 || ne00 == 512 || ne00 == 1024) {
+        if (k > 1 && k <= 64 && ne00 > 1024) {
+            float out_v[64];
+            top_k_scan(cur_values, ne00, k, out_v, cur_indices);
+        } else if (ne00 == 32 || ne00 == 64 || ne00 == 128 || ne00 == 256 || ne00 == 512 || ne00 == 1024) {
            switch (ne00) {
                case 32:   sort32_f32_hvx((uint8_t *) cur_values, (uint8_t *) cur_indices, GGML_SORT_ORDER_DESC); break;
                case 64:   sort64_f32_hvx((uint8_t *) cur_values, (uint8_t *) cur_indices, GGML_SORT_ORDER_DESC); break;
@@ -788,7 +802,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
            dma_queue_pop(dma_q);

            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);
-            compute_row_sort(cur_values, cur_indices, ne00, chunk_elems, is_top_k, order);
+            compute_row_sort(cur_values, cur_indices, ne00, chunk_elems, is_top_k, order, k);
            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);

            // Push writeback of current slot
@@ -817,7 +831,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
            dma_sync_read(dma_q, values_buf, src_addr, src_bytes);

            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);
-            compute_row_sort(values_buf, indices_buf, ne00, chunk_elems, is_top_k, order);
+            compute_row_sort(values_buf, indices_buf, ne00, chunk_elems, is_top_k, order, k);
            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);

            dma_sync_write(dma_q, dst_addr, indices_buf, dst_bytes);
@@ -87,11 +87,11 @@ static inline HVX_Vector conv2d_hmx_rt_restore(const __fp16 * scales, uint32_t r
    return Q6_V_vsplat_R((127 + e) << 23);
 }

-// Output rows [*row_base, *row_base + *n_rows) of N * OH for this device.
+// Output rows [*row_base, *row_base + *n_rows) of total for this device.
 // The output is written by DMA only, so device ranges need no cache line alignment.
-static void conv2d_mdev_rows(const struct htp_ops_context * octx, uint32_t * row_base, uint32_t * n_rows) {
+static void conv2d_mdev_rows(const struct htp_ops_context * octx, uint32_t total, uint32_t * row_base, uint32_t * n_rows) {
    *row_base = 0;
-    *n_rows   = octx->dst->ne[3] * octx->dst->ne[1];
+    *n_rows   = total;
    if (octx->ctx->mdev.count > 1) {
        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
            *n_rows, 1, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
@@ -134,9 +134,51 @@ static void conv2d_hmx_worker(void * data) {
    }
 }

+// Tiled path values whose source differs between conv_2d and conv_3d.
+// conv_3d runs as conv_2d with virtual KH = KD * KHr and OH = OD * OHr, only the input stage unfolds them.
+struct htp_conv_geom {
+    uint32_t KH, KHr, Cin, Cout;
+    uint32_t OH, OHr, N, ID;
+    int32_t  s0, s1, s2, p0, p1, p2, d0, d1, d2;
+    size_t   x_nb_d, x_nb_c, x_nb_n;   // input strides for d, cin, n
+    size_t   y_nb_c, y_nb_n;           // dst strides for cout, n
+    size_t   w_nb_co;                  // kernel cout stride
+};
+
+// conv_3d tensors are [W, H, D, C * N] with the channel as the inner index of ne[3]
+static void conv_geom_init(struct htp_conv_geom * g, const struct htp_ops_context * octx) {
+    const struct htp_tensor * kernel = octx->src[0];
+    const struct htp_tensor * input  = octx->src[1];
+    const struct htp_tensor * dst    = octx->dst;
+    const int32_t *           op     = octx->op_params;
+
+    if (octx->op == HTP_OP_CONV_3D) {
+        const uint32_t IC = op[9];
+        const uint32_t OC = op[11];
+        *g = (struct htp_conv_geom) {
+            .KH = kernel->ne[2] * kernel->ne[1], .KHr = kernel->ne[1], .Cin = IC, .Cout = OC,
+            .OH = dst->ne[2] * dst->ne[1], .OHr = dst->ne[1], .N = op[10], .ID = input->ne[2],
+            .s0 = op[0], .s1 = op[1], .s2 = op[2], .p0 = op[3], .p1 = op[4], .p2 = op[5], .d0 = op[6], .d1 = op[7], .d2 = op[8],
+            .x_nb_d = input->nb[2], .x_nb_c = input->nb[3], .x_nb_n = (size_t) IC * input->nb[3],
+            .y_nb_c = dst->nb[3], .y_nb_n = (size_t) OC * dst->nb[3],
+            .w_nb_co = (size_t) IC * kernel->nb[3],
+        };
+    } else {
+        *g = (struct htp_conv_geom) {
+            .KH = kernel->ne[1], .KHr = kernel->ne[1], .Cin = kernel->ne[2], .Cout = kernel->ne[3],
+            .OH = dst->ne[1], .OHr = dst->ne[1], .N = dst->ne[3], .ID = 1,
+            .s0 = op[0], .s1 = op[1], .s2 = 1, .p0 = op[2], .p1 = op[3], .p2 = 0, .d0 = op[4], .d1 = op[5], .d2 = 1,
+            .x_nb_d = 0, .x_nb_c = input->nb[2], .x_nb_n = input->nb[3],
+            .y_nb_c = dst->nb[2], .y_nb_n = dst->nb[3],
+            .w_nb_co = kernel->nb[3],
+        };
+    }
+}
+
 struct htp_conv2d_hmx_context {
    struct htp_ops_context *                octx;
    const struct htp_conv2d_kernel_params * kp;
+    const struct htp_conv_geom *            g;
    struct htp_conv2d_vtcm_layout           layout;
    __fp16 *                                scales;  // [rt] HMX scales
    __fp16 *                                w_tiles;
@@ -169,7 +211,7 @@ static void conv2d_hmx_wpack_thread(unsigned int nth, unsigned int ith, void * d
    const struct htp_conv2d_kernel_params * kp   = hc->kp;
    const struct htp_tensor *               kernel = octx->src[0];

-    const uint32_t Cout   = kernel->ne[3];
+    const uint32_t Cout   = hc->g->Cout;
    const uint32_t K      = kp->tiled.k_size;
    const uint32_t k_pad  = kp->tiled.n_k_tiles * HMX_FP16_TILE_N_COLS;
    const uint32_t WK     = kp->tiled.wk_chunk;
@@ -200,8 +242,8 @@ static void conv2d_hmx_wpack_thread(unsigned int nth, unsigned int ith, void * d
        for (uint32_t i = 0; i < 2 && i * WK < k_pad; i++) {
            const uint32_t kc = i * WK;
            const uint32_t kv = MIN(WK, K - kc);
-            dma_queue_push(dmaq, dma_make_data(wdma[i], kernel->data + (size_t) co0 * kernel->nb[3] + (size_t) kc * es),
-                           WK * es, kernel->nb[3], kv * es, nco);
+            dma_queue_push(dmaq, dma_make_data(wdma[i], kernel->data + (size_t) co0 * hc->g->w_nb_co + (size_t) kc * es),
+                           WK * es, hc->g->w_nb_co, kv * es, nco);
        }

        for (uint32_t kc = 0; kc < k_pad; kc += WK) {
@@ -234,42 +276,50 @@ static void conv2d_hmx_wpack_thread(unsigned int nth, unsigned int ith, void * d
            const uint32_t next_kc = kc + 2 * WK;
            if (next_kc < k_pad) {
                const uint32_t next_kv = MIN(WK, K - next_kc);
-                dma_queue_push(dmaq, dma_make_data(wdma[bi], kernel->data + (size_t) co0 * kernel->nb[3] + (size_t) next_kc * es),
-                               WK * es, kernel->nb[3], next_kv * es, nco);
+                dma_queue_push(dmaq, dma_make_data(wdma[bi], kernel->data + (size_t) co0 * hc->g->w_nb_co + (size_t) next_kc * es),
+                               WK * es, hc->g->w_nb_co, next_kv * es, nco);
            }
            bi ^= 1;
        }
    }
 }

-// Stage Cin*KH input rows for G col tiles: stage[(cin*KH + kh)*SR + j] = x[n][cin][ih(kh)][ow0 - p0 + j], zero outside
+// Stage Cin*KH input rows for G col tiles: stage[(cin*KH + kh)*SR + j] = x[n][cin][ih(kh)][ow0*s0 - p0 + j], zero outside
+// For conv_3d, kh = kd * KHr + kh' and oh = od * OHr + oh' are unfolded here
 static void conv2d_hmx_stage(const struct htp_ops_context * octx, const struct htp_conv2d_kernel_params * kp,
-                             dma_queue * dmaq, float * stage, uint32_t n, uint32_t oh, uint32_t ow0, uint32_t G) {
+                             const struct htp_conv_geom * g, dma_queue * dmaq, float * stage, uint32_t n, uint32_t oh,
+                             uint32_t ow0, uint32_t G) {
    const struct htp_tensor * kernel = octx->src[0];
    const struct htp_tensor * input  = octx->src[1];

    const uint32_t KW  = kernel->ne[0];
-    const uint32_t KH  = kernel->ne[1];
-    const uint32_t Cin = kernel->ne[2];
+    const uint32_t KH  = g->KH;
+    const uint32_t Cin = g->Cin;
    const int32_t  IW  = input->ne[0];
    const int32_t  IH  = input->ne[1];
-    const int32_t  s1  = octx->op_params[1];
-    const int32_t  p0  = octx->op_params[2];
-    const int32_t  p1  = octx->op_params[3];
-    const int32_t  d0  = octx->op_params[4];
-    const int32_t  d1  = octx->op_params[5];
+    const int32_t  s0  = g->s0;
+    const int32_t  s1  = g->s1;
+    const int32_t  p0  = g->p0;
+    const int32_t  p1  = g->p1;
+    const int32_t  d0  = g->d0;
+    const int32_t  d1  = g->d1;
    const uint32_t SR  = kp->tiled.stage_row;

-    const int32_t Wb  = 32 * G + (KW - 1) * d0;
-    const int32_t iw0 = (int32_t) ow0 - p0;
+    const int32_t Wb  = 32 * G * s0 + (KW - 1) * d0;
+    const int32_t iw0 = (int32_t) ow0 * s0 - p0;
    const int32_t lo  = MAX(0, -iw0);
    const int32_t hi  = MIN(Wb, IW - iw0);

+    const uint32_t od = fastdiv(oh, &kp->tiled.div_ohr);
+    oh -= od * g->OHr;
+
    for (uint32_t kh = 0; kh < KH; kh++) {
-        const int32_t ih   = (int32_t) oh * s1 + (int32_t) kh * d1 - p1;
+        const int32_t kd   = fastdiv(kh, &kp->tiled.div_khr);
+        const int32_t id   = (int32_t) od * g->s2 + kd * g->d2 - g->p2;
+        const int32_t ih   = (int32_t) oh * s1 + (int32_t) (kh - kd * g->KHr) * d1 - p1;
        float *       rows = stage + (size_t) kh * SR;

-        if (ih < 0 || ih >= IH || hi <= lo) {
+        if (id < 0 || id >= (int32_t) g->ID || ih < 0 || ih >= IH || hi <= lo) {
            for (uint32_t c = 0; c < Cin; c++) {
                hvx_splat_f32_u(rows + (size_t) c * KH * SR, 0.0f, Wb);
            }
@@ -277,8 +327,8 @@ static void conv2d_hmx_stage(const struct htp_ops_context * octx, const struct h
            continue;
        }

-        const dma_addr_t src = input->data + (size_t) n * input->nb[3] + (size_t) ih * input->nb[1] + (size_t) (iw0 + lo) * sizeof(float);
-        dma_queue_push(dmaq, dma_make_data(rows + lo, src), (size_t) KH * SR * sizeof(float), input->nb[2],
+        const dma_addr_t src = input->data + (size_t) n * g->x_nb_n + (size_t) id * g->x_nb_d + (size_t) ih * input->nb[1] + (size_t) (iw0 + lo) * sizeof(float);
+        dma_queue_push(dmaq, dma_make_data(rows + lo, src), (size_t) KH * SR * sizeof(float), g->x_nb_c,
                       (size_t) (hi - lo) * sizeof(float), Cin);

        if (lo > 0 || hi < Wb) {
@@ -295,9 +345,24 @@ static void conv2d_hmx_stage(const struct htp_ops_context * octx, const struct h
    }
 }

+static inline HVX_Vector conv2d_hmx_load_strided_f32(const float * src, uint32_t s0, HVX_Vector offsets, HVX_Vector * scratch) {
+    if (s0 == 1) {
+        return hvx_vmemu(src);
+    }
+    if (s0 == 2) {
+        const HVX_Vector lo = hvx_vmemu(src);
+        const HVX_Vector hi = hvx_vmemu(src + 32);
+        return Q6_V_lo_W(Q6_W_vdeal_VVR(hi, lo, -4));
+    }
+    Q6_vgather_ARMVw(scratch, (size_t) src, 32 * s0 * sizeof(float) - 1, offsets);
+    hvx_gather_sync(scratch);
+    return *scratch;
+}
+
 // One weight-side tile column: vector i holds k rows 2i and 2i+1 interleaved over 32 pixels
 static inline void conv2d_hmx_pack_tile(__fp16 * restrict dst, const float * restrict src, uint32_t K, uint32_t KW,
-                                        uint32_t d0, uint32_t SR, uint32_t n_vec) {
+                                        uint32_t s0, uint32_t d0, uint32_t SR, uint32_t n_vec, HVX_Vector offsets,
+                                        HVX_Vector * scratch) {
    HVX_Vector * restrict out  = (HVX_Vector *) dst;
    const HVX_Vector      zero = Q6_V_vzero();
    const uint32_t        wrap = SR - KW * d0;
@@ -307,13 +372,13 @@ static inline void conv2d_hmx_pack_tile(__fp16 * restrict dst, const float * res
    uint32_t      i  = 0;

    for (; i < K / 2; i++) {
-        const HVX_Vector a = hvx_vmemu(p);
+        const HVX_Vector a = conv2d_hmx_load_strided_f32(p, s0, offsets, scratch);
        p += d0;
        if (++kw == KW) {
            kw = 0;
            p += wrap;
        }
-        const HVX_Vector b = hvx_vmemu(p);
+        const HVX_Vector b = conv2d_hmx_load_strided_f32(p, s0, offsets, scratch);
        p += d0;
        if (++kw == KW) {
            kw = 0;
@@ -322,14 +387,14 @@ static inline void conv2d_hmx_pack_tile(__fp16 * restrict dst, const float * res
        out[i] = conv2d_f32_to_f16_shuff(a, b);
    }
    if (K & 1) {
-        out[i++] = conv2d_f32_to_f16_shuff(hvx_vmemu(p), zero);
+        out[i++] = conv2d_f32_to_f16_shuff(conv2d_hmx_load_strided_f32(p, s0, offsets, scratch), zero);
    }
    for (; i < n_vec; i++) {
        out[i] = zero;
    }
 }

-// CWHN: stage[kh*SR + j*Cin + cin] = x[n][ih(kh)][ow0 - p0 + j][cin], zero outside
+// CWHN: stage[kh*SR + j*Cin + cin] = x[n][ih(kh)][ow0*s0 - p0 + j][cin], zero outside
 static void conv2d_hmx_stage_cwhn(const struct htp_ops_context * octx, const struct htp_conv2d_kernel_params * kp,
                                  dma_queue * dmaq, float * stage, uint32_t n, uint32_t oh, uint32_t ow0, uint32_t G) {
    const struct htp_tensor * kernel = octx->src[0];
@@ -340,6 +405,7 @@ static void conv2d_hmx_stage_cwhn(const struct htp_ops_context * octx, const str
    const uint32_t Cin = kernel->ne[2];
    const int32_t  IW  = input->ne[0];
    const int32_t  IH  = input->ne[1];
+    const int32_t  s0  = octx->op_params[0];
    const int32_t  s1  = octx->op_params[1];
    const int32_t  p0  = octx->op_params[2];
    const int32_t  p1  = octx->op_params[3];
@@ -347,8 +413,8 @@ static void conv2d_hmx_stage_cwhn(const struct htp_ops_context * octx, const str
    const int32_t  d1  = octx->op_params[5];
    const uint32_t SR  = kp->tiled.stage_row;

-    const int32_t Wb  = 32 * G + (KW - 1) * d0;
-    const int32_t iw0 = (int32_t) ow0 - p0;
+    const int32_t Wb  = 32 * G * s0 + (KW - 1) * d0;
+    const int32_t iw0 = (int32_t) ow0 * s0 - p0;
    const int32_t lo  = MAX(0, -iw0);
    const int32_t hi  = MIN(Wb, IW - iw0);

@@ -383,8 +449,8 @@ static void conv2d_hmx_cwhn_f16(__fp16 * restrict dst, const float * restrict sr

 // Stage the segment of col tiles starting at t, returns its number of col tiles
 static uint32_t conv2d_hmx_issue_seg(const struct htp_ops_context * octx, const struct htp_conv2d_kernel_params * kp,
-                                     dma_queue * dmaq, float * stage, uint32_t t, uint32_t t1) {
-    const uint32_t OH    = octx->dst->ne[1];
+                                     const struct htp_conv_geom * g, dma_queue * dmaq, float * stage, uint32_t t, uint32_t t1) {
+    const uint32_t OH    = g->OH;
    const uint32_t n_owt = kp->tiled.n_ow_tiles;
    const uint32_t row   = fastdiv(t, &kp->tiled.div_ow_tiles);
    const uint32_t owt   = t - row * n_owt;
@@ -395,7 +461,7 @@ static uint32_t conv2d_hmx_issue_seg(const struct htp_ops_context * octx, const
    if (kp->cwhn) {
        conv2d_hmx_stage_cwhn(octx, kp, dmaq, stage, n, oh, owt * 32, G);
    } else {
-        conv2d_hmx_stage(octx, kp, dmaq, stage, n, oh, owt * 32, G);
+        conv2d_hmx_stage(octx, kp, g, dmaq, stage, n, oh, owt * 32, G);
    }
    return G;
 }
@@ -487,7 +553,7 @@ static inline void conv2d_transpose_strided(const __fp16 * restrict base, uint32
 // the second one rotated into the upper half.
 static void conv2d_hmx_pack_tile_cwhn(__fp16 * restrict dst, const __fp16 * restrict src, HVX_Vector * restrict vin,
                                      HVX_Vector * restrict scratch, const struct htp_conv2d_kernel_params * kp,
-                                      uint32_t gpr, uint32_t Cin, uint32_t d0) {
+                                      uint32_t gpr, uint32_t Cin, uint32_t s0, uint32_t d0) {
    const uint32_t k_pad = kp->tiled.n_k_tiles * HMX_FP16_TILE_N_COLS;
    const uint32_t KG    = kp->tiled.k_group;
    const uint32_t KGp   = kp->tiled.k_group_pad;
@@ -523,15 +589,15 @@ static void conv2d_hmx_pack_tile_cwhn(__fp16 * restrict dst, const __fp16 * rest
        }

        if (va[0] == 32 && va[1] == 32 && pa[1] == pa[0] + 32) {
-            conv2d_transpose_strided(pa[0], Cin, scratch, out + 32 * b, out + 32 * b + 16);
+            conv2d_transpose_strided(pa[0], Cin * s0, scratch, out + 32 * b, out + 32 * b + 16);
            continue;
        }

        const HVX_VectorPred q0 = Q6_Q_vsetq_R(va[0] * sizeof(__fp16));
        const HVX_VectorPred q1 = Q6_Q_vsetq_R(va[1] * sizeof(__fp16));
        for (uint32_t j = 0; j < 32; j++) {
-            const HVX_Vector v0 = Q6_V_vand_QV(q0, hvx_vmemu(pa[0] + j * Cin));
-            const HVX_Vector v1 = Q6_V_vand_QV(q1, hvx_vmemu(pa[1] + j * Cin));
+            const HVX_Vector v0 = Q6_V_vand_QV(q0, hvx_vmemu(pa[0] + j * Cin * s0));
+            const HVX_Vector v1 = Q6_V_vand_QV(q1, hvx_vmemu(pa[1] + j * Cin * s0));
            vin[j] = Q6_V_vor_VV(v0, Q6_V_vror_VR(v1, 64));
        }

@@ -718,9 +784,10 @@ static void conv2d_hmx_pack(const struct htp_conv2d_hmx_context * hc, uint32_t i
    const struct htp_tensor *               kernel = octx->src[0];

    const uint32_t KW    = kernel->ne[0];
-    const uint32_t KH    = kernel->ne[1];
-    const uint32_t Cin   = kernel->ne[2];
-    const uint32_t d0    = octx->op_params[4];
+    const uint32_t KH    = hc->g->KH;
+    const uint32_t Cin   = hc->g->Cin;
+    const uint32_t s0    = (uint32_t) hc->g->s0;
+    const uint32_t d0    = (uint32_t) hc->g->d0;
    const uint32_t n_vec = kp->tiled.n_k_tiles * (HMX_FP16_TILE_N_ELMS / 64);

    dma_queue * dmaq  = octx->ctx->dma[ith];
@@ -747,10 +814,14 @@ static void conv2d_hmx_pack(const struct htp_conv2d_hmx_context * hc, uint32_t i
        return;
    }

+    // Reuse output staging after unpack has drained its DMA.
+    HVX_Vector *     gather_tmp = (HVX_Vector *) (thr + (size_t) n_stage * kp->stage_size);
+    const HVX_Vector offsets    = kp->cwhn || s0 <= 2 ? Q6_V_vzero() : hvx_vec_gather_offsets_w(s0 * sizeof(float));
+
    uint32_t ng[2] = { 0, 0 };
    uint32_t tq    = t0;
    for (uint32_t i = 0; i < n_stage && tq < t1; i++) {
-        ng[i] = conv2d_hmx_issue_seg(octx, kp, dmaq, stages[i], tq, t1);
+        ng[i] = conv2d_hmx_issue_seg(octx, kp, hc->g, dmaq, stages[i], tq, t1);
        tq   += ng[i];
    }

@@ -765,22 +836,29 @@ static void conv2d_hmx_pack(const struct htp_conv2d_hmx_context * hc, uint32_t i

        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, (uint16_t) t);
        if (kp->cwhn) {
-            const uint32_t Wb = 32 * G + (KW - 1) * d0;
+            const uint32_t Wb = 32 * G * s0 + (KW - 1) * d0;
            conv2d_hmx_cwhn_f16(stage16, stage, KH, kp->tiled.stage_row, Wb * Cin);
            for (uint32_t g = 0; g < G; g++) {
                __fp16 * tile = act + (size_t) (t - cs + g) * n_vec * 64;
-                conv2d_hmx_pack_tile_cwhn(tile, stage16 + (size_t) g * 32 * Cin, vin, scratch, kp, gpr, Cin, d0);
+                conv2d_hmx_pack_tile_cwhn(tile, stage16 + (size_t) g * 32 * Cin * s0, vin, scratch, kp, gpr, Cin, s0, d0);
            }
        } else {
            for (uint32_t g = 0; g < G; g++) {
-                __fp16 * tile = act + (size_t) (t - cs + g) * n_vec * 64;
-                conv2d_hmx_pack_tile(tile, stage + g * 32, kp->tiled.k_size, KW, d0, kp->tiled.stage_row, n_vec);
+                __fp16 *      tile   = act + (size_t) (t - cs + g) * n_vec * 64;
+                const float * pixels = stage + (size_t) g * 32 * s0;
+                if (s0 == 1) {
+                    conv2d_hmx_pack_tile(tile, pixels, kp->tiled.k_size, KW, 1, d0, kp->tiled.stage_row, n_vec, offsets, gather_tmp);
+                } else if (s0 == 2) {
+                    conv2d_hmx_pack_tile(tile, pixels, kp->tiled.k_size, KW, 2, d0, kp->tiled.stage_row, n_vec, offsets, gather_tmp);
+                } else {
+                    conv2d_hmx_pack_tile(tile, pixels, kp->tiled.k_size, KW, s0, d0, kp->tiled.stage_row, n_vec, offsets, gather_tmp);
+                }
            }
        }
        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, (uint16_t) t);

        if (tq < t1) {
-            ng[bi] = conv2d_hmx_issue_seg(octx, kp, dmaq, stage, tq, t1);
+            ng[bi] = conv2d_hmx_issue_seg(octx, kp, hc->g, dmaq, stage, tq, t1);
            tq    += ng[bi];
        }
        t += G;
@@ -794,8 +872,8 @@ static void conv2d_hmx_unpack(const struct htp_conv2d_hmx_context * hc, uint32_t
    const struct htp_tensor *               dst  = octx->dst;

    const uint32_t OW    = dst->ne[0];
-    const uint32_t OH    = dst->ne[1];
-    const uint32_t Cout  = dst->ne[2];
+    const uint32_t OH    = hc->g->OH;
+    const uint32_t Cout  = hc->g->Cout;
    const uint32_t n_owt = kp->tiled.n_ow_tiles;
    const uint32_t OS    = kp->tiled.seg_tiles * 32;

@@ -850,8 +928,8 @@ static void conv2d_hmx_unpack(const struct htp_conv2d_hmx_context * hc, uint32_t
            }
            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) t);

-            const dma_addr_t ddr = dst->data + (size_t) n * dst->nb[3] + (size_t) oh * dst->nb[1] + (size_t) co0 * dst->nb[2] + (size_t) ow0 * sizeof(float);
-            dma_queue_push(dmaq, dma_make_data(ddr, ob), dst->nb[2], (size_t) OS * sizeof(float), (size_t) nv * sizeof(float), nco);
+            const dma_addr_t ddr = dst->data + (size_t) n * hc->g->y_nb_n + (size_t) oh * dst->nb[1] + (size_t) co0 * hc->g->y_nb_c + (size_t) ow0 * sizeof(float);
+            dma_queue_push(dmaq, dma_make_data(ddr, ob), hc->g->y_nb_c, (size_t) OS * sizeof(float), (size_t) nv * sizeof(float), nco);
            n_pushed++;
            b ^= 1;
        }
@@ -904,8 +982,11 @@ static void conv2d_hmx_run_phase(struct htp_conv2d_hmx_context * hc, int32_t pac
 }

 static int conv2d_hmx_tiled(struct htp_ops_context * octx, const struct htp_conv2d_kernel_params * kp) {
+    struct htp_conv_geom g;
+    conv_geom_init(&g, octx);
+
    uint32_t row_base, n_rows;
-    conv2d_mdev_rows(octx, &row_base, &n_rows);
+    conv2d_mdev_rows(octx, g.N * g.OH, &row_base, &n_rows);
    if (n_rows == 0) {
        return HTP_STATUS_OK;
    }
@@ -923,6 +1004,7 @@ static int conv2d_hmx_tiled(struct htp_ops_context * octx, const struct htp_conv
    struct htp_conv2d_hmx_context hc = {
        .octx     = octx,
        .kp       = kp,
+        .g        = &g,
        .layout   = layout,
        .scales   = VTCM_LAYOUT_PTR(__fp16, base, layout.off_scales),
        .w_tiles  = VTCM_LAYOUT_PTR(__fp16, base, layout.off_w),
@@ -2153,7 +2235,7 @@ static int conv2d_hmx_win(struct htp_ops_context * octx, const struct htp_conv2d
    const uint32_t d1  = octx->op_params[5];

    uint32_t row_base, n_rows;
-    conv2d_mdev_rows(octx, &row_base, &n_rows);
+    conv2d_mdev_rows(octx, dst->ne[3] * dst->ne[1], &row_base, &n_rows);
    if (n_rows == 0) {
        return HTP_STATUS_OK;
    }
@@ -35,6 +35,8 @@ struct htp_conv2d_tiled_params {
    uint32_t wk_chunk;
    struct fastdiv_values div_ow_tiles;
    struct fastdiv_values div_oh;
+    struct fastdiv_values div_khr;  // kernel ne[1], unfolds conv_3d kh = kd * KHr + kh'
+    struct fastdiv_values div_ohr;  // dst ne[1], unfolds conv_3d oh = od * OHr + oh'
 };

 struct htp_conv2d_window_params {
@@ -0,0 +1,675 @@
+#pragma clang diagnostic ignored "-Wunused-variable"
+#pragma clang diagnostic ignored "-Wunused-function"
+
+#include <HAP_farf.h>
+
+#include <stdint.h>
+#include <string.h>
+#include <math.h>
+
+#include "dma-queue.h"
+#include "hex-common.h"
+#include "hex-profile.h"
+#include "htp-ctx.h"
+#include "htp-tensor.h"
+#include "hvx-utils.h"
+#include "dsv4-hc-ops.h"
+
+// ---------------------------------------------------------------------------
+// DSV4_HC_PRE
+// ---------------------------------------------------------------------------
+
+// hvx_vec_fast_sigmoid_f32_guard_2it with fewer live constants and predicates
+static inline HVX_Vector dsv4_hc_sigmoid_f32(HVX_Vector v, HVX_Vector max_exp, HVX_Vector min_exp) {
+    v = Q6_Vsf_vmin_VsfVsf(Q6_Vsf_vmax_VsfVsf(v, min_exp), max_exp);
+    v = Q6_Vqf32_vmpy_VsfVsf(v, hvx_vec_splat_f32(0.5f * 1.442695022f));
+
+    HVX_Vector in_int = Q6_Vw_equals_Vsf(Q6_Vsf_equals_Vqf32(v));
+    HVX_Vector x      = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
+    HVX_Vector xx     = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x);
+
+    HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2));
+    v1            = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
+
+    HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1));
+    v2            = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx);
+    v2            = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x);
+
+    HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
+    v3            = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
+
+    HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
+    HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
+
+    // Newton-Raphson with 2 iterations
+    HVX_Vector r = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(0x7EEEEBB3), v5);
+    for (int i = 0; i < 2; ++i) {
+        const HVX_Vector rv = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(r, v5));
+        r = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(Q6_Vqf32_vadd_VsfVsf(r, r), Q6_Vqf32_vmpy_VsfVsf(r, rv)));
+    }
+
+    return Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, r));
+}
+
+struct dsv4_hc_pre_context {
+    struct htp_ops_context * octx;
+    const struct htp_dsv4_hc_pre_kernel_params * kparams;
+    uint32_t token_start;
+    uint32_t n_chunks;
+    uint32_t total_jobs;
+    uint32_t jobs_per_thread;
+    float scale;
+    uint8_t * vtcm;
+    uint32_t dma_block_chunks;
+    uint32_t dma_row_stride;
+    uint32_t dma_buffer_size;
+};
+
+struct dsv4_hc_block {
+    uint32_t it;
+    uint32_t chunk;
+    uint32_t count;
+};
+
+static inline struct dsv4_hc_block dsv4_hc_pre_block_make(
+        const struct dsv4_hc_pre_context * dctx,
+        uint32_t job,
+        uint32_t end) {
+    const uint32_t local_token = fastdiv(job, &dctx->kparams->div_n_chunks);
+    const uint32_t chunk       = job - local_token*dctx->n_chunks;
+    const uint32_t count       = hex_smin(dctx->dma_block_chunks, hex_smin(dctx->n_chunks - chunk, end - job));
+    return (struct dsv4_hc_block) {
+        .it    = dctx->token_start + local_token,
+        .chunk = chunk,
+        .count = count,
+    };
+}
+
+static inline void dsv4_hc_pre_push_plain(
+        dma_queue * dma_q,
+        const struct dsv4_hc_pre_context * dctx,
+        const struct dsv4_hc_block * block,
+        uint8_t * spad) {
+    const struct htp_tensor * x       = dctx->octx->src[0];
+    const struct htp_tensor * weights = dctx->octx->src[1];
+    const uint32_t row_stride         = dctx->dma_row_stride;
+    const uint32_t row_size           = block->count * VLEN;
+
+    dma_queue_push(dma_q, dma_make_data(spad, x->data + (block->it*x->nb[2] + block->chunk*VLEN)), row_stride, x->nb[1], row_size, 4);
+    dma_queue_push(dma_q, dma_make_data(spad + 5*row_stride, weights->data + (block->it*weights->nb[1])), 16, 16, 16, 1);
+}
+
+static inline void dsv4_hc_pre_push_gated(
+        dma_queue * dma_q,
+        const struct dsv4_hc_pre_context * dctx,
+        const struct dsv4_hc_block * block,
+        uint8_t * spad) {
+    const struct htp_tensor * x       = dctx->octx->src[0];
+    const struct htp_tensor * g       = dctx->octx->src[1];
+    const uint32_t row_stride         = dctx->dma_row_stride;
+    const uint32_t row_size           = block->count * VLEN;
+
+    dma_queue_push(dma_q, dma_make_data(spad, x->data + (block->it*x->nb[2] + block->chunk*VLEN)), row_stride, x->nb[1], row_size, 4);
+    dma_queue_push(dma_q, dma_make_data(spad + 4*row_stride, g->data + (block->it*g->nb[2] + block->chunk*VLEN)), row_stride, g->nb[1], row_size, 4);
+}
+
+static __attribute__((noinline)) void dsv4_hc_pre_plain_block_f32(
+        const struct dsv4_hc_pre_context * dctx,
+        const struct dsv4_hc_block * block,
+        const uint8_t * x_spad,
+        uint8_t * d_spad,
+        const float * wb) {
+    const uint32_t row_stride = dctx->dma_row_stride;
+    const HVX_Vector scale_v  = hvx_vec_splat_f32(dctx->scale);
+    const HVX_Vector w0       = hvx_vec_splat_f32(wb[0]);
+    const HVX_Vector w1       = hvx_vec_splat_f32(wb[1]);
+    const HVX_Vector w2       = hvx_vec_splat_f32(wb[2]);
+    const HVX_Vector w3       = hvx_vec_splat_f32(wb[3]);
+
+    for (uint32_t c = 0; c < block->count; ++c) {
+        const uint8_t * xb = x_spad + c*VLEN;
+
+        HVX_Vector x0 = *(const HVX_Vector *) xb;
+        HVX_Vector x1 = *(const HVX_Vector *) (xb + row_stride);
+        HVX_Vector x2 = *(const HVX_Vector *) (xb + 2*row_stride);
+        HVX_Vector x3 = *(const HVX_Vector *) (xb + 3*row_stride);
+
+        HVX_Vector acc = hvx_vec_mul_f32_f32(x0, w0);
+        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(x1, w1));
+        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(x2, w2));
+        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(x3, w3));
+
+        *(HVX_Vector *) (d_spad + c*VLEN) = hvx_vec_mul_f32_f32(acc, scale_v);
+    }
+}
+
+static __attribute__((noinline)) void dsv4_hc_pre_gated_block_f32(
+        const struct dsv4_hc_pre_context * dctx,
+        const struct dsv4_hc_block * block,
+        const uint8_t * x_spad,
+        const uint8_t * g_spad,
+        uint8_t * d_spad) {
+    const uint32_t row_stride = dctx->dma_row_stride;
+    const HVX_Vector scale_v = hvx_vec_splat_f32(dctx->scale);
+    const HVX_Vector max_exp = hvx_vec_splat_f32(87.0f);
+    const HVX_Vector min_exp = hvx_vec_splat_f32(-87.0f);
+
+    for (uint32_t c = 0; c < block->count; ++c) {
+        const uint8_t * xb = x_spad + c*VLEN;
+        const uint8_t * gb = g_spad + c*VLEN;
+
+        HVX_Vector xv = *(const HVX_Vector *) xb;
+        HVX_Vector gv = dsv4_hc_sigmoid_f32(*(const HVX_Vector *) gb, max_exp, min_exp);
+        HVX_Vector acc = hvx_vec_mul_f32_f32(xv, gv);
+
+        xv = *(const HVX_Vector *) (xb + row_stride);
+        gv = dsv4_hc_sigmoid_f32(*(const HVX_Vector *) (gb + row_stride), max_exp, min_exp);
+        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(xv, gv));
+
+        xv = *(const HVX_Vector *) (xb + 2*row_stride);
+        gv = dsv4_hc_sigmoid_f32(*(const HVX_Vector *) (gb + 2*row_stride), max_exp, min_exp);
+        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(xv, gv));
+
+        xv = *(const HVX_Vector *) (xb + 3*row_stride);
+        gv = dsv4_hc_sigmoid_f32(*(const HVX_Vector *) (gb + 3*row_stride), max_exp, min_exp);
+        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(xv, gv));
+
+        *(HVX_Vector *) (d_spad + c*VLEN) = hvx_vec_mul_f32_f32(acc, scale_v);
+    }
+}
+
+static void dsv4_hc_pre_thread_plain_f32(unsigned int nth, unsigned int ith, void * data) {
+    struct dsv4_hc_pre_context * dctx = (struct dsv4_hc_pre_context *) data;
+    struct htp_ops_context *     octx = dctx->octx;
+
+    const struct htp_tensor * dst    = octx->dst;
+    const uint64_t            d_data = dst->data;
+
+    const uint32_t row_stride  = dctx->dma_row_stride;
+    const uint32_t buffer_size = dctx->dma_buffer_size;
+    const uint32_t start       = dctx->jobs_per_thread * ith;
+    const uint32_t end         = hex_smin(start + dctx->jobs_per_thread, dctx->total_jobs);
+    if (start >= end) {
+        return;
+    }
+
+    uint8_t * spad[2] = {
+        dctx->vtcm + 2*buffer_size*ith,
+        dctx->vtcm + 2*buffer_size*ith + buffer_size,
+    };
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    uint32_t prefetch_job = start;
+    for (uint32_t i = 0; i < 2 && prefetch_job < end; ++i) {
+        const struct dsv4_hc_block block = dsv4_hc_pre_block_make(dctx, prefetch_job, end);
+        // empty dst writeback keeps pop order dst, x, weights in loop
+        dma_queue_push(dma_q, dma_make_data(d_data, spad[i] + 4*row_stride), block.count*VLEN, block.count*VLEN, 0, 0);
+        dsv4_hc_pre_push_plain(dma_q, dctx, &block, spad[i]);
+        prefetch_job += block.count;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    uint32_t job  = start;
+    uint32_t curr = 0;
+    while (job < end) {
+        const struct dsv4_hc_block block = dsv4_hc_pre_block_make(dctx, job, end);
+
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+
+        const float * wb = (const float *) (spad[curr] + 5*row_stride);
+        uint8_t * d_spad = spad[curr] + 4*row_stride;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+        dsv4_hc_pre_plain_block_f32(dctx, &block, spad[curr], d_spad, wb);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+
+        const uint64_t d_dst = d_data + (block.it*dst->nb[1] + block.chunk*VLEN);
+        dma_queue_push(dma_q, dma_make_data(d_dst, d_spad), block.count*VLEN, block.count*VLEN, block.count*VLEN, 1);
+
+        if (prefetch_job < end) {
+            const struct dsv4_hc_block next = dsv4_hc_pre_block_make(dctx, prefetch_job, end);
+            dsv4_hc_pre_push_plain(dma_q, dctx, &next, spad[curr]);
+            prefetch_job += next.count;
+        }
+
+        job += block.count;
+        curr ^= 1;
+    }
+
+    dma_queue_flush(dma_q);
+
+    (void) nth;
+}
+
+static void dsv4_hc_pre_thread_gated_f32(unsigned int nth, unsigned int ith, void * data) {
+    struct dsv4_hc_pre_context * dctx = (struct dsv4_hc_pre_context *) data;
+    struct htp_ops_context *     octx = dctx->octx;
+
+    const struct htp_tensor * dst    = octx->dst;
+    const uint64_t            d_data = dst->data;
+
+    const uint32_t row_stride  = dctx->dma_row_stride;
+    const uint32_t buffer_size = dctx->dma_buffer_size;
+    const uint32_t start       = dctx->jobs_per_thread * ith;
+    const uint32_t end         = hex_smin(start + dctx->jobs_per_thread, dctx->total_jobs);
+    if (start >= end) {
+        return;
+    }
+
+    uint8_t * spad[2] = {
+        dctx->vtcm + 2*buffer_size*ith,
+        dctx->vtcm + 2*buffer_size*ith + buffer_size,
+    };
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    uint32_t prefetch_job = start;
+    for (uint32_t i = 0; i < 2 && prefetch_job < end; ++i) {
+        const struct dsv4_hc_block block = dsv4_hc_pre_block_make(dctx, prefetch_job, end);
+        // empty dst writeback keeps pop order dst, x, g in loop
+        dma_queue_push(dma_q, dma_make_data(d_data, spad[i] + 8*row_stride), block.count*VLEN, block.count*VLEN, 0, 0);
+        dsv4_hc_pre_push_gated(dma_q, dctx, &block, spad[i]);
+        prefetch_job += block.count;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    uint32_t job  = start;
+    uint32_t curr = 0;
+    while (job < end) {
+        const struct dsv4_hc_block block = dsv4_hc_pre_block_make(dctx, job, end);
+
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+
+        uint8_t * d_spad = spad[curr] + 8*row_stride;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+        dsv4_hc_pre_gated_block_f32(dctx, &block, spad[curr], spad[curr] + 4*row_stride, d_spad);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+
+        const uint64_t d_dst = d_data + (block.it*dst->nb[1] + block.chunk*VLEN);
+        dma_queue_push(dma_q, dma_make_data(d_dst, d_spad), block.count*VLEN, block.count*VLEN, block.count*VLEN, 1);
+
+        if (prefetch_job < end) {
+            const struct dsv4_hc_block next = dsv4_hc_pre_block_make(dctx, prefetch_job, end);
+            dsv4_hc_pre_push_gated(dma_q, dctx, &next, spad[curr]);
+            prefetch_job += next.count;
+        }
+
+        job += block.count;
+        curr ^= 1;
+    }
+
+    dma_queue_flush(dma_q);
+
+    (void) nth;
+}
+
+int op_dsv4_hc_pre(struct htp_ops_context * octx) {
+    const struct htp_tensor * x       = octx->src[0];
+    const struct htp_tensor * weights = octx->src[1];
+    const struct htp_tensor * dst     = octx->dst;
+
+    const struct htp_dsv4_hc_pre_kernel_params * kparams =
+        (const struct htp_dsv4_hc_pre_kernel_params *) octx->kernel_params;
+
+    const uint32_t n_threads = kparams->n_threads;
+    if (!htp_ops_context_set_n_threads_div(octx, n_threads, kparams->div_n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    uint32_t token_start     = 0;
+    uint32_t n_tokens        = x->ne[2];
+    uint32_t total_jobs      = kparams->total_jobs;
+    uint32_t jobs_per_thread = kparams->jobs_per_thread;
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_mdev_data_aligned(dst) && ((dst->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(n_tokens, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        token_start     = range.start;
+        n_tokens        = range.count;
+        total_jobs      = n_tokens * kparams->n_chunks;
+        jobs_per_thread = fastdiv(total_jobs + n_threads - 1, &kparams->div_n_threads);
+    }
+
+    if (n_tokens == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    float scale;
+    memcpy(&scale, octx->op_params, sizeof(scale));
+
+    struct dsv4_hc_pre_context dctx = {
+        .octx             = octx,
+        .kparams          = kparams,
+        .token_start      = token_start,
+        .n_chunks         = kparams->n_chunks,
+        .total_jobs       = total_jobs,
+        .jobs_per_thread  = jobs_per_thread,
+        .scale            = scale,
+        .vtcm             = (uint8_t *) octx->ctx->vtcm_base,
+        .dma_block_chunks = kparams->dma_block_chunks,
+        .dma_row_stride   = kparams->dma_row_stride,
+        .dma_buffer_size  = kparams->dma_buffer_size,
+    };
+
+    work_queue_func_t worker = NULL;
+    switch (kparams->kernel_type) {
+        case HTP_DSV4_HC_PRE_KERNEL_PLAIN: worker = dsv4_hc_pre_thread_plain_f32; break;
+        case HTP_DSV4_HC_PRE_KERNEL_GATED: worker = dsv4_hc_pre_thread_gated_f32; break;
+        default: return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    FARF(HIGH, "dsv4-hc-pre: (%ux%ux%u) -> (%ux%u) : ktype %u tokens %u:%u nth %u\n",
+         x->ne[0], x->ne[1], x->ne[2], dst->ne[0], dst->ne[1], kparams->kernel_type,
+         token_start, token_start + n_tokens, n_threads);
+
+    work_queue_run(octx->ctx->work_queue, worker, &dctx, n_threads);
+    return HTP_STATUS_OK;
+}
+
+// ---------------------------------------------------------------------------
+// DSV4_HC_POST
+// ---------------------------------------------------------------------------
+
+struct dsv4_hc_post_context {
+    struct htp_ops_context * octx;
+    const struct htp_dsv4_hc_post_kernel_params * kparams;
+    uint32_t token_start;
+    uint32_t total_jobs;
+    uint32_t jobs_per_thread;
+    uint8_t * vtcm;
+};
+
+static inline struct dsv4_hc_block dsv4_hc_post_block_make(
+        const struct dsv4_hc_post_context * dctx,
+        uint32_t job,
+        uint32_t end,
+        uint32_t block_chunks) {
+    const uint32_t local_token = fastdiv(job, &dctx->kparams->div_n_chunks);
+    const uint32_t chunk       = job - local_token*dctx->kparams->n_chunks;
+    const uint32_t count       = hex_smin(block_chunks, hex_smin(dctx->kparams->n_chunks - chunk, end - job));
+    return (struct dsv4_hc_block) {
+        .it    = dctx->token_start + local_token,
+        .chunk = chunk,
+        .count = count,
+    };
+}
+
+static __attribute__((noinline)) void dsv4_hc_post_identity_block_f32(
+        const uint8_t * xb,
+        const uint8_t * rb,
+        size_t r_stride,
+        uint8_t * db,
+        size_t d_stride,
+        uint32_t count,
+        const float * p) {
+    for (uint32_t id = 0; id < 4; ++id) {
+        const HVX_Vector pv = hvx_vec_splat_f32(p[id]);
+        const uint8_t *  rr = rb + id*r_stride;
+        uint8_t *        dd = db + id*d_stride;
+
+        for (uint32_t i = 0; i < count; ++i) {
+            const HVX_Vector xv = *(const HVX_Vector *) (xb + i*VLEN);
+            const HVX_Vector rv = *(const HVX_Vector *) (rr + i*VLEN);
+            *(HVX_Vector *) (dd + i*VLEN) = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(xv, pv), rv);
+        }
+    }
+}
+
+static __attribute__((noinline)) void dsv4_hc_post_comb_block_f32(
+        const uint8_t * xb,
+        const uint8_t * rb,
+        size_t r_stride,
+        uint8_t * db,
+        size_t d_stride,
+        uint32_t count,
+        const float * p,
+        const float * c) {
+    for (uint32_t id = 0; id < 4; ++id) {
+        const HVX_Vector pv = hvx_vec_splat_f32(p[id]);
+        const HVX_Vector c0 = hvx_vec_splat_f32(c[id]);
+        const HVX_Vector c1 = hvx_vec_splat_f32(c[4 + id]);
+        const HVX_Vector c2 = hvx_vec_splat_f32(c[8 + id]);
+        const HVX_Vector c3 = hvx_vec_splat_f32(c[12 + id]);
+        uint8_t * dd = db + id*d_stride;
+
+        for (uint32_t i = 0; i < count; ++i) {
+            const uint8_t * r = rb + i*VLEN;
+            HVX_Vector acc = hvx_vec_mul_f32_f32(*(const HVX_Vector *) (xb + i*VLEN), pv);
+            acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(*(const HVX_Vector *) r, c0));
+            acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(*(const HVX_Vector *) (r + r_stride), c1));
+            acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(*(const HVX_Vector *) (r + 2*r_stride), c2));
+            acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(*(const HVX_Vector *) (r + 3*r_stride), c3));
+            *(HVX_Vector *) (dd + i*VLEN) = acc;
+        }
+    }
+}
+
+static inline void dsv4_hc_post_push_identity(
+        dma_queue * dma_q,
+        const struct dsv4_hc_post_context * dctx,
+        const struct dsv4_hc_block * block,
+        uint8_t * spad) {
+    const struct htp_tensor * x        = dctx->octx->src[0];
+    const struct htp_tensor * residual = dctx->octx->src[1];
+    const struct htp_tensor * post     = dctx->octx->src[2];
+    const uint32_t row_stride          = dctx->kparams->dma_row_stride;
+    const uint32_t row_size            = block->count * VLEN;
+
+    dma_queue_push(dma_q, dma_make_data(spad, residual->data + (block->it*residual->nb[2] + block->chunk*VLEN)),
+                   row_stride, residual->nb[1], row_size, 4);
+    dma_queue_push(dma_q, dma_make_data(spad + 4*row_stride, x->data + (block->it*x->nb[1] + block->chunk*VLEN)),
+                   row_stride, row_size, row_size, 1);
+    dma_queue_push(dma_q, dma_make_data(spad + 9*row_stride, post->data + (block->it*post->nb[1])),
+                   16, 16, 16, 1);
+}
+
+static inline void dsv4_hc_post_push_comb(
+        dma_queue * dma_q,
+        const struct dsv4_hc_post_context * dctx,
+        const struct dsv4_hc_block * block,
+        uint8_t * spad) {
+    const struct htp_tensor * comb = dctx->octx->src[3];
+    const uint32_t row_stride      = dctx->kparams->dma_row_stride;
+
+    dsv4_hc_post_push_identity(dma_q, dctx, block, spad);
+    dma_queue_push(dma_q, dma_make_data(spad + 9*row_stride + 64, comb->data + (block->it*comb->nb[2])),
+                   16, comb->nb[1], 16, 4);
+}
+
+static void dsv4_hc_post_thread_identity_f32(unsigned int nth, unsigned int ith, void * data) {
+    struct dsv4_hc_post_context * dctx = (struct dsv4_hc_post_context *) data;
+    struct htp_ops_context *      octx = dctx->octx;
+
+    const struct htp_tensor * dst    = octx->dst;
+    const uint64_t            d_data = dst->data;
+
+    const uint32_t block_chunks = dctx->kparams->dma_block_chunks;
+    const uint32_t row_stride   = dctx->kparams->dma_row_stride;
+    const uint32_t buffer_size  = dctx->kparams->dma_buffer_size;
+    const uint32_t start        = dctx->jobs_per_thread * ith;
+    const uint32_t end          = hex_smin(start + dctx->jobs_per_thread, dctx->total_jobs);
+    if (start >= end) {
+        return;
+    }
+
+    uint8_t * spad[2] = {
+        dctx->vtcm + 2*buffer_size*ith,
+        dctx->vtcm + 2*buffer_size*ith + buffer_size,
+    };
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    uint32_t prefetch_job = start;
+    for (uint32_t i = 0; i < 2 && prefetch_job < end; ++i) {
+        const struct dsv4_hc_block block = dsv4_hc_post_block_make(dctx, prefetch_job, end, block_chunks);
+        // empty dst writeback keeps pop order dst, residual, x, post in loop
+        dma_queue_push(dma_q, dma_make_data(d_data, spad[i] + 5*row_stride), dst->nb[1], row_stride, 0, 0);
+        dsv4_hc_post_push_identity(dma_q, dctx, &block, spad[i]);
+        prefetch_job += block.count;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    uint32_t job  = start;
+    uint32_t curr = 0;
+    while (job < end) {
+        const struct dsv4_hc_block block = dsv4_hc_post_block_make(dctx, job, end, block_chunks);
+
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+
+        const float * p      = (const float *) (spad[curr] + 9*row_stride);
+        uint8_t *     d_spad = spad[curr] + 5*row_stride;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+        dsv4_hc_post_identity_block_f32(spad[curr] + 4*row_stride, spad[curr], row_stride, d_spad, row_stride, block.count, p);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+
+        const uint64_t d_dst = d_data + (block.it*dst->nb[2] + block.chunk*VLEN);
+        dma_queue_push(dma_q, dma_make_data(d_dst, d_spad), dst->nb[1], row_stride, block.count*VLEN, 4);
+
+        if (prefetch_job < end) {
+            const struct dsv4_hc_block next = dsv4_hc_post_block_make(dctx, prefetch_job, end, block_chunks);
+            dsv4_hc_post_push_identity(dma_q, dctx, &next, spad[curr]);
+            prefetch_job += next.count;
+        }
+
+        job += block.count;
+        curr ^= 1;
+    }
+
+    dma_queue_flush(dma_q);
+
+    (void) nth;
+}
+
+static void dsv4_hc_post_thread_comb_f32(unsigned int nth, unsigned int ith, void * data) {
+    struct dsv4_hc_post_context * dctx = (struct dsv4_hc_post_context *) data;
+    struct htp_ops_context *      octx = dctx->octx;
+
+    const struct htp_tensor * dst    = octx->dst;
+    const uint64_t            d_data = dst->data;
+
+    const uint32_t block_chunks = dctx->kparams->dma_block_chunks;
+    const uint32_t row_stride   = dctx->kparams->dma_row_stride;
+    const uint32_t buffer_size  = dctx->kparams->dma_buffer_size;
+    const uint32_t start        = dctx->jobs_per_thread * ith;
+    const uint32_t end          = hex_smin(start + dctx->jobs_per_thread, dctx->total_jobs);
+    if (start >= end) {
+        return;
+    }
+
+    uint8_t * spad[2] = {
+        dctx->vtcm + 2*buffer_size*ith,
+        dctx->vtcm + 2*buffer_size*ith + buffer_size,
+    };
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    uint32_t prefetch_job = start;
+    for (uint32_t i = 0; i < 2 && prefetch_job < end; ++i) {
+        const struct dsv4_hc_block block = dsv4_hc_post_block_make(dctx, prefetch_job, end, block_chunks);
+        // empty dst writeback keeps pop order dst, residual, x, post, comb in loop
+        dma_queue_push(dma_q, dma_make_data(d_data, spad[i] + 5*row_stride), dst->nb[1], row_stride, 0, 0);
+        dsv4_hc_post_push_comb(dma_q, dctx, &block, spad[i]);
+        prefetch_job += block.count;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    uint32_t job  = start;
+    uint32_t curr = 0;
+    while (job < end) {
+        const struct dsv4_hc_block block = dsv4_hc_post_block_make(dctx, job, end, block_chunks);
+
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+
+        const float * p      = (const float *) (spad[curr] + 9*row_stride);
+        const float * c      = (const float *) (spad[curr] + 9*row_stride + 64);
+        uint8_t *     d_spad = spad[curr] + 5*row_stride;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+        dsv4_hc_post_comb_block_f32(spad[curr] + 4*row_stride, spad[curr], row_stride, d_spad, row_stride, block.count, p, c);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+
+        const uint64_t d_dst = d_data + (block.it*dst->nb[2] + block.chunk*VLEN);
+        dma_queue_push(dma_q, dma_make_data(d_dst, d_spad), dst->nb[1], row_stride, block.count*VLEN, 4);
+
+        if (prefetch_job < end) {
+            const struct dsv4_hc_block next = dsv4_hc_post_block_make(dctx, prefetch_job, end, block_chunks);
+            dsv4_hc_post_push_comb(dma_q, dctx, &next, spad[curr]);
+            prefetch_job += next.count;
+        }
+
+        job += block.count;
+        curr ^= 1;
+    }
+
+    dma_queue_flush(dma_q);
+
+    (void) nth;
+}
+
+int op_dsv4_hc_post(struct htp_ops_context * octx) {
+    const struct htp_tensor * x        = octx->src[0];
+    const struct htp_tensor * residual = octx->src[1];
+    const struct htp_tensor * dst      = octx->dst;
+
+    const struct htp_dsv4_hc_post_kernel_params * kparams =
+        (const struct htp_dsv4_hc_post_kernel_params *) octx->kernel_params;
+
+    const uint32_t n_threads = kparams->n_threads;
+    if (!htp_ops_context_set_n_threads_div(octx, n_threads, kparams->div_n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    uint32_t token_start     = 0;
+    uint32_t n_tokens        = x->ne[1];
+    uint32_t total_jobs      = kparams->total_jobs;
+    uint32_t jobs_per_thread = kparams->jobs_per_thread;
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_mdev_data_aligned(dst) && ((dst->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(n_tokens, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        token_start     = range.start;
+        n_tokens        = range.count;
+        total_jobs      = n_tokens * kparams->n_chunks;
+        jobs_per_thread = fastdiv(total_jobs + n_threads - 1, &kparams->div_n_threads);
+    }
+
+    if (n_tokens == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    struct dsv4_hc_post_context dctx = {
+        .octx            = octx,
+        .kparams         = kparams,
+        .token_start     = token_start,
+        .total_jobs      = total_jobs,
+        .jobs_per_thread = jobs_per_thread,
+        .vtcm            = (uint8_t *) octx->ctx->vtcm_base,
+    };
+
+    work_queue_func_t worker = NULL;
+    switch (kparams->kernel_type) {
+        case HTP_DSV4_HC_POST_KERNEL_IDENTITY: worker = dsv4_hc_post_thread_identity_f32; break;
+        case HTP_DSV4_HC_POST_KERNEL_COMB:     worker = dsv4_hc_post_thread_comb_f32;     break;
+        default: return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    FARF(HIGH, "dsv4-hc-post: (%ux%u) (%ux%ux%u) -> (%ux%ux%u) : ktype %u tokens %u:%u nth %u\n",
+         x->ne[0], x->ne[1], residual->ne[0], residual->ne[1], residual->ne[2],
+         dst->ne[0], dst->ne[1], dst->ne[2], kparams->kernel_type,
+         token_start, token_start + n_tokens, n_threads);
+
+    work_queue_run(octx->ctx->work_queue, worker, &dctx, n_threads);
+    return HTP_STATUS_OK;
+}
@@ -0,0 +1,77 @@
+#ifndef HTP_DSV4_HC_OPS_H
+#define HTP_DSV4_HC_OPS_H
+
+#include <stdint.h>
+
+#include "hex-fastdiv.h"
+#include "htp-ops.h"
+
+#define HTP_DSV4_HC_CHUNK_ELEMS 32
+#define HTP_DSV4_HC_MAX_CHUNKS 128
+#define HTP_DSV4_HC_PRE_GATED_ROWS 9
+#define HTP_DSV4_HC_PRE_PLAIN_ROWS 6
+
+enum htp_dsv4_hc_pre_kernel_type {
+    HTP_DSV4_HC_PRE_KERNEL_PLAIN = 0,
+    HTP_DSV4_HC_PRE_KERNEL_GATED = 1,
+};
+
+struct htp_dsv4_hc_pre_kernel_params {
+    uint32_t kernel_type;
+    uint32_t n_threads;
+    uint32_t n_chunks;
+    uint32_t total_jobs;
+    uint32_t jobs_per_thread;
+
+    uint32_t dma_block_chunks;
+    uint32_t dma_row_stride;
+    uint32_t dma_buffer_size;
+    uint32_t vtcm_size;
+
+    struct fastdiv_values div_n_chunks;
+    struct fastdiv_values div_n_threads;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_dsv4_hc_pre_kernel_params) <= 128, "htp_dsv4_hc_pre_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_dsv4_hc_pre_kernel_params) <= 128, "htp_dsv4_hc_pre_kernel_params is too large for kernel_params blob");
+#endif
+
+// dma buffer holds residual rows 0..3, x at row 4, dst rows 5..8, post/comb at row 9
+#define HTP_DSV4_HC_POST_ROWS 10
+
+enum htp_dsv4_hc_post_kernel_type {
+    HTP_DSV4_HC_POST_KERNEL_IDENTITY = 0,
+    HTP_DSV4_HC_POST_KERNEL_COMB     = 1,
+};
+
+struct htp_dsv4_hc_post_kernel_params {
+    uint32_t kernel_type;
+    uint32_t n_threads;
+    uint32_t n_chunks;
+    uint32_t total_jobs;
+    uint32_t jobs_per_thread;
+    uint32_t has_comb;
+
+    uint32_t dma_block_chunks;
+    uint32_t dma_row_stride;
+    uint32_t dma_buffer_size;
+    uint32_t vtcm_size;
+
+    struct fastdiv_values div_n_chunks;
+    struct fastdiv_values div_n_threads;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_dsv4_hc_post_kernel_params) <= 128, "htp_dsv4_hc_post_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_dsv4_hc_post_kernel_params) <= 128, "htp_dsv4_hc_post_kernel_params is too large for kernel_params blob");
+#endif
+
+struct htp_ops_context;
+
+int op_dsv4_hc_pre(struct htp_ops_context * octx);
+int op_dsv4_hc_post(struct htp_ops_context * octx);
+
+#endif // HTP_DSV4_HC_OPS_H
@@ -186,53 +186,6 @@ static void swiglu_clamp_f32(const float * restrict src0,
 static const float GELU_COEF_A     = 0.044715f;
 static const float SQRT_2_OVER_PI  = 0.79788456080286535587989211986876f;

-static inline HVX_Vector hvx_vec_fast_sigmoid_f32_2it(HVX_Vector v) {
-    v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
-    v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3));
-
-    HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
-    HVX_Vector x      = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
-    HVX_Vector xx     = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x);
-
-    HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2));
-    v1            = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
-
-    HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1));
-    v2            = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx);
-    v2            = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x);
-
-    HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
-    v3            = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
-
-    HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
-    HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
-
-    // Newton-Raphson with 2 iterations
-    HVX_Vector two_sf = hvx_vec_splat_f32(2.0f);
-    HVX_Vector i_sf   = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(0x7EEEEBB3), v5);
-    HVX_Vector r_qf   = Q6_Vqf32_vmpy_VsfVsf(
-        i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
-    r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
-        r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
-    HVX_Vector res = Q6_Vsf_equals_Vqf32(r_qf);
-
-    res = Q6_Vqf32_vmpy_VsfVsf(v3, res);
-
-    return Q6_Vsf_equals_Vqf32(res);
-}
-
-static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard_2it(HVX_Vector v,
-                                                            HVX_Vector one,
-                                                            HVX_Vector max_exp,
-                                                            HVX_Vector min_exp) {
-    const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v);
-    const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp);
-
-    HVX_Vector out = hvx_vec_fast_sigmoid_f32_2it(v);
-    out            = Q6_V_vmux_QVV(pred_max, out, one);
-    return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero());
-}
-
 static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
    assert((unsigned long) dst  % 128 == 0);
    assert((unsigned long) src0 % 128 == 0);
@@ -131,21 +131,34 @@ static inline void htp_ops_context_reset(struct htp_ops_context * octx, struct h
    octx->n_threads_div = ctx->n_threads_div;
 }

-static inline bool htp_ops_context_set_n_threads(struct htp_ops_context * octx, uint32_t n_threads) {
+static inline bool htp_ops_context_set_n_threads_div(
+        struct htp_ops_context * octx,
+        uint32_t n_threads,
+        struct fastdiv_values n_threads_div) {
    if (n_threads == 0 || n_threads > octx->ctx->n_threads) {
        return false;
    }

-    if (n_threads != octx->n_threads) {
-        octx->n_threads = n_threads;
-        octx->n_threads_div = n_threads == octx->ctx->n_threads
-            ? octx->ctx->n_threads_div
-            : init_fastdiv_values(n_threads);
-    }
+    octx->n_threads     = n_threads;
+    octx->n_threads_div = n_threads_div;

    return true;
 }

+static inline bool htp_ops_context_set_n_threads(struct htp_ops_context * octx, uint32_t n_threads) {
+    if (n_threads == 0 || n_threads > octx->ctx->n_threads) {
+        return false;
+    }
+    if (n_threads == octx->n_threads) {
+        return true;
+    }
+
+    const struct fastdiv_values n_threads_div = n_threads == octx->ctx->n_threads
+        ? octx->ctx->n_threads_div
+        : init_fastdiv_values(n_threads);
+    return htp_ops_context_set_n_threads_div(octx, n_threads, n_threads_div);
+}
+
 static inline void htp_ops_context_set_status(struct htp_ops_context * octx, int status) {
    if (status > HTP_STATUS_OK && octx->status == HTP_STATUS_OK) {
        octx->status = status;
@@ -174,6 +187,7 @@ int op_argsort(struct htp_ops_context * octx);
 int op_top_k(struct htp_ops_context * octx);
 int op_ssm_conv(struct htp_ops_context * octx);
 int op_cumsum(struct htp_ops_context * octx);
+int op_top_k_mask(struct htp_ops_context * octx);
 int op_fill(struct htp_ops_context * octx);
 int op_concat(struct htp_ops_context * octx);
 int op_diag(struct htp_ops_context * octx);
@@ -189,5 +203,6 @@ int op_pool_1d(struct htp_ops_context * octx);
 int op_allreduce(struct htp_ops_context * octx);
 int op_roll(struct htp_ops_context * octx);
 int op_upscale(struct htp_ops_context * octx);
+int op_lightning_indexer(struct htp_ops_context * octx);

 #endif /* HTP_CTX_H */
@@ -51,6 +51,7 @@ enum htp_op_code {
    HTP_OP_MUL_MAT_ADD,
    HTP_OP_RMS_NORM,
    HTP_OP_RMS_NORM_MUL,
+    HTP_OP_RMS_NORM_MUL_ADD,
    HTP_OP_UNARY_SILU,
    HTP_OP_UNARY_GELU,
    HTP_OP_UNARY_SIGMOID,
@@ -117,6 +118,11 @@ enum htp_op_code {
    HTP_OP_UNARY_CEIL,
    HTP_OP_UNARY_SGN,
    HTP_OP_UPSCALE,
+    HTP_OP_CONV_3D,
+    HTP_OP_DSV4_HC_PRE,
+    HTP_OP_DSV4_HC_POST,
+    HTP_OP_TOP_K_MASK,
+    HTP_OP_LIGHTNING_INDEXER,

    HTP_OP_INVALID
 };
@@ -141,6 +141,71 @@ static inline void hvx_fast_rms_norm_mul_f32(const uint8_t * restrict src,
    }
 }

+static inline void hvx_fast_rms_norm_mul_add_f32(const uint8_t * restrict src,
+                                                 const uint8_t * restrict weight,
+                                                 const uint8_t * restrict add,
+                                                 uint8_t * restrict dst,
+                                                 const int num_elems,
+                                                 float     epsilon,
+                                                 float     scale) {
+    const HVX_Vector * restrict v_src    = (const HVX_Vector *) src;
+    const HVX_Vector * restrict v_weight = (const HVX_Vector *) weight;
+    const HVX_Vector * restrict v_add    = (const HVX_Vector *) add;
+    HVX_Vector * restrict v_dst          = (HVX_Vector *) dst;
+
+    const int nvec = num_elems / VLEN_FP32;
+    const int nloe = num_elems % VLEN_FP32;
+
+    HVX_Vector sum_v     = Q6_V_vsplat_R(0x00000000);
+    HVX_Vector epsilon_v = hvx_vec_splat_f32(epsilon);
+
+    #pragma unroll(4)
+    for (int i = 0; i < nvec; i++) {
+        HVX_Vector v1 = v_src[i];
+        HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
+        sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
+    }
+
+    if (nloe > 0) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+        HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+        HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, v1);
+        sum_v = Q6_Vqf32_vadd_Vqf32Vqf32(sum_v, v2);
+    }
+
+    sum_v = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(sum_v));
+
+    HVX_Vector t_v            = hvx_vec_splat_f32((float) num_elems);
+    HVX_Vector denom_v        = hvx_vec_inverse_f32(t_v);
+    HVX_Vector mean_v         = Q6_Vqf32_vmpy_VsfVsf(sum_v, denom_v);
+    HVX_Vector mean_epsilon_v = Q6_Vqf32_vadd_Vqf32Vsf(mean_v, epsilon_v);
+
+    HVX_Vector scale_v = hvx_vec_rsqrt_f32(Q6_Vsf_equals_Vqf32(mean_epsilon_v));
+    if (scale != 0.0f) {
+        scale_v = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(scale_v, hvx_vec_splat_f32(scale)));
+    }
+
+    #pragma unroll(4)
+    for (int i = 0; i < nvec; i++) {
+        HVX_Vector v1 = v_src[i];
+        HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
+        HVX_Vector v3 = Q6_Vsf_equals_Vqf32(v2);
+        HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, v_weight[i]));
+        v_dst[i] = hvx_vec_add_f32_f32(v4, v_add[i]);
+    }
+
+    if (nloe > 0) {
+        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
+        HVX_Vector v1 = Q6_V_vand_QV(bmask, v_src[nvec]);
+        HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(v1, scale_v);
+        HVX_Vector v3 = Q6_Vsf_equals_Vqf32(v2);
+        HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, v_weight[nvec]));
+        HVX_Vector res_v = hvx_vec_add_f32_f32(v4, v_add[nvec]);
+
+        hvx_vec_store_a(&v_dst[nvec], nloe * 4, res_v);
+    }
+}
+
 static inline void hvx_fast_norm_f32(const uint8_t * restrict src,
                                     uint8_t * restrict dst,
                                     const int num_elems,
@@ -52,6 +52,53 @@ static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard(HVX_Vector v,
    return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero());
 }

+static inline HVX_Vector hvx_vec_fast_sigmoid_f32_2it(HVX_Vector v) {
+    v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
+    v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3));
+
+    HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
+    HVX_Vector x      = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
+    HVX_Vector xx     = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x);
+
+    HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2));
+    v1            = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
+
+    HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1));
+    v2            = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx);
+    v2            = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x);
+
+    HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
+    v3            = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
+
+    HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
+    HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
+
+    // Newton-Raphson with 2 iterations
+    HVX_Vector two_sf = hvx_vec_splat_f32(2.0f);
+    HVX_Vector i_sf   = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(0x7EEEEBB3), v5);
+    HVX_Vector r_qf   = Q6_Vqf32_vmpy_VsfVsf(
+        i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
+    r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
+        r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
+    HVX_Vector res = Q6_Vsf_equals_Vqf32(r_qf);
+
+    res = Q6_Vqf32_vmpy_VsfVsf(v3, res);
+
+    return Q6_Vsf_equals_Vqf32(res);
+}
+
+static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard_2it(HVX_Vector v,
+                                                            HVX_Vector one,
+                                                            HVX_Vector max_exp,
+                                                            HVX_Vector min_exp) {
+    const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v);
+    const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp);
+
+    HVX_Vector out = hvx_vec_fast_sigmoid_f32_2it(v);
+    out            = Q6_V_vmux_QVV(pred_max, out, one);
+    return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero());
+}
+
 static inline HVX_Vector hvx_vec_tanh_f32(HVX_Vector x) {
    // tanh(x) = 2 * sigmoid(2x) - 1
    HVX_Vector two = hvx_vec_splat_f32(2.0f);
@@ -0,0 +1,213 @@
+#ifndef HTP_LIGHTNING_INDEXER_OPS_H
+#define HTP_LIGHTNING_INDEXER_OPS_H
+
+#include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
+
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#define HTP_LID_KV_BLOCK   64 // kv rows per block (HVX and HMX Bc)
+#define HTP_LID_TOKENS_MAX 32 // max tokens per Q tile (HVX splat path)
+#define HTP_LID_HMX_TOKENS_MAX 512 // HMX can keep more tokens; VTCM decides
+#define HTP_LID_HMX_KV_BLOCK_MAX 2048
+#define HTP_LID_HMX_MAX_QSETS    8    // streams in one HMX pipeline run
+#define HTP_LID_N_BUF      4  // K / mask / output buffers per thread (two blocks computed, two in flight)
+
+#ifndef HMX_FP16_TILE_N_COLS
+#define HMX_FP16_TILE_N_COLS 32
+#endif
+#ifndef HMX_FP16_TILE_SIZE
+#define HMX_FP16_TILE_SIZE 2048
+#endif
+
+enum htp_lid_kernel_type {
+    HTP_LID_KERNEL_HVX = 0,
+    HTP_LID_KERNEL_HMX = 1,
+};
+
+struct htp_lid_kernel_params {
+    uint8_t  kernel_type;       // enum htp_lid_kernel_type
+    uint32_t n_threads;
+    uint32_t n_tokens_tile;     // tokens per Q tile
+    uint32_t n_kv_blocks;       // HVX: per stream, HMX: per run of n_qsets sets
+    uint32_t blocks_per_thread;
+    uint32_t vtcm_size;
+    uint32_t k_row_size;        // K row bytes
+    uint32_t kv_block;          // HMX: kv rows per block (multiple of 64)
+    uint32_t n_qsets;           // HMX: Q sets kept in VTCM (streams run in one pipeline)
+    uint32_t n_qbuf;            // HMX: Q buffers, 2: the next run's Q is prepped during the current run
+    uint32_t kv_seg;            // HMX: rows per set, n_kv padded to 64
+    uint32_t n_tt;              // HMX: token tiles per stream
+
+    struct fastdiv_values div_n_head;
+    struct fastdiv_values div_m_ne3;  // mask streams (broadcast)
+    struct fastdiv_values div_kv_seg; // HMX
+    struct fastdiv_values div_n_tt;   // HMX
+    struct fastdiv_values div_n_ch;   // HMX: channel tiles per Q set
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_lid_kernel_params) <= 128, "htp_lid_kernel_params is too large");
+#else
+_Static_assert(sizeof(struct htp_lid_kernel_params) <= 128, "htp_lid_kernel_params is too large");
+#endif
+
+// VTCM layout shared by host (fit check) and device (buffer offsets).
+// Shared regions: Q splats (one f16 vector per q element) and head weight splats for the current token tile.
+// Per-thread regions: Q prep staging, K / mask / output buffers, two transposed K blocks and transpose scratch.
+struct htp_lid_vtcm_layout {
+    size_t off_q;
+    size_t off_w;
+    size_t off_qs;
+    size_t off_k;
+    size_t off_kt;
+    size_t off_m;
+    size_t off_o;
+
+    size_t size_q_row;   // f16 Q row, padded
+    size_t size_qs_row;  // f32 Q staging row, padded
+    size_t size_k_row;   // f16 K row, padded
+    size_t size_k_raw_row;
+    size_t size_m_row;   // f16 mask row of one kv block, padded
+    size_t size_o_row;   // f32 output row of one kv block, padded
+
+    size_t size_qs;      // per thread
+    size_t size_k_buf;   // one K buffer
+    size_t size_kt;      // per thread: two transposed K blocks, two 64x64 scratch tiles, Q8_0 block scales
+    size_t size_m_buf;   // one mask buffer
+    size_t size_o_buf;   // one output buffer
+
+    size_t total_bytes;
+};
+
+static inline void htp_lid_vtcm_layout_build(struct htp_lid_vtcm_layout * L, size_t D, size_t n_head, size_t n_tok, size_t n_threads, size_t k_row_size) {
+    L->size_q_row  = hex_round_up(D * sizeof(uint16_t), 128);
+    L->size_qs_row = hex_round_up(D * sizeof(float), 128);
+    L->size_k_row  = hex_round_up(D * sizeof(uint16_t), 128);
+    L->size_k_raw_row = L->size_k_row > hex_round_up(k_row_size, 128) ? L->size_k_row : hex_round_up(k_row_size, 128);
+    L->size_m_row  = hex_round_up(HTP_LID_KV_BLOCK * sizeof(uint16_t), 128);
+    L->size_o_row  = hex_round_up(HTP_LID_KV_BLOCK * sizeof(float), 128);
+
+    L->size_qs    = 2 * (n_head * L->size_qs_row + hex_round_up(n_head * sizeof(float), 128)) + L->size_q_row;
+    L->size_k_buf = HTP_LID_KV_BLOCK * L->size_k_raw_row;
+    L->size_kt    = (2 * hex_round_up(D, HTP_LID_KV_BLOCK) + 2 * HTP_LID_KV_BLOCK + D / 32) * 128;
+    L->size_m_buf = n_tok * L->size_m_row;
+    L->size_o_buf = n_tok * L->size_o_row;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_q,  n_tok * n_head * D * 128);
+    VTCM_LAYOUT_ALLOC(off, off_w,  n_tok * n_head * 128);
+    VTCM_LAYOUT_ALLOC(off, off_qs, n_threads * L->size_qs);
+    VTCM_LAYOUT_ALLOC(off, off_k,  n_threads * HTP_LID_N_BUF * L->size_k_buf);
+    VTCM_LAYOUT_ALLOC(off, off_kt, n_threads * L->size_kt);
+    VTCM_LAYOUT_ALLOC(off, off_m,  n_threads * HTP_LID_N_BUF * L->size_m_buf);
+    VTCM_LAYOUT_ALLOC(off, off_o,  n_threads * HTP_LID_N_BUF * L->size_o_buf);
+
+    L->total_bytes = off;
+}
+
+static inline size_t htp_lid_compute_vtcm_usage(size_t D, size_t n_head, size_t n_tok, size_t n_threads, size_t k_row_size) {
+    struct htp_lid_vtcm_layout L;
+    htp_lid_vtcm_layout_build(&L, D, n_head, n_tok, n_threads, k_row_size);
+    return L.total_bytes;
+}
+
+// HMX path. K is the activation (kv rows), Q the weights (one output channel per (token, head)).
+// mm1: S = relu(K Q^T) * |w| per channel (HMX convert scale + pos clamp)
+// mm2: O = S G, G[(t,h), t] = sign(w[t,h]). O is [kv, token], f16.
+// Q weights, G, scales and w are per Q set (one per stream when all tokens fit one tile).
+struct htp_lid_hmx_vtcm_layout {
+    size_t off_qw;       // Q weight tiles [n_qsets][n_ch][n_dot]
+    size_t off_g;        // G weight tiles [n_qsets][n_ch]
+    size_t off_sc;       // per channel scales, 256 bytes per channel tile, [n_qsets][n_ch]
+    size_t off_sc_id;    // identity scales
+    size_t off_s;        // S tiles of one block and token tile, [n_rt][n_sg]
+    size_t off_w;        // w of the token tile, [n_qsets][n_tok][n_head] f32
+    size_t off_qs;       // per thread: f32 and f16 staging for 32 Q rows, gather scratch
+    size_t off_k_raw[2]; // K rows, [Bc][size_k_row]
+    size_t off_k_t[2];   // K tiles [n_rt][n_dot]
+    size_t off_o[2];     // O tiles [n_rt][n_u]
+    size_t off_m[4];     // mask rows [n_tok][size_m_row]
+    size_t off_out[2];   // f32 output rows [n_tok][size_o_row]
+
+    size_t size_qs_row;  // f32 Q row, padded
+    size_t size_qh_row;  // f16 Q row, padded
+    size_t size_k_row;   // f16 K row, padded
+    size_t size_k_raw_row;
+    size_t size_m_row;
+    size_t size_o_row;
+    size_t size_qs;      // per thread
+    size_t size_w;       // per Q set
+
+    uint32_t n_ch;       // channel tiles: ceil(n_tok * n_head / 32)
+    uint32_t n_u;        // token tiles: ceil(n_tok / 32)
+    uint32_t n_dot;      // D / 32
+    uint32_t n_rt;       // kv row tiles per block: Bc / 32
+    uint32_t n_sg;       // S tiles per (row tile, token tile): min(n_head, n_ch)
+
+    size_t total_bytes;
+};
+
+static inline void htp_lid_hmx_vtcm_layout_build(struct htp_lid_hmx_vtcm_layout * L,
+                                                  size_t D, size_t n_head, size_t n_tok, size_t Bc, size_t n_qsets, size_t n_threads, size_t k_row_size) {
+    L->n_ch  = (uint32_t) ((n_tok * n_head + 31) / 32);
+    L->n_u   = (uint32_t) ((n_tok + 31) / 32);
+    L->n_dot = (uint32_t) (D / 32);
+    L->n_rt  = (uint32_t) (Bc / 32);
+    L->n_sg  = (uint32_t) (n_head < L->n_ch ? n_head : L->n_ch);
+
+    L->size_qs_row = hex_round_up(D * sizeof(float), 128);
+    L->size_qh_row = hex_round_up(D * sizeof(uint16_t), 128);
+    L->size_k_row  = hex_round_up(D * sizeof(uint16_t), 128);
+    L->size_k_raw_row = L->size_k_row > hex_round_up(k_row_size, 128) ? L->size_k_row : hex_round_up(k_row_size, 128);
+    L->size_m_row  = hex_round_up(Bc * sizeof(uint16_t), 128);
+    L->size_o_row  = hex_round_up(Bc * sizeof(float), 128);
+    L->size_qs     = 2 * 32 * L->size_qs_row + 32 * L->size_qh_row + hex_round_up(Bc / 64 * 128, 128);
+    L->size_w      = hex_round_up(n_tok * n_head * sizeof(float), 128);
+
+    const size_t tile = HMX_FP16_TILE_SIZE;
+
+    // tile regions first: they need 2 KB alignment
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_qw,     n_qsets * L->n_ch * L->n_dot * tile);
+    VTCM_LAYOUT_ALLOC(off, off_g,      n_qsets * L->n_ch * tile);
+    VTCM_LAYOUT_ALLOC(off, off_s,      (size_t) L->n_rt * L->n_sg * tile);
+    VTCM_LAYOUT_ALLOC(off, off_k_t[0], (size_t) L->n_rt * L->n_dot * tile);
+    VTCM_LAYOUT_ALLOC(off, off_k_t[1], (size_t) L->n_rt * L->n_dot * tile);
+    VTCM_LAYOUT_ALLOC(off, off_o[0],   (size_t) L->n_rt * L->n_u * tile);
+    VTCM_LAYOUT_ALLOC(off, off_o[1],   (size_t) L->n_rt * L->n_u * tile);
+    VTCM_LAYOUT_ALLOC(off, off_sc,     n_qsets * L->n_ch * 256);
+    VTCM_LAYOUT_ALLOC(off, off_sc_id,  256);
+    VTCM_LAYOUT_ALLOC(off, off_w,      n_qsets * L->size_w);
+    VTCM_LAYOUT_ALLOC(off, off_qs,     n_threads * L->size_qs);
+    for (int i = 0; i < 2; i++) {
+        VTCM_LAYOUT_ALLOC(off, off_k_raw[i], Bc * L->size_k_raw_row);
+    }
+    for (int i = 0; i < 4; i++) {
+        VTCM_LAYOUT_ALLOC(off, off_m[i], n_tok * L->size_m_row);
+    }
+    for (int i = 0; i < 2; i++) {
+        VTCM_LAYOUT_ALLOC(off, off_out[i], n_tok * L->size_o_row);
+    }
+
+    L->total_bytes = off;
+}
+
+static inline size_t htp_lid_hmx_compute_vtcm_usage(size_t D, size_t n_head, size_t n_tok, size_t Bc, size_t n_qsets, size_t n_threads, size_t k_row_size) {
+    struct htp_lid_hmx_vtcm_layout L;
+    htp_lid_hmx_vtcm_layout_build(&L, D, n_head, n_tok, Bc, n_qsets, n_threads, k_row_size);
+    return L.total_bytes;
+}
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* HTP_LIGHTNING_INDEXER_OPS_H */
@@ -34,6 +34,7 @@
 #include "work-queue.h"
 #include "hex-profile.h"
 #include "allreduce-ops.h"
+#include "dsv4-hc-ops.h"
 #include "htp-fence.h"

 #define HMX_QUEUE_CAPACITY     128
@@ -930,6 +931,7 @@ static int execute_op(struct htp_ops_context * octx) {
        case HTP_OP_NORM:
        case HTP_OP_RMS_NORM:
        case HTP_OP_RMS_NORM_MUL:
+        case HTP_OP_RMS_NORM_MUL_ADD:
        case HTP_OP_SCALE:
        case HTP_OP_CLAMP:
        case HTP_OP_LEAKY_RELU:
@@ -1007,6 +1009,9 @@ static int execute_op(struct htp_ops_context * octx) {
        case HTP_OP_CUMSUM:
            return op_cumsum(octx);

+        case HTP_OP_TOP_K_MASK:
+            return op_top_k_mask(octx);
+
        case HTP_OP_FILL:
            return op_fill(octx);

@@ -1026,6 +1031,7 @@ static int execute_op(struct htp_ops_context * octx) {
            return op_im2col_3d(octx);

        case HTP_OP_CONV_2D:
+        case HTP_OP_CONV_3D:
            return op_conv_2d(octx);

        case HTP_OP_CONV_2D_DW:
@@ -1040,12 +1046,21 @@ static int execute_op(struct htp_ops_context * octx) {
        case HTP_OP_POOL_1D:
            return op_pool_1d(octx);

+        case HTP_OP_DSV4_HC_PRE:
+            return op_dsv4_hc_pre(octx);
+
+        case HTP_OP_DSV4_HC_POST:
+            return op_dsv4_hc_post(octx);
+
        case HTP_OP_CONCAT:
            return op_concat(octx);

        case HTP_OP_GATED_DELTA_NET:
            return op_gated_delta_net(octx);

+        case HTP_OP_LIGHTNING_INDEXER:
+            return op_lightning_indexer(octx);
+
        case HTP_OP_TRI:
            return op_unary(octx);

@@ -0,0 +1,665 @@
+#pragma clang diagnostic ignored "-Wunused-variable"
+#pragma clang diagnostic ignored "-Wunused-function"
+#pragma clang diagnostic ignored "-Wunused-but-set-variable"
+
+#include <HAP_farf.h>
+#include <HAP_perf.h>
+#include <stdatomic.h>
+
+#define GGML_COMMON_DECL_C
+#include "ggml-common.h"
+#include "hex-common.h"
+#include "hex-profile.h"
+#include "htp-ctx.h"
+#include "htp-ops.h"
+#include "htp-tensor.h"
+#include "htp-vtcm.h"
+#include "hvx-types.h"
+#include "hvx-utils.h"
+#include "dma-queue.h"
+#include "top-k-mask-ops.h"
+
+#define TOP_K_MASK_NINF_F32 0xff800000u // f32 score compare
+
+struct htp_top_k_mask_context {
+    struct htp_ops_context *                    octx;
+    const struct htp_top_k_mask_kernel_params * kparams;
+
+    uint32_t rows_per_thread;
+    uint32_t row0, row1;           // rows this device handles
+    uint32_t off_extra, off_score; // parts inside one row input slot
+    uint32_t col0, ncols_out;      // PICKS: columns this device writes back
+    atomic_uint barrier;           // PICKS: arrivals, barrier k waits for k * n_threads
+
+    uint8_t * blocks;
+    uint8_t * mask;
+    uint8_t * rowin;
+    uint8_t * tiles;
+    uint8_t * mtiles;
+    uint8_t * scratch;
+};
+
+// DMA queue pop with a trace event around the wait
+static inline void top_k_mask_pop(dma_queue * dma_q, struct htp_thread_trace * tr, uint16_t info) {
+    htp_trace_event_start(tr, HTP_TRACE_EVT_DMA, info);
+    dma_queue_pop(dma_q);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_DMA, info);
+}
+
+// Lanes 0..31 as words, in registers: inclusive scan of ones, minus one
+static inline HVX_Vector top_k_mask_lane_iota(void) {
+    const HVX_Vector zero = Q6_V_vzero();
+    const HVX_Vector one  = Q6_V_vsplat_R(1);
+    HVX_Vector       v    = one;
+    v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 4));
+    v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 8));
+    v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 16));
+    v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 32));
+    v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 64));
+    return Q6_Vw_vsub_VwVw(v, one);
+}
+
+static __attribute__((noinline)) void top_k_mask_fill(uint8_t * restrict dst, uint32_t nbytes, uint32_t value) {
+    const HVX_Vector v    = Q6_V_vsplat_R(value);
+    HVX_Vector *     vdst = (HVX_Vector *) dst;
+    const uint32_t   nvec = nbytes / VLEN;
+
+    #pragma unroll(4)
+    for (uint32_t i = 0; i < nvec; i++) {
+        vdst[i] = v;
+    }
+}
+
+#define TOP_K_MASK_DEAD 0x7fffffffu // cell value of a dead lane: c - c0 >=u ncols for every tile, signed max
+
+// Byte offsets of the cells in [c0, c0 + ncols) relative to c0, 0 for the others
+static inline HVX_Vector top_k_mask_keep_off(HVX_Vector c, HVX_Vector v_c0, HVX_Vector v_ncols, const bool is_f16) {
+    const HVX_Vector cr = Q6_Vw_vsub_VwVw(c, v_c0);
+    return Q6_Vw_vasl_VwR(Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VuwVuw(v_ncols, cr), cr, Q6_V_vzero()), is_f16 ? 1 : 2);
+}
+
+// Mask value of each cell lane in [c0, c0 + ncols) into tmp (predicated gather), returns the scatter offsets
+static inline HVX_Vector top_k_mask_keep_gather(HVX_Vector c, const uint8_t * restrict mrow, HVX_Vector v_c0, HVX_Vector v_ncols,
+                                                uint32_t mu, HVX_Vector * restrict tmp, const bool is_f16) {
+    const HVX_VectorPred q   = Q6_Q_vcmp_gt_VuwVuw(v_ncols, Q6_Vw_vsub_VwVw(c, v_c0));
+    const HVX_Vector     off = top_k_mask_keep_off(c, v_c0, v_ncols, is_f16);
+
+    // Halfword data with word offsets: both halves of the pair hold the same offsets, the word predicate covers both halfwords of a lane
+    if (is_f16) {
+        Q6_vgather_AQRMWw(tmp, q, (size_t) mrow, mu, Q6_W_vcombine_VV(off, off));
+    } else {
+        Q6_vgather_AQRMVw(tmp, q, (size_t) mrow, mu, off);
+    }
+    return off;
+}
+
+static inline void top_k_mask_keep_scatter(HVX_Vector c, HVX_Vector off, uint8_t * restrict tile, HVX_Vector v_c0, HVX_Vector v_ncols,
+                                           uint32_t mu, HVX_Vector val, const bool is_f16) {
+    const HVX_VectorPred q = Q6_Q_vcmp_gt_VuwVuw(v_ncols, Q6_Vw_vsub_VwVw(c, v_c0));
+    if (is_f16) {
+        Q6_vscatter_QRMWwV(q, (size_t) tile, mu, Q6_W_vcombine_VV(off, off), val);
+    } else {
+        Q6_vscatter_QRMVwV(q, (size_t) tile, mu, off, val);
+    }
+}
+
+// Keep 4 cell vectors (dead lanes = TOP_K_MASK_DEAD): tile[c - c0] = mrow[c - c0] for c in [c0, c0 + ncols). One gather sync.
+// flag: no mask, tile[c - c0] = all ones (keep flag)
+static inline void top_k_mask_keep4(HVX_Vector c0v, HVX_Vector c1v, HVX_Vector c2v, HVX_Vector c3v, uint8_t * restrict tile,
+                                    const uint8_t * restrict mrow, HVX_Vector v_c0, HVX_Vector v_ncols, uint32_t mu,
+                                    HVX_Vector * restrict tmp, const bool is_f16, const bool flag) {
+    if (flag) {
+        const HVX_Vector ones = Q6_V_vsplat_R(0xffffffffu);
+        top_k_mask_keep_scatter(c0v, top_k_mask_keep_off(c0v, v_c0, v_ncols, is_f16), tile, v_c0, v_ncols, mu, ones, is_f16);
+        top_k_mask_keep_scatter(c1v, top_k_mask_keep_off(c1v, v_c0, v_ncols, is_f16), tile, v_c0, v_ncols, mu, ones, is_f16);
+        top_k_mask_keep_scatter(c2v, top_k_mask_keep_off(c2v, v_c0, v_ncols, is_f16), tile, v_c0, v_ncols, mu, ones, is_f16);
+        top_k_mask_keep_scatter(c3v, top_k_mask_keep_off(c3v, v_c0, v_ncols, is_f16), tile, v_c0, v_ncols, mu, ones, is_f16);
+        return;
+    }
+    const HVX_Vector o0 = top_k_mask_keep_gather(c0v, mrow, v_c0, v_ncols, mu, tmp + 0, is_f16);
+    const HVX_Vector o1 = top_k_mask_keep_gather(c1v, mrow, v_c0, v_ncols, mu, tmp + 1, is_f16);
+    const HVX_Vector o2 = top_k_mask_keep_gather(c2v, mrow, v_c0, v_ncols, mu, tmp + 2, is_f16);
+    const HVX_Vector o3 = top_k_mask_keep_gather(c3v, mrow, v_c0, v_ncols, mu, tmp + 3, is_f16);
+    hvx_gather_sync(tmp + 3);
+    top_k_mask_keep_scatter(c0v, o0, tile, v_c0, v_ncols, mu, tmp[0], is_f16);
+    top_k_mask_keep_scatter(c1v, o1, tile, v_c0, v_ncols, mu, tmp[1], is_f16);
+    top_k_mask_keep_scatter(c2v, o2, tile, v_c0, v_ncols, mu, tmp[2], is_f16);
+    top_k_mask_keep_scatter(c3v, o3, tile, v_c0, v_ncols, mu, tmp[3], is_f16);
+}
+
+// Pick vector j: picks below j1 and n_pick are valid, returns them with dead lanes marked. v_rem = j1 - lane index.
+static inline HVX_Vector top_k_mask_load_picks(const int32_t * restrict tk, uint32_t j, HVX_Vector v_rem, HVX_Vector v_npick) {
+    const HVX_Vector     p_raw = hvx_vmem(tk + j);
+    // lane + j <u j1 is the tail, p <u n_pick drops p < 0 and p >= n_pick
+    const HVX_VectorPred valid = Q6_Q_vcmp_gtand_QVuwVuw(Q6_Q_vcmp_gt_VwVw(v_rem, Q6_V_vsplat_R(j)), v_npick, p_raw);
+    return Q6_V_vmux_QVV(valid, p_raw, Q6_V_vsplat_R(TOP_K_MASK_DEAD));
+}
+
+// Scatter the kept cells of picks [j0, j1) (and the extras) into the tile of columns [c0, c0 + ncols). slot holds the row inputs, mrow the mask columns.
+// Gathers are issued in batches with one sync per batch: tmp[0..3] scores/blocks, tmp[4] score with blocks, tmp[5..8] mask values.
+static inline __attribute__((always_inline)) void top_k_mask_scatter_tile(
+        const struct htp_top_k_mask_context * tctx, const uint8_t * restrict slot, const uint8_t * restrict mrow, uint8_t * restrict tile,
+        HVX_Vector * restrict tmp, uint32_t c0, uint32_t ncols, uint32_t j0, uint32_t j1, bool do_extra,
+        const bool is_f16, const bool flag) {
+    const struct htp_top_k_mask_kernel_params * kp = tctx->kparams;
+
+    const int32_t * restrict tk  = (const int32_t *) slot;
+    const int32_t * restrict ex  = (const int32_t *) (slot + tctx->off_extra);
+    const uint8_t * restrict sc  = slot + tctx->off_score;
+    const uint8_t * restrict blk = tctx->blocks;
+
+    const bool     has_score  = kp->flags & HTP_TOP_K_MASK_HAS_SCORE;
+    const bool     has_blocks = kp->flags & HTP_TOP_K_MASK_HAS_BLOCKS;
+    const uint32_t n_blk      = kp->n_blk;
+    const uint32_t n_extra    = do_extra ? kp->n_extra : 0;
+    const uint32_t mu         = ncols * kp->elem_size - 1;
+    const uint32_t sc_mu      = kp->score_bytes - 1;
+    const uint32_t blk_mu     = kp->blocks_bytes - 1;
+    const uint32_t nb1_h      = ((uint32_t) kp->blocks_nb1 << 16) | kp->blocks_nb1;
+
+    const HVX_Vector v_lane  = top_k_mask_lane_iota(); // in the HVX_COMP scope of the caller
+    const HVX_Vector v_c0    = Q6_V_vsplat_R(c0);
+    const HVX_Vector v_ncols = Q6_V_vsplat_R(ncols);
+    const HVX_Vector v_npick = Q6_V_vsplat_R(kp->n_pick);
+    const HVX_Vector v_pmax  = Q6_V_vsplat_R(kp->n_pick - 1); // dead picks gather the last entry, every gather stays in its region
+    const HVX_Vector v_rem   = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(j1), v_lane);
+    const HVX_Vector v_ninf  = Q6_V_vsplat_R(TOP_K_MASK_NINF_F32);
+    const HVX_Vector v_dead  = Q6_V_vsplat_R(TOP_K_MASK_DEAD);
+
+    HVX_Vector * restrict mtmp = tmp + 5;
+
+    if (has_blocks) {
+        for (uint32_t j = j0; j < j1; j += VLEN_FP32) {
+            HVX_Vector       cell = top_k_mask_load_picks(tk, j, v_rem, v_npick);
+            const HVX_Vector p    = Q6_Vw_vmin_VwVw(cell, v_pmax);
+            if (has_score) {
+                Q6_vgather_ARMVw(tmp + 4, (size_t) sc, sc_mu, Q6_Vw_vasl_VwR(p, 2));
+            }
+            const HVX_Vector boff = Q6_Vw_vmpyi_VwRh(p, nb1_h);
+            for (uint32_t k = 0; k < n_blk; k += 4) {
+                // blocks k..k+3 of every pick; the missing ones (k + i >= n_blk) are dead
+                const uint32_t   nk = MIN(n_blk - k, 4);
+                const HVX_Vector bo = Q6_Vw_vadd_VwVw(boff, Q6_V_vsplat_R(k * 4));
+                for (uint32_t i = 0; i < nk; i++) {
+                    Q6_vgather_ARMVw(tmp + i, (size_t) blk, blk_mu, Q6_Vw_vadd_VwVw(bo, Q6_V_vsplat_R(i * 4)));
+                }
+                hvx_gather_sync(tmp + nk - 1);
+                if (has_score && k == 0) {
+                    // IEEE compare: NaN and -inf scores are dead
+                    cell = Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VsfVsf(tmp[4], v_ninf), cell, v_dead);
+                }
+                const HVX_VectorPred dead = Q6_Q_vcmp_eq_VwVw(cell, v_dead);
+                top_k_mask_keep4(Q6_V_vmux_QVV(dead, v_dead, tmp[0]),
+                                 nk > 1 ? Q6_V_vmux_QVV(dead, v_dead, tmp[1]) : v_dead,
+                                 nk > 2 ? Q6_V_vmux_QVV(dead, v_dead, tmp[2]) : v_dead,
+                                 nk > 3 ? Q6_V_vmux_QVV(dead, v_dead, tmp[3]) : v_dead,
+                                 tile, mrow, v_c0, v_ncols, mu, mtmp, is_f16, flag);
+            }
+        }
+    } else {
+        // 4 pick vectors per batch, the vectors at or past j1 are dead and not loaded
+        for (uint32_t j = j0; j < j1; j += 4 * VLEN_FP32) {
+            HVX_Vector c0v = top_k_mask_load_picks(tk, j, v_rem, v_npick);
+            HVX_Vector c1v = j + 1 * VLEN_FP32 < j1 ? top_k_mask_load_picks(tk, j + 1 * VLEN_FP32, v_rem, v_npick) : v_dead;
+            HVX_Vector c2v = j + 2 * VLEN_FP32 < j1 ? top_k_mask_load_picks(tk, j + 2 * VLEN_FP32, v_rem, v_npick) : v_dead;
+            HVX_Vector c3v = j + 3 * VLEN_FP32 < j1 ? top_k_mask_load_picks(tk, j + 3 * VLEN_FP32, v_rem, v_npick) : v_dead;
+            if (has_score) {
+                Q6_vgather_ARMVw(tmp + 0, (size_t) sc, sc_mu, Q6_Vw_vasl_VwR(Q6_Vw_vmin_VwVw(c0v, v_pmax), 2));
+                Q6_vgather_ARMVw(tmp + 1, (size_t) sc, sc_mu, Q6_Vw_vasl_VwR(Q6_Vw_vmin_VwVw(c1v, v_pmax), 2));
+                Q6_vgather_ARMVw(tmp + 2, (size_t) sc, sc_mu, Q6_Vw_vasl_VwR(Q6_Vw_vmin_VwVw(c2v, v_pmax), 2));
+                Q6_vgather_ARMVw(tmp + 3, (size_t) sc, sc_mu, Q6_Vw_vasl_VwR(Q6_Vw_vmin_VwVw(c3v, v_pmax), 2));
+                hvx_gather_sync(tmp + 3);
+                c0v = Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VsfVsf(tmp[0], v_ninf), c0v, v_dead);
+                c1v = Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VsfVsf(tmp[1], v_ninf), c1v, v_dead);
+                c2v = Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VsfVsf(tmp[2], v_ninf), c2v, v_dead);
+                c3v = Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VsfVsf(tmp[3], v_ninf), c3v, v_dead);
+            }
+            top_k_mask_keep4(c0v, c1v, c2v, c3v, tile, mrow, v_c0, v_ncols, mu, mtmp, is_f16, flag);
+        }
+    }
+
+    const HVX_Vector v_rem_ex = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(n_extra), v_lane);
+    for (uint32_t j = 0; j < n_extra; j += VLEN_FP32) {
+        const HVX_VectorPred tail = Q6_Q_vcmp_gt_VwVw(v_rem_ex, Q6_V_vsplat_R(j));
+        top_k_mask_keep4(Q6_V_vmux_QVV(tail, hvx_vmem(ex + j), v_dead), v_dead, v_dead, v_dead, tile, mrow, v_c0, v_ncols, mu, mtmp, is_f16, flag);
+    }
+
+    // release the scatters before DMA reads the tile
+    hvx_gather_sync(tile);
+}
+
+static __attribute__((noinline)) void top_k_mask_scatter_tile_f32(const struct htp_top_k_mask_context * tctx, const uint8_t * restrict slot,
+                                        const uint8_t * restrict mrow, uint8_t * restrict tile, HVX_Vector * restrict tmp,
+                                        uint32_t c0, uint32_t ncols, uint32_t j0, uint32_t j1, bool do_extra) {
+    top_k_mask_scatter_tile(tctx, slot, mrow, tile, tmp, c0, ncols, j0, j1, do_extra, false, false);
+}
+
+static __attribute__((noinline)) void top_k_mask_scatter_tile_f16(const struct htp_top_k_mask_context * tctx, const uint8_t * restrict slot,
+                                        const uint8_t * restrict mrow, uint8_t * restrict tile, HVX_Vector * restrict tmp,
+                                        uint32_t c0, uint32_t ncols, uint32_t j0, uint32_t j1, bool do_extra) {
+    top_k_mask_scatter_tile(tctx, slot, mrow, tile, tmp, c0, ncols, j0, j1, do_extra, true, false);
+}
+
+// PICKS: keep flags (all ones) of the picks [j0, j1) into the flag tile of the whole row, the element size picks the flag width
+static __attribute__((noinline)) void top_k_mask_flag_tile_f32(const struct htp_top_k_mask_context * tctx, const uint8_t * restrict slot,
+                                        uint8_t * restrict flags, HVX_Vector * restrict tmp, uint32_t j0, uint32_t j1, bool do_extra) {
+    top_k_mask_scatter_tile(tctx, slot, NULL, flags, tmp, 0, tctx->kparams->n_kv, j0, j1, do_extra, false, true);
+}
+
+static __attribute__((noinline)) void top_k_mask_flag_tile_f16(const struct htp_top_k_mask_context * tctx, const uint8_t * restrict slot,
+                                        uint8_t * restrict flags, HVX_Vector * restrict tmp, uint32_t j0, uint32_t j1, bool do_extra) {
+    top_k_mask_scatter_tile(tctx, slot, NULL, flags, tmp, 0, tctx->kparams->n_kv, j0, j1, do_extra, true, true);
+}
+
+// out = flag ? out : -inf over nbytes (a multiple of 128) of the row tile
+static __attribute__((noinline)) void top_k_mask_select(uint8_t * restrict out, const uint8_t * restrict flags, uint32_t nbytes, uint32_t ninf) {
+    const HVX_Vector     v_ninf = Q6_V_vsplat_R(ninf);
+    HVX_Vector *         vout   = (HVX_Vector *) out;
+    const HVX_Vector *   vflag  = (const HVX_Vector *) flags;
+    const uint32_t       nvec   = nbytes / VLEN;
+
+    #pragma unroll(4)
+    for (uint32_t i = 0; i < nvec; i++) {
+        vout[i] = Q6_V_vxor_VV(v_ninf, Q6_V_vand_VV(Q6_V_vxor_VV(vout[i], v_ninf), vflag[i]));
+    }
+}
+
+// Push the row inputs of flat row r into slot: top_k, extra, score. Returns the push count.
+static inline uint32_t top_k_mask_push_rowin(const struct htp_top_k_mask_context * tctx, dma_queue * dma_q, uint32_t r, uint8_t * slot) {
+    const struct htp_ops_context *              octx  = tctx->octx;
+    const struct htp_top_k_mask_kernel_params * kp    = tctx->kparams;
+    const struct htp_tensor *                   top_k = octx->src[1];
+    const struct htp_tensor *                   score = octx->src[2];
+    const struct htp_tensor *                   extra = octx->src[4];
+
+    const uint32_t tk_bytes = kp->n_top * sizeof(int32_t);
+    dma_queue_push(dma_q, dma_make_data(slot, top_k->data + r * kp->u.rows.rs_top), tk_bytes, tk_bytes, tk_bytes, 1);
+    uint32_t n = 1;
+
+    if (kp->flags & HTP_TOP_K_MASK_HAS_EXTRA) {
+        const uint32_t ex_bytes = kp->n_extra * sizeof(int32_t);
+        dma_queue_push(dma_q, dma_make_data(slot + tctx->off_extra, extra->data + r * kp->u.rows.rs_extra), ex_bytes, ex_bytes, ex_bytes, 1);
+        n++;
+    }
+
+    if (kp->flags & HTP_TOP_K_MASK_HAS_SCORE) {
+        dma_queue_push(dma_q, dma_make_data(slot + tctx->off_score, score->data + r * kp->u.rows.rs_score), kp->score_bytes, kp->score_bytes, kp->score_bytes, 1);
+        n++;
+    }
+
+    return n;
+}
+
+static inline struct fastdiv_values top_k_mask_fd(const struct htp_top_k_mask_kernel_params * kp, uint32_t k) {
+    const struct fastdiv_values fd = { kp->u.rows.m.dec.mp[k], kp->u.rows.m.dec.l[k] };
+    return fd;
+}
+
+// Mask row of flat row r: (i1, i2, i3) of dst, each taken modulo the mask dims, rows dense above the mask ne1
+static inline uint32_t top_k_mask_mask_row(const struct htp_top_k_mask_kernel_params * kp, uint32_t r) {
+    const uint16_t *            ne  = kp->u.rows.m.dec.ne;
+    const struct fastdiv_values fd1 = top_k_mask_fd(kp, 0);
+    const struct fastdiv_values fd2 = top_k_mask_fd(kp, 1);
+    const struct fastdiv_values fm1 = top_k_mask_fd(kp, 2);
+    const struct fastdiv_values fm2 = top_k_mask_fd(kp, 3);
+    const struct fastdiv_values fm3 = top_k_mask_fd(kp, 4);
+
+    const uint32_t q1 = fastdiv(r, &fd1);
+    const uint32_t i1 = r - q1 * ne[0];
+    const uint32_t i3 = fastdiv(q1, &fd2);
+    const uint32_t i2 = q1 - i3 * ne[1];
+
+    const uint32_t m1 = fastmodulo(i1, ne[2], &fm1);
+    const uint32_t m2 = fastmodulo(i2, ne[3], &fm2);
+    const uint32_t m3 = fastmodulo(i3, ne[4], &fm3);
+    return (m3 * ne[3] + m2) * ne[2] + m1;
+}
+
+// Push the inputs of item (row r, col tile ct): row inputs on the first tile of a row, then the mask tile. Returns the push count.
+static inline uint32_t top_k_mask_push_item(const struct htp_top_k_mask_context * tctx, dma_queue * dma_q, uint32_t r, uint32_t ct,
+                                            uint8_t * slot, uint8_t * mslot) {
+    const struct htp_ops_context *              octx = tctx->octx;
+    const struct htp_top_k_mask_kernel_params * kp   = tctx->kparams;
+    const struct htp_tensor *                   mask = octx->src[0];
+
+    const uint32_t c0 = ct * kp->tile_cols;
+
+    uint32_t n = 0;
+    if (ct == 0) {
+        n += top_k_mask_push_rowin(tctx, dma_q, r, slot);
+    }
+
+    if (mslot) {
+        const uint32_t m_off = top_k_mask_mask_row(kp, r) * kp->u.rows.mask_nb1 + c0 * kp->elem_size;
+        const uint32_t nb    = MIN(kp->tile_cols, kp->n_kv - c0) * kp->elem_size;
+        dma_queue_push(dma_q, dma_make_data(mslot, mask->data + m_off), nb, nb, nb, 1);
+        n++;
+    }
+
+    return n;
+}
+
+static inline void top_k_mask_next(uint32_t * r, uint32_t * ct, uint32_t n_col_tiles) {
+    if (++(*ct) == n_col_tiles) {
+        *ct = 0;
+        ++(*r);
+    }
+}
+
+// ROWS: the thread count follows the device rows (mdev), so the slice of an n byte copy is sized here
+static inline uint32_t top_k_mask_slice(const struct htp_ops_context * octx, unsigned int nth, uint32_t n) {
+    return hex_round_up(fastdiv(n + nth - 1, &octx->n_threads_div), HTP_TOP_K_MASK_UNIT);
+}
+
+// Push this thread's slice (128B aligned) of a contiguous n byte copy. Returns the push count.
+static inline uint32_t top_k_mask_push_slice(dma_queue * dma_q, unsigned int ith, dma_addr_t dst, dma_addr_t src, uint32_t n, uint32_t slice) {
+    const uint32_t b0    = ith * slice;
+    if (b0 >= n) {
+        return 0;
+    }
+    const uint32_t nb = MIN(slice, n - b0);
+    dma_queue_push(dma_q, dma_make_data(dst + b0, src + b0), nb, nb, nb, 1);
+    return 1;
+}
+
+// Pass 1 (ROWS): load the blocks table and the resident mask row, split across threads. Fill the tiles of this thread with -inf meanwhile.
+static void top_k_mask_preload(unsigned int nth, unsigned int ith, void * data) {
+    struct htp_top_k_mask_context *             tctx  = (struct htp_top_k_mask_context *) data;
+    const struct htp_top_k_mask_kernel_params * kp    = tctx->kparams;
+    struct htp_ops_context *                    octx  = tctx->octx;
+    dma_queue *                                 dma_q = octx->ctx->dma[ith];
+    struct htp_thread_trace *                   tr    = &octx->ctx->trace[ith];
+
+    uint32_t n_pushed = 0;
+
+    if (kp->blocks_bytes) {
+        n_pushed += top_k_mask_push_slice(dma_q, ith, (dma_addr_t) tctx->blocks, octx->src[3]->data, kp->blocks_bytes, top_k_mask_slice(octx, nth, kp->blocks_bytes));
+    }
+    if (kp->flags & HTP_TOP_K_MASK_MASK_RESIDENT) {
+        const uint32_t mask_bytes = kp->u.rows.m.mask_bytes;
+        n_pushed += top_k_mask_push_slice(dma_q, ith, (dma_addr_t) tctx->mask, octx->src[0]->data, mask_bytes, top_k_mask_slice(octx, nth, mask_bytes));
+    }
+
+    const uint32_t tile_bytes = kp->tile_units * HTP_TOP_K_MASK_UNIT;
+
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ith);
+    top_k_mask_fill(tctx->tiles + ith * 2 * tile_bytes, 2 * tile_bytes, kp->ninf);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ith);
+
+    for (uint32_t i = 0; i < n_pushed; i++) {
+        top_k_mask_pop(dma_q, tr, (uint16_t) i);
+    }
+}
+
+// Pass 2 (ROWS): depth-2 DMA pipeline over the items of rows [row_first, row_last), item = (row, col tile).
+// The row inputs alternate between 2 slots by row.
+static void top_k_mask_run_items(const struct htp_top_k_mask_context * tctx, unsigned int ith, uint32_t row_first, uint32_t row_last) {
+    struct htp_ops_context *                    octx  = tctx->octx;
+    const struct htp_top_k_mask_kernel_params * kp    = tctx->kparams;
+    const struct htp_tensor *                   dst   = octx->dst;
+    dma_queue *                                 dma_q = octx->ctx->dma[ith];
+    struct htp_thread_trace *                   tr    = &octx->ctx->trace[ith];
+
+    const bool     is_f16    = kp->flags & HTP_TOP_K_MASK_F16;
+    const uint32_t elem_size = kp->elem_size;
+
+    const uint32_t rowin_bytes = kp->rowin_units * HTP_TOP_K_MASK_UNIT;
+    const uint32_t tile_bytes  = kp->tile_units * HTP_TOP_K_MASK_UNIT;
+    const uint32_t tile_cols   = kp->tile_cols;
+    const uint32_t n_col_tiles = kp->n_col_tiles;
+    const uint32_t n_items     = (row_last - row_first) * n_col_tiles;
+
+    // 2 row input slots per thread, alternating by row
+    uint8_t * rowin  = tctx->rowin + ith * 2 * rowin_bytes;
+    uint8_t * tiles  = tctx->tiles + ith * 2 * tile_bytes;
+    uint8_t * mtiles = (kp->flags & HTP_TOP_K_MASK_MASK_RESIDENT) ? NULL : tctx->mtiles + ith * 2 * tile_bytes;
+
+    HVX_Vector * tmp = (HVX_Vector *) (tctx->scratch + ith * HTP_TOP_K_MASK_SCRATCH_BYTES);
+
+    const dma_addr_t dst_data = dst->data;
+
+    uint32_t n_in[2] = { 0, 0 };
+
+    // prefetch cursor
+    uint32_t pr = row_first, pct = 0;
+
+    for (uint32_t it = 0; it < MIN(n_items, 2); it++) {
+        // Dummy dst writeback to establish queue ordering
+        dma_queue_push(dma_q, dma_make_data(dst_data, tiles + it * tile_bytes), 0, 0, 0, 0);
+        n_in[it] = top_k_mask_push_item(tctx, dma_q, pr, pct, rowin + ((pr - row_first) & 1) * rowin_bytes, mtiles ? mtiles + it * tile_bytes : NULL);
+        top_k_mask_next(&pr, &pct, n_col_tiles);
+    }
+
+    uint32_t r = row_first, ct = 0;
+
+    for (uint32_t it = 0; it < n_items; it++) {
+        const uint32_t s    = it & 1;
+        uint8_t *      tile = tiles + s * tile_bytes;
+
+        // writeback of item it - 2 (or the dummy), then the tile is free again
+        top_k_mask_pop(dma_q, tr, (uint16_t) it);
+        if (it >= 2) {
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) it);
+            top_k_mask_fill(tile, tile_bytes, kp->ninf);
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) it);
+        }
+
+        for (uint32_t i = 0; i < n_in[s]; i++) {
+            top_k_mask_pop(dma_q, tr, (uint16_t) it);
+        }
+
+        const uint32_t c0    = ct * tile_cols;
+        const uint32_t ncols = MIN(tile_cols, kp->n_kv - c0);
+
+        const uint8_t * slot = rowin + ((r - row_first) & 1) * rowin_bytes;
+        const uint8_t * mrow = mtiles ? mtiles + s * tile_bytes : tctx->mask + c0 * elem_size;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) it);
+        if (is_f16) {
+            top_k_mask_scatter_tile_f16(tctx, slot, mrow, tile, tmp, c0, ncols, 0, kp->n_top, true);
+        } else {
+            top_k_mask_scatter_tile_f32(tctx, slot, mrow, tile, tmp, c0, ncols, 0, kp->n_top, true);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) it);
+
+        const uint32_t d_off = r * kp->u.rows.rs_dst + c0 * elem_size;
+        const uint32_t nb    = ncols * elem_size;
+        dma_queue_push(dma_q, dma_make_data(dst_data + d_off, tile), nb, nb, nb, 1);
+
+        if (it + 2 < n_items) {
+            n_in[s] = top_k_mask_push_item(tctx, dma_q, pr, pct, rowin + ((pr - row_first) & 1) * rowin_bytes, mtiles ? mtiles + s * tile_bytes : NULL);
+            top_k_mask_next(&pr, &pct, n_col_tiles);
+        }
+        top_k_mask_next(&r, &ct, n_col_tiles);
+    }
+
+    const uint32_t n_wb = MIN(n_items, 2);
+    for (uint32_t i = 0; i < n_wb; i++) {
+        top_k_mask_pop(dma_q, tr, (uint16_t) i);
+    }
+}
+
+// Pass 2 (ROWS): whole rows split across threads, the col tiles of a row in order
+static void top_k_mask_thread_rows(unsigned int nth, unsigned int ith, void * data) {
+    const struct htp_top_k_mask_context * tctx = (const struct htp_top_k_mask_context *) data;
+
+    const uint32_t ir0 = tctx->row0 + ith * tctx->rows_per_thread;
+    const uint32_t ir1 = MIN(ir0 + tctx->rows_per_thread, tctx->row1);
+    if (ir0 >= ir1) {
+        return;
+    }
+
+    top_k_mask_run_items(tctx, ith, ir0, ir1);
+}
+
+static inline void top_k_mask_barrier(atomic_uint * b, uint32_t target, struct htp_thread_trace * tr, uint16_t info) {
+    htp_trace_event_start(tr, HTP_TRACE_EVT_FENCE, info);
+    atomic_fetch_add_explicit(b, 1, memory_order_release);
+    while (atomic_load_explicit(b, memory_order_acquire) < target) {
+        hex_pause();
+    }
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_FENCE, info);
+}
+
+// PICKS, one pass over n_virt = kp->n_threads virtual workers (host slices); real worker ith runs the virtual workers ith, ith + nth, ...
+// A: shared preload and a slice of the flag clear per virtual worker, the mask slice of this worker into the row tile;
+// barrier; B: scatter keep flags of the picks of each virtual worker; barrier; C: select mask or -inf in this worker's slice and write it back
+static void top_k_mask_thread_picks(unsigned int nth, unsigned int ith, void * data) {
+    struct htp_top_k_mask_context *             tctx  = (struct htp_top_k_mask_context *) data;
+    const struct htp_top_k_mask_kernel_params * kp    = tctx->kparams;
+    struct htp_ops_context *                    octx  = tctx->octx;
+    dma_queue *                                 dma_q = octx->ctx->dma[ith];
+    struct htp_thread_trace *                   tr    = &octx->ctx->trace[ith];
+
+    const struct htp_tensor * top_k = octx->src[1];
+    const struct htp_tensor * score = octx->src[2];
+    const struct htp_tensor * extra = octx->src[4];
+
+    uint8_t *      slot       = tctx->rowin;
+    const uint32_t tile_bytes = kp->tile_units * HTP_TOP_K_MASK_UNIT;
+    uint8_t *      out        = tctx->tiles;
+    uint8_t *      flags      = tctx->tiles + tile_bytes;
+
+    const uint32_t n_virt = kp->n_threads;
+
+    // this worker's slice of the device columns: the mask goes straight into the row tile, it is popped after the scan.
+    // The device share of the columns depends on mdev, so this slice is sized here.
+    const uint32_t out_b0 = tctx->col0 * kp->elem_size;
+    const uint32_t out_nb = tctx->ncols_out * kp->elem_size;
+    const uint32_t wslice = hex_round_up(fastdiv(out_nb + nth - 1, &octx->n_threads_div), 128);
+    const uint32_t wb0    = out_b0 + ith * wslice;
+    const uint32_t wnb    = ith * wslice < out_nb ? MIN(wslice, out_nb - ith * wslice) : 0;
+
+    const uint32_t fslice = kp->u.picks.flags_slice_u * HTP_TOP_K_MASK_UNIT;
+
+    for (uint32_t v = ith; v < n_virt; v += nth) {
+        uint32_t n_pushed = 0;
+        if (kp->blocks_bytes) {
+            n_pushed += top_k_mask_push_slice(dma_q, v, (dma_addr_t) tctx->blocks, octx->src[3]->data, kp->blocks_bytes, kp->u.picks.blocks_slice_u * HTP_TOP_K_MASK_UNIT);
+        }
+        n_pushed += top_k_mask_push_slice(dma_q, v, (dma_addr_t) slot, top_k->data, kp->n_top * sizeof(int32_t), kp->u.picks.top_slice_u * HTP_TOP_K_MASK_UNIT);
+        if (kp->flags & HTP_TOP_K_MASK_HAS_EXTRA) {
+            n_pushed += top_k_mask_push_slice(dma_q, v, (dma_addr_t) (slot + tctx->off_extra), extra->data, kp->n_extra * sizeof(int32_t), kp->u.picks.extra_slice_u * HTP_TOP_K_MASK_UNIT);
+        }
+        if (kp->flags & HTP_TOP_K_MASK_HAS_SCORE) {
+            n_pushed += top_k_mask_push_slice(dma_q, v, (dma_addr_t) (slot + tctx->off_score), score->data, kp->score_bytes, kp->u.picks.score_slice_u * HTP_TOP_K_MASK_UNIT);
+        }
+
+        // the mask slice follows the last preload in the queue, so the pops below leave it in flight
+        if (wnb && v + nth >= n_virt) {
+            dma_queue_push(dma_q, dma_make_data(out + wb0, octx->src[0]->data + wb0), wnb, wnb, wnb, 1);
+        }
+
+        const uint32_t fb0 = v * fslice;
+        if (fb0 < tile_bytes) {
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) v);
+            top_k_mask_fill(flags + fb0, MIN(fslice, tile_bytes - fb0), 0); // clear the flags
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) v);
+        }
+
+        for (uint32_t i = 0; i < n_pushed; i++) {
+            top_k_mask_pop(dma_q, tr, (uint16_t) i);
+        }
+    }
+
+    top_k_mask_barrier(&tctx->barrier, nth, tr, 0);
+
+    HVX_Vector * tmp = (HVX_Vector *) (tctx->scratch + ith * HTP_TOP_K_MASK_SCRATCH_BYTES);
+
+    for (uint32_t v = ith; v < n_virt; v += nth) {
+        const uint32_t j0       = MIN(v * kp->u.picks.picks_per_thread, kp->n_top);
+        const uint32_t j1       = MIN(j0 + kp->u.picks.picks_per_thread, kp->n_top);
+        const bool     do_extra = v == n_virt - 1;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) v);
+        if (kp->flags & HTP_TOP_K_MASK_F16) {
+            top_k_mask_flag_tile_f16(tctx, slot, flags, tmp, j0, j1, do_extra);
+        } else {
+            top_k_mask_flag_tile_f32(tctx, slot, flags, tmp, j0, j1, do_extra);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) v);
+    }
+
+    top_k_mask_barrier(&tctx->barrier, 2 * nth, tr, 1);
+
+    if (wnb) {
+        top_k_mask_pop(dma_q, tr, (uint16_t) ith);
+        // wb0 is 128B aligned (col0 is a multiple of 64 columns), the select may run past the row inside the tile
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ith);
+        top_k_mask_select(out + wb0, flags + wb0, hex_round_up(wnb, 128), kp->ninf);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ith);
+        dma_queue_push(dma_q, dma_make_data(octx->dst->data + wb0, out + wb0), wnb, wnb, wnb, 1);
+        top_k_mask_pop(dma_q, tr, (uint16_t) ith);
+    }
+}
+
+int op_top_k_mask(struct htp_ops_context * octx) {
+    const struct htp_top_k_mask_kernel_params * kp = (const struct htp_top_k_mask_kernel_params *) octx->kernel_params;
+
+    if (kp->vtcm_size > octx->ctx->vtcm_size) {
+        return HTP_STATUS_VTCM_TOO_SMALL;
+    }
+
+    const bool picks = kp->kernel_type == HTP_TOP_K_MASK_KERNEL_PICKS;
+
+    // mdev: ROWS gives this device a contiguous share of the rows, DMA writes are exempt from the 128B rule.
+    // PICKS: every device scans the row, the write back is split by 64 column chunks.
+    uint32_t dev_row0 = 0;
+    uint32_t dev_rows = picks ? 1 : kp->u.rows.total_rows;
+    uint32_t dev_c0   = 0;
+    uint32_t dev_cols = kp->n_kv;
+    if (octx->ctx->mdev.count > 1) {
+        const struct htp_tensor_mdev_range range = picks ?
+            htp_tensor_mdev_partition(kp->n_kv, 64, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div) :
+            htp_tensor_mdev_partition(kp->u.rows.total_rows, 1, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        if (range.count == 0) {
+            return HTP_STATUS_OK;
+        }
+        if (picks) {
+            dev_c0   = range.start;
+            dev_cols = range.count;
+        } else {
+            dev_row0 = range.start;
+            dev_rows = range.count;
+        }
+    }
+
+    // fewer workers than kparams->n_threads always fit the VTCM layout: ROWS runs one row range per worker, PICKS loops over the host slices
+    const uint32_t max_threads = MIN((uint32_t) kp->n_threads, octx->ctx->n_threads);
+    if (!htp_ops_context_set_n_threads(octx, picks ? max_threads : MIN(max_threads, dev_rows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+    const uint32_t n_threads = octx->n_threads;
+
+    uint8_t * vtcm = octx->ctx->vtcm_base;
+
+    struct htp_top_k_mask_context tctx = {
+        .octx      = octx,
+        .kparams   = kp,
+        .off_extra = kp->off_extra_units * HTP_TOP_K_MASK_UNIT,
+        .off_score = kp->off_score_units * HTP_TOP_K_MASK_UNIT,
+        .row0      = dev_row0,
+        .row1      = dev_row0 + dev_rows,
+        .col0      = dev_c0,
+        .ncols_out = dev_cols,
+        .blocks    = VTCM_LAYOUT_PTR(uint8_t, vtcm, kp->off_blocks_u * HTP_TOP_K_MASK_UNIT),
+        .mask      = VTCM_LAYOUT_PTR(uint8_t, vtcm, kp->off_mask_u * HTP_TOP_K_MASK_UNIT),
+        .rowin     = VTCM_LAYOUT_PTR(uint8_t, vtcm, kp->off_rowin_u * HTP_TOP_K_MASK_UNIT),
+        .tiles     = VTCM_LAYOUT_PTR(uint8_t, vtcm, kp->off_tiles_u * HTP_TOP_K_MASK_UNIT),
+        .mtiles    = VTCM_LAYOUT_PTR(uint8_t, vtcm, kp->off_mtiles_u * HTP_TOP_K_MASK_UNIT),
+        .scratch   = VTCM_LAYOUT_PTR(uint8_t, vtcm, kp->off_scratch_u * HTP_TOP_K_MASK_UNIT),
+    };
+
+    if (picks) {
+        atomic_init(&tctx.barrier, 0);
+        work_queue_run(octx->ctx->work_queue, top_k_mask_thread_picks, &tctx, n_threads);
+        return HTP_STATUS_OK;
+    }
+
+    tctx.rows_per_thread = fastdiv(dev_rows + n_threads - 1, &octx->n_threads_div);
+
+    work_queue_run(octx->ctx->work_queue, top_k_mask_preload, &tctx, n_threads);
+    work_queue_run(octx->ctx->work_queue, top_k_mask_thread_rows, &tctx, n_threads);
+
+    return HTP_STATUS_OK;
+}
--- a/Show More
+++ b/Show More