llama : respect no_alloc in all memory caches (#31087 )

* llama : respect no_alloc in memory caches Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4-Flash-Vision-Exp-GGUF (Mac Studio, M2 Ultra, 192GB) * llama : TODO: move no_alloc out of llama_hparams Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4-Flash-Vision-Exp-GGUF (Mac Studio, M2 Ultra, 192GB) * llama : TODO: move no_alloc out of llama_hparams Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4-Flash-Vision-Exp-GGUF (Mac Studio, M2 Ultra, 192GB)
tools : reject truncated GGUF tensor data in merge (#31044 )
2026-10-25 05:57:36 -05:00 · 2026-10-25 11:56:04 +02:00 · 2026-10-25 11:25:10 +02:00 · 2026-10-25 11:24:38 +02:00 · 2026-10-25 11:15:45 +02:00 · 2026-10-25 11:13:48 +02:00
350 changed files with 37434 additions and 21707 deletions
@@ -0,0 +1,109 @@
+name: Behavior Check
+
+on:
+  workflow_dispatch: # allows manual triggering
+  push:
+    branches:
+      - master
+    paths: [
+      '.github/workflows/behavior-check.yml',
+      '**/CMakeLists.txt',
+      '**/.cmake',
+      '**/*.h',
+      '**/*.hpp',
+      '**/*.c',
+      '**/*.cpp',
+    ]
+
+  pull_request:
+    types: [opened, synchronize, reopened]
+    paths: [
+      '.github/workflows/behavior-check.yml',
+      '**/CMakeLists.txt',
+      '**/.cmake',
+      '**/*.h',
+      '**/*.hpp',
+      '**/*.c',
+      '**/*.cpp'
+    ]
+
+cache-mode: none
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
+  cancel-in-progress: true
+
+env:
+  GGML_NLOOP: 3
+  GGML_N_THREADS: 1
+  GGML_SCHED_DEBUG_REALLOC: 1
+  LLAMA_ARG_LOG_COLORS: 1
+  LLAMA_ARG_LOG_PREFIX: 1
+  LLAMA_ARG_LOG_TIMESTAMPS: 1
+
+jobs:
+  ubuntu:
+    runs-on: ubuntu-26.04-arm
+
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.24
+        with:
+          restore: false
+          save: false
+
+      - name: Build Dependencies
+        id: build_depends
+        run: |
+          sudo apt update
+          sudo apt install -y build-essential libssl-dev git-lfs libjpeg-dev python3 python3-venv python3-pip
+
+      - name: Python Dependencies
+        id: python_depends
+        run: |
+          python3 -m venv .venv
+          source .venv/bin/activate
+          pip install jinja2==3.1.6
+
+      - name: ccache-buckets-restore
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
+        with:
+          key: behavior-check-ubuntu-26.04-arm
+          folder: llama.cpp
+          hf_bucket: ggml-org/cache
+
+      - name: Build
+        id: cmake_build
+        run: |
+          cmake -B build \
+            -DGGML_NATIVE=OFF \
+            -DGGML_BACKEND_DL=OFF \
+            -DLLAMA_FATAL_WARNINGS=ON
+          time cmake --build build --config Release -j $(nproc)
+
+      - name: ccache-buckets-save
+        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+        uses: ./.github/actions/ccache-buckets
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
+        with:
+          key: behavior-check-ubuntu-26.04-arm
+          folder: llama.cpp
+          evict-old-files: 1d
+          hf_bucket: ggml-org/cache
+          save: true
+
+      - name: Test
+        id: cmake_test
+        run: |
+          source .venv/bin/activate
+          cd build
+          ctest -L behavior -LE 'model|backend' --verbose --timeout 900
@@ -105,7 +105,7 @@ jobs:
          cd build
          # Metal Paravirtual devices are difficult to support -> disable
          # ref: https://github.com/ggml-org/llama.cpp/pull/19802#issuecomment-4013704023
-          ctest -L main -E "test-llama-archs|test-llama-context" --verbose --timeout 900
+          ctest -L main -E "test-llama-archs|test-llama-context|test-speculative-rejection-sampling" --verbose --timeout 900

  macos-latest-x64:
    runs-on: macos-15-intel
@@ -9,7 +9,7 @@ permissions:

 jobs:
  linux:
-    runs-on: [self-hosted, Linux, CPU]
+    runs-on: [self-hosted, Linux, fast]
    steps:
      - uses: actions/checkout@v6
        with:
@@ -129,7 +129,7 @@ jobs:
        id: cmake_test
        run: |
          cd build
-          ctest -L 'main|python' --verbose --timeout 900
+          ctest -L main --verbose --timeout 900

      - name: Test llama2c conversion
        id: llama2c_test
@@ -49,7 +49,7 @@ jobs:
        - sanitizer: THREAD
          machine: ubuntu-24.04
        - sanitizer: UNDEFINED
-          machine: [self-hosted, X64, Linux]
+          machine: [self-hosted, X64, Linux, CPU]

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

@@ -46,7 +46,7 @@ env:

 jobs:
  cpu-x64-high-perf:
-    runs-on: [self-hosted, Linux, X64]
+    runs-on: [self-hosted, Linux, X64, CPU]

    steps:
      - name: Clone
@@ -56,7 +56,7 @@ jobs:
      - 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 ~/mnt/llama.cpp
+          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

  cpu-arm64-high-perf-graviton4:
    runs-on: ah-ubuntu_24_04-c8g_8x
@@ -60,4 +60,4 @@ jobs:
      - name: Test
        id: ggml-ci
        run: |
-          GG_BUILD_METAL=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
+          GG_BUILD_METAL=1 bash ./ci/run.sh ~/results/llama.cpp/$RUNNER_NAME ~/mnt/llama.cpp
@@ -79,7 +79,7 @@ jobs:
        id: ggml-ci
        run: |
          vulkaninfo --summary
-          GG_BUILD_VULKAN=1 GGML_VK_DISABLE_COOPMAT2=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
+          GG_BUILD_VULKAN=1 GGML_VK_DISABLE_COOPMAT2=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' }}
@@ -124,7 +124,7 @@ jobs:
        id: ggml-ci
        run: |
          vulkaninfo --summary
-          GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
+          GG_BUILD_VULKAN=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' }}
@@ -150,7 +150,7 @@ jobs:
        id: ggml-ci
        run: |
          vulkaninfo --summary
-          GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
+          GG_BUILD_VULKAN=1 bash ./ci/run.sh ~/results/llama.cpp/$RUNNER_NAME ~/mnt/llama.cpp

  gpu-vulkan-intel-linux:
    runs-on: [self-hosted, Linux, Intel]
@@ -131,4 +131,4 @@ jobs:
        id: ggml-ci
        run: |
          GG_BUILD_WEBGPU=1 GG_BUILD_WEBGPU_DAWN_PREFIX="$GITHUB_WORKSPACE/dawn" \
-            bash ./ci/run.sh ~/results/llama.cpp ~/mnt/llama.cpp
+            bash ./ci/run.sh ~/results/llama.cpp/$RUNNER_NAME ~/mnt/llama.cpp
@@ -185,18 +185,28 @@ fi

 ## helpers

-# download a file if it does not exist or if it is outdated
+# download a file unless already cached in $MNT
 function gg_wget {
    local out=$1
    local url=$2

    local cwd=`pwd`
+    local name=`basename $url`
+    local tmp=.gg-wget.$$.tmp

    mkdir -p $out
    cd $out

-    # should not re-download if file is the same
-    wget -nv -c -N $url
+    # cache first: HF artifacts are static, HF sends no Last-Modified for -N
+    if [ ! -s $name ]; then
+        # publish with mv: concurrent runners never read partial files
+        rm -rf $tmp && mkdir $tmp
+        (cd $tmp && wget -nv $url)
+        if [ -s $tmp/$name ]; then
+            mv -f $tmp/$name $name
+        fi
+        rm -rf $tmp
+    fi

    cd $cwd
 }
@@ -254,7 +264,7 @@ function gg_run_ctest_release {
    (time cmake --build . --config Release -j$(nproc)) 2>&1 | tee -a $OUT/${ci}-make.log

    if [ -z ${GG_BUILD_LOW_PERF} ]; then
-        (time ctest -C Release --output-on-failure -L 'main|python' ${CTEST_EXTRA}) 2>&1 | tee -a $OUT/${ci}-ctest.log
+        (time ctest -C Release --output-on-failure -L main ${CTEST_EXTRA}) 2>&1 | tee -a $OUT/${ci}-ctest.log
    else
        (time ctest -C Release --output-on-failure -L main -E test-opt ${CTEST_EXTRA}) 2>&1 | tee -a $OUT/${ci}-ctest.log
    fi
@@ -312,8 +322,8 @@ function gg_run_test_scripts {

    set -e

-    (cd ./tools/gguf-split && time bash tests.sh "$SRC/build-ci-release/bin" "$MNT/models") 2>&1 | tee -a $OUT/${ci}-scripts.log
-    (cd ./tools/quantize   && time bash tests.sh "$SRC/build-ci-release/bin" "$MNT/models") 2>&1 | tee -a $OUT/${ci}-scripts.log
+    (cd ./tools/gguf-split && time bash tests.sh "$SRC/build-ci-release/bin" "$OUT/scratch") 2>&1 | tee -a $OUT/${ci}-scripts.log
+    (cd ./tools/quantize   && time bash tests.sh "$SRC/build-ci-release/bin" "$OUT/scratch") 2>&1 | tee -a $OUT/${ci}-scripts.log

    set +e
 }
@@ -374,10 +384,14 @@ function gg_run_qwen3_0_6b {


    gg_wget models-mnt/wikitext/ https://huggingface.co/datasets/ggml-org/ci/resolve/main/wikitext-2-raw-v1.zip
-    unzip -o models-mnt/wikitext/wikitext-2-raw-v1.zip -d models-mnt/wikitext/
+
+    rm -rf $OUT/scratch/qwen3-0.6B $OUT/scratch/wikitext
+    mkdir -p $OUT/scratch/qwen3-0.6B
+    unzip -o models-mnt/wikitext/wikitext-2-raw-v1.zip -d $OUT/scratch/wikitext/

    path_models="../models-mnt/qwen3/0.6B"
-    path_wiki="../models-mnt/wikitext/wikitext-2-raw"
+    path_local="$OUT/scratch/qwen3-0.6B"
+    path_wiki="$OUT/scratch/wikitext/wikitext-2-raw"

    rm -rf build-ci-release && mkdir build-ci-release && cd build-ci-release

@@ -386,72 +400,15 @@ function gg_run_qwen3_0_6b {
    (cmake -G "${CMAKE_GENERATOR}" -DCMAKE_BUILD_TYPE=Release ${CMAKE_EXTRA} .. ) 2>&1 | tee -a $OUT/${ci}-cmake.log
    (time cmake --build . --config Release -j$(nproc)) 2>&1 | tee -a $OUT/${ci}-make.log

-    python3 ../convert_hf_to_gguf.py ${path_models} --outfile ${path_models}/ggml-model-f16.gguf  --outtype f16
-    python3 ../convert_hf_to_gguf.py ${path_models} --outfile ${path_models}/ggml-model-bf16.gguf --outtype bf16
+    python3 ../convert_hf_to_gguf.py ${path_models} --outfile ${path_local}/ggml-model-f16.gguf  --outtype f16
+    python3 ../convert_hf_to_gguf.py ${path_models} --outfile ${path_local}/ggml-model-bf16.gguf --outtype bf16

-    model_f16="${path_models}/ggml-model-f16.gguf"
-    model_bf16="${path_models}/ggml-model-bf16.gguf"
-    model_q8_0="${path_models}/ggml-model-q8_0.gguf"
-    model_q4_0="${path_models}/ggml-model-q4_0.gguf"
-    model_q4_1="${path_models}/ggml-model-q4_1.gguf"
-    model_q5_0="${path_models}/ggml-model-q5_0.gguf"
-    model_q5_1="${path_models}/ggml-model-q5_1.gguf"
-    model_q2_k="${path_models}/ggml-model-q2_k.gguf"
-    model_q3_k="${path_models}/ggml-model-q3_k.gguf"
-    model_q4_k="${path_models}/ggml-model-q4_k.gguf"
-    model_q5_k="${path_models}/ggml-model-q5_k.gguf"
-    model_q6_k="${path_models}/ggml-model-q6_k.gguf"
+    model_f16="${path_local}/ggml-model-f16.gguf"
+    model_bf16="${path_local}/ggml-model-bf16.gguf"
+    model_q4_0="${path_local}/ggml-model-q4_0.gguf"

    wiki_test="${path_wiki}/wiki.test.raw"

-    ./bin/llama-quantize ${model_bf16} ${model_q8_0} q8_0 $(nproc)
-    ./bin/llama-quantize ${model_bf16} ${model_q4_0} q4_0 $(nproc)
-    ./bin/llama-quantize ${model_bf16} ${model_q4_1} q4_1 $(nproc)
-    ./bin/llama-quantize ${model_bf16} ${model_q5_0} q5_0 $(nproc)
-    ./bin/llama-quantize ${model_bf16} ${model_q5_1} q5_1 $(nproc)
-    ./bin/llama-quantize ${model_bf16} ${model_q2_k} q2_k $(nproc)
-    ./bin/llama-quantize ${model_bf16} ${model_q3_k} q3_k $(nproc)
-    ./bin/llama-quantize ${model_bf16} ${model_q4_k} q4_k $(nproc)
-    ./bin/llama-quantize ${model_bf16} ${model_q5_k} q5_k $(nproc)
-    ./bin/llama-quantize ${model_bf16} ${model_q6_k} q6_k $(nproc)
-
-    (time ./bin/llama-fit-params --model ${model_f16} 2>&1 | tee -a $OUT/${ci}-fp-f16.log)
-
-    (time ./bin/llama-completion -no-cnv --model ${model_f16}  -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-f16.log
-    (time ./bin/llama-completion -no-cnv --model ${model_bf16} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-bf16.log
-    (time ./bin/llama-completion -no-cnv --model ${model_q8_0} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-q8_0.log
-    (time ./bin/llama-completion -no-cnv --model ${model_q4_0} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-q4_0.log
-    (time ./bin/llama-completion -no-cnv --model ${model_q4_1} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-q4_1.log
-    (time ./bin/llama-completion -no-cnv --model ${model_q5_0} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-q5_0.log
-    (time ./bin/llama-completion -no-cnv --model ${model_q5_1} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-q5_1.log
-    (time ./bin/llama-completion -no-cnv --model ${model_q2_k} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-q2_k.log
-    (time ./bin/llama-completion -no-cnv --model ${model_q3_k} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-q3_k.log
-    (time ./bin/llama-completion -no-cnv --model ${model_q4_k} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-q4_k.log
-    (time ./bin/llama-completion -no-cnv --model ${model_q5_k} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-q5_k.log
-    (time ./bin/llama-completion -no-cnv --model ${model_q6_k} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-q6_k.log
-
-    (time ./bin/llama-perplexity --model ${model_f16}  -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-f16.log
-    if [ -z ${GG_BUILD_NO_BF16} ]; then
-        (time ./bin/llama-perplexity --model ${model_bf16} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-bf16.log
-    fi
-    (time ./bin/llama-perplexity --model ${model_q8_0} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-q8_0.log
-    (time ./bin/llama-perplexity --model ${model_q4_0} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-q4_0.log
-    (time ./bin/llama-perplexity --model ${model_q4_1} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-q4_1.log
-    (time ./bin/llama-perplexity --model ${model_q5_0} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-q5_0.log
-    (time ./bin/llama-perplexity --model ${model_q5_1} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-q5_1.log
-    (time ./bin/llama-perplexity --model ${model_q2_k} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-q2_k.log
-    (time ./bin/llama-perplexity --model ${model_q3_k} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-q3_k.log
-    (time ./bin/llama-perplexity --model ${model_q4_k} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-q4_k.log
-    (time ./bin/llama-perplexity --model ${model_q5_k} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-q5_k.log
-    (time ./bin/llama-perplexity --model ${model_q6_k} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-q6_k.log
-
-    (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
-
-    (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 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 off                ) 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
-
    function check_ppl {
        qnt="$1"
        ppl=$(echo "$2" | grep -oE "[0-9]+\.[0-9]+" | tail -n 1)
@@ -465,20 +422,56 @@ function gg_run_qwen3_0_6b {
        return 0
    }

+    for qnt in q8_0 q4_0 q4_1 q5_0 q5_1 q2_k q3_k q4_k q5_k q6_k; do
+        model="${path_local}/ggml-model-${qnt}.gguf"
+
+        ./bin/llama-quantize ${model_bf16} ${model} ${qnt} $(nproc)
+
+        (time ./bin/llama-completion -no-cnv --model ${model} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-${qnt}.log
+        (time ./bin/llama-perplexity --model ${model} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-${qnt}.log
+
+        # note: ppl > 20.0 for q2_k with this model
+        if [ "${qnt}" != "q2_k" ]; then
+            check_ppl "${qnt}" "$(cat $OUT/${ci}-tg-${qnt}.log | grep "^\[1\]")"
+        fi
+
+        if [ "${qnt}" = "q4_0" ]; then
+            # q4_0 is kept for other tests
+            q4_0="${path_models}/ggml-model-q4_0.gguf"
+            ln ${model} "$q4_0" 2>/dev/null || { cp ${model} $q4_0.tmp.$$ && mv -f $q4_0.tmp.$$ $q4_0; }
+        else
+            rm -f ${model}
+        fi
+    done
+
+    (time ./bin/llama-fit-params --model ${model_f16} 2>&1 | tee -a $OUT/${ci}-fp-f16.log)
+
+    (time ./bin/llama-completion -no-cnv --model ${model_f16}  -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-f16.log
+    (time ./bin/llama-completion -no-cnv --model ${model_bf16} -ngl 99 -c 1024 -s 1234 -n 64 --ignore-eos -p "I believe the meaning of life is" ) 2>&1 | tee -a $OUT/${ci}-tg-bf16.log
+
+    (time ./bin/llama-perplexity --model ${model_f16}  -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-f16.log
+    if [ -z ${GG_BUILD_NO_BF16} ]; then
+        (time ./bin/llama-perplexity --model ${model_bf16} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-tg-bf16.log
+    fi
+
+    (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 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 off                ) 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
+    fi
+
    check_ppl "f16"  "$(cat $OUT/${ci}-tg-f16.log  | grep "^\[1\]")"
    if [ -z ${GG_BUILD_NO_BF16} ]; then
        check_ppl "bf16" "$(cat $OUT/${ci}-tg-bf16.log | grep "^\[1\]")"
    fi
-    check_ppl "q8_0" "$(cat $OUT/${ci}-tg-q8_0.log | grep "^\[1\]")"
-    check_ppl "q4_0" "$(cat $OUT/${ci}-tg-q4_0.log | grep "^\[1\]")"
-    check_ppl "q4_1" "$(cat $OUT/${ci}-tg-q4_1.log | grep "^\[1\]")"
-    check_ppl "q5_0" "$(cat $OUT/${ci}-tg-q5_0.log | grep "^\[1\]")"
-    check_ppl "q5_1" "$(cat $OUT/${ci}-tg-q5_1.log | grep "^\[1\]")"
-   #check_ppl "q2_k" "$(cat $OUT/${ci}-tg-q2_k.log | grep "^\[1\]")" # note: ppl > 20.0 for this quant and model
-    check_ppl "q3_k" "$(cat $OUT/${ci}-tg-q3_k.log | grep "^\[1\]")"
-    check_ppl "q4_k" "$(cat $OUT/${ci}-tg-q4_k.log | grep "^\[1\]")"
-    check_ppl "q5_k" "$(cat $OUT/${ci}-tg-q5_k.log | grep "^\[1\]")"
-    check_ppl "q6_k" "$(cat $OUT/${ci}-tg-q6_k.log | grep "^\[1\]")"
+
+    # scratch is rebuilt on every run; q4_0 is already published to $MNT
+    rm -rf $OUT/scratch/qwen3-0.6B $OUT/scratch/wikitext

    set +e
 }
@@ -501,6 +494,8 @@ function gg_run_embd_bge_small {
    gg_wget models-mnt/bge-small/1_Pooling https://huggingface.co/BAAI/bge-small-en-v1.5/raw/main/1_Pooling/config.json

    path_models="../models-mnt/bge-small"
+    path_local="$OUT/scratch/bge-small"
+    rm -rf $path_local && mkdir -p $path_local

    rm -rf build-ci-release && mkdir build-ci-release && cd build-ci-release

@@ -509,10 +504,10 @@ function gg_run_embd_bge_small {
    (cmake -G "${CMAKE_GENERATOR}" -DCMAKE_BUILD_TYPE=Release ${CMAKE_EXTRA} .. ) 2>&1 | tee -a $OUT/${ci}-cmake.log
    (time cmake --build . --config Release -j$(nproc)) 2>&1 | tee -a $OUT/${ci}-make.log

-    python3 ../convert_hf_to_gguf.py ${path_models} --outfile ${path_models}/ggml-model-f16.gguf
+    python3 ../convert_hf_to_gguf.py ${path_models} --outfile ${path_local}/ggml-model-f16.gguf

-    model_f16="${path_models}/ggml-model-f16.gguf"
-    model_q8_0="${path_models}/ggml-model-q8_0.gguf"
+    model_f16="${path_local}/ggml-model-f16.gguf"
+    model_q8_0="${path_local}/ggml-model-q8_0.gguf"

    ./bin/llama-quantize ${model_f16} ${model_q8_0} q8_0

@@ -521,6 +516,8 @@ function gg_run_embd_bge_small {
    (time ./bin/llama-embedding --model ${model_f16}  -p "I believe the meaning of life is" -ngl 99 -c 0 --no-op-offload) 2>&1 | tee -a $OUT/${ci}-tg-f16.log
    (time ./bin/llama-embedding --model ${model_q8_0} -p "I believe the meaning of life is" -ngl 99 -c 0 --no-op-offload) 2>&1 | tee -a $OUT/${ci}-tg-q8_0.log

+    rm -rf $path_local
+
    set +e
 }

@@ -537,6 +534,8 @@ function gg_run_rerank_tiny {
    gg_wget models-mnt/rerank-tiny/ https://huggingface.co/jinaai/jina-reranker-v1-tiny-en/raw/main/vocab.json

    path_models="../models-mnt/rerank-tiny"
+    path_local="$OUT/scratch/rerank-tiny"
+    rm -rf $path_local && mkdir -p $path_local

    rm -rf build-ci-release && mkdir build-ci-release && cd build-ci-release

@@ -545,9 +544,9 @@ function gg_run_rerank_tiny {
    (cmake -G "${CMAKE_GENERATOR}" -DCMAKE_BUILD_TYPE=Release ${CMAKE_EXTRA} .. ) 2>&1 | tee -a $OUT/${ci}-cmake.log
    (time cmake --build . --config Release -j$(nproc)) 2>&1 | tee -a $OUT/${ci}-make.log

-    python3 ../convert_hf_to_gguf.py ${path_models} --outfile ${path_models}/ggml-model-f16.gguf
+    python3 ../convert_hf_to_gguf.py ${path_models} --outfile ${path_local}/ggml-model-f16.gguf

-    model_f16="${path_models}/ggml-model-f16.gguf"
+    model_f16="${path_local}/ggml-model-f16.gguf"

    (time ./bin/llama-fit-params --model ${model_f16} 2>&1 | tee -a $OUT/${ci}-fp-f16.log)

@@ -577,6 +576,8 @@ function gg_run_rerank_tiny {
    check_score "rerank score 1" "$(cat $OUT/${ci}-rk-f16.log | grep "rerank score 1")" "0.00" "0.05" | tee -a $OUT/${ci}-rk-f16.log
    check_score "rerank score 2" "$(cat $OUT/${ci}-rk-f16.log | grep "rerank score 2")" "0.10" "0.30" | tee -a $OUT/${ci}-rk-f16.log

+    rm -rf $path_local
+
    set +e
 }

@@ -683,11 +684,11 @@ if [ -z ${GG_BUILD_LOW_PERF} ]; then
    ln -sfn ${mnt_models} ${SRC}/models-mnt

    # Create a fresh python3 venv and enter it
-    if ! python3 -m venv "$MNT/venv"; then
-        echo "Error: Failed to create Python virtual environment at $MNT/venv."
+    if ! python3 -m venv "$OUT/venv"; then
+        echo "Error: Failed to create Python virtual environment at $OUT/venv."
        exit 1
    fi
-    source "$MNT/venv/bin/activate"
+    source "$OUT/venv/bin/activate"

    pip install -r ${SRC}/requirements.txt --disable-pip-version-check
    pip install --editable gguf-py --disable-pip-version-check
@@ -78,6 +78,8 @@ add_library(${TARGET}
    download.h
    fit.cpp
    fit.h
+    file-lock.cpp
+    file-lock.h
    hf-cache.cpp
    hf-cache.h
    http.h
@@ -184,3 +186,8 @@ if (LLAMA_LLGUIDANCE)
 endif()

 target_link_libraries(${TARGET} PUBLIC llama Threads::Threads)
+
+if (CMAKE_SYSTEM_NAME MATCHES "AIX")
+    # AIX's flock() function comes from libbsd.a
+    target_link_libraries(${TARGET} PRIVATE -lbsd)
+endif()
@@ -470,6 +470,14 @@ void common_models_handler_apply(common_models_handler & handler, common_params
    auto opts = handler.opts; // copy
    opts.callback = callback;

+    // another process may hold the cache lock during finalization:
+    // let the download callback cancel that wait too
+    auto finalize = [callback](const hf_cache::hf_file & f) {
+        return hf_cache::finalize_file(f, [callback]() {
+            return !(callback && callback->is_cancelled());
+        });
+    };
+
    // handle plain "url" if needed
    auto handle_url = [&](common_params_model & model) {
        if (!model.url.empty()) {
@@ -518,7 +526,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params
    }

    // handle hf_plan tasks
-    auto add_tasks = [&opts, &tasks](const hf_cache::hf_files  & model_files,
+    auto add_tasks = [&opts, &tasks, &finalize](const hf_cache::hf_files  & model_files,
                                    const hf_cache::hf_file    & primary,
                                    common_params_model        & model) {
        for (size_t i = 0; i < model_files.size(); ++i) {
@@ -527,9 +535,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params
            tasks.emplace_back(model_file, opts, [&, is_primary]() {
                if (is_primary) {
                    // the primary file is the first split (00001-of), use it as model path
-                    model.path = hf_cache::finalize_file(model_file);
+                    model.path = finalize(model_file);
                } else {
-                    hf_cache::finalize_file(model_file);
+                    finalize(model_file);
                }
            });
        }
@@ -581,9 +589,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params
        tasks.emplace_back(plan_spec.mtp, opts, [&]() {
            // only use the discovered MTP head when no draft path is set yet
            if (params.speculative.draft.mparams.path.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.mtp);
+                params.speculative.draft.mparams.path = finalize(plan_spec.mtp);
            } else {
-                hf_cache::finalize_file(plan_spec.mtp);
+                finalize(plan_spec.mtp);
            }
        });
    }
@@ -591,9 +599,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params
        tasks.emplace_back(plan_spec.dflash, opts, [&]() {
            // only use the discovered DFlash sidecar when no draft path is set yet
            if (params.speculative.draft.mparams.path.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.dflash);
+                params.speculative.draft.mparams.path = finalize(plan_spec.dflash);
            } else {
-                hf_cache::finalize_file(plan_spec.dflash);
+                finalize(plan_spec.dflash);
            }
        });
    }
@@ -601,9 +609,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params
        tasks.emplace_back(plan_spec.eagle3, opts, [&]() {
            // only use the discovered Eagle3 sidecar when no draft path is set yet
            if (params.speculative.draft.mparams.path.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.eagle3);
+                params.speculative.draft.mparams.path = finalize(plan_spec.eagle3);
            } else {
-                hf_cache::finalize_file(plan_spec.eagle3);
+                finalize(plan_spec.eagle3);
            }
        });
    }
@@ -611,9 +619,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params
        tasks.emplace_back(plan_spec.dspark, opts, [&]() {
            // only use the discovered DSpark sidecar when no draft path is set yet
            if (params.speculative.draft.mparams.path.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.dspark);
+                params.speculative.draft.mparams.path = finalize(plan_spec.dspark);
            } else {
-                hf_cache::finalize_file(plan_spec.dspark);
+                finalize(plan_spec.dspark);
            }
        });
    }
@@ -634,16 +642,16 @@ void common_models_handler_apply(common_models_handler & handler, common_params
    }
    if (!plan.mmproj.local_path.empty()) {
        tasks.emplace_back(plan.mmproj, opts, [&]() {
-            params.mmproj.path = hf_cache::finalize_file(plan.mmproj);
+            params.mmproj.path = finalize(plan.mmproj);
        });
    }
    if (!plan.mtp.local_path.empty() && !had_spec_url) {
        tasks.emplace_back(plan.mtp, opts, [&]() {
            // only fall back to the discovered MTP head when no draft was explicitly provided
            if (params.speculative.draft.mparams.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.mtp);
+                params.speculative.draft.mparams.path = finalize(plan.mtp);
            } else {
-                hf_cache::finalize_file(plan.mtp);
+                finalize(plan.mtp);
            }
        });
    }
@@ -651,9 +659,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params
        tasks.emplace_back(plan.dflash, opts, [&]() {
            // only fall back to the discovered DFlash sidecar when no draft was explicitly provided
            if (params.speculative.draft.mparams.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.dflash);
+                params.speculative.draft.mparams.path = finalize(plan.dflash);
            } else {
-                hf_cache::finalize_file(plan.dflash);
+                finalize(plan.dflash);
            }
        });
    }
@@ -661,9 +669,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params
        tasks.emplace_back(plan.eagle3, opts, [&]() {
            // only fall back to the discovered Eagle3 sidecar when no draft was explicitly provided
            if (params.speculative.draft.mparams.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.eagle3);
+                params.speculative.draft.mparams.path = finalize(plan.eagle3);
            } else {
-                hf_cache::finalize_file(plan.eagle3);
+                finalize(plan.eagle3);
            }
        });
    }
@@ -671,9 +679,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params
        tasks.emplace_back(plan.dspark, opts, [&]() {
            // only fall back to the discovered DSpark sidecar when no draft was explicitly provided
            if (params.speculative.draft.mparams.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.dspark);
+                params.speculative.draft.mparams.path = finalize(plan.dspark);
            } else {
-                hf_cache::finalize_file(plan.dspark);
+                finalize(plan.dspark);
            }
        });
    }
@@ -681,7 +689,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params
        tasks.emplace_back(plan.preset, opts, [&]() {
            // if HF repo is a preset repo, we simply run server in router mode with the preset.ini file
            params.models_preset_hf = params.model.hf_repo; // only for showing a warning
-            params.models_preset    = hf_cache::finalize_file(plan.preset);
+            params.models_preset    = finalize(plan.preset);
            // clear the model so the server starts in router mode
            params.model.path.clear();
            params.model.hf_repo.clear();
@@ -428,6 +428,41 @@ std::vector<common_chat_msg> common_chat_msgs_parse_oaicompat(const json & messa
            if (message.contains("tool_call_id")) {
                msg.tool_call_id = message.at("tool_call_id");
            }
+            {
+                json extra = json::object();
+                for (const auto & [k, v] : message.items()) {
+                    if (k != "role" && k != "content" && k != "tool_calls" && k != "reasoning_content" &&
+                        k != "name" && k != "tool_call_id") {
+                        extra[k] = v;
+                    }
+                }
+                if (has_tool_calls) {
+                    // per tool call: keys besides type/function/id, and function keys besides name/arguments
+                    json tc_extra = json::array();
+                    bool any      = false;
+                    for (const auto & tool_call : message.at("tool_calls")) {
+                        json e = json::object();
+                        for (const auto & [k, v] : tool_call.items()) {
+                            if (k != "type" && k != "function" && k != "id") {
+                                e[k] = v;
+                            }
+                        }
+                        for (const auto & [k, v] : tool_call.at("function").items()) {
+                            if (k != "name" && k != "arguments") {
+                                e["function"][k] = v;
+                            }
+                        }
+                        any = any || !e.empty();
+                        tc_extra.push_back(e);
+                    }
+                    if (any) {
+                        extra["__tool_call_extras"] = tc_extra;
+                    }
+                }
+                if (!extra.empty()) {
+                    msg.extra_json = extra.dump();
+                }
+            }

            msgs.push_back(msg);
        }
@@ -516,6 +551,50 @@ static json render_message_to_json(const std::vector<common_chat_msg> & msgs, co
    return messages_inp_normalizer(c).normalize(messages);
 }

+// DeepSeek V4.1 reads non-OpenAI message keys and joins content blocks with "\n\n" around images,
+// so keep media markers as separate text parts and restore the extra keys
+static json render_message_to_json_deepseek_v4_1(const std::vector<common_chat_msg> & msgs, const jinja::caps & c) {
+    json messages = json::array();
+    for (const auto & msg : msgs) {
+        json jmsg;
+        if (msg.contains_media() && msg.content.empty()) {
+            common_chat_msg typed = msg;
+            for (auto & part : typed.content_parts) {
+                part.type = "text";
+            }
+            jmsg = typed.to_json_oaicompat(/* concat_typed_text= */ false);
+        } else {
+            jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ false);
+        }
+        if (!msg.extra_json.empty()) {
+            const json extra = json::parse(msg.extra_json);
+            for (const auto & [k, v] : extra.items()) {
+                if (k == "__tool_call_extras") {
+                    if (!jmsg.contains("tool_calls")) {
+                        continue;
+                    }
+                    auto & tcs = jmsg.at("tool_calls");
+                    for (size_t i = 0; i < tcs.size() && i < v.size(); i++) {
+                        for (const auto & [tk, tv] : v[i].items()) {
+                            if (tk == "function") {
+                                for (const auto & [fk, fv] : tv.items()) {
+                                    tcs[i]["function"][fk] = fv;
+                                }
+                            } else {
+                                tcs[i][tk] = tv;
+                            }
+                        }
+                    }
+                } else if (!jmsg.contains(k)) {
+                    jmsg[k] = v;
+                }
+            }
+        }
+        messages.push_back(std::move(jmsg));
+    }
+    return messages_inp_normalizer(c).normalize(messages);
+}
+
 // DEPRECATED: only used in tests
 json common_chat_msgs_to_json_oaicompat(const std::vector<common_chat_msg> & msgs, bool concat_typed_text) {
    jinja::caps c;
@@ -1208,11 +1287,13 @@ static common_chat_params_init_fn common_chat_template_detect_params_init(const
    // DeepSeek V3.2/V4 format detection: template defines dsml_token and uses it for tool calls.
    // The template source contains the token as a variable assignment, not as a literal in markup.
    // V3.2 names the tool call block "function_calls", V4 names it "tool_calls".
-    if (src.find("dsml_token") != std::string::npos &&
-        src.find("DSML") != std::string::npos &&
-        (src.find("function_calls") != std::string::npos ||
-         src.find("tool_calls") != std::string::npos)) {
-        LOG_DBG("Using specialized template: DeepSeek V3.2/V4\n");
+    // V4.1 tag names have a leading space and its template does not name the dsml_token variable.
+    if ((src.find("dsml_token") != std::string::npos &&
+         src.find("DSML") != std::string::npos &&
+         (src.find("function_calls") != std::string::npos ||
+          src.find("tool_calls") != std::string::npos)) ||
+        src.find("｜DSML｜ calls>") != std::string::npos) {
+        LOG_DBG("Using specialized template: DeepSeek V3.2/V4/V4.1\n");
        return common_chat_params_init_deepseek_v3_2;
    }

@@ -1313,7 +1394,9 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_
        workaround::trim_all_content(trimmed_messages);
        messages_to_render = &trimmed_messages;
    }
-    params.messages              = render_message_to_json(*messages_to_render, tmpl.original_caps());
+    params.messages              = src.find("｜DSML｜ calls>") != std::string::npos ?
+                                       render_message_to_json_deepseek_v4_1(*messages_to_render, tmpl.original_caps()) :
+                                       render_message_to_json(*messages_to_render, tmpl.original_caps());
    params.tool_choice           = inputs.tool_choice;
    params.reasoning_format      = inputs.reasoning_format;
    params.enable_thinking       = inputs.enable_thinking;
@@ -94,6 +94,8 @@ struct common_chat_msg {
    std::string                               reasoning_content;
    std::string                               tool_name;
    std::string                               tool_call_id;
+    // non-OpenAI message keys as a JSON object string, only for templates that read them (DeepSeek V4.1)
+    std::string                               extra_json;

    common_json to_json_oaicompat(bool concat_typed_text = false) const;

@@ -26,6 +26,7 @@
 #include <fstream>
 #include <iostream>
 #include <iterator>
+#include <random>
 #include <regex>
 #include <sstream>
 #include <string>
@@ -917,9 +918,13 @@ std::string fs_path_to_utf8(const std::filesystem::path & path) {
 }

 void fs_write_atomic(const std::filesystem::path & path, const std::string & data) {
+    // a random suffix gives every writer its own temporary file, across processes and containers
+    std::random_device rd;
+    const uint64_t id = ((uint64_t) rd() << 32) | rd();
+
    std::error_code ec;
    std::filesystem::path path_tmp = path;
-    path_tmp += ".tmp";
+    path_tmp += "." + std::to_string(id) + ".tmp";

    if (path.has_parent_path()) {
        std::filesystem::create_directories(path.parent_path(), ec);
@@ -4,6 +4,7 @@
 #include "common.h"
 #include "log.h"
 #include "download.h"
+#include "file-lock.h"
 #include "hf-cache.h"
 #include "json.h"

@@ -77,6 +78,18 @@ std::pair<std::string, std::string> common_download_split_repo_tag(const std::st
    return {hf_repo, tag};
 }

+// file name from url, without the query string
+static std::string url_get_filename(const std::string & url) {
+    std::string filename = url;
+    if (auto pos = filename.rfind('/'); pos != std::string::npos) {
+        filename = filename.substr(pos + 1);
+    }
+    if (auto pos = filename.find('?'); pos != std::string::npos) {
+        filename = filename.substr(0, pos);
+    }
+    return filename;
+}
+
 class ProgressBar : public common_download_callback {
    static inline std::mutex mutex;
    static inline std::map<const ProgressBar *, int> lines;
@@ -100,14 +113,8 @@ public:
    ProgressBar() = default;

    void on_start(const common_download_progress & p) override {
-        filename = p.url;
+        filename = url_get_filename(p.url);

-        if (auto pos = filename.rfind('/'); pos != std::string::npos) {
-            filename = filename.substr(pos + 1);
-        }
-        if (auto pos = filename.find('?'); pos != std::string::npos) {
-            filename = filename.substr(0, pos);
-        }
        for (size_t i = 0; i < filename.size(); ++i) {
            if ((filename[i] & 0xC0) != 0x80) {
                if (len++ == 39) {
@@ -257,15 +264,62 @@ static bool common_pull_file(httplib::Client & cli,
 static int common_download_file_single_online(const std::string & url,
                                              const std::string & path,
                                              const common_download_opts & opts,
-                                              bool skip_etag) {
+                                              bool skip_etag,
+                                              const std::string & lock_path,
+                                              const std::string & alt_path) {
    static const int max_attempts        = 3;
    static const int retry_delay_seconds = 2;

-    const bool file_exists = std::filesystem::exists(std::filesystem::u8path(path));
+    const std::string path_temporary = path + PARTIAL_SUFFIX;
+    const std::string path_lock = lock_path.empty() ? path + ".lock" : lock_path;
+
+    common_download_progress p;
+    p.url = url;
+
+    if (opts.callback) {
+        opts.callback->on_start(p);
+    }
+
+    // the completion callback may itself take the same lock (the HF cache
+    // finalizes blobs under it), so every exit goes through finish()
+    common_file_lock lock(path_lock);
+    auto finish = [&](int status) {
+        lock.close();
+        if (opts.callback) {
+            opts.callback->on_done(p, is_http_status_ok(status));
+        }
+        return status;
+    };
+
+    // 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)));
+    };
+
+    if (skip_etag && cache_hit()) {
+        LOG_DBG("%s: using cached file: %s\n", __func__, path.c_str());
+        return finish(304); // 304 Not Modified - fake cached response
+    }
+
+    // take the lock before touching the cache, so two processes cannot race on the same file
+    bool waited = false;
+    if (!lock.acquire([&] {
+            if (!waited) {
+                LOG_INF("%s: file '%s' is being downloaded by another process, waiting...\n", __func__, url_get_filename(url).c_str());
+                waited = true;
+            }
+            return !(opts.callback && opts.callback->is_cancelled());
+        })) {
+        return finish(-1);
+    }
+
+    // under the lock: another process may have finished while we waited
+    bool file_exists = cache_hit();

    if (file_exists && skip_etag) {
        LOG_DBG("%s: using cached file: %s\n", __func__, path.c_str());
-        return 304; // 304 Not Modified - fake cached response
+        return finish(304); // 304 Not Modified - fake cached response
    }

    auto [cli, parts] = common_http_client(url);
@@ -294,9 +348,9 @@ static int common_download_file_single_online(const std::string & url,
        LOG_TRC("%s: HEAD failed, status: %d\n", __func__, head ? head->status : -1);
        if (file_exists) {
            LOG_TRC("%s: using cached file (HEAD failed): %s\n", __func__, path.c_str());
-            return 304; // 304 Not Modified - fake cached response
+            return finish(304); // 304 Not Modified - fake cached response
        }
-        return head ? head->status : -1;
+        return finish(head ? head->status : -1);
    }

    std::string etag;
@@ -304,8 +358,6 @@ static int common_download_file_single_online(const std::string & url,
        etag = head->get_header_value("ETag");
    }

-    common_download_progress p;
-    p.url = url;
    if (head->has_header("Content-Length")) {
        try {
            p.total = std::stoull(head->get_header_value("Content-Length"));
@@ -322,16 +374,16 @@ static int common_download_file_single_online(const std::string & url,
    if (file_exists) {
        if (etag.empty()) {
            LOG_DBG("%s: using cached file (no server etag): %s\n", __func__, path.c_str());
-            return 304; // 304 Not Modified - fake cached response
+            return finish(304); // 304 Not Modified - fake cached response
        }
        if (!last_etag.empty() && last_etag == etag) {
            LOG_DBG("%s: using cached file (same etag): %s\n", __func__, path.c_str());
-            return 304; // 304 Not Modified - fake cached response
+            return finish(304); // 304 Not Modified - fake cached response
        }
        // pass this point, the file exists but is different from the server version, so we need to redownload it
        if (remove(path.c_str()) != 0) {
            LOG_ERR("%s: unable to delete file: %s\n", __func__, path.c_str());
-            return -1;
+            return finish(-1);
        }
    }

@@ -341,7 +393,6 @@ static int common_download_file_single_online(const std::string & url,
    }

    bool success = false;
-    const std::string path_temporary = path + PARTIAL_SUFFIX;
    int delay = retry_delay_seconds;

    if (!skip_etag && !etag.empty()) {
@@ -349,7 +400,7 @@ static int common_download_file_single_online(const std::string & url,
            LOG_WRN("%s: partial file does not match the server etag, discarding it: %s\n", __func__, path_temporary.c_str());
            if (remove(path_temporary.c_str()) != 0) {
                LOG_ERR("%s: unable to delete file: %s\n", __func__, path_temporary.c_str());
-                return -1;
+                return finish(-1);
            }
        }
        write_etag(path_temporary, etag);
@@ -359,7 +410,7 @@ static int common_download_file_single_online(const std::string & url,
        // report the resume offset, so a resumed file does not start at 0
        p.downloaded = supports_ranges && std::filesystem::exists(path_temporary) ?
                       std::filesystem::file_size(path_temporary) : 0;
-        opts.callback->on_start(p);
+        opts.callback->on_update(p);
    }

    for (int i = 0; i < max_attempts; ++i) {
@@ -403,16 +454,12 @@ static int common_download_file_single_online(const std::string & url,
        }
    }

-    if (opts.callback) {
-        opts.callback->on_done(p, success);
-    }
    // a stopped download keeps path_temporary, so the next attempt resumes from it
    if (!success) {
        LOG_ERR("%s: download failed after %d attempts\n", __func__, max_attempts);
-        return -1; // max attempts reached
    }

-    return head->status;
+    return finish(success ? head->status : -1);
 }

 std::pair<long, std::vector<char>> common_remote_get_content(const std::string          & url,
@@ -452,14 +499,16 @@ std::pair<long, std::vector<char>> common_remote_get_content(const std::string
 int common_download_file_single(const std::string & url,
                                const std::string & path,
                                const common_download_opts & opts,
-                                bool skip_etag) {
+                                bool skip_etag,
+                                const std::string & lock_path,
+                                const std::string & alt_path) {
    if (!opts.offline) {
        ProgressBar tty_cb;
        common_download_opts online_opts = opts;
        if (!online_opts.callback) {
            online_opts.callback = &tty_cb;
        }
-        return common_download_file_single_online(url, path, online_opts, skip_etag);
+        return common_download_file_single_online(url, path, online_opts, skip_etag, lock_path, alt_path);
    }

    if (!std::filesystem::exists(std::filesystem::u8path(path))) {
@@ -821,7 +870,7 @@ void common_download_run_tasks(const std::vector<common_download_task> & tasks)
    for (const auto & task : tasks) {
        futures.push_back(std::async(std::launch::async,
            [&task]() {
-                return common_download_file_single(task.url, task.local_path, task.opts, task.is_hf);
+                return common_download_file_single(task.url, task.local_path, task.opts, task.is_hf, task.lock_path, task.alt_path);
            }
        ));
    }
@@ -69,12 +69,15 @@ struct common_download_task {
    std::string local_path;
    std::function<void()> on_done;
    bool is_hf = false;
+    std::string lock_path;
+    std::string alt_path; // HF snapshot path, also treated as cached (the blob may be moved there)

    common_download_task() = default;
    common_download_task(hf_cache::hf_file f,
            const common_download_opts & opts,
            std::function<void()> on_done = nullptr)
-        : opts(opts), url(f.url), local_path(f.local_path), on_done(on_done), is_hf(true) {}
+        : opts(opts), url(f.url), local_path(f.local_path), on_done(on_done), is_hf(true),
+          lock_path(hf_cache::get_lock_path(f.repo_id, f.oid)), alt_path(f.final_path) {}
 };

 void common_download_run_tasks(const std::vector<common_download_task> & tasks);
@@ -92,10 +95,14 @@ std::string common_download_resolve_path(const std::string & hf_repo_with_tag, c
 // download single file from url to local path
 // returns status code or -1 on error
 // skip_etag: if true, don't read/write .etag files (for HF cache where filename is the hash)
+// lock_path: lock file used to serialize downloads of the same file; when empty, "<path>.lock" is used
+// alt_path: additional path treated as present in cache (see cache_hit in download.cpp)
 int common_download_file_single(const std::string & url,
                                const std::string & path,
                                const common_download_opts & opts = {},
-                                bool skip_etag = false);
+                                bool skip_etag = false,
+                                const std::string & lock_path = "",
+                                const std::string & alt_path = "");

 // resolve and download model from Docker registry
 // return local path to downloaded model file
@@ -0,0 +1,179 @@
+#include "file-lock.h"
+
+#include "log.h"
+
+#include <chrono>
+#include <filesystem>
+#include <thread>
+
+#include <fcntl.h>
+#include <sys/stat.h>
+
+#if defined(_WIN32)
+#define WIN32_LEAN_AND_MEAN
+#ifndef NOMINMAX
+#define NOMINMAX
+#endif
+#include <windows.h>
+#include <io.h>
+#else
+#include <sys/file.h>
+#include <unistd.h>
+#include <cerrno>
+#include <cstring>
+#endif
+
+common_file_lock::common_file_lock(const std::string & lock_path) {
+    this->lock_path = lock_path;
+}
+
+#if defined(_WIN32)
+// the \\?\ prefix (for paths longer than 255 chars) is only accepted on absolute paths
+// with backslash separators: normalize first, then check the length of the absolute form
+static std::wstring as_extended_path(const std::string & path) {
+    // convert from the narrow encoding std::filesystem uses on this system
+    int wlen = MultiByteToWideChar(CP_ACP, 0, path.c_str(), (int) path.size(), nullptr, 0);
+    if (wlen <= 0) {
+        return std::wstring(path.begin(), path.end()); // best effort for exotic encodings
+    }
+    std::wstring wide(wlen, L'\0');
+    MultiByteToWideChar(CP_ACP, 0, path.c_str(), (int) path.size(), wide.data(), wlen);
+
+    if (wide.rfind(L"\\\\?\\", 0) == 0) {
+        return wide;
+    }
+
+    // GetFullPathNameW resolves relative paths and normalizes separators
+    DWORD size = GetFullPathNameW(wide.c_str(), 0, nullptr, nullptr);
+    if (size == 0) {
+        return wide;
+    }
+    std::wstring full(size, L'\0');
+    DWORD written = GetFullPathNameW(wide.c_str(), size, full.data(), nullptr);
+    if (written == 0 || written >= size) {
+        return wide;
+    }
+    full.resize(written);
+
+    if (full.size() <= 255) {
+        return wide; // short enough: open the path exactly as given
+    }
+    if (full.rfind(L"\\\\", 0) == 0) {
+        return L"\\\\?\\UNC\\" + full.substr(2); // \\server\share -> \\?\UNC\server\share
+    }
+    return L"\\\\?\\" + full;
+}
+#endif
+
+common_file_lock::status common_file_lock::try_lock() {
+    if (lock_path.empty()) {
+        return status::acquired; // no lock path, nothing to lock
+    }
+#if defined(_WIN32)
+    if (fd < 0) {
+        // opened with a wide path: extended paths cannot go through the narrow CRT
+        std::wstring lock_path_w = as_extended_path(lock_path);
+        fd = _wopen(lock_path_w.c_str(), _O_RDWR | _O_CREAT | _O_BINARY | _O_NOINHERIT, _S_IREAD | _S_IWRITE);
+        if (fd < 0) {
+            LOG_ERR("%s: cannot open lock file: %s\n", __func__, lock_path.c_str());
+            return status::error;
+        }
+    }
+    // lock the first byte, the file itself stays empty
+    OVERLAPPED ov = {};
+    locked = LockFileEx((HANDLE) _get_osfhandle(fd),
+                        LOCKFILE_EXCLUSIVE_LOCK | LOCKFILE_FAIL_IMMEDIATELY, 0, 1, 0, &ov);
+    if (locked) {
+        return status::acquired;
+    }
+    if (GetLastError() == ERROR_LOCK_VIOLATION) {
+        return status::busy;
+    }
+    LOG_ERR("%s: file locking failed (%d)\n", __func__, (int) GetLastError());
+    return status::error;
+#else
+    if (fd < 0) {
+        int open_flags = O_RDWR | O_CREAT | O_CLOEXEC;
+#ifdef O_NOFOLLOW
+        open_flags |= O_NOFOLLOW;
+#endif
+        // group-writable, so users sharing a cache can wait on the lock
+        fd = open(lock_path.c_str(), open_flags, 0664);
+        if (fd < 0) {
+            LOG_ERR("%s: cannot open lock file: %s\n", __func__, lock_path.c_str());
+            return status::error;
+        }
+        // open() filters the mode through the umask
+        (void) fchmod(fd, 0664);
+    }
+
+    if (flock(fd, LOCK_EX | LOCK_NB) == 0) {
+        locked = true;
+        return status::acquired;
+    }
+    // EACCES is reported as contention on some filesystems
+    if (errno == EWOULDBLOCK || errno == EAGAIN || errno == EACCES) {
+        return status::busy;
+    }
+    // flock is not implemented (some NFS/FUSE mounts): report it instead of racing
+    LOG_ERR("%s: file locking failed (%s)\n", __func__, strerror(errno));
+    return status::error;
+#endif
+}
+
+bool common_file_lock::acquire(const std::function<bool()> & keep_waiting) {
+    // the lock file is opened by try_lock, its parent must exist by then
+    if (!lock_path.empty()) {
+        std::error_code ec;
+        auto lock_dir = std::filesystem::path(lock_path).parent_path();
+        if (!lock_dir.empty()) {
+            std::filesystem::create_directories(lock_dir, ec);
+            if (ec) {
+                LOG_ERR("%s: unable to create lock directory '%s': %s\n", __func__, lock_dir.string().c_str(), ec.message().c_str());
+                return false;
+            }
+        }
+    }
+
+    while (true) {
+        switch (try_lock()) {
+            case status::acquired:
+                return true;
+            case status::busy:
+                if (keep_waiting && !keep_waiting()) {
+                    return false;
+                }
+                std::this_thread::sleep_for(std::chrono::milliseconds(100));
+                continue;
+            case status::error:
+                return false;
+        }
+    }
+}
+
+// the lock file is never deleted: another process could then create a fresh
+// file at the same path and lock that
+void common_file_lock::close() {
+#if defined(_WIN32)
+    if (fd >= 0) {
+        if (locked) {
+            OVERLAPPED ov = {};
+            UnlockFileEx((HANDLE) _get_osfhandle(fd), 0, 1, 0, &ov);
+        }
+        _close(fd);
+    }
+#else
+    if (fd >= 0) {
+        if (locked) {
+            flock(fd, LOCK_UN);
+        }
+        ::close(fd);
+    }
+#endif
+    fd     = -1;
+    locked = false;
+}
+
+common_file_lock::~common_file_lock() {
+    close();
+}
@@ -0,0 +1,36 @@
+#pragma once
+
+#include <functional>
+#include <string>
+
+// Cross-process exclusive lock (flock on POSIX, LockFileEx on Windows).
+// Lazy: acquire() opens the lock file, so a read-only cache works if the lock is never acquired.
+class common_file_lock {
+public:
+    common_file_lock(const std::string & lock_path);
+    ~common_file_lock();
+
+    // take the lock, waiting while another process holds it;
+    // keep_waiting is polled before each retry, false cancels the wait.
+    // Also returns false when locking is unavailable.
+    bool acquire(const std::function<bool()> & keep_waiting = {});
+
+    void close();
+
+    common_file_lock(const common_file_lock &) = delete;
+    common_file_lock & operator=(const common_file_lock &) = delete;
+
+private:
+    enum class status {
+        acquired,
+        busy,  // another process holds the lock
+        error, // file locking is unavailable: no exclusion is possible
+    };
+
+    status try_lock();
+
+    std::string lock_path;
+
+    int  fd     = -1;
+    bool locked = false;
+};
@@ -2,6 +2,7 @@

 #include "build-info.h"
 #include "common.h"
+#include "file-lock.h"
 #include "log.h"
 #include "http.h"
 #include "json.h"
@@ -161,6 +162,14 @@ static bool is_valid_oid(const std::string & oid) {
    return is_hex_string(oid, 40) || is_hex_string(oid, 64);
 }

+std::string get_lock_path(const std::string & repo_id, const std::string & oid) {
+    if (!is_valid_repo_id(repo_id) || !is_valid_oid(oid)) {
+        return {};
+    }
+    fs::path lock_path = fs::path(get_cache_path()) / ".locks" / repo_to_folder_name(repo_id) / (oid + ".lock");
+    return lock_path.string();
+}
+
 static bool is_valid_subpath(const fs::path & path, const fs::path & subpath) {
    if (subpath.is_absolute()) {
        return false; // never do a / b with b absolute
@@ -437,7 +446,7 @@ hf_files get_cached_files(const std::string & repo_id) {
    return files;
 }

-std::string finalize_file(const hf_file & file) {
+std::string finalize_file(const hf_file & file, const std::function<bool()> & keep_waiting) {
    static std::atomic<bool> symlinks_disabled{false};

    std::error_code ec;
@@ -452,6 +461,24 @@ std::string finalize_file(const hf_file & file) {
        return file.final_path;
    }

+    // the lock file is shared with huggingface_hub: another process may be
+    // finalizing the same blob
+    common_file_lock lock(get_lock_path(file.repo_id, file.oid));
+    if (!lock.acquire(keep_waiting)) {
+        // without the lock, another process could move the blob while we check it
+        throw std::runtime_error("failed to lock cache for finalization: " + file.final_path);
+    }
+
+    // acquire() only polls keep_waiting while waiting for a busy lock
+    if (keep_waiting && !keep_waiting()) {
+        throw std::runtime_error("download cancelled before finalization: " + file.final_path);
+    }
+
+    // re-check under the lock: the winner may have created the link already
+    if (fs::exists(final_path, ec)) {
+        return file.final_path;
+    }
+
    fs::create_directories(final_path.parent_path(), ec);

    if (!symlinks_disabled) {
@@ -1,5 +1,6 @@
 #pragma once

+#include <functional>
 #include <string>
 #include <vector>

@@ -26,8 +27,14 @@ hf_files get_repo_files(

 hf_files get_cached_files(const std::string & repo_id = {});

-// Create snapshot path (link or move/copy) and return it
-std::string finalize_file(const hf_file & file);
+// Lock file path for a blob, shared with the hf CLI: <cache>/.locks/<repo folder>/<oid>.lock
+// Pure path computation: creates nothing. Returns an empty string on invalid repo or oid.
+std::string get_lock_path(const std::string & repo_id, const std::string & oid);
+
+// Create snapshot path (link or move/copy) and return it.
+// Throws when the lock cannot be taken or keep_waiting returns false.
+// A completed snapshot (final_path present) is returned without the lock, so a read-only cache works.
+std::string finalize_file(const hf_file & file, const std::function<bool()> & keep_waiting = {});

 // Remove the entire cached directory for a repo, returns true if removed
 bool remove_cached_repo(const std::string & repo_id);
@@ -6,6 +6,7 @@
 #include "json.h"

 #include <sstream>
+#include <stdexcept>
 #include <string>
 #include <cctype>
 #include <cfloat>
@@ -299,74 +300,53 @@ static value & get_attribute(const value & val, const value & attr, value & defa
    return default_val;
 }

-template<bool is_reject>
-static value selectattr(const func_args & args) {
-    args.ensure_count(2, 4);
-    args.ensure_vals<value_array, value_string, value_string, value_string>(true, true, false, false);
+template<bool is_reject, bool lookup_attr>
+static value select_or_reject(const func_args & args) {
+    // example: array | select("odd") -> test_is_odd(item)
+    // example: array | selectattr("status", "equalto", "active") -> test_is_equalto(item.status, "active")
+    constexpr size_t test_pos = lookup_attr ? 2 : 1;
+    args.ensure_count(test_pos);
+    args.ensure_val<value_array>(args.get_pos(0));
+    value attr;
+    if constexpr (lookup_attr) {
+        attr = args.get_pos(1);
+        args.ensure_val<value_string>(attr);
+    }
+
+    auto & builtins = global_builtins();
+    auto test_it = builtins.end();
+    if (args.count() > test_pos) {
+        args.ensure_val<value_string>(args.get_pos(test_pos));
+        std::string test_name = args.get_pos(test_pos)->as_string().str();
+        test_it = builtins.find("test_is_" + test_name);
+        if (test_it == builtins.end()) {
+            throw raised_exception(args.func_name + ": unknown test '" + test_name + "'");
+        }
+    }

    auto arr = args.get_pos(0)->as_array();
-    auto attribute = args.get_pos(1);
    auto out = mk_val<value_array>();
    value val_default = mk_val<value_undefined>();

-    if (args.count() == 2) {
-        // example: array | selectattr("active")
-        for (const auto & item : arr) {
-            value attr_val = get_attribute(item, attribute, val_default);
-            bool is_selected = attr_val->as_bool();
-            if constexpr (is_reject) is_selected = !is_selected;
-            if (is_selected) out->push_back(item);
+    for (const auto & item : arr) {
+        value val = item;
+        if constexpr (lookup_attr) {
+            val = get_attribute(item, attr, val_default);
        }
-        return out;
-
-    } else if (args.count() == 3) {
-        // example: array | selectattr("equalto", "text")
-        // translated to: test_is_equalto(item, "text")
-        std::string test_name = args.get_pos(1)->as_string().str();
-        value test_val = args.get_pos(2);
-        auto & builtins = global_builtins();
-        auto it = builtins.find("test_is_" + test_name);
-        if (it == builtins.end()) {
-            throw raised_exception("selectattr: unknown test '" + test_name + "'");
-        }
-        auto test_fn = it->second;
-        for (const auto & item : arr) {
+        bool is_selected;
+        if (test_it == builtins.end()) {
+            is_selected = val->as_bool();
+        } else {
            func_args test_args(args.ctx);
-            test_args.push_back(item); // current object
-            test_args.push_back(test_val); // extra argument
-            value test_result = test_fn(test_args);
-            bool is_selected = test_result->as_bool();
-            if constexpr (is_reject) is_selected = !is_selected;
-            if (is_selected) out->push_back(item);
+            test_args.push_back(val);
+            for (size_t i = test_pos + 1; i < args.count(); i++) {
+                test_args.push_back(args.get_pos(i));
+            }
+            is_selected = test_it->second(test_args)->as_bool();
        }
-        return out;
-
-    } else if (args.count() == 4) {
-        // example: array | selectattr("status", "equalto", "active")
-        // translated to: test_is_equalto(item.status, "active")
-        std::string test_name = args.get_pos(2)->as_string().str();
-        auto extra_arg = args.get_pos(3);
-        auto & builtins = global_builtins();
-        auto it = builtins.find("test_is_" + test_name);
-        if (it == builtins.end()) {
-            throw raised_exception("selectattr: unknown test '" + test_name + "'");
-        }
-        auto test_fn = it->second;
-        for (const auto & item : arr) {
-            value attr_val = get_attribute(item, attribute, val_default);
-            func_args test_args(args.ctx);
-            test_args.push_back(attr_val); // attribute value
-            test_args.push_back(extra_arg); // extra argument
-            value test_result = test_fn(test_args);
-            bool is_selected = test_result->as_bool();
-            if constexpr (is_reject) is_selected = !is_selected;
-            if (is_selected) out->push_back(item);
-        }
-        return out;
-    } else {
-        throw raised_exception("selectattr: invalid number of arguments");
+        if constexpr (is_reject) is_selected = !is_selected;
+        if (is_selected) out->push_back(item);
    }
-
    return out;
 }

@@ -425,7 +405,12 @@ const func_builtins & global_builtins() {
            throw raised_exception("Jinja Exception: " + msg);
        }},
        {"dict", toobject},
-        {"namespace", toobject},
+        {"namespace", [](const func_args & args) -> value {
+            // unlike a dict, a namespace has no methods: ns.items is the attribute, as in jinja2
+            value out = toobject(args);
+            cast_val<value_object>(out)->has_builtins = false;
+            return out;
+        }},
        {"strftime_now", [](const func_args & args) -> value {
            args.ensure_vals<value_string>();
            std::string format = args.get_pos(0)->as_string().str();
@@ -1090,10 +1075,10 @@ const func_builtins & value_array_t::get_builtins() const {
            auto arr = slice(val->as_array(), start, stop, step);
            return is_val<value_tuple>(val) ? mk_val<value_tuple>(std::move(arr)) : mk_val<value_array>(std::move(arr));
        }},
-        {"selectattr", selectattr<false>},
-        {"select", selectattr<false>},
-        {"rejectattr", selectattr<true>},
-        {"reject", selectattr<true>},
+        {"selectattr", select_or_reject<false, true>},
+        {"select", select_or_reject<false, false>},
+        {"rejectattr", select_or_reject<true, true>},
+        {"reject", select_or_reject<true, false>},
        {"join", [](const func_args & args) -> value {
            args.ensure_count(1, 3);
            if (!is_val<value_array>(args.get_pos(0))) {
@@ -1455,7 +1440,10 @@ const func_builtins & value_undefined_t::get_builtins() const {
 //////////////////////////////////


-static value from_json(const common_json & j, bool mark_input) {
+// the runtime walks values with recursion, a deeper input overflows the stack
+static constexpr int FROM_JSON_MAX_DEPTH = 128;
+
+static value from_json(const common_json & j, bool mark_input, int depth = 0) {
    if (j.is_null()) {
        return mk_val<value_none>();
    } else if (j.is_boolean()) {
@@ -1470,16 +1458,18 @@ static value from_json(const common_json & j, bool mark_input) {
            str->mark_input();
        }
        return str;
+    } else if ((j.is_array() || j.is_object()) && depth >= FROM_JSON_MAX_DEPTH) {
+        throw std::invalid_argument("JSON is nested more than " + std::to_string(FROM_JSON_MAX_DEPTH) + " levels deep");
    } else if (j.is_array()) {
        auto arr = mk_val<value_array>();
        for (const auto & item : j) {
-            arr->push_back(from_json(item, mark_input));
+            arr->push_back(from_json(item, mark_input, depth + 1));
        }
        return arr;
    } else if (j.is_object()) {
        auto obj = mk_val<value_object>();
        for (auto it = j.begin(); it != j.end(); ++it) {
-            obj->insert(it.key(), from_json(it.value(), mark_input));
+            obj->insert(it.key(), from_json(it.value(), mark_input, depth + 1));
        }
        return obj;
    } else {
@@ -224,7 +224,21 @@ class common_chat_schema_builder {
        return node;
    }

+    // the builder and the code that walks the built schema recurse once per level, a deeper schema overflows the stack
+    static constexpr int max_depth = 128;
+    int                  depth_    = 0;
+
    common_chat_schema_ptr build_node(const common_json & schema, const std::string & path) {
+        if (depth_ >= max_depth) {
+            throw std::invalid_argument("JSON schema is nested more than " + std::to_string(max_depth) + " levels deep");
+        }
+        depth_++;
+        auto node = build_node_impl(schema, path);
+        depth_--;
+        return node;
+    }
+
+    common_chat_schema_ptr build_node_impl(const common_json & schema, const std::string & path) {
        if (!schema.is_object()) {
            fail(path, "schema must be an object");
        }
@@ -1,4 +1,5 @@
 #include "parsers.h"
+#include "log.h"

 // The DeepSeek V4 reference implementation renders consecutive tool results into a single
 // user block, ordered by the tool call order of the preceding assistant message (matched
@@ -75,13 +76,30 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
    // non-thinking generation prompt ends with a bare </think> instead of an empty
    // <think></think> pair.
    const bool is_v4 = tmpl.source().find("function_calls") == std::string::npos;
+    // V4.1: DSML tag names carry a leading space, e.g. <｜DSML｜ calls>
+    const bool is_v41 = is_v4 && tmpl.source().find("｜DSML｜ calls>") != std::string::npos;

    std::optional<json> adjusted_messages;
    if (is_v4) {
        adjusted_messages = deepseek_v4_sort_tool_results(inputs.messages);
    }

-    auto has_tools           = inputs.tools.is_array() && !inputs.tools.empty();
+    // V4.1 can declare tools on a system message instead of top-level "tools": parse calls to those too.
+    json tools = inputs.tools;
+    if (is_v41 && !(tools.is_array() && !tools.empty()) && inputs.messages.is_array()) {
+        for (const auto & msg : inputs.messages) {
+            if (msg.value("role", "") == "system" && msg.contains("tools") && msg.at("tools").is_array()) {
+                if (!tools.is_array()) {
+                    tools = json::array();
+                }
+                for (const auto & tool : msg.at("tools")) {
+                    tools.push_back(tool);
+                }
+            }
+        }
+    }
+
+    auto has_tools           = tools.is_array() && !tools.empty();
    auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object();
    auto extract_reasoning   = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
    auto include_grammar     = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE);
@@ -90,17 +108,39 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
    if (is_v4 && has_response_format) {
        additional_context = json{ { "response_format", inputs.json_schema } };
    }
+    // V4.1 raises on anything but low/high/max or 1-100: map minimal -> low, xhigh -> max, medium -> default
+    if (is_v41 && inputs.extra_context.is_object() && inputs.extra_context.contains("reasoning_effort") &&
+        inputs.extra_context.at("reasoning_effort").is_string()) {
+        const std::string effort = inputs.extra_context.at("reasoning_effort").get<std::string>();
+        json mapped;
+        if (effort == "minimal") {
+            mapped = "low";
+        } else if (effort == "xhigh") {
+            mapped = "max";
+        } else if (effort == "medium") {
+            mapped = nullptr;
+        }
+        if (effort == "minimal" || effort == "xhigh" || effort == "medium") {
+            LOG_DBG("%s: DeepSeek V4.1 reasoning_effort '%s' -> %s\n", __func__, effort.c_str(),
+                    mapped.is_null() ? "unset (template default)" : mapped.get<std::string>().c_str());
+            if (!additional_context.has_value()) {
+                additional_context = json::object();
+            }
+            (*additional_context)["reasoning_effort"] = mapped;
+        }
+    }

    const std::string DSML         = "｜DSML｜";
    const std::string THINK_START  = "<think>";
    const std::string THINK_END    = "</think>";
-    const std::string TC_BLOCK     = is_v4 ? "tool_calls" : "function_calls";
-    const std::string FC_START     = "<" + DSML + TC_BLOCK + ">";
-    const std::string FC_END       = "</" + DSML + TC_BLOCK + ">";
-    const std::string INVOKE_START = "<" + DSML + "invoke";
-    const std::string INVOKE_END   = "</" + DSML + "invoke>";
-    const std::string PARAM_START  = "<" + DSML + "parameter";
-    const std::string PARAM_END    = "</" + DSML + "parameter>";
+    const std::string TAG_SP       = is_v41 ? " " : "";
+    const std::string TC_BLOCK     = is_v41 ? "calls" : is_v4 ? "tool_calls" : "function_calls";
+    const std::string FC_START     = "<" + DSML + TAG_SP + TC_BLOCK + ">";
+    const std::string FC_END       = "</" + DSML + TAG_SP + TC_BLOCK + ">";
+    const std::string INVOKE_START = "<" + DSML + TAG_SP + "invoke";
+    const std::string INVOKE_END   = "</" + DSML + TAG_SP + "invoke>";
+    const std::string PARAM_START  = "<" + DSML + TAG_SP + "parameter";
+    const std::string PARAM_END    = "</" + DSML + TAG_SP + "parameter>";
    const std::string GEN_PROMPT   = "<｜Assistant｜>";
    const std::string TC_SEPARATOR = "\n\n";

@@ -152,7 +192,7 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
        // build tool call section first since we might need it in reasoning
        auto tool_choice = p.choice();
        if (has_tool_calls) {
-            foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
+            foreach_function(tools, [&](size_t tool_index, const json & tool) {
                const auto & function = tool.at("function");
                std::string  name     = function.at("name");

@@ -223,7 +263,10 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
        }

        if (extract_reasoning && inputs.enable_thinking) {
-            reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END);
+            // V4.1 renders a continuation without reasoning as "<｜Assistant｜></think>" even in thinking mode,
+            // so consume the bare THINK_END too or it leaks into the content
+            reasoning = is_v41 ? p.optional((THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END) | p.literal(THINK_END))
+                               : p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END);
            reasoning_with_tc = THINK_START +
                p.reasoning(p.until_one_of({ TC_SEPARATOR + FC_START, FC_START, THINK_END })) +
                p.space() + obligatory_tool_calls;
@@ -19,6 +19,7 @@
 #include <fstream>
 #include <iomanip>
 #include <map>
+#include <random>
 #include <cinttypes>

 #define SPC_DBG(fmt, ...) LOG_DBG("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
@@ -1010,6 +1011,10 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
    bool    is_mrope       = false;
    int32_t selector_top_k = 0;

+    // DFlash2 probabilistic drafting: one generator per sequence and the seed it was last seeded with (-1: reseed)
+    std::vector<std::mt19937> dflash2_rng;
+    std::vector<int64_t>      dflash2_seed;
+
    // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout
    const bool is_dspark;

@@ -1060,6 +1065,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

        selector_top_k = llama_model_dflash_selector_top_k(model_dft);
        is_dflash2     = selector_top_k > 0;
+        dflash2_rng.resize(n_seq);
+        dflash2_seed.assign(n_seq, -1);
        char lilicorr_meta[32] = {};
        lilicorr_top_k = llama_model_meta_val_str(model_dft, "lilicorr.candidate_topk", lilicorr_meta, sizeof(lilicorr_meta)) >= 0
                ? std::atoi(lilicorr_meta) : 0;
@@ -1158,6 +1165,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
            return;
        }

+        // reseed per request rather than per round, or two identical requests differ
+        dflash2_seed[seq_id] = -1;
+
        const int32_t N = (int32_t) prompt.size();
        if (N <= 0) {
            return;
@@ -1338,6 +1348,21 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
                const float * lattice = llama_get_embeddings_nextn(ctx_dft);
                GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");

+                // DFlash2 can instead sample each step and hand q to the verifier for rejection sampling
+                const bool sample = is_dflash2 && params.probabilistic && dp.result_q;
+
+                // greedy keeps temperature 1, sampling uses the target's
+                const float temp = sample ? dp.temp : 1.0f;
+
+                if (sample && dflash2_seed[seq_id] != dp.seed) {
+                    dflash2_seed[seq_id] = dp.seed;
+                    // the default seed means random; mixed so it differs from the target's, as in spec_retune
+                    dflash2_rng[seq_id].seed(dp.seed == LLAMA_DEFAULT_SEED ? std::random_device{}() : dp.seed ^ 0x85ebca6bu);
+                }
+
+                std::vector<float>   probs;
+                std::vector<int32_t> sorted_idx;
+
                int32_t predecessor = 0;
                for (int32_t i = 1; i < (is_dflash2 ? n_block_tokens : params.n_max + 1); ++i) {
                    const float * row = lattice + (size_t) (beg + i) * n_embd_dec;
@@ -1346,21 +1371,53 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

                    predecessor = (int32_t) std::distance(scores,
                            std::max_element(scores, scores + top_k));
-                    if (params.p_min > 0.0f) {
-                        // softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max))
+                    if (params.p_min > 0.0f || sample) {
+                        // softmax(scores / temp) at the argmax, i.e. 1 / sum(exp((s_k - s_max) / temp))
+                        probs.resize(top_k);
                        float sum = 0.0f;
                        for (int32_t k = 0; k < top_k; ++k) {
-                            sum += std::exp(scores[k] - scores[predecessor]);
+                            probs[k] = std::exp((scores[k] - scores[predecessor]) / temp);
+                            sum += probs[k];
                        }
                        if (1.0f / sum < params.p_min) {
                            break;
                        }
+                        for (int32_t k = 0; k < top_k; ++k) {
+                            probs[k] /= sum;
+                        }
+                    }
+                    if (sample) {
+                        // ggml_top_k leaves candidates unordered, so walk them by token id to keep a seed's draft stable
+                        sorted_idx.resize(top_k);
+                        for (int32_t k = 0; k < top_k; ++k) {
+                            sorted_idx[k] = k;
+                        }
+                        std::sort(sorted_idx.begin(), sorted_idx.end(), [row](int32_t a, int32_t b) { return row[a] < row[b]; });
+
+                        // a round-off overshoot keeps the argmax, which has q > 0
+                        double u = std::uniform_real_distribution<double>(0.0, 1.0)(dflash2_rng[seq_id]);
+                        for (const int32_t k : sorted_idx) {
+                            u -= probs[k];
+                            if (u < 0.0) {
+                                predecessor = k;
+                                break;
+                            }
+                        }
+
+                        auto & q = dp.result_q->emplace_back();
+                        q.reserve(top_k);
+                        for (int32_t k = 0; k < top_k; ++k) {
+                            q.push_back({ (llama_token) row[k], scores[k], probs[k] });
+                        }
                    }
                    result.push_back((llama_token) row[predecessor]);
                }

                if (result.size() < (size_t) params.n_min) {
                    result.clear();
+                    if (sample) {
+                        dp.result_q->clear();
+                    }
                }
                continue;
            }
@@ -73,7 +73,7 @@ struct common_speculative_draft_params {
    // the generated draft from the last _draft() call
    llama_tokens * result;

-    // candidate distribution per drafted token; set it to make draft-simple and draft-mtp sample
+    // candidate distribution per drafted token; set it to make draft-simple, draft-mtp and draft-dflash (DFlash2 only) sample
    std::vector<std::vector<llama_token_data>> * result_q = nullptr;

    // the target's temp and seed, read only when the drafter samples probabilistically
@@ -17,8 +17,10 @@ __all__ = [


 TEXT_MODEL_MAP: dict[str, str] = {
+    "NEOChatModel": "sensenova_u1",
    "AfmoeForCausalLM": "afmoe",
    "LagunaForCausalLM": "laguna",
+    "FlintForCausalLM": "flint",
    "ApertusForCausalLM": "llama",
    "ArceeForCausalLM": "llama",
    "ArcticForCausalLM": "arctic",
@@ -68,6 +70,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "LingDSparkModel": "qwen",
    "DeepseekV4ForCausalLM": "deepseek",
    "DeepseekV4DSparkModel": "deepseek",
+    "DeepseekV41ForCausalLM": "deepseek41",
    "DistilBertForMaskedLM": "bert",
    "DistilBertForSequenceClassification": "bert",
    "DistilBertModel": "bert",
@@ -194,6 +197,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "MiniCPMForCausalLM": "minicpm",
    "MiniCPMV4_6ForConditionalGeneration": "minicpm",
    "MiniCPMV4_7ForConditionalGeneration": "minicpm",
+    "MiniCPMOcrForConditionalGeneration": "minicpm",
    "MiniMaxText01ForCausalLM": "minimax",
    "MiniMaxM1ForCausalLM": "minimax",
    "MiniMaxM2ForCausalLM": "minimax",
@@ -283,6 +287,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "SeedOssForCausalLM": "olmo",
    "SmallThinkerForCausalLM": "smallthinker",
    "SmolLM3ForCausalLM": "llama",
+    "SopranoTextModel": "soprano",
    "Spark2_5ForCausalLM": "spark2_5",
    "SolarOpenForCausalLM": "glm",
    "StableLMEpochForCausalLM": "stablelm",
@@ -322,6 +327,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "VllmD3Model": "vllm_d3",
    "DeepseekOCR2ForCausalLM": "deepseek",
    "DeepseekOCRForCausalLM": "deepseek",
+    "DeepseekV41ForCausalLM": "deepseek41",
    "DeepseekV4ForCausalLM": "deepseek",
    "Dots3NoteForCausalLM": "dots3",
    "Dots3NoteForConditionalGeneration": "dots3",
@@ -359,6 +365,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "MiniMaxM3SparseForConditionalGeneration": "minimax",
    "MiniCPMV4_6ForConditionalGeneration": "minicpm",
    "MiniCPMV4_7ForConditionalGeneration": "minicpm",
+    "MiniCPMOcrForConditionalGeneration": "minicpm",
    "Mistral3ForConditionalGeneration": "llava",
    "Mistral3Model": "llava",
    "NemotronH_Omni_Reasoning_V3": "nemotron",
@@ -367,6 +374,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "MuseSparkForConditionalGeneration": "muse_spark",
    "PaddleOCRVisionModel": "ernie",
    "Phi4ForCausalLMV": "phi",
+    "NEOVisionModel": "sensenova_u1",
    "Qwen2AudioForConditionalGeneration": "ultravox",
    "Qwen2VLForConditionalGeneration": "qwenvl",
    "Qwen2VLModel": "qwenvl",
@@ -375,6 +383,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "Qwen3ASRForConditionalGeneration": "qwen3vl",
    "Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
    "PocketTTSModel": "pockettts",
+    "SopranoModel": "soprano",
    "Qwen3TTSForConditionalGeneration": "qwen3tts",
    "Qwen3VLForConditionalGeneration": "qwen3vl",
    "Qwen3VLMoeForConditionalGeneration": "qwen3vl",
@@ -1701,7 +1701,7 @@ class TextModel(ModelBase):
            self.gguf_writer.add_expert_group_used_count(n_group_used)
            logger.info(f"gguf: expert groups used count = {n_group_used}")

-        if (score_func := self.find_hparam(["score_function", "scoring_func", "score_func", "moe_router_activation", "moe_router_activation_func", "expert_selection_fn", "router_score_func"], optional=True)) is not None:
+        if (score_func := self.find_hparam(["score_function", "scoring_func", "score_func", "moe_router_activation", "moe_router_activation_func", "expert_selection_fn", "router_score_func", "router_score_function"], optional=True)) is not None:
            if score_func == "sigmoid":
                self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
            elif score_func == "softmax":
@@ -1837,6 +1837,9 @@ class TextModel(ModelBase):
        if chkhsh == "1431a23e583c97432bc230bff598d103ddb5a1f89960c8f1d1051aaa944d0b35":
            # ref: https://huggingface.co/sapienzanlp/Minerva-7B-base-v1.0
            res = "minerva-7b"
+        if chkhsh == "6e4d72c6e5260b90606e26a0931fa15e4b42ec6e97719d12919f2b1f93478262":
+            # ref: https://huggingface.co/domyn/Domyn-Small-v1.0
+            res = "domyn-small"
        if chkhsh == "7e57df22b1fe23a7b1e1c7f3dc4e3f96d43a4eb0836d0c6bdc3436d7b2f1c664":
            # ref: https://huggingface.co/tencent/Hunyuan-A13B-Instruct
            res = "hunyuan"
@@ -1939,6 +1942,9 @@ class TextModel(ModelBase):
        if chkhsh == "52df12b4c8d4176e7481aab4b6e8454d1fd0a210a04a574f6d4e067d10e23c3e":
            # ref: https://huggingface.co/CohereLabs/North-Mini-Code-1.0
            res = "cohere2moe"
+        if chkhsh == "d5420e601e979078c670bc503d9ac411c311fa76b76a76267c31e0d6ba816637":
+            # ref: https://huggingface.co/RecursiveSI/flint
+            res = "flint"
        if chkhsh == "e636dc30a262dcc0d8c323492e32ae2b70728f4df7dfe9737d9f920a282b8aea":
            # ref: https://huggingface.co/Qwen/Qwen1.5-7B
            res = "qwen2"
@@ -3143,7 +3149,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 ("StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"):
+    if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"):
        return arch

    # if "architectures" is found in the sub-config, use that instead
@@ -0,0 +1,467 @@
+from __future__ import annotations
+
+import json
+import math
+import re
+from pathlib import Path
+from typing import Any, Callable, Iterable, TYPE_CHECKING, cast
+
+import numpy as np
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import LazyTorchTensor, ModelBase, TextModel, gguf
+from .deepseek import DeepseekV4FlashVisionModel, DeepseekV4Model
+
+
+# Engram hash constants from inference/engram.py, stored in the GGUF: the token map needs the HF tokenizers
+# normalizers and the multipliers numpy's PCG64 stream, neither of which the runtime has.
+
+def _is_prime(n: int) -> bool:
+    if n < 2 or n % 2 == 0:
+        return n == 2
+    return all(n % d for d in range(3, math.isqrt(n) + 1, 2))
+
+
+def engram_primes(n_layers: int, max_ngram_size: int, n_heads: int, vocab_size: int) -> list[list[list[int]]]:
+    """EngramLayout.from_args: [layer][ngram order 2..max][head] bucket moduli, never reused across layers."""
+    seen: set[int] = set()
+    out: list[list[list[int]]] = []
+    for _ in range(n_layers):
+        per_ngram = []
+        for _ in range(max_ngram_size - 1):
+            sizes, cur = [], vocab_size - 1
+            for _ in range(n_heads):
+                cand = cur + 1
+                while not _is_prime(cand) or cand in seen:
+                    cand += 1
+                cur = cand
+                seen.add(cur)
+                sizes.append(cur)
+            per_ngram.append(sizes)
+        out.append(per_ngram)
+    return out
+
+
+def engram_multipliers(layer_ids: list[int], max_ngram_size: int, compressed_vocab_size: int) -> list[list[int]]:
+    """compute_hash_multipliers: odd, bounded so that token_id * multiplier cannot overflow int64."""
+    bound = max(1, (np.iinfo(np.int64).max // compressed_vocab_size) // 2)
+    rows = []
+    for layer_id in layer_ids:
+        rng = np.random.default_rng(10007 * layer_id)
+        values = rng.integers(low=0, high=bound, size=(max_ngram_size,), dtype=np.int64)
+        rows.append([int(v) * 2 + 1 for v in values])
+    return rows
+
+
+def engram_compressed_token_map(tokenizer_dir: Path) -> tuple[list[int], int]:
+    """build_compressed_token_map: token id -> compressed id (tokens that normalize alike collapse)."""
+    from tokenizers import Regex, normalizers
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
+    sentinel = "\ue000"
+    normalizer = normalizers.Sequence([
+        normalizers.NFKC(),
+        normalizers.NFD(),
+        normalizers.StripAccents(),
+        normalizers.Lowercase(),
+        normalizers.Replace(Regex(r"[ \t\r\n]+"), " "),
+        normalizers.Replace(Regex(r"^ $"), sentinel),
+        normalizers.Strip(),
+        normalizers.Replace(sentinel, " "),
+    ])
+    backend = tokenizer.backend_tokenizer
+    key_to_new: dict[str, int] = {}
+    lookup = [0] * len(tokenizer)
+    for token_id in range(len(tokenizer)):
+        text = backend.decode([token_id], skip_special_tokens=False)
+        if "\ufffd" in text:
+            key = backend.id_to_token(token_id)
+        else:
+            normalized = normalizer.normalize_str(text)
+            key = normalized if normalized else text
+        lookup[token_id] = key_to_new.setdefault(key, len(key_to_new))
+    return lookup, len(key_to_new)
+
+
+class _StreamedBytes:
+    """Raw tensor bytes produced chunk by chunk at write time (an Engram table is ~200 GB in BF16)."""
+    def __init__(self, shape: tuple[int, ...], chunks: Callable[[], Iterable[np.ndarray]]):
+        self.shape = shape
+        self.dtype = np.dtype(np.uint8)
+        self.nbytes = int(np.prod(shape))
+        self._chunks = chunks
+
+    def tofile(self, fout) -> None:
+        written = 0
+        for chunk in self._chunks():
+            chunk = np.ascontiguousarray(chunk, dtype=np.uint8)
+            chunk.tofile(fout)
+            written += chunk.nbytes
+        if written != self.nbytes:
+            raise ValueError(f"streamed {written} bytes, expected {self.nbytes}")
+
+    def byteswap(self, inplace: bool = False):
+        raise NotImplementedError("streamed tensors are written little-endian only")
+
+
+def _e8m0_check(scale_u8: np.ndarray, name: str) -> None:
+    # e4m3 (|x| in [2^-9, 448]) * 2^(e - 127) is exact in BF16 for e in [3, 246]
+    if scale_u8.size and (int(scale_u8.min()) < 3 or int(scale_u8.max()) > 246):
+        raise ValueError(f"{name}: E8M0 scale exponents outside [3, 246], the dequantized values would not be exact")
+
+
+@ModelBase.register("DeepseekV41ForCausalLM")
+class DeepseekV41Model(DeepseekV4Model):
+    """DeepSeek-V4.1 text model, lossless: FP8 (e4m3, 32x32 ue8m0 blocks) linears dequantized to BF16 (F32 with
+    --outtype f32), routed experts copied bit-exact as MXFP4, Engram tables (e4m3 rows, one e8m0 scale per 32) as the
+    exact BF16 lookup (F32 with --outtype f32), router and Engram gate F32. DSpark (mtp.*) is not converted here
+    (--dspark, DeepseekV41DSparkModel); the vision tower goes to an mmproj (--mmproj, DeepseekV41VisionModel).
+    The Engram tables of all layers are written as one per_layer_token_embd (rows concatenated in layer order) with
+    the PLE keys of qwen4exp: absolute head offsets and bucket counts, so the runtime can read it lazily."""
+    model_arch = gguf.MODEL_ARCH.DEEPSEEK41
+    supports_mtp_export = False
+
+    _ENGRAM_EMBED_RE = re.compile(r"layers\.(\d+)\.engram\.embed\.(weight|scale)$")
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.hparams.setdefault("num_hash_layers", 0)  # no hash-routed layers in V4.1
+        self._f32 = self.ftype == gguf.LlamaFileType.ALL_F32
+        qc = self.hparams.get("quantization_config") or {}
+        self._fp8_block = tuple(qc.get("weight_block_size") or (32, 32))
+        if qc.get("scale_fmt", "ue8m0") != "ue8m0":
+            raise ValueError(f"unsupported FP8 scale format {qc.get('scale_fmt')!r}")
+        if self.hparams.get("gate_temp", 1.0) != 1.0:
+            raise ValueError("gate_temp != 1 is not supported (the router logits would need dividing by it)")
+
+        # V4.1 does not use the DeepSeek-V4 template; a template shipped with the model is still added by set_vocab
+        self.gguf_writer.remove_key(gguf.Keys.Tokenizer.CHAT_TEMPLATE)
+
+    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+        # num_hidden_layers lives in text_config, which is merged only after indexing; MTP stages are skipped anyway
+        return TextModel.index_tensors(self, remote_hf_model_id=remote_hf_model_id)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        if item[0].startswith("vision."):
+            return None
+        return super().filter_tensors(item)
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        hp = self.hparams
+        w = self.gguf_writer
+        w.add_attention_kv_source_layers(hp["kv_source_layer_ids"])
+        w.add_indexer_source_layers(hp["index_source_layer_ids"])
+        if hp.get("candidate_source_layer_id", -1) >= 0:
+            w.add_indexer_candidate_source_layer(hp["candidate_source_layer_id"])
+            w.add_indexer_candidate_top_k_blocks(hp["candidate_topk_blocks"])
+            w.add_indexer_candidate_block_size(hp["candidate_block_size"])
+
+        # GGUF cannot hold empty arrays: the Engram keys are omitted when the model has no Engram layers
+        layer_ids = list(hp.get("engram_layer_ids") or [])
+        if layer_ids:
+            max_ngram, n_heads = hp["engram_max_ngram_size"], hp["engram_n_heads"]
+            primes = engram_primes(len(layer_ids), max_ngram, n_heads, hp["engram_vocab_size"])
+            for li, layer_primes in enumerate(primes):
+                if sum(p for per in layer_primes for p in per) != hp["engram_num_embeddings"][li]:
+                    raise ValueError(f"engram layer {layer_ids[li]}: prime sum != num_embeddings")
+            token_map, cvocab = engram_compressed_token_map(self.dir_model)
+            if cvocab != hp["engram_compressed_vocab_size"] or len(token_map) != hp["vocab_size"]:
+                raise ValueError(f"compressed token map ({len(token_map)} -> {cvocab}) does not match the config")
+            # absolute row offsets into the concatenated table
+            sizes = [p for layer in primes for per in layer for p in per]
+            offsets = [sum(sizes[:i]) for i in range(len(sizes))]
+            w.add_key_value(gguf.Keys.PerLayerEmbedding.LAYERS.format(arch=w.arch), layer_ids,
+                            gguf.GGUFValueType.ARRAY, gguf.GGUFValueType.UINT32)
+            w.add_ple_ngram_size(max_ngram)
+            w.add_ple_heads_per_ngram(n_heads)
+            w.add_embedding_length_per_layer_input(hp["engram_head_dim"])
+            w.add_ple_layer_multipliers([m for row in engram_multipliers(layer_ids, max_ngram, cvocab) for m in row])
+            w.add_ple_head_offsets(offsets)
+            w.add_ple_head_vocab_sizes(sizes)
+            w.add_engram_pad_token_id(hp["engram_pad_token_id"])
+            w.add_engram_token_map(token_map)
+
+    def dequant_model(self):
+        bo, bi = self._fp8_block
+
+        def dequant(weight: Tensor, scale: Tensor) -> Tensor:
+            out_f, in_f = weight.shape
+            s = self._e8m0_to_float(scale).repeat_interleave(bo, 0)[:out_f].repeat_interleave(bi, 1)[:, :in_f]
+            return weight.float() * s
+
+        for name in [n for n in self.model_tensors if n.endswith(".scale") and not self._ENGRAM_EMBED_RE.match(n)]:
+            weight_name = name.removesuffix(".scale") + ".weight"
+            if weight_name not in self.model_tensors or self.model_tensors[weight_name]().dtype not in self._float8_dtypes():
+                continue
+            w, s = self.model_tensors[weight_name], self.model_tensors.pop(name)
+            out_f, in_f = w().shape
+            if tuple(s().shape) != (-(-out_f // bo), -(-in_f // bi)):
+                raise ValueError(f"{weight_name}: scale shape {tuple(s().shape)} does not match {(bo, bi)} blocks")
+            _e8m0_check(LazyTorchTensor.to_eager(s()).view(torch.uint8).numpy(), name)
+            self.model_tensors[weight_name] = lambda w=w, s=s: dequant(w(), s())
+            self._dsv4_fp8_dequantized.add(weight_name)
+
+    def _tensor_part(self, name: str) -> Path:
+        index = self.dir_model / "model.safetensors.index.json"
+        if not index.is_file():
+            return self.dir_model / "model.safetensors"
+        with open(index, "r", encoding="utf-8") as f:
+            return self.dir_model / json.load(f)["weight_map"][name]
+
+    def _write_mxfp4_experts(self, bid: int, proj: str, key: gguf.MODEL_TENSOR, n_expert: int | None = None) -> None:
+        n_expert = self.hparams["n_routed_experts"] if n_expert is None else n_expert
+        names = [f"layers.{bid}.ffn.experts.{e}.{proj}" for e in range(n_expert)]
+        pairs = [(self.model_tensors.pop(n + ".weight"), self.model_tensors.pop(n + ".scale")) for n in names]
+        rows, packed_cols = pairs[0][0]().shape
+
+        def chunks() -> Iterable[np.ndarray]:
+            for w, s in pairs:
+                yield self.repack_mxfp4_blocks(LazyTorchTensor.to_eager(w()), LazyTorchTensor.to_eager(s()))
+        data = cast(np.ndarray, _StreamedBytes((len(pairs), rows, packed_cols * 2 // 32 * 17), chunks))
+        self.gguf_writer.add_tensor(self.format_tensor_name(key, bid), data, raw_dtype=gguf.GGMLQuantizationType.MXFP4)
+
+    def _write_engram_tables(self, layer_ids: list[int]) -> None:
+        # official lookup e4m3 * 2^(e8m0 - 127), exact in BF16
+        parts, dim = [], None
+        for bid in layer_ids:
+            wn, sn = f"layers.{bid}.engram.embed.weight", f"layers.{bid}.engram.embed.scale"
+            rows, d = self.model_tensors.pop(wn)().shape
+            if d % 32 != 0 or tuple(self.model_tensors.pop(sn)().shape) != (rows, d // 32) or d != (dim or d):
+                raise ValueError(f"{wn}: expected one e8m0 scale per 32 elements and the same width in every layer")
+            if rows != self.hparams["engram_num_embeddings"][len(parts)]:
+                raise ValueError(f"{wn}: {rows} rows, config says {self.hparams['engram_num_embeddings'][len(parts)]}")
+            dim = d
+            parts.append((wn, sn, rows, self._tensor_part(wn), self._tensor_part(sn)))
+        assert dim is not None
+        name = self.format_tensor_name(gguf.MODEL_TENSOR.PER_LAYER_TOKEN_EMBD)
+        n_rows = sum(p[2] for p in parts)
+
+        def chunks() -> Iterable[np.ndarray]:
+            from safetensors import safe_open
+            for wn, sn, rows, wpart, spart in parts:
+                with safe_open(wpart, framework="pt") as fw, safe_open(spart, framework="pt") as fs:
+                    ws, ss = fw.get_slice(wn), fs.get_slice(sn)
+                    for r0 in range(0, rows, 1 << 18):
+                        q = ws[r0:r0 + (1 << 18)].view(torch.uint8).reshape(-1, 32)
+                        e = ss[r0:r0 + (1 << 18)].view(torch.uint8).reshape(-1, 1)
+                        _e8m0_check(e.numpy(), sn)
+                        v = (q.view(torch.float8_e4m3fn).float() * torch.exp2(e.float() - 127)).reshape(-1, dim)
+                        yield (v if self._f32 else v.bfloat16()).view(torch.uint8).numpy()
+        qt = gguf.GGMLQuantizationType.F32 if self._f32 else gguf.GGMLQuantizationType.BF16
+        data = cast(np.ndarray, _StreamedBytes((n_rows, dim * (4 if self._f32 else 2)), chunks))
+        self.gguf_writer.add_tensor(name, data, raw_dtype=qt)
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        if not self._dsv4_mxfp4_generated:
+            for bid in range(self.block_count):
+                self._write_mxfp4_experts(bid, "w1", gguf.MODEL_TENSOR.FFN_GATE_EXP)
+                self._write_mxfp4_experts(bid, "w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP)
+                self._write_mxfp4_experts(bid, "w3", gguf.MODEL_TENSOR.FFN_UP_EXP)
+            if self.hparams.get("engram_layer_ids"):
+                self._write_engram_tables(list(self.hparams["engram_layer_ids"]))
+            self._dsv4_mxfp4_generated = True
+        return ()
+
+    _LAYER_MAP: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
+        "attn.indexer.wk.weight": (gguf.MODEL_TENSOR.INDEXER_ATTN_K, ".weight"),
+        "attn.indexer.k_norm.weight": (gguf.MODEL_TENSOR.INDEXER_K_NORM, ".weight"),
+        "engram.wkv.weight": (gguf.MODEL_TENSOR.ENGRAM_KV, ".weight"),
+    }
+
+    def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_TENSOR, str]:
+        m = re.match(r"layers\.\d+\.(.+)$", name)
+        if m is not None and m.group(1) in self._LAYER_MAP:
+            return self._LAYER_MAP[m.group(1)]
+        return super()._map_dsv4_tensor_name(name, bid)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.endswith(".engram.k_weight"):
+            return []  # folded into engram_gate
+        if name.endswith(".engram.q_weight"):
+            # Engram gate weight = q_weight * k_weight (bf16 * bf16 is exact in f32)
+            k = self.model_tensors[name.removesuffix("q_weight") + "k_weight"]
+            return [(self.format_tensor_name(gguf.MODEL_TENSOR.ENGRAM_GATE, bid), data_torch.float() * k().float())]
+        return 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:
+        if new_name.endswith((".engram_gate.weight", ".ffn_gate_inp.weight")) or name in self._dsv4_f32_tensors:
+            return gguf.GGMLQuantizationType.F32
+        if n_dims >= 2 and (name in self._dsv4_fp8_dequantized or name in self._dsv4_bf16_tensors):
+            return gguf.GGMLQuantizationType.F32 if self._f32 else gguf.GGMLQuantizationType.BF16
+        return False
+
+
+class DeepseekV41DSparkModel(DeepseekV41Model):
+    """DeepSeek-V4.1 DSpark speculative draft (mtp.0..2 of the checkpoint) as a DFLASH GGUF (`--dspark`).
+
+    Layout = the V4 DSpark draft (conversion/deepseek.py DeepseekV4DSparkModel) with V4.1's differences flagged by
+    `dflash.hyper_connection.single_pass`: stages blk.0..2 (MXFP4 experts bit-exact, dense BF16), mtp.0.main_proj /
+    main_norm -> fc / enc.output_norm, mtp.2.norm -> output_norm, markov_head.embed / head -> markov_w1 / markov_w2,
+    confidence_head.proj -> conf_proj (F32), no hc_head. Token embeddings and the lm head are the target's (shared at
+    runtime through ctx_other). target_layers are the V4.1 ids as-is: the draft reads the attention input of those
+    layers, which is what llama_set_embeddings_layer_inp extracts.
+    """
+    model_arch = gguf.MODEL_ARCH.DFLASH
+    mtp_only = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        hp = self.hparams
+        self._n_main = hp["num_hidden_layers"]
+        self._n_stage = hp["num_nextn_predict_layers"]
+        if self._n_stage <= 0 or "dspark_block_size" not in hp:
+            raise ValueError("DeepSeek-V4.1 DSpark export needs num_nextn_predict_layers > 0 and the dspark_* config")
+        # stages keep main-model block ids (mtp.M -> layers.n_main+M); format_tensor_name() maps them back
+        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self._n_stage)
+
+    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+        hp = {**self.hparams, **self.hparams.get("text_config", {})}  # text_config is merged only after indexing
+        type(self)._dsv4_main_layers = hp["num_hidden_layers"]
+        type(self)._dsv4_nextn_layers = hp.get("num_nextn_predict_layers", 0)
+        return TextModel.index_tensors(self, remote_hf_model_id=remote_hf_model_id)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        m = re.match(r"mtp\.(\d+)\.(.+)$", name)
+        if m is None:
+            return None  # shared with the target at runtime
+        stage, rest = int(m.group(1)), m.group(2)
+        assert cls._dsv4_main_layers is not None
+        if stage >= cls._dsv4_nextn_layers:
+            raise ValueError(f"Unexpected DSpark stage {stage}")
+        if rest.startswith(("markov_head.", "confidence_head.")):
+            return f"dspark.{rest}", gen
+        return f"layers.{cls._dsv4_main_layers + stage}.{rest}", gen
+
+    def format_tensor_name(self, key: gguf.MODEL_TENSOR, bid: int | None = None, suffix: str = ".weight") -> str:
+        if bid is not None and bid >= self._n_main:
+            bid -= self._n_main
+        return super().format_tensor_name(key, bid, suffix)
+
+    def set_vocab(self):
+        super().set_vocab()
+        self.gguf_writer.add_mask_token_id(self.hparams["dspark_noise_token_id"])
+
+    def set_gguf_parameters(self):
+        hp = self.hparams
+        n = self._n_stage
+        w = self.gguf_writer
+        w.add_block_count(n)
+        w.add_context_length(hp["max_position_embeddings"])
+        w.add_embedding_length(hp["hidden_size"])
+        w.add_feed_forward_length(hp["moe_intermediate_size"])
+        w.add_head_count(hp["num_attention_heads"])
+        w.add_head_count_kv(1)
+        w.add_key_length(hp["head_dim"])
+        w.add_value_length(hp["head_dim"])
+        w.add_rope_dimension_count(hp["qk_rope_head_dim"])
+        # compress-ratio-0 layers: plain rope at rope_theta, no YaRN
+        w.add_rope_freq_base(float(hp["rope_theta"]))
+        w.add_layer_norm_rms_eps(hp["rms_norm_eps"])
+        w.add_q_lora_rank(hp["q_lora_rank"])
+        w.add_sliding_window(hp["sliding_window"])
+        w.add_expert_count(hp["dspark_n_routed_experts"])
+        w.add_expert_used_count(hp["dspark_num_experts_per_tok"])
+        w.add_expert_feed_forward_length(hp["moe_intermediate_size"])
+        w.add_expert_shared_count(hp["n_shared_experts"])
+        w.add_expert_weights_scale(hp["routed_scaling_factor"])
+        w.add_expert_weights_norm(hp["norm_topk_prob"])
+        w.add_expert_gating_func(gguf.ExpertGatingFuncType.SQRTSOFTPLUS)
+        w.add_swiglu_clamp_exp([hp["swiglu_limit"]] * n)
+        w.add_swiglu_clamp_shexp([hp["swiglu_limit"]] * n)
+        w.add_attention_output_group_count(hp["o_groups"])
+        w.add_attention_output_lora_rank(hp["o_lora_rank"])
+        w.add_attention_compress_ratios([0] * n)
+        w.add_hyper_connection_count(hp["hc_mult"])
+        w.add_hyper_connection_sinkhorn_iterations(hp["hc_sinkhorn_iters"])
+        w.add_hyper_connection_epsilon(hp["hc_eps"])
+        w.add_hyper_connection_single_pass(True)
+        w.add_block_size(hp["dspark_block_size"])
+        w.add_target_layers(hp["dspark_target_layer_ids"])
+        w.add_sample_from_anchor(True)
+        w.add_has_confidence_head(True)
+        w.add_file_type(self.ftype)
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # no Engram tables in the draft
+        if not self._dsv4_mxfp4_generated:
+            n_expert = self.hparams["dspark_n_routed_experts"]
+            for bid in range(self._n_main, self._n_main + self._n_stage):
+                self._write_mxfp4_experts(bid, "w1", gguf.MODEL_TENSOR.FFN_GATE_EXP, n_expert)
+                self._write_mxfp4_experts(bid, "w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP, n_expert)
+                self._write_mxfp4_experts(bid, "w3", gguf.MODEL_TENSOR.FFN_UP_EXP, n_expert)
+            self._dsv4_mxfp4_generated = True
+        return ()
+
+    _DSPARK_ROOT: dict[str, gguf.MODEL_TENSOR] = {
+        "main_proj.weight": gguf.MODEL_TENSOR.FC,
+        "main_norm.weight": gguf.MODEL_TENSOR.ENC_OUTPUT_NORM,
+        "norm.weight": gguf.MODEL_TENSOR.OUTPUT_NORM,
+        "dspark.markov_head.embed.weight": gguf.MODEL_TENSOR.DSPARK_MARKOV_W1,
+        "dspark.markov_head.head.weight": gguf.MODEL_TENSOR.DSPARK_MARKOV_W2,
+        "dspark.confidence_head.proj.weight": gguf.MODEL_TENSOR.DSPARK_CONF_PROJ,
+    }
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.endswith(".ffn.gate.bias_vl"):
+            return []
+        if name in self._DSPARK_ROOT:
+            return [(self.format_tensor_name(self._DSPARK_ROOT[name], None), data_torch)]
+        m = re.match(r"layers\.\d+\.(main_proj\.weight|main_norm\.weight|norm\.weight)$", name)
+        if m is not None:
+            return [(self.format_tensor_name(self._DSPARK_ROOT[m.group(1)], None), data_torch)]
+        return 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:
+        if new_name == "conf_proj.weight":
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+
+@ModelBase.register("DeepseekV41ForCausalLM")
+@ModelBase.example("deepseek-ai/DeepSeek-V4.1-Flash-1017")
+class DeepseekV41VisionModel(DeepseekV4FlashVisionModel):
+    """DeepSeek-V4.1 vision tower (projector deepseek41v): the V4-Flash-Vision ViT + 3x3 aligner, with the V4.1 token
+    block ([START] ([IMAGE]*w [NEWLINE])*h [END], no pads) and resize plan. ref: inference/vision.py,
+    inference/image_processor.py of deepseek-ai/DeepSeek-V4.1-Flash-1017. The checkpoint config is HF-style
+    (vision_config); image_start / image_end / image_newline are written like V4's sentinels, there is no image_pad."""
+
+    def get_vision_config(self) -> dict[str, Any] | None:
+        vc = self.global_config.get("vision_config")
+        if not vc or vc.get("num_hidden_layers", 0) == 0:
+            raise ValueError("DeepseekV41VisionModel requires vision_config.num_hidden_layers > 0")
+        return {
+            "num_hidden_layers":   vc["num_hidden_layers"],
+            "hidden_size":         vc["hidden_size"],
+            "num_attention_heads": vc["num_attention_heads"],
+            "intermediate_size":   vc["intermediate_size"],
+            "patch_size":          vc["patch_size"],
+            # dynamic resolution; only used for compat / warmup
+            "image_size":          vc["patch_size"] * vc["downsample_ratio"] * 16,
+            "rope_theta":          vc.get("rope_theta", 10000.0),
+            "downsample_ratio":    vc["downsample_ratio"],
+            "min_pixels":          vc["min_pixels"],
+            "max_image_tokens":    vc["max_image_tokens"],
+            "max_wh_ratio":        vc.get("max_wh_ratio"),
+        }
+
+    def set_gguf_parameters(self):
+        assert self.hparams_vision is not None
+        super(DeepseekV4FlashVisionModel, self).set_gguf_parameters()
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.DEEPSEEK41V)
+        # pytorch default eps (vision.py), not the LLM rms_norm_eps (1e-20)
+        self.gguf_writer.add_vision_attention_layernorm_eps(1e-6)
+        self.gguf_writer.add_vision_use_silu(True)
+        self.gguf_writer.add_vision_projector_scale_factor(self.hparams_vision["downsample_ratio"])
+        self.gguf_writer.add_vision_min_pixels(self.hparams_vision["min_pixels"])
+        # hardcoded in clip.cpp (PROJECTOR_TYPE_DEEPSEEK41V); add GGUF keys if a later model differs
+        assert self.hparams_vision["max_image_tokens"] == 1024, self.hparams_vision["max_image_tokens"]
+        assert self.hparams_vision["max_wh_ratio"] is None, self.hparams_vision["max_wh_ratio"]
+        assert self.hparams_vision["rope_theta"] == 10000, self.hparams_vision["rope_theta"]
@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from .base import ModelBase, gguf
+from .qwen import Qwen2MoeModel
+
+
+@ModelBase.register("FlintForCausalLM")
+@ModelBase.example("RecursiveSI/flint")
+class FlintModel(Qwen2MoeModel):
+    model_arch = gguf.MODEL_ARCH.FLINT
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        dense = self.hparams["mlp_only_layers"]
+        if dense != list(range(len(dense))):
+            raise ValueError("Flint dense layers must precede the MoE layers")
+        self.gguf_writer.add_leading_dense_block_count(len(dense))
+        self.gguf_writer.add_expert_shared_count(
+            self.hparams["shared_expert_intermediate_size"] // self.hparams["moe_intermediate_size"]
+        )
+        self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])
+        self.gguf_writer.add_expert_weights_scale(1.0)
+        self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])
+        self.gguf_writer.add_sliding_window_pattern([
+            kind == "sliding_attention" for kind in self.hparams["layer_types"]
+        ])
+        self.gguf_writer.add_rope_pattern([
+            i not in self.hparams["no_rope_layer_indices"] for i in range(self.block_count)
+        ])
+
+    def modify_tensors(self, data_torch, name, bid):
+        if name.endswith(".softmax_offset"):
+            name += ".weight"
+        if name.endswith(".gate.expert_bias"):
+            if not self.hparams["router_expert_bias"]:
+                return
+            name += ".bias"
+        yield from super().modify_tensors(data_torch, name, bid)
@@ -272,3 +272,18 @@ class MiniCPMV4_7VisionModel(MiniCPMV4_6VisionModel):
            yield self.format_tensor_name(gguf.MODEL_TENSOR.V_TOK_EMBD_SEP, suffix=""), data_torch[ids]
            return
        yield from super().modify_tensors(data_torch, name, bid)
+
+
+# MiniCPM-OCR 1.0: the 4.6 stack with model_type minicpm_ocr.
+
+@ModelBase.register("MiniCPMOcrForConditionalGeneration")
+@ModelBase.example("openbmb/MiniCPM-OCR-1.0")
+class MiniCPMOCR1_0TextModel(MiniCPMV4_6TextModel):
+    model_arch = gguf.MODEL_ARCH.QWEN35
+    no_mtp = True
+
+
+@ModelBase.register("MiniCPMOcrForConditionalGeneration")
+@ModelBase.example("openbmb/MiniCPM-OCR-1.0")
+class MiniCPMOCR1_0VisionModel(MiniCPMV4_6VisionModel):
+    projector_type = gguf.VisionProjectorType.MINICPM_OCR_1_0
@@ -39,6 +39,11 @@ class Mistral3Model(TextModel):
            self.gguf_writer.add_add_bos_token(False)
            self.gguf_writer.add_add_eos_token(False)

+        def get_vocab_base_pre(self, tokenizer) -> str:
+            if self.origin_hf_arch == "Mistral3Model":
+                return "nemotron-3.5-embed"
+            return super().get_vocab_base_pre(tokenizer)
+
        def set_gguf_parameters(self):
            super().set_gguf_parameters()
            self._try_set_pooling_type()
@@ -1,5 +1,6 @@
 from __future__ import annotations

+import re
 from typing import Any, Callable, Iterable, Sequence, TYPE_CHECKING

 import numpy as np
@@ -367,9 +368,21 @@ class NemotronModel(TextModel):
    model_arch = gguf.MODEL_ARCH.NEMOTRON

    def set_vocab(self):
-        self._set_vocab_sentencepiece()
-        self.gguf_writer.add_pad_token_id(0)
-        self.gguf_writer.add_unk_token_id(1)
+        if (self.dir_model / "tokenizer.model").exists():
+            self._set_vocab_sentencepiece()
+            self.gguf_writer.add_pad_token_id(0)
+            self.gguf_writer.add_unk_token_id(1)
+        else:
+            # SPM-style BPE (Metaspace + byte fallback): mark <0xXX> tokens as BYTE so they decode to raw bytes
+            tokens, toktypes, tokpre = self.get_vocab_base()
+            toktypes = [gguf.TokenType.BYTE if re.fullmatch(r"<0x[0-9A-F]{2}>", tok) else typ for tok, typ in zip(tokens, toktypes)]
+            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)
+
+            special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
+            special_vocab.add_to_gguf(self.gguf_writer)

    def set_gguf_parameters(self):
        super().set_gguf_parameters()
@@ -383,7 +396,10 @@ class NemotronModel(TextModel):
        rot_pct = self.rope_parameters["partial_rotary_factor"]
        n_embd = self.find_hparam(["hidden_size", "n_embd"])
        n_head = self.find_hparam(["num_attention_heads", "n_head"])
-        self.gguf_writer.add_rope_dimension_count(int(rot_pct * n_embd) // n_head)
+        # Use explicit head_dim when available; fallback keeps standard Nemotron working.
+        # Without this, models where n_head * head_dim != n_embd (e.g. 48 heads, head_dim=128) produce the wrong rope dimension count.
+        head_dim = self.hparams.get("head_dim") or n_embd // n_head
+        self.gguf_writer.add_rope_dimension_count(int(rot_pct * head_dim))

        # * RopeScaling for Nemotron
        factor = self.hparams.get("factor") or self.rope_parameters.get("factor")
@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+from .base import MmprojModel, ModelBase, gguf
+from .qwen import Qwen2Model
+
+
+@ModelBase.example("sensenova/SenseNova-U1.5-8B-MoT")
+@ModelBase.register("NEOChatModel")
+class SenseNovaU1Model(Qwen2Model):
+    model_arch = gguf.MODEL_ARCH.SENSENOVA_U1
+
+    @classmethod
+    def filter_tensors(cls, item):
+        name, _ = item
+        if not name.startswith("language_model.") or "_mot_gen" in name:
+            return None
+        return super().filter_tensors(item)
+
+    def set_gguf_parameters(self):
+        if int(self.hparams.get("num_experts") or 0) > 1:
+            raise ValueError("SenseNova U1 conversion currently supports the dense understanding branch")
+        if self.rope_parameters.get("rope_type", "default") not in ("default", None):
+            raise ValueError("SenseNova U1 scaled RoPE is not supported")
+        if self.hparams.get("attention_bias", False):
+            raise ValueError("SenseNova U1 attention bias is not supported")
+        if self.hparams.get("use_sliding_window", False):
+            raise ValueError("SenseNova U1 sliding attention is not supported")
+        super().set_gguf_parameters()
+        self.gguf_writer.add_rope_freq_base_spatial(self.hparams["rope_theta_hw"])
+
+    def get_tensors(self):
+        for name, gen in self.model_tensors.items():
+            if name.endswith((".q_norm_hw.weight", ".k_norm_hw.weight")):
+                continue
+            data = gen()
+            if name.endswith((".q_norm.weight", ".k_norm.weight")):
+                spatial = self.model_tensors[name.replace("_norm.weight", "_norm_hw.weight")]()
+                # Keep one weight vector; the graph normalizes its two halves independently.
+                data = torch.cat((data, spatial), dim=0)
+            yield name, data
+
+
+@ModelBase.register("NEOVisionModel")
+@ModelBase.example("sensenova/SenseNova-U1.5-8B-MoT")
+class SenseNovaU1VisionModel(MmprojModel):
+    def __init__(self, dir_model: Path, *args, **kwargs):
+        hparams = kwargs.pop("hparams", None)
+        if hparams is None:
+            hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
+        else:
+            hparams = dict(hparams)
+
+        vision_config = dict(hparams["vision_config"])
+        # NEOVisionModel has no transformer blocks
+        # these defaults satisfy the generic mmproj metadata contract without replacing checkpoint values
+        vision_config.setdefault("image_size", 0)
+        vision_config.setdefault("intermediate_size", 0)
+        vision_config.setdefault("num_hidden_layers", 0)
+        vision_config.setdefault("num_attention_heads", 1)
+        hparams["vision_config"] = vision_config
+        super().__init__(dir_model, *args, hparams=hparams, **kwargs)
+
+        self.preprocessor_config.setdefault("image_mean", [0.485, 0.456, 0.406])
+        self.preprocessor_config.setdefault("image_std", [0.229, 0.224, 0.225])
+
+    @classmethod
+    def filter_tensors(cls, item):
+        name, gen = item
+        if not name.startswith("vision_model."):
+            return None
+        return super().filter_tensors((name, gen))
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_clip_projector_type("sensenova_u1")
+        vision_config = self.hparams["vision_config"]
+        # These defaults match the known U1.5 checkpoint when older configs omit fields.
+        hidden_size = int(vision_config.get("hidden_size", 1024))
+        min_pixels = int(vision_config.get("min_pixels", 65536))
+        max_pixels = int(vision_config.get("max_pixels", 16777216))
+        downsample_ratio = float(vision_config.get("downsample_ratio", 0.5))
+        if downsample_ratio <= 0:
+            raise ValueError(f"SenseNova U1 downsample_ratio must be positive, got {downsample_ratio}")
+        merge_size = round(1.0 / downsample_ratio)
+        if abs((1.0 / downsample_ratio) - merge_size) > 1e-6:
+            raise ValueError(f"SenseNova U1 downsample_ratio must have an integer reciprocal, got {downsample_ratio}")
+        norm_eps = float(vision_config.get("layer_norm_eps", vision_config.get("rms_norm_eps", 1e-6)))
+        self.gguf_writer.add_vision_head_dim(hidden_size)
+        self.gguf_writer.add_vision_attention_layernorm_eps(norm_eps)
+        self.gguf_writer.add_vision_min_pixels(min_pixels)
+        self.gguf_writer.add_vision_max_pixels(max_pixels)
+        self.gguf_writer.add_vision_spatial_merge_size(merge_size)
+
+    def modify_tensors(self, data_torch, name, bid):
+        tensor_names = {
+            "vision_model.embeddings.patch_embedding.weight": "v.patch_embd.weight",
+            "vision_model.embeddings.patch_embedding.bias":   "v.patch_embd.bias",
+            "vision_model.embeddings.dense_embedding.weight": self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ, 0),
+            "vision_model.embeddings.dense_embedding.bias": self.format_tensor_name(
+                gguf.MODEL_TENSOR.V_MMPROJ, 0, suffix=".bias"
+            ),
+        }
+        if name not in tensor_names:
+            raise ValueError(f"Unexpected SenseNova U1 vision tensor: {name!r}")
+        return [(tensor_names[name], data_torch)]
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if new_name == self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ, 0):
+            return gguf.GGMLQuantizationType.F16 if self.ftype == gguf.LlamaFileType.MOSTLY_F16 else gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from .base import ModelBase, MmprojModel, gguf, logger
+from .qwen import Qwen3Model
+
+
+def _is_soprano_checkpoint(dir_model: Path) -> bool:
+    if not (dir_model / "decoder.pth").is_file() or not (dir_model / "config.json").is_file():
+        return False
+    with open(dir_model / "config.json", encoding="utf-8") as f:
+        config = json.load(f)
+    if config.get("model_type") != "qwen3" or config.get("bos_token_id") != 3 or config.get("eos_token_id") != 3:
+        return False
+    if not (dir_model / "tokenizer.json").is_file():
+        return False
+    with open(dir_model / "tokenizer.json", encoding="utf-8") as f:
+        added = {t["id"]: t["content"] for t in json.load(f).get("added_tokens", [])}
+    return [added.get(i) for i in range(4)] == ["[UNK]", "[TEXT]", "[START]", "[STOP]"]
+
+
+@ModelBase.register_hparams_loader(_is_soprano_checkpoint)
+def _load_soprano_hparams(dir_model: Path) -> dict[str, Any]:
+    logger.info("gguf: detected Soprano checkpoint")
+    hparams = ModelBase.load_hparams(dir_model, False, guess=False)
+    # the text model stays Qwen3, the mmproj holds the decoder
+    hparams["text_config"] = {"architectures": ["SopranoTextModel"]}
+    hparams["architectures"] = ["SopranoModel"]
+    return hparams
+
+
+@ModelBase.register("SopranoTextModel")
+class SopranoTextModel(Qwen3Model):
+    model_arch = gguf.MODEL_ARCH.QWEN3
+
+    def get_vocab_base_pre(self, tokenizer) -> str:
+        return "soprano"
+
+    def set_vocab(self):
+        super().set_vocab()
+        self.gguf_writer.add_unk_token_id(0)
+
+
+@ModelBase.register("SopranoModel")
+@ModelBase.example("ekwek/Soprano-80M", "ekwek/Soprano-1.1-80M")
+class SopranoModel(MmprojModel):
+    has_vision_encoder = False
+    has_audio_encoder = False
+
+    def _load_decoder(self) -> dict[str, torch.Tensor]:
+        state = torch.load(self.dir_model / "decoder.pth", map_location="cpu", mmap=True, weights_only=True)
+        return state.get("state_dict", state)
+
+    def get_audio_config(self):
+        state = self._load_decoder()
+        n_embd, n_inp, _ = state["decoder.embed.weight"].shape
+        n_ff = state["decoder.convnext.0.pwconv1.weight"].shape[0]
+        n_layer = sum(1 for name in state if name.endswith(".pwconv1.weight"))
+        if n_inp != self.n_embd_text:
+            raise ValueError(f"Soprano decoder input dim {n_inp} does not match the backbone dim {self.n_embd_text}")
+        return {"num_hidden_layers": n_layer, "hidden_size": n_embd, "intermediate_size": n_ff}
+
+    def set_gguf_parameters(self):
+        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.SOPRANO)
+        self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text)
+        self.gguf_writer.add_gen_audio_embedding_length(self.hparams["hidden_size"])
+        self.gguf_writer.add_gen_audio_feed_forward_length(self.hparams["intermediate_size"])
+        self.gguf_writer.add_gen_audio_block_count(self.hparams["num_hidden_layers"])
+        self.gguf_writer.add_gen_audio_head_count(1)
+        self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-6)
+
+    def get_tensors(self):
+        yield from self._load_decoder().items()
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if "dwconv.weight" in new_name or (name == "decoder.embed.weight" and n_dims == 3):
+            return gguf.GGMLQuantizationType.F16
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    def modify_tensors(self, data_torch, name, bid):
+        if name == "head.istft.window":
+            # the ISTFT in mtmd builds its own window
+            if not torch.equal(data_torch.float(), torch.hann_window(data_torch.shape[0])):
+                raise ValueError("Soprano requires a periodic Hann window")
+            return
+        if name == "decoder.embed.weight" and data_torch.shape[-1] == 1:
+            data_torch = data_torch.squeeze(-1)
+
+        T = gguf.MODEL_TENSOR
+        if name.startswith("decoder.convnext."):
+            _, _, blk, key, *suffix = name.split(".")
+            block_map = {
+                "dwconv":  T.A_GEN_WAV_UP_DWCONV,
+                "norm":    T.A_GEN_WAV_UP_NORM,
+                "pwconv1": T.A_GEN_WAV_UP_PW1,
+                "pwconv2": T.A_GEN_WAV_UP_PW2,
+                "gamma":   T.A_GEN_WAV_UP_GAMMA,
+            }
+            yield self.format_tensor_name(block_map[key], int(blk), suffix="".join("." + s for s in suffix)), data_torch
+            return
+
+        tensor_map = {
+            "decoder.embed":            T.A_GEN_WAV_INPUT,
+            "decoder.norm":             T.A_GEN_WAV_NORM,
+            "decoder.final_layer_norm": T.A_GEN_WAV_OUTPUT_NORM,
+            "head.out":                 T.A_GEN_WAV_OUTPUT,
+        }
+        key, suffix = name.rsplit(".", 1)
+        yield self.format_tensor_name(tensor_map[key], suffix="." + suffix), data_torch
@@ -267,11 +267,15 @@ def main() -> None:
            sys.exit(1)

        if args.dspark:
-            if is_mistral_format or model_architecture != "DeepseekV4ForCausalLM":
-                logger.error("--dspark is only supported for DeepseekV4ForCausalLM")
+            if is_mistral_format or model_architecture not in ("DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"):
+                logger.error("--dspark is only supported for DeepseekV4ForCausalLM and DeepseekV41ForCausalLM")
                sys.exit(1)
-            from conversion.deepseek import DeepseekV4DSparkModel
-            model_class = DeepseekV4DSparkModel
+            if model_architecture == "DeepseekV41ForCausalLM":
+                from conversion.deepseek41 import DeepseekV41DSparkModel
+                model_class = DeepseekV41DSparkModel
+            else:
+                from conversion.deepseek import DeepseekV4DSparkModel
+                model_class = DeepseekV4DSparkModel

        if args.mtp or args.no_mtp:
            if not model_class.supports_mtp_export:
@@ -101,6 +101,7 @@ models = [
    {"name": "command-r",        "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/CohereForAI/c4ai-command-r-v01", },
    {"name": "tiny_aya",         "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/CohereLabs/tiny-aya-base", },
    {"name": "cohere2moe",       "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/CohereLabs/North-Mini-Code-1.0", },
+    {"name": "flint",            "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/RecursiveSI/flint", },
    {"name": "qwen2",            "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/Qwen/Qwen1.5-7B", },
    {"name": "olmo",             "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/allenai/OLMo-1.7-7B-hf", },
    {"name": "dbrx",             "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/databricks/dbrx-base", },
@@ -177,6 +178,7 @@ pre_computed_hashes = [
    {"name": "glm4", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/zai-org/GLM-4.5-Air", "chkhsh": "9ca2dd618e8afaf09731a7cf6e2105b373ba6a1821559f258b272fe83e6eb902"},
    {"name": "glm4", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/zai-org/GLM-4.7-Flash", "chkhsh": "cdf5f35325780597efd76153d4d1c16778f766173908894c04afc20108536267"},
    {"name": "minerva-7b", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/sapienzanlp/Minerva-7B-base-v1.0", "chkhsh": "1431a23e583c97432bc230bff598d103ddb5a1f89960c8f1d1051aaa944d0b35"},
+    {"name": "domyn-small", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/domyn/Domyn-Small-v1.0", "chkhsh": "6e4d72c6e5260b90606e26a0931fa15e4b42ec6e97719d12919f2b1f93478262"},
    {"name": "hunyuan", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tencent/Hunyuan-A13B-Instruct", "chkhsh": "7e57df22b1fe23a7b1e1c7f3dc4e3f96d43a4eb0836d0c6bdc3436d7b2f1c664"},
    {"name": "hunyuan-dense", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tencent/Hunyuan-4B-Instruct", "chkhsh": "bba3b3366b646dbdded5dbc42d59598b849371afc42f7beafa914afaa5b70aa6"},
    {"name": "hy_v4", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/tencent/Hy4-preview", "chkhsh": "e6ddf9c6686791c12d698d34c31ab9be1fea9af5a3d9a6909783ab382198ae1c"},
@@ -783,7 +783,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL          | ON (mandatory)                        | Enable build with SYCL code path.           |
 | GGML_SYCL_TARGET   | INTEL *(default)*                     | Set the SYCL target device type.            |
 | GGML_SYCL_DEVICE_ARCH | Optional                           | Set the SYCL device architecture. Setting the device architecture can improve the performance. See the table [--offload-arch](https://github.com/intel/llvm/blob/sycl/sycl/doc/design/OffloadDesign.md#--offload-arch) for a list of valid architectures. |
-| GGML_SYCL_F16      | OFF *(default)* \|ON *(optional)*     | Enable FP16 build with SYCL code path. (1.) |
+| GGML_SYCL_F16      | OFF *(default)* \|ON *(optional)*     | Enable FP16 build with SYCL code path. [1.](#sycl-f16-footnote) |
 | GGML_SYCL_GRAPH    | ON *(default)* \|OFF *(Optional)*     | Enable build with [SYCL Graph extension](https://github.com/intel/llvm/blob/sycl/sycl/doc/extensions/experimental/sycl_ext_oneapi_graph.asciidoc). |
 | GGML_SYCL_DNN      | ON *(default)* \|OFF *(Optional)*     | Enable build with oneDNN.                   |
 | GGML_SYCL_HOST_MEM_FALLBACK | ON *(default)* \|OFF *(Optional)* | Allow host memory fallback when device memory is full during quantized weight reorder. Enables inference to continue at reduced speed (reading over PCIe) instead of failing. Requires Linux kernel 6.8+. |
@@ -791,7 +791,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | CMAKE_C_COMPILER   | `icx` *(Linux)*, `icx/cl` *(Windows)* | Set `icx` compiler for SYCL code path.      |
 | CMAKE_CXX_COMPILER | `icpx` *(Linux)*, `icx` *(Windows)*   | Set `icpx/icx` compiler for SYCL code path. |

-1. FP32 or FP16 have different performance impact to LLM. Recommended to test them for better prompt processing performance on your models. You need to rebuild the code after change `GGML_SYCL_F16=OFF/ON`.
+1. <a id="sycl-f16-footnote"></a>FP32 or FP16 have different performance impact to LLM. Recommended to test them for better prompt processing performance on your models. You need to rebuild the code after change `GGML_SYCL_F16=OFF/ON`.

 ### Runtime

@@ -802,7 +802,8 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_DEV2DEV_MEMCPY | 0 (default), 1, 2 | Choose the method of dev2dev memory copy.<br>Value: <br>*  0: SYCL API (default), only support dGPUs.<br>* 1: L0 API -- Better performance, only support dGPUs, found to lead to abnormal crash in some case. <br>* 2: Host Forward -- Most stable method for all cases (including iGPU + dGPU*N), but with lower performance (-2% to -5%).<br>SYCL & L0 API are easy to be impacted by Intel GPU driver issue. When you meet the garbled output or crash issues in multiple GPUs case, try with this debug flag to work around or check the issue.|
 | GGML_SYCL_ENABLE_FLASH_ATTN | 1 (default) or 0| Enable Flash-Attention. It can reduce memory usage. The performance impact depends on the LLM.|
 | GGML_SYCL_ENABLE_OPT | 0 or 1 (default)| Enable optimize features for Intel GPUs. (Recommended to 0 for Intel devices older than Gen 10) |
-| GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. |
+| GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still in development, performance depends on model and device. [1.](#sycl-graph-footnote) |
+| GGML_SYCL_GRAPH_EVICTION_TIMEOUT | 10 (default) or positive integer | Seconds a cached SYCL graph can go unused before it is evicted. Checked on a sweep that runs every half of this timeout. |
 | GGML_SYCL_ENABLE_HOST_PINNED_MEM | 0 or 1 (default) | Enable host pinned memory to speed up copy data from host to device. When disable it, host memory will common malloc() on CPU. Disable it when use `--load-model mlock`.|
 | GGML_SYCL_HOST_PINNED_MEM_2G | 0 (default) or 1 | Limit the max memory allocation to be no more than 2GB when enable host pinned memory. USM allocations above 2 GiB take the relaxed/large-allocation path, which serializes H2D copies with compute and prevents copy/compute overlap. It will impact the startup time. Need more test. Depend on `GGML_SYCL_ENABLE_HOST_PINNED_MEM=1`.|
 | GGML_SYCL_UPLOAD_STAGING_SLOTS | 4 (default) or non-negative integer | Number of 8 MiB pinned host slots used to stage tensor uploads (model loading), so the host copy of one slot overlaps the transfer of the previous one. Set to 0 to use the old path: a malloc'd bounce buffer and a blocking copy per tensor. |
@@ -819,6 +820,9 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. |
 | GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. Unsupported types and layouts fall back to the standalone op kernels. See `ggml_sycl_can_fuse()`. |
 | GGML_SYCL_ENABLE_ESIMD | 0 or 1 (default)| Enable ESIMD kernels when available. |
+| GGML_SYCL_ENABLE_XMX_ESIMD | 0 or 1 (default) | Use the XMX (DPAS) multi-column mat-vec kernels for reordered Q2_K to Q6_K and Q8_0 weights on Xe2 GPUs. Speeds up speculative decoding and short prompt batches. |
+| GGML_SYCL_XMX_MAX_COLS | non-negative integer, 80 (default) | Most activation columns (batch size) the `GGML_SYCL_ENABLE_XMX_ESIMD` kernels take; wider batches use the other mat-mul kernels. A larger value keeps XMX for bigger batches (as repeated launches), `0` never uses it. |
+| GGML_SYCL_XMX_GLU_MAX_COLS | 0 to 16, 16 (default) | Most columns the fused gate and up XMX kernel takes; wider batches run gate and up as separate mat-vecs. 16 is the most the kernel is built for, larger values are clamped to 16 and `0` turns the fused kernel off. |
 | GGML_SYCL_XMX_GATHER_TYPES | decimal bitmask, all bits set (default) | Weight formats that may use the XMX dequant-GEMM paths, which dequantize weights straight into the XMX tiles. This speeds up prompt processing of MoE models on GPUs with XMX units (Arc A- and B-series, Arc Pro, Data Center GPU Max), for example pp512 of Qwen3-30B-A3B UD-IQ3_XXS by about 50% on an Arc Pro B60. Bits:<br>* 1: IQ4_NL, 2: IQ3_S, 4: IQ4_XS, 8: IQ3_XXS, 16: IQ2_XXS, 32: IQ2_XS, 64: IQ2_S, 128: IQ1_S, 256: IQ1_M<br>* 512: Q8_0, 1024: Q4_K, 2048: Q5_K, 4096: Q6_K (MoE `MUL_MAT_ID` only)<br>Add values to combine them, for example `3` for IQ4_NL and IQ3_S; `0` disables the paths. A set bit does not force the path: batches of more than 64 tokens per expert or row lengths that are not a multiple of 256 (32 for IQ4_NL and Q8_0) use the library GEMM. |
 | GGML_SYCL_XMX_GATHER_SHAPES | decimal bitmask, 255 (default) | XMX `joint_matrix` combinations the paths of `GGML_SYCL_XMX_GATHER_TYPES` may use; the operand type comes from `GGML_SYCL_DYNAMIC_PRECISION` and the best supported combination is picked automatically (logged as `fg_pick_combo`). Bits:<br>* Xe2, Xe3, Xe-HPC: 1: f16 8x16x16, 2: f16 16x16x16, 4: f16 32x64x16, 8: f16 32x64x32, 32: tf32 8x16x8, 64: bf16 8x16x16<br>* Xe-HPG (Arc A770, ARL-H): 16: f16 8x8x16, 128: bf16 8x8x16<br>Clear a bit to exclude a combination, or set a single bit to force one for testing. |
 | GGML_SYCL_DYNAMIC_PRECISION | `F16` (default with `GGML_SYCL_F16=ON`), `BF16`, `TF32` or `F32` (default otherwise) | Operand type of the XMX dequant-GEMM paths (`GGML_SYCL_XMX_GATHER_TYPES`); accumulation is always f32. `F16` is the fastest, but activations above 65504 overflow. `BF16` keeps the f32 range at a 7-bit mantissa, `TF32` keeps the range and the f16 mantissa but is about 30% slower and needs Xe2, Xe3 or Xe-HPC, and `F32` turns the XMX paths off. Ops that request a higher src1 precision ([TAG_GGML_PREC]) get it regardless of this setting. |
@@ -831,6 +835,17 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. |
 | GGML_SYCL_USM_SYSTEM | 0 (default) or 1 | Enable experimental support for [USM system allocations](https://github.khronos.org/SYCL_Reference/iface/usm_basic_concept.html#system-allocations) for large GPU buffers. This requires enough host memory for model weights and caches, an Intel Xe2+ GPU such as BMG or newer and supported on Linux only, with CONFIG_DRM_XE_GPUSVM enabled. |

+1. <a id="sycl-graph-footnote"></a>SYCL Graph support (`GGML_SYCL_ENABLE_GRAPH=1`) is limited to single-GPU workloads and may also depend on the model architecture and kernels in use.
+   Performance depends on how often a recorded graph can be reused, which will vary by workload.
+   For this reason, it is recommended to measure performance both with and without graphs enabled on your workloads.
+
+   To check whether graphs are supported for your workload, enable INFO logging (`-lv 4`) and look for `enabling SYCL graphs` or `disabling SYCL graphs`.
+   Support may differ between prefill and decode, so check the status for the phases that matter to you.
+   DEBUG logging (`-lv 5`) reports why graphs are disabled.
+
+   Then benchmark with your preferred tool, such as `llama-bench`.
+   For graph activity details (including reuse), set `GGML_SYCL_DEBUG=1` and look for `[SYCL-GRAPH]` messages.
+
 ## Compile-time Flags

 Pass these via `CXXFLAGS` or add a one-off `#define` to enable a flag on the spot.
@@ -14,7 +14,7 @@ Legend:

 | Operation | BLAS | CANN | CPU | CUDA | ET | HTP | MTL | OpenCL | SYCL | Vulkan | WebGPU | ZenDNN | zDNN |
 |-----------|------|------|------|------|------|------|------|------|------|------|------|------|------|
-|                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              ACC | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             ADD1 | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
@@ -22,12 +22,12 @@ Legend:
 |                           ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            CLAMP | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        COL2IM_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                       CONV_1D_DW | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ |
+|                       CONV_1D_DW | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ |
 |                          CONV_2D | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
 |                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          CONV_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -45,7 +45,7 @@ Legend:
 |                     DSV4_HC_COMB | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                     DSV4_HC_POST | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                      DSV4_HC_PRE | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                              DUP | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                              DUP | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              ELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              EXP | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            EXPM1 | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -55,10 +55,10 @@ Legend:
 |                  GATED_DELTA_NET | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                GATED_LINEAR_ATTN | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                            GEGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        GEGLU_ERF | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                        GEGLU_ERF | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                      GEGLU_QUICK | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             GELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                         GELU_ERF | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         GELU_ERF | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                       GELU_QUICK | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         GET_ROWS | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ❌ | ❌ |
 |                    GET_ROWS_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ |
@@ -66,7 +66,7 @@ Legend:
 |                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | 🟡 | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                LIGHTNING_INDEXER | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
@@ -74,12 +74,12 @@ Legend:
 |                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ | ❌ |
 |                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
-|                 MUL_MAT_HADAMARD | ❌ | ❌ | ✅ | ❌ | ✅ | 🟡 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 |
+|                 MUL_MAT_HADAMARD | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 |
 |                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | 🟡 | ❌ |
-|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                     MUL_MAT_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|                     MUL_MAT_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -87,9 +87,9 @@ Legend:
 |                         OUT_PROD | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ |
 |                              PAD | ❌ | 🟡 | ✅ | 🟡 | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ |
 |                   PAD_REFLECT_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                          POOL_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                          POOL_2D | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                          POOL_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                          POOL_2D | ❌ | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             RELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                      REPEAT_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -101,7 +101,7 @@ Legend:
 |                            ROUND | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        RWKV_WKV6 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                        RWKV_WKV7 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            SCALE | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            SCALE | ❌ | 🟡 | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              SET | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ |
 |                         SET_ROWS | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
 |                              SGN | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -117,12 +117,12 @@ Legend:
 |                             SQRT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         SSM_CONV | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         SSM_SCAN | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ❌ | ❌ |
-|                             STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              SUB | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              SUM | ❌ | 🟡 | ✅ | 🟡 | ❌ | 🟡 | 🟡 | ❌ | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
 |                         SUM_ROWS | ❌ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ |
 |                           SWIGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                     SWIGLU_CLAMP | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     SWIGLU_CLAMP | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                       SWIGLU_OAI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |               TIMESTEP_EMBEDDING | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -926,6 +926,7 @@ extern "C" {
            enum   ggml_type      type);

    // dst[i0, i1, i2] = a[i0, i1, i2] + b[i0, ids[i1, i2]]
+    // an id of -1 adds nothing: dst[i0, i1, i2] = a[i0, i1, i2]
    GGML_API struct ggml_tensor * ggml_add_id(
            struct ggml_context * ctx,
            struct ggml_tensor  * a,
@@ -1477,6 +1478,16 @@ extern "C" {
            enum ggml_prec       prec,
            int                  idx);

+    // [TAG_GGML_PREC]
+    // get the accumulator / src[idx] precision set by the functions above
+    // return GGML_PREC_UNDEFINED when the op does not carry the requested precision
+    GGML_API enum ggml_prec ggml_prec_get_acc(
+            const struct ggml_tensor * a);
+
+    GGML_API enum ggml_prec ggml_prec_get_src(
+            const struct ggml_tensor * a,
+            int                        idx);
+
    // A: k columns, n rows => [ne03, ne02, n, k]
    // B: k columns, m rows  (i.e. we transpose it internally) => [ne03 * x, ne02 * y, m, k]
    // result is n columns, m rows => [ne03 * x, ne02 * y, m, n]
@@ -1498,6 +1509,7 @@ extern "C" {
            enum ggml_op_hint    hint);

    // indirect matrix multiplication
+    // an id of -1 skips the slot and zeroes the matching row of the result
    //
    // [TAG_MUL_MAT_ID_DUPLICATE_ID]
    // note: this op has undefined behaviour if a given `a` is referenced more than once for a given `b`
@@ -774,13 +774,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
        return src_ss[0];
    };

-    auto handle_pool_1d = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
-        if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) {
-            return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
-        }
-        return src_ss[0];
-    };
-
    auto handle_flash_attn_ext = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
        GGML_ASSERT(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);

@@ -997,14 +990,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
            case GGML_OP_CONV_3D:
            case GGML_OP_CONV_2D_DW:
            case GGML_OP_CONV_TRANSPOSE_2D:
+            case GGML_OP_POOL_1D:
            case GGML_OP_POOL_2D:
            case GGML_OP_POOL_2D_BACK:
            case GGML_OP_UPSCALE: {
                split_state = handle_generic(src_ss, /*scalar_only =*/ true);
            } break;
-            case GGML_OP_POOL_1D: {
-                split_state = handle_pool_1d(src_ss);
-            } break;
            case GGML_OP_PAD: {
                split_state = handle_pad(src_ss);
            } break;
@@ -920,7 +920,7 @@ static void * ggml_backend_cann_buffer_get_base(ggml_backend_buffer_t buffer) {
 * @param dst Pointer to the destination buffer where transformed data will be
 * stored.
 */
-static void ggml_backend_cann_transform_q4_0(ggml_tensor * tensor, const void * src, void * dst) {
+static void ggml_backend_cann_transform_q4_0(const ggml_tensor * tensor, const void * src, void * dst) {
    int64_t n_elems     = ggml_nelements(tensor);
    int64_t groups      = n_elems / QK4_0;
    size_t  quant_bytes = n_elems * sizeof(uint8_t) / 2;
@@ -1014,7 +1014,7 @@ static void ggml_backend_cann_transform_back_q4_0(const ggml_tensor * tensor, vo
 * @param dst Pointer to the destination buffer where transformed data will be
 * stored.
 */
-static void ggml_backend_cann_transform_q8_0(ggml_tensor * tensor, const void * src, void * dst) {
+static void ggml_backend_cann_transform_q8_0(const ggml_tensor * tensor, const void * src, void * dst) {
    int64_t n_elems     = ggml_nelements(tensor);
    int64_t groups      = n_elems / QK8_0;
    size_t  quant_bytes = n_elems * sizeof(uint8_t);
@@ -1075,7 +1075,7 @@ static void ggml_backend_cann_transform_back_q8_0(const ggml_tensor * tensor, co
 * @param dst Pointer to the destination buffer where transformed data will be
 * stored.
 */
-static void ggml_backend_cann_transform(ggml_tensor * tensor, const void * src, void * dst) {
+static void ggml_backend_cann_transform(const ggml_tensor * tensor, const void * src, void * dst) {
    switch (tensor->type) {
        case GGML_TYPE_Q4_0:
            ggml_backend_cann_transform_q4_0(tensor, src, dst);
@@ -650,9 +650,9 @@ function(ggml_add_cpu_backend_variant_impl tag_name)

        # Fetch KleidiAI sources:
        include(FetchContent)
-        set(KLEIDIAI_COMMIT_TAG "v1.31.0")
+        set(KLEIDIAI_COMMIT_TAG "v1.32.0")
        set(KLEIDIAI_DOWNLOAD_URL "https://github.com/ARM-software/kleidiai/releases/download/${KLEIDIAI_COMMIT_TAG}/kleidiai-${KLEIDIAI_COMMIT_TAG}-src.tar.gz")
-        set(KLEIDIAI_RELEASE_ARCHIVE_MD5  "930495f3aeb16b25dc18a2bde45924f4")
+        set(KLEIDIAI_RELEASE_ARCHIVE_MD5  "05f1953428ffe8e8c7f585dc6d508887")

        set(KLEIDIAI_FETCH_ARGS
            URL ${KLEIDIAI_DOWNLOAD_URL}
@@ -1652,7 +1652,12 @@ static void ggml_compute_forward_mul_mat_id(
            for (int id = 0; id < n_ids; ++id) {
                const int32_t i02 = *(const int32_t *) ((const char *) ids->data + iid1*ids->nb[1] + id*ids->nb[0]);

-                assert(i02 >= 0 && i02 < n_as);
+                assert(i02 == -1 || (i02 >= 0 && i02 < n_as));
+
+                if (i02 == -1) {
+                    memset((char *) dst->data + id*nb1 + iid1*nb2, 0, ne0*sizeof(float));
+                    continue;
+                }

                MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = (struct mmid_row_mapping) {id, iid1};
                matrix_row_counts[i02] += 1;
@@ -201,85 +201,117 @@ static struct kai_matmul_uker_config q4_k_matmul_config() {
    return config;
 }

-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_m_step() {
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    return api.get_step(&config).m;
+static struct kai_matmul_uker_config default_matmul_config() {
+    return {};
 }

-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_n_step() {
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    return api.get_step(&config).n;
-}
+template<
+        struct kai_matmul_uker_api (*GetApi)(void),
+        struct kai_matmul_uker_config (*GetConfig)(void),
+        size_t MR, size_t NR, size_t KR, size_t SR>
+struct matmul_uker_adapter {
+    static size_t get_m_step() {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        return api.get_step(&config).m;
+    }

-template<size_t Value>
-static size_t q4_k_get_constant() {
-    return Value;
-}
+    static size_t get_n_step() {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        return api.get_step(&config).n;
+    }

-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_dst_offset(size_t m_idx, size_t n_idx, size_t dst_stride) {
-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_dst_dim_args index = { m_idx, n_idx };
-    const struct kai_matmul_uker_dst_stride_args stride = { dst_stride };
-    return api.get_dst_offset(&config, &index, &stride);
-}
+    static size_t get_mr() {
+        if constexpr (MR == 0) {
+            return get_m_step();
+        }
+        return MR;
+    }

-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_lhs_packed_offset(size_t m_idx, size_t k, size_t bl) {
-    GGML_ASSERT(bl == KLEIDIAI_Q4_K_BLOCK_LENGTH);
+    static size_t get_nr() {
+        if constexpr (NR == 0) {
+            return get_n_step();
+        }
+        return NR;
+    }

-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_lhs_dim_args shape = { m_idx, k };
-    const struct kai_matmul_uker_lhs_dim_args index = { m_idx, 0 };
-    const struct kai_matmul_uker_lhs_stride_args stride = api.get_lhs_stride(&config, &shape);
-    return api.get_lhs_offset(&config, &index, &stride);
-}
+    static size_t get_kr() {
+        return KR;
+    }

-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static size_t q4_k_get_rhs_packed_offset(size_t n_idx, size_t k, size_t bl) {
-    GGML_ASSERT(bl == KLEIDIAI_Q4_K_BLOCK_LENGTH);
+    static size_t get_sr() {
+        return SR;
+    }

-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_rhs_dim_args shape = { n_idx, k };
-    const struct kai_matmul_uker_rhs_dim_args index = { n_idx, 0 };
-    const struct kai_matmul_uker_rhs_stride_args stride = api.get_rhs_stride(&config, &shape);
-    return api.get_rhs_offset(&config, &index, &stride);
-}
+    static size_t get_dst_offset(size_t m_idx, size_t n_idx, size_t dst_stride) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        const struct kai_matmul_uker_dst_dim_args index = { m_idx, n_idx };
+        const struct kai_matmul_uker_dst_stride_args stride = { dst_stride };
+        return api.get_dst_offset(&config, &index, &stride);
+    }

-template<struct kai_matmul_uker_api (*GetApi)(void)>
-static void q4_k_run_matmul(
-        size_t m, size_t n, size_t k, size_t bl,
-        const void * lhs_packed, const void * rhs_packed,
-        void * dst, size_t dst_stride_row, size_t dst_stride_col,
-        float clamp_min, float clamp_max) {
-    GGML_ASSERT(bl == KLEIDIAI_Q4_K_BLOCK_LENGTH);
-    GGML_ASSERT(dst_stride_col == sizeof(float));
+    static size_t get_lhs_packed_offset(size_t m_idx, size_t k, size_t bl) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        GGML_ASSERT(bl == config.format.bl);
+        const struct kai_matmul_uker_lhs_dim_args shape = { m_idx, k };
+        const struct kai_matmul_uker_lhs_dim_args index = { m_idx, 0 };
+        const struct kai_matmul_uker_lhs_stride_args stride = api.get_lhs_stride(&config, &shape);
+        return api.get_lhs_offset(&config, &index, &stride);
+    }

-    const struct kai_matmul_uker_api api = GetApi();
-    const struct kai_matmul_uker_config config = q4_k_matmul_config();
-    const struct kai_matmul_uker_lhs_dim_args lhs_shape = { m, k };
-    const struct kai_matmul_uker_rhs_dim_args rhs_shape = { n, k };
+    static size_t get_rhs_packed_offset(size_t n_idx, size_t k, size_t bl) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        GGML_ASSERT(bl == config.format.bl);
+        const struct kai_matmul_uker_rhs_dim_args shape = { n_idx, k };
+        const struct kai_matmul_uker_rhs_dim_args index = { n_idx, 0 };
+        const struct kai_matmul_uker_rhs_stride_args stride = api.get_rhs_stride(&config, &shape);
+        return api.get_rhs_offset(&config, &index, &stride);
+    }

-    struct kai_matmul_uker_args args = {};
-    args.flags = KAI_MATMUL_UKER_FLAGS_ARGS_CLAMP;
-    args.shape = { m, n, k };
-    args.operand.lhs.ptr = lhs_packed;
-    args.operand.lhs.stride = api.get_lhs_stride(&config, &lhs_shape);
-    args.operand.rhs.ptr = rhs_packed;
-    args.operand.rhs.stride = api.get_rhs_stride(&config, &rhs_shape);
-    args.operand.dst.ptr = dst;
-    args.operand.dst.stride.m = dst_stride_row;
-    args.activation.clamp.min_ptr = &clamp_min;
-    args.activation.clamp.max_ptr = &clamp_max;
-    api.run(&config, &args);
-}
+    static void run(
+            size_t m, size_t n, size_t k, size_t bl,
+            const void * lhs_packed, const void * rhs_packed,
+            void * dst, size_t dst_stride_row, size_t dst_stride_col,
+            float clamp_min, float clamp_max) {
+        const struct kai_matmul_uker_api api = GetApi();
+        const struct kai_matmul_uker_config config = GetConfig();
+        GGML_ASSERT(bl == config.format.bl);
+        GGML_ASSERT(dst_stride_col == sizeof(float));
+        const struct kai_matmul_uker_lhs_dim_args lhs_shape = { m, k };
+        const struct kai_matmul_uker_rhs_dim_args rhs_shape = { n, k };
+
+        struct kai_matmul_uker_args args = {};
+        args.flags = KAI_MATMUL_UKER_FLAGS_ARGS_CLAMP;
+        args.shape = { m, n, k };
+        args.operand.lhs.ptr = lhs_packed;
+        args.operand.lhs.stride = api.get_lhs_stride(&config, &lhs_shape);
+        args.operand.rhs.ptr = rhs_packed;
+        args.operand.rhs.stride = api.get_rhs_stride(&config, &rhs_shape);
+        args.operand.dst.ptr = dst;
+        args.operand.dst.stride.m = dst_stride_row;
+        args.activation.clamp.min_ptr = &clamp_min;
+        args.activation.clamp.max_ptr = &clamp_max;
+        api.run(&config, &args);
+    }
+};
+
+using q4_k_gemm_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa,
+    q4_k_matmul_config, 0, 0, 4, 2>;
+using q4_k_gemv_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot,
+    q4_k_matmul_config, 0, 0, 4, 2>;
+// GEMM and GEMV share the RHS layout only at SVE256 (NR=8, KR=8, SR=1).
+using sve_q8_gemm_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_qai8dxp4x8sf32_qsi8cxp4vsx8sf32bf32_16x4vs_sve_i8mm,
+    default_matmul_config, 4, 0, 8, 1>;
+using sve_q8_gemv_adapter = matmul_uker_adapter<
+    kai_matmul_clamp_f32_qai8dxp1x8_qsi8cxp8x8_1x8_sve_dot,
+    default_matmul_config, 1, 8, 8, 1>;

 static struct kai_matmul_pack_lhs_uker_config q4_k_lhs_pack_config() {
    struct kai_matmul_pack_lhs_uker_config config = {};
@@ -628,7 +660,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
            /* .packed_size_ex        = */ &rhs_ps_fn5<kai_get_rhs_packed_size_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
            /* .pack_func_ex          = */ &rhs_pack_fn12<kai_run_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0>,
        },
-        /* .required_cpu       = */ CPU_FEATURE_SVE | CPU_FEATURE_I8MM | CPU_FEATURE_DOTPROD,
+        /* .required_cpu       = */ CPU_FEATURE_SVE | CPU_FEATURE_SVEI8MM | CPU_FEATURE_DOTPROD,
        /* .lhs_type           = */ GGML_TYPE_F32,
        /* .rhs_type           = */ GGML_TYPE_Q4_0,
        /* .op_type            = */ GGML_TYPE_F32,
@@ -737,16 +769,16 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q4_k[] = {
    {
        /* SME2 GEMM */
        {
-            /* .get_m_step            = */ &q4_k_get_m_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_n_step            = */ &q4_k_get_n_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_mr                = */ &q4_k_get_m_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_nr                = */ &q4_k_get_n_step<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_kr                = */ &q4_k_get_constant<4>,
-            /* .get_sr                = */ &q4_k_get_constant<2>,
-            /* .get_dst_offset        = */ &q4_k_get_dst_offset<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_lhs_offset_ex     = */ &q4_k_get_lhs_packed_offset<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .get_rhs_packed_offset_ex = */ &q4_k_get_rhs_packed_offset<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
-            /* .run_kernel_ex         = */ &q4_k_run_matmul<kai_matmul_clamp_f32_f16p4vsx2_qai4c32p16vsx4s1s0sf16_4vsx16vs_sme2_mopa>,
+            /* .get_m_step            = */ q4_k_gemm_adapter::get_m_step,
+            /* .get_n_step            = */ q4_k_gemm_adapter::get_n_step,
+            /* .get_mr                = */ q4_k_gemm_adapter::get_mr,
+            /* .get_nr                = */ q4_k_gemm_adapter::get_nr,
+            /* .get_kr                = */ q4_k_gemm_adapter::get_kr,
+            /* .get_sr                = */ q4_k_gemm_adapter::get_sr,
+            /* .get_dst_offset        = */ q4_k_gemm_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ q4_k_gemm_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ q4_k_gemm_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ q4_k_gemm_adapter::run,
        },
        {
            /* .get_offset            = */ kai_get_lhs_offset_lhs_pack_f16pmrx2_f32_neon,
@@ -756,16 +788,16 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q4_k[] = {
        },
        /* SME2 GEMV */
        {
-            /* .get_m_step            = */ &q4_k_get_m_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_n_step            = */ &q4_k_get_n_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_mr                = */ &q4_k_get_m_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_nr                = */ &q4_k_get_n_step<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_kr                = */ &q4_k_get_constant<4>,
-            /* .get_sr                = */ &q4_k_get_constant<2>,
-            /* .get_dst_offset        = */ &q4_k_get_dst_offset<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_lhs_offset_ex     = */ &q4_k_get_lhs_packed_offset<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .get_rhs_packed_offset_ex = */ &q4_k_get_rhs_packed_offset<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
-            /* .run_kernel_ex         = */ &q4_k_run_matmul<kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p16vsx4s1s0sf16_1x16vs_sme2_dot>,
+            /* .get_m_step            = */ q4_k_gemv_adapter::get_m_step,
+            /* .get_n_step            = */ q4_k_gemv_adapter::get_n_step,
+            /* .get_mr                = */ q4_k_gemv_adapter::get_mr,
+            /* .get_nr                = */ q4_k_gemv_adapter::get_nr,
+            /* .get_kr                = */ q4_k_gemv_adapter::get_kr,
+            /* .get_sr                = */ q4_k_gemv_adapter::get_sr,
+            /* .get_dst_offset        = */ q4_k_gemv_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ q4_k_gemv_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ q4_k_gemv_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ q4_k_gemv_adapter::run,
        },
        {
            /* .get_offset            = */ q4_k_lhs_get_offset,
@@ -883,6 +915,54 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
        /* .rhs_type           = */ GGML_TYPE_Q8_0,
        /* .op_type            = */ GGML_TYPE_F32,
    },
+    {
+        /* SVE I8MM GEMM */
+        {
+            /* .get_m_step            = */ sve_q8_gemm_adapter::get_m_step,
+            /* .get_n_step            = */ sve_q8_gemm_adapter::get_n_step,
+            /* .get_mr                = */ sve_q8_gemm_adapter::get_mr,
+            /* .get_nr                = */ sve_q8_gemm_adapter::get_nr,
+            /* .get_kr                = */ sve_q8_gemm_adapter::get_kr,
+            /* .get_sr                = */ sve_q8_gemm_adapter::get_sr,
+            /* .get_dst_offset        = */ sve_q8_gemm_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ sve_q8_gemm_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ sve_q8_gemm_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ sve_q8_gemm_adapter::run,
+        },
+        /* .gemm_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
+            /* .pack_func_ex          = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
+        },
+        /* SVE DOT GEMV */
+        {
+            /* .get_m_step            = */ sve_q8_gemv_adapter::get_m_step,
+            /* .get_n_step            = */ sve_q8_gemv_adapter::get_n_step,
+            /* .get_mr                = */ sve_q8_gemv_adapter::get_mr,
+            /* .get_nr                = */ sve_q8_gemv_adapter::get_nr,
+            /* .get_kr                = */ sve_q8_gemv_adapter::get_kr,
+            /* .get_sr                = */ sve_q8_gemv_adapter::get_sr,
+            /* .get_dst_offset        = */ sve_q8_gemv_adapter::get_dst_offset,
+            /* .get_lhs_offset_ex     = */ sve_q8_gemv_adapter::get_lhs_packed_offset,
+            /* .get_rhs_packed_offset_ex = */ sve_q8_gemv_adapter::get_rhs_packed_offset,
+            /* .run_kernel_ex         = */ sve_q8_gemv_adapter::run,
+        },
+        /* .gemv_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
+            /* .pack_func_ex          = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
+        },
+        /* .rhs_info = */ {
+            /* .packed_size_ex        = */ &rhs_ps_fn5<kai_get_rhs_packed_size_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+            /* .pack_func_ex          = */ &rhs_pack_scale_fn12<kai_run_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+        },
+        /* .required_cpu       = */ CPU_FEATURE_SVE | CPU_FEATURE_SVEI8MM | CPU_FEATURE_DOTPROD,
+        /* .lhs_type           = */ GGML_TYPE_F32,
+        /* .rhs_type           = */ GGML_TYPE_Q8_0,
+        /* .op_type            = */ GGML_TYPE_F32,
+    },
    {
        /* I8MM GEMM */
        {
@@ -13,7 +13,8 @@ enum cpu_feature {
    CPU_FEATURE_SVE     = 4,
    CPU_FEATURE_SME     = 8,
    CPU_FEATURE_SME2    = 16,
-    CPU_FEATURE_FP16    = 32
+    CPU_FEATURE_FP16    = 32,
+    CPU_FEATURE_SVEI8MM = 64
 };

 inline cpu_feature& operator|=(cpu_feature& lhs, cpu_feature rhs) {
@@ -319,8 +319,10 @@ static void init_kleidiai_context(void) {

        size_t detected_smcus = 0;

+        // The fixed-width SVE kernels require 32-byte vectors, including the Q8 shared RHS layout.
        ctx.features  = (runtime_feat.has_dotprod  ? CPU_FEATURE_DOTPROD : CPU_FEATURE_NONE) |
                        (runtime_feat.has_i8mm     ? CPU_FEATURE_I8MM    : CPU_FEATURE_NONE) |
+                        (runtime_feat.has_svei8mm  ? CPU_FEATURE_SVEI8MM : CPU_FEATURE_NONE) |
                        (runtime_feat.has_fp16     ? CPU_FEATURE_FP16    : CPU_FEATURE_NONE) |
                        (runtime_feat.sve_cnt == QK8_0 ? CPU_FEATURE_SVE : CPU_FEATURE_NONE);

@@ -1616,6 +1618,11 @@ public:
            const size_t kr = kernel->get_kr();
            const size_t sr = kernel->get_sr();
            const ggml_type rhs_type = kernels->rhs_type;
+            if (rhs_type == GGML_TYPE_Q8_0 && (kernels->required_cpu & CPU_FEATURE_SVE)) {
+                GGML_ASSERT(nr == kernels->gemv.get_nr());
+                GGML_ASSERT(kr == kernels->gemv.get_kr());
+                GGML_ASSERT(sr == kernels->gemv.get_sr());
+            }
            const size_t block_len = rhs_type == GGML_TYPE_Q8_0 ? QK8_0 :
                                     rhs_type == GGML_TYPE_Q4_0 ? QK4_0 :
                                     rhs_type == GGML_TYPE_Q4_K ? GGML_KLEIDIAI_Q4_K_BL :
@@ -743,12 +743,17 @@ static void ggml_compute_forward_add_id_f32(
        // src1 indices
        const int i11 = *(int32_t *) ((char *) src2->data + i1*nb20 + i2*nb21);

-        GGML_ASSERT(i11 >= 0 && i11 < ne11);
+        GGML_ASSERT(i11 == -1 || (i11 >= 0 && i11 < ne11));

-        ggml_vec_add_f32(ne0,
-                (float *) ((char *) dst->data  + i3*nb3  + i2*nb2  + i1*nb1 ),
-                (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01),
-                (float *) ((char *) src1->data + i11*nb11));
+        float       * dst_row  = (float *) ((char *) dst->data  + i3*nb3  + i2*nb2  + i1*nb1 );
+        const float * src0_row = (float *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01);
+
+        if (i11 == -1) {
+            memcpy(dst_row, src0_row, ne0*sizeof(float));
+            continue;
+        }
+
+        ggml_vec_add_f32(ne0, dst_row, src0_row, (float *) ((char *) src1->data + i11*nb11));
    }
 }

@@ -4855,7 +4855,12 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
                    const int32_t i02 =
                        *(const int32_t *) ((const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);

-                    GGML_ASSERT(i02 >= 0 && i02 < n_as);
+                    GGML_ASSERT(i02 == -1 || (i02 >= 0 && i02 < n_as));
+
+                    if (i02 == -1) {
+                        memset((char *) dst->data + id*nb1 + iid1*nb2, 0, ne0*sizeof(float));
+                        continue;
+                    }

                    MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = { id, iid1 };
                    matrix_row_counts[i02] += 1;
@@ -682,7 +682,12 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
                    const int32_t i02 =
                        *(const int32_t *) ((const char *) ids->data + iid1 * ids->nb[1] + id * ids->nb[0]);

-                    GGML_ASSERT(i02 >= 0 && i02 < n_as);
+                    GGML_ASSERT(i02 == -1 || (i02 >= 0 && i02 < n_as));
+
+                    if (i02 == -1) {
+                        memset((char *) dst->data + id * nb1 + iid1 * nb2, 0, ne0 * sizeof(float));
+                        continue;
+                    }

                    MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = { id, iid1 };
                    matrix_row_counts[i02] += 1;
@@ -78,7 +78,7 @@ static inline float op_expm1(float x) {
 }

 static inline float op_softplus(float x) {
-    return (x > 20.0f) ? x : logf(1.0f + expf(x));
+    return (x > 20.0f) ? x : log1pf(expf(x));
 }

 static inline float op_floor(float x) {
@@ -1264,15 +1264,12 @@ struct ggml_tensor_extra_gpu {
 struct ggml_cuda_graph {
 #ifdef USE_CUDA_GRAPH
    ~ggml_cuda_graph() {
-        if (instance != nullptr) {
+        for (cudaGraphExec_t instance : instances) {
            CUDA_CHECK(cudaGraphExecDestroy(instance));
        }
-        if (graph != nullptr) {
-            CUDA_CHECK(cudaGraphDestroy(graph));
-        }
    }
-    cudaGraph_t graph = nullptr;
-    cudaGraphExec_t instance = nullptr;
+    // the ggml graph is captured as a series of CUDA graphs, so the GPU can run one chunk while the next is enqueued
+    std::vector<cudaGraphExec_t> instances;
    size_t num_nodes = 0;
    std::vector<cudaGraphNode_t> nodes;
    bool disable_due_to_gpu_arch = false;
@@ -1501,7 +1498,7 @@ struct ggml_backend_cuda_context {
    // Check if any CUDA graph has an instance for this context
    bool any_cuda_graph_has_instance() const {
        for (const auto & [key, graph] : cuda_graphs) {
-            if (graph && graph->instance != nullptr) {
+            if (graph && !graph->instances.empty()) {
                return true;
            }
        }
@@ -496,7 +496,7 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool {
 #ifdef USE_CUDA_GRAPH
        int n_captured = 0;
        for (const auto & [key, graph] : ctx.cuda_graphs) {
-            n_captured += graph->instance != nullptr;
+            n_captured += graph->instances.size();
        }
        if (n_captured > 0) {
            GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: memory freed, %d captured CUDA graphs must be captured again\n", device, n_captured);
@@ -1973,6 +1973,61 @@ static bool ggml_cuda_match_shared_expert(const ggml_cgraph * graph, int routed_
    return true;
 }

+// 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) {
+    const ggml_tensor * x_up = up->src[0];
+    const ggml_tensor * y    = up->src[1];
+
+    if (up->op != GGML_OP_MUL_MAT || !ggml_cuda_mmq_fusion_supported(x_up->type) || y->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    // The fused kernel is only compiled for weights whose rows fill whole tiles (no fallback).
+    if (ggml_cuda_mmq_needs_fallback(x_up->ne[1])) {
+        return false;
+    }
+
+    if (!ggml_cuda_should_fuse_mul_mat(up, gate, glu)) {
+        return false;
+    }
+
+    // The fused write-back applies SWIGLU_OAI with the default alpha and limit of ggml_cuda_op_swiglu_oai_single.
+    if (ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU_OAI && (ggml_get_op_params_f32(glu, 2) != 1.702f || ggml_get_op_params_f32(glu, 3) != 7.0f)) {
+        return false;
+    }
+
+    // NVIDIA only: the fused write-back relies on the MMA accumulator layout.
+    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) {
+    // 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;
+    }
+
+    // The padding of a compute-buffer weight that is a view cannot be cleared, see ggml_cuda_mul_mat.
+    const auto bad_padding_clear = [](const ggml_tensor * x) {
+        return ggml_backend_buffer_get_usage(x->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE &&
+            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])) {
+        return false;
+    }
+
+    // The batch size checks are done here, not in ggml_cuda_mul_mat_q_fusion_matches.
+    const ggml_tensor * x_up = up->src[0];
+    const ggml_tensor * y    = up->src[1];
+    const int           cc   = ggml_cuda_info().devices[device].cc;
+    if (ggml_cuda_should_use_mmvq(x_up->type, cc, y->ne[1]) || !ggml_cuda_should_use_mmq(x_up->type, cc, y->ne[1], 0)) {
+        return false;
+    }
+
+    // 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], y->ne[1], y->ne[2]*y->ne[3], device);
+}
+
 static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
    GGML_TENSOR_BINARY_OP_LOCALS

@@ -2803,16 +2858,14 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx
    return res;
 }

-static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) {
-    ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key);
-
+static void ggml_cuda_graph_update_executable(cudaGraphExec_t & instance, cudaGraph_t cuda_graph) {
 #if CUDART_VERSION >= 12000
    cudaGraphExecUpdateResultInfo result_info;
-    cudaError_t stat = cudaGraphExecUpdate(graph->instance, graph->graph, &result_info);
+    cudaError_t stat = cudaGraphExecUpdate(instance, cuda_graph, &result_info);
 #else
    cudaGraphNode_t errorNode;
    cudaGraphExecUpdateResult result_info;
-    cudaError_t stat = cudaGraphExecUpdate(graph->instance, graph->graph, &errorNode, &result_info);
+    cudaError_t stat = cudaGraphExecUpdate(instance, cuda_graph, &errorNode, &result_info);
 #endif // CUDART_VERSION >= 12000

    if (stat == cudaErrorGraphExecUpdateFailure) {
@@ -2823,13 +2876,39 @@ static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_c
        // The pre-existing graph exec cannot be updated due to violated constraints
        // so instead clear error and re-instantiate
        (void)cudaGetLastError();
-        CUDA_CHECK(cudaGraphExecDestroy(graph->instance));
-        graph->instance = nullptr;
-        CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0));
+        CUDA_CHECK(cudaGraphExecDestroy(instance));
+        instance = nullptr;
+        CUDA_CHECK(cudaGraphInstantiate(&instance, cuda_graph, nullptr, nullptr, 0));
    } else {
        GGML_ASSERT(stat == cudaSuccess);
    }
 }
+
+// number of ggml nodes per CUDA graph chunk, 0 captures the whole ggml graph as one CUDA graph
+static int ggml_cuda_graph_chunk_size() {
+    static const int chunk_size = [] {
+        const char * env = getenv("GGML_CUDA_GRAPH_CHUNK_SIZE");
+        return env != nullptr ? atoi(env) : 512;
+    }();
+    return chunk_size;
+}
+
+// end the capture of chunk `idx`, update or create its executable and launch it
+static void ggml_cuda_graph_launch_chunk(ggml_backend_cuda_context * cuda_ctx, ggml_cuda_graph * graph, size_t idx) {
+    cudaGraph_t cuda_graph = nullptr;
+    CUDA_CHECK(cudaStreamEndCapture(cuda_ctx->stream(), &cuda_graph));
+
+    if (idx < graph->instances.size()) {
+        ggml_cuda_graph_update_executable(graph->instances[idx], cuda_graph);
+    } else {
+        GGML_ASSERT(idx == graph->instances.size());
+        graph->instances.emplace_back(nullptr);
+        CUDA_CHECK(cudaGraphInstantiate(&graph->instances[idx], cuda_graph, nullptr, nullptr, 0));
+    }
+    CUDA_CHECK(cudaGraphDestroy(cuda_graph));
+
+    CUDA_CHECK(cudaGraphLaunch(graph->instances[idx], cuda_ctx->stream()));
+}
 #endif // USE_CUDA_GRAPH

 static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope,
@@ -2857,9 +2936,10 @@ static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope,
        return false;
    }

-    // Only norm/neox shaders have the fusion code
+    // Only norm/neox/multi shaders have the fusion code
    const int mode = ((const int32_t *) rope->op_params)[2];
-    if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX) {
+    if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX &&
+        mode != GGML_ROPE_TYPE_MROPE && mode != GGML_ROPE_TYPE_IMROPE) {
        return false;
    }

@@ -4353,6 +4433,18 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
                fused_node_count  = 3;
                break;
            }
+
+            if (ggml_cuda_should_fuse_mul_mat_q(up, gate, glu)) {
+                ggml_cuda_mm_fusion_args_host fusion_data{};
+                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);
+
+                ggml_cuda_mul_mat_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data);
+                fused_mul_mat_vec = true;
+                fused_node_count  = 3;
+                break;
+            }
        }
    }

@@ -4583,6 +4675,13 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
    ggml_cuda_concurrent_event * concurrent_event           = nullptr;
    bool                         should_launch_concurrent_events = false;

+#ifdef USE_CUDA_GRAPH
+    ggml_cuda_graph * graph      = use_cuda_graph ? cuda_ctx->cuda_graph(graph_key) : nullptr;
+    const int         chunk_size = ggml_cuda_graph_chunk_size();
+    int               chunk_end  = chunk_size;
+    size_t            n_chunks   = 0;
+#endif // USE_CUDA_GRAPH
+
    const auto try_launch_concurrent_event = [&](const ggml_tensor * node) {
        if (stream_ctx.concurrent_events.find(node) != stream_ctx.concurrent_events.end()) {
            concurrent_event = &stream_ctx.concurrent_events[node];
@@ -4672,6 +4771,15 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud

            for (int i = 0; i < cgraph->n_nodes; i++) {
                ggml_tensor * node = cgraph->nodes[i];
+#ifdef USE_CUDA_GRAPH
+                // fused nodes are skipped together, so a cut never splits a fusion
+                // a concurrent region must join before its capture can end
+                if (use_cuda_graph && chunk_size > 0 && i >= chunk_end && !is_concurrent_event_active) {
+                    ggml_cuda_graph_launch_chunk(cuda_ctx, graph, n_chunks++);
+                    CUDA_CHECK(cudaStreamBeginCapture(cuda_ctx->stream(), cudaStreamCaptureModeRelaxed));
+                    chunk_end = i + chunk_size;
+                }
+#endif // USE_CUDA_GRAPH
                if (is_concurrent_event_active) {
                    GGML_ASSERT(concurrent_event);

@@ -4754,14 +4862,14 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
        }

 #ifdef USE_CUDA_GRAPH
-        ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key);
        if (use_cuda_graph && cuda_graph_update_required) { // End CUDA graph capture
-            if (graph->graph != nullptr) {
-                CUDA_CHECK(cudaGraphDestroy(graph->graph));
-                graph->graph = nullptr;
-            }
+            ggml_cuda_graph_launch_chunk(cuda_ctx, graph, n_chunks++);

-            CUDA_CHECK(cudaStreamEndCapture(cuda_ctx->stream(), &graph->graph));
+            // drop the chunks left over from a longer graph
+            for (size_t i = n_chunks; i < graph->instances.size(); ++i) {
+                CUDA_CHECK(cudaGraphExecDestroy(graph->instances[i]));
+            }
+            graph->instances.resize(n_chunks);
            graph_evaluated_or_captured = true; // CUDA graph has been captured

            std::lock_guard<std::mutex> lock(ggml_cuda_lock);
@@ -4773,16 +4881,11 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
        }
    }

-    if (use_cuda_graph) {
-        ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key);
-        if (graph->instance == nullptr) { // Create executable graph from captured graph.
-            CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0));
+    if (use_cuda_graph && !cuda_graph_update_required) {
+        // the chunks were launched during capture, otherwise launch them here
+        for (cudaGraphExec_t instance : graph->instances) {
+            CUDA_CHECK(cudaGraphLaunch(instance, cuda_ctx->stream()));
        }
-        if (cuda_graph_update_required) { // Update graph executable
-            ggml_cuda_graph_update_executable(cuda_ctx, graph_key);
-        }
-        // Launch graph
-        CUDA_CHECK(cudaGraphLaunch(graph->instance, cuda_ctx->stream()));
 #else
        GGML_UNUSED(graph_key);
        graph_evaluated_or_captured = true;
@@ -4794,7 +4897,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
 static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) {
    ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key);

-    if (graph->graph == nullptr) {
+    if (graph->instances.empty()) {
        if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_VOLTA) {
            if (!graph->disable_due_to_gpu_arch && ggml_cuda_graph_debug_enabled()) {
                GGML_LOG_DEBUG("%s: disabling CUDA graphs due to GPU architecture\n", __func__);
@@ -4848,7 +4951,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend,
                    }
                } else {
                    use_cuda_graph = true;
-                    cuda_graph_update_required = graph->instance == nullptr;
+                    cuda_graph_update_required = graph->instances.empty();
                }
            }
        }
@@ -4938,6 +5041,19 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph
            }
        }
        for (int i = 0; i < cgraph->n_nodes; ++i) {
+            // Keep the inputs of the fused gate/up/GLU MMQ kernel alive until the GLU so that they cannot alias its output (under op offload the weights are buffer copies).
+            if (i + 2 < cgraph->n_nodes && cgraph->nodes[i + 2]->op == GGML_OP_GLU) {
+                ggml_tensor * glu  = cgraph->nodes[i + 2];
+                ggml_tensor * gate = glu->src[0];
+                ggml_tensor * up   = glu->src[1];
+                const bool ok = (gate == cgraph->nodes[i] && up == cgraph->nodes[i + 1]) || (gate == cgraph->nodes[i + 1] && up == cgraph->nodes[i]);
+                if (ok && ggml_cuda_mul_mat_q_fusion_matches(up, gate, glu, cuda_ctx->device)) {
+                    add_alloc_deps(i, i + 2);
+                    i += 2;
+                    continue;
+                }
+            }
+
            ggml_cuda_moe_weighted_reduction_match match;
            if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) {
                params->add_alloc_dep(params->user_data, const_cast<ggml_tensor *>(match.experts), match.dst);
@@ -1409,14 +1409,14 @@ namespace ggml_cuda_mma {
            tile<16, 16, float, dl_d> & D, const tile<16, 8, nv_bfloat162, dl_ab> & A, const tile<16, 8, nv_bfloat162, dl_ab> & B) {
 #if defined(AMD_WMMA_AVAILABLE)
 #if defined(RDNA4)
-        using bf16x8_t = __attribute__((ext_vector_type(8))) __bf16;
+        using bf16x8_t = __attribute__((ext_vector_type(8))) short;
        using floatx8_t = __attribute__((ext_vector_type(8))) float;
        floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]);
        const bf16x8_t& a_frag = reinterpret_cast<const bf16x8_t&>(A.x[0]);
        const bf16x8_t& b_frag = reinterpret_cast<const bf16x8_t&>(B.x[0]);
        acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12(a_frag, b_frag, acc_frag);
 #elif defined(RDNA3)
-        using bf16x16_t = __attribute__((ext_vector_type(16))) __bf16;
+        using bf16x16_t = __attribute__((ext_vector_type(16))) short;
        using floatx8_t = __attribute__((ext_vector_type(8))) float;
        floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]);
        const bf16x16_t& a_frag = reinterpret_cast<const bf16x16_t&>(A.x[0]);
@@ -5,7 +5,9 @@
 #include "mmq.cuh"

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q1_0(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -60,10 +62,10 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            const uint32_t q2_indices = (q2_bits & 0x01) | ((q2_bits & 0x02) << 7) | ((q2_bits & 0x04) << 14) | ((q2_bits & 0x08) << 21);
            const uint32_t q3_indices = (q3_bits & 0x01) | ((q3_bits & 0x02) << 7) | ((q3_bits & 0x04) << 14) | ((q3_bits & 0x08) << 21);

-            const int v0 = __builtin_amdgcn_perm(0x000001FF, 0x000001FF, q0_indices);
-            const int v1 = __builtin_amdgcn_perm(0x000001FF, 0x000001FF, q1_indices);
-            const int v2 = __builtin_amdgcn_perm(0x000001FF, 0x000001FF, q2_indices);
-            const int v3 = __builtin_amdgcn_perm(0x000001FF, 0x000001FF, q3_indices);
+            const int v0 = ggml_hip_perm(0x000001FF, 0x000001FF, q0_indices);
+            const int v1 = ggml_hip_perm(0x000001FF, 0x000001FF, q1_indices);
+            const int v2 = ggml_hip_perm(0x000001FF, 0x000001FF, q2_indices);
+            const int v3 = ggml_hip_perm(0x000001FF, 0x000001FF, q3_indices);
 #else
            // unpack crumbs into nibble indices
            const int n0 = __byte_perm(0x11100100, 0x11100100, q >> 0); // [0, 1, 4, 5] [ 8,  9, 12, 13]
@@ -116,7 +118,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_0(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -162,8 +166,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            const uint32_t qx_indices = (q & 0x03) | ((q & 0x0C) << 6) | ((q & 0x30) << 12) | ((q & 0xC0) << 18);
            const uint32_t qy_bits    = q >> 8;
            const uint32_t qy_indices = (qy_bits & 0x03) | ((qy_bits & 0x0C) << 6) | ((qy_bits & 0x30) << 12) | ((qy_bits & 0xC0) << 18);
-            const int qx = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qx_indices);
-            const int qy = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qy_indices);
+            const int      qx         = ggml_hip_perm(0x020100FF, 0x020100FF, qx_indices);
+            const int      qy         = ggml_hip_perm(0x020100FF, 0x020100FF, qy_indices);
 #else
            // unpack even and odd crumbs into byte values
            const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0);
@@ -205,7 +209,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_0(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -268,7 +274,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_1(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -331,7 +339,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_0(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -411,7 +421,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_1(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -488,8 +500,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    }
 }

-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q8_0(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q8_0(
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -519,7 +531,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

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

 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
        x_qs[i*sram_stride + 0             + txi] = get_int_b2(bxi[0].qs,                   kqsx);
@@ -542,7 +554,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

-        const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbxd;
+        const block_q8_0 * bxi = ggml_cuda_mmq_x_row<block_q8_0, has_fusion>(x, x_gate, kbx0, i, stride) + kbxd;

 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
        x_df[i*sram_stride                           + kbxd] = bxi->d;
@@ -555,7 +567,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 // ---------------------------------------------------------------------------------------------

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_K(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -616,7 +630,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q3_K(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -728,8 +744,8 @@ static __device__ __forceinline__ int unpack_scales_q45_K(const int * scales, co
           ((scales[ksc/2]              >> (2 * (ksc % 2)))       & 0x30303030);  // upper 2 bits
 }

-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K(
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -757,7 +773,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

-        const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+        const block_q4_K * bxi = ggml_cuda_mmq_x_row<block_q4_K, has_fusion>(x, x_gate, kbx0, i, stride);
        const int qs0 = get_int_b4(bxi->qs, txi);

 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
@@ -786,7 +802,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
                i = min(i, i_max);
            }

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

            const int * scales = (const int *) bxi->scales;
            const int ksc = threadIdx.x % 2;
@@ -814,7 +830,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

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

        x_dm[i] = bxi->dm;
    }
@@ -827,7 +843,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

-        const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / (QI4_K/8);
+        const block_q4_K * bxi = ggml_cuda_mmq_x_row<block_q4_K, has_fusion>(x, x_gate, kbx0, i, stride) + (threadIdx.x % (MMQ_TILE_NE_K/8)) / (QI4_K/8);

        const int * scales = (const int *) bxi->scales;

@@ -839,8 +855,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
 }

-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K(
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -868,7 +884,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

-        const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+        const block_q5_K * bxi = ggml_cuda_mmq_x_row<block_q5_K, has_fusion>(x, x_gate, kbx0, i, stride);
        const int ky = QR5_K*txi;

        const int ql = get_int_b4(bxi->qs, txi);
@@ -909,7 +925,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
                i = min(i, i_max);
            }

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

            const int * scales = (const int *) bxi->scales;
            const int ksc = threadIdx.x % 2;
@@ -937,7 +953,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

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

        x_dm[i] = bxi->dm;
    }
@@ -951,7 +967,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

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

        const int * scales = (const int *) bxi->scales;

@@ -963,8 +979,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
 }

-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q6_K(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q6_K(
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -993,7 +1009,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

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

        const int ql = get_int_b2(bxi->ql, txi);
        const int ql0 = (ql >> 0) & 0x0F0F0F0F;
@@ -1023,7 +1039,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

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

 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
        x_df[i*sram_stride]                     = bxi->d;
@@ -1041,7 +1057,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

-        const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / 4;
+        const block_q6_K * bxi = ggml_cuda_mmq_x_row<block_q6_K, has_fusion>(x, x_gate, kbx0, i, stride) + (threadIdx.x % (MMQ_TILE_NE_K/8)) / 4;

 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
        x_sc[i*sram_stride + threadIdx.x%4] = get_int_b2(bxi->scales, threadIdx.x % (MMQ_TILE_NE_K/8));
@@ -1054,7 +1070,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 // ---------------------------------------------------------------------------------------------

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq1_s(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1116,7 +1134,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_xxs(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1180,7 +1200,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_xs(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1245,7 +1267,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_s(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1313,7 +1337,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq3_xxs(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1377,7 +1403,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq3_s(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1445,8 +1473,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
    }
 }

-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_xs(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_xs(
+        const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
    constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
    constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1473,7 +1501,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

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

        const int aux_q4 = get_int_b4(bxi->qs, kqsx);
        const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl);
@@ -1497,7 +1525,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            i = min(i, i_max);
        }

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

        const float d = __half2float(bxi->d);

@@ -1513,7 +1541,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_nl(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1582,7 +1612,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 // ---------------------------------------------------------------------------------------------

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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_Q8) / warp_size;
    constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1649,7 +1681,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 }

 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4_fp4(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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);
@@ -1691,7 +1725,9 @@ 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(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kb0, const int i_max, const int stride) {
+        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);
@@ -1747,7 +1783,9 @@ 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(
-        const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+        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);
@@ -131,12 +131,51 @@ static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggm
    return GGML_PREC_Q4;
 }

+// Tile width in ne11 (dense) / ne12 (MoE) direction with the fewest tiles that fits into shared memory, 0 if there is none.
+static int ggml_cuda_mmq_get_J_best(
+        const ggml_type type, const int64_t ncols, const bool fallback, const int cc, const size_t smpbo, const ggml_prec prec_src1,
+        int * nthreads_best = nullptr) {
+    int J_best        = 0;
+    int ntiles_J_best = INT_MAX;
+
+    for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) {
+        const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1);
+        if (config.type == GGML_TYPE_COUNT) {
+            continue;
+        }
+
+        if (mmq_get_nbytes_shared(config, cc) > smpbo) {
+            continue;
+        }
+
+        const int ntiles_x = (ncols + config.J - 1) / config.J;
+
+        if (ntiles_x < ntiles_J_best) {
+            J_best = J;
+            if (nthreads_best) {
+                *nthreads_best = config.nthreads;
+            }
+            ntiles_J_best = ntiles_x;
+        }
+    }
+    return J_best;
+}
+
 void ggml_cuda_mul_mat_q(
-        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) {
+        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
+        const ggml_cuda_mm_fusion_args_host * fusion) {
    GGML_ASSERT(        src1->type == GGML_TYPE_F32);
    GGML_ASSERT(        dst->type  == GGML_TYPE_F32);
    GGML_ASSERT(!ids || ids->type  == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID.

+    // Fused gate/up/GLU: src0 is the up weight, dst the GLU output.
+    const ggml_tensor * gate = fusion ? fusion->gate : nullptr;
+    if (fusion) {
+        // TODO MoE support
+        GGML_ASSERT(gate && !ids && !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));
+    }
+
    GGML_TENSOR_BINARY_OP_LOCALS;

    cudaStream_t stream = ctx.stream();
@@ -158,14 +197,16 @@ void ggml_cuda_mul_mat_q(
    const float * src1_d = (const float *) src1->data;
    float       *  dst_d = (float       *)  dst->data;

-    // If src0 is a temporary compute buffer, clear any potential padding.
-    if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
-        const size_t size_data  = ggml_nbytes(src0);
-        const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0);
-        if (size_alloc > size_data) {
-            GGML_ASSERT(ggml_is_contiguously_allocated(src0));
-            GGML_ASSERT(!src0->view_src);
-            CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream));
+    // If src0 (or the fused gate weight) is a temporary compute buffer, clear any potential padding.
+    for (const ggml_tensor * t : { src0, gate }) {
+        if (t && ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+            const size_t size_data  = ggml_nbytes(t);
+            const size_t size_alloc = ggml_backend_buffer_get_alloc_size(t->buffer, t);
+            if (size_alloc > size_data) {
+                GGML_ASSERT(ggml_is_contiguously_allocated(t));
+                GGML_ASSERT(!t->view_src);
+                CUDA_CHECK(cudaMemsetAsync((char *) t->data + size_data, 0, size_alloc - size_data, stream));
+            }
        }
    }

@@ -180,7 +221,9 @@ void ggml_cuda_mul_mat_q(

    const bool fallback = ggml_cuda_mmq_needs_fallback(ne01);

-    const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0, dst, cc);
+    // 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);

    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);
@@ -201,26 +244,7 @@ void ggml_cuda_mul_mat_q(
            }
        }

-        int ntiles_J_best = INT_MAX;
-
-        for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) {
-            const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(src0->type, J, fallback, cc, prec_src1);
-            if (config.type == GGML_TYPE_COUNT) {
-                continue;
-            }
-
-            if (mmq_get_nbytes_shared(config, cc) > smpbo) {
-                continue;
-            }
-
-            const int ntiles_x = (ncols_opt + config.J - 1) / config.J;
-
-            if (ntiles_x < ntiles_J_best) {
-                J_best = J;
-                nthreads_best = config.nthreads;
-                ntiles_J_best = ntiles_x;
-            }
-        }
+        J_best = ggml_cuda_mmq_get_J_best(src0->type, ncols_opt, fallback, cc, smpbo, prec_src1, &nthreads_best);
    }
    GGML_ASSERT(J_best > 0);

@@ -262,13 +286,20 @@ void ggml_cuda_mul_mat_q(
                                ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
        const int64_t s13 = ne12*s12;

+        ggml_cuda_mm_fusion_args_device fusion_device = {};
+        if (gate) {
+            fusion_device.gate      = gate->data;
+            fusion_device.glu_op    = fusion->glu_op;
+            fusion_device.glu_limit = fusion->glu_limit;
+        }
+
        const mmq_args args = {
            src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d,
            src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr,
            ne00, ne01, ne1, s01, ne11, s1,
            ne02, ne12, s02, s12, s2,
            ne03, ne13, s03, s13, s3,
-            ne1, J_best};
+            ne1, J_best, fusion_device};
        ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
        return;
    }
@@ -485,3 +516,22 @@ 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) {
+    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);
+    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 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;
+    const int64_t tiles_efficiency_percent = 100 * ntiles / (nsm*tiles_nwaves);
+    return tiles_efficiency_percent >= 90;
+}
@@ -1,6 +1,7 @@
 #pragma once

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

 #include <climits>
 #include <cstdint>
@@ -10,10 +11,12 @@
 #define MMQ_ITER_K_FP4         512
 #define MMQ_NWARPS               8

-typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride);
+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);
+    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);

 enum mmq_q8_1_ds_layout {
    MMQ_Q8_1_DS_LAYOUT_D4,
@@ -426,15 +429,41 @@ 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;
+}
+
+// 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.
+static constexpr int GGML_CUDA_MMQ_FUSION_ROWS = 8;
+
+// Source of weight tile row i; with has_fusion the up or gate row that the interleaved tile holds there.
+template <typename block_t, bool has_fusion>
+static __device__ __forceinline__ const block_t * ggml_cuda_mmq_x_row(
+        const char * __restrict__ x, const char * __restrict__ x_gate, const int kbx0, const int i, const int stride) {
+    if constexpr (has_fusion) {
+        constexpr int R = GGML_CUDA_MMQ_FUSION_ROWS;
+        const char * x_src = (i / R) % 2 == 0 ? x : x_gate;
+        return (const block_t *) x_src + kbx0 + ((i / (2*R))*R + i % R)*stride;
+    } else {
+        GGML_UNUSED(x_gate);
+        return (const block_t *) x + kbx0 + i*stride;
+    }
+}
+
 #include "mmq-load-tiles.cuh"
 #include "mmq-vec-dot.cuh"

 template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a(
        const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst,
-        const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
+        const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+        const ggml_glu_op glu_op, const float glu_limit) {
+    GGML_UNUSED_VARS(glu_op, glu_limit);
+
    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;

@@ -468,16 +497,19 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
    }
 }

-template<ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+// 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.
+template<ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
            const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst,
-            const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
+            const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+            const ggml_glu_op glu_op, const float glu_limit) {

 #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
    typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
 #else
    typedef tile<16,  8, int> tile_C;
 #endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    static_assert(!has_fusion || tile_C::I == 2*GGML_CUDA_MMQ_FUSION_ROWS, "each C tile must hold one chunk of up rows and one of gate rows");

    constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, prec_src1);
    constexpr int ntx           = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -491,20 +523,23 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
 #pragma unroll
        for (int n = 0; n < ntx; ++n) {
 #pragma unroll
-            for (int l = 0; l < tile_C::ne; ++l) {
+            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) {
                    continue;
                }

-                const int i = i0 + n*tile_C::I + tile_C::get_i(l);
+                const int i = (i0 + n*tile_C::I)/(has_fusion ? 2 : 1) + tile_C::get_i(l);

                if (fallback && i > i_max) {
                    continue;
                }

-                if constexpr (type == GGML_TYPE_NVFP4) {
+                if constexpr (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 {
@@ -517,6 +552,8 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
            }
        }
    }
+
+    GGML_UNUSED_VARS(y_scale_used, glu_op, glu_limit);
 }

 // -------------------------------------------------------------------------------------------------------------------------------------
@@ -533,8 +570,9 @@ struct ggml_cuda_mmq_util_funcs {
        vdr(vdr), load_tiles(load_tiles), vec_dot(vec_dot), write_back(write_back) {}
 };

-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_funcs() {
+    static_assert(!has_fusion || ggml_cuda_mmq_fusion_supported(type), "no fused loader for this type");
    if (!ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).use_mma_data_layout()) {
        switch (type) {
            case GGML_TYPE_Q1_0:
@@ -576,7 +614,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
            case GGML_TYPE_Q8_0:
                return ggml_cuda_mmq_util_funcs(
                    VDR_Q8_0_Q8_1_MMQ,
-                    ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback>,
+                    ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback, false>,
                    ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
 // ---------------------------------------------------------------------------------------------
@@ -595,19 +633,19 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
            case GGML_TYPE_Q4_K:
                return ggml_cuda_mmq_util_funcs(
                    VDR_Q4_K_Q8_1_MMQ,
-                    ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback>,
+                    ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback, false>,
                    ggml_cuda_mmq_vec_dot_q4_K_q8_1_dp4a<type, J, fallback>,
                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
            case GGML_TYPE_Q5_K:
                return ggml_cuda_mmq_util_funcs(
                    VDR_Q5_K_Q8_1_MMQ,
-                    ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback>,
+                    ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback, false>,
                    ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a<type, J, fallback>,
                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
            case GGML_TYPE_Q6_K:
                return ggml_cuda_mmq_util_funcs(
                    VDR_Q6_K_Q8_1_MMQ,
-                    ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback>,
+                    ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback, false>,
                    ggml_cuda_mmq_vec_dot_q6_K_q8_1_dp4a<type, J, fallback>,
                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
 // ---------------------------------------------------------------------------------------------
@@ -650,7 +688,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
            case GGML_TYPE_IQ4_XS:
                return ggml_cuda_mmq_util_funcs(
                    VDR_IQ4_XS_Q8_1_MMQ,
-                    ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback>,
+                    ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback, false>,
                    ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
                    ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
            case GGML_TYPE_IQ4_NL:
@@ -687,7 +725,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
                    -1,
                    ggml_cuda_mmq_load_tiles_mxfp4_fp4<type, J, fallback>,
                    ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
-                    ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                    ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
            }
            break;
        case GGML_TYPE_NVFP4:
@@ -696,7 +734,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
                    -1,
                    ggml_cuda_mmq_load_tiles_nvfp4_nvfp4<type, J, fallback>,
                    ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
-                    ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                    ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
            }
            break;
        default:
@@ -712,136 +750,136 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
                -1,
                ggml_cuda_mmq_load_tiles_q1_0<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_Q2_0:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_q2_0<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_Q4_0:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_q4_0<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_DS4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_Q4_1:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_q4_1<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_Q5_0:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_q5_0<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_Q5_1:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_q5_1<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_Q8_0:
            return ggml_cuda_mmq_util_funcs(
                -1,
-                ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback>,
+                ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback, has_fusion>,
                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
 // ---------------------------------------------------------------------------------------------
        case GGML_TYPE_Q2_K:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_q2_K<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q2_K_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_Q3_K:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_q3_K<type, J, fallback>,
                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>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_Q4_K:
            return ggml_cuda_mmq_util_funcs(
                -1,
-                ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback>,
+                ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback, has_fusion>,
                ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
        case GGML_TYPE_Q5_K:
            return ggml_cuda_mmq_util_funcs(
                -1,
-                ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback>,
+                ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback, has_fusion>,
                ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
        case GGML_TYPE_Q6_K:
            return ggml_cuda_mmq_util_funcs(
                -1,
-                ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback>,
+                ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback, has_fusion>,
                ggml_cuda_mmq_vec_dot_q6_K_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
 // ---------------------------------------------------------------------------------------------
        case GGML_TYPE_IQ1_S:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_iq1_s<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_IQ2_XXS:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_iq2_xxs<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_IQ2_XS:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_iq2_xs<type, J, fallback>,
                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>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_IQ2_S:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_iq2_s<type, J, fallback>,
                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>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_IQ3_XXS:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_iq3_xxs<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_IQ3_S:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_iq3_s<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        case GGML_TYPE_IQ4_XS:
            return ggml_cuda_mmq_util_funcs(
                -1,
-                ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback>,
+                ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback, has_fusion>,
                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
        case GGML_TYPE_IQ4_NL:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_iq4_nl<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
 // ---------------------------------------------------------------------------------------------
        case GGML_TYPE_MXFP4:
            return ggml_cuda_mmq_util_funcs(
                -1,
                ggml_cuda_mmq_load_tiles_mxfp4<type, J, fallback>,
                ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        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_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
-                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+                ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
        default:
            return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr);
    }
@@ -849,41 +887,44 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func

 template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
 static constexpr __device__ int ggml_cuda_mmq_get_vdr() {
-    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().vdr;
+    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1, false>().vdr;
 }

-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 static constexpr __device__ ggml_cuda_mmq_load_tiles_t ggml_cuda_mmq_get_load_tiles() {
-    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().load_tiles;
+    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1, has_fusion>().load_tiles;
 }

 template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
 static constexpr __device__ ggml_cuda_mmq_vec_dot_t ggml_cuda_mmq_get_vec_dot() {
-    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().vec_dot;
+    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1, false>().vec_dot;
 }

-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_back() {
-    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().write_back;
+    return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1, has_fusion>().write_back;
 }

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

-template <ggml_type type, int J, bool fallback, bool fixup, ggml_prec prec_src1>
+// has_fusion: x and x_gate are the up and gate weights.
+template <ggml_type type, int J, bool fallback, bool fixup, ggml_prec prec_src1, bool has_fusion>
 static __device__ __forceinline__ void mul_mat_q_process_tile(
-        const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
+        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 float * __restrict__ y_scale, const ggml_glu_op glu_op, const float glu_limit,
        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) {
+    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");

    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              qk         = ggml_cuda_type_traits<type>::qk;
    constexpr int              I          = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
-    constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles<type, J, fallback, prec_src1>();
+    constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles<type, J, fallback, prec_src1, has_fusion>();
    constexpr ggml_cuda_mmq_vec_dot_t    vec_dot    = ggml_cuda_mmq_get_vec_dot<type, J, fallback, prec_src1>();
-    constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back<type, J, fallback, prec_src1>();
+    constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back<type, J, fallback, prec_src1, has_fusion>();

    extern __shared__ int data_mul_mat_q[];
    int * tile_y = data_mul_mat_q + J;
@@ -906,7 +947,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, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x);
+        load_tiles(x, x_gate, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x);
        {
            const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz;
 #pragma unroll
@@ -941,28 +982,34 @@ 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);
+        write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J, glu_op, glu_limit);
    } else {
-        write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
+        write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j, glu_op, glu_limit);
    }
 }


 // The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598

-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 __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 int * __restrict__ y, const int32_t * __restrict__ ids_dst,
+        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 float * __restrict__ y_scale,
+        const float * __restrict__ y_scale, const ggml_glu_op glu_op, const float glu_limit,
        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) {

+#ifdef TURING_MMA_AVAILABLE
+    constexpr bool fusion_available = true;
+#else
+    constexpr bool fusion_available = false;
+#endif // TURING_MMA_AVAILABLE
+
    // Skip unused template specializations for faster compilation:
-    if (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type == GGML_TYPE_COUNT) {
+    if (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type == GGML_TYPE_COUNT || (has_fusion && !fusion_available)) {
        NO_DEVICE_CODE;
        return;
    }
@@ -990,12 +1037,16 @@ static __global__ void mul_mat_q(
    }
    __syncthreads();

-    if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback, prec_src1)) {
+    // The fused kernel needs the whole K range of a tile in one block, so it always uses tiling.
+    if constexpr (has_fusion || !ggml_cuda_mmq_get_stream_k(type, J, fallback, prec_src1)) {
+        // The fused kernel writes the GLU of I/2 up rows and I/2 gate rows per tile.
+        constexpr int I_dst = has_fusion ? I/2 : I;
+
        const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y);
        const int wt = tmp2.x;
        const int zt = tmp2.y;
-        const int jt = blockIdx.y;
-        const int it = blockIdx.x;
+        const int jt = blockIdx.x;
+        const int it = blockIdx.y;

        // Defaults for regular matrix multiplication:
        int col_low    = 0;
@@ -1040,22 +1091,22 @@ static __global__ void mul_mat_q(
        }

        offset_y   += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int));
-        offset_dst += it*I;
+        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;
        }

-        const int tile_x_max_i = nrows_x  - it*I - 1;
+        const int tile_x_max_i = nrows_x  - it*I_dst - 1;
        const int tile_y_max_j = col_diff - jt*J - 1;

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

        constexpr bool fixup = false;
-        mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
-            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
-             stride_row_x, ncols_y, stride_col_dst,
+        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);
        return;
    }
@@ -1147,9 +1198,9 @@ static __global__ void mul_mat_q(
        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>
-            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
-             stride_row_x, ncols_y, stride_col_dst,
+        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);

        kbc += blocks_per_ne00.z;
@@ -1231,9 +1282,9 @@ static __global__ void mul_mat_q(
    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>
-        (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
-         stride_row_x, ncols_y, stride_col_dst,
+    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);
 }

@@ -1383,6 +1434,7 @@ struct mmq_args {
    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;
    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
 };

 static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) {
@@ -1392,7 +1444,7 @@ static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const i
    return nbs_ids + nbs_x + GGML_PAD(nbs_y, config.nthreads*sizeof(int));
 }

-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
 static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
    const int id = ggml_cuda_get_device();
    const int cc = ggml_cuda_info().devices[id].cc;
@@ -1406,13 +1458,17 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a

    const dim3 block_dims(warp_size, nwarps, 1);

-    CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, false, prec_src1>), nbytes_shared);
-    CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J,  true, prec_src1>), nbytes_shared);
+    CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, /*fallback =*/ false, prec_src1, has_fusion>), nbytes_shared);
+    if constexpr (!has_fusion) {
+        CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, /*fallback =*/ true, prec_src1, false>), nbytes_shared);
+    }

-    const int nty  = (args.nrows_x   + config.I - 1) / config.I;
+    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 dim3 block_nums_xy_tiling(nty, ntx, ntzw);
+    // 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);

    GGML_ASSERT(args.nchannels_y % args.nchannels_x == 0);
    GGML_ASSERT(args.nsamples_y  % args.nsamples_x  == 0);
@@ -1426,9 +1482,10 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
    const uint3 channel_ratio_fd   = init_fastdiv_values(channel_ratio);
    const uint3 sample_ratio_fd    = init_fastdiv_values(sample_ratio);

-    if (!config.stream_k) {
-        mul_mat_q<type, J, fallback, prec_src1><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
-            (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale,
+    if (has_fusion || !config.stream_k) {
+        GGML_ASSERT(nty <= 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,
             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,
@@ -1458,8 +1515,8 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
    const dim3 block_nums_fixup(block_nums_stream_k.x, config.I/warp_size, 1);
    const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);

-    mul_mat_q<type, J, fallback, prec_src1><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
-        (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale,
+    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,
         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,
@@ -1476,56 +1533,56 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
         ntx_fd);
 }

-template <ggml_type type, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, bool fallback, ggml_prec prec_src1, bool has_fusion>
 void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
    switch (args.J_best) {
        case   8:
-            launch_mul_mat_q<type,   8, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,   8, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case  16:
-            launch_mul_mat_q<type,  16, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  16, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case  24:
-            launch_mul_mat_q<type,  24, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  24, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case  32:
-            launch_mul_mat_q<type,  32, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  32, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case  40:
-            launch_mul_mat_q<type,  40, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  40, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case  48:
-            launch_mul_mat_q<type,  48, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  48, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case  56:
-            launch_mul_mat_q<type,  56, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  56, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case  64:
-            launch_mul_mat_q<type,  64, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  64, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case  72:
-            launch_mul_mat_q<type,  72, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  72, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case  80:
-            launch_mul_mat_q<type,  80, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  80, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case  88:
-            launch_mul_mat_q<type,  88, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  88, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case  96:
-            launch_mul_mat_q<type,  96, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type,  96, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case 104:
-            launch_mul_mat_q<type, 104, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type, 104, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case 112:
-            launch_mul_mat_q<type, 112, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type, 112, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case 120:
-            launch_mul_mat_q<type, 120, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type, 120, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        case 128:
-            launch_mul_mat_q<type, 128, fallback, prec_src1>(ctx, args, stream);
+            launch_mul_mat_q<type, 128, fallback, prec_src1, has_fusion>(ctx, args, stream);
            break;
        default:
            fprintf(stderr, "J_best=%d\n", args.J_best);
@@ -1533,14 +1590,26 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
    }
 }

+template <ggml_type type, bool fallback, ggml_prec prec_src1>
+static void mul_mat_q_switch_fusion(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
+    if constexpr (!fallback && ggml_cuda_mmq_fusion_supported(type)) {
+        if (args.fusion.gate) {
+            mul_mat_q_switch_J<type, fallback, prec_src1, true>(ctx, args, stream);
+            return;
+        }
+    }
+    GGML_ASSERT(!args.fusion.gate);
+    mul_mat_q_switch_J<type, fallback, prec_src1, false>(ctx, args, stream);
+}
+
 template <ggml_type type, ggml_prec prec_src1>
 void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
    if (ggml_cuda_mmq_needs_fallback(args.nrows_x)) {
        constexpr bool fallback = true;
-        mul_mat_q_switch_J<type, fallback, prec_src1>(ctx, args, stream);
+        mul_mat_q_switch_fusion<type, fallback, prec_src1>(ctx, args, stream);
    } else {
        constexpr bool fallback = false;
-        mul_mat_q_switch_J<type, fallback, prec_src1>(ctx, args, stream);
+        mul_mat_q_switch_fusion<type, fallback, prec_src1>(ctx, args, stream);
    }
 }

@@ -1582,6 +1651,9 @@ extern DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4);
 // -------------------------------------------------------------------------------------------------------------------------

 void ggml_cuda_mul_mat_q(
-        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst);
+        ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
+        const ggml_cuda_mm_fusion_args_host * fusion = nullptr);

 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);
@@ -196,9 +196,9 @@ static __global__ void rope_neox(const T *            x,
    dst[idst + n_offs/2 + n_dims / 2] = ggml_cuda_cast<D>(x0 * sin_theta + x1 * cos_theta);
 }

-template <bool forward, bool has_ff, typename T>
+template <bool forward, bool has_ff, typename T, typename D>
 static __global__ void rope_multi(const T *            x,
-                                  T *                  dst,
+                                  D *                  dst,
                                  const int            ne00,
                                  const int            ne01,
                                  const int            ne02,
@@ -219,6 +219,8 @@ static __global__ void rope_multi(const T *            x,
                                  const float *        freq_factors,
                                  const mrope_sections sections,
                                  const bool           is_imrope,
+                                  const int64_t *      row_indices,
+                                  const int            set_rows_stride,
                                  const bool           inplace) {
    const int i0 = 2 * (blockDim.y * blockIdx.y + threadIdx.y);

@@ -236,12 +238,20 @@ static __global__ void rope_multi(const T *            x,
    const int ix   = i0 / 2 + i1 * s01 + i2 * s02 + i3 * s03;

    ggml_cuda_pdl_sync();
+
+    // Fusion optimization: ROPE + VIEW + SET_ROWS.
+    // The rope output is viewed as a 1D tensor and offset based on a row index in row_indices.
+    if (set_rows_stride != 0) {
+        idst = i1 * s1 + i0 / 2;
+        idst += row_indices[i2] * set_rows_stride;
+    }
+
    if (i0 < n_offs || i0 >= n_offs + n_dims) {
        if (inplace) {
            return;
        }
-        dst[idst + i0/2 + 0] = x[ix + i0/2 + 0];
-        dst[idst + i0/2 + 1] = x[ix + i0/2 + 1];
+        dst[idst + i0/2 + 0] = ggml_cuda_cast<D>(x[ix + i0/2 + 0]);
+        dst[idst + i0/2 + 1] = ggml_cuda_cast<D>(x[ix + i0/2 + 1]);

        return;
    }
@@ -286,8 +296,8 @@ static __global__ void rope_multi(const T *            x,
    const float x0 = x[ix + n_offs/2 + 0];
    const float x1 = x[ix + n_offs/2 + n_dims/2];

-    dst[idst + n_offs/2 + 0]        = x0*cos_theta - x1*sin_theta;
-    dst[idst + n_offs/2 + n_dims/2] = x0*sin_theta + x1*cos_theta;
+    dst[idst + n_offs/2 + 0]        = ggml_cuda_cast<D>(x0*cos_theta - x1*sin_theta);
+    dst[idst + n_offs/2 + n_dims/2] = ggml_cuda_cast<D>(x0*sin_theta + x1*cos_theta);
 }

 template <bool forward, bool has_ff, typename T>
@@ -443,9 +453,9 @@ static void rope_neox_cuda(const T *            x,
    }
 }

-template <bool forward, typename T>
+template <bool forward, typename T, typename D>
 static void rope_multi_cuda(const T *            x,
-                            T *                  dst,
+                            D *                  dst,
                            const int            ne00,
                            const int            ne01,
                            const int            ne02,
@@ -467,6 +477,8 @@ static void rope_multi_cuda(const T *            x,
                            const float *        freq_factors,
                            const mrope_sections sections,
                            const bool           is_imrope,
+                            const int64_t *      row_indices,
+                            const int            set_rows_stride,
                            const bool           inplace,
                            cudaStream_t         stream) {
    GGML_ASSERT(ne00 % 2 == 0);
@@ -478,14 +490,14 @@ static void rope_multi_cuda(const T *            x,

    if (freq_factors == nullptr) {
        const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
-        ggml_cuda_kernel_launch(rope_multi<forward, false, T>, launch_params,
+        ggml_cuda_kernel_launch(rope_multi<forward, false, T, D>, launch_params,
            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor,
-            attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, inplace);
+            attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, row_indices, set_rows_stride, inplace);
    } else {
        const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
-        ggml_cuda_kernel_launch(rope_multi<forward, true, T>, launch_params,
+        ggml_cuda_kernel_launch(rope_multi<forward, true, T, D>, launch_params,
            x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor,
-            attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, inplace);
+            attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, row_indices, set_rows_stride, inplace);
    }
 }

@@ -648,14 +660,21 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
            GGML_ABORT("fatal error");
        }
    } else if (is_mrope && !is_vision) {
-        if (src0->type == GGML_TYPE_F32) {
-            rope_multi_cuda<forward>((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1,
-                                     s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor,
-                                     corr_dims, freq_factors, sections, is_imrope, inplace, stream);
-        } else if (src0->type == GGML_TYPE_F16) {
-            rope_multi_cuda<forward>((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1,
-                                     s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor,
-                                     corr_dims, freq_factors, sections, is_imrope, inplace, stream);
+        if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) {
+            rope_multi_cuda<forward, float, float>((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02,
+                                                   s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
+                                                   ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope,
+                                                   row_indices, set_rows_stride, inplace, stream);
+        } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) {
+            rope_multi_cuda<forward, float, half>((const float *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02,
+                                                  s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
+                                                  ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope,
+                                                  row_indices, set_rows_stride, inplace, stream);
+        } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) {
+            rope_multi_cuda<forward, half, half>((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02,
+                                                 s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base,
+                                                 ext_factor, attn_factor, corr_dims, freq_factors, sections, is_imrope,
+                                                 row_indices, set_rows_stride, inplace, stream);
        } else {
            GGML_ABORT("fatal error");
        }
@@ -76,7 +76,7 @@ __device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], co
    for (int i = 0; i < experts_per_thread; i++) {
        const int  idx    = lane + i * WARP_SIZE;
        const bool active = !use_limit || (idx < limit);
-        vals[i]           = active ? sqrtf(vals[i] > 20.0f ? vals[i] : logf(1.0f + expf(vals[i]))) : -INFINITY;
+        vals[i]           = active ? sqrtf(vals[i] > 20.0f ? vals[i] : log1pf(expf(vals[i]))) : -INFINITY;
    }
 }

@@ -91,7 +91,7 @@ static __device__ __forceinline__ float op_expm1(float x) {
 }

 static __device__ __forceinline__ float op_softplus(float x) {
-    return (x > 20.0f) ? x : logf(1.0f + expf(x));
+    return (x > 20.0f) ? x : log1pf(expf(x));
 }

 static __device__ __forceinline__ float op_elu(float x) {
@@ -121,3 +121,18 @@ __device__ __forceinline__ float ggml_cuda_op_swiglu_clamp_single(float gate, fl

    return ggml_cuda_op_silu_single(gate) * up;
 }
+
+__device__ __forceinline__ float ggml_cuda_op_glu_single(ggml_glu_op op, float gate, float up, float limit) {
+    switch (op) {
+        case GGML_GLU_OP_SWIGLU:
+            return up * ggml_cuda_op_silu_single(gate);
+        case GGML_GLU_OP_GEGLU:
+            return up * ggml_cuda_op_gelu_single(gate);
+        case GGML_GLU_OP_SWIGLU_OAI:
+            return ggml_cuda_op_swiglu_oai_single(gate, up);
+        case GGML_GLU_OP_SWIGLU_CLAMP:
+            return ggml_cuda_op_swiglu_clamp_single(gate, up, limit);
+        default:
+            return up * gate;
+    }
+}
@@ -40,18 +40,18 @@ static __device__ __forceinline__ int2 get_int_from_table_16(const int & q4, con
    const uint32_t q_odd  = (q4 >> 4);

    // Perform lookups in the lower half of the table (indices 0-7).
-    uint32_t v_even_low = __builtin_amdgcn_perm(values[1], values[0], q_even & 0x07070707);
-    uint32_t v_odd_low = __builtin_amdgcn_perm(values[1], values[0], q_odd & 0x07070707);
+    uint32_t v_even_low = ggml_hip_perm(values[1], values[0], q_even & 0x07070707);
+    uint32_t v_odd_low  = ggml_hip_perm(values[1], values[0], q_odd & 0x07070707);

    // Perform lookups in the upper half of the table (indices 8-15).
-    uint32_t v_even_high = __builtin_amdgcn_perm(values[3], values[2], q_even & 0x07070707);
-    uint32_t v_odd_high = __builtin_amdgcn_perm(values[3], values[2], q_odd & 0x07070707);
+    uint32_t v_even_high = ggml_hip_perm(values[3], values[2], q_even & 0x07070707);
+    uint32_t v_odd_high  = ggml_hip_perm(values[3], values[2], q_odd & 0x07070707);

    // Select between the low and high results based on the MSB of each index nibble.
    uint32_t mask_even = 0x03020100 | ((q_even & 0x08080808) >> 1);
-    uint32_t res_x = __builtin_amdgcn_perm(v_even_high, v_even_low, mask_even);
+    uint32_t res_x     = ggml_hip_perm(v_even_high, v_even_low, mask_even);
    uint32_t mask_odd = 0x03020100 | ((q_odd & 0x08080808) >> 1);
-    uint32_t res_y = __builtin_amdgcn_perm(v_odd_high, v_odd_low, mask_odd);
+    uint32_t res_y     = ggml_hip_perm(v_odd_high, v_odd_low, mask_odd);

    return make_int2(res_x, res_y);
 #elif !defined(GGML_USE_MUSA)
@@ -720,10 +720,10 @@ static __device__ __forceinline__ float vec_dot_q1_0_q8_1(
        const uint32_t q2_indices = (q2_bits & 0x01) | ((q2_bits & 0x02) << 7) | ((q2_bits & 0x04) << 14) | ((q2_bits & 0x08) << 21);
        const uint32_t q3_indices = (q3_bits & 0x01) | ((q3_bits & 0x02) << 7) | ((q3_bits & 0x04) << 14) | ((q3_bits & 0x08) << 21);

-        const int v0 = __builtin_amdgcn_perm(0x000001FF, 0x000001FF, q0_indices);
-        const int v1 = __builtin_amdgcn_perm(0x000001FF, 0x000001FF, q1_indices);
-        const int v2 = __builtin_amdgcn_perm(0x000001FF, 0x000001FF, q2_indices);
-        const int v3 = __builtin_amdgcn_perm(0x000001FF, 0x000001FF, q3_indices);
+        const int v0 = ggml_hip_perm(0x000001FF, 0x000001FF, q0_indices);
+        const int v1 = ggml_hip_perm(0x000001FF, 0x000001FF, q1_indices);
+        const int v2 = ggml_hip_perm(0x000001FF, 0x000001FF, q2_indices);
+        const int v3 = ggml_hip_perm(0x000001FF, 0x000001FF, q3_indices);
 #else
        // unpack crumbs into nibble indices
        const int n0 = __byte_perm(0x11100100, 0x11100100, q >> 0); // [0, 1, 4, 5] [ 8,  9, 12, 13]
@@ -777,8 +777,8 @@ static __device__ __forceinline__ float vec_dot_q2_0_q8_1(
        const uint32_t qx_indices = (q & 0x03) | ((q & 0x0C) << 6) | ((q & 0x30) << 12) | ((q & 0xC0) << 18);
        const uint32_t qy_bits    = q >> 8;
        const uint32_t qy_indices = (qy_bits & 0x03) | ((qy_bits & 0x0C) << 6) | ((qy_bits & 0x30) << 12) | ((qy_bits & 0xC0) << 18);
-        const int qx = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qx_indices);
-        const int qy = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qy_indices);
+        const int      qx         = ggml_hip_perm(0x020100FF, 0x020100FF, qx_indices);
+        const int      qy         = ggml_hip_perm(0x020100FF, 0x020100FF, qy_indices);
 #else
        // unpack even and odd crumbs into byte values
        const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0);
@@ -311,3 +311,22 @@ static __device__ __forceinline__ unsigned int __vcmpne4(unsigned int a, unsigne

    return (ne_any_bit >> 7) * 0xff;
 }
+
+// v_perm_b32 with selector bytes in [0, 7]: 0-3 pick bytes of lo, 4-7 pick bytes of hi.
+static __device__ __forceinline__ uint32_t ggml_hip_perm(const uint32_t hi, const uint32_t lo, const uint32_t sel) {
+#if __has_builtin(__builtin_amdgcn_is_invocable)
+    if (__builtin_amdgcn_is_invocable(__builtin_amdgcn_perm)) {
+        return __builtin_amdgcn_perm(hi, lo, sel);
+    }
+    const uint64_t src = ((uint64_t) hi << 32) | lo;
+    uint32_t res = 0;
+#pragma unroll
+    for (int i = 0; i < 4; ++i) {
+        const uint32_t byte_sel = (sel >> (8 * i)) & 0x07;
+        res |= (uint32_t) ((src >> (8 * byte_sel)) & 0xFF) << (8 * i);
+    }
+    return res;
+#else
+    return __builtin_amdgcn_perm(hi, lo, sel);
+#endif // __has_builtin(__builtin_amdgcn_is_invocable)
+}
@@ -22,6 +22,10 @@
 #define HWCAP2_I8MM (1ULL << 13)
 #endif

+#if !defined(HWCAP2_SVEI8MM)
+#define HWCAP2_SVEI8MM (1ULL << 9)
+#endif
+
 #if !defined(HWCAP_ASIMDDP)
 #define HWCAP_ASIMDDP (1 << 20)
 #endif
@@ -63,6 +67,10 @@
 #define PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE 47
 #endif

+#if !defined(PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE)
+#define PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE 57
+#endif
+
 #if !defined(PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE)
 #define PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE 66
 #endif
@@ -87,6 +95,7 @@ typedef struct ggml_feats_arch64_runtime {
    bool has_sve;
    bool has_sve2;
    bool has_i8mm;
+    bool has_svei8mm;
    bool has_sme;
    bool has_sme2;
    int sve_cnt;
@@ -104,6 +113,7 @@ static inline ggml_feats_arch64_runtime_t ggml_feats_get_arch64_runtime(void) {
    runtime_feat.has_sve     = !!(hwcap & HWCAP_SVE);
    runtime_feat.has_sve2    = !!(hwcap2 & HWCAP2_SVE2);
    runtime_feat.has_i8mm    = !!(hwcap2 & HWCAP2_I8MM);
+    runtime_feat.has_svei8mm = !!(hwcap2 & HWCAP2_SVEI8MM);
    runtime_feat.has_sme     = !!(hwcap2 & HWCAP2_SME);
    runtime_feat.has_sme2    = !!(hwcap2 & HWCAP2_SME2);

@@ -153,6 +163,7 @@ static inline ggml_feats_arch64_runtime_t ggml_feats_get_arch64_runtime(void) {
    runtime_feat.has_sve     = IsProcessorFeaturePresent(PF_ARM_SVE_INSTRUCTIONS_AVAILABLE) != 0;
    runtime_feat.has_sve2    = IsProcessorFeaturePresent(PF_ARM_SVE2_INSTRUCTIONS_AVAILABLE) != 0;
    runtime_feat.has_i8mm    = IsProcessorFeaturePresent(PF_ARM_V82_I8MM_INSTRUCTIONS_AVAILABLE) != 0;
+    runtime_feat.has_svei8mm = IsProcessorFeaturePresent(PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE) != 0;
    runtime_feat.has_sme     = IsProcessorFeaturePresent(PF_ARM_SME_INSTRUCTIONS_AVAILABLE) != 0;
    runtime_feat.has_sme2    = IsProcessorFeaturePresent(PF_ARM_SME2_INSTRUCTIONS_AVAILABLE) != 0;

@@ -51,6 +51,13 @@ endif()
 set(TARGET_NAME ggml-hexagon)
 ggml_add_backend_library(${TARGET_NAME}
    ggml-hexagon.cpp
+    htp-common.h
+    htp-kparams.cpp
+    htp-kparams.h
+    htp-opnode.cpp
+    htp-opnode.h
+    htp-fusion.cpp
+    htp-fusion.h
    htp-repack.cpp
    htp-repack.h
    htp-drv.cpp
@@ -0,0 +1,106 @@
+#ifndef HTP_COMMON_H
+#define HTP_COMMON_H
+
+#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;
+extern int    opt_nhmx;
+extern size_t opt_vmem;
+extern int    opt_etm;
+extern int    opt_verbose;
+extern int    opt_profile;
+extern bool   opt_hostbuf;
+extern bool   opt_dma64;
+extern int    opt_power_mode;
+
+extern size_t opt_mbuf_dyn;
+extern size_t opt_mbuf_static;
+extern size_t opt_mbuf_total;
+
+extern int    opt_mm_select;
+extern int    opt_fa_select;
+extern int    opt_fa_split;
+extern int    opt_gdn_select;
+extern int    opt_ar_select;
+extern int    opt_ar_scatter;
+
+extern int    opt_opbatch;
+extern int    opt_opqueue;
+extern int    opt_optrace;
+extern int    opt_oppoll;
+extern int    opt_opfusion;
+extern int    opt_rt_repack;
+
+#ifndef HEX_VERBOSE
+#define HEX_VERBOSE(...) \
+    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 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 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 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)) {
+        return false;
+    }
+    return htp_can_rt_repack(t);
+}
+
+static inline bool htp_tensors_overlap(const struct ggml_tensor * a, const struct ggml_tensor * b) {
+    const uintptr_t a0 = (uintptr_t) a->data;
+    const uintptr_t b0 = (uintptr_t) b->data;
+    const uintptr_t a1 = a0 + ggml_nbytes(a);
+    const uintptr_t b1 = b0 + ggml_nbytes(b);
+
+    return a0 < b1 && b0 < a1;
+}
+
+static inline bool htp_can_row_partition(const struct ggml_tensor * t) {
+    if (t->ne[1] > 1 && (t->nb[1] & 127) != 0) return false;
+    if (t->ne[2] > 1 && (t->nb[2] & 127) != 0) return false;
+    if (t->ne[3] > 1 && (t->nb[3] & 127) != 0) return false;
+    return true;
+}
+
+struct ggml_hexagon_session;
+
+const char * htp_session_name(const struct ggml_hexagon_session * sess);
+size_t       htp_session_vtcm_size(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_n_threads(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_n_hmx(const struct ggml_hexagon_session * sess);
+uint32_t     htp_session_mdev_count(const struct ggml_hexagon_session * sess);
+
+htp_op_code  htp_op_remap(const struct ggml_tensor * t);
+
+#endif // HTP_COMMON_H
@@ -0,0 +1,649 @@
+#include "htp-fusion.h"
+#include "htp-kparams.h"
+#include "ggml-impl.h"
+
+#include <algorithm>
+#include <cstring>
+
+static inline bool htp_same_tensor_data(const ggml_tensor * a, const ggml_tensor * b) {
+    return a && b && (a == b || a->data == b->data);
+}
+
+static inline bool htp_same_shape_and_contiguity(const ggml_tensor * a, const ggml_tensor * b) {
+    return ggml_are_same_shape(a, b) && (ggml_is_contiguous(a) == ggml_is_contiguous(b));
+}
+
+static inline bool htp_inplace_overlap_unsafe(const ggml_tensor * dst, const ggml_tensor * in) {
+    return htp_tensors_overlap(dst, in) && (dst->data != in->data || !ggml_are_same_stride(dst, in));
+}
+
+static inline bool htp_inputs_overlap_unsafe(const htp_opnode & prev, const ggml_tensor * dst, const ggml_tensor * extra = nullptr) {
+    if (extra && htp_inplace_overlap_unsafe(dst, extra)) return true;
+    for (const auto * in : prev.inputs) {
+        if (htp_inplace_overlap_unsafe(dst, in)) return true;
+    }
+    return false;
+}
+
+static inline bool htp_match_consumer(
+    const htp_opnode & prev,
+    const htp_opnode & cur,
+    bool commutative,
+    const ggml_tensor ** other_src = nullptr
+) {
+    if (prev.outputs.size() != 1) {
+        return false;
+    }
+    const ggml_tensor * prev_out = prev.dst();
+    if (!htp_tensor_is_fuseable(prev_out)) {
+        return false;
+    }
+    if (htp_same_tensor_data(cur.src0(), prev_out)) {
+        if (other_src) *other_src = cur.src1();
+        return true;
+    }
+    if (commutative && htp_same_tensor_data(cur.src1(), prev_out)) {
+        if (other_src) *other_src = cur.src0();
+        return true;
+    }
+    return false;
+}
+
+static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams) {
+    if (kparams->n_hmx) {
+        return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D;
+    }
+    return htp_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) &&
+           (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK);
+}
+
+// --- Fusion rules ---
+
+static bool fuse_allreduce_add(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (opt_ar_select != 2 || prev.opcode != HTP_OP_ALLREDUCE) return false;
+
+    const auto * ar_kparams = prev.kparams<htp_allreduce_kernel_params>();
+    const uint32_t rank    = (uint32_t) ar_kparams->rank;
+    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;
+
+    const ggml_tensor * add_dst    = cur.dst();
+    const ggml_tensor * res_tensor = nullptr;
+    if (htp_same_tensor_data(cur.src0(), ar_local)) {
+        res_tensor = cur.src1();
+    } else if (htp_same_tensor_data(cur.src1(), ar_local)) {
+        res_tensor = cur.src0();
+    } else {
+        return false;
+    }
+
+    if (ar_local->type != res_tensor->type) return false;
+
+    const bool is_same_shape = ggml_are_same_shape(ar_local, res_tensor);
+    const bool is_row_bcast  = !is_same_shape && (ar_local->ne[0] == res_tensor->ne[0] &&
+                                                  res_tensor->ne[1] == 1 && res_tensor->ne[2] == 1 && res_tensor->ne[3] == 1);
+    if (!is_same_shape && !is_row_bcast) return false;
+
+    if (is_same_shape && !ggml_are_same_stride(ar_local, res_tensor)) {
+        return false;
+    }
+    if (ggml_is_contiguous(ar_local) != ggml_is_contiguous(add_dst)) return false;
+
+    const bool is_shard_ok = n_ranks <= HTP_OP_MAX_OUTPUTS && add_dst->data == ar_local->data;
+
+    struct htp_allreduce_kernel_params kparams;
+    if (!htp_kparams_solve_allreduce(sess, add_dst, rank, n_ranks, true, is_row_bcast, is_shard_ok, &kparams)) {
+        HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: solver failed\n", htp_session_name(sess));
+        return false;
+    }
+
+    const bool scatter_ok = kparams.n_dsts > 1;
+    kparams.mode = scatter_ok ? HTP_ALLREDUCE_SHARDED_FANOUT : HTP_ALLREDUCE_FULL;
+
+    for (uint32_t r = 0; r < n_ranks; r++) {
+        const ggml_tensor * ar_src = prev.inputs[r];
+        if (!htp_tensors_overlap(add_dst, ar_src)) continue;
+        if (scatter_ok && r == rank) continue;
+        HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: dst overlaps allreduce src %u\n", htp_session_name(sess), r);
+        return false;
+    }
+
+    prev.opcode = HTP_OP_ALLREDUCE_ADD;
+    prev.inputs.push_back(res_tensor);
+    if (kparams.mode != HTP_ALLREDUCE_SHARDED_FANOUT) {
+        prev.outputs = { add_dst };
+    }
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_rms_norm_mul(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_RMS_NORM) return false;
+
+    const ggml_tensor * weight = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &weight)) return false;
+
+    const ggml_tensor * src0 = prev.src0();
+    const ggml_tensor * dst  = cur.dst();
+    if (src0->type != GGML_TYPE_F32 || src0->ne[0] != weight->ne[0] || !htp_same_shape_and_contiguity(prev.dst(), dst)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst, weight)) return false;
+
+    const bool is_row_bcast  = (weight->ne[1] == 1 && weight->ne[2] == 1 && weight->ne[3] == 1);
+    const bool is_same_shape = ggml_are_same_shape(src0, weight);
+    if (!is_row_bcast && !is_same_shape) return false;
+
+    struct htp_unary_kernel_params kparams;
+    if (!htp_kparams_solve_rms_norm_mul(sess, src0, weight, dst, &kparams)) {
+        HEX_VERBOSE("ggml-hex: %s skip RMS_NORM_MUL fusion: VTCM needed (%d) > budget (%zu)\n",
+                    htp_session_name(sess), kparams.vtcm_size, htp_session_vtcm_size(sess));
+        return false;
+    }
+
+    prev.opcode  = HTP_OP_RMS_NORM_MUL;
+    prev.inputs  = { src0, weight };
+    prev.outputs = { 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;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false)) return false;
+
+    const ggml_tensor * src0 = prev.src0();
+    const ggml_tensor * dst  = cur.dst();
+    if (src0->type != GGML_TYPE_F32 || !htp_same_shape_and_contiguity(prev.dst(), dst)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst)) return false;
+
+    float scale = 0.0f;
+    float bias  = 0.0f;
+    memcpy(&scale, (const float *) cur.params + 0, sizeof(float));
+    memcpy(&bias,  (const float *) cur.params + 1, sizeof(float));
+    if (bias != 0.0f || scale == 0.0f) return false;
+
+    float prev_scale = 0.0f;
+    memcpy(&prev_scale, &prev.params[1], sizeof(float));
+    if (prev_scale != 0.0f) {
+        scale *= prev_scale;
+    }
+    memcpy(&prev.params[1], &scale, sizeof(float));
+
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_mul_mat_add(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL_MAT) return false;
+
+    const ggml_tensor * src2 = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &src2)) return false;
+    if (src2->type != GGML_TYPE_F32) return false;
+
+    const ggml_tensor * src0 = prev.node->src[0];
+    const ggml_tensor * src1 = prev.node->src[1];
+    const ggml_tensor * dst  = cur.node;
+    if (!ggml_are_same_shape(prev.node, dst)) return false;
+
+    const bool prev_collapsed = htp_matmul_can_collapse(src0, src1, prev.node);
+    const bool cur_collapsed  = htp_matmul_add_can_collapse(src0, src1, src2, dst);
+    if (prev_collapsed != cur_collapsed) {
+        return false;
+    }
+
+    const auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
+    struct htp_mm_kernel_params kparams;
+    if (!htp_kparams_solve_matmul_add(sess, src0, src1, src2, dst, &kparams)) {
+        return false;
+    }
+
+    const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3];
+    if (kparams.n_hmx == 0 && src1_nrows != 1) return false;
+
+    if (kparams.n_hmx > 0 && orig_kparams->n_hmx > 0) {
+        if (kparams.m_chunk < orig_kparams->m_chunk ||
+            kparams.n_chunk < orig_kparams->n_chunk ||
+            kparams.n_act_threads < orig_kparams->n_act_threads) {
+            HEX_VERBOSE("ggml-hex: %s skip MUL_MAT_ADD fusion: HMX efficiency reduced (m %d->%d, n %d->%d, th %d->%d)\n",
+                        htp_session_name(sess), orig_kparams->m_chunk, kparams.m_chunk,
+                        orig_kparams->n_chunk, kparams.n_chunk,
+                        orig_kparams->n_act_threads, kparams.n_act_threads);
+            return false;
+        }
+    }
+
+    prev.opcode  = HTP_OP_MUL_MAT_ADD;
+    prev.inputs  = { src0, src1, src2 };
+    prev.outputs = { dst };
+    htp_opnode_collapse(sess, prev);
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_mul_mat_nx(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL_MAT && prev.opcode != HTP_OP_MUL_MAT_NX) return false;
+
+    const uint32_t curr_n = (uint32_t) prev.outputs.size();
+    if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 1 >= HTP_OP_MAX_INPUTS) return false;
+
+    const ggml_tensor * w0     = prev.inputs[0];
+    const ggml_tensor * x_orig = prev.node->src[1];
+    const ggml_tensor * w_in   = cur.node->src[0];
+    const ggml_tensor * x_in   = cur.node->src[1];
+    const ggml_tensor * d_in   = cur.node;
+
+    auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
+    const auto * cur_kparams = cur.kparams<htp_mm_kernel_params>();
+
+    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) ||
+        orig_kparams->n_hmx != cur_kparams->n_hmx) {
+        return false;
+    }
+
+    const bool collapse = htp_matmul_can_collapse(w0, x_orig, prev.node) && htp_matmul_can_collapse(w_in, x_in, d_in);
+    if ((x_orig->ne[2] > 1 || x_orig->ne[3] > 1) && !collapse) return false;
+
+    if (prev.opcode == HTP_OP_MUL_MAT_NX) {
+        orig_kparams->n_weights = curr_n + 1;
+    } else {
+        struct htp_mm_kernel_params kparams;
+        const ggml_tensor x_collapsed = collapse ? htp_tensor_collapse_rows(x_orig) : *x_orig;
+        if (!htp_kparams_solve_mmnx(sess, w0, &x_collapsed, curr_n + 1, &kparams)) return false;
+        if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) return false;
+
+        prev.set_kparams(kparams);
+        prev.opcode = HTP_OP_MUL_MAT_NX;
+    }
+
+    prev.inputs.insert(prev.inputs.begin() + curr_n, w_in);
+    prev.outputs.push_back(collapse ? prev.add_dummy(htp_tensor_collapse_rows(d_in)) : d_in);
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_mul_mat_id_nx(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL_MAT_ID && prev.opcode != HTP_OP_MUL_MAT_ID_NX) return false;
+
+    const uint32_t curr_n = (uint32_t) prev.outputs.size();
+    if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 2 >= HTP_OP_MAX_INPUTS) return false;
+
+    const ggml_tensor * w0     = prev.inputs[0];
+    const ggml_tensor * x      = prev.inputs[curr_n];
+    const ggml_tensor * ids    = prev.inputs[curr_n + 1];
+    const ggml_tensor * w_in   = cur.src0();
+    const ggml_tensor * x_in   = cur.src1();
+    const ggml_tensor * ids_in = cur.node->src[2];
+    const ggml_tensor * d_in   = cur.dst();
+
+    auto * orig_kparams = prev.kparams<htp_mm_kernel_params>();
+    const auto * cur_kparams = cur.kparams<htp_mm_kernel_params>();
+
+    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) ||
+        orig_kparams->n_hmx != cur_kparams->n_hmx) {
+        return false;
+    }
+
+    struct htp_mm_kernel_params kparams;
+    if (d_in->nb[1] == prev.node->nb[1]) {
+        kparams = *orig_kparams;
+        kparams.n_weights = curr_n + 1;
+    } else if (!htp_kparams_solve_mmidnx(sess, w0, x, d_in, curr_n + 1, &kparams)) {
+        return false;
+    }
+    if (!is_supported_mul_mat_id_nx_kernel(w0, &kparams)) return false;
+
+    prev.opcode = HTP_OP_MUL_MAT_ID_NX;
+    prev.inputs.insert(prev.inputs.begin() + curr_n, w_in);
+    prev.outputs.push_back(d_in);
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_gdn_cpy(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    GGML_UNUSED(sess);
+    if (prev.opcode != HTP_OP_GATED_DELTA_NET || prev.outputs.size() != 1) return false;
+
+    const ggml_tensor * gdn_out  = prev.dst();
+    const ggml_tensor * cpy_node = cur.node;
+    const ggml_tensor * cpy_src  = cur.src0();
+    const ggml_tensor * cpy_dst  = cur.dst();
+
+    if (gdn_out->type != GGML_TYPE_F32 || cpy_src->type != GGML_TYPE_F32 || cpy_dst->type != GGML_TYPE_F32) return false;
+    if ((gdn_out->flags & GGML_TENSOR_FLAG_OUTPUT) || (cpy_node->flags & GGML_TENSOR_FLAG_OUTPUT)) return false;
+
+    const ggml_tensor * v = prev.node->src[2];
+    if (!v) return false;
+
+    const int64_t S_v       = v->ne[0];
+    const int64_t H         = v->ne[1];
+    const int64_t n_tokens  = v->ne[2];
+    const int64_t n_seqs    = v->ne[3];
+    const int64_t K         = ggml_get_op_params_i32(prev.node, 0);
+    const size_t  tail_off  = (size_t) S_v * H * n_tokens * n_seqs * sizeof(float);
+    const int64_t D         = S_v * S_v * H;
+    const int64_t n_written = std::min<int64_t>(n_tokens, K);
+
+    if (cpy_src->op != GGML_OP_VIEW || !htp_same_tensor_data(cpy_src->view_src, gdn_out) ||
+        cpy_src->view_offs != tail_off || !ggml_is_contiguous(cpy_src)) {
+        return false;
+    }
+
+    if (cpy_dst->ne[0] != D || cpy_dst->ne[1] != n_seqs || cpy_dst->ne[2] != n_written || !ggml_is_contiguous(cpy_dst)) {
+        return false;
+    }
+
+    prev.outputs.push_back(cpy_dst);
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_ssm_conv_silu(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    GGML_UNUSED(sess);
+    if (prev.opcode != HTP_OP_SSM_CONV) return false;
+
+    auto * kparams = prev.kparams<htp_ssm_conv_kernel_params>();
+    if (kparams->silu) return false;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false)) return false;
+
+    const ggml_tensor * conv_dst = prev.dst();
+    const ggml_tensor * dst      = cur.dst();
+    if (!ggml_are_same_layout(conv_dst, dst)) return false;
+    if (htp_tensors_overlap(dst, prev.src0()) || htp_tensors_overlap(dst, prev.src1())) return false;
+
+    kparams->silu   = 1;
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    return true;
+}
+
+static bool fuse_repeat_get_rows(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_REPEAT) return false;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false)) return false;
+
+    const ggml_tensor * repeat_out  = prev.dst();
+    const ggml_tensor * repeat_src0 = prev.src0();
+    const ggml_tensor * gr_src1     = cur.src1();
+    const ggml_tensor * gr_dst      = cur.dst();
+
+    if (repeat_src0->ne[0] != repeat_out->ne[0] || repeat_src0->ne[1] != repeat_out->ne[1]) return false;
+    if (repeat_src0->type != repeat_out->type) return false;
+
+    struct htp_get_rows_kernel_params kparams;
+    if (!htp_kparams_solve_repeat_get_rows(sess, repeat_src0, gr_src1, gr_dst, &kparams) || kparams.vtcm_size == 0) return false;
+
+    prev.opcode  = HTP_OP_GET_ROWS;
+    prev.inputs  = { repeat_src0, gr_src1 };
+    prev.outputs = { gr_dst };
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_unary_mul_glu(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    htp_op_code glu_op = HTP_OP_INVALID;
+    switch (prev.opcode) {
+        case HTP_OP_UNARY_SILU:       glu_op = HTP_OP_GLU_SWIGLU;      break;
+        case HTP_OP_UNARY_GELU:       glu_op = HTP_OP_GLU_GEGLU;       break;
+        case HTP_OP_UNARY_GELU_QUICK: glu_op = HTP_OP_GLU_GEGLU_QUICK; break;
+        case HTP_OP_UNARY_GELU_ERF:   glu_op = HTP_OP_GLU_GEGLU_ERF;   break;
+        case HTP_OP_UNARY_RELU:       glu_op = HTP_OP_GLU_REGLU;       break;
+        case HTP_OP_UNARY_SIGMOID:    glu_op = HTP_OP_GLU_SIGMOID;     break;
+        default: return false;
+    }
+
+    const ggml_tensor * up = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &up)) return false;
+
+    const ggml_tensor * gate = prev.src0();
+    const ggml_tensor * dst  = cur.dst();
+    if (htp_inputs_overlap_unsafe(prev, dst, up)) return false;
+
+    prev.opcode  = glu_op;
+    prev.inputs  = { gate, up };
+    prev.outputs = { dst };
+    if (!htp_kparams_solve(sess, prev)) return false;
+
+    prev.fused.push_back(cur.node);
+    memset(prev.params, 0, sizeof(prev.params));
+    return true;
+}
+
+static bool fuse_mul_mul_scalar(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != HTP_OP_MUL || prev.inputs.size() != 2) return false;
+
+    const auto * k0 = prev.kparams<htp_binary_kernel_params>();
+    const auto * k1 = cur.kparams<htp_binary_kernel_params>();
+    if (k0->kernel_type != HTP_BINARY_KERNEL_SCALAR || k1->kernel_type != HTP_BINARY_KERNEL_SCALAR) return false;
+    if (k0->src1_strided || k1->src1_strided) return false;
+
+    const ggml_tensor * src2 = nullptr;
+    if (!htp_match_consumer(prev, cur, /*commutative=*/ false, &src2)) return false;
+
+    const ggml_tensor * src1 = prev.src1();
+    const ggml_tensor * dst  = cur.dst();
+    if (!ggml_are_same_shape(src1, src2) || !htp_same_shape_and_contiguity(prev.dst(), dst)) return false;
+    if (htp_inputs_overlap_unsafe(prev, dst, src2)) return false;
+
+    const size_t vtcm_size = htp_session_vtcm_size(sess);
+    struct htp_binary_kernel_params kparams = *k0;
+    kparams.src2_size       = kparams.src1_size;
+    kparams.rows_per_buffer = 0;
+    if (kparams.src1_size > 65536 || kparams.src1_size + kparams.src2_size > vtcm_size) return false;
+
+    struct htp_binary_vtcm_layout L;
+    htp_binary_vtcm_layout_build(&L, &kparams, vtcm_size);
+    if (L.rows_per_buffer == 0 || L.total_bytes > vtcm_size) return false;
+
+    kparams.rows_per_buffer = L.rows_per_buffer;
+    kparams.vtcm_size       = L.total_bytes;
+
+    prev.inputs.push_back(src2);
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+static bool fuse_binary_nx(const ggml_hexagon_session * sess, htp_opnode & prev, const htp_opnode & cur) {
+    if (prev.opcode != cur.opcode) return false;
+
+    const uint32_t curr_n = (uint32_t) prev.inputs.size();
+    if (curr_n < 2 || curr_n >= HTP_OP_MAX_INPUTS) return false;
+
+    const auto * k0 = prev.kparams<htp_binary_kernel_params>();
+    const auto * k1 = cur.kparams<htp_binary_kernel_params>();
+    const bool k0_ok = (k0->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) ||
+                       (k0->kernel_type == HTP_BINARY_KERNEL_CHUNKED && !k0->is_scalar);
+    const bool k1_ok = (k1->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) ||
+                       (k1->kernel_type == HTP_BINARY_KERNEL_CHUNKED && !k1->is_scalar);
+    if (!k0_ok || !k1_ok) return false;
+
+    const bool is_comm = (cur.opcode == HTP_OP_ADD || cur.opcode == HTP_OP_MUL);
+    const ggml_tensor * new_src = nullptr;
+    if (!htp_match_consumer(prev, cur, is_comm, &new_src)) return false;
+
+    const ggml_tensor * src0 = prev.inputs[0];
+    const ggml_tensor * dst  = cur.dst();
+    if (!ggml_are_same_shape(src0, prev.inputs[1]) ||
+        !ggml_are_same_shape(src0, new_src) ||
+        !htp_same_shape_and_contiguity(prev.dst(), dst)) {
+        return false;
+    }
+    if (htp_inputs_overlap_unsafe(prev, dst, new_src)) return false;
+
+    const size_t vtcm_size = htp_session_vtcm_size(sess);
+    struct htp_binary_kernel_params kparams = *k0;
+    struct htp_binary_vtcm_layout L;
+    memset(&L, 0, sizeof(L));
+
+    const uint32_t n_srcs = curr_n + 1;
+    if (k0->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE || k1->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) {
+        if (k0->kernel_type != HTP_BINARY_KERNEL_SAME_SHAPE) {
+            kparams = *k1;
+        }
+        kparams.n_srcs          = (uint8_t) n_srcs;
+        kparams.rows_per_buffer = 0;
+
+        auto check_split = [&](const ggml_tensor * t) {
+            if (src0->ne[2] > 1 && t->nb[2] != t->ne[1] * t->nb[1]) kparams.split_at_ne01 = 1;
+            if (src0->ne[3] > 1 && t->nb[3] != t->ne[2] * t->nb[2]) kparams.split_at_ne02 = 1;
+        };
+        for (const auto * in : prev.inputs) check_split(in);
+        check_split(new_src);
+        check_split(dst);
+
+        htp_binary_vtcm_layout_build(&L, &kparams, vtcm_size);
+    }
+
+    if (L.rows_per_buffer == 0 || L.total_bytes == 0 || L.total_bytes > vtcm_size) {
+        bool all_contig = ggml_is_contiguous(new_src) && ggml_is_contiguous(dst);
+        for (const auto * in : prev.inputs) {
+            all_contig = all_contig && ggml_is_contiguous(in);
+        }
+        if (!all_contig) return false;
+
+        const uint32_t total_elems = (uint32_t) ggml_nelements(src0);
+        const uint32_t elem_size   = (uint32_t) ggml_type_size(src0->type);
+        const uint32_t max_threads = htp_session_n_threads(sess);
+        if (!htp_binary_solve_chunked(&L, &kparams, total_elems, elem_size, max_threads, n_srcs, false, vtcm_size)) {
+            return false;
+        }
+    }
+
+    kparams.rows_per_buffer = L.rows_per_buffer;
+    kparams.vtcm_size       = (uint32_t) L.total_bytes;
+
+    prev.inputs.push_back(new_src);
+    prev.outputs[0] = dst;
+    prev.fused.push_back(cur.node);
+    prev.set_kparams(kparams);
+    return true;
+}
+
+bool htp_fusion_can_stack(const struct ggml_tensor * a) {
+    if (!a || !a->src[0] || !a->src[1]) return false;
+    switch (a->op) {
+        case GGML_OP_MUL_MAT:
+            if (!htp_fusion_enabled(HTP_FUSE_MUL_MAT_NX)) return false;
+            if (htp_is_hadamard(a)) return false;
+            if (a->src[0]->ne[2] != 1 || a->src[0]->ne[3] != 1) return false;
+            return ggml_is_quantized(a->src[0]->type);
+        case GGML_OP_MUL_MAT_ID:
+            if (!htp_fusion_enabled(HTP_FUSE_MUL_MAT_ID_NX)) return false;
+            return a->src[2] != nullptr && ggml_is_quantized(a->src[0]->type);
+        default:
+            return false;
+    }
+}
+
+bool htp_fusion_can_stack(const struct ggml_tensor * a, const struct ggml_tensor * b) {
+    if (!htp_fusion_can_stack(b) || a->op != b->op || a->src[1] != b->src[1]) {
+        return false;
+    }
+    if (a->op == GGML_OP_MUL_MAT_ID && a->src[2] != b->src[2]) {
+        return false;
+    }
+    const ggml_tensor * w0 = a->src[0];
+    const ggml_tensor * w1 = b->src[0];
+    return w0->type == w1->type && w0->ne[0] == w1->ne[0] && w0->ne[2] == w1->ne[2];
+}
+
+int htp_fusion_max(const struct ggml_cgraph * graph, int idx) {
+    const int n = graph->n_nodes;
+    const ggml_op op0 = graph->nodes[idx]->op;
+    if (op0 != GGML_OP_ADD && op0 != GGML_OP_NORM && op0 != GGML_OP_RMS_NORM) {
+        return 1;
+    }
+
+    enum ggml_op ops[HTP_OP_MAX_FUSED];
+    ops[0] = op0;
+
+    int f = idx + 1;
+    while (f < n && f < idx + HTP_OP_MAX_FUSED) {
+        const ggml_op op = graph->nodes[f]->op;
+        if (op != GGML_OP_ADD && op != GGML_OP_MUL && op != GGML_OP_SCALE &&
+            op != GGML_OP_NORM && op != GGML_OP_RMS_NORM) {
+            break;
+        }
+        ops[f - idx] = op;
+        f++;
+    }
+
+    f -= idx;
+    for (; f > 1; f--) {
+        if (ggml_can_fuse(graph, idx, ops, f)) {
+            break;
+        }
+    }
+    return f;
+}
+
+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;
+
+        extra->flags &= ~HTP_TENSOR_FUSEABLE;
+
+        const ggml_op op = graph->nodes[i]->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]))) &&
+            ggml_node_has_n_uses(graph, i, 1)) {
+            extra->flags |= HTP_TENSOR_FUSEABLE;
+        }
+    }
+}
+
+bool htp_fusion_try_fuse(const struct ggml_hexagon_session * sess, struct htp_opnode & prev, const struct htp_opnode & cur) {
+    if (!opt_opfusion) return false;
+
+    switch (cur.opcode) {
+        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_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))   ||
+                   (htp_fusion_enabled(HTP_FUSE_UNARY_MUL_GLU)  && fuse_unary_mul_glu(sess, prev, cur))  ||
+                   (htp_fusion_enabled(HTP_FUSE_MUL_MUL_SCALAR) && fuse_mul_mul_scalar(sess, prev, cur)) ||
+                   (htp_fusion_enabled(HTP_FUSE_BINARY_NX)      && fuse_binary_nx(sess, prev, cur));
+        case HTP_OP_SUB:
+        case HTP_OP_DIV:
+            return htp_fusion_enabled(HTP_FUSE_BINARY_NX) && fuse_binary_nx(sess, prev, cur);
+        case HTP_OP_SCALE:
+            return htp_fusion_enabled(HTP_FUSE_RMS_NORM_SCALE) && fuse_rms_norm_scale(sess, prev, cur);
+        case HTP_OP_UNARY_SILU:
+            return htp_fusion_enabled(HTP_FUSE_SSM_CONV_SILU) && fuse_ssm_conv_silu(sess, prev, cur);
+        case HTP_OP_GET_ROWS:
+            return htp_fusion_enabled(HTP_FUSE_REPEAT_GET_ROWS) && fuse_repeat_get_rows(sess, prev, cur);
+        case HTP_OP_CPY:
+            return htp_fusion_enabled(HTP_FUSE_GDN_CPY) && fuse_gdn_cpy(sess, prev, cur);
+        case HTP_OP_MUL_MAT:
+            return htp_fusion_enabled(HTP_FUSE_MUL_MAT_NX) && fuse_mul_mat_nx(sess, prev, cur);
+        case HTP_OP_MUL_MAT_ID:
+            return htp_fusion_enabled(HTP_FUSE_MUL_MAT_ID_NX) && fuse_mul_mat_id_nx(sess, prev, cur);
+        default:
+            return false;
+    }
+}
@@ -0,0 +1,33 @@
+#ifndef HTP_FUSION_H
+#define HTP_FUSION_H
+
+#include "htp-common.h"
+
+enum htp_fusion_flags {
+    HTP_FUSE_ALLREDUCE_ADD   = (1 << 1),  // 2
+    HTP_FUSE_RMS_NORM_MUL    = (1 << 2),  // 4
+    HTP_FUSE_MUL_MAT_ADD     = (1 << 3),  // 8
+    HTP_FUSE_MUL_MAT_NX      = (1 << 4),  // 16
+    HTP_FUSE_MUL_MAT_ID_NX   = (1 << 5),  // 32
+    HTP_FUSE_GDN_CPY         = (1 << 6),  // 64
+    HTP_FUSE_REPEAT_GET_ROWS = (1 << 7),  // 128
+    HTP_FUSE_UNARY_MUL_GLU   = (1 << 8),  // 256
+    HTP_FUSE_MUL_MUL_SCALAR  = (1 << 9),  // 512
+    HTP_FUSE_BINARY_NX       = (1 << 10), // 1024
+    HTP_FUSE_RMS_NORM_SCALE  = (1 << 11), // 2048
+    HTP_FUSE_SSM_CONV_SILU   = (1 << 12), // 4096
+};
+
+static inline bool htp_fusion_enabled(int flag) {
+    if (opt_opfusion <= 0) return false;
+    if (opt_opfusion == 1) return true; // 1 enables all
+    return (opt_opfusion & flag) != 0;
+}
+
+bool htp_fusion_can_stack(const struct ggml_tensor * a);
+bool htp_fusion_can_stack(const struct ggml_tensor * a, const struct ggml_tensor * b);
+int  htp_fusion_max(const struct ggml_cgraph * graph, int idx);
+void htp_fusion_tag_graph(struct ggml_cgraph * graph);
+bool htp_fusion_try_fuse(const struct ggml_hexagon_session * sess, struct htp_opnode & prev, const struct htp_opnode & cur);
+
+#endif // HTP_FUSION_H
@@ -0,0 +1,136 @@
+#ifndef HTP_KPARAMS_H
+#define HTP_KPARAMS_H
+
+#include "htp-common.h"
+
+#include "htp/allreduce-ops.h"
+#include "htp/argsort-ops.h"
+#include "htp/binary-ops.h"
+#include "htp/concat-ops.h"
+#include "htp/cpy-ops.h"
+#include "htp/cumsum-ops.h"
+#include "htp/diag-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/matmul-ops.h"
+#include "htp/pad-ops.h"
+#include "htp/pool-ops.h"
+#include "htp/repeat-ops.h"
+#include "htp/roll-ops.h"
+#include "htp/rope-ops.h"
+#include "htp/set-rows-ops.h"
+#include "htp/softmax-ops.h"
+#include "htp/solve-tri-ops.h"
+#include "htp/ssm-conv.h"
+#include "htp/sum-rows-ops.h"
+#include "htp/unary-ops.h"
+
+#include <type_traits>
+#include <utility>
+
+bool        htp_rows_stride(const int64_t * ne, const size_t * nb, size_t * stride);
+bool        htp_matmul_can_collapse(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst);
+bool        htp_matmul_add_can_collapse(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * src2, const struct ggml_tensor * dst);
+ggml_tensor htp_tensor_collapse_rows(const struct ggml_tensor * t);
+
+template <typename... Args>
+auto htp_kparams_type(bool (*)(Args...)) -> std::remove_pointer_t<std::remove_reference_t<decltype((std::declval<Args>(), ...))>>;
+
+template <auto SolveFn, typename... Args>
+static inline bool htp_kparams_check(Args &&... args) {
+    decltype(htp_kparams_type(SolveFn)) kparams;
+    return SolveFn(std::forward<Args>(args)..., &kparams);
+}
+
+struct htp_opnode;
+
+bool htp_kparams_solve(
+    const struct ggml_hexagon_session * sess,
+    struct htp_opnode & node
+);
+
+void htp_opnode_collapse(
+    const struct ggml_hexagon_session * sess,
+    struct htp_opnode & node
+);
+
+bool htp_kparams_solve(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    void * kparams = nullptr
+);
+
+static inline bool htp_kparams_check(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op
+) {
+    return htp_kparams_solve(sess, op, nullptr);
+}
+
+bool htp_kparams_solve_allreduce(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * dst,
+    uint32_t rank,
+    uint32_t n_ranks,
+    bool has_add,
+    bool is_row_bcast,
+    bool is_shard_ok,
+    struct htp_allreduce_kernel_params * kparams
+);
+
+bool htp_kparams_solve_rms_norm_mul(
+    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,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * src2,
+    const struct ggml_tensor * dst,
+    struct htp_mm_kernel_params * kparams
+);
+
+bool htp_kparams_solve_mmnx(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * act,
+    int32_t n_weights,
+    struct htp_mm_kernel_params * kparams
+);
+
+bool htp_kparams_solve_mmidnx(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * act,
+    const struct ggml_tensor * dst,
+    int32_t n_weights,
+    struct htp_mm_kernel_params * kparams
+);
+
+bool htp_kparams_solve_repeat_get_rows(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    struct htp_get_rows_kernel_params * kparams
+);
+
+bool htp_kparams_solve_unary_mul_glu(
+    const struct ggml_hexagon_session * sess,
+    uint32_t op,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    struct htp_glu_kernel_params * kparams
+);
+
+#endif // HTP_KPARAMS_H
@@ -0,0 +1,260 @@
+#include "htp-opnode.h"
+
+#include "htp/allreduce-ops.h"
+#include "htp/argsort-ops.h"
+#include "htp/binary-ops.h"
+#include "htp/cumsum-ops.h"
+#include "htp/diag-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/matmul-ops.h"
+#include "htp/pad-ops.h"
+#include "htp/repeat-ops.h"
+#include "htp/roll-ops.h"
+#include "htp/softmax-ops.h"
+#include "htp/solve-tri-ops.h"
+#include "htp/ssm-conv.h"
+#include "htp/sum-rows-ops.h"
+#include "htp/unary-ops.h"
+
+#include <stdio.h>
+
+std::string htp_opnode::name() const {
+    switch (opcode) {
+        case HTP_OP_MDEV_GROUP:    return "MDEV_GROUP";
+        case HTP_OP_CPY_FENCE:     return "CPY+FENCE";
+        case HTP_OP_FENCE:         return (node && node->op_params[1] == 0) ? "FENCE_WAIT" : "FENCE_SIGNAL";
+        case HTP_OP_ALLREDUCE:     return "ALLREDUCE";
+        case HTP_OP_ALLREDUCE_ADD: return "ALLREDUCE+ADD";
+        case HTP_OP_HADAMARD:      return "MUL_MAT_HADAMARD";
+        case HTP_OP_MUL_MAT_NX:    return "MUL_MAT_NX";
+        case HTP_OP_MUL_MAT_ID_NX: return "MUL_MAT_ID_NX";
+        case HTP_OP_ADD:
+        case HTP_OP_SUB:
+        case HTP_OP_MUL:
+        case HTP_OP_DIV:
+            if (!fused.empty()) {
+                const auto * kp = kparams<htp_binary_kernel_params>();
+                if (kp->kernel_type != HTP_BINARY_KERNEL_SCALAR) {
+                    switch (opcode) {
+                        case HTP_OP_ADD: return "ADD_NX";
+                        case HTP_OP_SUB: return "SUB_NX";
+                        case HTP_OP_MUL: return "MUL_NX";
+                        case HTP_OP_DIV: return "DIV_NX";
+                        default: break;
+                    }
+                }
+            }
+            break;
+        default:
+            break;
+    }
+
+    if (!node) {
+        return "";
+    }
+    if (htp_is_hadamard(node)) {
+        return "MUL_MAT_HADAMARD";
+    }
+
+    std::string res = ggml_op_desc(node);
+    for (const auto * t : fused) {
+        res += "+";
+        res += ggml_op_desc(t);
+    }
+    return res;
+}
+
+static const char * format_tensor_name(const struct ggml_tensor * t) {
+    return t ? t->name : "NONE";
+}
+
+static std::string format_tensor_dims(const struct ggml_tensor * t) {
+    if (!t) {
+        return "NONE";
+    }
+    char buf[64];
+    if (t->ne[2] == 1 && t->ne[3] == 1) {
+        snprintf(buf, sizeof(buf), "%d:%d", (int) t->ne[0], (int) t->ne[1]);
+    } else {
+        snprintf(buf, sizeof(buf), "%d:%d:%d:%d", (int) t->ne[0], (int) t->ne[1], (int) t->ne[2], (int) t->ne[3]);
+    }
+    return buf;
+}
+
+static const char * format_tensor_type(const struct ggml_tensor * t) {
+    return t ? ggml_type_name(t->type) : "NONE";
+}
+
+static std::string format_tensor_strides(const struct ggml_tensor * t) {
+    if (!t) {
+        return "NONE";
+    }
+    char buf[64];
+    const char * c = ggml_is_contiguous(t) ? "" : "!";
+    if (t->ne[2] == 1 && t->ne[3] == 1) {
+        snprintf(buf, sizeof(buf), "%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], c);
+    } else {
+        snprintf(buf, sizeof(buf), "%zu:%zu:%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], (size_t) t->nb[2], (size_t) t->nb[3], c);
+    }
+    return buf;
+}
+
+static const char * format_tensor_buff(const struct ggml_tensor * t) {
+    if (t && t->buffer) {
+        return ggml_backend_buffer_name(t->buffer);
+    }
+    return "NONE";
+}
+
+template <typename Fn>
+static std::string format_op_tensors(const htp_opnode & node, Fn fn) {
+    std::string res;
+    const auto & inputs  = node.get_inputs();
+    const auto & outputs = node.get_outputs();
+
+    if (!inputs.empty()) {
+        for (size_t i = 0; i < inputs.size(); i++) {
+            if (i > 0) res += " x ";
+            res += fn(inputs[i]);
+        }
+        res += " -> ";
+    }
+
+    if (outputs.empty()) {
+        res += fn(node.dst());
+    } else {
+        for (size_t i = 0; i < outputs.size(); i++) {
+            if (i > 0) res += " x ";
+            res += fn(outputs[i]);
+        }
+    }
+    return res;
+}
+
+static std::string format_kernel_params(const htp_opnode & node) {
+    char str[128];
+    const size_t max_size = sizeof(str);
+
+    if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID ||
+        node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ID_NX ||
+        node.opcode == HTP_OP_MUL_MAT_ADD) {
+        const auto * kparams = node.kparams<htp_mm_kernel_params>();
+        const char * path = "unknown";
+        int32_t type = kparams->kernel_type;
+        if (type == HTP_MM_KERNEL_HMX_2D || type == HTP_MM_KERNEL_HMX_BATCHED) {
+            path = (kparams->tile_size > 0) ? "hmx-tiled" : "hmx-flat";
+        } else if (type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
+                   type == HTP_MM_KERNEL_HVX_BF16_F32_VTCM) {
+            path = "hvx-flat";
+        } else if (type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
+            path = "hvx-tiled";
+        }
+        snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
+        const auto * kparams = node.kparams<htp_fa_kernel_params>();
+        const char * path = "unknown";
+        int32_t type = kparams->kernel_type;
+        if (type == HTP_FA_KERNEL_HMX) {
+            path = kparams->u.hmx.pipeline ? "hmx-pipe" : "hmx-seq";
+        } else if (type == HTP_FA_KERNEL_HVX) {
+            path = "hvx";
+        }
+        snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
+    } else if (htp_op_is_unary(node.opcode)) {
+        const auto * kparams = node.kparams<htp_unary_kernel_params>();
+        snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_MDEV_GROUP && node.node) {
+        snprintf(str, max_size, "idx %d count %d", (int) node.node->op_params[0], (int) node.dst()->ne[1]);
+    } else if ((node.opcode == HTP_OP_FENCE || node.opcode == HTP_OP_CPY_FENCE) && node.node) {
+        snprintf(str, max_size, "seq 0x%x", (uint32_t) node.node->op_params[0]);
+    } else if (node.opcode == HTP_OP_ALLREDUCE && node.node) {
+        snprintf(str, max_size, "seq 0x%x -> 0x%x", (uint32_t) node.node->op_params[0], (uint32_t) node.node->op_params[1]);
+    } else if (node.opcode == HTP_OP_SSM_CONV) {
+        const auto * kparams = node.kparams<htp_ssm_conv_kernel_params>();
+        snprintf(str, max_size, "%s vtcm %d", kparams->n_t == 1 ? "decode" : "prefill", (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_SOFTMAX) {
+        const auto * kparams = node.kparams<htp_softmax_kernel_params>();
+        snprintf(str, max_size, "k%d nth %d vtcm %d", (int) kparams->kernel_type, (int) kparams->n_threads, (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_GATED_DELTA_NET) {
+        const auto * kparams = node.kparams<htp_gdn_kernel_params>();
+        const char * path = (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) ? "hmx-chunked" : "hvx-recurrent";
+        snprintf(str, max_size, "%s-%s vtcm %u", path, kparams->kda ? "kda" : "scalar",
+                 (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_MUL || node.opcode == HTP_OP_ADD || node.opcode == HTP_OP_ADD_ID ||
+               node.opcode == HTP_OP_SUB || node.opcode == HTP_OP_DIV) {
+        const auto * kparams = node.kparams<htp_binary_kernel_params>();
+        snprintf(str, max_size, "vtcm %u", (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
+        const auto * kparams = node.kparams<htp_sort_kernel_params>();
+        snprintf(str, max_size, "%s nth %d nchk %d chk %d vtcm %d",
+                 node.opcode == HTP_OP_TOP_K ? "top_k" : "argsort",
+                 (int) kparams->n_threads, (int) kparams->n_chunks,
+                 (int) kparams->chunk_elems, (int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_GET_ROWS) {
+        const auto * kparams = node.kparams<htp_get_rows_kernel_params>();
+        const char * ktype_str = "unknown";
+        switch (kparams->kernel_type) {
+            case HTP_GET_ROWS_KERNEL_SAMETYPE: ktype_str = "sametype"; break;
+            case HTP_GET_ROWS_KERNEL_TILED:    ktype_str = "tiled"; break;
+            case HTP_GET_ROWS_KERNEL_FLAT:     ktype_str = "flat"; break;
+        }
+        snprintf(str, max_size, "%s%s vtcm %u", ktype_str, kparams->n_threads > 1 ? "-multi" : "", (unsigned int) kparams->vtcm_size);
+    } else if (htp_op_is_glu(node.opcode)) {
+        const auto * kparams = node.kparams<htp_glu_kernel_params>();
+        snprintf(str, max_size, "nth %d vtcm %u", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_PAD) {
+        const auto * kparams = node.kparams<htp_pad_kernel_params>();
+        snprintf(str, max_size, "k%d nth %d vtcm %u", (int) kparams->kernel_type, (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } 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_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);
+    } else if (node.opcode == HTP_OP_ROLL) {
+        const auto * kparams = node.kparams<htp_roll_kernel_params>();
+        snprintf(str, max_size, "k%d nth %d", (int) kparams->kernel_type, (int) kparams->n_threads);
+    } else if (node.opcode == HTP_OP_DIAG) {
+        const auto * kparams = node.kparams<htp_diag_kernel_params>();
+        snprintf(str, max_size, "%s nth %d vtcm %u", kparams->use_dma ? "dma" : "hvx", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_SOLVE_TRI) {
+        const auto * kparams = node.kparams<htp_solve_tri_kernel_params>();
+        snprintf(str, max_size, "%s nth %d", kparams->batched ? "batched" : "chunked", (int) kparams->n_threads);
+    } else {
+        snprintf(str, max_size, "----");
+    }
+    return str;
+}
+
+static void append_section(std::string & res, const std::string & sec) {
+    if (!res.empty()) {
+        res += "|";
+    }
+    res += sec;
+}
+
+std::string htp_format_tensor(const ggml_tensor * t, uint32_t flags) {
+    std::string res;
+    if (flags & HTP_OPFMT_NAMES)   append_section(res, format_tensor_name(t));
+    if (flags & HTP_OPFMT_DIMS)    append_section(res, format_tensor_dims(t));
+    if (flags & HTP_OPFMT_TYPES)   append_section(res, format_tensor_type(t));
+    if (flags & HTP_OPFMT_STRIDES) append_section(res, format_tensor_strides(t));
+    if (flags & HTP_OPFMT_BUFFS)   append_section(res, format_tensor_buff(t));
+    return res;
+}
+
+std::string htp_format_op(const htp_opnode & node, uint32_t flags) {
+    std::string res;
+    if (flags & HTP_OPFMT_OP)      append_section(res, node.name());
+    if (flags & HTP_OPFMT_NAMES)   append_section(res, format_op_tensors(node, format_tensor_name));
+    if (flags & HTP_OPFMT_DIMS)    append_section(res, format_op_tensors(node, format_tensor_dims));
+    if (flags & HTP_OPFMT_TYPES)   append_section(res, format_op_tensors(node, format_tensor_type));
+    if (flags & HTP_OPFMT_STRIDES) append_section(res, format_op_tensors(node, format_tensor_strides));
+    if (flags & HTP_OPFMT_BUFFS)   append_section(res, format_op_tensors(node, format_tensor_buff));
+    if (flags & HTP_OPFMT_KPARAMS) append_section(res, format_kernel_params(node));
+    return res;
+}
@@ -1,38 +1,74 @@
 #ifndef HTP_OPNODE_H
 #define HTP_OPNODE_H

-#define GGML_COMMON_IMPL_CPP
 #include "ggml-backend-impl.h"
-#include "ggml-common.h"
+#include "ggml-impl.h"
+#include "htp/htp-ops.h"

-#include <algorithm>
+#include <cstring>
+#include <initializer_list>
+#include <memory>
 #include <string>
 #include <vector>
-#include <memory>
-#include <stdio.h>
-#include "htp-ops.h"
-#include "htp/matmul-ops.h"
-#include "htp/flash-attn-ops.h"
-#include "htp/unary-ops.h"
-#include "htp/binary-ops.h"
-#include "htp/allreduce-ops.h"
-#include "htp/ssm-conv.h"
-#include "htp/gated-delta-net-ops.h"
-#include "htp/softmax-ops.h"
-#include "htp/argsort-ops.h"
-#include "htp/get-rows-ops.h"
+
+#define HTP_OP_MAX_FUSED 16
+
+// local std::inplace_vector replacement
+template <typename T, size_t MAX_N>
+struct htp_opvec {
+    T      data[MAX_N] {};
+    size_t n = 0;
+
+    htp_opvec() = default;
+    htp_opvec(std::initializer_list<T> init) {
+        GGML_ASSERT(init.size() <= MAX_N);
+        for (const auto & v : init) {
+            data[n++] = v;
+        }
+    }
+
+    size_t    size()  const { return n; }
+    bool      empty() const { return n == 0; }
+    void      clear()       { n = 0; }
+    T *       begin()       { return data; }
+    T *       end()         { return data + n; }
+    const T * begin() const { return data; }
+    const T * end()   const { return data + n; }
+    T &       back()        { return data[n - 1]; }
+    const T & back()  const { return data[n - 1]; }
+    T &       operator[](size_t i)       { return data[i]; }
+    const T & operator[](size_t i) const { return data[i]; }
+
+    void push_back(T val) {
+        GGML_ASSERT(n < MAX_N);
+        data[n++] = val;
+    }
+
+    void insert(T * pos, T val) {
+        GGML_ASSERT(n < MAX_N);
+        for (T * p = data + n; p > pos; --p) {
+            *p = *(p - 1);
+        }
+        *pos = val;
+        n++;
+    }
+};
+
+static inline bool htp_is_hadamard(const struct ggml_tensor * t) {
+    return t && t->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(t, 1) == GGML_HINT_SRC0_IS_HADAMARD;
+}

 struct htp_opnode {
-    ggml_tensor * node   { nullptr };
-    htp_op_code   opcode { HTP_OP_INVALID };
-    int32_t       kernel_params[HTP_OP_MAX_KERN_PARAMS] {0};
+    const ggml_tensor * node   { nullptr };
+    htp_op_code         opcode { HTP_OP_INVALID };
+    int32_t             params[HTP_OP_MAX_PARAMS] {0};
+    int32_t             kernel_params[HTP_OP_MAX_KERN_PARAMS] {0};

-    std::vector<ggml_tensor *>                fused;
-    std::vector<std::shared_ptr<ggml_tensor>> dummy;
+    std::vector<std::shared_ptr<ggml_tensor>>          dummy;

-    std::vector<const ggml_tensor *> inputs;
-    std::vector<const ggml_tensor *> outputs;
-    std::string                      name;
+    htp_opvec<const ggml_tensor *, HTP_OP_MAX_FUSED>   fused;
+    htp_opvec<const ggml_tensor *, HTP_OP_MAX_INPUTS>  inputs;
+    htp_opvec<const ggml_tensor *, HTP_OP_MAX_OUTPUTS> outputs;

    int n_active_src(const ggml_tensor * t) const {
        if (!t) return 0;
@@ -44,34 +80,36 @@ struct htp_opnode {
        return 0;
    }

-    void init(ggml_tensor * node) {
+    void init(const ggml_tensor * node) {
+        static_assert(sizeof(this->params) == sizeof(this->node->op_params), "params size mismatch");
+
        this->node = node;
-        if (this->node) {
-            if (this->node->op == GGML_OP_MUL_MAT &&
-                ggml_get_op_params_i32(this->node, 1) == GGML_HINT_SRC0_IS_HADAMARD) {
-                this->name = "MUL_MAT_HADAMARD";
-                this->inputs.push_back(this->node->src[1]);
-            } else {
-                this->name = ggml_op_desc(this->node);
+        if (!node) return;

-                // Build inputs (preserving optional nullptrs)
-                int n_inputs = n_active_src(this->node);
-                this->inputs.resize(n_inputs, nullptr);
-                for (int i = 0; i < n_inputs; i++) {
-                    this->inputs[i] = this->node->src[i];
-                }
+        memcpy(this->params, this->node->op_params, sizeof(this->params));
+
+        // Build inputs (preserving optional nullptrs)
+        if (htp_is_hadamard(this->node)) {
+            this->inputs.push_back(this->node->src[1]);
+        } else {
+            int n_inputs = n_active_src(this->node);
+            for (int i = 0; i < n_inputs; i++) {
+                this->inputs.push_back(this->node->src[i]);
            }
-
-            // Build outputs
-            this->outputs.push_back(this->dst());
        }
+
+        // Build outputs
+        this->outputs.push_back(this->dst());
    }

-    htp_opnode(htp_op_code opcode = HTP_OP_INVALID, ggml_tensor * node = nullptr) : opcode(opcode) {
+    htp_opnode(htp_op_code opcode = HTP_OP_INVALID, const ggml_tensor * node = nullptr) : opcode(opcode) {
+        init(node);
+    }
+
+    htp_opnode(const ggml_tensor * node) : opcode(HTP_OP_INVALID) {
        init(node);
    }

-    ggml_op             op()   const { return node->op; }
    const ggml_tensor * src0() const { return node->src[0]; }
    const ggml_tensor * src1() const { return node->src[1]; }
    const ggml_tensor * dst()  const { return outputs.empty() ? node : outputs.back(); }
@@ -81,343 +119,54 @@ struct htp_opnode {
        return dummy.back().get();
    }

-    void add_fused(ggml_tensor * t, bool extra_dst = false) {
-        fused.push_back(t);
-
-        name += "+";
-        name += ggml_op_desc(t);
-
-        if (extra_dst) {
-            outputs.push_back(t);
-        } else {
-            outputs.clear();
-            outputs.push_back(t);
-        }
-
-        // Remove the newly fused intermediate output tensor t from inputs (if it was there)
-        inputs.erase(std::remove(inputs.begin(), inputs.end(), t), inputs.end());
-
-        // Append new inputs from t, preserving middle nullptrs
-        int n_inputs = n_active_src(t);
-        for (int i = 0; i < n_inputs; i++) {
-            const auto * src = t->src[i];
-            if (!src) {
-                inputs.push_back(nullptr);
-            } else if (src != node &&
-                       std::find(fused.begin(), fused.end(), src) == fused.end() &&
-                       std::find(inputs.begin(), inputs.end(), src) == inputs.end()) {
-                inputs.push_back(src);
-            }
-        }
-    }
-
-    const std::vector<const ggml_tensor *> & get_inputs() const {
+    const htp_opvec<const ggml_tensor *, HTP_OP_MAX_INPUTS> & get_inputs() const {
        return inputs;
    }

-    const std::vector<const ggml_tensor *> & get_outputs() const {
+    const htp_opvec<const ggml_tensor *, HTP_OP_MAX_OUTPUTS> & get_outputs() const {
        return outputs;
    }

-    std::string op_name() const {
-        return name;
+    template <typename T>
+    T * kparams() {
+        static_assert(sizeof(T) <= sizeof(kernel_params), "kparams size exceeds HTP_OP_MAX_KERN_PARAMS");
+        static_assert(alignof(T) <= alignof(int32_t), "kparams alignment exceeds int32_t");
+        return reinterpret_cast<T *>(kernel_params);
    }

-    bool is_empty() const {
-        return ggml_op_is_empty(node->op);
+    template <typename T>
+    const T * kparams() const {
+        static_assert(sizeof(T) <= sizeof(kernel_params), "kparams size exceeds HTP_OP_MAX_KERN_PARAMS");
+        static_assert(alignof(T) <= alignof(int32_t), "kparams alignment exceeds int32_t");
+        return reinterpret_cast<const T *>(kernel_params);
    }

-    bool stackable() const {
-        switch (this->op()) {
-            case GGML_OP_MUL_MAT:
-                if (ggml_get_op_params_i32(this->node, 1) == GGML_HINT_SRC0_IS_HADAMARD) {
-                    return false;
-                }
-                return ggml_is_quantized(this->src0()->type);
-            case GGML_OP_MUL_MAT_ID:
-                return ggml_is_quantized(this->src0()->type);
-            default:
-                return false;
-        }
+    template <typename T>
+    void set_kparams(const T & kp) {
+        static_assert(sizeof(T) <= sizeof(kernel_params), "kparams size exceeds HTP_OP_MAX_KERN_PARAMS");
+        static_assert(alignof(T) <= alignof(int32_t), "kparams alignment exceeds int32_t");
+        memcpy(kernel_params, &kp, sizeof(T));
    }

-    bool same_input(const htp_opnode& n) const {
-        return n.src1() == this->src1();
-    }
+    std::string name() const;
 };

-struct htp_opformat {
-    char strides[64 * GGML_MAX_SRC];
-    char dims[64 * GGML_MAX_SRC];
-    char types[16 * GGML_MAX_SRC];
-    char buffs[64 * GGML_MAX_SRC];
-    char names[64 * GGML_MAX_SRC];
-    char kparams[128];
+enum htp_opformat_flags {
+    HTP_OPFMT_OP      = (1 << 0),
+    HTP_OPFMT_NAMES   = (1 << 1),
+    HTP_OPFMT_DIMS    = (1 << 2),
+    HTP_OPFMT_TYPES   = (1 << 3),
+    HTP_OPFMT_STRIDES = (1 << 4),
+    HTP_OPFMT_BUFFS   = (1 << 5),
+    HTP_OPFMT_KPARAMS = (1 << 6),

-    int format_tensor_dims(char * str, size_t max_size, const struct ggml_tensor * t) {
-        if (!t) {
-            return snprintf(str, max_size, "NONE");
-        }
-        if (t->ne[2] == 1 && t->ne[3] == 1) {
-            return snprintf(str, max_size, "%d:%d", (int) t->ne[0], (int) t->ne[1]);
-        } else {
-            return snprintf(str, max_size, "%d:%d:%d:%d", (int) t->ne[0], (int) t->ne[1], (int) t->ne[2], (int) t->ne[3]);
-        }
-    }
-
-    void format_op_dims(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            p += std::min((size_t)format_tensor_dims(p, p_end - p, inputs[0]), (size_t)(p_end - p));
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)format_tensor_dims(p, p_end - p, inputs[i]), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        char self[64];
-        format_tensor_dims(self, sizeof(self), node.dst());
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", self), (size_t)(p_end - p));
-        }
-    }
-
-    int format_tensor_strides(char * str, size_t max_size, const struct ggml_tensor * t) {
-        if (!t) {
-            return snprintf(str, max_size, "NONE");
-        }
-        const char * c = ggml_is_contiguous(t) ? "" : "!";
-
-        if (t->ne[2] == 1 && t->ne[3] == 1) {
-            return snprintf(str, max_size, "%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], c);
-        } else {
-            return snprintf(str, max_size, "%zu:%zu:%zu:%zu%s", (size_t) t->nb[0], (size_t) t->nb[1], (size_t) t->nb[2], (size_t) t->nb[3], c);
-        }
-    }
-
-    void format_op_strides(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            p += std::min((size_t)format_tensor_strides(p, p_end - p, inputs[0]), (size_t)(p_end - p));
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)format_tensor_strides(p, p_end - p, inputs[i]), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        char self[64];
-        format_tensor_strides(self, sizeof(self), node.dst());
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", self), (size_t)(p_end - p));
-        }
-    }
-
-    void format_op_types(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[0] ? ggml_type_name(inputs[0]->type) : "NONE"), (size_t)(p_end - p));
-            }
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[i] ? ggml_type_name(inputs[i]->type) : "NONE"), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", ggml_type_name(node.dst()->type)), (size_t)(p_end - p));
-        }
-    }
-
-    const char * tensor_buff_name(const struct ggml_tensor * t) {
-        if (t && t->buffer) {
-            return ggml_backend_buffer_name(t->buffer);
-        }
-        return "NONE";
-    }
-
-    void format_op_buffs(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(inputs[0])), (size_t)(p_end - p));
-            }
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(inputs[i])), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", tensor_buff_name(node.dst())), (size_t)(p_end - p));
-        }
-    }
-
-    void format_op_names(char * str, size_t max_size, const htp_opnode & node) {
-        char * p = str;
-        char * p_end = str + max_size;
-        auto inputs = node.get_inputs();
-
-        if (!inputs.empty()) {
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[0] ? inputs[0]->name : "NONE"), (size_t)(p_end - p));
-            }
-
-            for (size_t i = 1; i < inputs.size(); i++) {
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, " x "), (size_t)(p_end - p));
-                }
-                if (p < p_end) {
-                    p += std::min((size_t)snprintf(p, p_end - p, "%s", inputs[i] ? inputs[i]->name : "NONE"), (size_t)(p_end - p));
-                }
-            }
-
-            if (p < p_end) {
-                p += std::min((size_t)snprintf(p, p_end - p, " -> "), (size_t)(p_end - p));
-            }
-        }
-
-        if (p < p_end) {
-            p += std::min((size_t)snprintf(p, p_end - p, "%s", node.dst()->name), (size_t)(p_end - p));
-        }
-    }
-    void format_kernel_params(char * str, size_t max_size, const htp_opnode & node) {
-        if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID ||
-            node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ID_NX ||
-            node.opcode == HTP_OP_MUL_MAT_ADD) {
-            const auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params;
-            const char * path = "unknown";
-            int32_t type = kparams->kernel_type;
-            if (type == HTP_MM_KERNEL_HMX_2D || type == HTP_MM_KERNEL_HMX_BATCHED) {
-                path = (kparams->tile_size > 0) ? "hmx-tiled" : "hmx-flat";
-            } else if (type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
-                       type == HTP_MM_KERNEL_HVX_BF16_F32_VTCM) {
-                path = "hvx-flat";
-            } else if (type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
-                path = "hvx-tiled";
-            }
-            snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
-            const auto * kparams = (const struct htp_fa_kernel_params *) node.kernel_params;
-            const char * path = "unknown";
-            int32_t type = kparams->kernel_type;
-            if (type == HTP_FA_KERNEL_HMX) {
-                path = kparams->u.hmx.pipeline ? "hmx-pipe" : "hmx-seq";
-            } else if (type == HTP_FA_KERNEL_HVX) {
-                path = "hvx";
-            }
-            snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size);
-        } else if (htp_op_is_unary(node.opcode)) {
-            const auto * kparams = (const struct htp_unary_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "%s vtcm %d", kparams->col_tile ? "wide-row" : "row-block", (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_MDEV_GROUP && node.node) {
-            snprintf(str, max_size, "idx %d count %d", (int) node.node->op_params[0], (int) node.dst()->ne[1]);
-        } else if ((node.opcode == HTP_OP_FENCE || node.opcode == HTP_OP_CPY_FENCE) && node.node) {
-            snprintf(str, max_size, "seq 0x%x", (uint32_t) node.node->op_params[0]);
-        } else if (node.opcode == HTP_OP_ALLREDUCE && node.node) {
-            snprintf(str, max_size, "seq 0x%x -> 0x%x", (uint32_t) node.node->op_params[0], (uint32_t) node.node->op_params[1]);
-        } else if (node.opcode == HTP_OP_SSM_CONV) {
-            const auto * kparams = (const struct htp_ssm_conv_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "%s vtcm %d", kparams->n_t == 1 ? "decode" : "prefill", (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_SOFTMAX) {
-            const auto * kparams = (const struct htp_softmax_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "k%d nth %d vtcm %d", (int) kparams->kernel_id, (int) kparams->n_threads, (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_GATED_DELTA_NET) {
-            const auto * kparams = (const struct htp_gdn_kernel_params *) node.kernel_params;
-            const char * path = (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) ? "hmx-chunked" : "hvx-recurrent";
-            snprintf(str, max_size, "%s-%s vtcm %u", path, kparams->kda ? "kda" : "scalar",
-                     (unsigned int) (kparams->vtcm_size ? kparams->vtcm_size : kparams->vtcm_per_thread * kparams->n_threads));
-        } 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 = (const struct htp_binary_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "vtcm %u", (unsigned int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
-            const auto * kparams = (const struct htp_sort_kernel_params *) node.kernel_params;
-            snprintf(str, max_size, "%s nth %d nchk %d chk %d vtcm %d",
-                     node.opcode == HTP_OP_TOP_K ? "top_k" : "argsort",
-                     (int) kparams->n_threads, (int) kparams->n_chunks,
-                     (int) kparams->chunk_elems, (int) kparams->vtcm_size);
-        } else if (node.opcode == HTP_OP_GET_ROWS) {
-            const auto * kparams = (const struct htp_get_rows_kernel_params *) node.kernel_params;
-            const char * ktype_str = "unknown";
-            switch (kparams->kernel_type) {
-                case HTP_GET_ROWS_KERNEL_SAMETYPE: ktype_str = "sametype"; break;
-                case HTP_GET_ROWS_KERNEL_TILED:    ktype_str = "tiled"; break;
-                case HTP_GET_ROWS_KERNEL_FLAT:     ktype_str = "flat"; break;
-            }
-            snprintf(str, max_size, "%s%s vtcm %u", ktype_str, kparams->n_threads > 1 ? "-multi" : "", (unsigned int) kparams->vtcm_size);
-        } else {
-            snprintf(str, max_size, "----");
-        }
-
-    }
-
-    void format(const htp_opnode & node) {
-        format_op_dims(dims, sizeof(dims), node);
-        format_op_strides(strides, sizeof(strides), node);
-        format_op_types(types, sizeof(types), node);
-        format_op_buffs(buffs, sizeof(buffs), node);
-        format_op_names(names, sizeof(names), node);
-        format_kernel_params(kparams, sizeof(kparams), node);
-    }
-
-    htp_opformat() {
-        strides[0] = '\0';
-        dims[0]    = '\0';
-        types[0]   = '\0';
-        buffs[0]   = '\0';
-        names[0]   = '\0';
-        kparams[0] = '\0';
-    }
-    htp_opformat(const htp_opnode & node) { format(node); }
+    HTP_OPFMT_EXEC    = HTP_OPFMT_OP | HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES | HTP_OPFMT_BUFFS | HTP_OPFMT_KPARAMS,
+    HTP_OPFMT_SUPP    = HTP_OPFMT_OP | HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES | HTP_OPFMT_BUFFS,
+    HTP_OPFMT_PROF    = HTP_OPFMT_OP | HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES | HTP_OPFMT_KPARAMS,
+    HTP_OPFMT_TENSOR  = HTP_OPFMT_NAMES | HTP_OPFMT_DIMS | HTP_OPFMT_TYPES | HTP_OPFMT_STRIDES,
 };

+std::string htp_format_op(const htp_opnode & node, uint32_t flags = HTP_OPFMT_EXEC);
+std::string htp_format_tensor(const ggml_tensor * t, uint32_t flags = HTP_OPFMT_TENSOR);
+
 #endif // HTP_OPNODE_H
@@ -28,7 +28,7 @@ add_library(${HTP_LIB} SHARED
    unary-ops.c
    sum-rows-ops.c
    softmax-ops.c
-    act-ops.c
+    glu-ops.c
    rope-ops.c
    set-rows-ops.c
    get-rows-ops.c
@@ -32,6 +32,7 @@ static inline size_t htp_allreduce_vtcm_buffer_count(
 }

 struct htp_allreduce_kernel_params {
+    uint8_t kernel_type;
    int32_t rank;
    int32_t n_ranks;
    int32_t n_threads;
@@ -39,7 +40,6 @@ struct htp_allreduce_kernel_params {
    int32_t elems_per_thread;     // 1D: nelem_per_thread, 2D: nrows_per_thread
    int32_t vtcm_size_per_thread;
    int32_t vtcm_size;
-    int32_t kernel_type;
    int32_t ne0;
    int32_t ne1;
    int32_t row_size_aligned;
@@ -50,6 +50,12 @@ struct htp_allreduce_kernel_params {
    int32_t mode;                 // enum htp_allreduce_mode: FULL or SHARDED_FANOUT
 };

+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_allreduce_kernel_params) <= 128, "htp_allreduce_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_allreduce_kernel_params) <= 128, "htp_allreduce_kernel_params is too large for kernel_params blob");
+#endif
+
 #ifdef __cplusplus
 }
 #endif
@@ -453,7 +453,10 @@ static void merge_runs(
 struct htp_sort_chunk_ctx {
    struct htp_ops_context *              octx;
    const struct htp_sort_kernel_params * kparams;
-    uint8_t *                             vtcm_base;
+    uint8_t *                             phase1_base;
+    size_t                                phase1_slot_size;
+    float *                               merge_values;
+    int32_t *                             merge_indices;
    uint32_t                              real_count[HTP_MAX_NTHREADS];
 };

@@ -462,26 +465,27 @@ static void htp_sort_chunk_job(unsigned int n, unsigned int i, void * data) {
    struct htp_ops_context * octx = cctx->octx;
    const struct htp_sort_kernel_params * kparams = cctx->kparams;
    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;

    uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
-    uint32_t ne00        = (uint32_t) kparams->ne00;
+    uint32_t ne00        = src0->ne[0];
    uint32_t chunk_base  = i * chunk_elems;

    uint32_t real_count = (chunk_base < ne00) ? MIN(chunk_elems, ne00 - chunk_base) : 0;
    cctx->real_count[i] = real_count;

-    uint8_t * spad = cctx->vtcm_base + (size_t) kparams->phase1_slot_size * i;
+    uint8_t * spad = cctx->phase1_base + cctx->phase1_slot_size * i;
    float *   values_buf  = (float *) spad;
    int32_t * indices_buf = (int32_t *) (spad + chunk_elems * sizeof(float));

-    float *   merge_values  = (float *)   (cctx->vtcm_base + kparams->merge_values_off);
-    int32_t * merge_indices = (int32_t *) (cctx->vtcm_base + kparams->merge_indices_off);
+    float *   merge_values  = cctx->merge_values;
+    int32_t * merge_indices = cctx->merge_indices;

    struct htp_thread_trace * tr = &octx->ctx->trace[i];
    const dma_addr_t src_addr = src0->data + (size_t) chunk_base * sizeof(float);

    if (kparams->is_top_k) {
-        uint32_t k = (uint32_t) kparams->k;
+        uint32_t k = dst->ne[0];
        uint32_t local_k = MIN(k, chunk_elems);

        if (real_count == 0) {
@@ -559,17 +563,17 @@ static void htp_argsort_merge_l1_job(unsigned int n, unsigned int i, void * data
    uint32_t count0 = cctx->real_count[c0];
    uint32_t count1 = cctx->real_count[c1];

-    size_t slot_size = (size_t) kparams->phase1_slot_size;
+    size_t slot_size = cctx->phase1_slot_size;
    uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;

-    const float *   val0 = (const float *)   (cctx->vtcm_base + slot_size * c0);
-    const int32_t * idx0 = (const int32_t *) (cctx->vtcm_base + slot_size * c0 + chunk_elems * sizeof(float));
-    const float *   val1 = (const float *)   (cctx->vtcm_base + slot_size * c1);
-    const int32_t * idx1 = (const int32_t *) (cctx->vtcm_base + slot_size * c1 + chunk_elems * sizeof(float));
+    const float *   val0 = (const float *)   (cctx->phase1_base + slot_size * c0);
+    const int32_t * idx0 = (const int32_t *) (cctx->phase1_base + slot_size * c0 + chunk_elems * sizeof(float));
+    const float *   val1 = (const float *)   (cctx->phase1_base + slot_size * c1);
+    const int32_t * idx1 = (const int32_t *) (cctx->phase1_base + slot_size * c1 + chunk_elems * sizeof(float));

    uint32_t out_offset = (2 * i) * chunk_elems;
-    float *   out_val = (float *)   (cctx->vtcm_base + kparams->merge_values_off)  + out_offset;
-    int32_t * out_idx = (int32_t *) (cctx->vtcm_base + kparams->merge_indices_off) + out_offset;
+    float *   out_val = cctx->merge_values  + out_offset;
+    int32_t * out_idx = cctx->merge_indices + out_offset;

    struct htp_thread_trace * tr = &cctx->octx->ctx->trace[i];
    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
@@ -587,21 +591,21 @@ static void htp_argsort_merge_l2_job(unsigned int n, unsigned int i, void * data
    struct htp_sort_chunk_ctx * cctx = mctx->cctx;
    const struct htp_sort_kernel_params * kparams = cctx->kparams;

-    size_t slot_size = (size_t) kparams->phase1_slot_size;
+    size_t slot_size = cctx->phase1_slot_size;
    uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;

    uint32_t off0 = (4 * i) * chunk_elems;
    uint32_t off1 = (4 * i + 2) * chunk_elems;

-    const float *   val0 = (const float *)   (cctx->vtcm_base + kparams->merge_values_off)  + off0;
-    const int32_t * idx0 = (const int32_t *) (cctx->vtcm_base + kparams->merge_indices_off) + off0;
-    const float *   val1 = (const float *)   (cctx->vtcm_base + kparams->merge_values_off)  + off1;
-    const int32_t * idx1 = (const int32_t *) (cctx->vtcm_base + kparams->merge_indices_off) + off1;
+    const float *   val0 = cctx->merge_values  + off0;
+    const int32_t * idx0 = cctx->merge_indices + off0;
+    const float *   val1 = cctx->merge_values  + off1;
+    const int32_t * idx1 = cctx->merge_indices + off1;

    uint32_t count0 = cctx->real_count[4*i + 0] + cctx->real_count[4*i + 1];
    uint32_t count1 = cctx->real_count[4*i + 2] + cctx->real_count[4*i + 3];

-    uint8_t * slot_base = cctx->vtcm_base + (4 * i) * slot_size;
+    uint8_t * slot_base = cctx->phase1_base + (4 * i) * slot_size;
    float *   out_val   = (float *) slot_base;
    int32_t * out_idx   = (int32_t *) (slot_base + 4 * chunk_elems * sizeof(float));

@@ -614,7 +618,10 @@ static void htp_argsort_merge_l2_job(unsigned int n, unsigned int i, void * data
 struct htp_sort_multi_row_ctx {
    struct htp_ops_context *              octx;
    const struct htp_sort_kernel_params * kparams;
-    uint8_t *                             vtcm_base;
+    uint8_t *                             phase1_base;
+    size_t                                phase1_slot_size;
+    uint32_t                              row_start;
+    uint32_t                              row_end;
    uint32_t                              nrows_per_thread;
 };

@@ -624,12 +631,18 @@ struct htp_tensor_row_coord {
    uint32_t i03;
 };

-static inline struct htp_tensor_row_coord init_row_coord(uint32_t r, uint32_t ne01, uint32_t ne02) {
+static inline struct htp_tensor_row_coord init_row_coord(
+    uint32_t r,
+    uint32_t ne01,
+    uint32_t ne02,
+    const struct fastdiv_values * div_ne01,
+    const struct fastdiv_values * div_ne01_ne02
+) {
    uint32_t ne01_ne02 = ne01 * ne02;
-    uint32_t i03 = (ne01_ne02 > 0) ? (r / ne01_ne02) : 0;
-    uint32_t rem = (ne01_ne02 > 0) ? (r % ne01_ne02) : 0;
-    uint32_t i02 = (ne01 > 0) ? (rem / ne01) : 0;
-    uint32_t i01 = (ne01 > 0) ? (rem % ne01) : 0;
+    uint32_t i03 = (ne01_ne02 > 0) ? fastdiv(r, div_ne01_ne02) : 0;
+    uint32_t rem = (ne01_ne02 > 0) ? (r - i03 * ne01_ne02) : 0;
+    uint32_t i02 = (ne01 > 0) ? fastdiv(rem, div_ne01) : 0;
+    uint32_t i01 = (ne01 > 0) ? (rem - i02 * ne01) : 0;
    return (struct htp_tensor_row_coord){ i01, i02, i03 };
 }

@@ -701,7 +714,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
    const struct htp_tensor * src0 = octx->src[0];
    const struct htp_tensor * dst  = octx->dst;

-    uint32_t ne00 = (uint32_t) kparams->ne00;
+    uint32_t ne00 = src0->ne[0];
    uint32_t ne01 = src0->ne[1];
    uint32_t ne02 = src0->ne[2];

@@ -713,16 +726,16 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
    uint32_t nb2 = dst->nb[2];
    uint32_t nb3 = dst->nb[3];

-    uint32_t row_start = (uint32_t) kparams->row_start;
-    uint32_t row_end   = (uint32_t) kparams->row_end;
+    uint32_t row_start = mctx->row_start;
+    uint32_t row_end   = mctx->row_end;
    uint32_t r_start   = row_start + ith * mctx->nrows_per_thread;
    uint32_t r_end     = MIN(r_start + mctx->nrows_per_thread, row_end);

    if (r_start >= r_end) return;

-    size_t spad_offset = ith * (size_t) (kparams->phase1_slot_size * kparams->n_slots);
-    uint8_t * thread_spad = mctx->vtcm_base + spad_offset;
-    size_t slot_size = (size_t) kparams->phase1_slot_size;
+    size_t slot_size = mctx->phase1_slot_size;
+    size_t spad_offset = ith * slot_size * (size_t) kparams->n_slots;
+    uint8_t * thread_spad = mctx->phase1_base + spad_offset;

    uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
    float *   values_buf  = (float *) thread_spad;
@@ -732,14 +745,14 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];

    bool is_top_k = kparams->is_top_k != 0;
-    uint32_t k = (uint32_t) kparams->k;
+    uint32_t k = dst->ne[0];
    enum ggml_sort_order order = (enum ggml_sort_order) kparams->order;

    const uint32_t src_bytes = ne00 * sizeof(float);
    const uint32_t dst_bytes = (is_top_k ? k : ne00) * sizeof(int32_t);

    if (kparams->n_slots == 2) {
-        struct htp_tensor_row_coord prime_coord = init_row_coord(r_start, ne01, ne02);
+        struct htp_tensor_row_coord prime_coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
        for (uint32_t r = r_start, spad_idx = 0; r < r_end && spad_idx < 2; r++, spad_idx++) {
            uint8_t * cur_spad = thread_spad + spad_idx * slot_size;
            float *   cur_values  = (float *) cur_spad;
@@ -759,8 +772,8 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
            step_row_coord(&prime_coord, ne01, ne02);
        }

-        struct htp_tensor_row_coord dst_coord = init_row_coord(r_start, ne01, ne02);
-        struct htp_tensor_row_coord next_src_coord = init_row_coord(r_start + 2, ne01, ne02);
+        struct htp_tensor_row_coord dst_coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
+        struct htp_tensor_row_coord next_src_coord = init_row_coord(r_start + 2, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);

        for (uint32_t r = r_start; r < r_end; r++) {
            uint32_t cur_slot = (r - r_start) & 1;
@@ -796,7 +809,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data

        dma_queue_flush(dma_q);
    } else {
-        struct htp_tensor_row_coord coord = init_row_coord(r_start, ne01, ne02);
+        struct htp_tensor_row_coord coord = init_row_coord(r_start, ne01, ne02, &kparams->div_ne01, &kparams->div_ne01_ne02);
        for (uint32_t r = r_start; r < r_end; r++) {
            dma_addr_t src_addr = src0->data + coord_offset(coord, nb01, nb02, nb03);
            dma_addr_t dst_addr = dst->data + coord_offset(coord, nb1, nb2, nb3);
@@ -847,52 +860,39 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
    uint32_t ne00 = src0->ne[0];
    uint32_t k    = dst->ne[0];

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

-    if (kparams->n_threads == 0) {
-        struct htp_sort_vtcm_layout layout;
-        if (!htp_sort_solve_layout(&layout, ne00, total_rows, k, octx->n_threads, octx->ctx->vtcm_size, is_top_k)) {
-            return HTP_STATUS_VTCM_TOO_SMALL;
-        }
-        memset(&kparams_local, 0, sizeof(kparams_local));
-        kparams_local.n_threads         = (int32_t) layout.n_threads;
-        kparams_local.total_rows        = (int32_t) total_rows;
-        kparams_local.row_start         = (int32_t) row_start;
-        kparams_local.row_end           = (int32_t) row_end;
-        kparams_local.ne00              = (int32_t) ne00;
-        kparams_local.k                 = (int32_t) k;
-        kparams_local.order             = is_top_k ? GGML_SORT_ORDER_DESC : (int32_t) octx->op_params[0];
-        kparams_local.is_top_k          = is_top_k ? 1 : 0;
-        kparams_local.use_dma           = 1;
-        kparams_local.chunk_elems       = (int32_t) layout.chunk_elems;
-        kparams_local.n_chunks          = (int32_t) layout.n_chunks;
-        kparams_local.vtcm_size         = (int32_t) layout.total_bytes;
-        kparams_local.phase1_slot_size  = (int32_t) layout.phase1_slot_size;
-        kparams_local.merge_values_off  = (int32_t) layout.merge_values_off;
-        kparams_local.merge_indices_off = (int32_t) layout.merge_indices_off;
-        kparams_local.merge_elems       = (int32_t) layout.merge_elems;
-        kparams_local.n_slots           = (int32_t) layout.n_slots;
-        kparams = &kparams_local;
-    } else if (octx->ctx->mdev.count > 1) {
-        memcpy(&kparams_local, kparams, sizeof(kparams_local));
-        kparams_local.row_start = (int32_t) row_start;
-        kparams_local.row_end   = (int32_t) row_end;
-        kparams = &kparams_local;
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
    }

-    if (octx->ctx->vtcm_size < (size_t) kparams->vtcm_size) {
-        FARF(ERROR, "sort: VTCM size too small. Needed %d, have %zu", kparams->vtcm_size, octx->ctx->vtcm_size);
+    struct htp_sort_vtcm_layout layout;
+    htp_sort_vtcm_layout_build(&layout,
+                               total_rows,
+                               (uint32_t) kparams->chunk_elems,
+                               (uint32_t) kparams->n_chunks,
+                               (uint32_t) kparams->merge_elems,
+                               (uint32_t) kparams->n_threads,
+                               (uint32_t) kparams->n_slots);
+
+    if (octx->ctx->vtcm_size < layout.total_bytes) {
+        FARF(ERROR, "sort: VTCM size too small. Needed %zu, have %zu", layout.total_bytes, octx->ctx->vtcm_size);
        return HTP_STATUS_VTCM_TOO_SMALL;
    }

-    uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
+    uint8_t * vtcm_base     = (uint8_t *) octx->ctx->vtcm_base;
+    uint8_t * phase1_base   = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_phase1);
+    float *   merge_values  = VTCM_LAYOUT_PTR(float,   vtcm_base, layout.off_merge_values);
+    int32_t * merge_indices = VTCM_LAYOUT_PTR(int32_t, vtcm_base, layout.off_merge_indices);

    if (total_rows == 1 && kparams->n_chunks > 1) {
        struct htp_sort_chunk_ctx cctx;
-        cctx.octx      = octx;
-        cctx.kparams   = kparams;
-        cctx.vtcm_base = vtcm_base;
+        cctx.octx             = octx;
+        cctx.kparams          = kparams;
+        cctx.phase1_base      = phase1_base;
+        cctx.phase1_slot_size = layout.phase1_slot_size;
+        cctx.merge_values     = merge_values;
+        cctx.merge_indices    = merge_indices;
        memset(cctx.real_count, 0, sizeof(cctx.real_count));

        work_queue_run(octx->ctx->work_queue, htp_sort_chunk_job, &cctx, (uint32_t) kparams->n_chunks);
@@ -902,8 +902,8 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
            uint32_t local_k     = MIN(k, chunk_elems);

            if (k == 1 && ne00 >= 128*1024) {
-                float *   cand_vals = (float *)   (vtcm_base + kparams->merge_values_off);
-                int32_t * cand_idxs = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+                float *   cand_vals = merge_values;
+                int32_t * cand_idxs = merge_indices;
                float best_val = cand_vals[0];
                int32_t best_idx = cand_idxs[0];
                for (int32_t c = 1; c < kparams->n_chunks; c++) {
@@ -918,8 +918,8 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
            }

            uint32_t total_candidates = (uint32_t) kparams->n_chunks * local_k;
-            float *   cand_vals = (float *)   (vtcm_base + kparams->merge_values_off);
-            int32_t * cand_idxs = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+            float *   cand_vals = merge_values;
+            int32_t * cand_idxs = merge_indices;
            uint32_t merge_elems = (uint32_t) kparams->merge_elems;

            if (merge_elems > total_candidates) {
@@ -943,18 +943,18 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                struct htp_argsort_merge_l2_ctx l2ctx = { &cctx, asc };
                work_queue_run(octx->ctx->work_queue, htp_argsort_merge_l2_job, &l2ctx, 2);

-                size_t slot_size = (size_t) kparams->phase1_slot_size;
+                size_t slot_size = layout.phase1_slot_size;
                uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;

-                const float *   val0 = (const float *) vtcm_base;
-                const int32_t * idx0 = (const int32_t *) (vtcm_base + 4 * chunk_elems * sizeof(float));
-                const float *   val1 = (const float *) (vtcm_base + 4 * slot_size);
-                const int32_t * idx1 = (const int32_t *) (vtcm_base + 4 * slot_size + 4 * chunk_elems * sizeof(float));
+                const float *   val0 = (const float *) phase1_base;
+                const int32_t * idx0 = (const int32_t *) (phase1_base + 4 * chunk_elems * sizeof(float));
+                const float *   val1 = (const float *) (phase1_base + 4 * slot_size);
+                const int32_t * idx1 = (const int32_t *) (phase1_base + 4 * slot_size + 4 * chunk_elems * sizeof(float));

                uint32_t n01 = cctx.real_count[0] + cctx.real_count[1] + cctx.real_count[2] + cctx.real_count[3];
                uint32_t n23 = cctx.real_count[4] + cctx.real_count[5] + cctx.real_count[6] + cctx.real_count[7];

-                int32_t * final_idx = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+                int32_t * final_idx = merge_indices;
                struct htp_thread_trace * tr = &octx->ctx->trace[0];
                htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                merge_runs(val0, idx0, n01, val1, idx1, n23, NULL, final_idx, asc);
@@ -968,27 +968,27 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                uint32_t n01 = cctx.real_count[0] + cctx.real_count[1];
                uint32_t n23 = cctx.real_count[2] + cctx.real_count[3];

-                const float *   val01 = (const float *)   (vtcm_base + kparams->merge_values_off);
-                const int32_t * idx01 = (const int32_t *) (vtcm_base + kparams->merge_indices_off);
-                const float *   val23 = (const float *)   (vtcm_base + kparams->merge_values_off)  + 2 * chunk_elems;
-                const int32_t * idx23 = (const int32_t *) (vtcm_base + kparams->merge_indices_off) + 2 * chunk_elems;
+                const float *   val01 = merge_values;
+                const int32_t * idx01 = merge_indices;
+                const float *   val23 = merge_values  + 2 * chunk_elems;
+                const int32_t * idx23 = merge_indices + 2 * chunk_elems;

-                int32_t * final_idx = (int32_t *) vtcm_base;
+                int32_t * final_idx = (int32_t *) phase1_base;
                struct htp_thread_trace * tr = &octx->ctx->trace[0];
                htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                merge_runs(val01, idx01, n01, val23, idx23, n23, NULL, final_idx, asc);
                htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
            } else if (kparams->n_chunks == 2) {
-                size_t slot_size = (size_t) kparams->phase1_slot_size;
+                size_t slot_size = layout.phase1_slot_size;
                uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;

-                const float *   val0 = (const float *)   (vtcm_base);
-                const int32_t * idx0 = (const int32_t *) (vtcm_base + chunk_elems * sizeof(float));
-                const float *   val1 = (const float *)   (vtcm_base + slot_size);
-                const int32_t * idx1 = (const int32_t *) (vtcm_base + slot_size + chunk_elems * sizeof(float));
+                const float *   val0 = (const float *)   (phase1_base);
+                const int32_t * idx0 = (const int32_t *) (phase1_base + chunk_elems * sizeof(float));
+                const float *   val1 = (const float *)   (phase1_base + slot_size);
+                const int32_t * idx1 = (const int32_t *) (phase1_base + slot_size + chunk_elems * sizeof(float));

-                int32_t * final_idx = (int32_t *) (vtcm_base + kparams->merge_indices_off);
+                int32_t * final_idx = merge_indices;
                struct htp_thread_trace * tr = &octx->ctx->trace[0];
                htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
                merge_runs(val0, idx0, cctx.real_count[0], val1, idx1, cctx.real_count[1], NULL, final_idx, asc);
@@ -996,7 +996,7 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
                dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
            } else {
                uint32_t chunk_elems = (uint32_t) kparams->chunk_elems;
-                int32_t * final_idx = (int32_t *) (vtcm_base + chunk_elems * sizeof(float));
+                int32_t * final_idx = (int32_t *) (phase1_base + chunk_elems * sizeof(float));
                dma_sync_write(octx->ctx->dma[0], dst->data, final_idx, ne00 * sizeof(int32_t));
            }
        }
@@ -1006,7 +1006,10 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
    struct htp_sort_multi_row_ctx mctx;
    mctx.octx             = octx;
    mctx.kparams          = kparams;
-    mctx.vtcm_base        = vtcm_base;
+    mctx.phase1_base      = phase1_base;
+    mctx.phase1_slot_size = layout.phase1_slot_size;
+    mctx.row_start        = row_start;
+    mctx.row_end          = row_end;

    uint32_t dst_row_bytes = (uint32_t) dst_row_size;
    uint32_t rows_per_chunk = 0;
@@ -1014,13 +1017,13 @@ static int op_sort_common(struct htp_ops_context * octx, bool is_top_k) {
    if (rows_per_chunk == 0) {
        rows_per_chunk = 1;
    }
-    uint32_t dr = (nrows + (uint32_t) kparams->n_threads - 1) / (uint32_t) kparams->n_threads;
+    uint32_t dr = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div);
    if (rows_per_chunk > 1) {
        dr = hex_round_up(dr, rows_per_chunk);
    }
    mctx.nrows_per_thread = dr;

-    work_queue_run(octx->ctx->work_queue, htp_sort_multi_row_job, &mctx, (uint32_t) kparams->n_threads);
+    work_queue_run(octx->ctx->work_queue, htp_sort_multi_row_job, &mctx, octx->n_threads);

    return HTP_STATUS_OK;
 }
@@ -7,33 +7,27 @@
 #include <string.h>

 #include "hex-fastdiv.h"
+#include "htp-vtcm.h"

 struct htp_sort_kernel_params {
-    int32_t  n_threads;
-    int32_t  total_rows;
-    int32_t  row_start;
-    int32_t  row_end;
-    int32_t  ne00;
-    int32_t  k;
-    int32_t  order;             // GGML_SORT_ORDER_ASC (0) or GGML_SORT_ORDER_DESC (1)
-    int32_t  is_top_k;          // 1 if TOP_K, 0 if ARGSORT
-    int32_t  use_dma;           // 1 if DMA enabled
-    int32_t  chunk_elems;
-    int32_t  n_chunks;
-    int32_t  vtcm_size;
-    int32_t  phase1_slot_size;
-    int32_t  merge_values_off;
-    int32_t  merge_indices_off;
-    int32_t  merge_elems;
-    int32_t  n_slots;
-    int32_t  pad[15];
+    uint32_t n_threads;
+    uint32_t order;             // GGML_SORT_ORDER_ASC (0) or GGML_SORT_ORDER_DESC (1)
+    uint32_t is_top_k;          // 1 if TOP_K, 0 if ARGSORT
+    uint32_t chunk_elems;
+    uint32_t n_chunks;
+    uint32_t vtcm_size;
+    uint32_t merge_elems;
+    uint32_t n_slots;
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne01_ne02;
 };

 struct htp_sort_vtcm_layout {
    size_t   total_bytes;
    size_t   phase1_slot_size;
-    size_t   merge_values_off;
-    size_t   merge_indices_off;
+    size_t   off_phase1;
+    size_t   off_merge_values;
+    size_t   off_merge_indices;
    uint32_t chunk_elems;
    uint32_t n_chunks;
    uint32_t merge_elems;
@@ -41,6 +35,39 @@ struct htp_sort_vtcm_layout {
    uint32_t n_slots;
 };

+static inline void htp_sort_vtcm_layout_build(
+    struct htp_sort_vtcm_layout * L,
+    uint32_t total_rows,
+    uint32_t chunk_elems,
+    uint32_t n_chunks,
+    uint32_t merge_elems,
+    uint32_t n_threads,
+    uint32_t n_slots) {
+
+    memset(L, 0, sizeof(*L));
+    L->chunk_elems = chunk_elems;
+    L->n_chunks    = n_chunks;
+    L->merge_elems = merge_elems;
+    L->n_threads   = n_threads;
+    L->n_slots     = n_slots;
+
+    const size_t values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
+    const size_t indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
+    L->phase1_slot_size       = ((values_size + indices_size) + 255) & ~255;
+
+    size_t off = 0;
+    if (total_rows > 1) {
+        VTCM_LAYOUT_ALLOC(off, off_phase1, L->phase1_slot_size * n_slots * n_threads);
+    } else {
+        VTCM_LAYOUT_ALLOC(off, off_phase1, L->phase1_slot_size * n_chunks * n_slots);
+        const size_t merge_values_size  = ((merge_elems * sizeof(float)) + 127) & ~127;
+        const size_t merge_indices_size = ((merge_elems * sizeof(int32_t)) + 127) & ~127;
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_merge_values,  merge_values_size,  n_chunks > 1);
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_merge_indices, merge_indices_size, n_chunks > 1);
+    }
+    L->total_bytes = off;
+}
+
 static inline bool htp_sort_solve_layout(
    struct htp_sort_vtcm_layout * layout,
    uint32_t ne00,
@@ -75,18 +102,8 @@ static inline bool htp_sort_solve_layout(
            n_threads--;
        }

-        size_t total_bytes = spad_per_thread * n_threads;
-        if (total_bytes > vtcm_budget) {
-            return false;
-        }
-
-        layout->total_bytes      = total_bytes;
-        layout->phase1_slot_size = spad_per_slot;
-        layout->chunk_elems      = ne00_padded;
-        layout->n_chunks         = 1;
-        layout->n_threads        = n_threads;
-        layout->n_slots          = n_slots;
-        return true;
+        htp_sort_vtcm_layout_build(layout, total_rows, ne00_padded, 1, 0, n_threads, n_slots);
+        return layout->total_bytes <= vtcm_budget;
    }

    uint32_t n_vec = (ne00 + 31) / 32;
@@ -102,17 +119,7 @@ static inline bool htp_sort_solve_layout(

    uint32_t chunk_n_vec = n_vec_pow2 / n_chunks;
    uint32_t chunk_elems = chunk_n_vec * 32;
-
-    size_t phase1_values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
-    size_t phase1_indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
-    size_t phase1_slot_size    = ((phase1_values_size + phase1_indices_size) + 255) & ~255;
-    size_t phase1_total_size   = phase1_slot_size * n_chunks;
-
-    size_t merge_values_size  = 0;
-    size_t merge_indices_size = 0;
-    size_t merge_values_off   = phase1_total_size;
-    size_t merge_indices_off  = 0;
-    uint32_t merge_elems      = 0;
+    uint32_t merge_elems = 0;

    if (n_chunks > 1) {
        if (is_top_k) {
@@ -125,47 +132,25 @@ static inline bool htp_sort_solve_layout(
        } else {
            merge_elems = n_vec_pow2 * 32;
        }
-        merge_values_size  = ((merge_elems * sizeof(float)) + 127) & ~127;
-        merge_indices_size = ((merge_elems * sizeof(int32_t)) + 127) & ~127;
-        merge_indices_off  = merge_values_off + merge_values_size;
    }

-    size_t total_bytes = phase1_total_size + merge_values_size + merge_indices_size;
+    htp_sort_vtcm_layout_build(layout, 1, chunk_elems, n_chunks, merge_elems, n_chunks, 1);

-    if (total_bytes > vtcm_budget && n_chunks > 1) {
-        n_chunks = 1;
+    if (layout->total_bytes > vtcm_budget && n_chunks > 1) {
+        n_chunks    = 1;
        chunk_elems = n_vec_pow2 * 32;
-        phase1_values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
-        phase1_indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
-        phase1_slot_size    = ((phase1_values_size + phase1_indices_size) + 255) & ~255;
-        phase1_total_size   = phase1_slot_size;
-        merge_values_size   = 0;
-        merge_indices_size  = 0;
-        merge_values_off    = phase1_total_size;
-        merge_indices_off   = 0;
-        merge_elems         = 0;
-        total_bytes         = phase1_total_size;
+        merge_elems = 0;
+        htp_sort_vtcm_layout_build(layout, 1, chunk_elems, n_chunks, merge_elems, 1, 1);
    }

-    if (total_bytes > vtcm_budget) {
+    if (layout->total_bytes > vtcm_budget) {
        return false;
    }

-    uint32_t n_slots = 1;
-    if (n_chunks == 1 && phase1_slot_size * 2 <= vtcm_budget) {
-        n_slots = 2;
-        total_bytes = phase1_slot_size * 2;
+    if (n_chunks == 1 && layout->phase1_slot_size * 2 <= vtcm_budget) {
+        htp_sort_vtcm_layout_build(layout, 1, chunk_elems, 1, 0, 1, 2);
    }

-    layout->total_bytes       = total_bytes;
-    layout->phase1_slot_size  = phase1_slot_size;
-    layout->merge_values_off  = merge_values_off;
-    layout->merge_indices_off = merge_indices_off;
-    layout->chunk_elems       = chunk_elems;
-    layout->n_chunks          = n_chunks;
-    layout->merge_elems       = merge_elems;
-    layout->n_threads         = n_chunks;
-    layout->n_slots           = n_slots;
    return true;
 }

@@ -42,6 +42,7 @@ struct htp_binary_context {
    uint32_t nrows_per_thread;
    uint32_t total_rows;
    uint32_t row_start;
+    uint32_t n_srcs;
    size_t   src0_row_size_aligned;
    size_t   src1_row_size_aligned;
    size_t   dst_row_size_aligned;
@@ -380,6 +381,7 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
    size_t dst_spad_half  = layout->dst_spad_half_size;

    dma_queue * dma_q = octx->ctx->dma[ith];
+    const uint32_t n_srcs = bctx->n_srcs;
    uint32_t ir_prefetch = start_row;
    int spad_idx = 0;

@@ -391,12 +393,7 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
        i02 = fastdiv(rem, &bctx->src0_dim1_div);
        i01 = rem - i02 * ne01;

-        uint32_t i13 = (ne13 == 1) ? 0 : i03;
-        uint32_t i12 = (ne12 == 1) ? 0 : i02;
-        uint32_t i11 = (ne11 == 1) ? 0 : i01;
-
        dma_addr_t src0_curr = src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01;
-        dma_addr_t src1_curr = src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11;
        dma_addr_t dst_curr  = dst->data  + i03 * nb3  + i02 * nb2  + i01 * nb1;

        uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half;
@@ -405,7 +402,15 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,

        dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0);
        dma_queue_push(dma_q, dma_make_data(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size);
-        dma_queue_push(dma_q, dma_make_data(s1_spad, src1_curr), bctx->src1_row_size_aligned, nb11, row_size_bytes, current_block_size);
+        for (uint32_t s = 1; s < n_srcs; s++) {
+            const struct htp_tensor * src_s = octx->src[s];
+            uint32_t is3 = (src_s->ne[3] == 1) ? 0 : i03;
+            uint32_t is2 = (src_s->ne[2] == 1) ? 0 : i02;
+            uint32_t is1 = (src_s->ne[1] == 1) ? 0 : i01;
+            dma_addr_t s_curr = src_s->data + is3 * src_s->nb[3] + is2 * src_s->nb[2] + is1 * src_s->nb[1];
+            uint8_t * s_spad = s1_spad + (s - 1) * 2 * src1_spad_half;
+            dma_queue_push(dma_q, dma_make_data(s_spad, s_curr), bctx->src1_row_size_aligned, src_s->nb[1], row_size_bytes, current_block_size);
+        }
        ir_prefetch += current_block_size;
        spad_idx ^= 1;
    }
@@ -422,6 +427,11 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
        compute(d_spad, s0_spad, s1_spad, current_block_size,
                bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, bctx->src1_row_size_aligned, ne00);
+        for (uint32_t s = 2; s < n_srcs; s++) {
+            uint8_t * ss_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
+            compute(d_spad, d_spad, ss_spad, current_block_size,
+                    bctx->dst_row_size_aligned, bctx->dst_row_size_aligned, bctx->src1_row_size_aligned, ne00);
+        }
        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);

        uint32_t i03, i02, i01, rem;
@@ -440,15 +450,18 @@ static void binary_thread_vector_same_shape(unsigned int nth, unsigned int ith,
            p02 = fastdiv(prem, &bctx->src0_dim1_div);
            p01 = prem - p02 * ne01;

-            uint32_t p13 = (ne13 == 1) ? 0 : p03;
-            uint32_t p12 = (ne12 == 1) ? 0 : p02;
-            uint32_t p11 = (ne11 == 1) ? 0 : p01;
-
            dma_addr_t s0_next = src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01;
-            dma_addr_t s1_next = src1->data + p13 * nb13 + p12 * nb12 + p11 * nb11;
-
            dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size);
-            dma_queue_push(dma_q, dma_make_data(s1_spad, s1_next), bctx->src1_row_size_aligned, nb11, row_size_bytes, next_block_size);
+
+            for (uint32_t s = 1; s < n_srcs; s++) {
+                const struct htp_tensor * src_s = octx->src[s];
+                uint32_t ps3 = (src_s->ne[3] == 1) ? 0 : p03;
+                uint32_t ps2 = (src_s->ne[2] == 1) ? 0 : p02;
+                uint32_t ps1 = (src_s->ne[1] == 1) ? 0 : p01;
+                dma_addr_t s_next = src_s->data + ps3 * src_s->nb[3] + ps2 * src_s->nb[2] + ps1 * src_s->nb[1];
+                uint8_t * s_spad = s1_spad + (s - 1) * 2 * src1_spad_half;
+                dma_queue_push(dma_q, dma_make_data(s_spad, s_next), bctx->src1_row_size_aligned, src_s->nb[1], row_size_bytes, next_block_size);
+            }

            ir_prefetch += next_block_size;
        }
@@ -859,6 +872,7 @@ struct binary_chunked_context {
    uint32_t                 chunk_size;
    uint32_t                 chunks_per_thread;
    uint32_t                 total_chunks;
+    uint32_t                 n_srcs;
    bool                     is_scalar;
    float                    scalar_f32;
    _Float16                 scalar_f16;
@@ -872,7 +886,6 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
    struct binary_chunked_context * ctx = (struct binary_chunked_context *) data;
    struct htp_ops_context * octx = ctx->octx;
    const struct htp_tensor * src0 = octx->src[0];
-    const struct htp_tensor * src1 = octx->src[1];
    const struct htp_tensor * dst  = octx->dst;

    const uint32_t start_chunk = ctx->chunks_per_thread * ith;
@@ -895,6 +908,7 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);

    dma_queue * dma_q = octx->ctx->dma[ith];
+    const uint32_t n_srcs = ctx->n_srcs;
    uint32_t prefetch_chunk = start_chunk;
    int spad_idx = 0;

@@ -913,9 +927,12 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
        dma_queue_push(dma_q, dma_make_data(d_curr, d_spad), chunk_bytes, chunk_bytes, cur_bytes, 0);
        dma_queue_push(dma_q, dma_make_data(s0_spad, s0_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
        if (!ctx->is_scalar) {
-            dma_addr_t s1_curr = src1->data + (size_t) c_start * elem_size;
            uint8_t * s1_spad = src1_spad_base + spad_idx * chunk_bytes;
-            dma_queue_push(dma_q, dma_make_data(s1_spad, s1_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
+            for (uint32_t s = 1; s < n_srcs; s++) {
+                dma_addr_t s_curr = octx->src[s]->data + (size_t) c_start * elem_size;
+                uint8_t * s_spad  = s1_spad + (s - 1) * 2 * chunk_bytes;
+                dma_queue_push(dma_q, dma_make_data(s_spad, s_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
+            }
        }

        prefetch_chunk++;
@@ -943,6 +960,10 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
            }
        } else {
            ctx->compute(d_spad, s0_spad, s1_spad, cur_elems);
+            for (uint32_t s = 2; s < n_srcs; s++) {
+                uint8_t * s_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
+                ctx->compute(d_spad, d_spad, s_spad, cur_elems);
+            }
        }
        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) c);

@@ -958,8 +979,11 @@ static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * dat
            dma_addr_t s0_next = src0->data + (size_t) pc_start * elem_size;
            dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), chunk_bytes, chunk_bytes, p_bytes, 1);
            if (!ctx->is_scalar) {
-                dma_addr_t s1_next = src1->data + (size_t) pc_start * elem_size;
-                dma_queue_push(dma_q, dma_make_data(s1_spad, s1_next), chunk_bytes, chunk_bytes, p_bytes, 1);
+                for (uint32_t s = 1; s < n_srcs; s++) {
+                    dma_addr_t s_next = octx->src[s]->data + (size_t) pc_start * elem_size;
+                    uint8_t * s_spad  = s1_spad + (s - 1) * 2 * chunk_bytes;
+                    dma_queue_push(dma_q, dma_make_data(s_spad, s_next), chunk_bytes, chunk_bytes, p_bytes, 1);
+                }
            }

            prefetch_chunk++;
@@ -974,6 +998,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
    const struct htp_tensor * src1 = octx->src[1];
    const struct htp_tensor * dst  = octx->dst;
    const struct htp_binary_kernel_params * kparams = (const struct htp_binary_kernel_params *) octx->kernel_params;
+    const uint32_t n_srcs = kparams->n_srcs > 2 ? kparams->n_srcs : 2;

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

@@ -985,7 +1010,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
    const size_t dst_row_size  = dst->ne[0]  * elem_size;

    if (kparams->kernel_type == HTP_BINARY_KERNEL_CHUNKED) {
-        const bool is_scalar = kparams->is_scalar || (src1->ne[0] == 1 && src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1);
+        const bool is_scalar = kparams->is_scalar;

        compute_binary_chunked_t compute = NULL;
        compute_binary_scalar_chunked_f32_t compute_scalar_f32 = NULL;
@@ -1045,12 +1070,16 @@ static int execute_op_binary(struct htp_ops_context * octx) {
        const uint32_t elems_per_line = (src0_type == HTP_TYPE_F32) ? 32 : 64;

        if (octx->ctx->mdev.count > 1) {
-            const bool can_split = htp_tensor_mdev_data_aligned(dst) &&
-                                   htp_tensor_mdev_data_aligned(src0) &&
-                                   (is_scalar || htp_tensor_mdev_data_aligned(src1)) &&
-                                   htp_tensor_is_contiguous(dst, elem_size) &&
-                                   htp_tensor_is_contiguous(src0, elem_size) &&
-                                   (is_scalar || htp_tensor_is_contiguous(src1, elem_size));
+            bool can_split = htp_tensor_mdev_data_aligned(dst) &&
+                             htp_tensor_mdev_data_aligned(src0) &&
+                             (is_scalar || htp_tensor_mdev_data_aligned(src1)) &&
+                             htp_tensor_is_contiguous(dst, elem_size) &&
+                             htp_tensor_is_contiguous(src0, elem_size) &&
+                             (is_scalar || htp_tensor_is_contiguous(src1, elem_size));
+            for (uint32_t s = 2; can_split && s < n_srcs; s++) {
+                can_split = htp_tensor_mdev_data_aligned(octx->src[s]) &&
+                            htp_tensor_is_contiguous(octx->src[s], elem_size);
+            }
            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
                total_elems, can_split ? elems_per_line : 0,
                octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
@@ -1072,7 +1101,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
            return HTP_STATUS_VTCM_TOO_SMALL;
        }

-        const uint32_t chunk_size = kparams->chunk_size > 0 ? kparams->chunk_size : (32768 / elem_size);
+        const uint32_t chunk_size = kparams->chunk_size;
        const uint32_t total_chunks = (nelem + chunk_size - 1) / chunk_size;
        const uint32_t n_threads = (total_chunks >= 2) ? MIN(octx->n_threads, total_chunks) : 1;
        const uint32_t chunks_per_thread = (total_chunks + n_threads - 1) / n_threads;
@@ -1101,6 +1130,7 @@ static int execute_op_binary(struct htp_ops_context * octx) {
            .chunk_size          = chunk_size,
            .chunks_per_thread   = chunks_per_thread,
            .total_chunks        = total_chunks,
+            .n_srcs              = n_srcs,
            .is_scalar           = is_scalar,
            .scalar_f32          = scalar_f32,
            .scalar_f16          = scalar_f16,
@@ -1146,65 +1176,6 @@ static int execute_op_binary(struct htp_ops_context * octx) {
        return HTP_STATUS_VTCM_TOO_SMALL;
    }

-    dma_queue * dma_q = octx->ctx->dma[0];
-    uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src1);
-    if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST) {
-        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1->ne[0] * elem_size, 1);
-    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR ||
-               kparams->kernel_type == HTP_BINARY_KERNEL_REPEAT ||
-               kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX) {
-        if (!kparams->src1_strided) {
-            const size_t src1_bytes = (size_t) src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3] * elem_size;
-            dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1_bytes, 1);
-        } else {
-            const uint32_t ne11 = src1->ne[1];
-            const uint32_t ne12 = src1->ne[2];
-            const uint32_t ne13 = src1->ne[3];
-            const size_t row_bytes = src1->ne[0] * elem_size;
-            for (uint32_t i13 = 0; i13 < ne13; i13++) {
-                for (uint32_t i12 = 0; i12 < ne12; i12++) {
-                    dma_addr_t src_addr = src1->data + i13 * src1->nb[3] + i12 * src1->nb[2];
-                    uint8_t * dst_ptr = vtcm_src1 + (i13 * ne12 * ne11 + i12 * ne11) * row_bytes;
-                    while (!dma_queue_push(dma_q, dma_make_data(dst_ptr, src_addr),
-                                           row_bytes, src1->nb[1], row_bytes, ne11)) {
-                        dma_queue_pop(dma_q);
-                    }
-                }
-            }
-        }
-    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
-        const struct htp_tensor * src2 = octx->src[2];
-        uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
-        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data),
-                       kparams->src1_row_size_aligned, src1->nb[1],
-                       src1->ne[0] * elem_size, src1->ne[1]);
-        const size_t src2_row_bytes = src2->ne[0] * sizeof(int32_t);
-        dma_queue_push(dma_q, dma_make_data(vtcm_src2, src2->data),
-                       src2_row_bytes, src2->nb[1],
-                       src2_row_bytes, src2->ne[1]);
-    }
-
-    bctx.octx                  = octx;
-    bctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
-    bctx.total_rows            = nrows;
-    bctx.row_start             = row_start;
-    bctx.block_max             = bctx.vtcm_layout.rows_per_buffer;
-    bctx.src0_row_size_aligned = src0_row_size_aligned;
-    bctx.src1_row_size_aligned = src1_row_size_aligned;
-    bctx.dst_row_size_aligned  = dst_row_size_aligned;
-    bctx.row_size_bytes        = src0_row_size;
-
-    bctx.src0_dim1_div  = init_fastdiv_values(src0->ne[1]);
-    bctx.src0_dim2_div  = init_fastdiv_values(src0->ne[2]);
-    bctx.src0_dim12_div = init_fastdiv_values(src0->ne[1] * src0->ne[2]);
-
-    bctx.src1_dim1_div  = init_fastdiv_values(src1->ne[1]);
-    bctx.src1_dim2_div  = init_fastdiv_values(src1->ne[2]);
-    bctx.src1_dim3_div  = init_fastdiv_values(src1->ne[3]);
-
-    bctx.split_at_ne01  = kparams->split_at_ne01;
-    bctx.split_at_ne02  = kparams->split_at_ne02;
-
    worker_callback_t worker_func = NULL;
    void * compute_func = NULL;

@@ -1322,7 +1293,70 @@ static int execute_op_binary(struct htp_ops_context * octx) {
        return HTP_STATUS_NO_SUPPORT;
    }

-    bctx.compute = compute_func;
+    dma_queue * dma_q = octx->ctx->dma[0];
+    uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src1);
+    if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST) {
+        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1->ne[0] * elem_size, 1);
+    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR ||
+               kparams->kernel_type == HTP_BINARY_KERNEL_REPEAT ||
+               kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX) {
+        if (!kparams->src1_strided) {
+            const size_t src1_bytes = (size_t) src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3] * elem_size;
+            dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1_bytes, 1);
+            if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR && octx->src[2] != NULL) {
+                uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
+                dma_queue_push(dma_q, dma_make_data(vtcm_src2, octx->src[2]->data), bctx.vtcm_layout.src2_size, 0, src1_bytes, 1);
+            }
+        } else {
+            const uint32_t ne11 = src1->ne[1];
+            const uint32_t ne12 = src1->ne[2];
+            const uint32_t ne13 = src1->ne[3];
+            const size_t row_bytes = src1->ne[0] * elem_size;
+            for (uint32_t i13 = 0; i13 < ne13; i13++) {
+                for (uint32_t i12 = 0; i12 < ne12; i12++) {
+                    dma_addr_t src_addr = src1->data + i13 * src1->nb[3] + i12 * src1->nb[2];
+                    uint8_t * dst_ptr = vtcm_src1 + (i13 * ne12 * ne11 + i12 * ne11) * row_bytes;
+                    while (!dma_queue_push(dma_q, dma_make_data(dst_ptr, src_addr),
+                                           row_bytes, src1->nb[1], row_bytes, ne11)) {
+                        dma_queue_pop(dma_q);
+                    }
+                }
+            }
+        }
+    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
+        const struct htp_tensor * src2 = octx->src[2];
+        uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
+        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data),
+                       kparams->src1_row_size_aligned, src1->nb[1],
+                       src1->ne[0] * elem_size, src1->ne[1]);
+        const size_t src2_row_bytes = src2->ne[0] * sizeof(int32_t);
+        dma_queue_push(dma_q, dma_make_data(vtcm_src2, src2->data),
+                       src2_row_bytes, src2->nb[1],
+                       src2_row_bytes, src2->ne[1]);
+    }
+
+    bctx.octx                  = octx;
+    bctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
+    bctx.total_rows            = nrows;
+    bctx.row_start             = row_start;
+    bctx.n_srcs                = n_srcs;
+    bctx.block_max             = bctx.vtcm_layout.rows_per_buffer;
+    bctx.src0_row_size_aligned = src0_row_size_aligned;
+    bctx.src1_row_size_aligned = src1_row_size_aligned;
+    bctx.dst_row_size_aligned  = dst_row_size_aligned;
+    bctx.row_size_bytes        = src0_row_size;
+
+    bctx.src0_dim1_div  = kparams->div_ne01;
+    bctx.src0_dim2_div  = kparams->div_ne02;
+    bctx.src0_dim12_div = kparams->div_ne012;
+
+    bctx.src1_dim1_div  = kparams->div_ne11;
+    bctx.src1_dim2_div  = kparams->div_ne12;
+    bctx.src1_dim3_div  = kparams->div_ne13;
+
+    bctx.split_at_ne01  = kparams->split_at_ne01;
+    bctx.split_at_ne02  = kparams->split_at_ne02;
+    bctx.compute        = compute_func;

    if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST ||
        kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR ||
@@ -1330,6 +1364,15 @@ static int execute_op_binary(struct htp_ops_context * octx) {
        kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX ||
        kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
        dma_queue_flush(dma_q);
+        if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR && octx->src[2] != NULL && octx->op == HTP_OP_MUL) {
+            uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
+            const uint32_t n_scalars = src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3];
+            if (src0_type == HTP_TYPE_F32) {
+                hvx_mul_f32_aaa(vtcm_src1, vtcm_src1, vtcm_src2, n_scalars);
+            } else if (src0_type == HTP_TYPE_F16) {
+                hvx_mul_f16_aaa(vtcm_src1, vtcm_src1, vtcm_src2, n_scalars);
+            }
+        }
    }

    work_queue_run(octx->ctx->work_queue, worker_func, &bctx, n_threads);
@@ -6,6 +6,7 @@
 #include <string.h>

 #include "hex-common.h"
+#include "hex-fastdiv.h"
 #include "htp-ops.h"
 #include "htp-vtcm.h"

@@ -20,7 +21,15 @@ enum htp_binary_kernel_type {
 };

 struct htp_binary_kernel_params {
-    uint32_t kernel_type;
+    uint8_t  kernel_type;
+    uint8_t  split_at_ne01;
+    uint8_t  split_at_ne02;
+    uint8_t  src1_strided;
+
+    uint8_t  n_srcs;
+    uint8_t  is_scalar;
+    uint8_t  pad[2];
+
    uint32_t n_threads;
    uint32_t rows_per_buffer;

@@ -34,12 +43,14 @@ struct htp_binary_kernel_params {

    uint32_t chunk_size;
    uint32_t chunk_bytes;
-    uint32_t is_scalar;

-    uint8_t  split_at_ne01;
-    uint8_t  split_at_ne02;
-    uint8_t  src1_strided;
-    uint8_t  pad;
+    struct fastdiv_values div_ne01;
+    struct fastdiv_values div_ne02;
+    struct fastdiv_values div_ne012;
+
+    struct fastdiv_values div_ne11;
+    struct fastdiv_values div_ne12;
+    struct fastdiv_values div_ne13;
 };

 #if defined(__cplusplus)
@@ -100,6 +111,8 @@ static inline void htp_binary_vtcm_layout_build(
        return;
    }

+    const uint32_t n_srcs = kparams->n_srcs > 2 ? kparams->n_srcs : 2;
+
    if (kparams->kernel_type == HTP_BINARY_KERNEL_CHUNKED) {
        const size_t chunk_bytes = kparams->chunk_bytes;
        if (chunk_bytes == 0) {
@@ -107,7 +120,7 @@ static inline void htp_binary_vtcm_layout_build(
        }

        L->src0_bytes_per_thread = 2 * chunk_bytes;
-        L->src1_bytes_per_thread = kparams->is_scalar ? 0 : (2 * chunk_bytes);
+        L->src1_bytes_per_thread = kparams->is_scalar ? 0 : (2 * (n_srcs - 1) * chunk_bytes);
        L->dst_bytes_per_thread  = 2 * chunk_bytes;

        L->src0_spad_half_size = chunk_bytes;
@@ -135,7 +148,7 @@ static inline void htp_binary_vtcm_layout_build(
    }

    const size_t spad_row_total = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
-        ? 2 * (kparams->src0_row_size_aligned + kparams->src1_row_size_aligned + kparams->dst_row_size_aligned)
+        ? 2 * (kparams->src0_row_size_aligned + (n_srcs - 1) * kparams->src1_row_size_aligned + kparams->dst_row_size_aligned)
        : 2 * (kparams->src0_row_size_aligned + kparams->dst_row_size_aligned);

    const size_t static_vtcm_size = kparams->src1_size + kparams->src2_size;
@@ -143,7 +156,9 @@ static inline void htp_binary_vtcm_layout_build(
        return;
    }

-    const size_t rows_per_buffer = (vtcm_size - static_vtcm_size) / (n_threads * spad_row_total);
+    const size_t rows_per_buffer = (kparams->rows_per_buffer > 0)
+        ? kparams->rows_per_buffer
+        : (vtcm_size - static_vtcm_size) / (n_threads * spad_row_total);
    if (rows_per_buffer == 0) {
        return;
    }
@@ -155,11 +170,13 @@ static inline void htp_binary_vtcm_layout_build(
    L->src0_bytes_per_thread = rows_per_buffer * 2 * kparams->src0_row_size_aligned;
    L->dst_bytes_per_thread  = rows_per_buffer * 2 * kparams->dst_row_size_aligned;
    L->src1_bytes_per_thread = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
-        ? rows_per_buffer * 2 * kparams->src1_row_size_aligned
+        ? rows_per_buffer * 2 * (n_srcs - 1) * kparams->src1_row_size_aligned
        : 0;

    L->src0_spad_half_size = L->src0_bytes_per_thread / 2;
-    L->src1_spad_half_size = L->src1_bytes_per_thread / 2;
+    L->src1_spad_half_size = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
+        ? rows_per_buffer * kparams->src1_row_size_aligned
+        : 0;
    L->dst_spad_half_size  = L->dst_bytes_per_thread / 2;

    const size_t src0_total = n_threads * L->src0_bytes_per_thread;
@@ -178,4 +195,53 @@ static inline void htp_binary_vtcm_layout_build(
    L->total_bytes = off;
 }

+static inline bool htp_binary_solve_chunked(
+    struct htp_binary_vtcm_layout * L,
+    struct htp_binary_kernel_params * kparams,
+    uint32_t total_elems,
+    uint32_t elem_size,
+    uint32_t max_threads,
+    uint32_t n_srcs,
+    bool is_scalar,
+    size_t vtcm_size
+) {
+    const size_t n_bufs = is_scalar ? 2 : (size_t) n_srcs + 1;
+    const size_t scalar_bytes = is_scalar ? 128 : 0;
+    const size_t budget_per_buf = (max_threads > 0 && vtcm_size > scalar_bytes)
+        ? hex_align_down((vtcm_size - scalar_bytes) / (2 * max_threads * n_bufs), 128)
+        : 0;
+    const uint32_t max_chunk_bytes = (budget_per_buf > 0 && budget_per_buf < 32768) ? (uint32_t) budget_per_buf : 32768u;
+    const uint32_t max_chunk_elems = max_chunk_bytes / elem_size;
+    if (max_chunk_elems == 0) {
+        return false;
+    }
+
+    const uint32_t min_chunk_elems = 256 < max_chunk_elems ? 256 : max_chunk_elems;
+    const uint32_t target_chunk_elems = hex_round_up((total_elems + (2 * max_threads) - 1) / (2 * max_threads), 32);
+    const uint32_t chunk_size  = (target_chunk_elems < min_chunk_elems) ? min_chunk_elems
+                               : (target_chunk_elems > max_chunk_elems) ? max_chunk_elems
+                               : target_chunk_elems;
+    const uint32_t chunk_bytes  = hex_round_up(chunk_size * elem_size, 128);
+    const uint32_t total_chunks = (total_elems + chunk_size - 1) / chunk_size;
+    const uint32_t n_threads    = (total_chunks >= 2) ? (max_threads < total_chunks ? max_threads : total_chunks) : 1;
+
+    kparams->kernel_type           = HTP_BINARY_KERNEL_CHUNKED;
+    kparams->n_threads             = n_threads;
+    kparams->n_srcs                = (uint8_t) n_srcs;
+    kparams->rows_per_buffer       = 1;
+    kparams->src1_row_size_aligned = is_scalar ? 0 : kparams->src0_row_size_aligned;
+    kparams->src1_size             = 0;
+    kparams->chunk_size            = chunk_size;
+    kparams->chunk_bytes           = chunk_bytes;
+    kparams->is_scalar             = is_scalar ? 1 : 0;
+
+    htp_binary_vtcm_layout_build(L, kparams, vtcm_size);
+    if (L->total_bytes == 0 || L->total_bytes > vtcm_size) {
+        return false;
+    }
+
+    kparams->vtcm_size = (uint32_t) L->total_bytes;
+    return true;
+}
+
 #endif
@@ -294,22 +294,21 @@ static int concat_transposed(struct htp_ops_context * octx, const struct htp_con

    const uint32_t n_threads = octx->n_threads;

-    // layout precomputed on host; kept for reference:
-    // struct htp_concat_transposed_vtcm_layout layout;
-    // htp_concat_transposed_vtcm_layout_build(&layout, octx->src[0]->ne[0], octx->src[1]->ne[0], type_size, n_threads);
+    struct htp_concat_transposed_vtcm_layout layout;
+    htp_concat_transposed_vtcm_layout_build(&layout, octx->src[0]->ne[0], octx->src[1]->ne[0], type_size, n_threads);

    uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;

    struct htp_concat_context cctx;
    cctx.octx                  = octx;
-    cctx.spad0_base            = vtcm_base;
-    cctx.spad1_base            = vtcm_base + n_threads * kparams->u.transposed.spad0_size_per_thread;
-    cctx.spad0_size_per_thread = kparams->u.transposed.spad0_size_per_thread;
-    cctx.spad1_size_per_thread = kparams->u.transposed.spad1_size_per_thread;
+    cctx.spad0_base            = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_spad0);
+    cctx.spad1_base            = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, layout.off_spad1);
+    cctx.spad0_size_per_thread = layout.src0_spad_size_per_thread;
+    cctx.spad1_size_per_thread = layout.src1_spad_size_per_thread;
    cctx.row_start             = row_start;
    cctx.nrows                 = nrows;
    cctx.nplanes               = dst->ne[2] * dst->ne[3];
-    cctx.div_ne2               = init_fastdiv_values(dst->ne[2]);
+    cctx.div_ne2               = kparams->u.transposed.div_ne2;
    cctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);

    work_queue_func_t worker_func = (type_size == 4) ? concat_2d_f32_transposed : concat_2d_f16_transposed;
@@ -2,6 +2,8 @@
 #define HTP_CONCAT_OPS_H

 #include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 #include "gather-rows.h"
 #include <stdint.h>

@@ -13,8 +15,7 @@ enum htp_concat_kernel_type {
 };

 struct htp_concat_transposed_params {
-    uint32_t spad0_size_per_thread;
-    uint32_t spad1_size_per_thread;
+    struct fastdiv_values div_ne2;
 };

 struct htp_concat_kernel_params {
@@ -40,11 +41,13 @@ _Static_assert(sizeof(struct htp_concat_kernel_params) <= 128, "htp_concat_kerne
 struct htp_concat_transposed_vtcm_layout {
    uint32_t src0_spad_size_per_thread;
    uint32_t src1_spad_size_per_thread;
+    size_t   off_spad0;
+    size_t   off_spad1;
    uint32_t total_bytes;
 };

 static inline void htp_concat_transposed_vtcm_layout_build(
-    struct htp_concat_transposed_vtcm_layout * layout,
+    struct htp_concat_transposed_vtcm_layout * L,
    uint32_t src0_ne0,
    uint32_t src1_ne0,
    uint32_t type_size,
@@ -55,9 +58,13 @@ static inline void htp_concat_transposed_vtcm_layout_build(
    uint32_t src1_ne0_padded = hex_round_up(src1_ne0, block_i);
    uint32_t spad0_row_bytes = hex_round_up(src0_ne0 * type_size, 128) + src1_ne0_padded * type_size;

-    layout->src0_spad_size_per_thread = block_i * spad0_row_bytes;
-    layout->src1_spad_size_per_thread = src1_ne0_padded * spad1_stride;
-    layout->total_bytes = n_threads * (layout->src0_spad_size_per_thread + layout->src1_spad_size_per_thread);
+    L->src0_spad_size_per_thread = block_i * spad0_row_bytes;
+    L->src1_spad_size_per_thread = src1_ne0_padded * spad1_stride;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_spad0, (size_t) n_threads * L->src0_spad_size_per_thread);
+    VTCM_LAYOUT_ALLOC(off, off_spad1, (size_t) n_threads * L->src1_spad_size_per_thread);
+    L->total_bytes = (uint32_t) off;
 }

 #endif // HTP_CONCAT_OPS_H
@@ -32,6 +32,7 @@ struct htp_copy_context {
    uint32_t nelem;
    uint32_t elem_per_thread;

+    struct htp_copy_convert_vtcm_layout layout;
    uint8_t * vtcm_src0;
    uint8_t * vtcm_dst;
 };
@@ -176,13 +177,14 @@ static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, vo
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                   \
                                                                                             \
    const struct htp_copy_convert_params * cvt = &ct->kparams->u.convert;                    \
-    const uint32_t src0_buf_size = cvt->src0_buf_size;                                       \
-    const uint32_t dst_buf_size  = cvt->dst_buf_size;                                        \
-    const uint32_t src0_row_stride = cvt->src0_row_stride;                                   \
-    const uint32_t dst_row_stride  = cvt->dst_row_stride;                                    \
-    const uint32_t blk_rows        = cvt->blk_rows;                                          \
-    uint8_t * vtcm_src0_base = ct->vtcm_src0 + ith * cvt->spad0_size_per_thread;             \
-    uint8_t * vtcm_dst_base  = ct->vtcm_dst  + ith * cvt->spad1_size_per_thread;             \
+    const struct htp_copy_convert_vtcm_layout * L = &ct->layout;                             \
+    const uint32_t src0_buf_size   = L->src0_buf_size;                                       \
+    const uint32_t dst_buf_size    = L->dst_buf_size;                                        \
+    const uint32_t src0_row_stride = L->src0_row_stride;                                     \
+    const uint32_t dst_row_stride  = L->dst_row_stride;                                      \
+    const uint32_t blk_rows        = L->blk_rows;                                            \
+    uint8_t * vtcm_src0_base = ct->vtcm_src0 + ith * L->spad0_size_per_thread;               \
+    uint8_t * vtcm_dst_base  = ct->vtcm_dst  + ith * L->spad1_size_per_thread;               \
    const uint32_t src0_row_size = ne00 * ct->kparams->src0_type_size;                       \
    const uint32_t dst_row_size  = ne00 * ct->kparams->dst_type_size;                        \
                                                                                             \
@@ -254,17 +256,16 @@ static int cpy_scalar(struct htp_ops_context * octx, const struct htp_copy_kerne
    }

    dma_queue * dma_q = octx->ctx->dma[0];
-    dma_addr_t s_vtcm = (dma_addr_t)(uintptr_t) octx->ctx->vtcm_base;
-    dma_addr_t d_vtcm = s_vtcm + VLEN;
+    uint8_t * s_ptr = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, 0);
+    uint8_t * d_ptr = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, VLEN);
+    dma_addr_t s_vtcm = (dma_addr_t)(uintptr_t) s_ptr;
+    dma_addr_t d_vtcm = (dma_addr_t)(uintptr_t) d_ptr;
    const uint32_t s_size = kparams->src0_type_size;
    const uint32_t d_size = kparams->dst_type_size;

    dma_queue_push(dma_q, dma_make_data(s_vtcm, src0->data), s_size, s_size, s_size, 1);
    dma_queue_pop(dma_q);

-    uint8_t * s_ptr = (uint8_t *) octx->ctx->vtcm_base;
-    uint8_t * d_ptr = s_ptr + VLEN;
-
    const HVX_Vector v_src = hvx_vmem(s_ptr);
    HVX_Vector v_dst;

@@ -366,6 +367,8 @@ static int cpy_sameshape_convert(struct htp_ops_context * octx, const struct htp
    const struct htp_copy_convert_params * cvt = &kparams->u.convert;

    struct htp_copy_context ct;
+    htp_copy_convert_vtcm_layout_build(&ct.layout, src0->ne[0], kparams->src0_type_size, kparams->dst_type_size, cvt->blk_rows, n_threads);
+
    ct.octx                  = octx;
    ct.kparams               = kparams;
    ct.row_start             = row_start;
@@ -373,8 +376,8 @@ static int cpy_sameshape_convert(struct htp_ops_context * octx, const struct htp
    ct.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);

    uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
-    ct.vtcm_src0 = vtcm_base;
-    ct.vtcm_dst  = vtcm_base + (size_t) n_threads * cvt->spad0_size_per_thread;
+    ct.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, ct.layout.off_src0);
+    ct.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, ct.layout.off_dst);

    work_queue_func_t copy_fun = NULL;
    if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32) {
@@ -3,6 +3,7 @@

 #include "hex-common.h"
 #include "hex-fastdiv.h"
+#include "htp-vtcm.h"
 #include "gather-rows.h"
 #include <stdint.h>

@@ -17,12 +18,6 @@ enum htp_copy_kernel_type {
 };

 struct htp_copy_convert_params {
-    uint32_t              src0_buf_size;
-    uint32_t              dst_buf_size;
-    uint32_t              spad0_size_per_thread;
-    uint32_t              spad1_size_per_thread;
-    uint32_t              src0_row_stride;  // row pitch in the VTCM buffers
-    uint32_t              dst_row_stride;
    uint32_t              blk_rows;         // rows moved per DMA descriptor
    struct fastdiv_values div_ne01;
    struct fastdiv_values div_ne02_ne01;
@@ -60,6 +55,9 @@ struct htp_copy_kernel_params {
 #define HTP_COPY_VTCM_BLK_BYTES (32 * 1024)

 struct htp_copy_convert_vtcm_layout {
+    size_t   total_bytes;
+    size_t   off_src0;
+    size_t   off_dst;
    uint32_t src0_row_stride;
    uint32_t dst_row_stride;
    uint32_t blk_rows;
@@ -67,11 +65,35 @@ struct htp_copy_convert_vtcm_layout {
    uint32_t dst_buf_size;
    uint32_t spad0_size_per_thread;
    uint32_t spad1_size_per_thread;
-    uint32_t total_bytes;
 };

-// ne01 bounds a block: rows of one block share the dim-1 stride.
 static inline void htp_copy_convert_vtcm_layout_build(
+    struct htp_copy_convert_vtcm_layout * L,
+    uint32_t ne00,
+    uint32_t src_type_size,
+    uint32_t dst_type_size,
+    uint32_t blk_rows,
+    uint32_t n_threads) {
+
+    const uint32_t src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
+    const uint32_t dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);
+
+    L->src0_row_stride       = src0_row_stride;
+    L->dst_row_stride        = dst_row_stride;
+    L->blk_rows              = blk_rows;
+    L->src0_buf_size         = blk_rows * src0_row_stride;
+    L->dst_buf_size          = blk_rows * dst_row_stride;
+    L->spad0_size_per_thread = 2 * L->src0_buf_size;
+    L->spad1_size_per_thread = 2 * L->dst_buf_size;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, (size_t) n_threads * L->spad0_size_per_thread);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  (size_t) n_threads * L->spad1_size_per_thread);
+    L->total_bytes = off;
+}
+
+// ne01 bounds a block: rows of one block share the dim-1 stride.
+static inline void htp_copy_convert_solve_layout(
    struct htp_copy_convert_vtcm_layout * layout,
    uint32_t ne00,
    uint32_t ne01,
@@ -79,10 +101,10 @@ static inline void htp_copy_convert_vtcm_layout_build(
    uint32_t dst_type_size,
    uint32_t n_threads) {

-    layout->src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
-    layout->dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);
+    const uint32_t src0_row_stride = hex_round_up(ne00 * src_type_size, 256);
+    const uint32_t dst_row_stride  = hex_round_up(ne00 * dst_type_size, 256);

-    const uint32_t max_row = layout->src0_row_stride > layout->dst_row_stride ? layout->src0_row_stride : layout->dst_row_stride;
+    const uint32_t max_row = src0_row_stride > dst_row_stride ? src0_row_stride : dst_row_stride;
    uint32_t blk_rows = HTP_COPY_VTCM_BLK_BYTES / max_row;
    if (blk_rows > ne01) {
        blk_rows = ne01;
@@ -90,13 +112,8 @@ static inline void htp_copy_convert_vtcm_layout_build(
    if (blk_rows == 0) {
        blk_rows = 1;
    }
-    layout->blk_rows = blk_rows;

-    layout->src0_buf_size = blk_rows * layout->src0_row_stride;
-    layout->dst_buf_size  = blk_rows * layout->dst_row_stride;
-    layout->spad0_size_per_thread = 2 * layout->src0_buf_size;
-    layout->spad1_size_per_thread = 2 * layout->dst_buf_size;
-    layout->total_bytes = n_threads * (layout->spad0_size_per_thread + layout->spad1_size_per_thread);
+    htp_copy_convert_vtcm_layout_build(layout, ne00, src_type_size, dst_type_size, blk_rows, n_threads);
 }

 #if defined(__cplusplus)
@@ -12,9 +12,11 @@
 #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 "cumsum-ops.h"

 #define htp_cumsum_tensors_preamble                         \
    const struct htp_tensor * restrict src0 = octx->src[0]; \
@@ -89,8 +91,8 @@ static inline HVX_Vector hvx_splat_last_f32(HVX_Vector v) {
 }

 static inline void hvx_cumsum_row_f32(const float * restrict src, float * restrict dst, uint32_t n) {
-    const uint32_t nvec = n / VLEN_FP32;
-    const uint32_t nloe = n % VLEN_FP32;
+    const uint32_t nvec = n >> 5;
+    const uint32_t nloe = n & (VLEN_FP32 - 1);

    HVX_Vector carry = Q6_V_vsplat_R(0);

@@ -210,16 +212,19 @@ static void cumsum_thread_f32(unsigned int nth, unsigned int ith, void * data) {
 int op_cumsum_f32(struct htp_ops_context * octx) {
    const struct htp_tensor * src0 = octx->src[0];
    const struct htp_tensor * dst  = octx->dst;
+    const struct htp_cumsum_kernel_params * kparams =
+        (const struct htp_cumsum_kernel_params *) octx->kernel_params;

-    const uint32_t total_rows      = src0->ne[1] * src0->ne[2] * src0->ne[3];
-    const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
+    const uint32_t total_rows   = kparams->total_rows;
+    const size_t   src_row_size = src0->nb[1];
+    const size_t   dst_row_size = dst->nb[1];

    uint32_t row_start = 0;
    uint32_t nrows     = total_rows;

    if (octx->ctx->mdev.count > 1) {
        uint32_t rows_per_chunk = 0;
-        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_data_row_size, &rows_per_chunk);
+        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_row_size, &rows_per_chunk);
        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
        row_start = range.start;
        nrows     = range.count;
@@ -229,43 +234,45 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
    const uint32_t n_threads = octx->n_threads;

-    const size_t src_row_size         = src0->nb[1];
-    const size_t dst_row_size         = dst->nb[1];
-    const size_t src_row_size_aligned = hex_round_up(src_row_size, VLEN);
-    const size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN);
+    if (!kparams->use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
+        return HTP_STATUS_NO_SUPPORT;
+    }

-    // 2 ping-pong buffers per thread for src and dst
-    const size_t spad_per_thread = 2 * (src_row_size_aligned + dst_row_size_aligned);
+    struct htp_cumsum_vtcm_layout layout = { 0 };
+    if (kparams->use_dma) {
+        htp_cumsum_vtcm_layout_build(&layout, src_row_size, dst_row_size, n_threads);

-    octx->src0_spad.size_per_thread = src_row_size_aligned * 2;
-    octx->dst_spad.size_per_thread  = dst_row_size_aligned * 2;
-
-    octx->src0_spad.size  = n_threads * octx->src0_spad.size_per_thread;
-    octx->dst_spad.size   = n_threads * octx->dst_spad.size_per_thread;
-
-    octx->src0_spad.data  = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
-    octx->dst_spad.data   = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src  = NULL;
+        octx->src0_spad.size_per_thread = layout.src0_bytes_per_thread;
+        octx->dst_spad.size_per_thread  = layout.dst_bytes_per_thread;
+        octx->src0_spad.size = n_threads * layout.src0_bytes_per_thread;
+        octx->dst_spad.size  = n_threads * layout.dst_bytes_per_thread;
+        octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_src0);
+        octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_dst);
+        octx->src0_spad.src  = NULL;
+        octx->dst_spad.src   = NULL;
+    }

    struct htp_cumsum_context cctx = {
        .octx                 = octx,
        .src_row_size         = src_row_size,
        .dst_row_size         = dst_row_size,
-        .src_row_size_aligned = src_row_size_aligned,
-        .dst_row_size_aligned = dst_row_size_aligned,
+        .src_row_size_aligned = layout.src_row_size_aligned,
+        .dst_row_size_aligned = layout.dst_row_size_aligned,
        .rows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
        .total_rows           = nrows,
        .row_start            = row_start,
    };

-    if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
-        if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
-            return HTP_STATUS_NO_SUPPORT;
-        }
-        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32, &cctx, n_threads);
-    } else {
+    if (kparams->use_dma) {
        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32_dma, &cctx, n_threads);
+    } else {
+        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32, &cctx, n_threads);
    }

    return HTP_STATUS_OK;
@@ -0,0 +1,49 @@
+#ifndef HTP_CUMSUM_OPS_H
+#define HTP_CUMSUM_OPS_H
+
+#include "hex-common.h"
+#include "htp-vtcm.h"
+#include <stdint.h>
+
+struct htp_cumsum_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_dst;
+    size_t src_row_size_aligned;
+    size_t dst_row_size_aligned;
+    size_t src0_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+};
+
+static inline void htp_cumsum_vtcm_layout_build(
+    struct htp_cumsum_vtcm_layout * L,
+    size_t                          src_row_size,
+    size_t                          dst_row_size,
+    uint32_t                        n_threads
+) {
+    L->src_row_size_aligned  = hex_round_up(src_row_size, 128);
+    L->dst_row_size_aligned  = hex_round_up(dst_row_size, 128);
+    L->src0_bytes_per_thread = 2 * L->src_row_size_aligned;
+    L->dst_bytes_per_thread  = 2 * L->dst_row_size_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread * n_threads);
+
+    L->total_bytes = off;
+}
+
+struct htp_cumsum_kernel_params {
+    uint32_t n_threads;
+    uint32_t total_rows;
+    uint32_t use_dma;
+    uint32_t vtcm_size;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_cumsum_kernel_params) <= 128, "htp_cumsum_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_cumsum_kernel_params) <= 128, "htp_cumsum_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_CUMSUM_OPS_H
@@ -10,10 +10,12 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "htp-vtcm.h"
 #include "hvx-types.h"
 #include "hex-utils.h"
 #include "hvx-copy.h"
 #include "dma-queue.h"
+#include "diag-ops.h"

 #define htp_diag_tensors_preamble                           \
    const struct htp_tensor * restrict src0 = octx->src[0]; \
@@ -33,6 +35,7 @@

 struct htp_diag_context {
    struct htp_ops_context * octx;
+    const struct htp_diag_kernel_params * kparams;
    size_t          src_batch_size;
    size_t          dst_row_size;
    size_t          src_batch_size_aligned;
@@ -83,8 +86,8 @@ static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data)
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];

    for (uint32_t ib = ib0; ib < ib1; ib++) {
-        const uint32_t i3 = ib / ne02;
-        const uint32_t i2 = ib % ne02;
+        const uint32_t i3 = fastdiv(ib, &dctx->kparams->div_ne02);
+        const uint32_t i2 = ib - i3 * ne02;

        const dma_addr_t src_batch = src_data + i3 * nb03 + i2 * nb02;

@@ -134,8 +137,8 @@ static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ib0);

    for (uint32_t ib = ib0; ib < ib1; ib++) {
-        const uint32_t i3 = ib / ne02;
-        const uint32_t i2 = ib % ne02;
+        const uint32_t i3 = fastdiv(ib, &dctx->kparams->div_ne02);
+        const uint32_t i2 = ib - i3 * ne02;

        const float * restrict src_batch = (const float *)(src_data + i3 * nb03 + i2 * nb02);

@@ -155,29 +158,17 @@ static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
 int op_diag_f32(struct htp_ops_context * octx) {
    const struct htp_tensor * src0 = octx->src[0];
    const struct htp_tensor * dst  = octx->dst;
+    const struct htp_diag_kernel_params * kparams =
+        (const struct htp_diag_kernel_params *) octx->kernel_params;

-    const uint32_t total_batches = src0->ne[2] * src0->ne[3];
-    const size_t dst_batch_size  = dst->ne[1] * dst->nb[1];
+    const uint32_t total_batches = kparams->total_batches;

    uint32_t batch_start = 0;
    uint32_t nbatches    = total_batches;

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

-    const uint32_t n_threads     = octx->n_threads;
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nbatches))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }

-    const size_t src_batch_size         = src0->ne[0] * sizeof(float);
-    const size_t dst_row_size           = dst->ne[0] * sizeof(float);
-    const size_t src_batch_size_aligned = hex_round_up(src_batch_size, VLEN);
-    const size_t dst_row_size_aligned   = hex_round_up(dst_row_size, VLEN);
+    const uint32_t n_threads = octx->n_threads;

-    // 1 src buffer + 1 dst row buffer per thread
-    const size_t spad_per_thread = src_batch_size_aligned + dst_row_size_aligned;
+    if (!kparams->use_dma && (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst))) {
+        return HTP_STATUS_NO_SUPPORT;
+    }

-    octx->src0_spad.size_per_thread = src_batch_size_aligned;
-    octx->dst_spad.size_per_thread  = dst_row_size_aligned;
+    const size_t src_batch_size = src0->ne[0] * sizeof(float);
+    const size_t dst_row_size   = dst->ne[0] * sizeof(float);

-    octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread;
-    octx->dst_spad.size  = n_threads * octx->dst_spad.size_per_thread;
+    struct htp_diag_vtcm_layout layout = { 0 };
+    if (kparams->use_dma) {
+        htp_diag_vtcm_layout_build(&layout, src_batch_size, dst_row_size, n_threads);

-    octx->src0_spad.data = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
-    octx->dst_spad.data  = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src  = NULL;
+        octx->src0_spad.size_per_thread = layout.src0_bytes_per_thread;
+        octx->dst_spad.size_per_thread  = layout.dst_bytes_per_thread;
+        octx->src0_spad.size = n_threads * layout.src0_bytes_per_thread;
+        octx->dst_spad.size  = n_threads * layout.dst_bytes_per_thread;
+        octx->src0_spad.data = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_src0);
+        octx->dst_spad.data  = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, layout.off_dst);
+        octx->src0_spad.src  = NULL;
+        octx->dst_spad.src   = NULL;
+    }

    struct htp_diag_context dctx = {
        .octx                   = octx,
+        .kparams                = kparams,
        .src_batch_size         = src_batch_size,
        .dst_row_size           = dst_row_size,
-        .src_batch_size_aligned = src_batch_size_aligned,
-        .dst_row_size_aligned   = dst_row_size_aligned,
+        .src_batch_size_aligned = layout.src_batch_size_aligned,
+        .dst_row_size_aligned   = layout.dst_row_size_aligned,
        .batches_per_thread     = fastdiv(nbatches + n_threads - 1, &octx->n_threads_div),
        .total_batches          = nbatches,
        .batch_start            = batch_start,
    };

-    if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
-        if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
-            return HTP_STATUS_NO_SUPPORT;
-        }
-        work_queue_run(octx->ctx->work_queue, diag_thread_f32, &dctx, n_threads);
-    } else {
+    if (kparams->use_dma) {
        work_queue_run(octx->ctx->work_queue, diag_thread_f32_dma, &dctx, n_threads);
+    } else {
+        work_queue_run(octx->ctx->work_queue, diag_thread_f32, &dctx, n_threads);
    }

    return HTP_STATUS_OK;
@@ -0,0 +1,54 @@
+#ifndef HTP_DIAG_OPS_H
+#define HTP_DIAG_OPS_H
+
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
+#include <stdint.h>
+
+struct htp_diag_vtcm_layout {
+    size_t total_bytes;
+    size_t off_src0;
+    size_t off_dst;
+    size_t src_batch_size_aligned;
+    size_t dst_row_size_aligned;
+    size_t src0_bytes_per_thread;
+    size_t dst_bytes_per_thread;
+};
+
+static inline void htp_diag_vtcm_layout_build(
+    struct htp_diag_vtcm_layout * L,
+    size_t                        src_batch_size,
+    size_t                        dst_row_size,
+    uint32_t                      n_threads
+) {
+    L->src_batch_size_aligned = hex_round_up(src_batch_size, 128);
+    L->dst_row_size_aligned   = hex_round_up(dst_row_size, 128);
+    L->src0_bytes_per_thread  = L->src_batch_size_aligned;
+    L->dst_bytes_per_thread   = L->dst_row_size_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  L->dst_bytes_per_thread * n_threads);
+
+    L->total_bytes = off;
+}
+
+struct htp_diag_kernel_params {
+    uint32_t n_threads;
+    uint32_t total_batches;
+    uint32_t use_dma;
+    uint32_t can_split;
+    uint32_t batches_per_chunk;
+    uint32_t vtcm_size;
+
+    struct fastdiv_values div_ne02;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_diag_kernel_params) <= 128, "htp_diag_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_diag_kernel_params) <= 128, "htp_diag_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_DIAG_OPS_H
@@ -436,10 +436,10 @@ static inline void dma_sync_write(dma_queue * dma_q, dma_addr_t dst, const void
    }
 }

-#define DMA_CACHE_MAX_SIZE 256U
+#define DMA_CACHE_MAX_SIZE 128U

 // Fully assoc LRU cache
-typedef struct {
+typedef struct __attribute__((aligned(128))) {
    uint8_t *base;
    uint32_t line_size;
    uint32_t capacity;
@@ -485,7 +485,7 @@ static inline bool dma_cache_fa_push(dma_queue *q, dma_cache_fa *c, dma_addr_t s
 }

 // Direct mapped cache
-typedef struct {
+typedef struct __attribute__((aligned(128))) {
    uint8_t *base;
    uint32_t line_size;
    uint32_t capacity;
@@ -16,6 +16,7 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-tensor.h"
+#include "fill-ops.h"

 // ggml op_params layout for FILL:
 //   op_params[0] (as float) - the scalar fill value
@@ -30,21 +31,19 @@
                                               \
    const uint32_t nb1 = dst->nb[1];           \
    const uint32_t nb2 = dst->nb[2];           \
-    const uint32_t nb3 = dst->nb[3];           \
-                                               \
-    const uint32_t nr = ne1 * ne2 * ne3;
+    const uint32_t nb3 = dst->nb[3];

 struct htp_fill_context {
    struct htp_ops_context * octx;
+    const struct htp_fill_kernel_params * kparams;
    uint32_t nrows_per_thread;
    uint32_t total_rows;  // ne1 * ne2 * ne3
    uint32_t row_start;
-    bool     opt_path;
    HVX_Vector splat_vec;
    uint32_t   elem_size;
 };

-static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
+static void fill_thread_contiguous(unsigned int nth, unsigned int ith, void * data) {
    const struct htp_fill_context * fctx = (const struct htp_fill_context *) data;
    struct htp_ops_context        * octx = fctx->octx;
    fill_preamble;
@@ -60,45 +59,59 @@ static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);

-    if (fctx->opt_path) {
-        // Opt path: tensor is fully contiguous, treat as flat array
-        const uint32_t elem_start = ir0 * ne0;
-        const uint32_t elem_end = ir1 * ne0;
-        uint8_t * dst_ptr = (uint8_t *) dst->data + elem_start * fctx->elem_size;
-        hvx_splat_u(dst_ptr, fctx->splat_vec, elem_end - elem_start, fctx->elem_size);
-    } else {
-        // Non-contiguous path: must respect strides
-        for (uint32_t ir = ir0; ir < ir1; ++ir) {
-            const uint32_t i1 = ir % ne1;
-            const uint32_t i2 = (ir / ne1) % ne2;
-            const uint32_t i3 = ir / (ne1 * ne2);
-            uint8_t * dst_ptr = (uint8_t *) dst->data + i1*nb1 + i2*nb2 + i3*nb3;
-            hvx_splat_u(dst_ptr, fctx->splat_vec, ne0, fctx->elem_size);
-        }
+    const uint32_t elem_start = ir0 * ne0;
+    const uint32_t elem_end   = ir1 * ne0;
+    uint8_t * dst_ptr = (uint8_t *) dst->data + elem_start * fctx->elem_size;
+    hvx_splat_u(dst_ptr, fctx->splat_vec, elem_end - elem_start, fctx->elem_size);
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1);
+    FARF(HIGH, "fill-contig %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
+}
+
+static void fill_thread_strided(unsigned int nth, unsigned int ith, void * data) {
+    const struct htp_fill_context * fctx = (const struct htp_fill_context *) data;
+    struct htp_ops_context        * octx = fctx->octx;
+    const struct htp_fill_kernel_params * kparams = fctx->kparams;
+    fill_preamble;
+
+    // Parallelise over the flat row index spanning ne1*ne2*ne3
+    const uint32_t ir0 = fctx->row_start + fctx->nrows_per_thread * ith;
+    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_thread, fctx->row_start + fctx->total_rows);
+
+    if (ir0 >= ir1) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
+
+    for (uint32_t ir = ir0; ir < ir1; ++ir) {
+        const uint32_t i3  = fastdiv(ir, &kparams->div_ne1_ne2);
+        const uint32_t rem = ir - i3 * (ne1 * ne2);
+        const uint32_t i2  = fastdiv(rem, &kparams->div_ne1);
+        const uint32_t i1  = rem - i2 * ne1;
+        uint8_t * dst_ptr = (uint8_t *) dst->data + i1*nb1 + i2*nb2 + i3*nb3;
+        hvx_splat_u(dst_ptr, fctx->splat_vec, ne0, fctx->elem_size);
    }

    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1);
-    FARF(HIGH, "fill %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
+    FARF(HIGH, "fill-strided %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
 }

 int op_fill(struct htp_ops_context * octx) {
-    fill_preamble;
-
-    if (dst->type != HTP_TYPE_F32 && dst->type != HTP_TYPE_F16) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
+    const struct htp_tensor * dst = octx->dst;
+    const struct htp_fill_kernel_params * kparams =
+        (const struct htp_fill_kernel_params *) octx->kernel_params;

    if (htp_tensor_is_extended(dst)) {
        return HTP_STATUS_NO_SUPPORT;
    }

    uint32_t row_start = 0;
-    uint32_t nrows     = nr;
+    uint32_t nrows     = kparams->total_rows;

    if (octx->ctx->mdev.count > 1) {
-        const uint32_t row_size = nb1;
-        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
-        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(nr, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(kparams->total_rows, htp_tensor_mdev_data_aligned(dst) ? kparams->rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
        row_start = range.start;
        nrows     = range.count;
    }
@@ -107,40 +120,43 @@ int op_fill(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

-    // nr = ne1*ne2*ne3 (flat row count across all outer dims); parallelise over it.
+    if (!htp_ops_context_set_n_threads(octx, MIN(kparams->n_threads, nrows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
    const uint32_t n_threads = octx->n_threads;

-    // Optimize if fully contiguous: skip stride arithmetic, treat as flat array
-    const bool opt_path = (nb2 == nb1 * ne1) && (nb3 == nb2 * ne2);
-
-    FARF(HIGH, "fill: (%ux%ux%ux%u) type=%u opt=%d\n",
-         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], dst->type, (int) opt_path);
+    FARF(HIGH, "fill: (%ux%ux%ux%u) type=%u ktype=%u\n",
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], dst->type, kparams->kernel_type);

    float val_f32 = 0.f;
    memcpy(&val_f32, &octx->op_params[0], sizeof(float));

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

    switch (dst->type) {
    case HTP_TYPE_F32:
        fctx.splat_vec = hvx_vec_splat_f32(val_f32);
-        fctx.elem_size = sizeof(float);
        break;
    case HTP_TYPE_F16:
        fctx.splat_vec = hvx_vec_splat_f16((_Float16) val_f32);
-        fctx.elem_size = sizeof(_Float16);
        break;
    default:
        return HTP_STATUS_NO_SUPPORT;
    }

-    work_queue_run(octx->ctx->work_queue, fill_thread, &fctx, n_threads);
+    work_queue_func_t worker = (kparams->kernel_type == HTP_FILL_KERNEL_CONTIGUOUS)
+        ? fill_thread_contiguous
+        : fill_thread_strided;
+
+    work_queue_run(octx->ctx->work_queue, worker, &fctx, n_threads);

    return HTP_STATUS_OK;
 }
@@ -0,0 +1,29 @@
+#ifndef HTP_FILL_OPS_H
+#define HTP_FILL_OPS_H
+
+#include "hex-fastdiv.h"
+#include <stdint.h>
+
+enum htp_fill_kernel_type {
+    HTP_FILL_KERNEL_CONTIGUOUS = 0,
+    HTP_FILL_KERNEL_STRIDED    = 1,
+};
+
+struct htp_fill_kernel_params {
+    uint8_t  kernel_type;
+    uint32_t n_threads;
+    uint32_t total_rows;
+    uint32_t elem_size;
+    uint32_t rows_per_chunk;
+
+    struct fastdiv_values div_ne1;
+    struct fastdiv_values div_ne1_ne2;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_fill_kernel_params) <= 128, "htp_fill_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_fill_kernel_params) <= 128, "htp_fill_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_FILL_OPS_H
@@ -39,6 +39,26 @@
 // Must be multiple of 32
 #define FLASH_ATTN_BLOCK_SIZE (32 * 2)

+// Large scratch lives in ctx->ddr_spad to keep stack usage low.
+// Ops can run on the dspqueue callback thread, which has a small stack.
+struct hvx_fa_ddr_spad {
+    __fp16   slopes[512];
+    uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+    uint16_t thread_active_blocks[HTP_MAX_NTHREADS][FA_MAX_KV_BLOCKS] __attribute__((aligned(128)));
+};
+
+struct hmx_fa_ddr_spad {
+    uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+    __fp16   slopes[] __attribute__((aligned(128))); // [g_br]
+};
+
+static inline void * fa_get_ddr_spad(const struct htp_ops_context * octx, size_t size) {
+    if (!octx->ctx->ddr_spad_base || octx->ctx->ddr_spad_size < size) {
+        return NULL;
+    }
+    return octx->ctx->ddr_spad_base;
+}
+
 struct htp_fa_context {
    const struct htp_ops_context * octx;

@@ -61,7 +81,6 @@ struct htp_fa_context {
    uint32_t n_head_log2;
    float m0;
    float m1;
-    __fp16 slopes[512];

    uint32_t n_blocks;
    uint32_t n_active_blocks;
@@ -90,7 +109,9 @@ struct htp_fa_context {
    uint8_t *  spad_m;
    uint8_t *  spad_a;
    float *    spad_sinks;
-    uint16_t   active_blocks[FA_MAX_KV_BLOCKS];
+
+    struct hvx_fa_ddr_spad * ddr;
+    dma_cache_dm *           m_cache; // [n_threads]

    const struct htp_tensor * k;
    const struct htp_tensor * v;
@@ -156,6 +177,7 @@ struct hmx_fa_context {
    __fp16 *     vtcm_mask_buf;        // VTCM mask buffer [Br * m_line], DMA'd per KV block
    __fp16 *     vtcm_slopes;          // ALiBi slopes [g_br]
    float *      vtcm_sinks;           // Attention sinks
+    __fp16 *     ddr_slopes;           // ALiBi slopes staging buffer in DDR [g_br]
    size_t       row_buf_stride;       // HVX vectors per row buffer (Bc/64)
    size_t       mask_buf_row_stride;  // elements (__fp16) per row in mask buffer
    size_t       q_tile_bytes;
@@ -164,7 +186,7 @@ struct hmx_fa_context {
    size_t       col_vec_bytes;
    size_t       d_tile_bytes;
    bool         mask_broadcast;       // true when mask->ne[2] == 1 (head-independent, single 2D DMA)
-    dma_cache_fa m_cache;
+    dma_cache_fa * m_cache;
 };

 // Scan mask rows [q_start, q_start + n_rows_q) and build compact list of KV blocks with any entry != -inf.
@@ -308,11 +330,11 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
    uint8_t * spad_m = factx->spad_m + (mask ? factx->size_m_block * HVX_FA_DMA_CACHE_SIZE : 0) * ith;
    uint8_t * spad_a = factx->spad_a + factx->size_vkq_acc * ith;

-    dma_cache_dm m_cache;
-    dma_cache_dm_init(&m_cache, spad_m, factx->size_m_block, HVX_FA_DMA_CACHE_SIZE);
+    dma_cache_dm * m_cache = &factx->m_cache[ith];
+    dma_cache_dm_init(m_cache, spad_m, factx->size_m_block, HVX_FA_DMA_CACHE_SIZE);

-    uint16_t thread_active_blocks[FA_MAX_KV_BLOCKS];
-    const uint16_t * active_blocks = factx->shared_active_blocks ? factx->active_blocks : thread_active_blocks;
+    uint16_t * thread_active_blocks = factx->ddr->thread_active_blocks[ith];
+    const uint16_t * active_blocks = factx->shared_active_blocks ? factx->ddr->active_blocks : thread_active_blocks;
    uint32_t n_active_blocks       = factx->n_active_blocks;
    dma_addr_t prev_mp_base        = 0;

@@ -385,7 +407,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
            const uint32_t r_iq1 = (r - r_iq3*neq2*neq1 - r_iq2 * neq1);

            heads[g] = r_iq2;
-            const __fp16 slope = factx->slopes[r_iq2];
+            const __fp16 slope = factx->ddr->slopes[r_iq2];
            slope_vecs[g] = hvx_vec_splat_f16(slope);

            S_vec[g] = hvx_vec_splat_f32(0.0f);
@@ -422,7 +444,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
            // Mask
            if (mask) {
                const dma_addr_t m_src = mp_base + ic_start * sizeof(__fp16);
-                dma_cache_dm_push(dma_q, &m_cache, ib, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
+                dma_cache_dm_push(dma_q, m_cache, ib, m_src, current_block_size * 2, current_block_size * 2, current_block_size * 2, 1);
            }
        }

@@ -682,7 +704,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                    // Mask
                    if (mask) {
                        const dma_addr_t m_src = mp_base + next_ic_start * sizeof(__fp16);
-                        dma_cache_dm_push(dma_q, &m_cache, next_ib, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1);
+                        dma_cache_dm_push(dma_q, m_cache, next_ib, m_src, next_block_size * 2, next_block_size * 2, next_block_size * 2, 1);
                    }
                }
            } // end for g
@@ -1887,7 +1909,7 @@ static __attribute__((noinline)) void fa_compute_slopes(
    const float    m0          = factx->m0;
    const float    m1          = factx->m1;

-    __fp16 temp_slopes[512] __attribute__((aligned(128)));
+    __fp16 * temp_slopes = factx->ddr_slopes;
    if (G <= 32) {
        // Fast path: Compute G unique slope values in vector registers
        HVX_Vector v_val = hvx_alibi_slopes(kv_head, G, n_head_log2, m0, m1);
@@ -1905,14 +1927,13 @@ static __attribute__((noinline)) void fa_compute_slopes(
        }
    }

-    // Allocate stack buffer to avoid scalar writes to VTCM (which generates L2 misses)
-    __fp16 local_slopes[n_rows_g] __attribute__((aligned(128)));
-    for (size_t r = 0; r < n_rows_g; ++r) {
-        local_slopes[r] = temp_slopes[fastmodulo(r, G, &factx->div_G)];
+    // Replicate the G values in DDR to avoid scalar writes to VTCM (which generates L2 misses)
+    for (size_t r = G; r < n_rows_g; ++r) {
+        temp_slopes[r] = temp_slopes[r - G];
    }

    // Copy to VTCM slopes using HVX block copy (both are aligned to 128 bytes)
-    hvx_copy_f16_aa((uint8_t *)slopes, (const uint8_t *)local_slopes, n_rows_g);
+    hvx_copy_f16_aa((uint8_t *)slopes, (const uint8_t *)temp_slopes, n_rows_g);
 }

 static void fa_push_mask_dma_gqa(
@@ -1958,7 +1979,7 @@ static inline void fa_prefetch_block(dma_queue * dma_q, const struct htp_tensor
    if (mask) {
        if (__builtin_expect(factx->mask_broadcast, true)) {
            const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + prefetch_start * sizeof(__fp16);
-            dma_cache_fa_push(dma_q, &factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q);
+            dma_cache_fa_push(dma_q, factx->m_cache, ms_src, m_line_bytes, mask->nb[1], prefetch_rows * sizeof(__fp16), n_rows_q);
        } else {
            fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, prefetch_start, kv_head, G, m_line_bytes, prefetch_rows, n_rows_q, factx);
        }
@@ -2073,7 +2094,29 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
        const uint32_t dst_e_size = (dst->type == HTP_TYPE_F32) ? sizeof(float) : sizeof(__fp16);
        const bool can_split      = htp_tensor_can_row_partition(dst, dst_e_size);

-        if (kparams->head_split && can_split && n_kv_heads >= mdev_count && n_kv_heads % mdev_count == 0) {
+        if ((kparams->flags & HTP_FA_FLAG_TEAM_SPLIT) && can_split) {
+            // KV heads split across teams, Q-blocks split within a team
+            uint32_t n_teams, team_size;
+            htp_fa_mdev_teams(n_kv_heads, mdev_count, kparams->head_split != 0, &n_teams, &team_size);
+
+            const uint32_t team        = mdev_idx / team_size;
+            const uint32_t member      = mdev_idx - team * team_size;
+            const uint32_t kv_per_team = n_kv_heads / n_teams;
+            kv_head_min = team * kv_per_team;
+            kv_head_max = kv_head_min + kv_per_team;
+
+            const uint32_t base  = n_q_blocks / team_size;
+            const uint32_t rem   = n_q_blocks - base * team_size;
+            const uint32_t b0    = member * base + MIN(member, rem);
+            const uint32_t nblk  = base + (member < rem ? 1 : 0);
+
+            if (nblk == 0) {
+                return HTP_STATUS_OK;
+            }
+
+            q_start_min = b0 * Br;
+            q_start_max = MIN((b0 + nblk) * Br, neq1);
+        } else if (kparams->head_split && can_split && n_kv_heads >= mdev_count && n_kv_heads % mdev_count == 0) {
            const uint32_t kv_per_core = n_kv_heads / mdev_count;
            kv_head_min = mdev_idx * kv_per_core;
            kv_head_max = kv_head_min + kv_per_core;
@@ -2151,7 +2194,14 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {

    const size_t m_line_bytes = L.m_line_bytes;  // used by the mask DMAs in the KV loop

-    dma_cache_fa_init(&factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, L.m_buf_slot_bytes, HMX_FA_DMA_CACHE_SIZE);
+    struct hmx_fa_ddr_spad * ddr = (struct hmx_fa_ddr_spad *) fa_get_ddr_spad(octx, sizeof(struct hmx_fa_ddr_spad) + factx.g_br * sizeof(__fp16));
+    if (!ddr) {
+        return HTP_STATUS_INTERNAL_ERR;
+    }
+
+    factx.ddr_slopes = ddr->slopes;
+    factx.m_cache    = &octx->dma_cache_fa[0];
+    dma_cache_fa_init(factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, L.m_buf_slot_bytes, HMX_FA_DMA_CACHE_SIZE);

    // Head-dim padding: the K/V DMA staging buffers and the flat-Q buffer are laid out
    // with padded row strides (size_{k,v,q}_row_padded, covering D_pad columns) but the
@@ -2210,7 +2260,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
            const size_t   g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
            const size_t   n_row_tiles = g_br_actual / HMX_FP16_TILE_N_ROWS;

-            uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+            uint16_t * active_blocks = ddr->active_blocks;
            uint32_t n_active_blocks = factx.n_kv_blocks;
            if (mask && factx.mask_broadcast && factx.n_kv_blocks > 1) {
                const size_t scratch_bytes = factx.k_tile_bytes * (factx.pipeline ? 2 : 1);
@@ -2252,7 +2302,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                        if (factx.pipeline && mask) {
                            if (__builtin_expect(factx.mask_broadcast, true)) {
                                const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + b0_start * sizeof(__fp16);
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
+                                dma_cache_fa_push(dma_q, factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
                            } else {
                                fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, b0_start, kv_head, G, m_line_bytes, b0_rows, n_rows_q, &factx);
                            }
@@ -2453,7 +2503,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                        if (mask) {
                            if (__builtin_expect(factx.mask_broadcast, true)) {
                                const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
+                                dma_cache_fa_push(dma_q, factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
                            } else {
                                fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
                            }
@@ -2577,7 +2627,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                        if (factx.pipeline && mask) {
                            if (__builtin_expect(factx.mask_broadcast, true)) {
                                const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + b0_start * sizeof(__fp16);
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
+                                dma_cache_fa_push(dma_q, factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
                            } else {
                                fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, b0_start, next_kv_head, G, m_line_bytes, b0_rows, n_rows_q, &factx);
                            }
@@ -2680,8 +2730,15 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
    if (n_head > 512) {
        return HTP_STATUS_NO_SUPPORT;
    }
+
+    factx.ddr = (struct hvx_fa_ddr_spad *) fa_get_ddr_spad(octx, sizeof(struct hvx_fa_ddr_spad));
+    if (!factx.ddr) {
+        return HTP_STATUS_INTERNAL_ERR;
+    }
+    factx.m_cache = octx->dma_cache_dm;
+
    for (uint32_t h = 0; h < n_head; ++h) {
-        factx.slopes[h] = (__fp16) ((kparams->max_bias > 0.0f) ? alibi_slope(h, factx.n_head_log2, factx.m0, factx.m1) : 1.0f);
+        factx.ddr->slopes[h] = (__fp16) ((kparams->max_bias > 0.0f) ? alibi_slope(h, factx.n_head_log2, factx.m0, factx.m1) : 1.0f);
    }

    // total rows in q
@@ -2754,7 +2811,7 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
    if (has_mask && neq1 == 1 && neq3 == 1 && mask->ne[2] == 1 && factx.n_blocks > 1 && factx.n_blocks <= FA_MAX_KV_BLOCKS) {
        factx.n_active_blocks = fa_scan_active_blocks(
            octx->ctx->dma[0], mask, 0, 1, 0, 0, k->ne[1], FLASH_ATTN_BLOCK_SIZE,
-            factx.n_blocks, factx.spad_k, factx.size_k_block * 2, factx.active_blocks);
+            factx.n_blocks, factx.spad_k, factx.size_k_block * 2, factx.ddr->active_blocks);
        factx.shared_active_blocks = true;
    }

@@ -32,10 +32,20 @@ enum htp_fa_kernel_type {
    HTP_FA_KERNEL_HMX
 };

+// Multicore: gcd(n_kv_heads, N) teams own the KV heads, team members share the Q-blocks
+#define HTP_FA_FLAG_TEAM_SPLIT 0x01
+
+static inline void htp_fa_mdev_teams(uint32_t n_kv_heads, uint32_t mdev_count, bool head_split,
+                                     uint32_t * n_teams, uint32_t * team_size) {
+    const uint32_t t = (head_split && n_kv_heads > 0 && mdev_count > 1) ? hex_gcd_u32(n_kv_heads, mdev_count) : 1;
+    *n_teams   = t;
+    *team_size = mdev_count > 1 ? mdev_count / t : 1;
+}
+
 struct htp_fa_kernel_params {
    uint8_t  kernel_type;        // enum htp_fa_kernel_type
    uint8_t  head_split;         // 1 = partition by KV heads in multicore, 0 = token partition
-    uint8_t  flags;              // reserved
+    uint8_t  flags;              // HTP_FA_FLAG_*
    uint8_t  n_threads;          // Number of threads to run

    // Common parameters
@@ -49,9 +59,6 @@ struct htp_fa_kernel_params {
    float    logit_softcap;
    uint32_t vtcm_size;

-    uint32_t qrows;
-    uint32_t qrows_per_thread;
-    uint32_t qrow_start;
    float    m0;
    float    m1;
    uint32_t n_head_log2;
@@ -67,8 +74,6 @@ struct htp_fa_kernel_params {
    union {
        struct {
            uint32_t g_br;
-            uint32_t row_buf_stride;
-            uint32_t mask_buf_row_stride;
            int32_t  mask_broadcast;
            int32_t  pipeline;
            struct fastdiv_values div_G;
@@ -85,6 +90,8 @@ struct htp_fa_kernel_params {

 #if defined(__cplusplus)
 static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params is too large for kernel_params blob");
 #endif

 // VTCM region layout for the HMX flash-attention kernel.
@@ -325,7 +332,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
                                         size_t   n_threads,
                                         bool     is_q_fp32,
                                         bool     has_sinks,
-                                         size_t   n_heads) {
+                                         size_t   n_heads,
+                                         size_t   q_split) {
    const size_t T       = HMX_FP16_TILE_N_ROWS;  // 32
    const size_t br_unit = hmx_ceil_div(T, gqa_factor);
    const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2;  // 64
@@ -365,7 +373,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
                const size_t vecs_per_t     = n_use > 0 ? (n_row_vec_cnt + n_use - 1) / n_use : 1;

                const size_t c_iter_actual  = c_iter_base + c_softmax * vecs_per_t;
-                const size_t cost           = q_blocks * (c_q_fixed + kv_blocks * c_iter_actual);
+                const size_t q_blocks_nsp   = (q_blocks + q_split - 1) / (q_split ? q_split : 1);
+                const size_t cost           = q_blocks_nsp * (c_q_fixed + kv_blocks * c_iter_actual);
                const size_t mn             = Br * Bc;

                if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
@@ -2270,46 +2270,9 @@ static int gated_delta_net_f32_hmx_chunked(
 }

 int op_gated_delta_net(struct htp_ops_context * octx) {
-    const struct htp_tensor * q     = octx->src[0];
-    const struct htp_tensor * k     = octx->src[1];
-    const struct htp_tensor * v     = octx->src[2];
-    const struct htp_tensor * g     = octx->src[3];
-    const struct htp_tensor * beta  = octx->src[4];
-    const struct htp_tensor * state = octx->src[5];
-    const struct htp_tensor * dst   = octx->dst;
+    const struct htp_tensor * dst = octx->dst;

-    if (q->type != HTP_TYPE_F32 || k->type != HTP_TYPE_F32 || v->type != HTP_TYPE_F32 ||
-        g->type != HTP_TYPE_F32 || beta->type != HTP_TYPE_F32 || state->type != HTP_TYPE_F32 ||
-        dst->type != HTP_TYPE_F32) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    const uint32_t S_v      = v->ne[0];
-    const uint32_t H        = v->ne[1];
-    const uint32_t n_tokens = v->ne[2];
-    const uint32_t n_seqs   = v->ne[3];
-    const uint32_t K        = octx->op_params[0];
-
-    if (S_v == 0 || S_v > HTP_GDN_MAX_SV || H == 0 || n_tokens == 0 || n_seqs == 0) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    if ((g->ne[0] != 1 && g->ne[0] != S_v) || beta->ne[0] != 1) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    if (q->ne[0] != S_v || k->ne[0] != S_v || q->ne[1] == 0 || k->ne[1] == 0 ||
-        q->ne[2] != n_tokens || k->ne[2] != n_tokens || q->ne[3] == 0 || k->ne[3] == 0 ||
-        (n_seqs % q->ne[3]) != 0 || (n_seqs % k->ne[3]) != 0) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    // state holds s0 only: [S_v, S_v, H, n_seqs]
-    if (state->ne[0] != S_v || state->ne[1] != S_v || state->ne[2] != H || state->ne[3] != n_seqs) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-    if (dst->ne[0] != S_v * H || dst->ne[1] != n_tokens * n_seqs + S_v * n_seqs * K) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    for (int i = 0; i < 5; i++) {
+    for (int i = 0; i < 6; i++) {
        if (htp_tensor_is_extended(octx->src[i])) {
            return HTP_STATUS_NO_SUPPORT;
        }
@@ -2325,68 +2288,12 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
    }

    const struct htp_gdn_kernel_params * kparams = (const struct htp_gdn_kernel_params *) octx->kernel_params;
-    struct htp_gdn_kernel_params kparams_local;
-    if (!kparams || kparams->S_v == 0) {
-        const uint32_t rq3 = n_seqs / q->ne[3];
-        const uint32_t rk3 = n_seqs / k->ne[3];
-        const uint32_t total_rows = H * n_seqs;
-        uint32_t n_threads = (total_rows < octx->n_threads) ? total_rows : octx->n_threads;
-        if (n_threads == 0) {
-            n_threads = 1;
-        }
-
-        memset(&kparams_local, 0, sizeof(kparams_local));
-        kparams_local.n_threads           = n_threads;
-        kparams_local.S_v                 = S_v;
-        kparams_local.H                   = H;
-        kparams_local.n_tokens            = n_tokens;
-        kparams_local.n_seqs              = n_seqs;
-        kparams_local.K                   = K;
-        kparams_local.total_rows          = total_rows;
-        kparams_local.rows_per_thread     = (total_rows + n_threads - 1) / n_threads;
-        const bool can_use_hmx = (octx->ctx->hmx_enabled) &&
-                                 (S_v % 64 == 0) &&
-                                 (n_tokens >= HTP_GDN_MIN_TOKENS) &&
-                                 (g->ne[0] == 1) &&
-                                 (K == 1);
-
-        struct htp_gdn_hmx_vtcm_layout hmx_layout_local;
-        struct htp_gdn_vtcm_layout hvx_layout_local;
-        uint32_t n_heads_batch = 1;
-
-        if (can_use_hmx && htp_gdn_hmx_solve_layout(&hmx_layout_local, S_v, HTP_GDN_CHUNK_SIZE, total_rows, octx->ctx->vtcm_size, n_threads, true, &n_heads_batch)) {
-            kparams_local.kernel_type     = HTP_GDN_KERNEL_HMX_CHUNKED;
-            kparams_local.pipeline        = hmx_layout_local.pipeline ? 1 : 0;
-            kparams_local.chunk_size      = HTP_GDN_CHUNK_SIZE;
-            kparams_local.n_chunks        = (n_tokens + HTP_GDN_CHUNK_SIZE - 1) / HTP_GDN_CHUNK_SIZE;
-            kparams_local.n_heads_batch   = (uint16_t) n_heads_batch;
-            kparams_local.vtcm_size       = (uint32_t) hmx_layout_local.total_bytes;
-            kparams_local.state_aligned   = (uint32_t) hmx_layout_local.state_f32_bytes;
-            kparams_local.vtcm_per_thread = (uint32_t) (hmx_layout_local.total_bytes / (n_threads > 0 ? n_threads : 1));
-        } else {
-            htp_gdn_vtcm_layout_build(&hvx_layout_local, S_v, n_threads);
-            kparams_local.kernel_type     = HTP_GDN_KERNEL_HVX_RECURRENT;
-            kparams_local.pipeline        = 0;
-            kparams_local.n_heads_batch   = 1;
-            kparams_local.state_aligned   = (uint32_t) hvx_layout_local.state_aligned;
-            kparams_local.vtcm_per_thread = (uint32_t) hvx_layout_local.bytes_per_thread;
-            kparams_local.vtcm_size       = (uint32_t) hvx_layout_local.total_bytes;
-        }
-        kparams_local.kda                 = (g->ne[0] == S_v) ? 1 : 0;
-        kparams_local.scale               = 1.0f / sqrtf((float) S_v);
-        kparams_local.state_seq_stride    = (uint32_t) (state->nb[3] / sizeof(float));
-        kparams_local.state_size_per_snap = S_v * S_v * H * n_seqs;
-
-        kparams_local.div_H         = init_fastdiv_values(H);
-        kparams_local.div_q1        = init_fastdiv_values(q->ne[1]);
-        kparams_local.div_k1        = init_fastdiv_values(k->ne[1]);
-        kparams_local.div_rq3       = init_fastdiv_values(rq3);
-        kparams_local.div_rk3       = init_fastdiv_values(rk3);
-        kparams_local.div_n_threads = init_fastdiv_values(n_threads);
-
-        kparams = &kparams_local;
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
    }

+    const uint32_t S_v        = kparams->S_v;
+    const uint32_t n_tokens   = kparams->n_tokens;
    const uint32_t total_rows = kparams->total_rows;
    uint32_t row_start = 0;
    uint32_t nrows     = total_rows;
@@ -2416,16 +2323,16 @@ int op_gated_delta_net(struct htp_ops_context * octx) {
        return gated_delta_net_f32_hmx_chunked(octx, kparams, row_start, nrows);
    }

-    const uint32_t n_threads = (nrows < kparams->n_threads) ? nrows : kparams->n_threads;
+    const uint32_t n_threads = (nrows < octx->n_threads) ? nrows : octx->n_threads;

    struct htp_gdn_context gctx;
    gctx.octx      = octx;
    gctx.kparams   = kparams;
    gctx.row_start = row_start;
    gctx.nrows     = nrows;
-    gctx.vtcm_base = octx->ctx->vtcm_base;

    htp_gdn_vtcm_layout_build(&gctx.layout, S_v, n_threads);
+    gctx.vtcm_base = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, gctx.layout.off_state);

    if (gctx.layout.total_bytes > octx->ctx->vtcm_size) {
        return HTP_STATUS_VTCM_TOO_SMALL;
@@ -2433,10 +2340,10 @@ int op_gated_delta_net(struct htp_ops_context * octx) {

    FARF(HIGH, "gated-delta-net-f32: q(%ux%ux%ux%u) k(%ux%ux%ux%u) v(%ux%ux%ux%u) state(%ux%ux%ux%u) -> (%ux%ux%ux%u) : "
         "vtcm-size %zu n_threads %u\n",
-         q->ne[0], q->ne[1], q->ne[2], q->ne[3],
-         k->ne[0], k->ne[1], k->ne[2], k->ne[3],
-         v->ne[0], v->ne[1], v->ne[2], v->ne[3],
-         state->ne[0], state->ne[1], state->ne[2], state->ne[3],
+         octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
+         octx->src[1]->ne[0], octx->src[1]->ne[1], octx->src[1]->ne[2], octx->src[1]->ne[3],
+         octx->src[2]->ne[0], octx->src[2]->ne[1], octx->src[2]->ne[2], octx->src[2]->ne[3],
+         octx->src[5]->ne[0], octx->src[5]->ne[1], octx->src[5]->ne[2], octx->src[5]->ne[3],
         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
         gctx.layout.total_bytes, n_threads);

@@ -37,13 +37,7 @@ struct htp_gdn_kernel_params {
    uint32_t K;

    uint32_t total_rows;
-    uint32_t row_start;
-    uint32_t nrows;
-    uint32_t rows_per_thread;
-
    uint32_t kda;
-    uint32_t state_aligned;
-    uint32_t vtcm_per_thread;
    uint32_t vtcm_size;
    uint32_t state_seq_stride;
    uint32_t state_size_per_snap;
@@ -55,7 +49,6 @@ struct htp_gdn_kernel_params {
    struct fastdiv_values div_k1;
    struct fastdiv_values div_rq3;
    struct fastdiv_values div_rk3;
-    struct fastdiv_values div_n_threads;
 };

 #if defined(__cplusplus)
@@ -67,18 +60,21 @@ _Static_assert(sizeof(struct htp_gdn_kernel_params) <= 128, "htp_gdn_kernel_para
 struct htp_gdn_vtcm_layout {
    size_t state_aligned;
    size_t bytes_per_thread;
+    size_t off_state;
    size_t total_bytes;
 };

 static inline void htp_gdn_vtcm_layout_build(
-    struct htp_gdn_vtcm_layout * layout,
+    struct htp_gdn_vtcm_layout * L,
    uint32_t S_v,
    uint32_t n_threads
 ) {
    size_t state_bytes = (size_t) S_v * S_v * sizeof(float);
-    layout->state_aligned = hex_round_up(state_bytes, 128);
-    layout->bytes_per_thread = 2 * layout->state_aligned;
-    layout->total_bytes = layout->bytes_per_thread * n_threads;
+    L->state_aligned = hex_round_up(state_bytes, 128);
+    L->bytes_per_thread = 2 * L->state_aligned;
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_state, L->bytes_per_thread * n_threads);
+    L->total_bytes = off;
 }

 struct htp_gdn_hmx_vtcm_layout {
@@ -206,14 +206,22 @@ static void gather_rows_thread_f32(unsigned int nth, unsigned int ith, void * da

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

-    uint8_t * base   = task->octx->ctx->vtcm_base + ith * kparams->slice;
-    HVX_Vector * tab = (HVX_Vector *) base;
-    HVX_Vector * inc = tab + kparams->vg;
-    HVX_Vector * cur = inc + kparams->vg;
-    uint8_t * region_base = (uint8_t *) (cur + kparams->vg);
+    struct htp_gather_rows_vtcm_layout layout;
+    htp_gather_rows_vtcm_layout_build(&layout, kparams->vg, kparams->region_size, kparams->nr_max, kparams->ne, 4);

-    uint8_t * region[2] = { region_base, region_base + kparams->buf_stride };
-    uint8_t * out[2]    = { region[0] + kparams->out_off, region[1] + kparams->out_off };
+    uint8_t * base   = VTCM_LAYOUT_PTR(uint8_t, task->octx->ctx->vtcm_base, (size_t) ith * kparams->slice);
+    HVX_Vector * tab = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_tab);
+    HVX_Vector * inc = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_inc);
+    HVX_Vector * cur = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_cur);
+
+    uint8_t * region[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region1),
+    };
+    uint8_t * out[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out1),
+    };

    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, (uint16_t) r_beg);
    gather_rows_init_tables_f32(tab, inc, kparams, a, b);
@@ -279,14 +287,22 @@ static void gather_rows_thread_f16(unsigned int nth, unsigned int ith, void * da

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

-    uint8_t * base   = task->octx->ctx->vtcm_base + ith * kparams->slice;
-    HVX_Vector * tab = (HVX_Vector *) base;
-    HVX_Vector * inc = tab + kparams->vg;
-    HVX_Vector * cur = inc + kparams->vg;
-    uint8_t * region_base = (uint8_t *) (cur + kparams->vg);
+    struct htp_gather_rows_vtcm_layout layout;
+    htp_gather_rows_vtcm_layout_build(&layout, kparams->vg, kparams->region_size, kparams->nr_max, kparams->ne, 2);

-    uint8_t * region[2] = { region_base, region_base + kparams->buf_stride };
-    uint8_t * out[2]    = { region[0] + kparams->out_off, region[1] + kparams->out_off };
+    uint8_t * base   = VTCM_LAYOUT_PTR(uint8_t, task->octx->ctx->vtcm_base, (size_t) ith * kparams->slice);
+    HVX_Vector * tab = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_tab);
+    HVX_Vector * inc = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_inc);
+    HVX_Vector * cur = VTCM_LAYOUT_PTR(HVX_Vector, base, layout.off_cur);
+
+    uint8_t * region[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_region1),
+    };
+    uint8_t * out[2] = {
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out0),
+        VTCM_LAYOUT_PTR(uint8_t, base, layout.off_out1),
+    };

    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, (uint16_t) r_beg);
    gather_rows_init_tables_f16(tab, inc, kparams, a, b);
@@ -2,6 +2,7 @@
 #define HTP_GATHER_ROWS_H

 #include "hex-common.h"
+#include "htp-vtcm.h"
 #include <stdint.h>
 #include <stdbool.h>

@@ -19,8 +20,6 @@ struct htp_gather_rows_params {
    uint32_t v1;
    uint32_t span1;
    uint32_t region_size;
-    uint32_t out_off;
-    uint32_t buf_stride;
    uint32_t slice;
    uint32_t total_rows;
    uint32_t rows_per_thread;
@@ -28,6 +27,40 @@ struct htp_gather_rows_params {
    uint8_t  pad[3];
 };

+struct htp_gather_rows_vtcm_layout {
+    size_t off_tab;
+    size_t off_inc;
+    size_t off_cur;
+    size_t off_region0;
+    size_t off_out0;
+    size_t off_region1;
+    size_t off_out1;
+    size_t total_per_thread;
+};
+
+static inline void htp_gather_rows_vtcm_layout_build(
+    struct htp_gather_rows_vtcm_layout * L,
+    uint32_t vg,
+    uint32_t region_size,
+    uint32_t nr_max,
+    uint32_t ne,
+    uint32_t elem_size) {
+
+    const size_t vec_bytes      = (size_t) vg * 128;
+    const size_t region_aligned = hex_round_up(region_size, 128);
+    const size_t out_aligned    = hex_round_up(region_aligned + (size_t) nr_max * ne * elem_size, 128) - region_aligned;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_tab,     vec_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_inc,     vec_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_cur,     vec_bytes);
+    VTCM_LAYOUT_ALLOC(off, off_region0, region_aligned);
+    VTCM_LAYOUT_ALLOC(off, off_out0,    out_aligned);
+    VTCM_LAYOUT_ALLOC(off, off_region1, region_aligned);
+    VTCM_LAYOUT_ALLOC(off, off_out1,    out_aligned);
+    L->total_per_thread = off;
+}
+
 static inline bool htp_gather_rows_solve_layout(
    struct htp_gather_rows_params * kparams,
    uint32_t na, uint32_t a_nb1,
@@ -80,8 +113,6 @@ static inline bool htp_gather_rows_solve_layout(
    uint32_t v1          = 0;
    uint32_t span1       = 0;
    uint32_t region_size = 0;
-    uint32_t out_off     = 0;
-    uint32_t buf_stride  = 0;

    while (nr_max > 0) {
        v1 = hex_round_up((nr_max - 1) * a_nb1 + na * elem_size, 128);
@@ -98,10 +129,10 @@ static inline bool htp_gather_rows_solve_layout(
            region_size = v1;
        }

-        out_off    = hex_round_up(region_size, 128);
-        buf_stride = hex_round_up(out_off + nr_max * ne * elem_size, 128);
+        struct htp_gather_rows_vtcm_layout layout;
+        htp_gather_rows_vtcm_layout_build(&layout, vg, region_size, nr_max, ne, elem_size);

-        if ((elem_size != 2 || region_size <= 32768) && fixed + 2 * buf_stride <= slice) {
+        if ((elem_size != 2 || region_size <= 32768) && layout.total_per_thread <= slice) {
            break;
        }

@@ -121,8 +152,6 @@ static inline bool htp_gather_rows_solve_layout(
    kparams->v1              = v1;
    kparams->span1           = span1;
    kparams->region_size     = region_size;
-    kparams->out_off         = out_off;
-    kparams->buf_stride      = buf_stride;
    kparams->slice           = slice;
    kparams->total_rows      = total_rows;
    kparams->rows_per_thread = rows_per_thread;
@@ -404,7 +404,7 @@ int op_get_rows(struct htp_ops_context * octx) {

    if (octx->ctx->mdev.count > 1) {
        uint32_t tasks_per_chunk = 1;
-        htp_tensor_mdev_rows_per_chunk(dst, dst_row_size / dst->ne[0], dst_row_size, &tasks_per_chunk);
+        htp_tensor_mdev_rows_per_chunk(dst, (uint32_t) dst->nb[0], dst_row_size, &tasks_per_chunk);
        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_tasks, tasks_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
        task_start = range.start;
        tasks      = range.count;
@@ -414,7 +414,7 @@ int op_get_rows(struct htp_ops_context * octx) {
        return HTP_STATUS_OK;
    }

-    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
        return HTP_STATUS_INVAL_PARAMS;
    }

@@ -428,7 +428,7 @@ int op_get_rows(struct htp_ops_context * octx) {
    grctx.compute_tiled = NULL;
    grctx.task_start = task_start;
    grctx.tasks = tasks;
-    grctx.tasks_per_thread = octx->ctx->mdev.count == 1 ? kparams->tasks_per_thread : fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
+    grctx.tasks_per_thread = fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
    grctx.tile_size   = 0;
    grctx.tile_stride = 0;
    grctx.src0_row_size = src0_row_size;
@@ -15,11 +15,10 @@ enum htp_get_rows_kernel_type {
 };

 struct htp_get_rows_kernel_params {
-    int32_t  n_threads;
-    int32_t  kernel_type;
-    int32_t  total_tasks;
-    int32_t  tasks_per_thread;
-    int32_t  vtcm_size;
+    uint8_t  kernel_type;
+    uint32_t n_threads;
+    uint32_t total_tasks;
+    uint32_t vtcm_size;

    // Fastdiv helpers
    struct fastdiv_values div_ne10;
--- a/Show More
+++ b/Show More