tests : adjust default threads for test-llama-context (#31157 )

qwen: merge tensors on convert (#31028 )
* qwen: merge tensors on convert * convert: use generator functions * qwen4exp: remove conts and pre-scale
2026-10-25 14:07:36 -05:00 · 2026-10-25 19:46:44 +02:00 · 2026-10-25 15:23:22 +02:00 · 2026-10-25 13:53:55 +02:00 · 2026-10-25 12:31:09 +01:00 · 2026-10-25 13:28:38 +02:00
219 changed files with 10220 additions and 2181 deletions
@@ -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:
@@ -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
 }
@@ -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,76 +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
-
-    # 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
-
    function check_ppl {
        qnt="$1"
        ppl=$(echo "$2" | grep -oE "[0-9]+\.[0-9]+" | tail -n 1)
@@ -469,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 99 -c 1024 -fa off                ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+    fi
+    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa on  --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa on                 ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+
    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
 }
@@ -505,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

@@ -513,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

@@ -525,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
 }

@@ -541,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

@@ -549,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)

@@ -581,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
 }

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

@@ -1095,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))) {
@@ -1460,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()) {
@@ -1475,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;
@@ -70,6 +70,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "LingDSparkModel": "qwen",
    "DeepseekV4ForCausalLM": "deepseek",
    "DeepseekV4DSparkModel": "deepseek",
+    "DeepseekV41ForCausalLM": "deepseek41",
    "DistilBertForMaskedLM": "bert",
    "DistilBertForSequenceClassification": "bert",
    "DistilBertModel": "bert",
@@ -286,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",
@@ -325,6 +327,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "VllmD3Model": "vllm_d3",
    "DeepseekOCR2ForCausalLM": "deepseek",
    "DeepseekOCRForCausalLM": "deepseek",
+    "DeepseekV41ForCausalLM": "deepseek41",
    "DeepseekV4ForCausalLM": "deepseek",
    "Dots3NoteForCausalLM": "dots3",
    "Dots3NoteForConditionalGeneration": "dots3",
@@ -380,6 +383,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "Qwen3ASRForConditionalGeneration": "qwen3vl",
    "Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
    "PocketTTSModel": "pockettts",
+    "SopranoModel": "soprano",
    "Qwen3TTSForConditionalGeneration": "qwen3tts",
    "Qwen3VLForConditionalGeneration": "qwen3vl",
    "Qwen3VLMoeForConditionalGeneration": "qwen3vl",
@@ -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"]
@@ -41,6 +41,8 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
        self._ple_shards: dict[int, str] = {}
        self._ple_row_dim: int | None = None
        self._mtp_fc: dict[str, Tensor] = {}
+        # halves of the merged weights, keyed by (bid, merged tensor), see _merge_shared_input
+        self._merge_parts: dict[tuple[int, gguf.MODEL_TENSOR], dict[int, Tensor]] = {}

    @classmethod
    def filter_tensors(cls, item):
@@ -131,44 +133,79 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
        # int64 hash constants must stay exact; 1-D tensors force F32, so use KV
        if name.endswith("ple_embedding.layer_multipliers"):
            self._ple_multipliers = [int(x) for x in data_torch.tolist()]
-            return []
+            return
        if name.endswith("ple_embedding.ngram_heads_offsets"):
            self._ple_head_offsets = [int(x) for x in data_torch.tolist()]
-            return []
+            return
        if name.endswith("ple_embedding.ngram_heads_vocab_sizes"):
            self._ple_head_vocab_sizes = [int(x) for x in data_torch.tolist()]
-            return []
+            return

        if ".ngram_embedding.shard_" in name:
-            return self._place_ple_shard(data_torch, name)
+            yield from self._place_ple_shard(data_torch, name)
+            return

        # eh_proj([e ; h_s]) = fc_embedding(e) + fc_hidden(h_s) for every hc stream s
        if name.endswith((".nextn_fc_embedding.weight", ".nextn_fc_hidden.weight")):
            self._mtp_fc[name.rsplit(".", 2)[1]] = data_torch
            if len(self._mtp_fc) < 2:
-                return []
+                return
            eh = torch.cat([self._mtp_fc.pop("nextn_fc_embedding"), self._mtp_fc.pop("nextn_fc_hidden")], dim=1)
-            return [(self.format_tensor_name(gguf.MODEL_TENSOR.NEXTN_EH_PROJ, bid, ".weight"), eh)]
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.NEXTN_EH_PROJ, bid, ".weight"), eh)
+            return

        # one projection feeds indexer q and k; split it, as minimax-m3 does
        if ".indexer.index_qk_proj.weight" in name:
            n_q = self.hparams["indexer_n_heads"] * self.hparams["indexer_head_dim"]
-            q = data_torch[:n_q]
-            k = data_torch[n_q:]
-            return [
-                (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid, ".weight"), q),
-                (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid, ".weight"), k),
-            ]
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid, ".weight"), data_torch[:n_q])
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid, ".weight"), data_torch[n_q:])
+            return

        # Gemma zero-centred gammas the inherited norm.weight rule misses
        if name.endswith((".ple.norm_key.weight", ".ple.norm_query.weight", ".ple.norm_conv.weight",
                          ".indexer.q_layernorm.weight", ".indexer.k_layernorm.weight")):
-            return [(self.map_tensor_name(name), data_torch + 1)]
+            yield (self.map_tensor_name(name), data_torch + 1)
+            return

        if name.endswith(".ple.conv1d.weight"):
-            return [(self.map_tensor_name(name), data_torch.squeeze())]
+            yield (self.map_tensor_name(name), data_torch.squeeze())
+            return

-        return super().modify_tensors(data_torch, name, bid)
+        for new_name, data in super().modify_tensors(data_torch, name, bid):
+            yield from self._merge_shared_input(new_name, data, bid)
+
+    # matmuls that read the same input, written as one weight [a | b] so the runtime does one matmul:
+    # (merged tensor, (a, b), linear-attention layers only)
+    _SHARED_INPUT_MERGES = (
+        (gguf.MODEL_TENSOR.HC_ATTN_DOWN_INJECT, (gguf.MODEL_TENSOR.HC_ATTN_DOWN, gguf.MODEL_TENSOR.HC_ATTN_INJECT), False),
+        (gguf.MODEL_TENSOR.HC_FFN_DOWN_INJECT,  (gguf.MODEL_TENSOR.HC_FFN_DOWN,  gguf.MODEL_TENSOR.HC_FFN_INJECT),  False),
+        (gguf.MODEL_TENSOR.ATTN_QKV_GATE,       (gguf.MODEL_TENSOR.ATTN_QKV,     gguf.MODEL_TENSOR.ATTN_GATE),      True),
+        (gguf.MODEL_TENSOR.SSM_BETA_ALPHA_CAT,  (gguf.MODEL_TENSOR.SSM_BETA,     gguf.MODEL_TENSOR.SSM_ALPHA),      True),
+    )
+
+    # caches one half of a merged weight until the other half arrives, then yields [a | b]
+    def _merge_shared_input(self, new_name: str, data: Tensor, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # a pre-quantized checkpoint carries per-tensor scales, which one merged weight cannot hold
+        if bid is None or self.hparams.get("quantization_config"):
+            yield (new_name, data)
+            return
+
+        layer_types = self.hparams.get("layer_types") or []
+        is_linear = bid < len(layer_types) and layer_types[bid] == "linear_attention"
+
+        match = next(((merged, parts.index(part)) for merged, parts, linear_only in self._SHARED_INPUT_MERGES
+                      for part in parts
+                      if (is_linear or not linear_only) and self.match_model_tensor_name(new_name, part, bid)), None)
+        if match is None:
+            yield (new_name, data)
+            return
+
+        merged, idx = match
+        halves = self._merge_parts.setdefault((bid, merged), {})
+        halves[idx] = data
+        if len(halves) == 2:
+            del self._merge_parts[(bid, merged)]
+            yield (self.format_tensor_name(merged, bid), torch.cat([halves[0], halves[1]], dim=0))

    # the shards concatenate into a tensor of well over 100 GB
    # use LazyChunkedTensor here, a single shard resident at a time
@@ -214,6 +251,8 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
        super().prepare_tensors()
        if self._mtp_fc:
            raise ValueError(f"MTP projection missing its other half: {sorted(self._mtp_fc)}")
+        if self._merge_parts:
+            raise ValueError(f"merged weights missing a half: {sorted((bid, t.name) for bid, t in self._merge_parts)}")
        n_parts = self.hparams.get("split_ngram_parts", 0)
        if self._ple_shards and len(self._ple_shards) != n_parts:
            raise ValueError(
@@ -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:
@@ -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;
@@ -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;
@@ -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)
+}
@@ -133,7 +133,7 @@ size_t opt_mbuf_total  = 0;                        // total buffer space limit (

 int    opt_mm_select     = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
 int    opt_fa_select     = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
-int    opt_fa_head_split = 1; // 1 = partition flash_attn by KV heads in multicore (default on), 0 = token-based (original)
+int    opt_fa_split      = 1; // 1 = head-split + team-split when n_kv_heads < N_cores (default), 0 = off
 int    opt_gdn_select    = 2; // 2 = HMX -> HVX, 1 = HVX, 0 = CPU (unsupported)
 int    opt_ar_select     = 2; // 2 = fused ALLREDUCE+ADD (default), 1 = unfused ALLREDUCE, 0 = fallback to CPY+FENCE
 int    opt_ar_scatter    = 1; // 1 = reduce-scatter the fused ALLREDUCE+ADD (default), 0 = full reduction
@@ -2397,7 +2397,7 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg
                }
            }

-            GGML_ASSERT(htp_kparams_solve(sess, node.node, node.kernel_params));
+            GGML_ASSERT(htp_kparams_solve(sess, node));
            if (node.opcode == HTP_OP_CPY && node.kparams<htp_copy_kernel_params>()->total_elems == 0) {
                continue;
            }
@@ -3393,7 +3393,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
    const char * str_nhmx     = getenv("GGML_HEXAGON_NHMX");
    const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT");
    const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT");
-    const char * str_fa_head_split = getenv("GGML_HEXAGON_FA_HEAD_SPLIT");
+    const char * str_fa_split  = getenv("GGML_HEXAGON_FA_SPLIT");
    const char * str_gdn_select = getenv("GGML_HEXAGON_GDN_SELECT");
    const char * str_ar_select = getenv("GGML_HEXAGON_AR_SELECT");
    const char * str_ar_scatter = getenv("GGML_HEXAGON_AR_SCATTER");
@@ -3458,7 +3458,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
    opt_nhmx      = str_nhmx     ? atoi(str_nhmx)                         : opt_nhmx;
    opt_mm_select = str_mm_select ? atoi(str_mm_select)                   : opt_mm_select;
    opt_fa_select = str_fa_select ? atoi(str_fa_select)                   : opt_fa_select;
-    opt_fa_head_split = str_fa_head_split ? atoi(str_fa_head_split)       : opt_fa_head_split;
+    opt_fa_split  = str_fa_split  ? atoi(str_fa_split)                   : opt_fa_split;
    opt_gdn_select    = str_gdn_select    ? atoi(str_gdn_select)          : opt_gdn_select;
    opt_ar_select     = str_ar_select     ? atoi(str_ar_select)           : opt_ar_select;
    opt_ar_scatter    = str_ar_scatter    ? atoi(str_ar_scatter)          : opt_ar_scatter;
@@ -28,7 +28,7 @@ extern size_t opt_mbuf_total;

 extern int    opt_mm_select;
 extern int    opt_fa_select;
-extern int    opt_fa_head_split;
+extern int    opt_fa_split;
 extern int    opt_gdn_select;
 extern int    opt_ar_select;
 extern int    opt_ar_scatter;
@@ -192,15 +192,20 @@ static bool fuse_mul_mat_add(const ggml_hexagon_session * sess, htp_opnode & pre
    if (!htp_match_consumer(prev, cur, /*commutative=*/ true, &src2)) return false;
    if (src2->type != GGML_TYPE_F32) return false;

-    const ggml_tensor * src0 = prev.src0();
-    const ggml_tensor * src1 = prev.src1();
-    const ggml_tensor * dst  = cur.dst();
-    if (!ggml_are_same_shape(prev.dst(), dst)) 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) ||
-        orig_kparams->collapse != kparams.collapse) {
+    if (!htp_kparams_solve_matmul_add(sess, src0, src1, src2, dst, &kparams)) {
        return false;
    }

@@ -222,6 +227,7 @@ static bool fuse_mul_mat_add(const ggml_hexagon_session * sess, htp_opnode & pre
    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;
@@ -233,36 +239,31 @@ static bool fuse_mul_mat_nx(const ggml_hexagon_session * sess, htp_opnode & prev
    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    = prev.inputs[curr_n];
-    const ggml_tensor * w_in = cur.src0();
-    const ggml_tensor * x_in = cur.src1();
-    const ggml_tensor * d_in = cur.dst();
+    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 || w_in->type != w0->type || w_in->ne[0] != w0->ne[0] ||
+    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 = orig_kparams->collapse && htp_matmul_can_collapse(w_in, x, d_in);
-    if ((x->ne[2] > 1 || x->ne[3] > 1) && !collapse) 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;
-        if (collapse) {
-            const ggml_tensor x_collapsed = htp_tensor_collapse_rows(x);
-            if (!htp_kparams_solve_mmnx(sess, w0, &x_collapsed, curr_n + 1, &kparams)) return false;
-            kparams.collapse = 1;
-        } else if (!htp_kparams_solve_mmnx(sess, w0, x, curr_n + 1, &kparams)) {
-            return false;
-        }
+        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);
@@ -270,7 +271,7 @@ static bool fuse_mul_mat_nx(const ggml_hexagon_session * sess, htp_opnode & prev
    }

    prev.inputs.insert(prev.inputs.begin() + curr_n, w_in);
-    prev.outputs.push_back(d_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;
 }
@@ -415,15 +416,13 @@ static bool fuse_unary_mul_glu(const ggml_hexagon_session * sess, htp_opnode & p
    const ggml_tensor * dst  = cur.dst();
    if (htp_inputs_overlap_unsafe(prev, dst, up)) return false;

-    struct htp_glu_kernel_params kparams;
-    if (!htp_kparams_solve_unary_mul_glu(sess, glu_op, gate, up, dst, &kparams)) 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));
-    prev.set_kparams(kparams);
    return true;
 }

@@ -1,6 +1,7 @@
 #define _USE_MATH_DEFINES

 #include "htp-kparams.h"
+#include "htp-opnode.h"
 #include "ggml-impl.h"

 #include <algorithm>
@@ -73,6 +74,7 @@ htp_op_code htp_op_remap(const struct ggml_tensor * t) {
                case GGML_UNARY_OP_RELU:       return HTP_OP_UNARY_RELU;
                case GGML_UNARY_OP_STEP:       return HTP_OP_UNARY_STEP;
                case GGML_UNARY_OP_CEIL:       return HTP_OP_UNARY_CEIL;
+                case GGML_UNARY_OP_SGN:        return HTP_OP_UNARY_SGN;
                default:                       break;
            }
            break;
@@ -396,9 +398,20 @@ static bool htp_kparams_solve_flash_attn(
    kparams->scale         = scale;
    kparams->max_bias      = max_bias;
    kparams->logit_softcap = logit_softcap;
-    kparams->head_split    = (opt_fa_head_split != 0) ? 1 : 0;
+    kparams->head_split    = (opt_fa_split != 0) ? 1 : 0;
+    kparams->flags         = 0;
    kparams->G             = G;

+    // team-split only when a team has >1 member (n_kv_heads < N_cores)
+    uint32_t fa_teams = 1, fa_team_size = 1;
+    const uint32_t mdev_count = htp_session_mdev_count(sess);
+    if (opt_fa_split && mdev_count > 1) {
+        htp_fa_mdev_teams(n_kv_heads, mdev_count, kparams->head_split != 0, &fa_teams, &fa_team_size);
+        if (fa_team_size > 1) {
+            kparams->flags |= HTP_FA_FLAG_TEAM_SPLIT;
+        }
+    }
+
    const uint32_t n_head = q->ne[2];
    kparams->n_head_log2 = 1u << (uint32_t) std::floor(std::log2(n_head));
    kparams->m0 = std::pow(2.0f, -(max_bias) / kparams->n_head_log2);
@@ -408,7 +421,7 @@ static bool htp_kparams_solve_flash_attn(
        const uint32_t DK_pad = hex_round_up(DK, 64);
        const uint32_t DV_pad = hex_round_up(DV, 64);
        size_t Br = 0, Bc = 0;
-        int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, vtcm_size, n_threads, (q->type == GGML_TYPE_F32), sinks != nullptr, n_head);
+        int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, vtcm_size, n_threads, (q->type == GGML_TYPE_F32), sinks != nullptr, n_head, fa_team_size);
        if (ret == 0) {
            kparams->kernel_type = HTP_FA_KERNEL_HMX;
            if (logit_softcap == 0.0f) {
@@ -1011,14 +1024,12 @@ static bool htp_kparams_solve_matmul(
    if (htp_matmul_can_collapse(src0, src1, dst)) {
        const ggml_tensor src1_collapsed = htp_tensor_collapse_rows(src1);
        const ggml_tensor dst_collapsed  = htp_tensor_collapse_rows(dst);
-        const bool ok = htp_kparams_solve_matmul_impl(sess, src0, &src1_collapsed, &dst_collapsed, 0, 0, kparams);
-        kparams->collapse = 1;
-        return ok;
+        return htp_kparams_solve_matmul_impl(sess, src0, &src1_collapsed, &dst_collapsed, 0, 0, kparams);
    }
    return htp_kparams_solve_matmul_impl(sess, src0, src1, dst, 0, 0, kparams);
 }

-static bool htp_matmul_add_can_collapse(
+bool htp_matmul_add_can_collapse(
    const struct ggml_tensor * src0,
    const struct ggml_tensor * src1,
    const struct ggml_tensor * src2,
@@ -1056,9 +1067,7 @@ bool htp_kparams_solve_matmul_add(
                                           : (src2 ? *src2 : ggml_tensor{});
        const struct ggml_tensor * p_src2 = src2 ? &src2_collapsed : nullptr;
        const size_t src2_size = p_src2 ? hex_round_up(ggml_nbytes(p_src2), 128) : 0;
-        const bool ok = htp_kparams_solve_matmul_impl(sess, src0, &src1_collapsed, &dst_collapsed, p_src2 ? p_src2->nb[1] : 0, src2_size, kparams);
-        kparams->collapse = 1;
-        return ok;
+        return htp_kparams_solve_matmul_impl(sess, src0, &src1_collapsed, &dst_collapsed, p_src2 ? p_src2->nb[1] : 0, src2_size, kparams);
    }
    const size_t src2_size = src2 ? hex_round_up(ggml_nbytes(src2), 128) : 0;
    return htp_kparams_solve_matmul_impl(sess, src0, src1, dst, src2 ? src2->nb[1] : 0, src2_size, kparams);
@@ -1392,6 +1401,108 @@ static bool htp_kparams_solve_binary(
    return true;
 }

+// collapse contiguous dims of elementwise ops so that rows fill at least one HVX vector
+static void htp_collapse_dims(struct ggml_tensor ** tensors, size_t n_tensors, size_t elem_size, uint32_t n_threads) {
+    if (n_tensors == 0 || !tensors[0]) {
+        return;
+    }
+    for (size_t t = 0; t < n_tensors; t++) {
+        if (!tensors[t] || tensors[t]->nb[0] != elem_size) {
+            return;
+        }
+        if (t > 0 && !ggml_are_same_shape(tensors[0], tensors[t])) {
+            return;
+        }
+    }
+
+    struct ggml_tensor * t0 = tensors[0];
+
+    while (t0->ne[1] > 1 || t0->ne[2] > 1 || t0->ne[3] > 1) {
+        const size_t row_size = (size_t) t0->ne[0] * elem_size;
+        bool rows_contig = true;
+        for (size_t t = 0; t < n_tensors; t++) {
+            if (tensors[t]->nb[1] != (size_t) tensors[t]->ne[0] * elem_size) {
+                rows_contig = false;
+                break;
+            }
+        }
+        bool all_ne1_one = true;
+        for (size_t t = 0; t < n_tensors; t++) {
+            if (tensors[t]->ne[1] != 1) {
+                all_ne1_one = false;
+                break;
+            }
+        }
+        const bool can_merge = all_ne1_one ||
+                               (rows_contig && (t0->ne[0] == 1 || (row_size < 128 && row_size * t0->ne[1] <= 2048)));
+        if (!can_merge) {
+            break;
+        }
+        for (size_t t = 0; t < n_tensors; t++) {
+            tensors[t]->ne[0] *= tensors[t]->ne[1];
+            tensors[t]->ne[1]  = tensors[t]->ne[2];
+            tensors[t]->ne[2]  = tensors[t]->ne[3];
+            tensors[t]->ne[3]  = 1;
+            tensors[t]->nb[1]  = tensors[t]->nb[2];
+            tensors[t]->nb[2]  = tensors[t]->nb[3];
+            tensors[t]->nb[3]  = tensors[t]->nb[2] * tensors[t]->ne[2];
+        }
+    }
+
+    // ne1 too big to merge in full: merge it in chunks of k rows (k divides ne1)
+    const size_t row_size = (size_t) t0->ne[0] * elem_size;
+    bool rows_contig = true;
+    for (size_t t = 0; t < n_tensors; t++) {
+        if (tensors[t]->nb[1] != row_size) {
+            rows_contig = false;
+            break;
+        }
+    }
+    if (t0->ne[1] > 1 && row_size < 128 && rows_contig) {
+        int64_t k = 2048 / row_size;
+        while (k > 1 && t0->ne[1] % k != 0) {
+            k--;
+        }
+        if (k > 1) {
+            for (size_t t = 0; t < n_tensors; t++) {
+                tensors[t]->ne[0] *= k;
+                tensors[t]->ne[1] /= k;
+                tensors[t]->nb[1] *= k;
+                tensors[t]->nb[2]  = tensors[t]->nb[1] * tensors[t]->ne[1];
+                tensors[t]->nb[3]  = tensors[t]->nb[2] * tensors[t]->ne[2];
+            }
+        }
+    }
+
+    // nrows too small to saturate threads: split ne0 into k rows (k divides ne0)
+    const size_t   row_size_curr = (size_t) t0->ne[0] * elem_size;
+    const uint32_t nrows         = (uint32_t) (t0->ne[1] * t0->ne[2] * t0->ne[3]);
+    rows_contig = true;
+    for (size_t t = 0; t < n_tensors; t++) {
+        if (tensors[t]->ne[1] != 1 && tensors[t]->nb[1] != row_size_curr) {
+            rows_contig = false;
+            break;
+        }
+    }
+    if (nrows < n_threads && rows_contig && row_size_curr >= 4096) {
+        const size_t vec_elems = 128 / elem_size;
+        int64_t k = std::max<int64_t>(row_size_curr / 16384, (n_threads * 4 + nrows - 1) / nrows);
+        k = std::min<int64_t>(k, row_size_curr / 2048);
+        while (k > 1 && t0->ne[0] % (k * vec_elems) != 0) {
+            k--;
+        }
+        if (k > 1) {
+            for (size_t t = 0; t < n_tensors; t++) {
+                tensors[t]->ne[0] /= k;
+                tensors[t]->ne[1] *= k;
+                tensors[t]->nb[1]  = tensors[t]->ne[0] * elem_size;
+                tensors[t]->nb[2]  = tensors[t]->nb[1] * tensors[t]->ne[1];
+                tensors[t]->nb[3]  = tensors[t]->nb[2] * tensors[t]->ne[2];
+            }
+        }
+    }
+}
+
 static bool htp_kparams_solve_unary_impl(
    const struct ggml_hexagon_session * sess,
    uint32_t op,
@@ -1435,6 +1546,7 @@ static bool htp_kparams_solve_unary_impl(
                case HTP_OP_UNARY_ABS:
                case HTP_OP_UNARY_STEP:
                case HTP_OP_UNARY_CEIL:
+                case HTP_OP_UNARY_SGN:
                    break;
                default:
                    return false;
@@ -1448,14 +1560,6 @@ static bool htp_kparams_solve_unary_impl(

    const size_t elem_size = ggml_type_size(src0->type);

-    struct ggml_tensor src0_c = *src0;
-    struct ggml_tensor dst_c  = *dst;
-    if (htp_unary_is_elementwise(op)) {
-        HTP_UNARY_COLLAPSE_DIMS(src0_c.ne, src0_c.nb, dst_c.ne, dst_c.nb, elem_size);
-        src0 = &src0_c;
-        dst  = &dst_c;
-    }
-
    const size_t dst_data_row_size  = dst->ne[0]  * ggml_type_size(dst->type);

    const uint32_t src0_nrows          = src0->ne[1] * src0->ne[2] * src0->ne[3];
@@ -1469,9 +1573,14 @@ static bool htp_kparams_solve_unary_impl(
    const bool broadcast_weight = (op == HTP_OP_RMS_NORM_MUL) && (src1->ne[1] * src1->ne[2] * src1->ne[3] == 1);
    kparams->broadcast_weight   = broadcast_weight;

-    const bool     src0_contig = ggml_is_contiguous_1(src0);
-    const bool     src1_contig = (op == HTP_OP_RMS_NORM_MUL && !broadcast_weight) ? ggml_is_contiguous_1(src1) : true;
-    const uint32_t max_block   = (src0_contig && src1_contig) ? src0_nrows : (uint32_t) src0->ne[1];
+    const bool     src0_contig           = ggml_is_contiguous_1(src0);
+    const bool     src1_contig           = (op == HTP_OP_RMS_NORM_MUL && !broadcast_weight) ? ggml_is_contiguous_1(src1) : true;
+    const uint32_t src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads;
+    const uint32_t min_block_rows        = (uint32_t) (std::max<size_t>(1, 2048 / dst_data_row_size));
+    uint32_t       max_block             = (src0_contig && src1_contig) ? src0_nrows : (uint32_t) src0->ne[1];
+    if (src0_nrows_per_thread >= 2 * min_block_rows) {
+        max_block = (std::min)(max_block, (std::max)(min_block_rows, src0_nrows_per_thread / 4));
+    }

    kparams->src0_contig = src0_contig ? 1 : 0;
    kparams->src1_contig = src1_contig ? 1 : 0;
@@ -1936,11 +2045,17 @@ bool htp_kparams_solve_unary_mul_glu(
    const size_t   vtcm_size           = htp_session_vtcm_size(sess);

    const size_t   row_size_aligned    = hex_round_up(dst_row_size, 128);
-    const uint32_t vtcm_row_per_thread = (uint32_t) (vtcm_size / (n_threads * 3 * row_size_aligned));
+    uint32_t       vtcm_row_per_thread = (uint32_t) (vtcm_size / (n_threads * 3 * row_size_aligned));
    if (vtcm_row_per_thread < 2) {
        return false;
    }

+    const uint32_t src0_nrows_per_thread = (src0_nrows + n_threads - 1) / n_threads;
+    if (src0_nrows_per_thread >= 2 && vtcm_row_per_thread > src0_nrows_per_thread) {
+        uint32_t target_rpt = (src0_nrows_per_thread >= 4) ? (src0_nrows_per_thread / 2) & ~1u : 2;
+        vtcm_row_per_thread = (std::max)(2u, (std::min)(vtcm_row_per_thread, target_rpt));
+    }
+
    struct htp_glu_vtcm_layout L;
    htp_glu_vtcm_layout_build(&L, dst_row_size, n_threads, vtcm_row_per_thread);

@@ -3182,12 +3297,131 @@ static inline bool htp_kparams_invoke(
    return SolveFn(sess, op, kparams ? (KParams *) kparams : &tmp);
 }

+void htp_opnode_collapse(const struct ggml_hexagon_session * sess, htp_opnode & node) {
+    const uint32_t n_threads = htp_session_n_threads(sess);
+
+    if (htp_op_is_unary(node.opcode) && htp_unary_is_elementwise(node.opcode)) {
+        if (node.inputs.empty() || !node.inputs[0] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        if (!ggml_are_same_shape(node.inputs[0], node.outputs[0])) {
+            return;
+        }
+        const size_t elem_size = ggml_type_size(node.inputs[0]->type);
+        struct ggml_tensor src0_c = *node.inputs[0];
+        struct ggml_tensor dst_c  = *node.outputs[0];
+        struct ggml_tensor * tensors[2] = { &src0_c, &dst_c };
+        htp_collapse_dims(tensors, 2, elem_size, n_threads);
+        if (src0_c.ne[0] != node.inputs[0]->ne[0] || src0_c.ne[1] != node.inputs[0]->ne[1]) {
+            node.inputs[0]  = node.add_dummy(src0_c);
+            node.outputs[0] = node.add_dummy(dst_c);
+        }
+        return;
+    }
+
+    if (htp_op_is_glu(node.opcode)) {
+        if (node.inputs.size() < 2 || !node.inputs[0] || !node.inputs[1] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        if (!ggml_are_same_shape(node.inputs[0], node.outputs[0]) ||
+            !ggml_are_same_shape(node.inputs[1], node.outputs[0])) {
+            return;
+        }
+        const size_t elem_size = ggml_type_size(node.inputs[0]->type);
+        struct ggml_tensor src0_c = *node.inputs[0];
+        struct ggml_tensor src1_c = *node.inputs[1];
+        struct ggml_tensor dst_c  = *node.outputs[0];
+        struct ggml_tensor * tensors[3] = { &src0_c, &src1_c, &dst_c };
+        htp_collapse_dims(tensors, 3, elem_size, n_threads);
+        if (src0_c.ne[0] != node.inputs[0]->ne[0] || src0_c.ne[1] != node.inputs[0]->ne[1]) {
+            node.inputs[0]  = node.add_dummy(src0_c);
+            node.inputs[1]  = node.add_dummy(src1_c);
+            node.outputs[0] = node.add_dummy(dst_c);
+        }
+        return;
+    }
+
+    if (node.opcode == HTP_OP_MUL_MAT) {
+        if (node.inputs.size() < 2 || !node.inputs[0] || !node.inputs[1] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        const struct ggml_tensor * src0 = node.inputs[0];
+        const struct ggml_tensor * src1 = node.inputs[1];
+        const struct ggml_tensor * dst  = node.outputs[0];
+        if (htp_matmul_can_collapse(src0, src1, dst)) {
+            node.inputs[1]  = node.add_dummy(htp_tensor_collapse_rows(src1));
+            node.outputs[0] = node.add_dummy(htp_tensor_collapse_rows(dst));
+        }
+        return;
+    }
+
+    if (node.opcode == HTP_OP_MUL_MAT_ADD) {
+        if (node.inputs.size() < 2 || !node.inputs[0] || !node.inputs[1] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        const struct ggml_tensor * src0 = node.inputs[0];
+        const struct ggml_tensor * src1 = node.inputs[1];
+        const struct ggml_tensor * src2 = node.inputs.size() > 2 ? node.inputs[2] : nullptr;
+        const struct ggml_tensor * dst  = node.outputs[0];
+        if (htp_matmul_add_can_collapse(src0, src1, src2, dst)) {
+            node.inputs[1]  = node.add_dummy(htp_tensor_collapse_rows(src1));
+            node.outputs[0] = node.add_dummy(htp_tensor_collapse_rows(dst));
+            if (src2 && (src2->ne[1] * src2->ne[2] * src2->ne[3] > 1)) {
+                node.inputs[2] = node.add_dummy(htp_tensor_collapse_rows(src2));
+            }
+        }
+        return;
+    }
+}
+
+bool htp_kparams_solve(
+    const struct ggml_hexagon_session * sess,
+    htp_opnode & node
+) {
+    htp_opnode_collapse(sess, node);
+
+    const ggml_tensor * src0 = node.inputs.empty() ? nullptr : node.inputs[0];
+    const ggml_tensor * src1 = node.inputs.size() > 1 ? node.inputs[1] : nullptr;
+    const ggml_tensor * dst  = node.dst();
+
+    if (htp_op_is_unary(node.opcode)) {
+        return htp_kparams_solve_unary_impl(sess, node.opcode, src0, src1, dst, node.kparams<htp_unary_kernel_params>());
+    }
+    if (htp_op_is_glu(node.opcode)) {
+        return htp_kparams_solve_unary_mul_glu(sess, node.opcode, src0, src1, dst, node.kparams<htp_glu_kernel_params>());
+    }
+    if (node.opcode == HTP_OP_MUL_MAT) {
+        if (!htp_matmul_check_types(src0, src1, dst)) {
+            return false;
+        }
+        return htp_kparams_solve_matmul_impl(sess, src0, src1, dst, 0, 0, node.kparams<htp_mm_kernel_params>());
+    }
+    if (node.opcode == HTP_OP_MUL_MAT_ADD) {
+        const ggml_tensor * src2 = node.inputs.size() > 2 ? node.inputs[2] : nullptr;
+        const size_t src2_size = src2 ? hex_round_up(ggml_nbytes(src2), 128) : 0;
+        return htp_kparams_solve_matmul_impl(sess, src0, src1, dst, src2 ? src2->nb[1] : 0, src2_size, node.kparams<htp_mm_kernel_params>());
+    }
+
+    return htp_kparams_solve(sess, node.node, node.kernel_params);
+}
+
 bool htp_kparams_solve(
    const struct ggml_hexagon_session * sess,
    const struct ggml_tensor * op,
    void * kparams
 ) {
    const htp_op_code opcode = htp_op_remap(op);
+    if (htp_op_is_unary(opcode) || htp_op_is_glu(opcode)) {
+        htp_opnode tmp(opcode, op);
+        if (!htp_kparams_solve(sess, tmp)) {
+            return false;
+        }
+        if (kparams) {
+            memcpy(kparams, tmp.kernel_params, sizeof(tmp.kernel_params));
+        }
+        return true;
+    }
+
    switch (opcode) {
        case HTP_OP_MUL_MAT:
        case HTP_OP_MUL_MAT_ID:
@@ -34,6 +34,7 @@

 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>
@@ -45,6 +46,18 @@ static inline bool htp_kparams_check(Args &&... args) {
    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,
@@ -2094,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;
@@ -32,9 +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;              // HTP_FA_FLAG_*
    uint8_t  n_threads;          // Number of threads to run

    // Common parameters
@@ -321,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
@@ -361,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)) {
@@ -113,6 +113,7 @@ enum htp_op_code {
    HTP_OP_GLU_REGLU,
    HTP_OP_GLU_SIGMOID,
    HTP_OP_UNARY_CEIL,
+    HTP_OP_UNARY_SGN,

    HTP_OP_INVALID
 };
@@ -677,6 +677,56 @@ static inline void hvx_step_f32_aa(uint8_t * restrict dst, const uint8_t * restr
    }
 }

+static inline void hvx_sgn_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    assert((unsigned long) src % 128 == 0);
+
+    HVX_Vector * restrict vdst = (HVX_Vector *) dst;
+    HVX_Vector * restrict vsrc = (HVX_Vector *) src;
+
+    const uint32_t elem_size = sizeof(float);
+    const uint32_t epv       = 128 / elem_size;
+    const uint32_t nvec      = n / epv;
+    const uint32_t nloe      = n % epv;
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        vdst[i] = hvx_vec_sgn_f32(vsrc[i]);
+    }
+    if (nloe) {
+        HVX_Vector v = hvx_vec_sgn_f32(vsrc[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v);
+    }
+}
+
+static inline void hvx_sgn_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    assert((unsigned long) src % 128 == 0);
+
+    HVX_Vector * restrict vdst = (HVX_Vector *) dst;
+    HVX_Vector * restrict vsrc = (HVX_Vector *) src;
+
+    const uint32_t elem_size = sizeof(_Float16);
+    const uint32_t epv       = 128 / elem_size;
+    const uint32_t nvec      = n / epv;
+    const uint32_t nloe      = n % epv;
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        HVX_VectorPair p = hvx_vec_f16_to_f32(vsrc[i]);
+        vdst[i] = hvx_vec_f32_to_f16(hvx_vec_sgn_f32(Q6_V_lo_W(p)), hvx_vec_sgn_f32(Q6_V_hi_W(p)));
+    }
+    if (nloe) {
+        HVX_VectorPair p = hvx_vec_f16_to_f32(vsrc[i]);
+        HVX_Vector v = hvx_vec_f32_to_f16(hvx_vec_sgn_f32(Q6_V_lo_W(p)), hvx_vec_sgn_f32(Q6_V_hi_W(p)));
+        hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v);
+    }
+}
+
 static inline void hvx_step_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
    assert((unsigned long) dst % 128 == 0);
    assert((unsigned long) src % 128 == 0);
@@ -118,6 +118,17 @@ static inline HVX_Vector hvx_vec_step_f32(HVX_Vector v) {
    return Q6_V_vmux_QVV(q, one, zero);
 }

+static inline HVX_Vector hvx_vec_sgn_f32(HVX_Vector v) {
+    const HVX_Vector zero       = Q6_V_vzero();
+    const HVX_Vector one        = hvx_vec_splat_f32(1.0f);
+    const HVX_Vector sign_mask  = Q6_V_vsplat_R(0x80000000);
+    const HVX_Vector mag        = hvx_vec_abs_f32(v);
+    const HVX_Vector signed_one = Q6_V_vor_VV(one, Q6_V_vand_VV(v, sign_mask));
+    // The ordered magnitude comparison maps NaN and both signed zeros to zero.
+    const HVX_VectorPred nonzero = Q6_Q_vcmp_gt_VsfVsf(mag, zero);
+    return Q6_V_vmux_QVV(nonzero, signed_one, zero);
+}
+
 static inline HVX_Vector hvx_vec_step_f16(HVX_Vector v) {
    const HVX_Vector zero = Q6_V_vzero();
    const HVX_Vector one  = hvx_vec_splat_f16((_Float16) 1.0f);
@@ -946,6 +946,7 @@ static int execute_op(struct htp_ops_context * octx) {
        case HTP_OP_UNARY_RELU:
        case HTP_OP_UNARY_STEP:
        case HTP_OP_UNARY_CEIL:
+        case HTP_OP_UNARY_SGN:
        case HTP_OP_L2_NORM:
        case HTP_OP_HADAMARD:
            return op_unary(octx);
@@ -1413,10 +1413,10 @@ static void hvx_mv_id_tiled(unsigned int nth, unsigned int ith, void * data) {

    for (uint32_t ie1 = 0; ie1 < n_aids; ++ie1) {  // for each expert
        const int32_t eid = *(const int32_t *) ((const uint8_t *) src2->data + ie1 * src2->nb[0]);
-        if (eid < 0) {
+        if (eid == -1) {
            continue;
        }
-        assert(eid < (int32_t) n_ids);
+        assert(eid >= 0 && eid < (int32_t) n_ids);

        const dma_addr_t src0_row = src0->data + eid * nb02;
        const uint8_t * restrict act_col = (const uint8_t *) act_data;
@@ -3775,16 +3775,6 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k
    return HTP_STATUS_OK;
 }

-static inline void htp_mm_tensor_collapse_rows(struct htp_tensor * c, const struct htp_tensor * t, uint32_t stride) {
-    *c = *t;
-    c->ne[1] = t->ne[1] * t->ne[2] * t->ne[3];
-    c->ne[2] = 1;
-    c->ne[3] = 1;
-    c->nb[1] = stride;
-    c->nb[2] = c->nb[1] * c->ne[1];
-    c->nb[3] = c->nb[2];
-}
-
 static int op_matmul_impl(struct htp_ops_context * octx) {
    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;

@@ -3801,36 +3791,6 @@ static int op_matmul_impl(struct htp_ops_context * octx) {
 }

 int op_matmul(struct htp_ops_context * octx) {
-    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
-
-    if (kparams->collapse) {
-        const struct htp_tensor * act  = octx->src[1];
-        const struct htp_tensor * dst  = octx->dst;
-        const uint32_t s_act = (act->ne[1] > 1) ? act->nb[1] : ((act->ne[2] > 1) ? act->nb[2] : act->nb[3]);
-        const uint32_t sd    = (dst->ne[1] > 1) ? dst->nb[1] : ((dst->ne[2] > 1) ? dst->nb[2] : dst->nb[3]);
-        struct htp_tensor act_collapsed, dst_collapsed;
-        htp_mm_tensor_collapse_rows(&act_collapsed, act, s_act);
-        htp_mm_tensor_collapse_rows(&dst_collapsed, dst, sd);
-        octx->src[1] = &act_collapsed;
-        octx->dst    = &dst_collapsed;
-
-        struct htp_tensor src2_collapsed;
-        const struct htp_tensor * src2 = octx->src[2];
-        if (src2 && (src2->ne[1] * src2->ne[2] * src2->ne[3] > 1)) {
-            const uint32_t s2 = (src2->ne[1] > 1) ? src2->nb[1] : ((src2->ne[2] > 1) ? src2->nb[2] : src2->nb[3]);
-            htp_mm_tensor_collapse_rows(&src2_collapsed, src2, s2);
-            octx->src[2] = &src2_collapsed;
-        }
-
-        const int status = op_matmul_impl(octx);
-        octx->src[1] = act;
-        octx->dst    = dst;
-        if (src2) {
-            octx->src[2] = src2;
-        }
-        return status;
-    }
-
    return op_matmul_impl(octx);
 }

@@ -4166,10 +4126,10 @@ static inline void scan_expert_ids_n(
        const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
        for (uint32_t id = 0; id < n_ids; ++id) {
            const int32_t i02 = row_ptr[id];
-            if (i02 < 0) {
+            if (i02 == -1) {
                continue;
            }
-            assert(i02 < n_as);
+            assert(i02 >= 0 && i02 < n_as);

            if (matrix_rows) {
                matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
@@ -4179,6 +4139,28 @@ static inline void scan_expert_ids_n(
    }
 }

+// the matmuls only write the rows that an expert owns, so the rows of the skipped slots are zeroed here
+static inline void zero_skipped_dst_rows(
+    const struct htp_tensor * ids,
+    uint32_t n_ids,
+    const struct htp_tensor * dst,
+    uint32_t ne0,
+    size_t nb1,
+    size_t nb2
+) {
+    const uint8_t * ids_data = (const uint8_t *) ids->data;
+
+    for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
+        for (uint32_t id = 0; id < n_ids; ++id) {
+            const int32_t i02 = *(const int32_t *) (ids_data + iid1 * ids->nb[1] + id * ids->nb[0]);
+            if (i02 != -1) {
+                continue;
+            }
+            memset((uint8_t *) (dst->data + id * nb1 + iid1 * nb2), 0, ne0 * sizeof(float));
+        }
+    }
+}
+
 static inline void scan_expert_ids(
    const struct htp_tensor * ids,
    uint32_t n_ids,
@@ -4204,10 +4186,10 @@ static inline void scan_expert_ids(
            const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
            for (uint32_t id = 0; id < n_ids; ++id) {
                const int32_t i02 = *(const int32_t *) ((const uint8_t *) row_ptr + id * ids_nb0);
-                if (i02 < 0) {
+                if (i02 == -1) {
                    continue;
                }
-                assert(i02 < n_as);
+                assert(i02 >= 0 && i02 < n_as);

                if (matrix_rows) {
                    matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
@@ -4287,6 +4269,8 @@ int op_matmul_id(struct htp_ops_context * octx) {
        scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
    }

+    zero_skipped_dst_rows(ids, n_ids, dst, ne0, nb1, nb2);
+
    mmctx->matrix_row_counts    = matrix_row_counts;
    mmctx->matrix_rows          = matrix_rows;
    mmctx->mapping_stride       = mapping_stride;
@@ -4538,35 +4522,5 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) {
 }

 int op_matmul_nx(struct htp_ops_context * octx) {
-    const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
-
-    if (kparams->collapse) {
-        const uint32_t n_weights = kparams->n_weights;
-        const struct htp_tensor * act = octx->src[n_weights];
-        const uint32_t s1 = (act->ne[1] > 1) ? act->nb[1] : ((act->ne[2] > 1) ? act->nb[2] : act->nb[3]);
-        struct htp_tensor act_collapsed;
-        struct htp_tensor dsts_collapsed[HTP_OP_MAX_OUTPUTS];
-        const struct htp_tensor * orig_dsts[HTP_OP_MAX_OUTPUTS];
-
-        htp_mm_tensor_collapse_rows(&act_collapsed, act, s1);
-        octx->src[n_weights] = &act_collapsed;
-
-        for (uint32_t p = 0; p < n_weights; p++) {
-            orig_dsts[p] = octx->dsts[p];
-            const struct htp_tensor * d = octx->dsts[p];
-            const uint32_t sd = (d->ne[1] > 1) ? d->nb[1] : ((d->ne[2] > 1) ? d->nb[2] : d->nb[3]);
-            htp_mm_tensor_collapse_rows(&dsts_collapsed[p], d, sd);
-            octx->dsts[p] = &dsts_collapsed[p];
-        }
-
-        const int status = op_matmul_nx_impl(octx);
-
-        octx->src[n_weights] = act;
-        for (uint32_t p = 0; p < n_weights; p++) {
-            octx->dsts[p] = orig_dsts[p];
-        }
-        return status;
-    }
-
    return op_matmul_nx_impl(octx);
 }
@@ -61,7 +61,7 @@ enum htp_mm_kernel_type {
 struct htp_mm_kernel_params {
    uint8_t  kernel_type;        // enum htp_mm_kernel_type
    uint8_t  pipeline;           // 1 = pipelined execution, 0 = standard
-    uint8_t  collapse;           // 1 = collapse outer dims into 2D, 0 = standard
+    uint8_t  unused;             // reserved
    uint8_t  n_hmx;              // 1 = use HMX, 0 = use HVX

    uint8_t  n_threads;          // Number of threads to spawn
@@ -440,6 +440,20 @@ static void ceil_f16(const void * restrict src,
    }
 }

+static void sgn_f16(const void * restrict src,
+                    void * restrict dst,
+                    const uint32_t num_rows,
+                    const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_sgn_f16_aa((uint8_t *) dst_local, (const uint8_t *) src_local, ne0);
+    }
+}
+
 static void l2_norm_f16(const void * restrict src,
                        void * restrict dst,
                        const uint32_t num_rows,
@@ -976,6 +990,20 @@ static void ceil_f32(const void * restrict src,
    }
 }

+static void sgn_f32(const void * restrict src,
+                    void * restrict dst,
+                    const uint32_t num_rows,
+                    const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_sgn_f32_aa(dst_local, src_local, ne0);
+    }
+}
+
 static void log_f32(const void * restrict src,
                    void * restrict dst,
                    const uint32_t num_rows,
@@ -1094,6 +1122,11 @@ static void tile_ceil_f32(void * restrict dst, const void * restrict src, uint32
    hvx_ceil_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw);
 }

+static void tile_sgn_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) {
+    (void) uctx;
+    hvx_sgn_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw);
+}
+
 static void tri_apply_tile_f32(const void * restrict src, void * restrict dst,
                               uint32_t tile_elems, uint32_t col_start, uint32_t i01,
                               uint32_t ne0, int32_t ttype) {
@@ -1738,6 +1771,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
        case HTP_OP_UNARY_RELU:      op_type = "relu-f32";                                   break;
        case HTP_OP_UNARY_STEP:      op_type = is_f16 ? "step-f16"     : "step-f32";         break;
        case HTP_OP_UNARY_CEIL:      op_type = is_f16 ? "ceil-f16"     : "ceil-f32";         break;
+        case HTP_OP_UNARY_SGN:       op_type = is_f16 ? "sgn-f16"      : "sgn-f32";          break;
        case HTP_OP_L2_NORM:         op_type = is_f16 ? "l2norm-f16"   : "l2norm-f32";       break;
        case HTP_OP_TRI:             op_type = "tri-f32";                                    break;
        case HTP_OP_HADAMARD:        op_type = is_f16 ? "hadamard-f16" : "hadamard-f32";     break;
@@ -1752,9 +1786,9 @@ static int execute_op_unary(struct htp_ops_context * octx) {
        return HTP_STATUS_INVAL_PARAMS;
    }

-    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
    const size_t src0_elem_size = is_f16 ? sizeof(_Float16) : sizeof(float);
    const size_t dst_elem_size  = (dst->type == HTP_TYPE_F16) ? sizeof(_Float16) : sizeof(float);
+    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
    const size_t src0_data_row_size = src0->ne[0] * src0_elem_size;
    const size_t dst_data_row_size  = dst->ne[0]  * dst_elem_size;

@@ -1872,6 +1906,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
            case HTP_OP_UNARY_RELU:      compute_func = (void *) tile_relu_f32;           break;
            case HTP_OP_UNARY_STEP:      compute_func = (void *) tile_step_f32;           break;
            case HTP_OP_UNARY_CEIL:      compute_func = (void *) tile_ceil_f32;           break;
+            case HTP_OP_UNARY_SGN:       compute_func = (void *) tile_sgn_f32;            break;
            case HTP_OP_TRI:
                task_func    = unary_thread_tiled_tri_f32;
                compute_func = (void *) tri_apply_tile_f32;
@@ -1892,6 +1927,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
            case HTP_OP_UNARY_LOG:       compute_func = (void *) log_f16;                 break;
            case HTP_OP_UNARY_STEP:      compute_func = (void *) step_f16;                break;
            case HTP_OP_UNARY_CEIL:      compute_func = (void *) ceil_f16;                break;
+            case HTP_OP_UNARY_SGN:       compute_func = (void *) sgn_f16;                 break;
            case HTP_OP_HADAMARD:        compute_func = (void *) hadamard_f16;            break;
            default:                     break;
        }
@@ -1923,6 +1959,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
            case HTP_OP_UNARY_RELU:      compute_func = (void *) relu_f32;                break;
            case HTP_OP_UNARY_STEP:      compute_func = (void *) step_f32;                break;
            case HTP_OP_UNARY_CEIL:      compute_func = (void *) ceil_f32;                break;
+            case HTP_OP_UNARY_SGN:       compute_func = (void *) sgn_f32;                 break;
            case HTP_OP_L2_NORM:         compute_func = (void *) l2_norm_f32;             break;
            case HTP_OP_HADAMARD:        compute_func = (void *) hadamard_f32;            break;
            case HTP_OP_TRI:
@@ -1947,25 +1984,8 @@ static int execute_op_unary(struct htp_ops_context * octx) {
 int op_unary(struct htp_ops_context * octx) {
    switch (octx->src[0]->type) {
        case HTP_TYPE_F32:
-        case HTP_TYPE_F16: {
-            if (!htp_unary_is_elementwise(octx->op)) {
-                return execute_op_unary(octx);
-            }
-
-            struct htp_tensor src0_c = *octx->src[0];
-            struct htp_tensor dst_c  = *octx->dst;
-            const size_t elem_size = (src0_c.type == HTP_TYPE_F16) ? sizeof(_Float16) : sizeof(float);
-            HTP_UNARY_COLLAPSE_DIMS(src0_c.ne, src0_c.nb, dst_c.ne, dst_c.nb, elem_size);
-
-            const struct htp_tensor * orig_src0 = octx->src[0];
-            const struct htp_tensor * orig_dst  = octx->dst;
-            octx->src[0] = &src0_c;
-            octx->dst    = &dst_c;
-            const int status = execute_op_unary(octx);
-            octx->src[0] = orig_src0;
-            octx->dst    = orig_dst;
-            return status;
-        }
+        case HTP_TYPE_F16:
+            return execute_op_unary(octx);

        default:
            return HTP_STATUS_NO_SUPPORT;
@@ -56,6 +56,7 @@ static inline bool htp_op_is_unary(uint32_t opcode) {
        case HTP_OP_UNARY_RELU:
        case HTP_OP_UNARY_STEP:
        case HTP_OP_UNARY_CEIL:
+        case HTP_OP_UNARY_SGN:
        case HTP_OP_L2_NORM:
        case HTP_OP_TRI:
        case HTP_OP_HADAMARD:
@@ -66,6 +67,9 @@ static inline bool htp_op_is_unary(uint32_t opcode) {
 }

 static inline bool htp_unary_is_elementwise(uint32_t opcode) {
+    if (!htp_op_is_unary(opcode)) {
+        return false;
+    }
    switch (opcode) {
        case HTP_OP_NORM:
        case HTP_OP_RMS_NORM:
@@ -79,31 +83,6 @@ static inline bool htp_unary_is_elementwise(uint32_t opcode) {
    }
 }

-#define HTP_UNARY_COLLAPSE_DIMS(src_ne, src_nb, dst_ne, dst_nb, elem_size)                   \
-    do {                                                                                     \
-        while (((src_ne)[1] > 1 || (src_ne)[2] > 1 || (src_ne)[3] > 1) &&                    \
-               (size_t) (src_nb)[0] == (elem_size) &&                                        \
-               (size_t) (dst_nb)[0] == (elem_size) &&                                        \
-               (((src_ne)[1] == 1 && (dst_ne)[1] == 1) ||                                    \
-                ((size_t) (src_nb)[1] == (size_t) (src_ne)[0] * (size_t) (src_nb)[0] &&      \
-                 (size_t) (dst_nb)[1] == (size_t) (dst_ne)[0] * (size_t) (dst_nb)[0] &&      \
-                 ((src_ne)[0] == 1 ||                                                        \
-                  (uint64_t) (src_ne)[0] * (uint64_t) (src_ne)[1] * (elem_size) <= 128)))) { \
-            (src_ne)[0] *= (src_ne)[1];                                                      \
-            (src_ne)[1]  = (src_ne)[2];                                                      \
-            (src_ne)[2]  = (src_ne)[3];                                                      \
-            (src_ne)[3]  = 1;                                                                \
-            (dst_ne)[0] *= (dst_ne)[1];                                                      \
-            (dst_ne)[1]  = (dst_ne)[2];                                                      \
-            (dst_ne)[2]  = (dst_ne)[3];                                                      \
-            (dst_ne)[3]  = 1;                                                                \
-            (src_nb)[1]  = (src_nb)[2];                                                      \
-            (src_nb)[2]  = (src_nb)[3];                                                      \
-            (dst_nb)[1]  = (dst_nb)[2];                                                      \
-            (dst_nb)[2]  = (dst_nb)[3];                                                      \
-        }                                                                                    \
-    } while (0)
-
 struct htp_unary_vtcm_layout {
    size_t total_bytes;
    size_t off_src0;
@@ -1225,40 +1225,6 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
    return res;
 }

-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax_part(ggml_metal_library_t lib) {
-    char base[256];
-    char name[256];
-
-    snprintf(base, 256, "kernel_mul_mm_id_amax_part_f32");
-    snprintf(name, 256, "%s", base);
-
-    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
-    if (!res.pipeline) {
-        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
-    }
-
-    res.smem = 32*sizeof(float);
-
-    return res;
-}
-
-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax(ggml_metal_library_t lib) {
-    char base[256];
-    char name[256];
-
-    snprintf(base, 256, "kernel_mul_mm_id_amax_f32");
-    snprintf(name, 256, "%s", base);
-
-    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
-    if (!res.pipeline) {
-        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
-    }
-
-    res.smem = 32*sizeof(float);
-
-    return res;
-}
-
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(ggml_metal_library_t lib, int ne02, int ne20) {
    char base[256];
    char name[256];
@@ -1283,6 +1249,15 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(g
    return res;
 }

+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_zero(ggml_metal_library_t lib) {
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, "kernel_mul_mm_id_zero");
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, "kernel_mul_mm_id_zero", "kernel_mul_mm_id_zero", nullptr);
+    }
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id(ggml_metal_library_t lib, const ggml_tensor * op) {
    char base[256];
    char name[256];
@@ -1292,18 +1267,14 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id(ggml_m

    const bool bc_inp = op->src[0]->ne[0] % 32 != 0;

-    // src1 prec [TAG_GGML_PREC]
-    const bool amax = ggml_get_op_params_i32(op, 3) == GGML_PREC_F32;
-
    snprintf(base, 256, "kernel_mul_mm_id_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1));
-    snprintf(name, 256, "%s_bci=%d_amax=%d", base, bc_inp, amax);
+    snprintf(name, 256, "%s_bci=%d", base, bc_inp);

    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
    if (!res.pipeline) {
        ggml_metal_cv_t cv = ggml_metal_cv_init();

        ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0);
-        ggml_metal_cv_set_bool(cv, amax,   FC_MUL_MM + 6);

        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);

@@ -140,8 +140,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_mm
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm            (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv            (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0    (ggml_metal_library_t lib, int ne02, int ne20);
-struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax(ggml_metal_library_t lib);
-struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax_part(ggml_metal_library_t lib);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_zero    (ggml_metal_library_t lib);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax            (ggml_metal_library_t lib, const struct ggml_tensor * op);
@@ -14,8 +14,6 @@
 #define N_MM_SIMD_GROUP_X 2
 #define N_MM_SIMD_GROUP_Y 2

-#define N_MM_NPART_AMAX 256
-
 #define N_ARGMAX_NPART_MAX 256
 #define N_ARGMAX_CHUNK 4096

@@ -573,12 +571,11 @@ typedef struct {
 } ggml_metal_kargs_mul_mm_id_map0;

 typedef struct {
-    int32_t  ne00;
-    int32_t  ne01;
-    int32_t  ne02;
-    uint64_t nb01;
-    uint64_t nb02;
-} ggml_metal_kargs_mul_mm_id_amax;
+    int32_t  ne0;
+    uint64_t nb1;
+    uint64_t nb2;
+    uint64_t nb21;
+} ggml_metal_kargs_mul_mm_id_zero;

 typedef struct {
    int32_t  ne00;
@@ -2812,15 +2812,6 @@ size_t ggml_metal_op_mul_mat_id_extra_ids(const ggml_tensor * op) {
    return ggml_type_size(GGML_TYPE_I32)*(ne20*ne21 + ne02);
 }

-size_t ggml_metal_op_mul_mat_id_extra_amax(const ggml_tensor * op) {
-    assert(op->op == GGML_OP_MUL_MAT_ID);
-
-    GGML_UNUSED(op);
-
-    // 2 scaling factors (8 bytes) + N_MM_NPART_AMAX per-threadgroup scales for stage-1
-    return 8 + N_MM_NPART_AMAX*sizeof(float);
-}
-
 int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
    ggml_tensor * op = ctx->node(idx);

@@ -2872,39 +2863,6 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
        ggml_metal_buffer_id bid_ids = bid_tpe;
        bid_ids.offs += ggml_metal_op_mul_mat_id_extra_tpe(op);

-        ggml_metal_buffer_id bid_amax = bid_ids;
-        bid_amax.offs += ggml_metal_op_mul_mat_id_extra_ids(op);
-
-        // src1 prec [TAG_GGML_PREC]
-        const bool use_amax = ggml_get_op_params_i32(op, 3) == GGML_PREC_F32;
-
-        // src1 rescale factors, computed before the matmul
-        // ref: https://github.com/ggml-org/llama.cpp/pull/26223
-        if (use_amax) {
-            ggml_metal_kargs_mul_mm_id_amax args = {
-                /*.ne00 =*/ ne10,
-                /*.ne01 =*/ ne11,
-                /*.ne02 =*/ ne12,
-                /*.nb01 =*/ nb11,
-                /*.nb02 =*/ nb12,
-            };
-
-            auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_amax_part(lib);
-
-            const size_t smem = pipeline.smem;
-
-            GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size);
-
-            ggml_metal_encoder_set_pipeline(enc, pipeline);
-            ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
-            ggml_metal_encoder_set_buffer  (enc, bid_src1, 1);
-            ggml_metal_encoder_set_buffer  (enc, bid_amax, 2);
-
-            ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
-
-            ggml_metal_encoder_dispatch_threadgroups(enc, N_MM_NPART_AMAX, 1, 1, 256, 1, 1);
-        }
-
        {
            ggml_metal_kargs_mul_mm_id_map0 args = {
                ne10,
@@ -2935,20 +2893,28 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
            ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, ne02, 1, 1);
        }

-        ggml_metal_op_concurrency_reset(ctx);
+        // zero the rows that no expert owns - disjoint from what kernel_mul_mm_id writes
+        {
+            ggml_metal_kargs_mul_mm_id_zero args = {
+                /*.ne0  =*/ ne0,
+                /*.nb1  =*/ nb1,
+                /*.nb2  =*/ nb2,
+                /*.nb21 =*/ nb21,
+            };

-        if (use_amax) {
-            auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_amax(lib);
+            auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_zero(lib);

            ggml_metal_encoder_set_pipeline(enc, pipeline);
-            ggml_metal_encoder_set_buffer  (enc, bid_amax, 0);
+            ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+            ggml_metal_encoder_set_buffer  (enc, bid_src2, 1);
+            ggml_metal_encoder_set_buffer  (enc, bid_dst,  2);

-            ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, 32, 1, 1);
-
-            // the next kernel has to wait for the amax data
-            ggml_metal_op_concurrency_reset(ctx);
+            ggml_metal_encoder_dispatch_threadgroups(enc, ne20, ne21, 1, 32, 1, 1);
        }

+        // this barrier is always needed because the next kernel has to wait for the id maps to be computed
+        ggml_metal_op_concurrency_reset(ctx);
+
        {
            auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id(lib, op);

@@ -2977,7 +2943,6 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
            ggml_metal_encoder_set_buffer  (enc, bid_tpe,  3);
            ggml_metal_encoder_set_buffer  (enc, bid_ids,  4);
            ggml_metal_encoder_set_buffer  (enc, bid_dst,  5);
-            ggml_metal_encoder_set_buffer  (enc, bid_amax, 6);

            const size_t smem = pipeline.smem;

@@ -36,7 +36,6 @@ size_t ggml_metal_op_mul_mat_id_extra_tpe(const struct ggml_tensor * op);

 // id map [n_tokens, n_expert]
 size_t ggml_metal_op_mul_mat_id_extra_ids(const struct ggml_tensor * op);
-size_t ggml_metal_op_mul_mat_id_extra_amax(const struct ggml_tensor * op);

 // return true if we should use the FA vector kernel for this op
 bool ggml_metal_op_flash_attn_ext_use_vec(const struct ggml_tensor * op);
@@ -226,7 +226,6 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_
            {
                res += ggml_metal_op_mul_mat_id_extra_tpe(tensor);
                res += ggml_metal_op_mul_mat_id_extra_ids(tensor);
-                res += ggml_metal_op_mul_mat_id_extra_amax(tensor);
            } break;
        case GGML_OP_FLASH_ATTN_EXT:
            {
@@ -185,6 +185,14 @@ kernel void kernel_add_id(

    device       float * dst_row  = (device       float *)((device char *)dst  +  i1*nb1       + i2*nb2);
    device const float * src0_row = (device const float *)((device char *)src0 +  i1*args.nb01 + i2*args.nb02);
+
+    if (i11 == -1) {
+        for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+            dst_row[i0] = src0_row[i0];
+        }
+        return;
+    }
+
    device const float * src1_row = (device const float *)((device char *)src1 + i11*args.nb11);

    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
@@ -7,7 +7,6 @@ constant short FC_mul_mm_ne12  [[function_constant(FC_MUL_MM + 2)]];
 constant short FC_mul_mm_ne13  [[function_constant(FC_MUL_MM + 3)]];
 constant short FC_mul_mm_r2    [[function_constant(FC_MUL_MM + 4)]];
 constant short FC_mul_mm_r3    [[function_constant(FC_MUL_MM + 5)]];
-constant bool FC_mul_mm_id_amax [[function_constant(FC_MUL_MM + 6)]];
 constant int32_t FC_mul_mm_id_map0_ne02 [[function_constant(FC_MUL_MM + 7)]];

 // each block_q contains 16*nl weights
@@ -391,8 +390,10 @@ kernel void kernel_mul_mm_id_map0(
        if (i21 + tpitg < args.ne21) {
            device const int32_t * src2_i32 = (device const int32_t *) (src2 + (i21 + tpitg)*args.nb21);

+            // 0xFFFF marks a skipped slot, it never matches an expert since ne02 fits in a threadgroup
            FOR_UNROLL (short i20 = 0; i20 < ne20; i20++) {
-                sids[tpitg*ne20 + i20] = src2_i32[i20];
+                const int32_t id = src2_i32[i20];
+                sids[tpitg*ne20 + i20] = id == -1 ? 0xFFFF : (uint16_t) id;
            }
        }

@@ -506,82 +507,27 @@ kernel void kernel_mul_mm_id_map0(
    }
 }

-kernel void kernel_mul_mm_id_amax_part_f32(
-        constant ggml_metal_kargs_mul_mm_id_amax & args,
-        device   const char * src1,
-        device         char * dst,
-        threadgroup    char * shmem [[threadgroup(0)]],
-        uint  tgpig[[threadgroup_position_in_grid]],
-        ushort tiitg[[thread_index_in_threadgroup]],
-        ushort tiisg[[thread_index_in_simdgroup]],
-        ushort sgitg[[simdgroup_index_in_threadgroup]],
-        ushort   ntg[[threads_per_threadgroup]]) {
-    const int nrow = args.ne01*args.ne02;
+// kernel_mul_mm_id only writes the rows that an expert owns, so the skipped rows are zeroed here
+kernel void kernel_mul_mm_id_zero(
+        constant ggml_metal_kargs_mul_mm_id_zero & args,
+        device const char * src2,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+    const int i20 = tgpig.x; // slot
+    const int i21 = tgpig.y; // token

-    float lmax = 0.0f;
+    const int32_t id = ((device const int32_t *) (src2 + i21*args.nb21))[i20];

-    for (int ir = tgpig; ir < nrow; ir += N_MM_NPART_AMAX) {
-        const int i01 = ir % args.ne01;
-        const int i02 = ir / args.ne01;
-
-        device const float * row = (device const float *) (src1 + i02*args.nb02 + i01*args.nb01);
-
-        for (int i00 = tiitg; i00 < args.ne00; i00 += ntg) {
-            lmax = max(lmax, fabs(row[i00]));
-        }
+    if (id != -1) {
+        return;
    }

-    float amax = simd_max(lmax);
+    device float * dst_row = (device float *) (dst + i20*args.nb1 + i21*args.nb2);

-    threadgroup float * shared_amax = (threadgroup float *) shmem;
-
-    if (ntg > N_SIMDWIDTH) {
-        if (sgitg == 0) {
-            shared_amax[tiisg] = 0.0f;
-        }
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        if (tiisg == 0) {
-            shared_amax[sgitg] = amax;
-        }
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-
-        amax = shared_amax[tiisg];
-        amax = simd_max(amax);
-    }
-
-    if (tiitg == 0) {
-        ((device float *) (dst + 8))[tgpig] = amax;
-    }
-}
-
-kernel void kernel_mul_mm_id_amax_f32(
-        device char * dst,
-        ushort tiitg[[thread_index_in_threadgroup]]) {
-    device const float * part = (device const float *) (dst + 8);
-
-    float amax = 0.0f;
-
-    for (int i = tiitg; i < N_MM_NPART_AMAX; i += N_SIMDWIDTH) {
-        amax = max(amax, part[i]);
-    }
-
-    amax = simd_max(amax);
-
-    if (tiitg == 0) {
-        // leave a comfortable margin below the f16 max of 65504
-        float scale = 1.0f;
-
-        // isfinite: src1 already inf/nan is not ours to fix - keep the
-        // scale at 1.0 instead of turning it into a different failure
-        if (isfinite(amax) && amax > 32768.0f) {
-            scale = exp2(ceil(log2(amax)) - 15.0f);
-        }
-
-        device float * d = (device float *) dst;
-
-        d[0] = 1.0f/scale; // exact: scale is a power of two
-        d[1] = scale;
+    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+        dst_row[i0] = 0.0f;
    }
 }

@@ -607,7 +553,6 @@ kernel void kernel_mul_mm_id(
        device const char * htpe,
        device const char * hids,
        device       char * dst,
-        device const char * amax,
        threadgroup  char * shmem [[threadgroup(0)]],
        uint3  tgpig[[threadgroup_position_in_grid]],
        ushort tiitg[[thread_index_in_threadgroup]],
@@ -663,10 +608,6 @@ kernel void kernel_mul_mm_id(
    // skip the upper half of the token tile when the expert did not fill it
    constexpr short NR1H = NR1/2;

-    // power-of-two rescaling
-    const float s1_inv   = FC_mul_mm_id_amax ? ((device const float *) amax)[0] : 1.0f;
-    const float s1_scale = FC_mul_mm_id_amax ? ((device const float *) amax)[1] : 1.0f;
-
 #ifndef GGML_METAL_HAS_TENSOR
    S0_8x8 ma[4];
    S1_8x8 mb[2];
@@ -772,7 +713,7 @@ kernel void kernel_mul_mm_id(

                const short ib = 4*sx + sy;

-                *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) (*((device T1 *) y + i) * (T1) s1_inv) : 0;
+                *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
            }
        } else {
            const short sx = (tiitg%NL1);
@@ -785,7 +726,7 @@ kernel void kernel_mul_mm_id(

            const short ib = 4*sx + sy;

-            *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)((*((device T1_2x4 *) y)) * (T1) s1_inv);
+            *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
        }
 #else
        // load data and store to threadgroup memory
@@ -833,7 +774,7 @@ kernel void kernel_mul_mm_id(
                //const short lx = (tiitg/NL1)%8;
                //const short ly = i;

-                *(sb + NK*(8*sy + ly) + 8*sx + lx) = loop_k + iy + i < args.ne00 ? (S1) (*((device T1 *) y + i) * (T1) s1_inv) : 0;
+                *(sb + NK*(8*sy + ly) + 8*sx + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
            }
        } else {
            const short sx = (tiitg%NL1);
@@ -844,7 +785,7 @@ kernel void kernel_mul_mm_id(
            //const short lx = (tiitg/NL1)%8;
            //const short ly = i;

-            *(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) = (S1_2x4)((*((device T1_2x4 *) y)) * (T1) s1_inv);
+            *(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) = (S1_2x4)(*((device T1_2x4 *) y));
        }
 #endif

@@ -935,12 +876,12 @@ kernel void kernel_mul_mm_id(

        int i = tiisg;
        for (; i < nr0/4; i += 32) {
-            *(D4 + i) = *(C4 + i) * s1_scale;
+            *(D4 + i) = *(C4 + i);
        }

        i = (4*(nr0/4)) + tiisg;
        for (; i < nr0; i += 32) {
-            *(D + i) = *(C + i) * s1_scale;
+            *(D + i) = *(C + i);
        }
    }
 }
@@ -3302,7 +3302,8 @@ kernel void kernel_mul_mv_id(
        uint3  tgpig[[threadgroup_position_in_grid]],
        ushort tiitg[[thread_index_in_threadgroup]],
        ushort tiisg[[thread_index_in_simdgroup]],
-        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+        ushort sgitg[[simdgroup_index_in_threadgroup]],
+        ushort3 ntg [[threads_per_threadgroup]]) {
    const int iid1 = tgpig.z/args.nei0;
    const int idx  = tgpig.z%args.nei0;

@@ -3316,11 +3317,20 @@ kernel void kernel_mul_mv_id(
    const int64_t i1 = idx;
    const int64_t i2 = i12;

+    device char * dst_cur = dst + (i1*args.ne0 + i2*args.ne1*args.ne0)*sizeof(float);
+
+    if (i02 == -1) {
+        device float * dst_f32 = (device float *) dst_cur;
+        const ushort nth = ntg.x*ntg.y*ntg.z;
+        for (int i0 = tiitg; i0 < args.ne0; i0 += nth) {
+            dst_f32[i0] = 0.0f;
+        }
+        return;
+    }
+
    device const char * src0_cur = src0s + i02*args.nb02;
    device const char * src1_cur = src1  + i11*args.nb11 + i12*args.nb12;

-    device char * dst_cur = dst + (i1*args.ne0 + i2*args.ne1*args.ne0)*sizeof(float);
-
    ggml_metal_kargs_mul_mv args0 = {
        /*.ne00 =*/ args.ne00,
        /*.ne01 =*/ args.ne01,
@@ -66,6 +66,8 @@ set(GGML_OPENCL_KERNELS
    moe_add_id_glu
    argmax
    argsort
+    top_k
+    lightning_indexer
    tri
    fill
    clamp
@@ -99,6 +101,12 @@ set(GGML_OPENCL_KERNELS
    mul_mv_q4_0_f32_1d_16x_flat
    mul_mv_q4_1_f32
    mul_mv_q4_1_f32_flat
+    gemm_noshuffle_q2_k_q8_1_dp4a
+    gemv_noshuffle_q2_k_f32
+    gemm_noshuffle_q2_k_f32
+    mul_mv_q2_k_f32
+    mul_mv_q2_k_f32_flat
+    mul_mm_q2_k_f32_l4_lm
    mul_mv_q4_k_f32
    mul_mv_q4_k_f32_flat
    mul_mv_q5_0_f32
@@ -34,6 +34,14 @@ kernel void kernel_add_id(

    global float * dst_row  = (global float *)((global char *)dst  + i1*nb1 + i2*nb2);
    global float * src0_row = (global float *)((global char *)src0 + i1*nb01 + i2*nb02);
+
+    if (i11 == -1) { // skipped slot, add nothing
+        for (int i0 = get_local_id(0); i0 < ne0; i0 += get_local_size(0)) {
+            dst_row[i0] = src0_row[i0];
+        }
+        return;
+    }
+
    global float * src1_row = (global float *)((global char *)src1 + i11*nb11);

    for (int i0 = get_local_id(0); i0 < ne0; i0 += get_local_size(0)) {
@@ -2661,3 +2661,115 @@ kernel void kernel_moe_expand_scale_q5_K(
        dst_min[((long)e*ne01 + row)*nblk32 + sub] = (half)(dm * (float)mn);
    }
 }
+
+struct block_q2_K {
+    uint8_t scales[QK_K/16]; // low nibble scale, high nibble min, one per 16 weights
+    uint8_t qs[QK_K/4];      // quants, 2 bits each
+    half d;                  // super-block scale for the quantized scales
+    half dmin;               // super-block scale for the quantized mins
+};
+
+kernel void kernel_convert_block_q2_K(
+    global struct block_q2_K * src0,
+    global uchar * dst_qs,
+    global uchar * dst_sc,
+    global half * dst_dm,
+    ulong n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_q2_K * b = src0 + i;
+    for (int j = 0; j < 64; ++j) {
+        dst_qs[64*i + j] = b->qs[j];
+    }
+    for (int j = 0; j < 16; ++j) {
+        dst_sc[16*i + j] = b->scales[j];
+    }
+    dst_dm[2*i + 0] = b->d;
+    dst_dm[2*i + 1] = b->dmin;
+}
+
+kernel void kernel_restore_block_q2_K(
+    global uchar * src_qs,
+    global uchar * src_sc,
+    global half * src_dm,
+    global struct block_q2_K * dst,
+    ulong n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_q2_K * b = dst + i;
+    for (int j = 0; j < 64; ++j) {
+        b->qs[j] = src_qs[64*i + j];
+    }
+    for (int j = 0; j < 16; ++j) {
+        b->scales[j] = src_sc[16*i + j];
+    }
+    b->d = src_dm[2*i + 0];
+    b->dmin = src_dm[2*i + 1];
+}
+
+kernel void kernel_convert_block_q2_k_noshuffle(
+    global struct block_q2_K * src0,
+    global uchar * dst_qs,      // QK_K/4 per block
+    global uchar * dst_sc,      // QK_K/16 per block
+    global half  * dst_dm,      // 2 per block (d, dmin)
+    ulong          n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_q2_K * b = (global struct block_q2_K *) src0 + i;
+
+    dst_dm[2*i + 0] = b->d;
+    dst_dm[2*i + 1] = b->dmin;
+
+    for (int j = 0; j < QK_K/16; ++j) { dst_sc[(QK_K/16) * i + j] = b->scales[j]; }
+
+    for (int g = 0; g < QK_K/4; ++g) {
+        uint pk = 0;
+        for (int t = 0; t < 4; ++t) {
+            const int e  = 4*g + t;
+            const int sb = e >> 5;
+            const int ee = e & 31;
+            pk |= (uint)(((b->qs[32*(sb >> 2) + ee] >> (2*(sb & 3))) & 3) << (2*t));
+        }
+        dst_qs[(QK_K/4) * i + g] = (uchar)pk;
+    }
+}
+
+kernel void kernel_restore_block_q2_k_noshuffle(
+    global uchar * src_qs,
+    global uchar * src_sc,
+    global half  * src_dm,
+    global struct block_q2_K * dst,
+    ulong          n_blk
+) {
+    const ulong i = get_global_id(0);
+    if (i >= n_blk) {
+        return;
+    }
+    global struct block_q2_K * b = (global struct block_q2_K *) dst + i;
+
+    b->d    = src_dm[2*i + 0];
+    b->dmin = src_dm[2*i + 1];
+
+    for (int j = 0; j < QK_K/16; ++j) { b->scales[j] = src_sc[(QK_K/16) * i + j]; }
+
+    for (int n2 = 0; n2 < 2; ++n2) {
+        for (int ee = 0; ee < 32; ++ee) {
+            uint v = 0;
+            for (int j = 0; j < 4; ++j) {
+                const int sb = 4*n2 + j;
+                const int e  = 32*sb + ee;
+                v |= ((uint)(src_qs[(QK_K/4) * i + (e >> 2)] >> (2*(e & 3))) & 3u) << (2*j);
+            }
+            b->qs[32*n2 + ee] = (uchar)v;
+        }
+    }
+}
@@ -0,0 +1,83 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_qcom_reqd_sub_group_size
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+__attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+kernel void kernel_gemm_noshuffle_q2_k_f32(
+    global const uchar * src0_qs,
+    global const uchar * src0_sc,
+    global const half  * src0_dm,
+    read_only image1d_buffer_t src1,
+    global float * dst,
+    ulong offsetd,
+    int m,
+    int n,
+    int k,
+    int n_no_padding
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+    const int gy = get_global_id(0);
+    const int row = get_global_id(1) * 4;
+    const int n_4 = n / 4;
+    if (row >= m) {
+        return;
+    }
+
+    half8 c0 = 0, c1 = 0, c2 = 0, c3 = 0;
+    for (int i = 0; i < k; i += 16) {
+        const half8 dm = vload8(0, src0_dm + 2 * (row + (i / 256) * m));
+        const uchar4 sc = vload4(0, src0_sc + row + (i / 16) * m);
+        const half4 scale = convert_half4(convert_float4(dm.s0246) * convert_float4(sc & (uchar4)(15)));
+        const half4 mval  = convert_half4(convert_float4(dm.s1357) * convert_float4(sc >> (uchar4)(4)));
+
+        for (int l = 0; l < 16; l += 4) {
+            const int ki = i + l;
+            const uchar4 bits = vload4(0, src0_qs + row + (ki / 4) * m);
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                half8 b;
+                b.s0123 = read_imageh(src1, gy * 2     + (ki + j) * n_4);
+                b.s4567 = read_imageh(src1, gy * 2 + 1 + (ki + j) * n_4);
+                const half4 w = convert_half4((bits >> (uchar4)(2 * j)) & (uchar4)(3)) * scale - mval;
+                c0 += b * w.s0;
+                c1 += b * w.s1;
+                c2 += b * w.s2;
+                c3 += b * w.s3;
+            }
+        }
+    }
+
+    int idx = gy * 8 * m + row;
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s0, c1.s0, c2.s0, c3.s0), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s1, c1.s1, c2.s1, c3.s1), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s2, c1.s2, c2.s2, c3.s2), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s3, c1.s3, c2.s3, c3.s3), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s4, c1.s4, c2.s4, c3.s4), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s5, c1.s5, c2.s5, c3.s5), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s6, c1.s6, c2.s6, c3.s6), 0, dst + idx);
+        idx += m;
+    }
+    if (idx + 3 < m * n_no_padding) {
+        vstore4((float4)(c0.s7, c1.s7, c2.s7, c3.s7), 0, dst + idx);
+    }
+}
@@ -0,0 +1,146 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#ifdef cl_khr_integer_dot_product
+#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
+#endif
+
+#define QK_K 256
+
+#ifndef TILESIZE_N
+#define TILESIZE_N 32
+#endif
+
+inline uint q2_k_pack(uint pk) {
+    return ( (pk      ) & 3u)
+         | (((pk >> 2) & 3u) <<  8)
+         | (((pk >> 4) & 3u) << 16)
+         | (((pk >> 6) & 3u) << 24);
+}
+
+// Pass staged operands by value; local dp4a operands miscompile on X2.
+#define KQ_STAGE(p) vload4(0, (p))
+
+inline int dot4_q8a_v(uint4 qw, uint4 a) {
+    int r = 0;
+    r = dot_acc_sat_4x8packed_ss_int(qw.s0, a.x, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s1, a.y, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s2, a.z, r);
+    r = dot_acc_sat_4x8packed_ss_int(qw.s3, a.w, r);
+    return r;
+}
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_gemm_noshuffle_q2_k_q8_1_dp4a(
+        __global const uchar  * src0_qs,
+        __global const uchar  * src0_sc,
+        __global const half   * src0_dm,   // half pairs: d, dmin
+        __global const uint   * src1_qa,   // q8_1 activations int8 (as uint, 4/elem) [N, K]
+        __global const half   * src1_da,   // q8_1 per-block scale [N, K/32]
+        __global       float  * dst,
+        ulong  offsetd,
+        int    m,
+        int    n_no_padding,
+        int    k
+) {
+    dst = (global float *)((global char *)dst + offsetd);
+
+    const uint lid = get_local_id(0);
+    const uint block_id_m = get_global_id(1);
+    const uint block_id_n = get_global_id(2);
+
+    const uint row      = block_id_m * 64 + lid;
+    const uint col_base = block_id_n * TILESIZE_N;
+    const bool row_valid = row < (uint)m;
+    const uint rrow     = row_valid ? row : 0;
+
+    const uint k_u = (uint)k >> 2;
+    const uint k_b = (uint)k >> 5;
+
+    __local uint4 sh_qa4[TILESIZE_N][2];
+    __local float sh_s [TILESIZE_N][2];   // per-16 activation sums, in qa units
+    __local half  sh_d [TILESIZE_N];
+
+#define NGROUPS (TILESIZE_N / 4)
+    float4 acc[NGROUPS];
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) acc[g] = (float4)(0.0f);
+
+    for (uint step = 0; step < (uint)k; step += 32) {
+        const uint sub = step >> 5;
+        const uint ib  = sub >> 3;
+
+        const half2 dm  = vload2(rrow + ib * (uint)m, src0_dm);
+        const float dv  = (float)dm.s0;
+        const float dmv = (float)dm.s1;
+
+        const uint  scb = rrow + (sub << 1) * (uint)m;
+        const uint  sc0 = (uint)src0_sc[scb + 0u * (uint)m];
+        const uint  sc1 = (uint)src0_sc[scb + 1u * (uint)m];
+        const float dl0 = dv  * (float)( sc0       & 0xFu);
+        const float ml0 = dmv * (float)( sc0 >> 4);
+        const float dl1 = dv  * (float)( sc1       & 0xFu);
+        const float ml1 = dmv * (float)( sc1 >> 4);
+
+        const uint qsb = rrow + (step >> 2) * (uint)m;
+        uint4 qlo, qhi;
+        qlo.s0 = q2_k_pack((uint)src0_qs[qsb + 0u * (uint)m]);
+        qlo.s1 = q2_k_pack((uint)src0_qs[qsb + 1u * (uint)m]);
+        qlo.s2 = q2_k_pack((uint)src0_qs[qsb + 2u * (uint)m]);
+        qlo.s3 = q2_k_pack((uint)src0_qs[qsb + 3u * (uint)m]);
+        qhi.s0 = q2_k_pack((uint)src0_qs[qsb + 4u * (uint)m]);
+        qhi.s1 = q2_k_pack((uint)src0_qs[qsb + 5u * (uint)m]);
+        qhi.s2 = q2_k_pack((uint)src0_qs[qsb + 6u * (uint)m]);
+        qhi.s3 = q2_k_pack((uint)src0_qs[qsb + 7u * (uint)m]);
+
+        for (uint idx = lid; idx < TILESIZE_N * 2u; idx += 64u) {
+            const uint t  = idx >> 1;
+            const uint h  = idx & 1u;
+            const uint c  = col_base + t;
+            const bool ok = c < (uint)n_no_padding;
+            const uint4 w = ok ? KQ_STAGE(src1_qa + c * k_u + (step >> 2) + (h << 2))
+                               : (uint4)(0u);
+            sh_qa4[t][h] = w;
+            int s = 0;
+            s = dot_acc_sat_4x8packed_ss_int(w.x, 0x01010101u, s);
+            s = dot_acc_sat_4x8packed_ss_int(w.y, 0x01010101u, s);
+            s = dot_acc_sat_4x8packed_ss_int(w.z, 0x01010101u, s);
+            s = dot_acc_sat_4x8packed_ss_int(w.w, 0x01010101u, s);
+            sh_s[t][h] = (float)s;
+            if (h == 0u) {
+                sh_d[t] = ok ? src1_da[c * k_b + sub] : (half)0;
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+#define LD4(arr, b) ((float4)((float)arr[(b)+0], (float)arr[(b)+1], (float)arr[(b)+2], (float)arr[(b)+3]))
+#define Q2_K_COL(b) (dl0 * (float)dot4_q8a_v(qlo, (uint4)(sh_qa4[b][0])) - ml0 * sh_s[b][0] + \
+                    dl1 * (float)dot4_q8a_v(qhi, (uint4)(sh_qa4[b][1])) - ml1 * sh_s[b][1])
+        #pragma unroll
+        for (int g = 0; g < NGROUPS; ++g) {
+            const int b = g * 4;
+            float4 rf;
+            rf.s0 = Q2_K_COL(b+0);  rf.s1 = Q2_K_COL(b+1);
+            rf.s2 = Q2_K_COL(b+2);  rf.s3 = Q2_K_COL(b+3);
+            acc[g] += LD4(sh_d, b) * rf;
+        }
+#undef Q2_K_COL
+#undef LD4
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (!row_valid) {
+        return;
+    }
+
+    #pragma unroll
+    for (int g = 0; g < NGROUPS; ++g) {
+        const uint b = (uint)(g * 4);
+        const float4 a = acc[g];
+        const uint c0 = col_base + b;
+        if (c0 + 0 < (uint)n_no_padding) dst[(c0 + 0) * (uint)m + row] = a.s0;
+        if (c0 + 1 < (uint)n_no_padding) dst[(c0 + 1) * (uint)m + row] = a.s1;
+        if (c0 + 2 < (uint)n_no_padding) dst[(c0 + 2) * (uint)m + row] = a.s2;
+        if (c0 + 3 < (uint)n_no_padding) dst[(c0 + 3) * (uint)m + row] = a.s3;
+    }
+#undef NGROUPS
+}
@@ -89,6 +89,9 @@ __kernel void kernel_gemv_moe_mxfp4_f32_ns(
    uint i11 = i20 % ne11;

    uint expert_id = src2[i20];
+    if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }
    uint expert_offset = expert_id * ne00 * ne01 / 32;

    __private float sum = 0.0f; // each thread calculate partial sum of one output
@@ -188,6 +191,9 @@ __kernel void kernel_gemv_moe_mxfp4_f32_ns_wimg(
    uint i11 = i20 % ne11;

    uint expert_id = src2[i20];
+    if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }
    uint expert_offset = expert_id * ne00 * ne01 / 32;

    __private float sum = 0.0f;
@@ -44,6 +44,9 @@ __kernel void kernel_gemv_moe_q4_0_f32_ns(
    uint i11 = i20 % ne11;

    uint expert_id = src2[i20];
+    if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }
    uint expert_offset = expert_id * ne00 * ne01 / 32;

    __private float sum = 0.0f; // each thread calculate partial sum of one output
@@ -45,6 +45,9 @@ __kernel void kernel_gemv_moe_q4_1_f32_ns(
    uint i11 = i20 % ne11;

    uint expert_id = src2[i20];
+    if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }
    uint expert_offset = expert_id * ne00 * ne01 / 32;

    __private float sum = 0.0f; // each thread calculate partial sum of one output
@@ -61,6 +61,9 @@ __kernel void kernel_gemv_moe_q4_k_f32_ns(
    uint i11 = i20 % ne11;

    uint expert_id = src2[i20];
+    if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }

    int num_superblocks = ne00 / QK_K;
    int num_subblocks = ne00 / 32;
@@ -180,6 +183,9 @@ __kernel void kernel_gemv_moe_q4_k_f32_ns_wimg(
    uint i11 = i20 % ne11;

    uint expert_id = src2[i20];
+    if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }

    int num_superblocks = ne00 / QK_K;
    int num_subblocks = ne00 / 32;
@@ -45,6 +45,9 @@ __kernel void kernel_gemv_moe_q5_0_f32_ns(
    uint i11 = i20 % ne11;

    uint expert_id = src2[i20];
+    if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }
    uint expert_offset = expert_id * ne00 * ne01 / 32;

    __private float sum = 0.0f; // each thread calculate partial sum of one output
@@ -46,6 +46,9 @@ __kernel void kernel_gemv_moe_q5_1_f32_ns(
    uint i11 = i20 % ne11;

    uint expert_id = src2[i20];
+    if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }
    uint expert_offset = expert_id * ne00 * ne01 / 32;

    __private float sum = 0.0f; // each thread calculate partial sum of one output
@@ -62,6 +62,9 @@ __kernel void kernel_gemv_moe_q5_k_f32_ns(
    uint i11 = i20 % ne11;

    uint expert_id = src2[i20];
+    if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }

    int num_superblocks = ne00 / QK_K;
    int num_subblocks = ne00 / 32;
@@ -45,6 +45,9 @@ __kernel void kernel_gemv_moe_q6_k_f32_ns(
    uint i11 = i20 % ne11;

    uint expert_id = src2[i20];
+    if ((int)expert_id == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }

    int num_superblocks = ne00 / QK_K;
    int num_subblocks = ne00 / 32;  // 8 sub-blocks of 32 per super-block
@@ -0,0 +1,231 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#define QK_K 256
+
+#ifndef Q2_K_MV_NSG
+#define Q2_K_MV_NSG 8
+#endif
+
+#ifndef Q2_K_MV_R
+#define Q2_K_MV_R 4
+#endif
+
+inline float4 q2_k_vals(uint pk) {
+    return (float4)((float)( pk       & 3u), (float)((pk >> 2) & 3u),
+                    (float)((pk >> 4) & 3u), (float)((pk >> 6) & 3u));
+}
+
+kernel void kernel_gemv_noshuffle_q2_k_f32(
+        global const uchar * src0_qs,
+        global const uchar * src0_sc,
+        global const half  * src0_dm,
+        global const float * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,      // K
+        int ne01,      // M
+        int ne10,      // activation row stride, == K
+        int ne0        // dst row stride
+) {
+    src1 = (global const float *)((global const char *)src1 + offset1);
+    dst  = (global float       *)((global char       *)dst  + offsetd);
+
+    const uint m   = (uint)ne01;
+    const uint K   = (uint)ne00;
+    const uint nsb = K / QK_K;
+
+    const uint lid = get_local_id(0);
+    const uint sgi = get_local_id(1);
+    const uint col = get_group_id(1);
+
+    global const float * y = src1 + (ulong)col * (uint)ne10;
+
+    const uint mr  = m / Q2_K_MV_R;                  // row groups
+    const uint j   = get_group_id(0) * 64u + lid;   // row group this lane owns
+    const uint row = j * Q2_K_MV_R;
+
+#if Q2_K_MV_R == 4
+    float sumf0 = 0.f;
+    float sumf1 = 0.f;
+    float sumf2 = 0.f;
+    float sumf3 = 0.f;
+
+    if (j < mr) {
+        global const uint * qsu = (global const uint *)src0_qs;
+        global const uint * scu = (global const uint *)src0_sc;
+
+        for (uint ib = sgi; ib < nsb; ib += Q2_K_MV_NSG) {
+            const half4 dma = vload4(2u*(j + ib*mr) + 0u, src0_dm);  // d0, dmin0, d1, dmin1
+            const half4 dmb = vload4(2u*(j + ib*mr) + 1u, src0_dm);  // d2, dmin2, d3, dmin3
+
+            float ad0 = 0.f, am0 = 0.f, ad1 = 0.f, am1 = 0.f;
+            float ad2 = 0.f, am2 = 0.f, ad3 = 0.f, am3 = 0.f;
+
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint grp = ib * 64u + sb * 8u;
+                const uint qsb = j + grp * mr;
+
+                for (uint h = 0; h < 2u; ++h) {
+                    const uint scv = scu[j + (2u * (ib * 8u + sb) + h) * mr];
+
+                    float a0 = 0.f, a1 = 0.f, a2 = 0.f, a3 = 0.f;
+                    float4 as = (float4)(0.f);
+                    for (uint u = 0; u < 4u; ++u) {
+                        const uint gg  = 4u*h + u;
+                        const uint qsv = qsu[qsb + gg * mr];   // four rows, one load
+                        const float4 yv = vload4(grp + gg, y);
+                        as += yv;
+                        a0 += dot(yv, q2_k_vals( qsv        & 0xFFu));
+                        a1 += dot(yv, q2_k_vals((qsv >>  8) & 0xFFu));
+                        a2 += dot(yv, q2_k_vals((qsv >> 16) & 0xFFu));
+                        a3 += dot(yv, q2_k_vals((qsv >> 24) & 0xFFu));
+                    }
+                    const float asum = as.s0 + as.s1 + as.s2 + as.s3;
+
+                    const uint s0 =  scv        & 0xFFu;
+                    const uint s1 = (scv >>  8) & 0xFFu;
+                    const uint s2 = (scv >> 16) & 0xFFu;
+                    const uint s3 = (scv >> 24) & 0xFFu;
+                    ad0 += (float)(s0 & 0xFu) * a0;   am0 += (float)(s0 >> 4) * asum;
+                    ad1 += (float)(s1 & 0xFu) * a1;   am1 += (float)(s1 >> 4) * asum;
+                    ad2 += (float)(s2 & 0xFu) * a2;   am2 += (float)(s2 >> 4) * asum;
+                    ad3 += (float)(s3 & 0xFu) * a3;   am3 += (float)(s3 >> 4) * asum;
+                }
+            }
+            sumf0 += (float)dma.s0 * ad0 - (float)dma.s1 * am0;
+            sumf1 += (float)dma.s2 * ad1 - (float)dma.s3 * am1;
+            sumf2 += (float)dmb.s0 * ad2 - (float)dmb.s1 * am2;
+            sumf3 += (float)dmb.s2 * ad3 - (float)dmb.s3 * am3;
+        }
+    }
+
+#if Q2_K_MV_NSG > 1
+    __local float4 part[Q2_K_MV_NSG][64];
+    part[sgi][lid] = (float4)(sumf0, sumf1, sumf2, sumf3);
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < Q2_K_MV_NSG; ++s) {
+        const float4 p = part[s][lid];
+        sumf0 += p.s0;
+        sumf1 += p.s1;
+        sumf2 += p.s2;
+        sumf3 += p.s3;
+    }
+#endif
+
+    if (j < mr) {
+        global float * o = dst + (ulong)col * (uint)ne0 + row;
+        o[0] = sumf0; o[1] = sumf1; o[2] = sumf2; o[3] = sumf3;
+    }
+
+#elif Q2_K_MV_R == 2
+    float sumf0 = 0.f;
+    float sumf1 = 0.f;
+
+    if (j < mr) {
+        global const ushort * qsu = (global const ushort *)src0_qs;
+        global const ushort * scu = (global const ushort *)src0_sc;
+
+        for (uint ib = sgi; ib < nsb; ib += Q2_K_MV_NSG) {
+            const half4 dm = vload4(j + ib * mr, src0_dm);  // d0, dmin0, d1, dmin1
+
+            float ad0 = 0.f, am0 = 0.f, ad1 = 0.f, am1 = 0.f;
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint grp = ib * 64u + sb * 8u;
+                const uint qsb = j + grp * mr;
+
+                for (uint h = 0; h < 2u; ++h) {
+                    const uint scv = (uint)scu[j + (2u * (ib * 8u + sb) + h) * mr];
+                    const uint s0  = scv & 0xFFu;
+                    const uint s1  = scv >> 8;
+
+                    float a0 = 0.f, a1 = 0.f;
+                    float4 as = (float4)(0.f);
+                    for (uint u = 0; u < 4u; ++u) {
+                        const uint gg  = 4u*h + u;
+                        const uint qsv = (uint)qsu[qsb + gg * mr];
+                        const float4 yv = vload4(grp + gg, y);
+                        as += yv;
+                        a0 += dot(yv, q2_k_vals( qsv       & 0xFFu));
+                        a1 += dot(yv, q2_k_vals((qsv >> 8) & 0xFFu));
+                    }
+                    const float asum = as.s0 + as.s1 + as.s2 + as.s3;
+                    ad0 += (float)(s0 & 0xFu) * a0;   am0 += (float)(s0 >> 4) * asum;
+                    ad1 += (float)(s1 & 0xFu) * a1;   am1 += (float)(s1 >> 4) * asum;
+                }
+            }
+            sumf0 += (float)dm.s0 * ad0 - (float)dm.s1 * am0;
+            sumf1 += (float)dm.s2 * ad1 - (float)dm.s3 * am1;
+        }
+    }
+
+#if Q2_K_MV_NSG > 1
+    __local float2 part[Q2_K_MV_NSG][64];
+    part[sgi][lid] = (float2)(sumf0, sumf1);
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < Q2_K_MV_NSG; ++s) {
+        const float2 p = part[s][lid];
+        sumf0 += p.s0;
+        sumf1 += p.s1;
+    }
+#endif
+
+    if (j < mr) {
+        vstore2((float2)(sumf0, sumf1), 0, dst + (ulong)col * (uint)ne0 + row);
+    }
+
+#else
+    float sumf0 = 0.f;
+
+    if (j < mr) {
+        for (uint ib = sgi; ib < nsb; ib += Q2_K_MV_NSG) {
+            const half2 dm = vload2(row + ib * m, src0_dm);
+
+            float ad = 0.f, am = 0.f;
+            for (uint sb = 0; sb < 8u; ++sb) {
+                const uint grp = ib * 64u + sb * 8u;
+                const uint qsb = row + grp * m;
+
+                for (uint h = 0; h < 2u; ++h) {
+                    const uint sc = (uint)src0_sc[row + (2u * (ib * 8u + sb) + h) * m];
+
+                    float a = 0.f;
+                    float4 as = (float4)(0.f);
+                    for (uint u = 0; u < 4u; ++u) {
+                        const uint gg = 4u*h + u;
+                        const float4 yv = vload4(grp + gg, y);
+                        as += yv;
+                        const uint pkv = (uint)src0_qs[qsb + gg * m];
+                        a += dot(yv, q2_k_vals(pkv));
+                    }
+                    const float asum = as.s0 + as.s1 + as.s2 + as.s3;
+                    ad += (float)(sc & 0xFu) * a;   am += (float)(sc >> 4) * asum;
+                }
+            }
+            sumf0 += (float)dm.s0 * ad - (float)dm.s1 * am;
+        }
+    }
+
+#if Q2_K_MV_NSG > 1
+    __local float part[Q2_K_MV_NSG][64];
+    part[sgi][lid] = sumf0;
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (sgi != 0) {
+        return;
+    }
+    for (uint s = 1; s < Q2_K_MV_NSG; ++s) {
+        sumf0 += part[s][lid];
+    }
+#endif
+
+    if (j < mr) {
+        dst[(ulong)col * (uint)ne0 + row] = sumf0;
+    }
+#endif
+}
@@ -0,0 +1,111 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#ifdef cl_khr_subgroup_shuffle
+#pragma OPENCL EXTENSION cl_khr_subgroup_shuffle : enable
+#endif
+
+kernel void kernel_lightning_indexer(
+        global const char * q,
+        ulong offset_q,
+        global const char * k,
+        ulong offset_k,
+        global const char * w,
+        ulong offset_w,
+        global const char * mask,
+        ulong offset_m,
+        global float * dst,
+        ulong offset_d,
+        int head_size,
+        int n_heads,
+        int n_tokens,
+        int n_keys,
+        int n_masks,
+        int type_k,
+        ulong nbq1,
+        ulong nbq2,
+        ulong nbq3,
+        ulong nbk2,
+        ulong nbk3,
+        ulong nbw1,
+        ulong nbw3,
+        ulong nbm1,
+        ulong nbm3,
+        int use_shuffle
+) {
+    const int tid = get_local_id(0);
+    const int lane = tid & (use_shuffle ? 3 : 15);
+    const int key = use_shuffle ? get_group_id(0)*16 + tid/4 : get_group_id(0)*4 + tid/16;
+    const int token = get_group_id(1);
+    const int stream = get_group_id(2);
+
+    q += offset_q + token*nbq2 + stream*nbq3;
+    k += offset_k + (ulong) key*nbk2 + stream*nbk3;
+    w += offset_w + token*nbw1 + stream*nbw3;
+    mask += offset_m + token*nbm1 + (stream % n_masks)*nbm3;
+    dst = (global float *) ((global char *) dst + offset_d);
+
+#ifdef cl_khr_subgroup_shuffle
+    if (use_shuffle) {
+        // Each group of four lanes keeps one key for all query heads.
+        float4 cached[8];
+        #pragma unroll
+        for (int j = 0; j < 8; ++j) {
+            const int i = 4*(lane + j*4);
+            cached[j] = key < n_keys ? (type_k == 0 ? vload4(0, (global const float *) k + i) : convert_float4(vload4(0, (global const half *) k + i))) : (float4) (0.0f);
+        }
+
+        float score = 0.0f;
+        for (int h = 0; h < n_heads; ++h) {
+            global const float * q_head = (global const float *) (q + h*nbq1);
+            float4 sum = (float4) (0.0f);
+            #pragma unroll
+            for (int j = 0; j < 8; ++j) {
+                sum = mad(vload4(0, q_head + 4*(lane + j*4)), cached[j], sum);
+            }
+            float value = sum.s0 + sum.s1 + sum.s2 + sum.s3;
+            value += sub_group_shuffle_xor(value, 2);
+            value += sub_group_shuffle_xor(value, 1);
+            if (lane == 0) {
+                score = mad(fmax(value, 0.0f), ((global const float *) w)[h], score);
+            }
+        }
+        if (lane == 0 && key < n_keys) {
+            dst[((ulong) stream*n_tokens + token)*n_keys + key] = score + (float) ((global const half *) mask)[key];
+        }
+        return;
+    }
+#endif
+
+    local float sums[64];
+    float score = 0.0f;
+    for (int h = 0; h < n_heads; ++h) {
+        global const float * q_head = (global const float *) (q + h*nbq1);
+        float sum = 0.0f;
+        if (key < n_keys) {
+            for (int i = lane; i < head_size; i += 16) {
+                float kv;
+                if (type_k == 0) {
+                    kv = ((global const float *) k)[i];
+                } else {
+                    kv = (float) ((global const half *) k)[i];
+                }
+                sum = mad(q_head[i], kv, sum);
+            }
+        }
+        sums[tid] = sum;
+        barrier(CLK_LOCAL_MEM_FENCE);
+        for (int stride = 8; stride > 0; stride /= 2) {
+            if (lane < stride) {
+                sums[tid] += sums[tid + stride];
+            }
+            barrier(CLK_LOCAL_MEM_FENCE);
+        }
+        if (lane == 0) {
+            score = mad(fmax(sums[tid], 0.0f), ((global const float *) w)[h], score);
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (lane == 0 && key < n_keys) {
+        dst[((ulong) stream*n_tokens + token)*n_keys + key] = score + (float) ((global const half *) mask)[key];
+    }
+}
@@ -18,6 +18,9 @@ __kernel void kernel_moe_histogram(
    }

    int expert_id = input[n * n_experts + k];
+    if (expert_id == -1) { // skipped slot, owned by no expert
+        return;
+    }
    atomic_inc(&hist[expert_id]);
 }

@@ -63,6 +66,9 @@ __kernel void kernel_moe_scatter(
    }

    int val = input[n * n_experts + k];
+    if (val == -1) { // skipped slot, owned by no expert
+        return;
+    }

    int local_slot = atomic_inc(&slot_counter[val]);

@@ -74,6 +80,32 @@ __kernel void kernel_moe_scatter(
    emap[tile_idx] = val;
 }

+// the gemm only writes the rows an expert owns, so the skipped rows (negative ids) are zeroed here
+__kernel void kernel_moe_zero_dst(
+    __global const int * input,
+    __global float * dst,
+    uint N,
+    uint topK,
+    uint n_experts,
+    uint ne0
+) {
+    uint n = get_global_id(0);
+    uint k = get_global_id(1);
+
+    if (n >= N || k >= topK) {
+        return;
+    }
+
+    if (input[n * n_experts + k] != -1) {
+        return;
+    }
+
+    __global float * dst_row = dst + ((ulong)n * topK + k) * ne0;
+    for (uint i = get_local_id(2); i < ne0; i += get_local_size(2)) {
+        dst_row[i] = 0.0f;
+    }
+}
+
 // Deterministic replacement for kernel_moe_scatter.
 //
 // kernel_moe_scatter takes each token's slot from atomic_inc(slot_counter[expert]),
@@ -0,0 +1,203 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#define INTEL_GPU 1
+#endif
+
+#define QK_K 256
+
+struct block_q2_K {
+    uchar scales[QK_K/16];
+    uchar qs[QK_K/4];
+    half  d;
+    half  dmin;
+};
+
+#define LOAD_VEC_A 4
+#define LOAD_VEC_B 4
+
+#define BM 64
+#define BN 64
+#ifndef BK
+#define BK 32
+#endif
+#ifdef INTEL_GPU
+#define TM 8
+#else
+#define TM 4
+#endif
+#define TN 8
+
+#ifndef LM_HALF
+#define LM_HALF 0
+#endif
+#if LM_HALF
+typedef half lm_st;
+#define LM_LD4(p, i) convert_float4(vload4((i), (p)))
+#else
+typedef float lm_st;
+#define LM_LD4(p, i) vload4((i), (p))
+#endif
+
+kernel void kernel_mul_mm_q2_k_f32_l4_lm(
+    global uchar * src0_qs,
+    global uchar * src0_sc,
+    global half * src0_dm,
+    ulong offset0,
+    global float4 * src1,
+    ulong offset1,
+    global float  * dst,
+    ulong offsetd,
+
+    int ne00,
+    int ne01,
+    int ne02,
+    int ne11,
+    int ne12,
+
+    int stride_a,
+    int stride_b,
+    int stride_d,
+
+    int batch_stride_a,
+    int batch_stride_b,
+    int batch_stride_d,
+
+    int r2,
+    int r3
+) {
+    src1 = (global float4*)((global char*)src1 + offset1);
+    dst  = (global float *)((global char*)dst  + offsetd);
+
+    local lm_st buf_a[BM * BK];
+    local lm_st buf_b[BN * BK];
+
+    const int batch_idx = get_global_id(2);
+
+    const int i13 = batch_idx / ne12;
+    const int i12 = batch_idx % ne12;
+
+    const int i03 = i13 / r3;
+    const int i02 = i12 / r2;
+
+    const int batch_idx_a = i03 * ne02 + i02;
+
+    const int ir = get_group_id(0);
+    const int ic = get_group_id(1);
+
+    const int tid = get_local_id(0);
+    const int th_r  = tid % (BM / TM);
+    const int th_c  = tid / (BM / TM);
+
+    const int loadr_a = get_local_id(0) % (BK / LOAD_VEC_A);
+    const int loadc_a = get_local_id(0) / (BK / LOAD_VEC_A);
+    const int loadr_b = get_local_id(0) % (BK / LOAD_VEC_B);
+    const int loadc_b = get_local_id(0) / (BK / LOAD_VEC_B);
+
+    const int loadstride_a = get_local_size(0) * LOAD_VEC_A / BK;
+    const int loadstride_b = get_local_size(0) * LOAD_VEC_B / BK;
+
+    // pos_a counts elements, not blocks - one q2_K super block holds QK_K of them
+    int pos_a = batch_idx_a * batch_stride_a + ir * BM * stride_a;
+    int pos_b = (batch_idx   * batch_stride_b + ic * BN * stride_b) / LOAD_VEC_B;
+
+    float4 sums4[(TM/4) * TN];
+    float4 cache_a4[TM/4];
+
+    for (int i = 0; i < (TM/4) * TN; i++) {
+        sums4[i] = (float4)(0.0f);
+    }
+
+    for (int block = 0; block < ne00; block += BK) {
+        for (int l = 0; l < BM; l += loadstride_a) {
+            if (ir*BM + loadc_a + l < ne01) {
+                int idx = pos_a + (loadc_a + l) * stride_a + loadr_a * LOAD_VEC_A;
+                int ib  = idx / QK_K;
+                ib += offset0/sizeof(struct block_q2_K);
+                int e   = idx % QK_K;
+
+                // dequantize_row_q2_K walks 2 halves of 128, each as 4 shifts of
+                // 2 groups of 16. Recover that position from the element index.
+                int n     = e >> 7;
+                int rem   = e & 127;
+                int j     = rem >> 5;
+                int s     = (rem >> 4) & 1;
+                int l16   = rem & 15;
+                int isc   = 8*n + 2*j + s;
+                int shift = 2*j;
+
+                uchar sc = src0_sc[16*ib + isc];
+                float d  =  (float)src0_dm[2*ib]    * (float)(sc & 0xF);
+                float m  = -(float)src0_dm[2*ib + 1] * (float)(sc >> 4);
+
+                uchar4 q = vload4(0, src0_qs + 64*ib + 32*n + 16*s + l16);
+                float4 v1 = convert_float4((uchar4)((q.s0 >> shift) & 3,
+                                                    (q.s1 >> shift) & 3,
+                                                    (q.s2 >> shift) & 3,
+                                                    (q.s3 >> shift) & 3)) * d + m;
+
+                buf_a[(loadr_a * LOAD_VEC_A + 0) * BM + loadc_a + l] = v1.s0;
+                buf_a[(loadr_a * LOAD_VEC_A + 1) * BM + loadc_a + l] = v1.s1;
+                buf_a[(loadr_a * LOAD_VEC_A + 2) * BM + loadc_a + l] = v1.s2;
+                buf_a[(loadr_a * LOAD_VEC_A + 3) * BM + loadc_a + l] = v1.s3;
+            } else {
+                buf_a[(loadr_a * LOAD_VEC_A + 0) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 1) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 2) * BM + loadc_a + l] = 0.0f;
+                buf_a[(loadr_a * LOAD_VEC_A + 3) * BM + loadc_a + l] = 0.0f;
+            }
+        }
+
+        for (int l = 0; l < BN; l += loadstride_b) {
+            if (ic*BN + loadc_b + l < ne11) {
+                int idx = pos_b + (loadc_b + l) * stride_b / LOAD_VEC_B + loadr_b;
+                buf_b[(loadr_b * LOAD_VEC_B + 0) * BN + loadc_b + l] = src1[idx].s0;
+                buf_b[(loadr_b * LOAD_VEC_B + 1) * BN + loadc_b + l] = src1[idx].s1;
+                buf_b[(loadr_b * LOAD_VEC_B + 2) * BN + loadc_b + l] = src1[idx].s2;
+                buf_b[(loadr_b * LOAD_VEC_B + 3) * BN + loadc_b + l] = src1[idx].s3;
+            } else {
+                buf_b[(loadr_b * LOAD_VEC_B + 0) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 1) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 2) * BN + loadc_b + l] = 0.0f;
+                buf_b[(loadr_b * LOAD_VEC_B + 3) * BN + loadc_b + l] = 0.0f;
+            }
+        }
+
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+        pos_a += BK;
+        pos_b += BK / LOAD_VEC_B;
+
+        for (int i = 0; i < BK; i++) {
+            for (int a = 0; a < TM/4; a++) {
+                cache_a4[a] = LM_LD4(buf_a + (i) * BM + th_r * TM, a);
+            }
+
+            for (int cc = 0; cc < TN; cc++) {
+                const float cache_b = buf_b[(i) * BN + th_c * TN + cc];
+                for (int a = 0; a < TM/4; a++) {
+                    const int sums_idx = cc*(TM/4) + a;
+                    sums4[sums_idx] = mad(cache_a4[a], (float4)cache_b, sums4[sums_idx]);
+                }
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    const int dr = ir * BM + th_r * TM;
+    const int dc = ic * BN + th_c * TN;
+
+    const ulong offsets = batch_idx * batch_stride_d;
+
+    for (int cc = 0; cc < TN; cc++) {
+        for (int a = 0; a < TM/4; a++) {
+            const float4 v = sums4[cc * (TM/4) + a];
+            const float  vs[4] = { v.s0, v.s1, v.s2, v.s3 };
+            for (int k = 0; k < 4; k++) {
+                if (dr + 4*a + k < ne01 && dc + cc < ne11) {
+                    dst[offsets + (dc + cc) * stride_d + dr + 4*a + k] = vs[k];
+                }
+            }
+        }
+    }
+}
@@ -172,6 +172,9 @@ kernel void kernel_mul_mv_id_mxfp4_f32(
    const int idx  = get_group_id(2)%ne20;

    int i02 = ((global int *) (src2 + iid1*nb21))[idx];
+    if (i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }

    int i11 = idx % ne11;
    int i12 = iid1;
@@ -101,6 +101,9 @@ kernel void kernel_mul_mv_id_mxfp4_f32_flat(
    const int idx  = get_group_id(2) % ne20;

    uint i02 = ((global uint *) (src2 + offset2 + iid1 * nb21))[idx];
+    if ((int)i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }

    int i11 = idx % ne11;

@@ -267,6 +267,9 @@ kernel void kernel_mul_mv_id_q4_0_f32_8x_flat(
    const int idx  = get_group_id(2)%ne20;

    const int i02 = ((global int *)(src2 + iid1*nb21))[idx];
+    if (i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }

    const int i11 = idx%ne11;
    const int i12 = iid1;
@@ -73,6 +73,9 @@ kernel void kernel_mul_mv_id_q8_0_f32(
    int idx  = get_group_id(2)%ne20;

    int i02 = ((global int *) (src2 + iid1*nb21))[idx];
+    if (i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }

    int i11_ = idx % ne11;
    int i12_ = iid1;
@@ -72,6 +72,9 @@ kernel void kernel_mul_mv_id_q8_0_f32_flat(
    int idx  = (int)get_group_id(2)%ne20;

    int i02 = ((global int *) (src2 + iid1*nb21))[idx];
+    if (i02 == -1) { // skipped slot, the dst row is zeroed by kernel_moe_zero_dst
+        return;
+    }

    int i11_ = idx % ne11;
    int i12_ = iid1;
@@ -0,0 +1,161 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
+#define INTEL_GPU 1
+#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
+#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
+#elif defined(cl_qcom_reqd_sub_group_size)
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define ADRENO_GPU 1
+#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
+#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+
+//------------------------------------------------------------------------------
+// block_q2_K
+//------------------------------------------------------------------------------
+#define QK_K 256
+
+// 16 blocks of 16 elements each
+// weight is represented as x = a * q + b
+typedef struct {
+    uchar scales[QK_K/16]; // scales and mins, quantized with 4 bits
+    uchar qs[QK_K/4];      // 2-bit quants
+    half  d;              // super-block scale for quantized scales
+    half  dmin;            // super-block scale for quantized mins
+} block_q2_K;
+
+#undef N_DST
+#undef N_SIMDGROUP
+#undef N_SIMDWIDTH
+
+#ifdef INTEL_GPU
+#define N_DST 4 // number of rows each SIMD group works on
+#define N_SIMDGROUP 1 // number of SIMD groups in a thread group
+#define N_SIMDWIDTH 16 // SIMD group size
+#elif defined (ADRENO_GPU)
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 64
+#endif
+
+#undef  BLOCK_STRIDE
+// number of (super) blocks each subgroup processes
+// 8 threads cover one super block, so a wave covers N_SIMDWIDTH/8 of them
+#define BLOCK_STRIDE (N_SIMDWIDTH/8)
+
+#ifdef INTEL_GPU
+REQD_SUBGROUP_SIZE_16
+#elif defined (ADRENO_GPU)
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mv_q2_K_f32(
+        global char * src0,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global char * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne12,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src0 = src0 + offset0;
+    src1 = src1 + offset1;
+    dst  = dst  + offsetd;
+
+    int ix = get_sub_group_local_id()/8;  // super block index
+    int it = get_sub_group_local_id()%8;  // block index (inside super block)
+    int iq = it/4;      // 0 or 1 - first or second half of the super block
+    int ir = it%4;      // 0...3 - block index in the half super block
+    int is = (8*ir)/16; // 0 or 1 - scale pair
+
+    int nb = ne00/QK_K;
+
+    int r0 = get_group_id(0);
+    int r1 = get_group_id(1);
+    int im = get_group_id(2);
+    int first_row = (r0 * N_SIMDGROUP + get_sub_group_id()) * N_DST;
+
+    int i12 = im%ne12;
+    int i13 = im/ne12;
+
+    ulong offset_src0 = first_row*nb01 + (i12/r2)*nb02 + (i13/r3)*nb03;
+    ulong offset_src1 =        r1*nb11 + (i12   )*nb12 + (i13   )*nb13;
+
+    global block_q2_K * x = (global block_q2_K *) (src0 + offset_src0);
+    global float      * y = (global float      *) (src1 + offset_src1);
+
+    float yl[32];
+    float sumf[N_DST] = {0.f};
+    float all_sum;
+
+    global float * y4 = y + ix * QK_K + 128 * iq + 8 * ir;
+
+    for (int ib = ix; ib < nb; ib += BLOCK_STRIDE) {
+        float4 sumy = {0.f, 0.f, 0.f, 0.f};
+        for (int i = 0; i < 8; ++i) {
+            yl[i+ 0] = y4[i+ 0]; sumy.s0 += yl[i+ 0];
+            yl[i+ 8] = y4[i+32]; sumy.s1 += yl[i+ 8];
+            yl[i+16] = y4[i+64]; sumy.s2 += yl[i+16];
+            yl[i+24] = y4[i+96]; sumy.s3 += yl[i+24];
+        }
+
+        global uchar  * sc = (global uchar  *)x[ib].scales + 8*iq + is;
+        global ushort * qs = (global ushort *)x[ib].qs + 16 * iq + 4 * ir;
+        global half   * dh = &x[ib].d;
+
+        // only ne01 output rows exist; reading past them can pull in an inf scale
+        for (int row = 0; row < N_DST && first_row + row < ne01; row++) {
+            float4 acc1 = {0.f, 0.f, 0.f, 0.f};
+            float4 acc2 = {0.f, 0.f, 0.f, 0.f};
+            for (int i = 0; i < 8; i += 2) {
+                acc1.s0 += yl[i+ 0] * (qs[i/2] & 0x0003);
+                acc2.s0 += yl[i+ 1] * (qs[i/2] & 0x0300);
+                acc1.s1 += yl[i+ 8] * (qs[i/2] & 0x000c);
+                acc2.s1 += yl[i+ 9] * (qs[i/2] & 0x0c00);
+                acc1.s2 += yl[i+16] * (qs[i/2] & 0x0030);
+                acc2.s2 += yl[i+17] * (qs[i/2] & 0x3000);
+                acc1.s3 += yl[i+24] * (qs[i/2] & 0x00c0);
+                acc2.s3 += yl[i+25] * (qs[i/2] & 0xc000);
+            }
+
+            float dall = dh[0];
+            float dmin = dh[1] * 1.f/16.f;
+            sumf[row] += dall * ((acc1.s0 + 1.f/256.f * acc2.s0) * (sc[0] & 0xF) * 1.f/ 1.f +
+                                 (acc1.s1 + 1.f/256.f * acc2.s1) * (sc[2] & 0xF) * 1.f/ 4.f +
+                                 (acc1.s2 + 1.f/256.f * acc2.s2) * (sc[4] & 0xF) * 1.f/16.f +
+                                 (acc1.s3 + 1.f/256.f * acc2.s3) * (sc[6] & 0xF) * 1.f/64.f) -
+                        dmin * (sumy.s0 * (sc[0] & 0xF0) + sumy.s1 * (sc[2] & 0xF0) +
+                                sumy.s2 * (sc[4] & 0xF0) + sumy.s3 * (sc[6] & 0xF0));
+
+            qs += nb01/2;
+            sc += nb01;
+            dh += nb01/2;
+        }
+
+        y4 += BLOCK_STRIDE * QK_K;
+    }
+
+    global float * dst_f32 = (global float *) dst + im*ne0*ne1 + r1*ne0;
+
+    for (int row = 0; row < N_DST; ++row) {
+        all_sum = sub_group_reduce_add(sumf[row]);
+        if (first_row + row < ne01) {
+            if (get_sub_group_local_id() == 0) {
+                dst_f32[first_row + row] = all_sum;
+            }
+        }
+    }
+}
@@ -0,0 +1,162 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+#ifdef cl_intel_required_subgroup_size
+#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
+#define INTEL_GPU 1
+#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
+#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
+#elif defined(cl_qcom_reqd_sub_group_size)
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define ADRENO_GPU 1
+#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
+#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
+#endif
+
+//------------------------------------------------------------------------------
+// struct block_q2_K
+//------------------------------------------------------------------------------
+#define QK_K 256
+
+// 16 blocks of 16 elements each
+// weight is represented as x = a * q + b
+struct block_q2_K {
+    uchar scales[QK_K/16]; // scales and mins, quantized with 4 bits
+    uchar qs[QK_K/4];      // 2-bit quants
+    half  d;              // super-block scale for quantized scales
+    half  dmin;            // super-block scale for quantized mins
+};
+
+#undef N_DST
+#undef N_SIMDGROUP
+#undef N_SIMDWIDTH
+
+#ifdef INTEL_GPU
+#define N_DST 4 // number of rows each SIMD group works on
+#define N_SIMDGROUP 1 // number of SIMD groups in a thread group
+#define N_SIMDWIDTH 16 // SIMD group size
+#elif defined (ADRENO_GPU)
+#define N_DST 4
+#define N_SIMDGROUP 1
+#define N_SIMDWIDTH 64
+#endif
+
+#undef  BLOCK_STRIDE
+// number of (super) blocks each subgroup processes
+// 8 threads cover one super block, so a wave covers N_SIMDWIDTH/8 of them
+#define BLOCK_STRIDE (N_SIMDWIDTH/8)
+
+#ifdef INTEL_GPU
+REQD_SUBGROUP_SIZE_16
+#elif defined (ADRENO_GPU)
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mv_q2_K_f32_flat(
+        global uchar * src0_qs,
+        global uchar * src0_sc,
+        global half * src0_dm,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global char * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne12,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src1 = src1 + offset1;
+    dst  = dst  + offsetd;
+
+    int ix = get_sub_group_local_id()/8;  // super block index
+    int it = get_sub_group_local_id()%8;  // block index (inside super block)
+    int iq = it/4;      // 0 or 1 - first or second half of the super block
+    int ir = it%4;      // 0...3 - block index in the half super block
+    int is = (8*ir)/16; // 0 or 1 - scale pair
+
+    int nb = ne00/QK_K;
+
+    int r0 = get_group_id(0);
+    int r1 = get_group_id(1);
+    int im = get_group_id(2);
+    int first_row = (r0 * N_SIMDGROUP + get_sub_group_id()) * N_DST;
+
+    int i12 = im%ne12;
+    int i13 = im/ne12;
+
+    ulong offset_src0 = (offset0 + first_row*nb01 + (i12/r2)*nb02 + (i13/r3)*nb03)/sizeof(struct block_q2_K);
+    ulong row_stride = nb01/sizeof(struct block_q2_K);
+    ulong offset_src1 =        r1*nb11 + (i12   )*nb12 + (i13   )*nb13;
+
+    global float      * y = (global float      *) (src1 + offset_src1);
+
+    float yl[32];
+    float sumf[N_DST] = {0.f};
+    float all_sum;
+
+    global float * y4 = y + ix * QK_K + 128 * iq + 8 * ir;
+
+    for (int ib = ix; ib < nb; ib += BLOCK_STRIDE) {
+        float4 sumy = {0.f, 0.f, 0.f, 0.f};
+        for (int i = 0; i < 8; ++i) {
+            yl[i+ 0] = y4[i+ 0]; sumy.s0 += yl[i+ 0];
+            yl[i+ 8] = y4[i+32]; sumy.s1 += yl[i+ 8];
+            yl[i+16] = y4[i+64]; sumy.s2 += yl[i+16];
+            yl[i+24] = y4[i+96]; sumy.s3 += yl[i+24];
+        }
+
+        global uchar  * sc = (src0_sc + 16*(offset_src0 + ib)) + 8*iq + is;
+        global ushort * qs = (global ushort *)(src0_qs + 64*(offset_src0 + ib)) + 16 * iq + 4 * ir;
+        global half   * dh = src0_dm + 2*(offset_src0 + ib);
+
+        // only ne01 output rows exist; reading past them can pull in an inf scale
+        for (int row = 0; row < N_DST && first_row + row < ne01; row++) {
+            float4 acc1 = {0.f, 0.f, 0.f, 0.f};
+            float4 acc2 = {0.f, 0.f, 0.f, 0.f};
+            for (int i = 0; i < 8; i += 2) {
+                acc1.s0 += yl[i+ 0] * (qs[i/2] & 0x0003);
+                acc2.s0 += yl[i+ 1] * (qs[i/2] & 0x0300);
+                acc1.s1 += yl[i+ 8] * (qs[i/2] & 0x000c);
+                acc2.s1 += yl[i+ 9] * (qs[i/2] & 0x0c00);
+                acc1.s2 += yl[i+16] * (qs[i/2] & 0x0030);
+                acc2.s2 += yl[i+17] * (qs[i/2] & 0x3000);
+                acc1.s3 += yl[i+24] * (qs[i/2] & 0x00c0);
+                acc2.s3 += yl[i+25] * (qs[i/2] & 0xc000);
+            }
+
+            float dall = dh[0];
+            float dmin = dh[1] * 1.f/16.f;
+            sumf[row] += dall * ((acc1.s0 + 1.f/256.f * acc2.s0) * (sc[0] & 0xF) * 1.f/ 1.f +
+                                 (acc1.s1 + 1.f/256.f * acc2.s1) * (sc[2] & 0xF) * 1.f/ 4.f +
+                                 (acc1.s2 + 1.f/256.f * acc2.s2) * (sc[4] & 0xF) * 1.f/16.f +
+                                 (acc1.s3 + 1.f/256.f * acc2.s3) * (sc[6] & 0xF) * 1.f/64.f) -
+                        dmin * (sumy.s0 * (sc[0] & 0xF0) + sumy.s1 * (sc[2] & 0xF0) +
+                                sumy.s2 * (sc[4] & 0xF0) + sumy.s3 * (sc[6] & 0xF0));
+
+            qs += row_stride*32;
+            sc += row_stride*16;
+            dh += row_stride*2;
+        }
+
+        y4 += BLOCK_STRIDE * QK_K;
+    }
+
+    global float * dst_f32 = (global float *) dst + im*ne0*ne1 + r1*ne0;
+
+    for (int row = 0; row < N_DST; ++row) {
+        all_sum = sub_group_reduce_add(sumf[row]);
+        if (first_row + row < ne01) {
+            if (get_sub_group_local_id() == 0) {
+                dst_f32[first_row + row] = all_sum;
+            }
+        }
+    }
+}
@@ -0,0 +1,207 @@
+#ifdef cl_khr_subgroups
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#endif
+
+uint top_k_key(float value) {
+    uint bits = as_uint(value);
+    if ((bits & 0x7fffffffu) == 0) {
+        bits = 0;
+    }
+    return bits & 0x80000000u ? ~bits : bits ^ 0x80000000u;
+}
+
+uint4 top_k_keys(global const char * src, int i, int n, ulong nb0) {
+    if (nb0 == sizeof(float) && i + 3 < n) {
+        const float4 v = vload4(0, (global const float *) src + i);
+        return (uint4) (top_k_key(v.s0), top_k_key(v.s1), top_k_key(v.s2), top_k_key(v.s3));
+    }
+    return (uint4) (
+        i     < n ? top_k_key(*(global const float *) (src + (ulong)  i     *nb0)) : 0,
+        i + 1 < n ? top_k_key(*(global const float *) (src + (ulong) (i + 1)*nb0)) : 0,
+        i + 2 < n ? top_k_key(*(global const float *) (src + (ulong) (i + 2)*nb0)) : 0,
+        i + 3 < n ? top_k_key(*(global const float *) (src + (ulong) (i + 3)*nb0)) : 0);
+}
+
+void top_k_store(global int * dst, global uint2 * candidates, global const char * src,
+        ulong nb0, int indexed, int base, int i, uint pos) {
+    const int index = indexed ? *(global const int *) (src + (ulong) i*nb0 + sizeof(float)) : base + i;
+    if (candidates) {
+        candidates[pos] = (uint2) (as_uint(*(global const float *) (src + (ulong) i*nb0)), index);
+    } else {
+        dst[pos] = index;
+    }
+}
+
+void top_k_select(
+        global const char * src,
+        global int * dst,
+        global uint2 * candidates,
+        ulong nb00,
+        int ne00,
+        int k,
+        int base,
+        int indexed,
+        local uint * hist
+) {
+    const int tid = get_local_id(0);
+    const int nth = get_local_size(0);
+    if (k == ne00) {
+        for (int i = tid; i < ne00; i += nth) {
+            top_k_store(dst, candidates, src, nb00, indexed, base, i, i);
+        }
+        return;
+    }
+
+    uint prefix = 0;
+    uint mask = 0;
+    uint rank = k;
+    for (int shift = 24; shift >= 0; shift -= 8) {
+        for (int i = tid; i < 256; i += nth) {
+            hist[i] = 0;
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+        for (int i = 4*tid; i < ne00; i += 4*nth) {
+            const uint4 v = top_k_keys(src, i, ne00, nb00);
+            const uint keys[4] = {v.s0, v.s1, v.s2, v.s3};
+            #pragma unroll
+            for (int j = 0; j < 4; ++j) {
+                const uint key = keys[j];
+                if (i + j < ne00 && (key & mask) == prefix) {
+                    atomic_inc(hist + ((key >> shift) & 255));
+                }
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+        if (tid == 0) {
+            uint above = 0;
+            for (int b = 255; b >= 0; --b) {
+                uint count = hist[b];
+                if (above + count >= rank) {
+                    hist[256] = b;
+                    hist[257] = rank - above;
+                    hist[258] = count;
+                    break;
+                }
+                above += count;
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+        prefix |= hist[256] << shift;
+        rank = hist[257];
+        const bool done = hist[258] == rank;
+        barrier(CLK_LOCAL_MEM_FENCE);
+        if (done) {
+            // The whole bucket is selected, so its low bits do not matter.
+            break;
+        }
+        mask |= 255u << shift;
+    }
+
+    uint greater = 0;
+    uint equal = 0;
+    for (int i = 4*tid; i < ne00; i += 4*nth) {
+        const uint4 v = top_k_keys(src, i, ne00, nb00);
+        const uint keys[4] = {v.s0, v.s1, v.s2, v.s3};
+        #pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            greater += i + j < ne00 && keys[j] > prefix;
+            equal += i + j < ne00 && keys[j] == prefix;
+        }
+    }
+    uint pos_greater;
+    uint pos_equal;
+#ifdef cl_khr_subgroups
+    const uint sg = get_sub_group_id();
+    const uint nsg = get_num_sub_groups();
+    const uint total_greater = sub_group_reduce_add(greater);
+    const uint total_equal = sub_group_reduce_add(equal);
+    pos_greater = sub_group_scan_exclusive_add(greater);
+    pos_equal = sub_group_scan_exclusive_add(equal);
+    if (get_sub_group_local_id() == 0) {
+        hist[sg] = total_greater;
+        hist[nsg + sg] = total_equal;
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+    for (uint i = 0; i < nsg; ++i) {
+        pos_equal += hist[i];
+        if (i < sg) {
+            pos_greater += hist[i];
+            pos_equal += hist[nsg + i];
+        }
+    }
+#else
+    hist[tid] = greater;
+    hist[nth + tid] = equal;
+    barrier(CLK_LOCAL_MEM_FENCE);
+    pos_greater = 0;
+    pos_equal = 0;
+    for (int i = 0; i < nth; ++i) {
+        pos_equal += hist[i];
+        if (i < tid) {
+            pos_greater += hist[i];
+            pos_equal += hist[nth + i];
+        }
+    }
+#endif
+    for (int i = 4*tid; i < ne00; i += 4*nth) {
+        const uint4 v = top_k_keys(src, i, ne00, nb00);
+        const uint keys[4] = {v.s0, v.s1, v.s2, v.s3};
+        #pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            if (i + j < ne00) {
+                if (keys[j] > prefix) {
+                    top_k_store(dst, candidates, src, nb00, indexed, base, i + j, pos_greater++);
+                } else if (keys[j] == prefix && pos_equal < k) {
+                    top_k_store(dst, candidates, src, nb00, indexed, base, i + j, pos_equal++);
+                }
+            }
+        }
+    }
+}
+
+kernel void kernel_top_k_f32(
+        global const char * src,
+        ulong offset0,
+        global int * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne02,
+        ulong nb00,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int k_max,
+        local uint * hist,
+        global uint2 * candidates,
+        int tile_width
+) {
+    const ulong row = get_group_id(1);
+    const int part = get_group_id(0);
+    const int parts = get_num_groups(0);
+    const int base = part*tile_width;
+    const int n = min(ne00 - base, tile_width);
+    const int k = min(k_max, n);
+    const int out_width = (parts - 1)*k_max + min(k_max, ne00 - (parts - 1)*tile_width);
+    src += offset0 + (row % ne01)*nb01 + (row / ne01 % ne02)*nb02 + (row / ne01 / ne02)*nb03 + (ulong) base*nb00;
+    if (candidates) {
+        candidates += row*out_width + part*k_max;
+    } else {
+        dst = (global int *) ((global char *) dst + offsetd) + row*out_width + part*k_max;
+    }
+    top_k_select(src, dst, candidates, nb00, n, k, base, 0, hist);
+}
+
+kernel void kernel_top_k_merge_f32(
+        global const uint2 * candidates,
+        global int * dst,
+        ulong offsetd,
+        int n_candidates,
+        int k,
+        local uint * hist
+) {
+    const ulong row = get_group_id(1);
+    candidates += row*n_candidates;
+    dst = (global int *) ((global char *) dst + offsetd) + row*k;
+    top_k_select((global const char *) candidates, dst, 0, sizeof(uint2), n_candidates, k, 0, 1, hist);
+}
@@ -9,8 +9,12 @@
 #include <openvino/op/constant.hpp>
 #include <openvino/op/convert.hpp>
 #include <openvino/op/gather.hpp>
+#include <openvino/op/greater_eq.hpp>
+#include <openvino/op/maximum.hpp>
+#include <openvino/op/multiply.hpp>
 #include <openvino/op/reshape.hpp>
 #include <openvino/op/shape_of.hpp>
+#include <openvino/op/unsqueeze.hpp>

 namespace ov {
 namespace frontend {
@@ -52,8 +56,16 @@ OutputVector translate_add_id(const NodeContext & context) {
        ids = std::make_shared<ov::op::v0::Convert>(ids, ov::element::i32);
    }

+    // an id of -1 marks a skipped slot: gather row 0 for it and zero the bias with a mask
+    auto zero = ov::op::v0::Constant::create(ids.get_element_type(), ov::Shape{}, {0});
+    ov::Output<ov::Node> keep = std::make_shared<ov::op::v1::GreaterEqual>(ids, zero);
+    keep = std::make_shared<ov::op::v0::Convert>(keep, bias.get_element_type());
+    keep = std::make_shared<ov::op::v0::Unsqueeze>(keep, ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, {-1}));
+    ids = std::make_shared<ov::op::v1::Maximum>(ids, zero);
+
    auto gather_axis = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{}, {0});
    ov::Output<ov::Node> selected_bias = std::make_shared<ov::op::v8::Gather>(bias, ids, gather_axis);
+    selected_bias = std::make_shared<ov::op::v1::Multiply>(selected_bias, keep);
    selected_bias = std::make_shared<ov::op::v1::Reshape>(
        selected_bias, std::make_shared<ov::op::v3::ShapeOf>(input, ov::element::i64), false);

@@ -517,6 +517,11 @@ bool socket_t::impl::rdma_recv(void * data, size_t size) {

        int slot = (int)wc.wr_id;
        size_t got = wc.byte_len;
+        if (got > rem) {
+            LOG_ERROR("RDMA recv overrun: got=%zu rem=%zu\n", got, rem);
+            c->post_rx(slot);
+            return false;
+        }
        memcpy(dst, c->rx_slot(slot), got);

        if (!c->post_rx(slot)) return false;
@@ -26,6 +26,15 @@ static void add_id_kernel(
  float* dst_row = (float*)((char*)dst + i1 * nb1 + i2 * nb2);
  const float* src0_row =
      (const float*)((const char*)src0 + i1 * nb01 + i2 * nb02);
+
+  if (i11 == -1) { // skipped slot, add nothing
+    for (int64_t i0 = item_ct1.get_local_id(2); i0 < ne0;
+         i0 += item_ct1.get_local_range(2)) {
+      dst_row[i0] = src0_row[i0];
+    }
+    return;
+  }
+
  const float* src1_row = (const float*)((const char*)src1 + i11 * nb11);

  for (int64_t i0 = item_ct1.get_local_id(2); i0 < ne0;
@@ -5700,6 +5700,23 @@ __dpct_inline__ static void k_copy_dst_from_contiguous(
    }
 }

+__dpct_inline__ static void k_zero_dst_rows(
+    char *__restrict__ dst_original,
+    const mmid_row_mapping *__restrict__ row_mapping, int64_t ne0, size_t nb1,
+    size_t nb2, const sycl::nd_item<3> &item_ct1) {
+    int32_t i = item_ct1.get_group(2);
+
+    const int32_t i1 = row_mapping[i].i1;
+    const int32_t i2 = row_mapping[i].i2;
+
+    float * dst_row_original = (float *)(dst_original + i1*nb1 + i2*nb2);
+
+    for (int j = item_ct1.get_local_id(2); j < ne0;
+         j += item_ct1.get_local_range(2)) {
+        dst_row_original[j] = 0.0f;
+    }
+}
+
 // returns true when ggml_sycl_mul_mat_id takes the fallback path that requires stream synchronization
 static bool ggml_sycl_mul_mat_id_needs_sync(const ggml_tensor * dst) {
    const ggml_tensor * src0 = dst->src[0];
@@ -5785,14 +5802,18 @@ static void mmid_counting_sort_rows(
        int64_t n_ids, int64_t n_as, int64_t n_routed_rows,
        std::vector<int64_t> & expert_counts,
        std::vector<int64_t> & expert_row_offsets,
-        std::vector<mmid_row_mapping> & routed_row_src) {
+        std::vector<mmid_row_mapping> & routed_row_src,
+        std::vector<mmid_row_mapping> & skipped_row_dst) {

    // frequencies: how many routed rows each expert "owns"
    expert_counts.assign(n_as, 0);
    for (int64_t iid1 = 0; iid1 < ids->ne[1]; iid1++) {
        for (int64_t id = 0; id < n_ids; id++) {
            const int32_t row_id_i = *(const int32_t *) (ids_host + iid1*ids->nb[1] + id*ids->nb[0]);
-            GGML_ASSERT(row_id_i >= 0 && row_id_i < n_as);
+            GGML_ASSERT(row_id_i == -1 || (row_id_i >= 0 && row_id_i < n_as));
+            if (row_id_i == -1) {
+                continue;
+            }
            expert_counts[row_id_i]++;
        }
    }
@@ -5805,10 +5826,17 @@ static void mmid_counting_sort_rows(

    std::vector<int64_t> expert_row_next = expert_row_offsets;
    routed_row_src.resize(n_routed_rows);
+    skipped_row_dst.clear();
    for (int64_t iid1 = 0; iid1 < ids->ne[1]; iid1++) {
        for (int64_t id = 0; id < n_ids; id++) {
            const int32_t row_id_i = *(const int32_t *) (ids_host + iid1*ids->nb[1] + id*ids->nb[0]);
-            GGML_ASSERT(row_id_i >= 0 && row_id_i < n_as);
+            GGML_ASSERT(row_id_i == -1 || (row_id_i >= 0 && row_id_i < n_as));
+
+            // a skipped slot has no expert, its dst row is zeroed instead
+            if (row_id_i == -1) {
+                skipped_row_dst.push_back({(int32_t) id, (int32_t) iid1});
+                continue;
+            }

            // find and validate the next free row for a given expert (row_id_i)
            const int64_t routed_row = expert_row_next[row_id_i]++;
@@ -5817,6 +5845,7 @@ static void mmid_counting_sort_rows(
            routed_row_src[routed_row] = {(int32_t) id, (int32_t) iid1};
        }
    }
+    routed_row_src.resize(expert_row_offsets[n_as]);
 }

 static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
@@ -5884,7 +5913,7 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
        for (int64_t iid1 = 0; iid1 < ids->ne[1]; iid1++) {
            for (int64_t id = 0; id < n_ids; id++) {
                const int32_t i02 = *(const int32_t *) (ids_host.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));

                const int64_t i11 = id % ne11;
                const int64_t i12 = iid1;
@@ -5892,6 +5921,12 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
                const int64_t i1 = id;
                const int64_t i2 = i12;

+            if (i02 == -1) { // skipped slot
+                SYCL_CHECK(CHECK_TRY_ERROR(
+                        stream->memset(dst_original + i1*nb1 + i2*nb2, 0, ne0*sizeof(float))));
+                continue;
+            }
+
            src0_row.data = src0_original + i02*nb02;
            src1_row.data = src1_original + i11*nb11 + i12*nb12;
            dst_row.data = dst_original + i1*nb1 + i2*nb2;
@@ -5913,20 +5948,43 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
        std::vector<int64_t> expert_row_offsets;
        // the sources (slot/token pairs) of contiguous rows to guide k_copy_src1_to_contiguous
        std::vector<mmid_row_mapping> & routed_row_src = ctx.mmid_row_mapping_host;
+        // the slot/token pairs of the skipped slots, whose dst rows are zeroed
+        std::vector<mmid_row_mapping> skipped_row_dst;

        mmid_counting_sort_rows(ids, ids_host.data(), n_ids, n_as, n_routed_rows,
-                                expert_row_counts, expert_row_offsets, routed_row_src);
+                                expert_row_counts, expert_row_offsets, routed_row_src, skipped_row_dst);
+
+        const int64_t n_valid_rows = (int64_t) routed_row_src.size();

        ggml_sycl_pool_alloc<mmid_row_mapping> dev_row_mapping(ctx.pool(), n_routed_rows);
        SYCL_CHECK(CHECK_TRY_ERROR(
-                stream->memcpy(dev_row_mapping.get(), routed_row_src.data(), n_routed_rows*sizeof(mmid_row_mapping))));
+                stream->memcpy(dev_row_mapping.get(), routed_row_src.data(), n_valid_rows*sizeof(mmid_row_mapping))));

        const unsigned int max_work_group_size = ggml_sycl_info().max_work_group_sizes[ctx.device];
        assert(max_work_group_size % (WARP_SIZE * WARP_SIZE) == 0);

-        {
+        if (!skipped_row_dst.empty()) {
+            ggml_sycl_pool_alloc<mmid_row_mapping> dev_skipped(ctx.pool(), skipped_row_dst.size());
+            SYCL_CHECK(CHECK_TRY_ERROR(
+                    stream->memcpy(dev_skipped.get(), skipped_row_dst.data(), skipped_row_dst.size()*sizeof(mmid_row_mapping))));
+
+            sycl::range<3> block_dims(1, 1, std::min((unsigned int)ne0, max_work_group_size));
+            sycl::range<3> grid_dims(1, 1, skipped_row_dst.size());
+            stream->submit([&](sycl::handler &cgh) {
+                mmid_row_mapping *__restrict dev_skipped_get = dev_skipped.get();
+
+                cgh.parallel_for(
+                    sycl::nd_range<3>(grid_dims * block_dims, block_dims),
+                    [=](sycl::nd_item<3> item_ct1) {
+                        k_zero_dst_rows(dst_original, dev_skipped_get, ne0, nb1, nb2, item_ct1);
+                    });
+            });
+            SYCL_CHECK(CHECK_TRY_ERROR(stream->wait()));
+        }
+
+        if (n_valid_rows > 0) {
            sycl::range<3> block_dims(1, 1, std::min((unsigned int)ne10, max_work_group_size));
-            sycl::range<3> grid_dims(1, 1, n_routed_rows);
+            sycl::range<3> grid_dims(1, 1, n_valid_rows);
            stream->submit([&](sycl::handler &cgh) {
                char *__restrict src1_contiguous_get =
                    src1_contiguous.get();
@@ -5992,9 +6050,9 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
            ggml_sycl_mul_mat(ctx, &src0_row, &src1_row, &dst_row);
        }

-        {
+        if (n_valid_rows > 0) {
            sycl::range<3> block_dims(1, 1, std::min((unsigned int)ne0, max_work_group_size));
-            sycl::range<3> grid_dims(1, 1, n_routed_rows);
+            sycl::range<3> grid_dims(1, 1, n_valid_rows);
            stream->submit([&](sycl::handler &cgh) {
                const char *__restrict dst_contiguous_get =
                    dst_contiguous.get();
--- a/Show More
+++ b/Show More