tests: tolerate a shared pool abort in test_completion_unified (#28759)

The expected success table holds when the four requests enter the shared
pool together. On a loaded runner they are admitted tens of milliseconds
apart, the slot lifetimes overlap differently and the pool overflows
while a short request is still resident. The decode failure aborts every
slot, so a request the table marks as successful comes back with the
context error instead of its generation.

Such a request now passes on that error too, while any other status, a
different error or a truncated generation still fails the test.

.github/workflows/fusion.yml

@@ -0,0 +1,67 @@
+name: Fusion
+
+on:
+  workflow_dispatch: # allows manual triggering
+  push:
+    branches:
+      - master
+    paths: [
+      '.github/workflows/fusion.yml',
+      'ggml/**',
+      'tests/fusion/**',
+      'tests/test-fusion.cpp'
+    ]
+
+  pull_request:
+    types: [opened, synchronize, reopened]
+    paths: [
+      '.github/workflows/fusion.yml',
+      'ggml/**',
+      'tests/fusion/**',
+      'tests/test-fusion.cpp'
+    ]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
+  cancel-in-progress: true
+
+env:
+  GGML_NLOOP: 3
+  GGML_N_THREADS: 1
+  LLAMA_ARG_LOG_COLORS: 1
+  LLAMA_ARG_LOG_PREFIX: 1
+  LLAMA_ARG_LOG_TIMESTAMPS: 1
+
+jobs:
+  # TODO: add jobs for other backends as they adopt the fusion debug API
+  metal:
+    runs-on: [self-hosted, macOS, ARM64]
+
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+
+      - name: Build
+        id: cmake_build
+        run: |
+          cmake -B build \
+            -DCMAKE_BUILD_TYPE=Release \
+            -DLLAMA_FATAL_WARNINGS=ON \
+            -DLLAMA_OPENSSL=OFF \
+            -DGGML_SCHED_NO_REALLOC=ON \
+            -DGGML_BLAS=OFF \
+            -DGGML_METAL=ON
+          time cmake --build build --config Release --target test-llama-archs -j $(sysctl -n hw.logicalcpu)
+          time cmake --build build --config Release --target test-fusion -j $(sysctl -n hw.logicalcpu)
+
+      - name: Generate models
+        id: generate_models
+        run: |
+          rm -rf build-ci-models && mkdir -p build-ci-models
+          ./build/bin/test-llama-archs -o build-ci-models
+
+      - name: Test fusion
+        id: test_fusion
+        run: |
+          ./build/bin/test-fusion --models build-ci-models --device MTL0 --check tests/fusion/MTL.csv

.github/workflows/server-sanitize.yml

@@ -75,7 +75,7 @@ jobs:
       - name: ccache-buckets-restore
         uses: ./.github/actions/ccache-buckets
         with:
-          key: server-sanitize
+          key: server-sanitize-${{ matrix.sanitizer }}
           folder: llama.cpp
           hf_bucket: ggml-org/cache
 
@@ -99,7 +99,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
         with:
-          key: server-sanitize
+          key: server-sanitize-${{ matrix.sanitizer }}
           folder: llama.cpp
           evict-old-files: 1d
           hf_bucket: ggml-org/cache

ci/run.sh

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

common/CMakeLists.txt

@@ -134,6 +134,8 @@ set_target_properties(${TARGET} PROPERTIES
 target_include_directories(${TARGET} PUBLIC .)
 target_link_libraries     (${TARGET} PUBLIC vendor::nlohmann vendor::sheredom)
 target_compile_features   (${TARGET} PUBLIC cxx_std_17)
+target_precompile_headers (${TARGET} PRIVATE common.h)
+target_precompile_headers (${TARGET} PRIVATE chat.h)
 
 if (LLAMA_SUBPROCESS)
     target_compile_definitions(${TARGET} PUBLIC LLAMA_SUBPROCESS)

common/speculative.cpp

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

common/speculative.h

@@ -61,7 +61,7 @@ struct common_speculative_draft_params {
     // can be used to constraint the max draft based on the remaining context size
     int32_t n_max = -1;
 
-    llama_pos   n_past;
+    llama_pos   pos0;
     llama_token id_last;
 
     // TODO: remove in the future by keeping track of the prompt from the _begin() call and the consecutive accept calls

docs/build-profiling.md

@@ -0,0 +1,122 @@
+## Build profiling
+This page is a working document for analyzing the current build and try to
+identify ways to improve the build time.
+
+### Requirements
+The profiling script requires clang to be used as the compiler tool chain and
+also requires that ClangBuildAnalyzer is installed.
+
+Mac:
+```console
+brew install clang-build-analyzer
+```
+
+Linux:
+```console
+git clone https://github.com/aras-p/ClangBuildAnalyzer.git
+cd ClangBuildAnalyzer
+cmake -B build -DCMAKE_BUILD_TYPE=Release
+cmake --build build -j$(nproc)
+sudo cp build/ClangBuildAnalyzer /usr/local/bin/
+```
+
+Windows: install LLVM/clang and Ninja (e.g. via the
+[LLVM releases page](https://github.com/llvm/llvm-project/releases) and
+`winget install Ninja-build.Ninja`), then build ClangBuildAnalyzer the same
+way as on Linux:
+```console
+git clone https://github.com/aras-p/ClangBuildAnalyzer.git
+cd ClangBuildAnalyzer
+cmake -B build -G Ninja -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_BUILD_TYPE=Release
+cmake --build build --config Release
+```
+Then add `ClangBuildAnalyzer\build` to `PATH`.
+
+### Usage
+Mac/Linux:
+```console
+$ ./scripts/build-profile.sh
+```
+
+Windows:
+```console
+> .\scripts\build-profile.ps1
+```
+
+Both accept `--full`/`-Full` (include Server, Tools, and Tests) and a jobs
+override (`-jN` / `-Jobs N`).
+
+Note: on Windows, `cmake` defaults to the Visual Studio generator, which
+ignores `CMAKE_C_COMPILER`/`CMAKE_CXX_COMPILER` and silently falls back to
+MSVC. `build-profile.ps1` passes `-G Ninja` so clang is actually used, this
+is required on ARM64.
+
+### Linux (Ubuntu 24.04)
+
+Environment:
+- Clang:     18.1.3 (Ubuntu clang version 18.1.3 (1ubuntu1))
+- libstdc++: GCC 13.3.0 (Ubuntu 13.3.0-6ubuntu2~24.04.1)
+- Target:    x86_64-pc-linux-gnu
+
+```console
++------------------------+-----+------------+------------+------------+
+| Build                  | TUs | Frontend   | Backend    | Total      |
++------------------------+-----+------------+------------+------------+
+| Minimal, master        | 249 |   468.2 s  |   270.3 s  |   738.5 s  |
+| Minimal, with PCH      | 253 |   177.1 s  |   265.8 s  |   442.9 s  |
+| Full,    master        | 396 |   811.0 s  |   692.2 s  | 1,503.2 s  |
+| Full,    with PCH      | 405 |   380.0 s  |   664.7 s  | 1,044.7 s  |
+| Full,    with PCH + UB | 264 |   357.7 s  |   635.7 s  |   993.4 s  |
++------------------------+-----+------------+------------+------------+
+
+PCH  = precompiled header.
+Full = includes building Server, Tools, and Tests.
+UB   = unity build for models
+```
+Note that the number of translation units (TUs) increases when using precompiled
+headers — each PCH target adds one extra TU for the precompilation step itself.
+
+### Mac (Apple M3)
+
+Environment:
+- Clang:  Apple clang version 17.0.0 (clang-1700.3.19.1)
+- libc++: ships with Apple clang 17.0.0 (Xcode toolchain)
+- Target: arm64-apple-macosx15.6
+
+```console
++------------------------+-----+------------+------------+------------+
+| Build                  | TUs | Frontend   | Backend    | Total      |
++------------------------+-----+------------+------------+------------+
+| Minimal, master        | 256 |   154.5 s  |    94.8 s  |   249.3 s  |
+| Minimal, with PCH      | 261 |    65.9 s  |    90.0 s  |   155.9 s  |
+| Full,    master        | 407 |   265.7 s  |   209.7 s  |   475.4 s  |
+| Full,    with PCH      | 414 |   154.6 s  |   197.5 s  |   352.1 s  |
+| Full,    with PCH + UB | 274 |   143.0 s  |   192.2 s  |   335.2 s  |
++------------------------+-----+------------+------------+------------+
+
+PCH = precompiled header.
+Full = includes building Server, Tools, and Tests.
+UB   = unity build for models
+```
+
+### Windows (ARM64)
+
+Environment:
+- Clang:  clang version 22.1.8 (LLVM, `C:\Program Files\LLVM`)
+- STL:    MSVC STL (Visual Studio 2022 Build Tools 14.44.35207)
+- Target: aarch64-pc-windows-msvc
+
+```console
++------------------------+-----+------------+------------+------------+
+| Build                  | TUs | Frontend   | Backend    | Total      |
++------------------------+-----+------------+------------+------------+
+| Minimal, master        | 249 |   159.4 s  |    82.2 s  |   241.6 s  |
+| Full,    master        | 373 |   337.2 s  |   167.4 s  |   504.6 s  |
+| Minimal, with PCH + UB | 113 |    62.3 s  |    82.4 s  |   144.7 s  |
+| Full,    with PCH + UB | 240 |   233.0 s  |   185.1 s  |   418.1 s  |
++------------------------+-----+------------+------------+------------+
+
+PCH = precompiled header.
+Full = includes building Server, Tools, and Tests.
+UB   = unity build for models
+```

examples/speculative-simple/speculative-simple.cpp

@@ -188,7 +188,7 @@ int main(int argc, char ** argv) {
             common_speculative_get_draft_params(spec, seq_id) = {
                 /* .drafting   = */ true,
                 /* .n_max      = */ n_draft_max,
-                /* .n_past     = */ n_past,
+                /* .pos0       = */ n_past,
                 /* .id_last    = */ id_last,
                 /* .prompt     = */ &prompt_tgt,
                 /* .result     = */ &draft, // output

ggml/src/ggml-backend.cpp

@@ -1705,6 +1705,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                     ggml_tensor * ids_tensor = node->src[2];
                     ggml_backend_t ids_backend = split_backend;
 
+                    if (ggml_nelements(ids_tensor) == 0) {
+                        continue;
+                    }
+
                     // if the ids tensor is also an input of the split, it may not have been copied yet to the split backend
                     // in that case, we use the original ids tensor
                     for (int i = input_id + 1; i < split->n_inputs; i++) {

ggml/src/ggml-cpu/CMakeLists.txt

@@ -675,6 +675,12 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
     target_compile_options(${GGML_CPU_NAME} PRIVATE ${ARCH_FLAGS})
     target_compile_definitions(${GGML_CPU_NAME} PRIVATE ${ARCH_DEFINITIONS})
 
+    if (CMAKE_C_COMPILER_ID STREQUAL "GNU" AND NOT GGML_SYSTEM_ARCH STREQUAL "x86")
+        message(STATUS "Skipping PCH for ${GGML_CPU_NAME}: GCC PCH is only enabled for x86 (arch: ${GGML_SYSTEM_ARCH})")
+    else()
+        target_precompile_headers(${GGML_CPU_NAME} PRIVATE ggml-impl.h)
+    endif()
+
     if (EMSCRIPTEN)
         set_target_properties(${GGML_CPU_NAME} PROPERTIES COMPILE_FLAGS "-msimd128")
     endif()

ggml/src/ggml-cpu/ops.h

@@ -18,7 +18,15 @@
 #endif
 #endif
 
+// -Winterference-size was introduced in GCC 12
+#if defined(__cplusplus) && defined(__GNUC__) && !defined(__clang__) && __GNUC__ >= 12
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Winterference-size"
+#endif
 static const size_t CACHE_LINE_SIZE_F32 = CACHE_LINE_SIZE/sizeof(float);
+#if defined(__cplusplus) && defined(__GNUC__) && !defined(__clang__) && __GNUC__ >= 12
+#pragma GCC diagnostic pop
+#endif
 
 // Work buffer size for im2col operations in CONV2D
 #define GGML_IM2COL_WORK_SIZE (16 * 1024 * 1024)

ggml/src/ggml-cuda/fattn-common.cuh

@@ -1133,12 +1133,21 @@ void launch_fattn(
 
     dim3 blocks_num;
     if (stream_k) {
-        // For short contexts it can be faster to have the SMs work on whole tiles because this lets us skip the fixup.
-        const int max_blocks = max_blocks_per_sm*nsm;
-        const int tiles_nwaves = (ntiles_dst + max_blocks - 1) / max_blocks;
-        const int tiles_efficiency_percent = 100 * ntiles_dst / (max_blocks*tiles_nwaves);
+        auto should_use_stream_k = [](const int cc, const int ntiles_dst, const int max_blocks, const int DKQ) {
+            const int tiles_nwaves             = (ntiles_dst + max_blocks - 1) / max_blocks;
+            const int tiles_efficiency_percent = 100 * ntiles_dst / (max_blocks*tiles_nwaves);
 
-        const bool use_stream_k = cc >= GGML_CUDA_CC_ADA_LOVELACE || amd_wmma_available(cc) || tiles_efficiency_percent < 75;
+            if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_ADA_LOVELACE) {
+                return true;
+            }
+            if (amd_wmma_available(cc) && DKQ == 64) {
+                return true; // TODO better configuration
+            }
+            return tiles_efficiency_percent < 75;
+        };
+
+        const int  max_blocks   = max_blocks_per_sm*nsm;
+        const bool use_stream_k = should_use_stream_k(cc, ntiles_dst, max_blocks, Q->ne[0]);
 
         blocks_num.x = ntiles_dst;
         blocks_num.y = 1;

ggml/src/ggml-cuda/fattn-mma-f16.cuh

@@ -158,8 +158,8 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co
 
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8,  64, 2,  32, 128, 128, 128, 1, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16,  64, 2,  32, 128, 128, 128, 1, true);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 128, 2,  64, 128, 128,  64, 1, true);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2,  64, 128, 128,  64, 1, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 256, 2,  64, 128, 128,  64, 1, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 256, 2,  64, 128, 128,  64, 1, true);
 
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 128, 2,  32, 160, 128, 128, 1, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 64, 128, 2,  32, 160, 128, 128, 1, true);
@@ -1826,7 +1826,7 @@ static __global__ void flash_attn_ext_f16(
 #endif // __CUDA_ARCH__ == GGML_CUDA_CC_TURING
 
 #if defined(AMD_WMMA_AVAILABLE)
-    if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 128) {
+    if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 256) {
         NO_DEVICE_CODE;
         return;
     }

ggml/src/ggml-cuda/fattn.cu

@@ -221,6 +221,24 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con
         }
     }
 
+    // On RDNA it is preferable to minimize wasted compute vs. duplicate I/O for the mask.
+    if (amd_wmma_available(cc)) {
+        if (use_gqa_opt && gqa_ratio % 8 == 0) {
+            ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 8>(ctx, dst);
+            return;
+        }
+
+        if (use_gqa_opt && gqa_ratio % 4 == 0) {
+            ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 4>(ctx, dst);
+            return;
+        }
+
+        if (use_gqa_opt && gqa_ratio % 2 == 0) {
+            ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 2>(ctx, dst);
+            return;
+        }
+    }
+
     if (use_gqa_opt && gqa_ratio > 4) {
         ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 8>(ctx, dst);
         return;
@@ -646,8 +664,9 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
         }
     }
 
-    // AMD WMMA is always faster than the tile kernel if the full tile width of 16 can be utilized.
-    if ((amd_wmma_available(cc) && gqa_opt_applies && Q->ne[0] <= 128) && Q->ne[0] != 40 && Q->ne[0] != 72 && Q->ne[1] * gqa_ratio_eff > 8) {
+    // AMD WMMA is faster than the tile kernel if the wide tiles with high arithmetic intensity can be utilized.
+    if ((amd_wmma_available(cc) && gqa_opt_applies && Q->ne[0] <= 256) && Q->ne[0] != 40 && Q->ne[0] != 72 &&
+            Q->ne[1] * gqa_ratio_eff > (Q->ne[0] <= 128 ? 8 : 16)) {
         return BEST_FATTN_KERNEL_MMA_F16;
     }
 

ggml/src/ggml-metal/CMakeLists.txt

@@ -10,6 +10,7 @@ ggml_add_backend_library(ggml-metal
                          ggml-metal-device.cpp
                          ggml-metal-common.cpp
                          ggml-metal-context.m
+                         ggml-metal-fusion.cpp
                          ggml-metal-ops.cpp
                          ggml-metal-tuning.cpp
                         )

ggml/src/ggml-metal/ggml-metal-common.cpp

@@ -1,4 +1,5 @@
 #include "ggml-metal-common.h"
+#include "ggml-metal-fusion.h"
 
 #include "ggml.h"
 #include "ggml-impl.h"
@@ -390,59 +391,31 @@ static std::vector<int> ggml_metal_graph_optimize_reorder(const std::vector<node
 }
 
 void ggml_graph_optimize(ggml_cgraph * gf) {
-    constexpr int MAX_FUSE = 16;
-
     const int n = gf->n_nodes;
 
-    enum ggml_op ops[MAX_FUSE];
-
     std::vector<node_info> nodes;
     nodes.reserve(gf->n_nodes);
 
     // fuse nodes:
     // we don't want to make reorders that break fusing, so we first pack all fusable tensors
     //   and perform the reorder over the fused nodes. after the reorder is done, we unfuse
+    //
+    // the fusable sequences are declared in the fusion table (ggml-metal-fuse.cpp), so the
+    // packing here is driven by the same patterns that the op encoders will later use
     for (int i = 0; i < n; i++) {
         node_info node = {
             /*.node =*/ gf->nodes[i],
             /*.fused =*/ {},
         };
 
-        // fuse only ops that start with these operations
-        // can be expanded when needed
-        if (node.op() == GGML_OP_ADD ||
-            node.op() == GGML_OP_NORM ||
-            node.op() == GGML_OP_RMS_NORM) {
-            ops[0] = node.op();
+        const int f = ggml_metal_fusion_max(gf, i);
 
-            int f = i + 1;
-            while (f < n && f < i + MAX_FUSE) {
-                // conservatively allow fusing only these ops
-                // can be expanded when needed
-                if (gf->nodes[f]->op != GGML_OP_ADD &&
-                    gf->nodes[f]->op != GGML_OP_MUL &&
-                    gf->nodes[f]->op != GGML_OP_NORM &&
-                    gf->nodes[f]->op != GGML_OP_RMS_NORM) {
-                    break;
-                }
-                ops[f - i] = gf->nodes[f]->op;
-                f++;
-            }
+        // add the fused tensors into the node info so we can unfuse them later
+        for (int k = 1; k < f; k++) {
+            ++i;
 
-            f -= i;
-            for (; f > 1; f--) {
-                if (ggml_can_fuse(gf, i, ops, f)) {
-                    break;
-                }
-            }
-
-            // add the fused tensors into the node info so we can unfuse them later
-            for (int k = 1; k < f; k++) {
-                ++i;
-
-                // the .dst() becomes the last fused tensor
-                node.add_fused(gf->nodes[i]);
-            }
+            // the .dst() becomes the last fused tensor
+            node.add_fused(gf->nodes[i]);
         }
 
         nodes.push_back(std::move(node));

ggml/src/ggml-metal/ggml-metal-context.h

@@ -33,6 +33,7 @@ ggml_metal_event_t ggml_metal_get_ev_cpy(ggml_metal_t ctx);
 
 void ggml_metal_set_n_cb            (ggml_metal_t ctx, int n_cb);
 void ggml_metal_set_abort_callback  (ggml_metal_t ctx, ggml_abort_callback abort_callback, void * user_data);
+
 bool ggml_metal_supports_family     (ggml_metal_t ctx, int family);
 void ggml_metal_capture_next_compute(ggml_metal_t ctx);
 

ggml/src/ggml-metal/ggml-metal-context.m

@@ -6,6 +6,7 @@
 #import "ggml-metal-impl.h"
 #import "ggml-metal-common.h"
 #import "ggml-metal-ops.h"
+#import "ggml-metal-fusion.h"
 
 #import <Foundation/Foundation.h>
 
@@ -36,15 +37,12 @@ struct ggml_metal {
     // additional, inference-time compiled pipelines
     ggml_metal_pipelines_t pipelines_ext;
 
-    bool use_fusion;
     bool use_concurrency;
     bool use_graph_optimize;
 
     int debug_graph;
-    int debug_fusion;
 
-    // how many times a given op was fused
-    uint64_t fuse_cnt[GGML_OP_COUNT];
+    struct ggml_metal_fusion_info * finfo;
 
     // capture state
     int capture_compute;
@@ -139,7 +137,6 @@ ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) {
 
         res->d_queue = dispatch_queue_create("ggml-metal", DISPATCH_QUEUE_CONCURRENT);
 
-        res->use_fusion      = getenv("GGML_METAL_FUSION_DISABLE") == nil;
         res->use_concurrency = getenv("GGML_METAL_CONCURRENCY_DISABLE") == nil;
 
         {
@@ -147,20 +144,19 @@ ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) {
             res->debug_graph = val ? atoi(val) : 0;
         }
 
-        {
-            const char * val = getenv("GGML_METAL_FUSION_DEBUG");
-            res->debug_fusion = val ? atoi(val) : 0;
-        }
-
         res->use_graph_optimize = true;
 
         if (getenv("GGML_METAL_GRAPH_OPTIMIZE_DISABLE") != NULL) {
             res->use_graph_optimize = false;
         }
 
-        memset(res->fuse_cnt, 0, sizeof(res->fuse_cnt));
+        res->finfo = ggml_metal_device_get_fusion_info(dev);
+        if (ggml_metal_fusion_info_stats(res->finfo)) {
+            ggml_metal_fusion_info_labels_init(res->finfo);
+            res->n_cb = 0;
+        }
 
-        GGML_LOG_INFO("%s: use fusion         = %s\n", __func__, res->use_fusion         ? "true" : "false");
+        GGML_LOG_INFO("%s: use fusion         = %s\n", __func__, ggml_metal_fusion_info_enabled(res->finfo) ? "true" : "false");
         GGML_LOG_INFO("%s: use concurrency    = %s\n", __func__, res->use_concurrency    ? "true" : "false");
         GGML_LOG_INFO("%s: use graph optimize = %s\n", __func__, res->use_graph_optimize ? "true" : "false");
 
@@ -222,15 +218,18 @@ void ggml_metal_free(ggml_metal_t ctx) {
         ctx->pipelines_ext = nil;
     }
 
-    if (ctx->debug_fusion > 0) {
+    if (ggml_metal_fusion_info_debug(ctx->finfo) > 0) {
         GGML_LOG_DEBUG("%s: fusion stats:\n", __func__);
-        for (int i = 0; i < GGML_OP_COUNT; i++) {
-            if (ctx->fuse_cnt[i] == 0) {
+
+        const int n_fusions = ggml_metal_fusion_info_n_fusions(ctx->finfo);
+        for (int i = 0; i < n_fusions; i++) {
+            const uint64_t count = ggml_metal_fusion_info_count(ctx->finfo, i);
+            if (count == 0) {
                 continue;
             }
 
             // note: cannot use ggml_log here
-            GGML_LOG_DEBUG("%s: - %s: %" PRIu64 "\n", __func__, ggml_op_name((enum ggml_op) i), ctx->fuse_cnt[i]);
+            GGML_LOG_DEBUG("%s: - %s: %" PRIu64 "\n", __func__, ggml_metal_fusion_info_label(ctx->finfo, i), count);
         }
     }
 
@@ -481,10 +480,17 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph *
     @autoreleasepool {
         ctx->gf = gf;
 
-        ctx->n_nodes_0 = MIN(n_main, gf->n_nodes);
-        ctx->n_nodes_1 = gf->n_nodes - ctx->n_nodes_0;
+        if (ctx->n_cb == 0) {
+            // single-threaded encoding: the whole graph is encoded by one command buffer
+            ctx->n_nodes_0      = gf->n_nodes;
+            ctx->n_nodes_1      = 0;
+            ctx->n_nodes_per_cb = 0;
+        } else {
+            ctx->n_nodes_0      = MIN(n_main, gf->n_nodes);
+            ctx->n_nodes_1      = gf->n_nodes - ctx->n_nodes_0;
 
-        ctx->n_nodes_per_cb = (ctx->n_nodes_1 + ctx->n_cb - 1) / ctx->n_cb;
+            ctx->n_nodes_per_cb = (ctx->n_nodes_1 + ctx->n_cb - 1) / ctx->n_cb;
+        }
 
         if (ctx->capture_compute >= 0) {
             ctx->capture_compute--;
@@ -682,6 +688,12 @@ ggml_metal_event_t ggml_metal_get_ev_cpy(ggml_metal_t ctx) {
 }
 
 void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) {
+    // when fusion stats are collected the graph must be encoded by a single thread so the
+    // counters are race-free; override whatever the caller requested
+    if (ggml_metal_fusion_info_stats(ctx->finfo)) {
+        n_cb = 0;
+    }
+
     if (ctx->n_cb != n_cb) {
         ctx->n_cb = MIN(n_cb, GGML_METAL_MAX_COMMAND_BUFFERS);
 
@@ -717,13 +729,12 @@ void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) {
             ctx->dev,
             cmd_buf,
             ctx->gf,
+            ctx->finfo,
             idx_start,
             idx_end,
-            ctx->use_fusion,
             ctx->use_concurrency,
             ctx->capture_compute,
-            ctx->debug_graph,
-            ctx->debug_fusion);
+            ctx->debug_graph);
 
         for (int idx = 0; idx < ggml_metal_op_n_nodes(ctx_op); ++idx) {
             const int res = ggml_metal_op_encode(ctx_op, idx);

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -932,12 +932,24 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
                 nsg = N_SG_IQ2_XXS;
                 nr0 = N_R0_IQ2_XXS;
                 smem = 256*8+128;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ2_XXS_SPLIT;
+                    split = true;
+                }
             } break;
         case GGML_TYPE_IQ2_XS:
             {
                 nsg = N_SG_IQ2_XS;
                 nr0 = N_R0_IQ2_XS;
                 smem = 512*8+128;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ2_XS_SPLIT;
+                    split = true;
+                }
             } break;
         case GGML_TYPE_IQ3_XXS:
             {
@@ -957,21 +969,45 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
                 nsg = N_SG_IQ3_S;
                 nr0 = N_R0_IQ3_S;
                 smem = 512*4;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ3_S_SPLIT;
+                    split = true;
+                }
             } break;
         case GGML_TYPE_IQ2_S:
             {
                 nsg = N_SG_IQ2_S;
                 nr0 = N_R0_IQ2_S;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ2_S_SPLIT;
+                    split = true;
+                }
             } break;
         case GGML_TYPE_IQ1_S:
             {
                 nsg = N_SG_IQ1_S;
                 nr0 = N_R0_IQ1_S;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ1_S_SPLIT;
+                    split = true;
+                }
             } break;
         case GGML_TYPE_IQ1_M:
             {
                 nsg = N_SG_IQ1_M;
                 nr0 = N_R0_IQ1_M;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ1_M_SPLIT;
+                    split = true;
+                }
             } break;
         case GGML_TYPE_IQ4_NL:
             {
@@ -1177,12 +1213,24 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m
                 nsg = N_SG_IQ2_XXS;
                 nr0 = N_R0_IQ2_XXS;
                 smem = 256*8+128;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ2_XXS_SPLIT;
+                    split = true;
+                }
             } break;
         case GGML_TYPE_IQ2_XS:
             {
                 nsg = N_SG_IQ2_XS;
                 nr0 = N_R0_IQ2_XS;
                 smem = 512*8+128;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ2_XS_SPLIT;
+                    split = true;
+                }
             } break;
         case GGML_TYPE_IQ3_XXS:
             {
@@ -1202,21 +1250,45 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m
                 nsg = N_SG_IQ3_S;
                 nr0 = N_R0_IQ3_S;
                 smem = 512*4;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ3_S_SPLIT;
+                    split = true;
+                }
             } break;
         case GGML_TYPE_IQ2_S:
             {
                 nsg = N_SG_IQ2_S;
                 nr0 = N_R0_IQ2_S;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ2_S_SPLIT;
+                    split = true;
+                }
             } break;
         case GGML_TYPE_IQ1_S:
             {
                 nsg = N_SG_IQ1_S;
                 nr0 = N_R0_IQ1_S;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ1_S_SPLIT;
+                    split = true;
+                }
             } break;
         case GGML_TYPE_IQ1_M:
             {
                 nsg = N_SG_IQ1_M;
                 nr0 = N_R0_IQ1_M;
+
+                const int nb32 = ne00/32;
+                if (nb32 < 32 && (32 % nb32) == 0) {
+                    nr0 = N_R0_IQ1_M_SPLIT;
+                    split = true;
+                }
             } break;
         case GGML_TYPE_IQ4_NL:
             {

ggml/src/ggml-metal/ggml-metal-device.h

@@ -325,6 +325,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
 
 const struct ggml_metal_device_props * ggml_metal_device_get_props(ggml_metal_device_t dev);
 
+struct ggml_metal_fusion_info;
+
+// the device-owned fusion debugging context (NULL unless fusion debugging is enabled)
+struct ggml_metal_fusion_info * ggml_metal_device_get_fusion_info(ggml_metal_device_t dev);
+
 //
 // device buffers
 //

ggml/src/ggml-metal/ggml-metal-device.m

@@ -1,4 +1,5 @@
 #import "ggml-metal-device.h"
+#import "ggml-metal-fusion.h"
 
 #import "ggml-impl.h"
 #import "ggml-backend-impl.h"
@@ -896,6 +897,9 @@ struct ggml_metal_device {
 
     struct ggml_metal_device_props props;
 
+    // shared fusion debugging context
+    struct ggml_metal_fusion_info * finfo;
+
     // virtual address for GPU memory allocations
     atomic_uintptr_t addr_virt;
 };
@@ -1274,6 +1278,13 @@ ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) {
                     dev->props.max_working_set_size   = dev->mtl_device.maxBufferLength;
                 }
 
+                {
+                    const char * val = getenv("GGML_METAL_FUSION_DEBUG");
+                    dev->finfo = ggml_metal_fusion_info_init(
+                            getenv("GGML_METAL_FUSION_DISABLE") == nil,
+                            val ? atoi(val) : 0);
+                }
+
                 snprintf(dev->props.name, sizeof(dev->props.name), "%s%d", "MTL", device);
                 const char * gpu_name = [[dev->mtl_device name] UTF8String];
                 if (n_devices > 1) {
@@ -1348,6 +1359,8 @@ void ggml_metal_device_free(ggml_metal_device_t dev) {
     assert(dev != NULL);
 
     @autoreleasepool {
+        ggml_metal_fusion_info_free(dev->finfo);
+
         ggml_metal_rsets_free(dev->rsets);
 
         ggml_metal_library_free(dev->library);
@@ -1935,6 +1948,10 @@ static void ggml_metal_device_disable_tensor(ggml_metal_device_t dev) {
     dev->props.has_tensor = false;
 }
 
+struct ggml_metal_fusion_info * ggml_metal_device_get_fusion_info(ggml_metal_device_t dev) {
+    return dev->finfo;
+}
+
 //
 // device buffers
 //

ggml/src/ggml-metal/ggml-metal-fusion.cpp

@@ -0,0 +1,502 @@
+#include "ggml-metal-fusion.h"
+
+#include "ggml-backend-impl.h"
+#include "ggml-metal-device.h"
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+// ---- helpers -------------------------------------------------------------
+
+// true if two tensors live in the same Metal buffer
+static bool ggml_metal_fusion_same_buffer(const ggml_tensor * a, const ggml_tensor * b) {
+    if (!a || !b) {
+        return false;
+    }
+
+    ggml_backend_buffer_t ba = a->view_src ? a->view_src->buffer : a->buffer;
+    ggml_backend_buffer_t bb = b->view_src ? b->view_src->buffer : b->buffer;
+
+    ggml_metal_buffer_t ca = (ggml_metal_buffer_t) ba->context;
+    ggml_metal_buffer_t cb = (ggml_metal_buffer_t) bb->context;
+
+    return ggml_metal_buffer_get_id(ca, a).metal == ggml_metal_buffer_get_id(cb, b).metal;
+}
+
+// ---- pattern checks ------------------------------------------------------
+
+// NORM/RMS_NORM + MUL + ADD: the weight/bias of each fused step must match the norm input
+// width, be contiguous rows, and the fused outputs must stay F32
+static bool ggml_metal_fusion_check_norm(
+        const ggml_metal_fusion      * fusion,
+        const ggml_tensor * const    * nodes,
+              ggml_metal_fusion_mode   mode) {
+    GGML_UNUSED(mode);
+
+    GGML_ASSERT(fusion->n_ops >= 2);
+
+    for (int j = 1; j < fusion->n_ops; j++) {
+        // the fused MUL/ADD must read the previous node as src0
+        if (nodes[j]->src[0] != nodes[j - 1]) {
+            return false;
+        }
+
+        // the weight/bias must have the same row width as the norm input
+        if (nodes[j]->src[1]->ne[0] != nodes[0]->ne[0]) {
+            return false;
+        }
+
+        if (!ggml_is_contiguous_rows(nodes[j]->src[1])) {
+            return false;
+        }
+
+        if (nodes[j]->type != GGML_TYPE_F32) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+// ADD x N: each ADD reads the previous ADD as src0, and all addends must share layout
+// (and, in FULL mode, live in the same Metal buffer)
+static bool ggml_metal_fusion_check_add_chain(
+        const ggml_metal_fusion      * fusion,
+        const ggml_tensor * const    * nodes,
+              ggml_metal_fusion_mode   mode) {
+    GGML_ASSERT(fusion->n_ops >= 2);
+
+    for (int j = 1; j < fusion->n_ops; j++) {
+        if (nodes[j]->src[0] != nodes[j - 1]) {
+            return false;
+        }
+
+        if (!ggml_are_same_layout(nodes[j]->src[1], nodes[j - 1]->src[1])) {
+            return false;
+        }
+
+        if (mode == GGML_METAL_FUSION_FULL) {
+            if (!ggml_metal_fusion_same_buffer(nodes[j]->src[1], nodes[0]->src[1])) {
+                return false;
+            }
+        }
+    }
+
+    return true;
+}
+
+// GATED_DELTA_NET + CPY: the trailing cpy scatters the gdn state snapshots into the recurrent
+// cache, so the gdn kernel writes them straight to the cache and the cpy is elided.
+// mirrors ggml_metal_op_can_fuse_gdn_cache (PR #25788). the gdn output has other consumers (the
+// attn scores view), so unlike the other patterns this is not an elision chain: the structural
+// checks live entirely in this callback (unsafe = true).
+static bool ggml_metal_fusion_check_gdn_cache(
+        const ggml_metal_fusion      * fusion,
+        const ggml_tensor * const    * nodes,
+              ggml_metal_fusion_mode    mode) {
+    GGML_UNUSED(fusion);
+
+    const ggml_tensor * gdn = nodes[0];
+    const ggml_tensor * cpy = nodes[1];
+
+    // the kernel skips the snapshot tail, so the gdn output must not be a graph output
+    if (gdn->type != GGML_TYPE_F32 || (gdn->flags & GGML_TENSOR_FLAG_OUTPUT)) {
+        return false;
+    }
+
+    if (cpy->op != GGML_OP_CPY || (cpy->flags & GGML_TENSOR_FLAG_OUTPUT)) {
+        return false;
+    }
+
+    const int64_t S_v      = gdn->src[2]->ne[0];
+    const int64_t H        = gdn->src[2]->ne[1];
+    const int64_t n_tokens = gdn->src[2]->ne[2];
+    const int64_t n_seqs   = gdn->src[2]->ne[3];
+    const int64_t K        = ggml_get_op_params_i32(gdn, 0);
+    const size_t  tail_off = ggml_row_size(GGML_TYPE_F32, S_v * H * n_tokens * n_seqs);
+
+    const int64_t D         = S_v * S_v * H;
+    const int64_t n_written = std::min<int64_t>(n_tokens, K);
+
+    const ggml_tensor * src = cpy->src[0]; // gdn snapshot tail view
+    const ggml_tensor * dst = cpy->src[1]; // cache view
+
+    // src must be this gdn's snapshot tail (contiguous, at the tail offset)
+    if (src->op != GGML_OP_VIEW || src->view_src != gdn ||
+        src->view_offs != tail_off || !ggml_is_contiguous(src)) {
+        return false;
+    }
+
+    const int64_t expected_ne[GGML_MAX_DIMS] = { D, n_seqs, n_written, 1 };
+    if (dst->type != GGML_TYPE_F32 ||
+        !std::equal(expected_ne, expected_ne + GGML_MAX_DIMS, dst->ne) ||
+        dst->nb[0] != ggml_type_size(GGML_TYPE_F32) ||
+        dst->nb[1] != ggml_row_size(GGML_TYPE_F32, D)) {
+        return false;
+    }
+
+    if (mode == GGML_METAL_FUSION_FULL) {
+        // the cache must be allocated so the kernel can write straight to its buffer
+        if (dst->data == nullptr) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+// MUL + SIN + SQR + MUL + ADD (snake activation)
+static bool ggml_metal_fusion_check_snake(
+        const ggml_metal_fusion      * fusion,
+        const ggml_tensor * const    * nodes,
+              ggml_metal_fusion_mode   mode) {
+    GGML_UNUSED(fusion);
+    GGML_UNUSED(mode);
+
+    const ggml_tensor * mul0     = nodes[0];
+    const ggml_tensor * sin_node = nodes[1];
+    const ggml_tensor * sqr      = nodes[2];
+    const ggml_tensor * mul1     = nodes[3];
+    const ggml_tensor * add      = nodes[4];
+
+    // x carries the full activation shape, a is the broadcast operand
+    const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1];
+    const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0];
+
+    // mul1 reads sqr and inv_b in either operand order
+    const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0];
+
+    // closure check: the trailing add reads the same x as the leading mul
+    const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0];
+
+    // x is in the supported whitelist and every chain intermediate shares x's type.
+    // a and inv_b bind as device const float * in the kernel, so they stay F32.
+    const bool types_ok =
+        (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) &&
+        (a->type    == GGML_TYPE_F32) && (inv_b->type    == GGML_TYPE_F32) &&
+        (mul0->type == x->type)       && (sin_node->type == x->type) &&
+        (sqr->type  == x->type)       && (mul1->type     == x->type) &&
+        (add->type  == x->type);
+
+    // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D
+    const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1];
+    const bool dim_ok =
+        (x->ne[2]     == 1) && (x->ne[3]     == 1) &&
+        (add->ne[2]   == 1) && (add->ne[3]   == 1) &&
+        (a->ne[2]     == 1) && (a->ne[3]     == 1) &&
+        (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1);
+
+    // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous
+    const bool contig_ok =
+        ggml_is_contiguous(x) && ggml_is_contiguous(add) &&
+        ggml_is_contiguous(a) && ggml_is_contiguous(inv_b);
+
+    return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x;
+}
+
+// ---- patterns ------------------------------------------------------------
+
+static const ggml_op ops_norm_mul[]         = { GGML_OP_NORM, GGML_OP_MUL };
+static const ggml_op ops_norm_mul_add[]     = { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD };
+static const ggml_op ops_rms_norm_mul[]     = { GGML_OP_RMS_NORM, GGML_OP_MUL };
+static const ggml_op ops_rms_norm_mul_add[] = { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD };
+
+static const ggml_op ops_add_2[] = { GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_add_3[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_add_4[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_add_5[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_add_6[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_add_7[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD };
+static const ggml_op ops_snake[] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD };
+
+static const ggml_op ops_gdn_cache[] = { GGML_OP_GATED_DELTA_NET, GGML_OP_CPY };
+
+static const ggml_metal_fusion ggml_metal_fusions[] = {
+    { GGML_METAL_FUSION_NORM_MUL,     ops_norm_mul,         2, false, ggml_metal_fusion_check_norm },
+    { GGML_METAL_FUSION_NORM_MUL_ADD, ops_norm_mul_add,     3, false, ggml_metal_fusion_check_norm },
+    { GGML_METAL_FUSION_NORM_MUL,     ops_rms_norm_mul,     2, false, ggml_metal_fusion_check_norm },
+    { GGML_METAL_FUSION_NORM_MUL_ADD, ops_rms_norm_mul_add, 3, false, ggml_metal_fusion_check_norm },
+    { GGML_METAL_FUSION_ADD_CHAIN,    ops_add_2,            2, false, ggml_metal_fusion_check_add_chain },
+    { GGML_METAL_FUSION_ADD_CHAIN,    ops_add_3,            3, false, ggml_metal_fusion_check_add_chain },
+    { GGML_METAL_FUSION_ADD_CHAIN,    ops_add_4,            4, false, ggml_metal_fusion_check_add_chain },
+    { GGML_METAL_FUSION_ADD_CHAIN,    ops_add_5,            5, false, ggml_metal_fusion_check_add_chain },
+    { GGML_METAL_FUSION_ADD_CHAIN,    ops_add_6,            6, false, ggml_metal_fusion_check_add_chain },
+    { GGML_METAL_FUSION_ADD_CHAIN,    ops_add_7,            7, false, ggml_metal_fusion_check_add_chain },
+    { GGML_METAL_FUSION_SNAKE,        ops_snake,            5, false, ggml_metal_fusion_check_snake },
+    { GGML_METAL_FUSION_GDN_CACHE,    ops_gdn_cache,        2, true,  ggml_metal_fusion_check_gdn_cache },
+};
+
+const ggml_metal_fusion * ggml_metal_fusion_all(int * n) {
+    *n = (int) sizeof(ggml_metal_fusions) / sizeof(ggml_metal_fusions[0]);
+
+    return ggml_metal_fusions;
+}
+
+// ---- shared fusion info ---------------------------------------------------
+
+static std::string ggml_metal_fusion_label(const ggml_metal_fusion * fusion) {
+    GGML_ASSERT(fusion != nullptr);
+
+    std::string label;
+    for (int j = 0; j < fusion->n_ops; j++) {
+        if (j > 0) {
+            label += '+';
+        }
+        label += ggml_op_name(fusion->ops[j]);
+    }
+    return label;
+}
+
+struct ggml_metal_fusion_info {
+    std::vector<std::string> labels;
+    std::vector<uint64_t>    counts;
+    bool enabled;
+    bool stats;
+    bool labels_set;
+    int  debug;
+};
+
+struct ggml_metal_fusion_info * ggml_metal_fusion_info_init(bool enabled, int debug) {
+    ggml_metal_fusion_info * finfo = new ggml_metal_fusion_info;
+    finfo->enabled    = enabled;
+    finfo->stats      = debug > 0;
+    finfo->labels_set = false;
+    finfo->debug      = debug;
+
+    if (finfo->stats) {
+        ggml_metal_fusion_info_labels_init(finfo);
+    }
+
+    return finfo;
+}
+
+void ggml_metal_fusion_info_free(struct ggml_metal_fusion_info * finfo) {
+    delete finfo;
+}
+
+bool ggml_metal_fusion_info_enabled(const struct ggml_metal_fusion_info * finfo) {
+    return finfo->enabled;
+}
+
+bool ggml_metal_fusion_info_stats(const struct ggml_metal_fusion_info * finfo) {
+    return finfo->stats;
+}
+
+int ggml_metal_fusion_info_debug(const struct ggml_metal_fusion_info * finfo) {
+    return finfo->debug;
+}
+
+int ggml_metal_fusion_info_n_fusions(const struct ggml_metal_fusion_info * finfo) {
+    return (int) finfo->labels.size();
+}
+
+const char * ggml_metal_fusion_info_label(const struct ggml_metal_fusion_info * finfo, int idx) {
+    GGML_ASSERT(idx >= 0 && idx < (int) finfo->labels.size());
+    return finfo->labels[idx].c_str();
+}
+
+uint64_t ggml_metal_fusion_info_count(const struct ggml_metal_fusion_info * finfo, int idx) {
+    GGML_ASSERT(idx >= 0 && idx < (int) finfo->counts.size());
+    return finfo->counts[idx];
+}
+
+void ggml_metal_fusion_info_count_fusion(struct ggml_metal_fusion_info * finfo, const struct ggml_metal_fusion * fusion) {
+    if (!finfo->stats || fusion == nullptr) {
+        return;
+    }
+
+    int n = 0;
+    const ggml_metal_fusion * all = ggml_metal_fusion_all(&n);
+
+    int idx = -1;
+    for (int i = 0; i < n; i++) {
+        if (&all[i] == fusion) {
+            idx = i;
+            break;
+        }
+    }
+
+    if (idx >= 0 && idx < (int) finfo->counts.size()) {
+        finfo->counts[idx]++;
+    }
+}
+
+void ggml_metal_fusion_info_set_enabled(struct ggml_metal_fusion_info * finfo, bool enabled) {
+    finfo->enabled = enabled;
+}
+
+void ggml_metal_fusion_info_labels_init(struct ggml_metal_fusion_info * finfo) {
+    if (finfo->labels_set) {
+        return;
+    }
+
+    int n = 0;
+    const ggml_metal_fusion * all = ggml_metal_fusion_all(&n);
+
+    finfo->labels.clear();
+    finfo->counts.assign(n, 0);
+    finfo->labels.reserve(n);
+
+    for (int i = 0; i < n; i++) {
+        finfo->labels.emplace_back(ggml_metal_fusion_label(&all[i]));
+    }
+
+    finfo->labels_set = true;
+}
+
+void ggml_metal_fusion_info_stats_init(struct ggml_metal_fusion_info * finfo) {
+    finfo->stats = true;
+    ggml_metal_fusion_info_labels_init(finfo);
+}
+
+void ggml_metal_fusion_info_stats_reset(struct ggml_metal_fusion_info * finfo) {
+    std::fill(finfo->counts.begin(), finfo->counts.end(), 0);
+}
+
+int ggml_metal_fusion_info_stats_get(const struct ggml_metal_fusion_info * finfo, const char ** labels, uint64_t * counts, int n) {
+    const int n_fusions = (int) finfo->labels.size();
+
+    if (labels == nullptr) {
+        return n_fusions;
+    }
+
+    const int n_fill = std::min(n, n_fusions);
+    for (int i = 0; i < n_fill; i++) {
+        labels[i] = finfo->labels[i].c_str();
+        if (counts != nullptr) {
+            counts[i] = finfo->counts[i];
+        }
+    }
+
+    return n_fill;
+}
+
+// ---- queries -------------------------------------------------------------
+
+// find the longest pattern matching the node sequence starting at idx
+// (idx is a position in node_idxs, which maps to graph node indices)
+const ggml_metal_fusion * ggml_metal_fusion_next(
+        const ggml_cgraph * gf,
+        const int * node_idxs,
+        int n_idxs,
+        int idx,
+        ggml_metal_fusion_mode mode,
+        int * n_out) {
+    int n = 0;
+    const ggml_metal_fusion * all = ggml_metal_fusion_all(&n);
+
+    const ggml_metal_fusion * res = nullptr;
+    int best = 1;
+
+    for (int i = 0; i < n; i++) {
+        const ggml_metal_fusion * fusion = &all[i];
+
+        // only look for a longer match than the current best
+        if (fusion->n_ops <= best) {
+            continue;
+        }
+        if (idx + fusion->n_ops > n_idxs) {
+            continue;
+        }
+
+        const ggml_tensor * nodes[GGML_METAL_FUSION_MAX];
+
+        // the op sequence must match exactly
+        bool ok = true;
+        for (int j = 0; j < fusion->n_ops; j++) {
+            nodes[j] = gf->nodes[node_idxs[idx + j]];
+            if (nodes[j]->op != fusion->ops[j]) {
+                ok = false;
+                break;
+            }
+        }
+        if (!ok) {
+            continue;
+        }
+
+        if (!fusion->unsafe) {
+            // common element-wise chain constraints: each node reads the previous one,
+            // and all nodes have the same shape
+            for (int j = 1; j < fusion->n_ops && ok; j++) {
+                if (nodes[j]->src[0] != nodes[j - 1] && nodes[j]->src[1] != nodes[j - 1]) {
+                    ok = false;
+                    break;
+                }
+                if (!ggml_are_same_shape(nodes[j], nodes[j - 1])) {
+                    ok = false;
+                    break;
+                }
+            }
+            if (!ok) {
+                continue;
+            }
+
+            // all current fusions are single-output elision chains, so the last node is the only output
+            // TODO: multi-output fusions: store pattern-relative offsets in the table and translate them here
+            int outputs_buf[1];
+            outputs_buf[0] = node_idxs[idx + fusion->n_ops - 1];
+
+            // structural subgraph checks (op sequence, elidable uses, view containment)
+            if (!ggml_can_fuse_subgraph_ext(gf, node_idxs + idx, fusion->n_ops, fusion->ops, outputs_buf, 1)) {
+                continue;
+            }
+        }
+
+        // pattern-specific checks (the sole validator for unsafe patterns)
+        if (fusion->check && !fusion->check(fusion, nodes, mode)) {
+            continue;
+        }
+
+        best = fusion->n_ops;
+        res = fusion;
+    }
+
+    *n_out = best;
+
+    return res;
+}
+
+// optimize phase: maximum number of nodes starting at idx (a raw sequential graph index) that
+// could be fused, chaining patterns back-to-back. matching runs on the same filtered (view
+// transparent) node sequence that the compute phase uses, so the returned count is the raw index
+// span from idx to the last matched node (intermediate views are packed along).
+int ggml_metal_fusion_max(const ggml_cgraph * gf, int idx) {
+    // an empty/view node cannot start a pattern - pack it alone
+    if (ggml_op_is_empty(gf->nodes[idx]->op) || ggml_is_empty(gf->nodes[idx])) {
+        return 1;
+    }
+
+    // collect the non-empty node indices starting at idx
+    int idxs[GGML_METAL_FUSION_MAX];
+    int n_idxs = 0;
+    for (int i = idx; i < gf->n_nodes && n_idxs < GGML_METAL_FUSION_MAX; i++) {
+        if (!ggml_op_is_empty(gf->nodes[i]->op) && !ggml_is_empty(gf->nodes[i])) {
+            idxs[n_idxs++] = i;
+        }
+    }
+    if (n_idxs == 0) {
+        return 1;
+    }
+
+    int total = 0;
+    int i_f = 0;
+
+    while (i_f < n_idxs && total < GGML_METAL_FUSION_MAX) {
+        int len = 1;
+        const ggml_metal_fusion * fusion = ggml_metal_fusion_next(gf, idxs, n_idxs, i_f, GGML_METAL_FUSION_STRUCTURAL, &len);
+        if (!fusion || total + len > GGML_METAL_FUSION_MAX) {
+            break;
+        }
+
+        total += len;
+        i_f += len;
+    }
+
+    if (i_f == 0) {
+        return 1;
+    }
+
+    // map the matched non-empty nodes back to the raw index span (views are included)
+    return std::min(GGML_METAL_FUSION_MAX, idxs[i_f - 1] - idx + 1);
+}

ggml/src/ggml-metal/ggml-metal-fusion.h

@@ -0,0 +1,104 @@
+// single source of truth for the fusions supported by the Metal backend
+//
+// every fusable subgraph is declared exactly once as a ggml_metal_fusion entry in
+// the table in ggml-metal-fusion.cpp. both the graph optimizer (ggml_metal_fusion_max)
+// and the op encoders (ggml_metal_fusion_next) consult this same table, so the two
+// phases can never disagree about what can be fused.
+
+#pragma once
+
+#include "ggml-impl.h"
+
+#include <stdbool.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// the maximum number of nodes that can be fused in a single kernel
+// (also the maximum length of a packed fusion group during graph optimization)
+#define GGML_METAL_FUSION_MAX 16
+
+typedef enum ggml_metal_fusion_mode {
+    // structural checks only; used by the graph optimizer, at which point the graph
+    // tensors are not allocated yet, so buffer placement cannot be verified
+    GGML_METAL_FUSION_STRUCTURAL = 0,
+    // full checks, including buffer placement; used by the op encoders
+    GGML_METAL_FUSION_FULL,
+} ggml_metal_fusion_mode;
+
+// identifier of each fusion pattern so the op encoders know which kernel to use
+typedef enum ggml_metal_fusion_id {
+    GGML_METAL_FUSION_NONE = 0,
+    GGML_METAL_FUSION_NORM_MUL,     // NORM/RMS_NORM + MUL
+    GGML_METAL_FUSION_NORM_MUL_ADD, // NORM/RMS_NORM + MUL + ADD
+    GGML_METAL_FUSION_ADD_CHAIN,    // ADD x N (N in [2, 7])
+    GGML_METAL_FUSION_SNAKE,        // MUL + SIN + SQR + MUL + ADD
+    GGML_METAL_FUSION_GDN_CACHE,    // GATED_DELTA_NET + CPY (write snapshots into the recurrent cache)
+} ggml_metal_fusion_id;
+
+struct ggml_metal_fusion {
+    ggml_metal_fusion_id id;
+
+    const enum ggml_op * ops;        // op sequence (fixed length)
+    int                  n_ops;      // number of ops
+
+    // if unsafe: the generic chain/shape + ggml_can_fuse_subgraph checks are skipped and the
+    // check callback below is the sole validator (used for patterns that are not elision chains,
+    // e.g. the gdn + cache-cpy write-through fusion)
+    bool unsafe;
+
+    // extra backend constraints on top of ggml_can_fuse_subgraph
+    // nodes[j] is the j-th node of the pattern
+    bool (*check)(const struct ggml_metal_fusion   * fusion,
+                  const struct ggml_tensor * const * nodes,
+                        ggml_metal_fusion_mode       mode);
+};
+
+typedef struct ggml_metal_fusion ggml_metal_fusion;
+
+// the single table of all fusions supported by the Metal backend
+const ggml_metal_fusion * ggml_metal_fusion_all(int * n);
+
+// ---- shared fusion info ---------------------------------------------------
+
+// shared fusion debugging context, owned by the device; newly created backend contexts for that
+// device register with it so the fusion counters are race-free and accumulate across contexts.
+struct ggml_metal_fusion_info; // defined in ggml-metal-fusion.cpp
+
+struct ggml_metal_fusion_info * ggml_metal_fusion_info_init(bool enabled, int debug);
+void ggml_metal_fusion_info_free(struct ggml_metal_fusion_info * finfo);
+
+bool ggml_metal_fusion_info_enabled(const struct ggml_metal_fusion_info * finfo);
+bool ggml_metal_fusion_info_stats  (const struct ggml_metal_fusion_info * finfo);
+int  ggml_metal_fusion_info_debug  (const struct ggml_metal_fusion_info * finfo);
+
+int          ggml_metal_fusion_info_n_fusions(const struct ggml_metal_fusion_info * finfo);
+const char * ggml_metal_fusion_info_label    (const struct ggml_metal_fusion_info * finfo, int idx);
+uint64_t     ggml_metal_fusion_info_count    (const struct ggml_metal_fusion_info * finfo, int idx);
+
+void ggml_metal_fusion_info_count_fusion(struct ggml_metal_fusion_info * finfo, const struct ggml_metal_fusion * fusion);
+void ggml_metal_fusion_info_set_enabled (struct ggml_metal_fusion_info * finfo, bool enabled);
+
+void ggml_metal_fusion_info_stats_init (      struct ggml_metal_fusion_info * finfo);
+void ggml_metal_fusion_info_stats_reset(      struct ggml_metal_fusion_info * finfo);
+int  ggml_metal_fusion_info_stats_get  (const struct ggml_metal_fusion_info * finfo, const char ** labels, uint64_t * counts, int n);
+void ggml_metal_fusion_info_labels_init(      struct ggml_metal_fusion_info * finfo);
+
+// compute phase: longest fusion starting at idx (a position in node_idxs) that matches in `mode`.
+// returns the matching pattern (nullptr if no fusion) and sets *n_out to the number of nodes consumed.
+const ggml_metal_fusion * ggml_metal_fusion_next(
+        const struct ggml_cgraph * gf,
+        const int * node_idxs,
+        int n_idxs,
+        int idx,
+        ggml_metal_fusion_mode mode,
+        int * n_out);
+
+// optimize phase: maximum number of nodes starting at idx (a raw sequential graph index) that
+// could be fused, chaining patterns back-to-back. returns at least 1.
+int ggml_metal_fusion_max(const struct ggml_cgraph * gf, int idx);
+
+#ifdef __cplusplus
+}
+#endif

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -62,18 +62,23 @@
 
 #define N_R0_IQ1_S 4
 #define N_SG_IQ1_S 2
+#define N_R0_IQ1_S_SPLIT 8
 
 #define N_R0_IQ1_M 4
 #define N_SG_IQ1_M 2
+#define N_R0_IQ1_M_SPLIT 8
 
 #define N_R0_IQ2_XXS 4
 #define N_SG_IQ2_XXS 2
+#define N_R0_IQ2_XXS_SPLIT 8
 
 #define N_R0_IQ2_XS 4
 #define N_SG_IQ2_XS 2
+#define N_R0_IQ2_XS_SPLIT 8
 
 #define N_R0_IQ2_S 4
 #define N_SG_IQ2_S 2
+#define N_R0_IQ2_S_SPLIT 8
 
 #define N_R0_IQ3_XXS 4
 #define N_SG_IQ3_XXS 2
@@ -81,6 +86,7 @@
 
 #define N_R0_IQ3_S 4
 #define N_SG_IQ3_S 2
+#define N_R0_IQ3_S_SPLIT 8
 
 #define N_R0_IQ4_NL 2
 #define N_SG_IQ4_NL 2
@@ -979,6 +985,7 @@ typedef struct {
     uint64_t nb1;
     uint64_t nb2;
     uint64_t nb3;
+    uint64_t nb_out; // 0 => snapshots are appended after the attn scores (unfused)
 } ggml_metal_kargs_gated_delta_net;
 
 typedef struct {

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -7,6 +7,7 @@
 #include "ggml-metal-impl.h"
 #include "ggml-metal-common.h"
 #include "ggml-metal-device.h"
+#include "ggml-metal-fusion.h"
 #include "ggml-metal-tuning.h"
 
 #include <cassert>
@@ -31,24 +32,22 @@ struct ggml_metal_op {
         ggml_metal_device_t dev,
         ggml_metal_cmd_buf_t cmd_buf,
         ggml_cgraph * gf,
+        ggml_metal_fusion_info * finfo,
         int  idx_start,
         int  idx_end,
-        bool use_fusion,
         bool use_concurrency,
         bool use_capture,
-        int  debug_graph,
-        int  debug_fusion) {
+        int  debug_graph) {
         this->dev             = dev;
         this->lib             = ggml_metal_device_get_library(dev);
         this->enc             = ggml_metal_encoder_init(cmd_buf, use_concurrency);
         this->mem_ranges      = ggml_mem_ranges_init(debug_graph);
+        this->finfo           = finfo;
         this->idx_start       = idx_start;
         this->idx_end         = idx_end;
-        this->use_fusion      = use_fusion;
         this->use_concurrency = use_concurrency;
         this->use_capture     = use_capture;
         this->debug_graph     = debug_graph;
-        this->debug_fusion    = debug_fusion;
         this->gf              = gf;
 
         idxs.reserve(gf->n_nodes);
@@ -78,15 +77,24 @@ struct ggml_metal_op {
         return ggml_graph_node(gf, idxs[i]);
     }
 
-    bool can_fuse(int i0, const ggml_op * ops, int n_ops) const {
-        assert(use_fusion);
+    // consult the fusion table for the longest pattern starting at i0
+    // returns the matching pattern (nullptr if no fusion) and sets *n_out to the number of nodes
+    const ggml_metal_fusion * can_fuse(int i0, enum ggml_metal_fusion_mode mode, int * n_out) const {
+        assert(use_fusion());
         assert(i0 >= 0 && i0 < n_nodes());
 
-        if (i0 + n_ops > n_nodes()) {
-            return false;
-        }
+        return ggml_metal_fusion_next(gf, idxs.data(), (int) idxs.size(), i0, mode, n_out);
+    }
 
-        return ggml_can_fuse_ext(gf, idxs.data() + i0, ops, n_ops);
+    // whether to attempt fusion; the toggle lives in the shared fusion debugging context owned
+    // by the device (initialized from GGML_METAL_FUSION_DISABLE, overridable by the test)
+    bool use_fusion() const {
+        return ggml_metal_fusion_info_enabled(finfo);
+    }
+
+    // record that a fusion fired, indexed by the matching table entry
+    void count_fusions(const ggml_metal_fusion * fusion) const {
+        ggml_metal_fusion_info_count_fusion(finfo, fusion);
     }
 
     ggml_metal_device_t  dev;
@@ -94,12 +102,13 @@ struct ggml_metal_op {
     ggml_metal_encoder_t enc;
     ggml_mem_ranges_t    mem_ranges;
 
-    bool use_fusion;
+    // shared fusion debugging context
+    ggml_metal_fusion_info * finfo;
+
     bool use_concurrency;
     bool use_capture;
 
     int debug_graph;
-    int debug_fusion;
 
 private:
     ggml_cgraph * gf;
@@ -115,24 +124,22 @@ ggml_metal_op_t ggml_metal_op_init(
         ggml_metal_device_t dev,
         ggml_metal_cmd_buf_t cmd_buf,
         ggml_cgraph * gf,
+        ggml_metal_fusion_info * finfo,
         int idx_start,
         int idx_end,
-        bool use_fusion,
         bool use_concurrency,
         bool use_capture,
-        int debug_graph,
-        int debug_fusion) {
+        int debug_graph) {
     ggml_metal_op_t res = new ggml_metal_op(
         dev,
         cmd_buf,
         gf,
+        finfo,
         idx_start,
         idx_end,
-        use_fusion,
         use_concurrency,
         use_capture,
-        debug_graph,
-        debug_fusion);
+        debug_graph);
 
     return res;
 }
@@ -1868,6 +1875,8 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) {
     ggml_metal_library_t lib = ctx->lib;
     ggml_metal_encoder_t enc = ctx->enc;
 
+    const bool use_fusion = ctx->use_fusion();
+    const int  debug_fusion = ggml_metal_fusion_info_debug(ctx->finfo);
 
     GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
     GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
@@ -1880,6 +1889,31 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) {
 
     auto pipeline = ggml_metal_library_get_pipeline_gated_delta_net(lib, op);
 
+    // when fused with the trailing cache cpy, the snapshots are written straight into the
+    // recurrent cache and the cpy is skipped (see GGML_METAL_FUSION_GDN_CACHE)
+    ggml_metal_buffer_id bid_out = ggml_metal_get_buffer_id(op);
+    uint64_t nb_out = 0;
+    int n_fuse = 1;
+
+    if (use_fusion) {
+        int n = 1;
+        const ggml_metal_fusion * fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n);
+
+        if (fusion && fusion->id == GGML_METAL_FUSION_GDN_CACHE) {
+            const ggml_tensor * dst_cache = ctx->node(idx + 1)->src[1]; // cache view
+
+            bid_out = ggml_metal_get_buffer_id(dst_cache);
+            nb_out  = dst_cache->nb[2]/sizeof(float);
+            n_fuse = 2;
+
+            ctx->count_fusions(fusion);
+
+            if (debug_fusion > 1) {
+                GGML_LOG_DEBUG("%s: fuse: GATED_DELTA_NET + CPY\n", __func__);
+            }
+        }
+    }
+
     int ida = 0;
 
     ggml_metal_kargs_gated_delta_net args = {
@@ -1918,23 +1952,25 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) {
         /*.nb1  =*/ nb1,
         /*.nb2  =*/ nb2,
         /*.nb3  =*/ nb3,
+        /*.nb_out =*/ nb_out,
     };
 
     ggml_metal_encoder_set_pipeline(enc, pipeline);
-    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args),                  ida++);
+    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args),                  ida++); // args
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), ida++); // q
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[1]), ida++); // k
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[2]), ida++); // v
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[3]), ida++); // gate
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[4]), ida++); // beta
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[5]), ida++); // state
-    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         ida++); // dst
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         ida++); // dst (attn)
+    ggml_metal_encoder_set_buffer  (enc, bid_out,                              ida++); // state_out
 
     const int nsg = pipeline.nsg;
 
     ggml_metal_encoder_dispatch_threadgroups(enc, op->src[2]->ne[0]/nsg, op->src[2]->ne[1], op->src[2]->ne[3], 32, nsg, 1);
 
-    return 1;
+    return n_fuse;
 }
 
 int ggml_metal_op_solve_tri(ggml_metal_op_t ctx, int idx) {
@@ -3718,56 +3754,20 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
     return 1;
 }
 
-// Snake activation autofuse: mul -> sin -> sqr -> mul -> add
-static bool ggml_metal_op_can_fuse_snake(ggml_metal_op_t ctx, int idx) {
-    static constexpr ggml_op snake_ops[5] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD };
-
-    if (ctx->node(idx)->op != GGML_OP_MUL || !ctx->can_fuse(idx, snake_ops, 5)) {
-        return false;
-    }
-
-    const ggml_tensor * mul0     = ctx->node(idx + 0);
-    const ggml_tensor * sin_node = ctx->node(idx + 1);
-    const ggml_tensor * sqr      = ctx->node(idx + 2);
-    const ggml_tensor * mul1     = ctx->node(idx + 3);
-    const ggml_tensor * add      = ctx->node(idx + 4);
-
-    // x carries the full activation shape, a is the broadcast operand
-    const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1];
-    const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0];
-
-    // mul1 reads sqr and inv_b in either operand order
-    const ggml_tensor * inv_b    = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0];
-
-    // closure check: the trailing add reads the same x as the leading mul
-    const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0];
-
-    // x is in the supported whitelist and every chain intermediate shares x's type.
-    // a and inv_b bind as device const float * in the kernel, so they stay F32.
-    const bool types_ok =
-        (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) &&
-        (a->type    == GGML_TYPE_F32) && (inv_b->type    == GGML_TYPE_F32) &&
-        (mul0->type == x->type)       && (sin_node->type == x->type) &&
-        (sqr->type  == x->type)       && (mul1->type     == x->type) &&
-        (add->type  == x->type);
-    // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D
-    const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1];
-    const bool dim_ok =
-        (x->ne[2]     == 1) && (x->ne[3]     == 1) &&
-        (add->ne[2]   == 1) && (add->ne[3]   == 1) &&
-        (a->ne[2]     == 1) && (a->ne[3]     == 1) &&
-        (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1);
-    // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous
-    const bool contig_ok =
-        ggml_is_contiguous(x) && ggml_is_contiguous(add) &&
-        ggml_is_contiguous(a) && ggml_is_contiguous(inv_b);
-
-    return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x;
-}
-
 int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
-    if (ctx->use_fusion && ggml_metal_op_can_fuse_snake(ctx, idx)) {
-        return ggml_metal_op_snake_fused(ctx, idx);
+    int n_fuse = 1;
+    const ggml_metal_fusion * fusion = nullptr;
+
+    if (ctx->use_fusion()) {
+        int n = 1;
+        fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n);
+        n_fuse = n;
+
+        // snake activation autofuse: mul -> sin -> sqr -> mul -> add
+        if (fusion && fusion->id == GGML_METAL_FUSION_SNAKE) {
+            ctx->count_fusions(fusion);
+            return ggml_metal_op_snake_fused(ctx, idx);
+        }
     }
 
     ggml_tensor * op = ctx->node(idx);
@@ -3775,9 +3775,9 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
     ggml_metal_library_t lib = ctx->lib;
     ggml_metal_encoder_t enc = ctx->enc;
 
-    const bool use_fusion = ctx->use_fusion;
+    const bool use_fusion = ctx->use_fusion();
 
-    const int debug_fusion = ctx->debug_fusion;
+    const int debug_fusion = ggml_metal_fusion_info_debug(ctx->finfo);
 
     GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
     GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
@@ -3822,57 +3822,19 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
         /*.o1   =*/ { bid_src1.offs },
     };
 
-    ggml_op fops[8];
-
-    int n_fuse = 1;
-
     // c[0] = add(a,    b[0])
     // c[1] = add(c[0], b[1])
     // c[2] = add(c[1], b[2])
     // ...
-    if (use_fusion) {
-        fops[0] = GGML_OP_ADD;
-        fops[1] = GGML_OP_ADD;
-        fops[2] = GGML_OP_ADD;
-        fops[3] = GGML_OP_ADD;
-        fops[4] = GGML_OP_ADD;
-        fops[5] = GGML_OP_ADD;
-        fops[6] = GGML_OP_ADD;
-        fops[7] = GGML_OP_ADD;
-
-        // note: in metal, we sometimes encode the graph in parallel so we have to avoid fusing ops
-        //       across splits. idx_end indicates the last node in the current split
-        for (n_fuse = 0; n_fuse <= 6; ++n_fuse) {
-            if (!ctx->can_fuse(idx + n_fuse, fops + n_fuse, 2)) {
-                break;
-            }
-
-            ggml_tensor * f0 = ctx->node(idx + n_fuse);
-            ggml_tensor * f1 = ctx->node(idx + n_fuse + 1);
-
-            if (f0 != f1->src[0]) {
-                break;
-            }
-
-            // b[0] === b[1] === ...
-            if (!ggml_are_same_layout(f0->src[1], f1->src[1])) {
-                break;
-            }
-
-            // only fuse ops if src1 is in the same Metal buffer
-            ggml_metal_buffer_id bid_fuse = ggml_metal_get_buffer_id(f1->src[1]);
-            if (bid_fuse.metal != bid_src1.metal) {
-                break;
-            }
-
-            //ctx->fuse_cnt[ops[n_fuse + 1]->op]++;
-
-            args.o1[n_fuse + 1] = bid_fuse.offs;
+    if (use_fusion && fusion && fusion->id == GGML_METAL_FUSION_ADD_CHAIN) {
+        // the offsets of the fused addends are relative to the start of the src1 buffer
+        for (int i = 1; i < n_fuse; i++) {
+            args.o1[i] = ggml_metal_get_buffer_id(ctx->node(idx + i)->src[1]).offs;
         }
 
-        ++n_fuse;
+        ctx->count_fusions(fusion);
 
-        if (debug_fusion > 1 && n_fuse > 1) {
+        if (debug_fusion > 1) {
             GGML_LOG_DEBUG("%s: fuse: ADD x %d\n", __func__, n_fuse);
         }
     }
@@ -4080,9 +4042,9 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) {
     ggml_metal_library_t lib = ctx->lib;
     ggml_metal_encoder_t enc = ctx->enc;
 
-    const bool use_fusion = ctx->use_fusion;
+    const bool use_fusion = ctx->use_fusion();
 
-    const int debug_fusion = ctx->debug_fusion;
+    const int debug_fusion = ggml_metal_fusion_info_debug(ctx->finfo);
 
     GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
     GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
@@ -4110,8 +4072,6 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) {
         /*.nbf3   =*/ { nb03 },
     };
 
-    ggml_op fops[8];
-
     int n_fuse = 1;
 
     ggml_metal_buffer_id bid_fuse[2] = { bid_src0, bid_src0 };
@@ -4120,55 +4080,35 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) {
     // d[1] = mul(d[0], b)
     // d[2] = add(d[1], c)
     if (use_fusion) {
-        fops[0] = op->op;
-        fops[1] = GGML_OP_MUL;
-        fops[2] = GGML_OP_ADD;
+        int n = 1;
+        const ggml_metal_fusion * fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n);
 
-        for (n_fuse = 0; n_fuse <= 1; ++n_fuse) {
-            if (!ctx->can_fuse(idx + n_fuse, fops + n_fuse, 2)) {
-                break;
+        if (fusion && (fusion->id == GGML_METAL_FUSION_NORM_MUL || fusion->id == GGML_METAL_FUSION_NORM_MUL_ADD)) {
+            n_fuse = n;
+
+            ctx->count_fusions(fusion);
+
+            for (int i = 1; i < n_fuse; i++) {
+                const ggml_tensor * fn = ctx->node(idx + i);
+
+                bid_fuse[i - 1] = ggml_metal_get_buffer_id(fn->src[1]);
+
+                args.nef1[i] = fn->src[1]->ne[1];
+                args.nef2[i] = fn->src[1]->ne[2];
+                args.nef3[i] = fn->src[1]->ne[3];
+
+                args.nbf1[i] = fn->src[1]->nb[1];
+                args.nbf2[i] = fn->src[1]->nb[2];
+                args.nbf3[i] = fn->src[1]->nb[3];
             }
 
-            ggml_tensor * f0 = ctx->node(idx + n_fuse);
-            ggml_tensor * f1 = ctx->node(idx + n_fuse + 1);
-
-            if (f0 != f1->src[0]) {
-                break;
-            }
-
-            if (f1->src[1]->ne[0] != op->ne[0]) {
-                break;
-            }
-
-            if (!ggml_is_contiguous_rows(f1->src[1])) {
-                break;
-            }
-
-            if (f1->type != GGML_TYPE_F32) {
-                break;
-            }
-
-            //ctx->fuse_cnt[f1->op]++;
-
-            bid_fuse[n_fuse] = ggml_metal_get_buffer_id(f1->src[1]);
-
-            args.nef1[n_fuse + 1] = f1->src[1]->ne[1];
-            args.nef2[n_fuse + 1] = f1->src[1]->ne[2];
-            args.nef3[n_fuse + 1] = f1->src[1]->ne[3];
-
-            args.nbf1[n_fuse + 1] = f1->src[1]->nb[1];
-            args.nbf2[n_fuse + 1] = f1->src[1]->nb[2];
-            args.nbf3[n_fuse + 1] = f1->src[1]->nb[3];
-        }
-
-        ++n_fuse;
-
-        if (debug_fusion > 1 && n_fuse > 1) {
-            if (n_fuse == 2) {
-                GGML_LOG_DEBUG("%s: fuse: %s + MUL\n", __func__, ggml_op_name(op->op));
-            }
-            if (n_fuse == 3) {
-                GGML_LOG_DEBUG("%s: fuse: %s + MUL + ADD\n", __func__, ggml_op_name(op->op));
+            if (debug_fusion > 1) {
+                if (n_fuse == 2) {
+                    GGML_LOG_DEBUG("%s: fuse: %s + MUL\n", __func__, ggml_op_name(op->op));
+                }
+                if (n_fuse == 3) {
+                    GGML_LOG_DEBUG("%s: fuse: %s + MUL + ADD\n", __func__, ggml_op_name(op->op));
+                }
             }
         }
     }

ggml/src/ggml-metal/ggml-metal-ops.h

@@ -8,17 +8,18 @@ extern "C" {
 
 typedef struct ggml_metal_op * ggml_metal_op_t;
 
+struct ggml_metal_fusion; // forward decl (ggml-metal-device.h)
+
 ggml_metal_op_t ggml_metal_op_init(
         ggml_metal_device_t dev,
         ggml_metal_cmd_buf_t cmd_buf,
         struct ggml_cgraph * gf,
+        struct ggml_metal_fusion_info * finfo,
         int  idx_start,
         int  idx_end,
-        bool use_fusion,
         bool use_concurrency,
         bool use_capture,
-        int  debug_graph,
-        int  debug_fusion);
+        int  debug_graph);
 
 void ggml_metal_op_free(ggml_metal_op_t ctx);
 

ggml/src/ggml-metal/ggml-metal.cpp

@@ -4,6 +4,7 @@
 #include "ggml-backend-impl.h"
 
 #include "ggml-metal-device.h"
+#include "ggml-metal-fusion.h"
 #include "ggml-metal-context.h"
 #include "ggml-metal-ops.h"
 #include "ggml-metal-tuning.h"
@@ -906,6 +907,30 @@ static const char * ggml_backend_metal_tuning_device_token(ggml_backend_dev_t de
     return ggml_metal_device_id_token(ggml_metal_device_get_props(ctx_dev)->device_id);
 }
 
+// generic fusion debugging API (ad-hoc proc-address mechanism): the test resolves the device
+// fusion context once and passes that opaque handle to the rest of the functions
+typedef void * ggml_backend_fusion_t;
+
+static ggml_backend_fusion_t ggml_backend_metal_fusion_get(ggml_backend_dev_t dev) {
+    return ggml_metal_device_get_fusion_info((ggml_metal_device_t)dev->context);
+}
+
+static void ggml_backend_metal_fusion_stats_init(ggml_backend_fusion_t finfo) {
+    ggml_metal_fusion_info_stats_init((struct ggml_metal_fusion_info *) finfo);
+}
+
+static void ggml_backend_metal_fusion_stats_reset(ggml_backend_fusion_t finfo) {
+    ggml_metal_fusion_info_stats_reset((struct ggml_metal_fusion_info *) finfo);
+}
+
+static int ggml_backend_metal_fusion_stats_get(ggml_backend_fusion_t finfo, const char ** labels, uint64_t * counts, int n) {
+    return ggml_metal_fusion_info_stats_get((struct ggml_metal_fusion_info *) finfo, labels, counts, n);
+}
+
+static void ggml_backend_metal_fusion_set_enabled(ggml_backend_fusion_t finfo, bool enabled) {
+    ggml_metal_fusion_info_set_enabled((struct ggml_metal_fusion_info *) finfo, enabled);
+}
+
 static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const char * name) {
     if (strcmp(name, "ggml_backend_get_features") == 0) {
         return (void *)ggml_backend_metal_get_features;
@@ -928,6 +953,23 @@ static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const
     if (strcmp(name, "ggml_backend_metal_tuning_device_token") == 0) {
         return (void *)ggml_backend_metal_tuning_device_token;
     }
+    // generic fusion debugging API (ad-hoc proc-address mechanism, not part of the official
+    // ggml backend interface yet; a backend that adopts it exports these exact names)
+    if (strcmp(name, "ggml_backend_fusion_get") == 0) {
+        return (void *)ggml_backend_metal_fusion_get;
+    }
+    if (strcmp(name, "ggml_backend_fusion_stats_init") == 0) {
+        return (void *)ggml_backend_metal_fusion_stats_init;
+    }
+    if (strcmp(name, "ggml_backend_fusion_stats_reset") == 0) {
+        return (void *)ggml_backend_metal_fusion_stats_reset;
+    }
+    if (strcmp(name, "ggml_backend_fusion_stats_get") == 0) {
+        return (void *)ggml_backend_metal_fusion_stats_get;
+    }
+    if (strcmp(name, "ggml_backend_fusion_set_enabled") == 0) {
+        return (void *)ggml_backend_metal_fusion_set_enabled;
+    }
 
     return NULL;
 

ggml/src/ggml-metal/kernels/gated_delta_net.metal

@@ -15,6 +15,7 @@ kernel void kernel_gated_delta_net_impl(
         device const char * b,
         device const char * s,
         device       char * dst,
+        device       char * dst_fuse,
         uint3 tgpig[[threadgroup_position_in_grid]],
         uint3 tpitg[[thread_position_in_threadgroup]],
         uint3   ntg[[threads_per_threadgroup]])  {
@@ -65,6 +66,12 @@ kernel void kernel_gated_delta_net_impl(
     // per-(seq,head) offset within a slot
     const uint state_out_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;
 
+    // when fused with the cache cpy, write the snapshots straight into the cache buffer using
+    // the slot stride; otherwise append them after the attn scores (nb_out == 0)
+    const bool fused = args.nb_out > 0;
+    const device float * state_out = fused ? (device float *)dst_fuse : (device float *)dst + attn_size;
+    const uint slot_stride = fused ? (uint)args.nb_out : state_size_per_snap;
+
     for (short t = 0; t < args.ne22; t++) {
         float s_k = 0.0f;
 
@@ -116,7 +123,7 @@ kernel void kernel_gated_delta_net_impl(
         if (K > 1) {
             const int target_slot = (int)args.ne22 - 1 - (int)t;
             if (target_slot >= 0 && target_slot < (int)K) {
-                device float * dst_state = (device float *) (dst) + attn_size + (uint)target_slot * state_size_per_snap + state_out_base;
+                device float * dst_state = (device float *)state_out + (uint)target_slot * slot_stride + state_out_base;
                 FOR_UNROLL (short j = 0; j < NSG; j++) {
                     const short is = tx*NSG + j;
                     dst_state[is] = ls[j];
@@ -126,7 +133,7 @@ kernel void kernel_gated_delta_net_impl(
     }
 
     if (K == 1) {
-        device float * dst_state = (device float *) (dst) + attn_size + state_out_base;
+        device float * dst_state = (device float *)state_out + state_out_base;
         FOR_UNROLL (short j = 0; j < NSG; j++) {
             const short is = tx*NSG + j;
             dst_state[is] = ls[j];
@@ -158,6 +165,7 @@ kernel void kernel_gated_delta_net_impl(
         device const char * b,
         device const char * s,
         device       char * dst,
+        device       char * dst_fuse,
         uint3 tgpig[[threadgroup_position_in_grid]],
         uint3 tpitg[[thread_position_in_threadgroup]],
         uint3   ntg[[threads_per_threadgroup]])  {
@@ -230,7 +238,13 @@ kernel void kernel_gated_delta_net_impl(
         dst_attn += args.ne21*S_v;
     }
 
-    device float * dst_state  = (device float *) (dst) + args.ne23*args.ne22*args.ne21*S_v + (i23*args.ne21 + i21)*S_v*S_v + i20;
+    // when fused with the cache cpy, write the snapshots straight into the cache buffer using
+    // the slot stride; otherwise append them after the attn scores (nb_out == 0)
+    const bool fused = args.nb_out > 0;
+    const device float * state_out = fused ? (device float *)dst_fuse : (device float *)dst + args.ne23*args.ne22*args.ne21*S_v;
+    const uint slot_stride = fused ? (uint)args.nb_out : S_v*S_v;
+
+    device float * dst_state  = (device float *)state_out + (i23*args.ne21 + i21)*slot_stride + i20;
     device T     * dstt_state = (device T     *) (dst_state);
 
     FOR_UNROLL (short j = 0; j < NSG; j++) {

ggml/src/ggml-metal/kernels/mul_mm.metal

@@ -496,6 +496,13 @@ kernel void kernel_mul_mm_id(
         + args.nb11*i11
         + args.nb10*iy);
 
+    // skip the upper half of the token tile when the expert did not fill it
+    constexpr short NR1H = NR1/2;
+
+    const bool has_hi = nr1 > NR1H;
+
+    const short lb1 = (short) tiitg/NL1; // 0 .. NR1-1, this thread's row of the B tile
+
 #ifndef GGML_METAL_HAS_TENSOR
     S0_8x8 ma[4];
     S1_8x8 mb[2];
@@ -505,15 +512,22 @@ kernel void kernel_mul_mm_id(
     for (short i = 0; i < 8; i++){
         mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
     }
+
+    // simdgroups 2,3 own rows NR1H..NR1-1
+    const bool sg_active = has_hi || sgitg < 2;
 #else
-    auto tA = tensor<threadgroup S0, dextents<int32_t, 2>, tensor_inline>(sa, dextents<int32_t, 2>(NK,  NR0));
-    auto tB = tensor<threadgroup S1, dextents<int32_t, 2>, tensor_inline>(sb, dextents<int32_t, 2>(NR1, NK ));
+    auto tA  = tensor<threadgroup S0, dextents<int32_t, 2>, tensor_inline>(sa, dextents<int32_t, 2>(NK, NR0));
+
+    // sb is [NR1][NK] row-major
+    auto tB0 = tensor<threadgroup S1, dextents<int32_t, 2>, tensor_inline>(sb,             dextents<int32_t, 2>(NK, NR1H));
+    auto tB1 = tensor<threadgroup S1, dextents<int32_t, 2>, tensor_inline>(sb + NR1H*NK,   dextents<int32_t, 2>(NK, NR1H));
 
     mpp::tensor_ops::matmul2d<
-        mpp::tensor_ops::matmul2d_descriptor(NR1, NR0, NK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
+        mpp::tensor_ops::matmul2d_descriptor(NR1H, NR0, NK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
         execution_simdgroups<4>> mm;
 
-    auto cT = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB), float>();
+    auto cT0 = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB0), float>();
+    auto cT1 = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB1), float>();
 #endif
 
     for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
@@ -656,37 +670,45 @@ kernel void kernel_mul_mm_id(
         threadgroup_barrier(mem_flags::mem_threadgroup);
 
 #ifndef GGML_METAL_HAS_TENSOR
-        // load matrices from threadgroup memory and conduct outer products
-        threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
-        threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
+        if (sg_active) {
+            // load matrices from threadgroup memory and conduct outer products
+            threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
+            threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
 
-        FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
-            simdgroup_barrier(mem_flags::mem_none);
+            FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
+                simdgroup_barrier(mem_flags::mem_none);
 
-            FOR_UNROLL (short i = 0; i < 4; i++) {
-                simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
+                FOR_UNROLL (short i = 0; i < 4; i++) {
+                    simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
+                }
+
+                simdgroup_barrier(mem_flags::mem_none);
+
+                FOR_UNROLL (short i = 0; i < 2; i++) {
+                    simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
+                }
+
+                simdgroup_barrier(mem_flags::mem_none);
+
+                FOR_UNROLL (short i = 0; i < 8; i++){
+                    simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
+                }
+
+                lsma += 8*64;
+                lsmb += 4*64;
             }
-
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 2; i++) {
-                simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
-            }
-
-            simdgroup_barrier(mem_flags::mem_none);
-
-            FOR_UNROLL (short i = 0; i < 8; i++){
-                simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
-            }
-
-            lsma += 8*64;
-            lsmb += 4*64;
         }
 #else
-        auto sA = tA.slice(0, 0);
-        auto sB = tB.slice(0, 0);
+        auto sA  = tA.slice(0, 0);
+        auto sB0 = tB0.slice(0, 0);
 
-        mm.run(sB, sA, cT);
+        mm.run(sB0, sA, cT0);
+
+        if (has_hi) {
+            auto sB1 = tB1.slice(0, 0);
+
+            mm.run(sB1, sA, cT1);
+        }
 #endif
     }
 
@@ -694,13 +716,20 @@ kernel void kernel_mul_mm_id(
     threadgroup_barrier(mem_flags::mem_threadgroup);
 
 #ifdef GGML_METAL_HAS_TENSOR
-    auto tC = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc, dextents<int32_t, 2>(NR0, NR1));
-    cT.store(tC);
-#else
-    threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
+    auto tC0 = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc,             dextents<int32_t, 2>(NR0, NR1H));
+    cT0.store(tC0);
 
-    for (short i = 0; i < 8; i++) {
-        simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
+    if (has_hi) {
+        auto tC1 = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc + NR1H*NR0, dextents<int32_t, 2>(NR0, NR1H));
+        cT1.store(tC1);
+    }
+#else
+    if (sg_active) {
+        threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
+
+        for (short i = 0; i < 8; i++) {
+            simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
+        }
     }
 #endif
 

ggml/src/ggml-metal/kernels/mul_mv.metal

@@ -1889,8 +1889,19 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
     const uint i12 = im%FC_mul_mv_ne12;
     const uint i13 = im/FC_mul_mv_ne12;
 
-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
 
     device const block_iq2_xxs * x = (device const block_iq2_xxs *) (src0 + offset0);
     device const float         * y = (device const float         *) (src1 + offset1);
@@ -1898,8 +1909,6 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
     float yl[32];
     float sumf[nr0]={0.f};
 
-    const int nb32 = nb * (QK_K / 32);
-
     threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem);
     threadgroup uint8_t  * ssigns  = (threadgroup uint8_t  *)(svalues + 256);
     {
@@ -1912,11 +1921,9 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
         threadgroup_barrier(mem_flags::mem_threadgroup);
     }
 
-    const int ix = tiisg;
-
     device const float * y4 = y + 32 * ix;
 
-    for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
         for (short i = 0; i < 32; ++i) {
             yl[i] = y4[i];
         }
@@ -1928,7 +1935,7 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
         device const uint16_t * q2 = xr->qs + 4 * ib;
         device const half * dh = &xr->d;
 
-        for (short row = 0; row < nr0; row++) {
+        for (short row = row0; row < row1; row++) {
             const float db = dh[0];
             device const uint8_t * aux8 = (device const uint8_t *)q2;
             const uint32_t aux32 = q2[2] | (q2[3] << 16);
@@ -1948,7 +1955,7 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
             q2 += args.nb01/2;
         }
 
-        y4 += 32 * 32;
+        y4 += 32 * ntx;
     }
 
     device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -1961,6 +1968,23 @@ void kernel_mul_mv_iq2_xxs_f32_impl(
     }
 }
 
+template<typename args_t>
+void kernel_mul_mv_iq2_xxs_f32_disp(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    if (FC_mul_mv_split) {
+        kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    }
+}
+
 [[host_name("kernel_mul_mv_iq2_xxs_f32")]]
 kernel void kernel_mul_mv_iq2_xxs_f32(
         constant ggml_metal_kargs_mul_mv & args,
@@ -1971,7 +1995,7 @@ kernel void kernel_mul_mv_iq2_xxs_f32(
         uint3  tgpig[[threadgroup_position_in_grid]],
         ushort tiisg[[thread_index_in_simdgroup]],
         ushort sgitg[[simdgroup_index_in_threadgroup]]) {
-    kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    kernel_mul_mv_iq2_xxs_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
 }
 
 template<int nr0, typename args_t>
@@ -1997,8 +2021,19 @@ void kernel_mul_mv_iq2_xs_f32_impl(
     const uint i12 = im%FC_mul_mv_ne12;
     const uint i13 = im/FC_mul_mv_ne12;
 
-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
 
     device const block_iq2_xs * x = (device const block_iq2_xs *) (src0 + offset0);
     device const float        * y = (device const float        *) (src1 + offset1);
@@ -2006,8 +2041,6 @@ void kernel_mul_mv_iq2_xs_f32_impl(
     float yl[32];
     float sumf[nr0]={0.f};
 
-    const int nb32 = nb * (QK_K / 32);
-
     threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem);
     threadgroup uint8_t  * ssigns  = (threadgroup uint8_t  *)(svalues + 512);
     {
@@ -2020,11 +2053,9 @@ void kernel_mul_mv_iq2_xs_f32_impl(
         threadgroup_barrier(mem_flags::mem_threadgroup);
     }
 
-    const int ix = tiisg;
-
     device const float * y4 = y + 32 * ix;
 
-    for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
         for (short i = 0; i < 32; ++i) {
             yl[i] = y4[i];
         }
@@ -2037,7 +2068,7 @@ void kernel_mul_mv_iq2_xs_f32_impl(
         device const uint8_t  * sc = xr->scales + ib;
         device const half * dh = &xr->d;
 
-        for (short row = 0; row < nr0; row++) {
+        for (short row = row0; row < row1; row++) {
             const float db = dh[0];
             const uint8_t ls1 = sc[0] & 0xf;
             const uint8_t ls2 = sc[0] >>  4;
@@ -2066,7 +2097,7 @@ void kernel_mul_mv_iq2_xs_f32_impl(
             sc += args.nb01;
         }
 
-        y4 += 32 * 32;
+        y4 += 32 * ntx;
     }
 
     device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2079,6 +2110,23 @@ void kernel_mul_mv_iq2_xs_f32_impl(
     }
 }
 
+template<typename args_t>
+void kernel_mul_mv_iq2_xs_f32_disp(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    if (FC_mul_mv_split) {
+        kernel_mul_mv_iq2_xs_f32_impl<N_R0_IQ2_XS_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_iq2_xs_f32_impl<N_R0_IQ2_XS, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    }
+}
+
 [[host_name("kernel_mul_mv_iq2_xs_f32")]]
 kernel void kernel_mul_mv_iq2_xs_f32(
         constant ggml_metal_kargs_mul_mv & args,
@@ -2090,7 +2138,7 @@ kernel void kernel_mul_mv_iq2_xs_f32(
         ushort tiisg[[thread_index_in_simdgroup]],
         ushort sgitg[[simdgroup_index_in_threadgroup]]) {
 
-    kernel_mul_mv_iq2_xs_f32_impl<N_R0_IQ2_XS, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    kernel_mul_mv_iq2_xs_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
 }
 
 // FC_mul_mv_split: for nb32 < 32 (nb32 divides 32), 32/nb32 threads share each chunk and each takes a slice of the rows
@@ -2117,8 +2165,19 @@ void kernel_mul_mv_iq3_xxs_f32_impl(
     const uint i12 = im%FC_mul_mv_ne12;
     const uint i13 = im/FC_mul_mv_ne12;
 
-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
 
     device const block_iq3_xxs * x = (device const block_iq3_xxs *) (src0 + offset0);
     device const float         * y = (device const float         *) (src1 + offset1);
@@ -2126,8 +2185,6 @@ void kernel_mul_mv_iq3_xxs_f32_impl(
     float yl[32];
     float sumf[nr0]={0.f};
 
-    const int nb32 = nb * (QK_K / 32);
-
     threadgroup uint32_t * svalues = (threadgroup uint32_t *)(shmem);
     threadgroup uint8_t  * ssigns  = (threadgroup uint8_t  *)(svalues + 256);
     {
@@ -2140,15 +2197,6 @@ void kernel_mul_mv_iq3_xxs_f32_impl(
         threadgroup_barrier(mem_flags::mem_threadgroup);
     }
 
-    const short ntx  = FC_mul_mv_split ? nb32 : 32;
-    const short nrep = 32 / ntx;
-
-    const short ix   = tiisg % ntx;
-    const short irep = tiisg / ntx;
-
-    const short row0 = (nr0 * irep      ) / nrep;
-    const short row1 = (nr0 * (irep + 1)) / nrep;
-
     device const float * y4 = y + 32 * ix;
 
     for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
@@ -2160,9 +2208,9 @@ void kernel_mul_mv_iq3_xxs_f32_impl(
         const int ib  = ib32 % (QK_K / 32);
 
         device const block_iq3_xxs * xr = x + ibl;
-        device const uint8_t  * q3 = xr->qs + 8 * ib + (uint64_t) row0*args.nb01;
-        device const uint16_t * gas = (device const uint16_t *)(xr->qs + QK_K/4) + 2 * ib + (uint64_t) row0*args.nb01/2;
-        device const half * dh = &xr->d + (uint64_t) row0*args.nb01/2;
+        device const uint8_t  * q3 = xr->qs + 8 * ib;
+        device const uint16_t * gas = (device const uint16_t *)(xr->qs + QK_K/4) + 2 * ib;
+        device const half * dh = &xr->d;
 
         for (short row = row0; row < row1; row++) {
             const float db = dh[0];
@@ -2253,8 +2301,19 @@ void kernel_mul_mv_iq3_s_f32_impl(
     const uint i12 = im%FC_mul_mv_ne12;
     const uint i13 = im/FC_mul_mv_ne12;
 
-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
 
     device const block_iq3_s * x = (device const block_iq3_s *) (src0 + offset0);
     device const float       * y = (device const float       *) (src1 + offset1);
@@ -2262,8 +2321,6 @@ void kernel_mul_mv_iq3_s_f32_impl(
     float yl[32];
     float sumf[nr0]={0.f};
 
-    const int nb32 = nb * (QK_K / 32);
-
     threadgroup uint32_t * svalues = (threadgroup uint32_t *) shmem;
     {
         int nval = 8;
@@ -2272,11 +2329,9 @@ void kernel_mul_mv_iq3_s_f32_impl(
         threadgroup_barrier(mem_flags::mem_threadgroup);
     }
 
-    const int ix = tiisg;
-
     device const float * y4 = y + 32 * ix;
 
-    for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
         for (short i = 0; i < 32; ++i) {
             yl[i] = y4[i];
         }
@@ -2291,7 +2346,7 @@ void kernel_mul_mv_iq3_s_f32_impl(
         device const uint8_t * signs = xr->signs + 4 * ib;
         device const half * dh = &xr->d;
 
-        for (short row = 0; row < nr0; row++) {
+        for (short row = row0; row < row1; row++) {
             const float db = dh[0];
             const float d = db * (1 + 2*((sc[0] >> 4*(ib%2)) & 0xf));
 
@@ -2315,7 +2370,7 @@ void kernel_mul_mv_iq3_s_f32_impl(
             signs += args.nb01;
         }
 
-        y4 += 32 * 32;
+        y4 += 32 * ntx;
     }
 
     device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2328,6 +2383,23 @@ void kernel_mul_mv_iq3_s_f32_impl(
     }
 }
 
+template<typename args_t>
+void kernel_mul_mv_iq3_s_f32_disp(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    if (FC_mul_mv_split) {
+        kernel_mul_mv_iq3_s_f32_impl<N_R0_IQ3_S_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_iq3_s_f32_impl<N_R0_IQ3_S, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    }
+}
+
 [[host_name("kernel_mul_mv_iq3_s_f32")]]
 kernel void kernel_mul_mv_iq3_s_f32(
         constant ggml_metal_kargs_mul_mv & args,
@@ -2339,7 +2411,7 @@ kernel void kernel_mul_mv_iq3_s_f32(
         ushort tiisg[[thread_index_in_simdgroup]],
         ushort sgitg[[simdgroup_index_in_threadgroup]]) {
 
-    kernel_mul_mv_iq3_s_f32_impl<N_R0_IQ3_S, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    kernel_mul_mv_iq3_s_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
 }
 
 template<int nr0, typename args_t>
@@ -2365,8 +2437,19 @@ void kernel_mul_mv_iq2_s_f32_impl(
     const uint i12 = im%FC_mul_mv_ne12;
     const uint i13 = im/FC_mul_mv_ne12;
 
-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
 
     device const block_iq2_s * x = (device const block_iq2_s *) (src0 + offset0);
     device const float       * y = (device const float       *) (src1 + offset1);
@@ -2374,8 +2457,6 @@ void kernel_mul_mv_iq2_s_f32_impl(
     float yl[32];
     float sumf[nr0]={0.f};
 
-    const int nb32 = nb * (QK_K / 32);
-
     //threadgroup uint64_t * svalues = (threadgroup uint64_t *) shmem;
     //{
     //    int nval = 32;
@@ -2384,11 +2465,9 @@ void kernel_mul_mv_iq2_s_f32_impl(
     //    threadgroup_barrier(mem_flags::mem_threadgroup);
     //}
 
-    const short ix = tiisg;
-
     device const float * y4 = y + 32 * ix;
 
-    for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
         for (short i = 0; i < 32; ++i) {
             yl[i] = y4[i];
         }
@@ -2403,7 +2482,7 @@ void kernel_mul_mv_iq2_s_f32_impl(
         device const uint8_t * signs = qs + QK_K/8;
         device const half * dh = &xr->d;
 
-        for (short row = 0; row < nr0; row++) {
+        for (short row = row0; row < row1; row++) {
             const float db = dh[0];
             const float d1 = db * (0.5f + (sc[0] & 0xf));
             const float d2 = db * (0.5f + (sc[0] >>  4));
@@ -2428,7 +2507,7 @@ void kernel_mul_mv_iq2_s_f32_impl(
             signs += args.nb01;
         }
 
-        y4 += 32 * 32;
+        y4 += 32 * ntx;
     }
 
     device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2441,6 +2520,23 @@ void kernel_mul_mv_iq2_s_f32_impl(
     }
 }
 
+template<typename args_t>
+void kernel_mul_mv_iq2_s_f32_disp(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    if (FC_mul_mv_split) {
+        kernel_mul_mv_iq2_s_f32_impl<N_R0_IQ2_S_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_iq2_s_f32_impl<N_R0_IQ2_S, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    }
+}
+
 [[host_name("kernel_mul_mv_iq2_s_f32")]]
 kernel void kernel_mul_mv_iq2_s_f32(
         constant ggml_metal_kargs_mul_mv & args,
@@ -2452,7 +2548,7 @@ kernel void kernel_mul_mv_iq2_s_f32(
         ushort tiisg[[thread_index_in_simdgroup]],
         ushort sgitg[[simdgroup_index_in_threadgroup]]) {
 
-    kernel_mul_mv_iq2_s_f32_impl<N_R0_IQ2_S, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    kernel_mul_mv_iq2_s_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
 }
 
 template<int nr0, typename args_t>
@@ -2478,8 +2574,19 @@ void kernel_mul_mv_iq1_s_f32_impl(
     const uint i12 = im%FC_mul_mv_ne12;
     const uint i13 = im/FC_mul_mv_ne12;
 
-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
 
     device const block_iq1_s * x = (device const block_iq1_s *) (src0 + offset0);
     device const float       * y = (device const float       *) (src1 + offset1);
@@ -2487,13 +2594,9 @@ void kernel_mul_mv_iq1_s_f32_impl(
     float yl[32];
     float sumf[nr0]={0.f};
 
-    const int nb32 = nb * (QK_K / 32);
-
-    const short ix = tiisg;
-
     device const float * y4 = y + 32 * ix;
 
-    for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
         float sumy = 0;
         for (short i = 0; i < 32; ++i) {
             yl[i] = y4[i];
@@ -2508,7 +2611,7 @@ void kernel_mul_mv_iq1_s_f32_impl(
         device const uint16_t * qh = xr->qh + ib;
         device const half     * dh = &xr->d;
 
-        for (short row = 0; row < nr0; row++) {
+        for (short row = row0; row < row1; row++) {
             constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700)));
             constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((qh[0] << 5) & 0x700)));
             constant uint8_t * grid3 = (constant uint8_t *)(iq1s_grid_gpu + (qs[2] | ((qh[0] << 2) & 0x700)));
@@ -2528,7 +2631,7 @@ void kernel_mul_mv_iq1_s_f32_impl(
             qh += args.nb01/2;
         }
 
-        y4 += 32 * 32;
+        y4 += 32 * ntx;
     }
 
     device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2541,6 +2644,23 @@ void kernel_mul_mv_iq1_s_f32_impl(
     }
 }
 
+template<typename args_t>
+void kernel_mul_mv_iq1_s_f32_disp(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    if (FC_mul_mv_split) {
+        kernel_mul_mv_iq1_s_f32_impl<N_R0_IQ1_S_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_iq1_s_f32_impl<N_R0_IQ1_S, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    }
+}
+
 [[host_name("kernel_mul_mv_iq1_s_f32")]]
 kernel void kernel_mul_mv_iq1_s_f32(
         constant ggml_metal_kargs_mul_mv & args,
@@ -2551,7 +2671,7 @@ kernel void kernel_mul_mv_iq1_s_f32(
         ushort tiisg[[thread_index_in_simdgroup]],
         ushort sgitg[[simdgroup_index_in_threadgroup]]) {
 
-    kernel_mul_mv_iq1_s_f32_impl<N_R0_IQ1_S, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
+    kernel_mul_mv_iq1_s_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
 }
 
 template<int nr0, typename args_t>
@@ -2577,8 +2697,19 @@ void kernel_mul_mv_iq1_m_f32_impl(
     const uint i12 = im%FC_mul_mv_ne12;
     const uint i13 = im/FC_mul_mv_ne12;
 
-    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
-    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
+    const int nb32 = nb * (QK_K / 32);
+
+    const short ntx  = FC_mul_mv_split ? nb32 : 32;
+    const short nrep = 32 / ntx;
+
+    const short ix   = tiisg % ntx;
+    const short irep = tiisg / ntx;
+
+    const short row0 = (nr0 * irep      ) / nrep;
+    const short row1 = (nr0 * (irep + 1)) / nrep;
+
+    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
 
     device const block_iq1_m * x = (device const block_iq1_m *) (src0 + offset0);
     device const float       * y = (device const float       *) (src1 + offset1);
@@ -2586,15 +2717,11 @@ void kernel_mul_mv_iq1_m_f32_impl(
     float yl[32];
     float sumf[nr0]={0.f};
 
-    const int nb32 = nb * (QK_K / 32);
-
-    const short ix = tiisg;
-
     device const float * y4 = y + 32 * ix;
 
     iq1m_scale_t scale;
 
-    for (int ib32 = ix; ib32 < nb32; ib32 += 32) {
+    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
         float4 sumy = {0.f};
         for (short i = 0; i < 8; ++i) {
             yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0];
@@ -2611,7 +2738,7 @@ void kernel_mul_mv_iq1_m_f32_impl(
         device const uint8_t  * qh = xr->qh + 2 * ib;
         device const uint16_t * sc = (device const uint16_t *)xr->scales;
 
-        for (short row = 0; row < nr0; row++) {
+        for (short row = row0; row < row1; row++) {
             scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
 
             constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700)));
@@ -2637,7 +2764,7 @@ void kernel_mul_mv_iq1_m_f32_impl(
             qh += args.nb01;
         }
 
-        y4 += 32 * 32;
+        y4 += 32 * ntx;
     }
 
     device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
@@ -2650,6 +2777,23 @@ void kernel_mul_mv_iq1_m_f32_impl(
     }
 }
 
+template<typename args_t>
+void kernel_mul_mv_iq1_m_f32_disp(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    if (FC_mul_mv_split) {
+        kernel_mul_mv_iq1_m_f32_impl<N_R0_IQ1_M_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    } else {
+        kernel_mul_mv_iq1_m_f32_impl<N_R0_IQ1_M, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
+    }
+}
+
 [[host_name("kernel_mul_mv_iq1_m_f32")]]
 kernel void kernel_mul_mv_iq1_m_f32(
         constant ggml_metal_kargs_mul_mv & args,
@@ -2660,7 +2804,7 @@ kernel void kernel_mul_mv_iq1_m_f32(
         ushort tiisg[[thread_index_in_simdgroup]],
         ushort sgitg[[simdgroup_index_in_threadgroup]]) {
 
-    kernel_mul_mv_iq1_m_f32_impl<N_R0_IQ1_M, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
+    kernel_mul_mv_iq1_m_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
 }
 
 template<int NR0, typename args_t>
@@ -3239,13 +3383,13 @@ template [[host_name("kernel_mul_mv_id_q3_K_f32")]]    kernel kernel_mul_mv_id_t
 template [[host_name("kernel_mul_mv_id_q4_K_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q4_K_f32_impl   <N_R0_Q4_K>>>;
 template [[host_name("kernel_mul_mv_id_q5_K_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q5_K_f32_impl   <N_R0_Q5_K>>>;
 template [[host_name("kernel_mul_mv_id_q6_K_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q6_K_f32_impl   <N_R0_Q6_K>>>;
-template [[host_name("kernel_mul_mv_id_iq1_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_s_f32_impl  <N_R0_IQ1_S>>>;
-template [[host_name("kernel_mul_mv_id_iq1_m_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_m_f32_impl  <N_R0_IQ1_M>>>;
-template [[host_name("kernel_mul_mv_id_iq2_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS>>>;
-template [[host_name("kernel_mul_mv_id_iq2_xs_f32")]]  kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xs_f32_impl <N_R0_IQ2_XS>>>;
+template [[host_name("kernel_mul_mv_id_iq1_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_s_f32_disp<ggml_metal_kargs_mul_mv>>>;
+template [[host_name("kernel_mul_mv_id_iq1_m_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_m_f32_disp<ggml_metal_kargs_mul_mv>>>;
+template [[host_name("kernel_mul_mv_id_iq2_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xxs_f32_disp<ggml_metal_kargs_mul_mv>>>;
+template [[host_name("kernel_mul_mv_id_iq2_xs_f32")]]  kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xs_f32_disp<ggml_metal_kargs_mul_mv>>>;
 template [[host_name("kernel_mul_mv_id_iq3_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq3_xxs_f32_disp<ggml_metal_kargs_mul_mv>>>;
-template [[host_name("kernel_mul_mv_id_iq3_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq3_s_f32_impl  <N_R0_IQ3_S>>>;
-template [[host_name("kernel_mul_mv_id_iq2_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_s_f32_impl  <N_R0_IQ2_S>>>;
+template [[host_name("kernel_mul_mv_id_iq3_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq3_s_f32_disp<ggml_metal_kargs_mul_mv>>>;
+template [[host_name("kernel_mul_mv_id_iq2_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_s_f32_disp<ggml_metal_kargs_mul_mv>>>;
 template [[host_name("kernel_mul_mv_id_iq4_nl_f32")]]  kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq4_nl_f32_impl <N_R0_IQ4_NL>>>;
 template [[host_name("kernel_mul_mv_id_iq4_xs_f32")]]  kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq4_xs_f32_impl <N_R0_IQ4_XS>>>;
 template [[host_name("kernel_mul_mv_id_tq2_0_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_tq2_0_f32_impl  <N_R0_TQ2_0>>>;

ggml/src/ggml-opencl/CMakeLists.txt

@@ -170,6 +170,7 @@ set(GGML_OPENCL_KERNELS
     gemv_noshuffle_q4_0_f32
     gemv_noshuffle_q4_0_f32_spec
     gemm_noshuffle_q4_0_f32
+    gemv_noshuffle_q4_0_f32_32b_trans
     gemv_noshuffle_q4_1_f32
     gemm_noshuffle_q4_1_f32
     gemv_noshuffle_q5_0_f32

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -1160,6 +1160,8 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_gemm_noshuffle_q4_0_f32;
     cl_kernel kernel_gemv_noshuffle_q4_0_f32;
     cl_kernel kernel_gemv_noshuffle_q4_0_f32_mc3;  // multi-column (N=3) verify GEMV (spec/MTP)
+    cl_kernel kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin;
+    cl_kernel kernel_gemv_noshuffle_q4_0_f32_32b_trans;
     cl_kernel kernel_gemv_noshuffle_q4_0_f32_4096_1_11008;
     cl_kernel kernel_gemv_noshuffle_q4_0_f32_4096_1_4096;
     cl_kernel kernel_gemv_noshuffle_q4_0_f32_11008_1_4096;
@@ -3787,6 +3789,43 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
         GGML_LOG_CONT(".");
     }
 
+    backend_ctx->kernel_gemv_noshuffle_q4_0_f32_32b_trans = nullptr;
+    backend_ctx->kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin = nullptr;
+    if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E) {
+        {
+            std::string opts = std::string("-cl-std=") + opencl_c_std +
+                                           " -cl-mad-enable "
+                                           " -DSIMDGROUP_WIDTH=" +
+                                           std::to_string(backend_ctx->adreno_wave_size);
+#ifdef GGML_OPENCL_EMBED_KERNELS
+            const std::string kernel_src {
+                #include "gemv_noshuffle_q4_0_f32_32b_trans.cl.h"
+            };
+#else
+            const std::string kernel_src = read_file("gemv_noshuffle_q4_0_f32_32b_trans.cl");
+#endif
+            cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), opts);
+            CL_CHECK((backend_ctx->kernel_gemv_noshuffle_q4_0_f32_32b_trans =
+                clCreateKernel(prog, "kernel_gemv_noshuffle_q4_0_f32_32b_trans", &err), err));
+            CL_CHECK(clReleaseProgram(prog));
+            GGML_LOG_CONT(".");
+        }
+
+        if (use_adreno_bin_kernels(backend_ctx)) {
+            size_t bin_size = 0;
+            const char * kernel_bin = (const char *)backend_ctx->get_adreno_bin_kernel("gemm_noshuffle_q4_0_f32_32b_trans_ila_a8", &bin_size);
+            if (kernel_bin && bin_size > 0) {
+                cl_program bin_prog =
+                    build_program_from_binary(backend_ctx->context, backend_ctx->device, kernel_bin, "", bin_size);
+
+                CL_CHECK((backend_ctx->kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin =
+                    clCreateKernel(bin_prog, "kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8", &err), err));
+                CL_CHECK(clReleaseProgram(bin_prog));
+                GGML_LOG_CONT(".");
+            }
+        }
+    }
+
     // gemm_noshuffle_q4_1_f32
     {
 #ifdef GGML_OPENCL_EMBED_KERNELS
@@ -6725,11 +6764,10 @@ struct ggml_tensor_extra_cl_q4_0 {
             CL_CHECK(clReleaseMemObject(q_img));
             q_img = nullptr;
         }
-        // Currently, q_img and d_img are only initialized when SMALL_ALLOC is
-        // enabled. They point to the images in ggml_backend_opencl_buffer_context.
-        // So, there is no need to release them here.
-        // TODO: initialize them for non SMALL_PATH path, or remove them.
-        d_img = nullptr;
+        if (d_img != nullptr) {
+            CL_CHECK(clReleaseMemObject(d_img));
+            d_img = nullptr;
+        }
         size_q = 0;
         size_d = 0;
     }
@@ -8311,6 +8349,20 @@ inline bool enable_adreno_trans_weight_q5_K(const ggml_backend_opencl_context *b
            qh_img_width <= backend_ctx->image_max_buffer_size;
 }
 
+inline bool use_q4_0_ila_kernels(const ggml_backend_opencl_context *backend_ctx, const ggml_tensor *tensor) {
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+    if (!backend_ctx->kernel_gemv_noshuffle_q4_0_f32_32b_trans ||
+        !backend_ctx->kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin) {
+        return false;
+    }
+    return (tensor->ne[0] % 32 == 0) && (tensor->ne[1] % 64 == 0);
+#else
+    GGML_UNUSED(backend_ctx);
+    GGML_UNUSED(tensor);
+    return false;
+#endif
+}
+
 // The flat-GEMV large-m escape is OPT-IN (GGML_OPENCL_FLAT_LARGE_M=1) because it
 // is SLOWER than the route it replaces, not because it is unsafe. It was first
 // parked on the theory that it out-of-bounds-writes at vocab-scale shapes; that
@@ -9573,10 +9625,34 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,
 
             GGML_ASSERT(K % 32 == 0);
 
-            // Transpose q as ushort
-            transpose_2d_as_16b(backend_ctx, extra->q, extra->q, size_q, K/4, M);
-            // Transpose d as ushort
-            transpose_2d_as_16b(backend_ctx, extra->d, extra->d, size_d, K/32, M);
+            if (use_q4_0_ila_kernels(backend_ctx, tensor)) {
+                cl_int err;
+                cl_image_format wimg_fmt;
+                cl_image_desc   wimg_desc;
+
+                // transpose quants as 32-bit words (M-first)
+                GGML_ASSERT(M % 64 == 0);
+                transpose_2d_as_32b(backend_ctx, extra->q, extra->q, size_q, K / 8, M);
+                transpose_2d_as_16b(backend_ctx, extra->d, extra->d, size_d, K / 32, M);
+
+                wimg_fmt = { CL_R, CL_UNSIGNED_INT32 };
+                memset(&wimg_desc, 0, sizeof(wimg_desc));
+                wimg_desc.image_type  = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+                wimg_desc.image_width = (size_t)M * K / 8;
+                wimg_desc.buffer      = extra->q;
+                CL_CHECK((extra->q_img = clCreateImage(context, CL_MEM_READ_ONLY, &wimg_fmt, &wimg_desc, NULL, &err), err));
+
+                wimg_fmt = { CL_R, CL_HALF_FLOAT };
+                memset(&wimg_desc, 0, sizeof(wimg_desc));
+                wimg_desc.image_type  = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+                wimg_desc.image_width = (size_t)M * K / 32;
+                wimg_desc.buffer      = extra->d;
+                CL_CHECK((extra->d_img = clCreateImage(context, CL_MEM_READ_ONLY, &wimg_fmt, &wimg_desc, NULL, &err), err));
+            } else {
+                // Transpose q and d as ushort
+                transpose_2d_as_16b(backend_ctx, extra->q, extra->q, size_q, K/4, M);
+                transpose_2d_as_16b(backend_ctx, extra->d, extra->d, size_d, K/32, M);
+            }
         }
 #endif // GGML_OPENCL_USE_ADRENO_KERNELS
         return;
@@ -11104,7 +11180,11 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
             buf_trans_d.allocate(backend_ctx->context, size_d);
             buf_unpacked.allocate(backend_ctx->context, ggml_nbytes(tensor));
 
-            transpose_2d_as_16b(backend_ctx, extra->q, buf_trans_q.buffer, size_q, M, K/4);
+            if (use_q4_0_ila_kernels(backend_ctx, tensor)) {
+                transpose_2d_as_32b(backend_ctx, extra->q, buf_trans_q.buffer, size_q, M, K / 8);
+            } else {
+                transpose_2d_as_16b(backend_ctx, extra->q, buf_trans_q.buffer, size_q, M, K / 4);
+            }
             transpose_2d_as_16b(backend_ctx, extra->d, buf_trans_d.buffer, size_d, M, K/32);
 
             cl_uchar mask_0F = 0x0F;
@@ -18347,6 +18427,166 @@ static void ggml_cl_mul_mat_q1_0_f32_adreno(ggml_backend_t backend, const ggml_t
 #endif
 }
 
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+static void ggml_cl_mul_mat_q4_0_f32_adreno_ila(ggml_backend_t backend, const ggml_tensor * src0,
+                                                const ggml_tensor * src1, ggml_tensor * dst) {
+    GGML_ASSERT(src0);
+    GGML_ASSERT(src0->extra);
+    GGML_ASSERT(src1);
+    GGML_ASSERT(src1->extra);
+    GGML_ASSERT(dst);
+    GGML_ASSERT(dst->extra);
+
+    ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context;
+
+    ggml_tensor_extra_cl * extra1 = (ggml_tensor_extra_cl *)src1->extra;
+    ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra;
+    ggml_tensor_extra_cl_q4_0 * extra0_q4_0 = (ggml_tensor_extra_cl_q4_0 *)src0->extra;
+
+    cl_ulong offset1 = extra1->offset + src1->view_offs;
+    cl_ulong offsetd = extrad->offset + dst->view_offs;
+
+    const int ne00 = src0->ne[0];
+    const int ne01 = src0->ne[1];
+
+    const int ne1 = dst->ne[1];
+
+    GGML_ASSERT(ne00 % ggml_blck_size(src0->type) == 0);
+
+    cl_context context = backend_ctx->context;
+    cl_kernel kernel;
+
+    cl_int              err;
+    cl_image_format     img_fmt;
+    cl_image_desc       img_desc;
+    cl_buffer_region    region;
+
+    int M = ne01;
+    int N = ne1;
+    int K = ne00;
+
+    if (ne1 == 1) {
+        cl_mem b_sub_buf = nullptr;
+        cl_mem b_img     = nullptr;
+
+        region.origin = offset1;
+        region.size   = (size_t)K * N * sizeof(float);
+        CL_CHECK((b_sub_buf = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, &region, &err), err));
+
+        img_fmt = { CL_RGBA, CL_FLOAT };
+        memset(&img_desc, 0, sizeof(img_desc));
+        img_desc.image_type  = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+        img_desc.image_width = (size_t)K * N / 4;
+        img_desc.buffer      = b_sub_buf;
+        CL_CHECK((b_img = clCreateImage(context, CL_MEM_READ_ONLY, &img_fmt, &img_desc, NULL, &err), err));
+
+        kernel = backend_ctx->kernel_gemv_noshuffle_q4_0_f32_32b_trans;
+        CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem),   &extra0_q4_0->q_img));
+        CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem),   &extra0_q4_0->d));
+        CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem),   &b_img));
+        CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_mem),   &extrad->data_device));
+        CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_ulong), &offsetd));
+        CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_int),   &K));
+        CL_CHECK(clSetKernelArg(kernel, 6, sizeof(cl_int),   &M));
+
+        size_t wavesize = backend_ctx->adreno_wave_size;
+        size_t local_work_size[3]  = { wavesize, 4, 1 };
+        size_t global_work_size[3] = { (size_t)CEIL_DIV(M, 64) * 64, 4, 1 };
+        backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+
+        CL_CHECK(clReleaseMemObject(b_sub_buf));
+        CL_CHECK(clReleaseMemObject(b_img));
+    } else {
+        const int gemm_tile_n = 64;
+        int N_pad = (N + gemm_tile_n - 1) & ~(gemm_tile_n - 1);
+
+        cl_mem a_img = extra0_q4_0->q_img;
+        cl_mem s_img = extra0_q4_0->d_img;
+        GGML_ASSERT(a_img && s_img && "ILA Q4_0 weight images missing; set_tensor should have built them");
+
+        // Pad B through a zero-filled scratch buffer when N needs
+        // padding, since the GEMM kernel always reads a full N-tile.
+        const bool need_pad = N_pad > N;
+        cl_mem b_sub_buf = nullptr;
+        cl_mem b_padded  = nullptr;
+        if (need_pad) {
+            CL_CHECK((b_padded = clCreateBuffer(context, CL_MEM_READ_WRITE,
+                (size_t)K * N_pad * sizeof(float), NULL, &err), err));
+            const float zero = 0.0f;
+            CL_CHECK(clEnqueueFillBuffer(backend_ctx->queue, b_padded, &zero, sizeof(zero),
+                0, (size_t)K * N_pad * sizeof(float), 0, NULL, NULL));
+            CL_CHECK(clEnqueueCopyBuffer(backend_ctx->queue, extra1->data_device, b_padded,
+                offset1, 0, (size_t)K * N * sizeof(float), 0, NULL, NULL));
+        } else {
+            region.origin = offset1;
+            region.size   = (size_t)K * N * sizeof(float);
+            CL_CHECK((b_sub_buf = clCreateSubBuffer(extra1->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, &region, &err), err));
+        }
+
+        img_fmt = { CL_R, CL_FLOAT };
+        memset(&img_desc, 0, sizeof(img_desc));
+        img_desc.image_type  = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+        img_desc.image_width = need_pad ? (size_t)K * N_pad : (size_t)K * N;
+        img_desc.buffer      = need_pad ? b_padded : b_sub_buf;
+        cl_mem b_img;
+        CL_CHECK((b_img = clCreateImage(context, CL_MEM_READ_ONLY, &img_fmt, &img_desc, NULL, &err), err));
+
+        region.origin = offsetd;
+        region.size   = (size_t)M * N * sizeof(float);
+        cl_mem d_sub_buf;
+        CL_CHECK((d_sub_buf = clCreateSubBuffer(extrad->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, &region, &err), err));
+
+        img_fmt = { CL_R, CL_FLOAT };
+        memset(&img_desc, 0, sizeof(img_desc));
+        img_desc.image_type  = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+        img_desc.image_width = (size_t)M * N;
+        img_desc.buffer      = d_sub_buf;
+        cl_mem d_img;
+        CL_CHECK((d_img = clCreateImage(context, CL_MEM_WRITE_ONLY, &img_fmt, &img_desc, NULL, &err), err));
+
+        int line_stride_matrix_A_in_bytes = M * 4;
+        int line_stride_matrix_S_in_bytes = M * 2;
+        int line_stride_matrix_B_in_bytes = K * 4;
+        int line_stride_matrix_C_in_bytes = M * 4;
+
+        int c_offset_for_kernel = 0;
+        int b_offset_for_kernel = 0;
+
+        kernel = backend_ctx->kernel_gemm_noshuffle_q4_0_f32_32b_trans_ila_a8_bin;
+
+        cl_uint k_arg = 0;
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &a_img));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &s_img));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &b_img));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int),    &b_offset_for_kernel));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &d_img));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &c_offset_for_kernel));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &K));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &line_stride_matrix_A_in_bytes));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &line_stride_matrix_S_in_bytes));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &line_stride_matrix_B_in_bytes));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &line_stride_matrix_C_in_bytes));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &M));
+        CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &N));
+
+        size_t local_work_size[3]  = { 64, 2, 2 };
+        size_t m_tiles = (size_t)CEIL_DIV(M, 64);
+        size_t global_work_size[3] = { 64, m_tiles, (size_t)CEIL_DIV(N_pad, gemm_tile_n) };
+        backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+
+        CL_CHECK(clReleaseMemObject(b_img));
+        if (b_sub_buf) {
+            CL_CHECK(clReleaseMemObject(b_sub_buf));
+        }
+        if (b_padded) {
+            CL_CHECK(clReleaseMemObject(b_padded));
+        }
+        CL_CHECK(clReleaseMemObject(d_img));
+        CL_CHECK(clReleaseMemObject(d_sub_buf));
+    }
+}
+#endif // GGML_OPENCL_USE_ADRENO_KERNELS
+
 static void ggml_cl_mul_mat_q4_0_f32_adreno(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
     GGML_ASSERT(src0);
@@ -18399,6 +18639,20 @@ static void ggml_cl_mul_mat_q4_0_f32_adreno(ggml_backend_t backend, const ggml_t
     static const bool q40_mc3 = (getenv("GGML_OPENCL_Q40_MC3") != nullptr);
     const bool use_q40_mc3 = q40_mc3 && (ne1 >= 2 && ne1 <= 4) && (ne01 < 32768);
 
+    const bool use_ila = use_q4_0_ila_kernels(backend_ctx, src0);
+
+    if (use_ila) {
+        if (use_q40_mc3) {
+            static bool warned = false;
+            if (!warned) {
+                GGML_LOG_WARN("ggml_opencl: GGML_OPENCL_Q40_MC3 is bypassed by Q4_0 binary kernels\n");
+                warned = true;
+            }
+        }
+        ggml_cl_mul_mat_q4_0_f32_adreno_ila(backend, src0, src1, dst);
+        return;
+    }
+
     if (ne1 == 1 || use_q40_mc3) {
         cl_mem q_img = nullptr;
         cl_mem b_sub_buf = nullptr;

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_0_f32_32b_trans.cl

@@ -0,0 +1,137 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+
+#ifdef 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")))
+#endif
+
+#define QK4_0 32
+#define N_SIMDGROUP 4
+
+#define dequantizeBlockAccum_ila_1row_hi(total_sum, bits4, scale, y) \
+    float shared_y; \
+    shared_y = sub_group_broadcast(y.s0, 0); \
+    total_sum += ((bits4.s0 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s1, 0); \
+    total_sum += (((bits4.s0 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s2, 0); \
+    total_sum += (((bits4.s0 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s3, 0); \
+    total_sum += (((bits4.s0 & 0xF000) >> 12) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s4, 0); \
+    total_sum += ((bits4.s1 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s5, 0); \
+    total_sum += (((bits4.s1 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s6, 0); \
+    total_sum += (((bits4.s1 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s7, 0); \
+    total_sum += (((bits4.s1 & 0xF000) >> 12) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s0, 1); \
+    total_sum += ((bits4.s2 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s1, 1); \
+    total_sum += (((bits4.s2 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s2, 1); \
+    total_sum += (((bits4.s2 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s3, 1); \
+    total_sum += (((bits4.s2 & 0xF000) >> 12) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s4, 1); \
+    total_sum += ((bits4.s3 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s5, 1); \
+    total_sum += (((bits4.s3 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s6, 1); \
+    total_sum += (((bits4.s3 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s7, 1); \
+    total_sum += (((bits4.s3 & 0xF000) >> 12) - 8) * scale * shared_y;
+
+#define dequantizeBlockAccum_ila_1row_lo(total_sum, bits4, scale, y) \
+    shared_y = sub_group_broadcast(y.s0, 2); \
+    total_sum += ((bits4.s4 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s1, 2); \
+    total_sum += (((bits4.s4 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s2, 2); \
+    total_sum += (((bits4.s4 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s3, 2); \
+    total_sum += (((bits4.s4 & 0xF000) >> 12) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s4, 2); \
+    total_sum += ((bits4.s5 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s5, 2); \
+    total_sum += (((bits4.s5 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s6, 2); \
+    total_sum += (((bits4.s5 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s7, 2); \
+    total_sum += (((bits4.s5 & 0xF000) >> 12) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s0, 3); \
+    total_sum += ((bits4.s6 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s1, 3); \
+    total_sum += (((bits4.s6 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s2, 3); \
+    total_sum += (((bits4.s6 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s3, 3); \
+    total_sum += (((bits4.s6 & 0xF000) >> 12) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s4, 3); \
+    total_sum += ((bits4.s7 & 0x000F) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s5, 3); \
+    total_sum += (((bits4.s7 & 0x00F0) >> 4) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s6, 3); \
+    total_sum += (((bits4.s7 & 0x0F00) >> 8) - 8) * scale * shared_y; \
+    shared_y = sub_group_broadcast(y.s7, 3); \
+    total_sum += (((bits4.s7 & 0xF000) >> 12) - 8) * scale * shared_y;
+
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+__kernel void kernel_gemv_noshuffle_q4_0_f32_32b_trans(
+        __read_only  image1d_buffer_t src0_q,
+        global half  * src0_d,
+        __read_only  image1d_buffer_t src1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01)
+{
+    uint groupId = get_local_id(1);
+    uint gid     = get_global_id(0);
+    ushort slid  = get_sub_group_local_id();
+
+    uint K = ne00;
+    uint M = ne01;
+
+    __private uint4  regA;
+    __private half   regS;
+    __private float8 regB;
+    __private float  totalSum = 0.0f;
+
+    for (uint k = groupId; k < (K / QK4_0); k += N_SIMDGROUP) {
+        regS = src0_d[k * M + gid];
+        if (slid < 4) {
+            regB.s0123 = read_imagef(src1, (slid * 2 + k * 8));
+            regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8));
+        }
+        regA.s0 = read_imageui(src0_q, ((k * 4 + 0) * M + gid)).x;
+        regA.s1 = read_imageui(src0_q, ((k * 4 + 1) * M + gid)).x;
+        regA.s2 = read_imageui(src0_q, ((k * 4 + 2) * M + gid)).x;
+        regA.s3 = read_imageui(src0_q, ((k * 4 + 3) * M + gid)).x;
+
+        dequantizeBlockAccum_ila_1row_hi(totalSum, as_ushort8(regA), regS, regB);
+        dequantizeBlockAccum_ila_1row_lo(totalSum, as_ushort8(regA), regS, regB);
+    }
+
+    __local float reduceLM[SIMDGROUP_WIDTH * 3];
+    if (groupId == 1) reduceLM[SIMDGROUP_WIDTH * 0 + slid] = totalSum;
+    if (groupId == 2) reduceLM[SIMDGROUP_WIDTH * 1 + slid] = totalSum;
+    if (groupId == 3) reduceLM[SIMDGROUP_WIDTH * 2 + slid] = totalSum;
+    barrier(CLK_LOCAL_MEM_FENCE);
+    if (groupId == 0) totalSum += reduceLM[SIMDGROUP_WIDTH * 0 + slid];
+    if (groupId == 0) totalSum += reduceLM[SIMDGROUP_WIDTH * 1 + slid];
+    if (groupId == 0) totalSum += reduceLM[SIMDGROUP_WIDTH * 2 + slid];
+
+    if (groupId == 0) {
+        dst = (global float*)((global char*)dst + offsetd);
+        if (gid < M) {
+            dst[gid] = totalSum;
+        }
+    }
+}

ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp

@@ -33,7 +33,11 @@ void argsort(bool needs_bounds_check, const uint row) {
     const uint row_offset = row * p.ncols;
 
     // initialize indices
-    dst_row[col] = ivec2(col, floatBitsToInt(data_a[row_offset + col]));
+    ivec2 value = ivec2(col, 0);
+    if (!needs_bounds_check || col < p.ncols) {
+        value.y = floatBitsToInt(data_a[row_offset + col]);
+    }
+    dst_row[col] = value;
     barrier();
 
     uint num_outer_loop_iters = NCOLS_PADDED_LOG2;
@@ -42,18 +46,20 @@ void argsort(bool needs_bounds_check, const uint row) {
         [[unroll]] for (uint j = k / 2, inner_idx = 0; inner_idx < num_inner_loop_iters; j /= 2, inner_idx++) {
             const int ixj = int(col ^ j);
 
-            int idx_0 = (col & k) == 0 ? col : ixj;
-            int idx_1 = (col & k) == 0 ? ixj : col;
+            if (ixj > col) {
+                int idx_0 = (col & k) == 0 ? col : ixj;
+                int idx_1 = (col & k) == 0 ? ixj : col;
 
-            ivec2 sh_idx_0 = dst_row[idx_0];
-            ivec2 sh_idx_1 = dst_row[idx_1];
-            bool idx_0_oob = needs_bounds_check ? sh_idx_0.x >= p.ncols : false;
-            bool idx_1_oob = needs_bounds_check ? sh_idx_1.x >= p.ncols : false;
+                ivec2 sh_idx_0 = dst_row[idx_0];
+                ivec2 sh_idx_1 = dst_row[idx_1];
+                bool idx_0_oob = needs_bounds_check ? sh_idx_0.x >= p.ncols : false;
+                bool idx_1_oob = needs_bounds_check ? sh_idx_1.x >= p.ncols : false;
 
-            if ((idx_0_oob ||
-                (!idx_1_oob && intBitsToFloat(sh_idx_0.y) > intBitsToFloat(sh_idx_1.y))) && (ixj > col)) {
-                dst_row[idx_0] = sh_idx_1;
-                dst_row[idx_1] = sh_idx_0;
+                if (idx_0_oob ||
+                    (!idx_1_oob && intBitsToFloat(sh_idx_0.y) > intBitsToFloat(sh_idx_1.y))) {
+                    dst_row[idx_0] = sh_idx_1;
+                    dst_row[idx_1] = sh_idx_0;
+                }
             }
 
             barrier();

ggml/src/ggml-vulkan/vulkan-shaders/argsort_large.comp

@@ -42,7 +42,10 @@ void argsort(bool needs_bounds_check, const uint row) {
         [[unroll]] for (int u = 0; u < WG_UNROLL_FACTOR; ++u) {
             uint c = u*BLOCK_SIZE + col;
             if (c < p.ncols_padded) {
-                ivec2 v = ivec2(c, floatBitsToInt(data_a[row_offset + c]));
+                ivec2 v = ivec2(c, 0);
+                if (!needs_bounds_check || c < p.ncols) {
+                    v.y = floatBitsToInt(data_a[row_offset + c]);
+                }
                 tmp_idx[idx_offset + c] = v;
             }
         }

scripts/build-profile.ps1

@@ -0,0 +1,136 @@
+# Compile-time profiling using clang -ftime-trace + ClangBuildAnalyzer.
+#
+# Usage:
+#   .\scripts\build-profile.ps1 [-Full] [-Jobs N]
+#
+# -Full : include Server, Tools, and Tests (default: minimal build)
+# -Jobs :  number of parallel jobs (default: all cores)
+#
+# Requires ClangBuildAnalyzer:
+#   https://github.com/aras-p/ClangBuildAnalyzer
+
+param(
+    [switch]$Full,
+    [int]$Jobs = [Environment]::ProcessorCount
+)
+
+$ErrorActionPreference = "Stop"
+
+$ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path
+$RootDir = Split-Path -Parent $ScriptDir
+
+if ($Full) {
+    $BuildDir = Join-Path $RootDir "build-profile-full"
+    $Report = Join-Path $BuildDir "profile-report-full.txt"
+} else {
+    $BuildDir = Join-Path $RootDir "build-profile-baseline"
+    $Report = Join-Path $BuildDir "profile-report.txt"
+}
+
+$OutputBin = Join-Path $BuildDir "clang_analysis.bin"
+
+if (-not (Get-Command clang++ -ErrorAction SilentlyContinue)) {
+    Write-Error "clang++ not found"
+    exit 1
+}
+
+if (-not (Get-Command ninja -ErrorAction SilentlyContinue)) {
+    Write-Error "ninja not found (required so cmake does not fall back to the Visual Studio/MSVC generator)"
+    exit 1
+}
+
+if (-not (Get-Command ClangBuildAnalyzer -ErrorAction SilentlyContinue)) {
+    Write-Error "ClangBuildAnalyzer not found`n  https://github.com/aras-p/ClangBuildAnalyzer/releases"
+    exit 1
+}
+
+$ClangVer = (clang++ --version | Select-Object -First 1)
+Write-Host "compiler : $ClangVer"
+Write-Host "build dir: $BuildDir"
+Write-Host "output   : $OutputBin"
+Write-Host "jobs     : $Jobs"
+Write-Host ""
+
+if (Get-Command ccache -ErrorAction SilentlyContinue) {
+    Write-Host "clearing ccache..."
+    ccache -C -z
+}
+
+$env:CCACHE_DISABLE = "1"
+
+$TestsFlag = if ($Full) { "ON" } else { "OFF" }
+$ToolsFlag = if ($Full) { "ON" } else { "OFF" }
+$ServerFlag = if ($Full) { "ON" } else { "OFF" }
+
+cmake --fresh `
+    -S $RootDir `
+    -B $BuildDir `
+    -G "Ninja" `
+    -DCMAKE_BUILD_TYPE=Release `
+    -DCMAKE_C_COMPILER=clang `
+    -DCMAKE_CXX_COMPILER=clang++ `
+    -DCMAKE_C_FLAGS="-ftime-trace" `
+    -DCMAKE_CXX_FLAGS="-ftime-trace" `
+    -DGGML_CCACHE=OFF `
+    -DGGML_OPENMP=ON `
+    -DGGML_NATIVE=OFF `
+    "-DLLAMA_BUILD_TESTS=$TestsFlag" `
+    -DLLAMA_BUILD_EXAMPLES=OFF `
+    "-DLLAMA_BUILD_TOOLS=$ToolsFlag" `
+    "-DLLAMA_BUILD_SERVER=$ServerFlag" `
+    -DLLAMA_BUILD_APP=OFF
+
+if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE }
+
+$StrayTrace = Join-Path $RootDir "-.json"
+if (Test-Path $StrayTrace) {
+    Remove-Item $StrayTrace -Force
+}
+
+Write-Host ""
+Write-Host "Initializing ClangBuildAnalyzer..."
+ClangBuildAnalyzer --start $BuildDir
+Write-Host ""
+
+Write-Host "building..."
+Write-Host ""
+
+$StartTime = Get-Date
+
+cmake --build $BuildDir --clean-first -j $Jobs
+
+if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE }
+
+$Elapsed = (Get-Date) - $StartTime
+
+Write-Host ""
+Write-Host ("build time: {0}s ({1}m {2}s)" -f [int]$Elapsed.TotalSeconds, [int]$Elapsed.TotalMinutes, $Elapsed.Seconds)
+Write-Host ""
+
+Write-Host "Aggregating profile metrics..."
+ClangBuildAnalyzer --stop $BuildDir $OutputBin | Out-Null
+
+Write-Host ""
+Write-Host ("=" * 80)
+
+$TUs = "?"
+if (Test-Path $Report) {
+    $Match = Select-String -Path $Report -Pattern "Compilation \((\d+)" | Select-Object -First 1
+    if ($Match) { $TUs = $Match.Matches[0].Groups[1].Value }
+}
+
+ClangBuildAnalyzer --analyze $OutputBin | Tee-Object -FilePath $Report
+
+Write-Host ""
+Write-Host "translation units: $TUs"
+Write-Host ""
+Write-Host "largest trace files (top 20 by size):"
+
+Get-ChildItem -Path $BuildDir -Recurse -Filter "*.json" |
+    Where-Object { $_.Name -ne "compile_commands.json" } |
+    Sort-Object Length -Descending |
+    Select-Object -First 20 |
+    ForEach-Object { "{0,8:F1} KB  {1}" -f ($_.Length / 1024), $_.FullName }
+
+Write-Host ""
+Write-Host "ClangBuildAnalyzer report was generated: $Report"

scripts/build-profile.sh

@@ -0,0 +1,122 @@
+#!/usr/bin/env bash
+# Compile-time profiling using clang -ftime-trace + ClangBuildAnalyzer.
+#
+# Usage:
+#   ./scripts/build-profile.sh [--full] [-jN]
+#
+# --full: include Server, Tools, and Tests (default: minimal build)
+# -jN   :  number of parallel jobs (default: all cores)
+#
+# Requires ClangBuildAnalyzer:
+#   macOS: brew install clang-build-analyzer
+#   Linux: https://github.com/aras-p/ClangBuildAnalyzer.git
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+ROOT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
+
+FULL=0
+JOBS="-j$(nproc 2>/dev/null || sysctl -n hw.ncpu)"
+
+for arg in "$@"; do
+    case "${arg}" in
+        --full) FULL=1 ;;
+        -j*)    JOBS="${arg}" ;;
+        *) echo "error: unknown argument: ${arg}" >&2; exit 1 ;;
+    esac
+done
+
+if [ "${FULL}" -eq 1 ]; then
+    BUILD_DIR="${ROOT_DIR}/build-profile-full"
+    REPORT="${BUILD_DIR}/profile-report-full.txt"
+else
+    BUILD_DIR="${ROOT_DIR}/build-profile-baseline"
+    REPORT="${BUILD_DIR}/profile-report.txt"
+fi
+
+OUTPUT_BIN="${BUILD_DIR}/clang_analysis.bin"
+
+if ! command -v clang++ &>/dev/null; then
+    echo "error: clang++ not found" >&2
+    exit 1
+fi
+
+if ! command -v ClangBuildAnalyzer &>/dev/null; then
+    echo "error: ClangBuildAnalyzer not found" >&2
+    echo "  brew install clangbuildanalyzer  (macOS)" >&2
+    echo "  or: https://github.com/aras-p/ClangBuildAnalyzer/releases" >&2
+    exit 1
+fi
+
+CLANG_VER=$(clang++ --version | head -1)
+echo "compiler : ${CLANG_VER}"
+echo "build dir: ${BUILD_DIR}"
+echo "output   : ${OUTPUT_BIN}"
+echo "jobs     : ${JOBS}"
+echo
+
+if command -v ccache &>/dev/null; then
+    echo "clearing ccache..."
+    ccache -C -z
+fi
+
+export CCACHE_DISABLE=1
+
+cmake --fresh \
+    -S "${ROOT_DIR}" \
+    -B "${BUILD_DIR}" \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCMAKE_C_COMPILER=clang \
+    -DCMAKE_CXX_COMPILER=clang++ \
+    -DCMAKE_C_FLAGS="-ftime-trace" \
+    -DCMAKE_CXX_FLAGS="-ftime-trace" \
+    -DGGML_CCACHE=OFF \
+    -DGGML_OPENMP=ON \
+    -DGGML_NATIVE=OFF \
+    -DLLAMA_BUILD_TESTS=$([ "${FULL}" -eq 1 ] && echo ON || echo OFF) \
+    -DLLAMA_BUILD_EXAMPLES=OFF \
+    -DLLAMA_BUILD_TOOLS=$([ "${FULL}" -eq 1 ] && echo ON || echo OFF) \
+    -DLLAMA_BUILD_SERVER=$([ "${FULL}" -eq 1 ] && echo ON || echo OFF) \
+    -DLLAMA_BUILD_APP=OFF
+
+echo
+
+echo "Initializing ClangBuildAnalyzer..."
+ClangBuildAnalyzer --start "${BUILD_DIR}"
+echo
+
+echo "building..."
+echo
+
+START=$(date +%s)
+
+cmake --build "${BUILD_DIR}" --clean-first "${JOBS}"
+
+END=$(date +%s)
+ELAPSED=$((END - START))
+
+echo
+printf "build time: %ds (%dm %ds)\n" "${ELAPSED}" "$((ELAPSED / 60))" "$((ELAPSED % 60))"
+echo
+
+echo "Aggregating profile metrics..."
+ClangBuildAnalyzer --stop "${BUILD_DIR}" "${OUTPUT_BIN}" > /dev/null
+
+echo
+echo "================================================================================"
+TUS=$(grep -oP "Compilation \(\K[0-9]+" "${REPORT}" 2>/dev/null || echo "?")
+ClangBuildAnalyzer --analyze "${OUTPUT_BIN}" | tee "${REPORT}"
+
+echo
+echo "translation units: ${TUS}"
+echo
+echo "largest trace files (top 20 by size):"
+find "${BUILD_DIR}" -name "*.json" ! -name "compile_commands.json" \
+    | xargs ls -l 2>/dev/null \
+    | awk 'NF>5 {print $5, $NF}' \
+    | sort -rn \
+    | awk 'NR<=20 {printf "%8.1f KB  %s\n", $1/1024, $2}'
+
+echo
+echo "ClangBuildAnalyzer report was generated: ${REPORT}"

src/CMakeLists.txt

@@ -8,40 +8,44 @@ llama_add_compile_flags()
 
 file(GLOB LLAMA_MODELS_SOURCES "models/*.cpp")
 
+set(LLAMA_CORE_SOURCES
+    llama.cpp
+    llama-adapter.cpp
+    llama-arch.cpp
+    llama-batch.cpp
+    llama-chat.cpp
+    llama-context.cpp
+    llama-cparams.cpp
+    llama-grammar.cpp
+    llama-graph.cpp
+    llama-hparams.cpp
+    llama-impl.cpp
+    llama-io.cpp
+    llama-kv-cache.cpp
+    llama-kv-cache-iswa.cpp
+    llama-kv-cache-dsa.cpp
+    llama-kv-cache-dsa-iswa.cpp
+    llama-kv-cache-msa.cpp
+    llama-kv-cache-dsv4.cpp
+    llama-memory.cpp
+    llama-memory-hybrid.cpp
+    llama-memory-hybrid-iswa.cpp
+    llama-memory-hybrid-idx.cpp
+    llama-memory-recurrent.cpp
+    llama-mmap.cpp
+    llama-model-loader.cpp
+    llama-model-saver.cpp
+    llama-model.cpp
+    llama-quant.cpp
+    llama-sampler.cpp
+    llama-vocab.cpp
+    unicode-data.cpp
+    unicode.cpp
+)
+
 add_library(llama
             ../include/llama.h
-            llama.cpp
-            llama-adapter.cpp
-            llama-arch.cpp
-            llama-batch.cpp
-            llama-chat.cpp
-            llama-context.cpp
-            llama-cparams.cpp
-            llama-grammar.cpp
-            llama-graph.cpp
-            llama-hparams.cpp
-            llama-impl.cpp
-            llama-io.cpp
-            llama-kv-cache.cpp
-            llama-kv-cache-iswa.cpp
-            llama-kv-cache-dsa.cpp
-            llama-kv-cache-dsa-iswa.cpp
-            llama-kv-cache-msa.cpp
-            llama-kv-cache-dsv4.cpp
-            llama-memory.cpp
-            llama-memory-hybrid.cpp
-            llama-memory-hybrid-iswa.cpp
-            llama-memory-hybrid-idx.cpp
-            llama-memory-recurrent.cpp
-            llama-mmap.cpp
-            llama-model-loader.cpp
-            llama-model-saver.cpp
-            llama-model.cpp
-            llama-quant.cpp
-            llama-sampler.cpp
-            llama-vocab.cpp
-            unicode-data.cpp
-            unicode.cpp
+            ${LLAMA_CORE_SOURCES}
             unicode.h
             ${LLAMA_MODELS_SOURCES}
             )
@@ -50,13 +54,20 @@ set_target_properties(llama PROPERTIES
     VERSION ${LLAMA_VERSION_BASE}
     SOVERSION ${LLAMA_VERSION_MAJOR}
     MACHO_CURRENT_VERSION 0 # keep macOS linker from seeing oversized version number
+    UNITY_BUILD ON
+    UNITY_BUILD_BATCH_SIZE 16
 )
 
+# exclude non-model sources from unity build
+set_source_files_properties(${LLAMA_CORE_SOURCES} ../include/llama.h unicode.h
+    PROPERTIES SKIP_UNITY_BUILD_INCLUSION ON)
+
 configure_file(llama-version.h.in ${CMAKE_CURRENT_BINARY_DIR}/llama-version.h @ONLY)
 
 target_include_directories(llama PRIVATE . ${CMAKE_CURRENT_BINARY_DIR})
 target_include_directories(llama PUBLIC ../include)
 target_compile_features   (llama PRIVATE cxx_std_17) # don't bump
+target_precompile_headers (llama PRIVATE models/models.h)
 
 target_link_libraries(llama PUBLIC ggml)
 

src/llama-context.cpp

@@ -666,7 +666,9 @@ void llama_context::sched_reserve() {
         //       need to implement a more robust mechanism that tries a few different inputs and analyzes the results
         ggml_cgraph * gf = nullptr;
         switch (model.arch) {
+            case LLM_ARCH_KIMI_LINEAR:
             case LLM_ARCH_MINIMAX_01:
+                // [TAG_RESERVE_DIAG_DECAY]
                 // the `inp_diag_decay` tensor size scales with `n_seq_tokens^2` which
                 // makes `n_seqs == 1` use more memory for the compute graph compared to `n_seqs > 1`
                 gf = graph_reserve(n_tokens, 1,      n_outputs_pp, mctx.get(), model.hparams.no_alloc);

src/llama-model.cpp

@@ -2644,9 +2644,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                         filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
                     }
 
-                    if ((arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_HY_V3 || arch == LLM_ARCH_GLM_DSA ||
-                            arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_DEEPSEEK32) &&
-                            hparams.n_layer_nextn > 0) {
+                    // don't filter when n_layer_nextn is repurposed for a router layer the trunk attends
+                    // or when a model is entirely n_layer_nextn layers and has no trunk
+                    if (hparams.n_layer_nextn > 0 && hparams.n_layer() > 0 && hparams.router_layer < 0) {
                         if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
                             filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
                         } else {

src/models/gemma3n.cpp

@@ -82,7 +82,7 @@ std::unique_ptr<llm_graph_context> llama_model_gemma3n::build_arch_graph(const l
 }
 
 // get 2D slice view from a 3D tensor, the idx corresponds to the 3rd dim
-static ggml_tensor * ggml_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, int idx) {
+static ggml_tensor * gemma3n_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, int idx) {
     GGML_ASSERT(idx < (int) x->ne[2]);
     return ggml_view_2d(ctx0, x, x->ne[0], x->ne[1], ggml_row_size(x->type, x->ne[0]),
                         idx * x->ne[0] * x->ne[1] * ggml_element_size(x));
@@ -139,7 +139,7 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par
         ggml_tensor * predictions = altup_predict(cur, il);  // [n_embd, n_tokens, n_altup]
 
         // predicted value will go through self-attention and laurel
-        ggml_tensor * active_prediction = ggml_view_2d_slice(ctx0, predictions, i_altup_act);  // [n_embd, n_tokens]
+        ggml_tensor * active_prediction = gemma3n_view_2d_slice(ctx0, predictions, i_altup_act);  // [n_embd, n_tokens]
         cur = active_prediction;
         cb(cur, "active_prediction", il);
 
@@ -236,13 +236,13 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par
 
         ggml_tensor * first_prediction;                                                   // [n_embd, n_tokens]
         {
-            first_prediction = ggml_view_2d_slice(ctx0, corrected, i_altup_act);          // [n_embd, n_tokens]
+            first_prediction = gemma3n_view_2d_slice(ctx0, corrected, i_altup_act);       // [n_embd, n_tokens]
             first_prediction = ggml_mul(ctx0, first_prediction, model.layers[il].altup_correct_scale);
             first_prediction = build_lora_mm(model.layers[il].per_layer_inp_gate, first_prediction);
             first_prediction = ggml_gelu(ctx0, first_prediction);                 // [n_embd_altup, n_tokens]
             cb(first_prediction, "first_prediction_gated", il);
 
-            ggml_tensor * inp_this_layer = ggml_view_2d_slice(ctx0, inp_per_layer, il);   // [n_embd_altup, n_tokens]
+            ggml_tensor * inp_this_layer = gemma3n_view_2d_slice(ctx0, inp_per_layer, il);   // [n_embd_altup, n_tokens]
             first_prediction = ggml_mul(ctx0, first_prediction, inp_this_layer);  // [n_embd_altup, n_tokens]
             cb(first_prediction, "first_prediction_scaled", il);
 
@@ -253,7 +253,7 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par
         }
         // equivalent to python code: corrected_predictions[1:] += first_prediction
         {
-            ggml_tensor * slice_first = ggml_view_2d_slice(ctx0, corrected, 0);
+            ggml_tensor * slice_first = gemma3n_view_2d_slice(ctx0, corrected, 0);
             ggml_tensor * slice_rest  = ggml_view_3d(
                 ctx0, corrected, n_embd, n_tokens, n_altup - 1, ggml_row_size(corrected->type, n_embd),
                 ggml_row_size(corrected->type, n_embd * n_tokens), n_embd * n_tokens * ggml_element_size(corrected));
@@ -271,7 +271,7 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par
 
     // cur now has multiple altup(s), we want to merge them back to 1 altup
     {
-        ggml_tensor * target_magnitude = calc_magnitude(ggml_view_2d_slice(ctx0, cur, i_altup_act));  // [n_embd, n_tokens]
+        ggml_tensor * target_magnitude = calc_magnitude(gemma3n_view_2d_slice(ctx0, cur, i_altup_act));  // [n_embd, n_tokens]
         // do a view to skip the first slice (active altup)
         ggml_tensor * alt_slice =
             ggml_view_3d(ctx0, cur, n_embd, n_tokens, n_altup - 1, ggml_row_size(cur->type, n_embd),
@@ -283,9 +283,9 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par
         cb(altup_unembd, "altup_unembd", -1);
 
         // equivalent to torch.mean(hidden_states, dim=0)
-        cur = ggml_view_2d_slice(ctx0, cur, 0);  // [n_embd, n_tokens]
+        cur = gemma3n_view_2d_slice(ctx0, cur, 0);  // [n_embd, n_tokens]
         for (int i = 0; i < n_altup - 1; ++i) {
-            cur = ggml_add(ctx0, cur, ggml_view_2d_slice(ctx0, altup_unembd, i));
+            cur = ggml_add(ctx0, cur, gemma3n_view_2d_slice(ctx0, altup_unembd, i));
         }
         cur = ggml_scale(ctx0, cur, 1.0f / float(n_altup));  // [n_embd, n_tokens]
         cb(cur, "unembd_merged", -1);
@@ -419,7 +419,7 @@ ggml_tensor * llama_model_gemma3n::graph::altup_compute_router_modalities(ggml_t
 // input cur shape: [n_embd, n_tokens, n_altup]
 // output    shape: [n_embd, n_tokens, n_altup]
 ggml_tensor * llama_model_gemma3n::graph::altup_predict(ggml_tensor * cur, int il) {
-    ggml_tensor * activated  = ggml_view_2d_slice(ctx0, cur, i_altup_act);      // [n_embd, n_tokens]
+    ggml_tensor * activated  = gemma3n_view_2d_slice(ctx0, cur, i_altup_act);   // [n_embd, n_tokens]
     ggml_tensor * modalities = altup_compute_router_modalities(activated, il);  // [n_altup, n_tokens]
     cb(modalities, "modalities", il);
 
@@ -447,7 +447,7 @@ ggml_tensor * llama_model_gemma3n::graph::altup_correct(ggml_tensor * prediction
     ggml_tensor * modalities = altup_compute_router_modalities(activated, il);  // [n_altup, n_tokens]
     cb(modalities, "modalities", il);
 
-    ggml_tensor * active_prediction = ggml_view_2d_slice(ctx0, predictions, i_altup_act);
+    ggml_tensor * active_prediction = gemma3n_view_2d_slice(ctx0, predictions, i_altup_act);
     ggml_tensor * innovation        = ggml_sub(ctx0, activated, active_prediction);  // [n_embd, n_tokens]
     cb(innovation, "innovation", il);
 

src/models/gemma4.cpp

@@ -145,7 +145,7 @@ std::unique_ptr<llm_graph_context> llama_model_gemma4::build_arch_graph(const ll
 }
 
 // get 2D slice view from a 3D tensor, the idx corresponds to the 3rd dim
-static ggml_tensor * ggml_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, int idx) {
+static ggml_tensor * gemma4_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, int idx) {
     GGML_ASSERT(idx < (int) x->ne[2]);
     return ggml_view_2d(ctx0, x, x->ne[0], x->ne[1], ggml_row_size(x->type, x->ne[0]),
                         idx * x->ne[0] * x->ne[1] * ggml_element_size(x));
@@ -372,7 +372,7 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
             cur = build_lora_mm(model.layers[il].per_layer_inp_gate, cur); // [n_embd_per_layer, n_tokens]
             cur = ggml_gelu(ctx0, cur);
 
-            ggml_tensor * inp_this_layer = ggml_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_per_layer, n_tokens]
+            ggml_tensor * inp_this_layer = gemma4_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_per_layer, n_tokens]
 
             // TODO @ngxson : improve this
             if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {

src/models/minimax-01.cpp

@@ -229,6 +229,7 @@ llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_
     ggml_set_input(inp->inp_k_decay);
     cb(inp->inp_k_decay, "k_decay_exp", -1);
 
+    // [TAG_RESERVE_DIAG_DECAY]
     inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs);
     ggml_set_input(inp->inp_diag_decay);
     cb(inp->inp_diag_decay, "diag_decay_exp", -1);

src/models/plamo2.cpp

@@ -142,6 +142,11 @@ llama_model_plamo2::graph::graph(const llama_model & model, const llm_graph_para
             cur = build_plamo2_attn_layer(inp_hybrid->get_attn(), inp_pos, cur, model, il);
         }
 
+        if (il == n_layer - 1 && inp_out_ids) {
+            cur      = ggml_get_rows(ctx0, cur, inp_out_ids);
+            residual = ggml_get_rows(ctx0, residual, inp_out_ids);
+        }
+
         // post_mixer_norm
         cur = build_norm(cur, model.layers[il].attn_post_norm, NULL, LLM_NORM_RMS, il);
         cb(cur, "attn_post_norm", il);
@@ -167,11 +172,6 @@ llama_model_plamo2::graph::graph(const llama_model & model, const llm_graph_para
         cur = build_norm(cur, model.layers[il].ffn_post_norm, NULL, LLM_NORM_RMS, il);
         cb(cur, "ffn_post_norm", il);
 
-        if (il == n_layer - 1 && inp_out_ids) {
-            cur      = ggml_get_rows(ctx0, cur, inp_out_ids);
-            residual = ggml_get_rows(ctx0, residual, inp_out_ids);
-        }
-
         // residual connection
         cur = ggml_add(ctx0, cur, residual);
         cb(cur, "ffn_residual", il);

src/models/qwen3vl.cpp

@@ -18,6 +18,7 @@ void llama_model_qwen3vl::load_arch_tensors(llama_model_loader &) {
 
     int64_t n_vocab_out = n_vocab;
     if (arch == LLM_ARCH_QWEN3TTS) {
+        // [TAG_LLAMA_N_VOCAB_OUT]
         n_vocab_out = 3072;
     }
 

tests/.gitignore

@@ -1,6 +1,7 @@
 *
 !*.*
 !snapshots/
+!fusion/
 *.o
 ggml-common.h
 **/*.swp

tests/CMakeLists.txt

@@ -196,7 +196,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
 
     # llama_build_and_test(test-double-float.cpp) # SLOW
 
-    llama_build_and_test(test-llama-archs.cpp)
+    llama_build(test-llama-archs.cpp)
 
     set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/")
     file(MAKE_DIRECTORY "${MODEL_DIR}")
@@ -255,6 +255,8 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
         ARGS --models "${MODEL_DIR}"
     )
     set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED generate-models)
+
+    llama_build(test-fusion.cpp)
 endif()
 
 llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp)
@@ -276,6 +278,8 @@ llama_build_and_test(
     peg-parser/test-unicode.cpp
     peg-parser/tests.h
 )
+target_precompile_headers(test-peg-parser PRIVATE peg-parser/tests.h)
+
 
 if (NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "s390x")
     set(MODEL_NAME "tinyllamas/stories15M-q4_0.gguf")

tests/fusion/MTL.csv

@@ -0,0 +1,154 @@
+# test-fusion baseline for device MTL
+# arch              ,moe ,mode    ,label                       ,  count
+arcee               ,0   ,any     ,RMS_NORM+MUL                ,      5
+arctic              ,0   ,any     ,RMS_NORM+MUL                ,      7
+baichuan            ,0   ,any     ,RMS_NORM+MUL                ,      5
+bailingmoe          ,1   ,any     ,ADD+ADD                     ,      2
+bailingmoe          ,1   ,any     ,RMS_NORM+MUL                ,      5
+bailingmoe2         ,1   ,any     ,ADD+ADD                     ,      1
+bailingmoe2         ,1   ,any     ,RMS_NORM+MUL                ,      9
+bailingmoe3         ,1   ,any     ,ADD+ADD                     ,      1
+bailingmoe3         ,1   ,any     ,GATED_DELTA_NET+CPY         ,      1
+bailingmoe3         ,1   ,any     ,RMS_NORM+MUL                ,      8
+bloom               ,0   ,any     ,NORM+MUL+ADD                ,      6
+chatglm             ,0   ,any     ,RMS_NORM+MUL                ,      5
+codeshell           ,0   ,any     ,NORM+MUL+ADD                ,      5
+cogvlm              ,0   ,any     ,RMS_NORM+MUL                ,      5
+command-r           ,0   ,any     ,NORM+MUL                    ,      3
+dbrx                ,0   ,any     ,NORM+MUL                    ,      5
+deci                ,0   ,any     ,RMS_NORM+MUL                ,      5
+deepseek            ,0   ,any     ,ADD+ADD                     ,      1
+deepseek            ,0   ,any     ,RMS_NORM+MUL                ,      5
+deepseek2           ,0   ,any     ,ADD+ADD                     ,      1
+deepseek2           ,0   ,any     ,RMS_NORM+MUL                ,      9
+deepseek32          ,0   ,any     ,ADD+ADD                     ,      1
+deepseek32          ,0   ,any     ,NORM+MUL+ADD                ,      2
+deepseek32          ,0   ,any     ,RMS_NORM+MUL                ,      9
+deepseek4           ,0   ,any     ,RMS_NORM+MUL                ,     20
+dots1               ,0   ,any     ,ADD+ADD                     ,      1
+dots1               ,0   ,any     ,RMS_NORM+MUL                ,      9
+dream               ,0   ,any     ,RMS_NORM+MUL                ,      5
+ernie4_5-moe        ,1   ,any     ,ADD+ADD                     ,      1
+ernie4_5-moe        ,1   ,any     ,RMS_NORM+MUL                ,      5
+ernie4_5            ,0   ,any     ,RMS_NORM+MUL                ,      5
+exaone              ,0   ,any     ,RMS_NORM+MUL                ,      5
+exaone4             ,0   ,any     ,RMS_NORM+MUL                ,      5
+exaone4             ,0   ,any     ,RMS_NORM+MUL+ADD            ,      4
+falcon              ,0   ,any     ,ADD+ADD                     ,      2
+falcon              ,0   ,any     ,NORM+MUL+ADD                ,      5
+falcon-h1           ,0   ,any     ,ADD+ADD                     ,      2
+falcon-h1           ,0   ,any     ,RMS_NORM+MUL                ,      9
+gemma               ,0   ,any     ,RMS_NORM+MUL                ,      5
+gemma2              ,0   ,any     ,RMS_NORM+MUL                ,      5
+gemma2              ,0   ,any     ,RMS_NORM+MUL+ADD            ,      4
+glm-dsa             ,0   ,any     ,ADD+ADD                     ,      1
+glm-dsa             ,0   ,any     ,NORM+MUL+ADD                ,      2
+glm-dsa             ,0   ,any     ,RMS_NORM+MUL                ,      9
+glm4                ,0   ,any     ,RMS_NORM+MUL                ,      5
+glm4                ,0   ,any     ,RMS_NORM+MUL+ADD            ,      4
+glm4moe             ,1   ,any     ,ADD+ADD                     ,      1
+glm4moe             ,1   ,any     ,RMS_NORM+MUL                ,      9
+gpt-oss             ,0   ,any     ,RMS_NORM+MUL                ,      5
+gpt2                ,0   ,any     ,NORM+MUL+ADD                ,      5
+gptneox             ,0   ,any     ,NORM+MUL+ADD                ,      5
+granite             ,0   ,any     ,RMS_NORM+MUL                ,      5
+granite             ,0   ,any     ,RMS_NORM+MUL                ,      5
+granitehybrid       ,0   ,any     ,RMS_NORM+MUL                ,      6
+granitemoe          ,1   ,any     ,RMS_NORM+MUL                ,      5
+granitemoe          ,1   ,any     ,RMS_NORM+MUL                ,      5
+grok                ,0   ,any     ,RMS_NORM+MUL                ,      5
+grok                ,0   ,any     ,RMS_NORM+MUL+ADD            ,      4
+grovemoe            ,1   ,any     ,ADD+ADD                     ,      2
+grovemoe            ,1   ,any     ,RMS_NORM+MUL                ,      9
+hunyuan-dense       ,0   ,any     ,RMS_NORM+MUL                ,      9
+hunyuan-moe         ,1   ,any     ,ADD+ADD                     ,      2
+hunyuan-moe         ,1   ,any     ,RMS_NORM+MUL                ,      9
+hunyuan_vl          ,0   ,any     ,RMS_NORM+MUL                ,      9
+hy_v3               ,0   ,any     ,ADD+ADD                     ,      2
+hy_v3               ,0   ,any     ,RMS_NORM+MUL                ,      9
+hy_v4               ,0   ,any     ,NORM+MUL+ADD                ,      1
+hy_v4               ,0   ,any     ,RMS_NORM+MUL                ,      9
+internlm2           ,0   ,any     ,RMS_NORM+MUL                ,      5
+jais                ,0   ,any     ,NORM+MUL+ADD                ,      5
+jais2               ,0   ,any     ,NORM+MUL+ADD                ,      5
+jamba               ,0   ,any     ,RMS_NORM+MUL                ,      8
+kimi-k3             ,0   ,any     ,GATED_DELTA_NET+CPY         ,      1
+kimi-k3             ,0   ,any     ,RMS_NORM+MUL                ,     17
+kimi-linear         ,0   ,any     ,ADD+ADD                     ,      1
+kimi-linear         ,0   ,any     ,GATED_DELTA_NET+CPY         ,      1
+kimi-linear         ,0   ,any     ,RMS_NORM+MUL                ,      7
+lfm2                ,0   ,any     ,RMS_NORM+MUL                ,      7
+lfm2moe             ,1   ,any     ,RMS_NORM+MUL                ,      7
+llada               ,0   ,any     ,RMS_NORM+MUL                ,      5
+llada-moe           ,1   ,any     ,RMS_NORM+MUL                ,      9
+llama               ,0   ,any     ,RMS_NORM+MUL                ,      5
+llama               ,0   ,any     ,RMS_NORM+MUL                ,      5
+llama4              ,0   ,any     ,ADD+ADD                     ,      2
+llama4              ,0   ,any     ,RMS_NORM+MUL                ,      9
+maincoder           ,0   ,any     ,RMS_NORM+MUL                ,      9
+mamba               ,0   ,any     ,RMS_NORM+MUL                ,      3
+mamba2              ,0   ,any     ,RMS_NORM+MUL                ,      5
+minicpm             ,0   ,any     ,RMS_NORM+MUL                ,      5
+minicpm             ,0   ,any     ,RMS_NORM+MUL                ,      5
+minicpm3            ,0   ,any     ,RMS_NORM+MUL                ,      9
+minimax-01          ,0   ,any     ,RMS_NORM+MUL                ,      6
+minimax-m2          ,0   ,any     ,RMS_NORM+MUL                ,      9
+minimax-m3          ,0   ,any     ,ADD+ADD                     ,      1
+minimax-m3          ,0   ,any     ,RMS_NORM+MUL                ,     11
+mistral3            ,0   ,any     ,RMS_NORM+MUL                ,      5
+mistral3            ,0   ,any     ,RMS_NORM+MUL                ,      5
+mistral4            ,0   ,any     ,ADD+ADD                     ,      1
+mistral4            ,0   ,any     ,RMS_NORM+MUL                ,      9
+mpt                 ,0   ,any     ,NORM+MUL+ADD                ,      5
+nanbeige            ,0   ,any     ,RMS_NORM+MUL                ,      5
+nemotron            ,0   ,any     ,NORM+MUL+ADD                ,      5
+nemotron_h          ,0   ,any     ,RMS_NORM+MUL                ,      5
+nemotron_h_moe      ,1   ,any     ,RMS_NORM+MUL                ,      5
+olmoe               ,1   ,any     ,RMS_NORM+MUL                ,      9
+openelm             ,0   ,any     ,RMS_NORM+MUL                ,      9
+orion               ,0   ,any     ,NORM+MUL+ADD                ,      5
+paddleocr           ,0   ,any     ,RMS_NORM+MUL                ,      5
+pangu-embedded      ,0   ,any     ,RMS_NORM+MUL                ,      5
+phi2                ,0   ,any     ,ADD+ADD                     ,      2
+phi2                ,0   ,any     ,NORM+MUL+ADD                ,      3
+phi3                ,0   ,any     ,RMS_NORM+MUL                ,      5
+phimoe              ,1   ,any     ,RMS_NORM+MUL+ADD            ,      5
+plamo               ,0   ,any     ,ADD+ADD                     ,      2
+plamo               ,0   ,any     ,RMS_NORM+MUL                ,      3
+plamo2              ,0   ,any     ,RMS_NORM+MUL                ,     10
+plamo2              ,0   ,any     ,RMS_NORM+MUL+ADD            ,      4
+pockettts           ,0   ,any     ,NORM+MUL+ADD                ,      5
+qwen                ,0   ,any     ,RMS_NORM+MUL                ,      5
+qwen2               ,0   ,any     ,RMS_NORM+MUL                ,      5
+qwen2moe            ,1   ,any     ,ADD+ADD                     ,      2
+qwen2moe            ,1   ,any     ,RMS_NORM+MUL                ,      5
+qwen2vl             ,0   ,any     ,RMS_NORM+MUL                ,      5
+qwen3               ,0   ,any     ,RMS_NORM+MUL                ,      9
+qwen35              ,0   ,any     ,GATED_DELTA_NET+CPY         ,      1
+qwen35              ,0   ,any     ,RMS_NORM+MUL                ,      8
+qwen35moe           ,1   ,any     ,ADD+ADD                     ,      2
+qwen35moe           ,1   ,any     ,GATED_DELTA_NET+CPY         ,      1
+qwen35moe           ,1   ,any     ,RMS_NORM+MUL                ,      8
+qwen3moe            ,1   ,any     ,RMS_NORM+MUL                ,      9
+qwen3next           ,0   ,any     ,ADD+ADD                     ,      2
+qwen3next           ,0   ,any     ,GATED_DELTA_NET+CPY         ,      1
+qwen3next           ,0   ,any     ,RMS_NORM+MUL                ,      8
+qwen3tts            ,0   ,any     ,RMS_NORM+MUL                ,      9
+qwen3vl             ,0   ,any     ,RMS_NORM+MUL                ,      9
+qwen3vlmoe          ,1   ,any     ,RMS_NORM+MUL                ,      9
+qwen4exp            ,0   ,any     ,ADD+ADD+ADD                 ,      5
+qwen4exp            ,0   ,any     ,ADD+ADD+ADD+ADD+ADD+ADD+ADD ,      9
+qwen4exp            ,0   ,any     ,GATED_DELTA_NET+CPY         ,      1
+qwen4exp            ,0   ,any     ,RMS_NORM+MUL                ,      5
+refact              ,0   ,any     ,RMS_NORM+MUL                ,      5
+refact              ,0   ,any     ,RMS_NORM+MUL                ,      5
+rnd1                ,0   ,any     ,RMS_NORM+MUL                ,      9
+seed_oss            ,0   ,any     ,RMS_NORM+MUL                ,      5
+smallthinker        ,0   ,any     ,RMS_NORM+MUL                ,      5
+smollm3             ,0   ,any     ,RMS_NORM+MUL                ,      5
+stablelm            ,0   ,any     ,NORM+MUL                    ,      4
+stablelm            ,0   ,any     ,NORM+MUL+ADD                ,      5
+starcoder           ,0   ,any     ,NORM+MUL+ADD                ,      5
+starcoder2          ,0   ,any     ,NORM+MUL+ADD                ,      5
+talkie              ,0   ,any     ,ADD+ADD                     ,      2
+xverse              ,0   ,any     ,RMS_NORM+MUL                ,      5

tests/test-backend-ops.cpp

@@ -1219,6 +1219,11 @@ struct test_case {
         }
     }
 
+    // re-draw data-dependent inputs between timed perf iterations
+    virtual void reinit_perf_iter(ggml_context * ctx) {
+        GGML_UNUSED(ctx);
+    }
+
     virtual size_t op_size(ggml_tensor * t) {
         size_t size = ggml_nbytes(t);
         // add source tensors
@@ -1653,6 +1658,9 @@ struct test_case {
             total_time_us += end_time - start_time;
             total_mem += mem;
             total_runs += n_runs;
+
+            // re-draw any data-dependent inputs (expert ids) outside the timed region
+            reinit_perf_iter(ctx.get());
         } while (total_time_us < 1000*1000); // run for at least 1 second
 
         // Create test result
@@ -4638,6 +4646,122 @@ struct test_gated_delta_net : public test_case {
     }
 };
 
+// GGML_OP_GATED_DELTA_NET + GGML_OP_CPY (recurrent cache fusion)
+struct test_gated_delta_net_cache_fusion : public test_case {
+    const ggml_type type;
+
+    const int64_t head_count;
+    const int64_t head_size;
+    const int64_t n_seq_tokens;
+    const int64_t n_seqs;
+    const int64_t K; // snapshot slot count (>1)
+
+    ggml_tensor * cpy_node = nullptr;
+
+    std::string vars() override {
+        return VARS_TO_STR6(type, head_count, head_size, n_seq_tokens, n_seqs, K);
+    }
+
+    test_gated_delta_net_cache_fusion(ggml_type type = GGML_TYPE_F32,
+            int64_t head_count = 4, int64_t head_size = 32, int64_t n_seq_tokens = 2, int64_t n_seqs = 1,
+            int64_t K = 2)
+        : type(type), head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), K(K) {}
+
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        const int64_t S_v = head_size;
+        const int64_t H_v = head_count;
+        const int64_t H_k = head_count;
+        const int64_t D   = S_v * S_v * H_v;
+        const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
+
+        ggml_tensor * q     = ggml_new_tensor_4d(ctx, type, head_size, H_k, n_seq_tokens, n_seqs);
+        ggml_tensor * k     = ggml_new_tensor_4d(ctx, type, head_size, H_k, n_seq_tokens, n_seqs);
+        ggml_tensor * v     = ggml_new_tensor_4d(ctx, type, head_size, H_v, n_seq_tokens, n_seqs);
+        ggml_set_name(q, "q");
+        ggml_set_name(k, "k");
+        ggml_set_name(v, "v");
+        ggml_tensor * g     = ggml_new_tensor_4d(ctx, type, 1, H_v, n_seq_tokens, n_seqs);
+        ggml_tensor * beta  = ggml_new_tensor_4d(ctx, type, 1, H_v, n_seq_tokens, n_seqs);
+        ggml_tensor * state = ggml_new_tensor_4d(ctx, type, head_size, head_size, H_v, n_seqs);
+        ggml_set_name(g,     "g");
+        ggml_set_name(beta,  "beta");
+        ggml_set_name(state, "state");
+
+        q = ggml_l2_norm(ctx, q, 1e-6f);
+        k = ggml_l2_norm(ctx, k, 1e-6f);
+
+        ggml_tensor * gdn_out = ggml_gated_delta_net(ctx, q, k, v, g, beta, state, K);
+        ggml_set_name(gdn_out, "gdn_out");
+
+        // attn scores view (first part of the gdn output)
+        ggml_tensor * attn = ggml_view_4d(ctx, gdn_out,
+                S_v, H_v, n_seq_tokens, n_seqs,
+                ggml_row_size(gdn_out->type, S_v),
+                ggml_row_size(gdn_out->type, S_v * H_v),
+                ggml_row_size(gdn_out->type, S_v * H_v * n_seq_tokens), 0);
+        ggml_set_name(attn, "attn");
+
+        // snapshot tail view [D, n_seqs, n_written]
+        const int64_t attn_score_elems = S_v * H_v * n_seq_tokens * n_seqs;
+        ggml_tensor * src = ggml_view_3d(ctx, gdn_out,
+                D, n_seqs, n_written,
+                ggml_row_size(gdn_out->type, D),
+                ggml_row_size(gdn_out->type, D * n_seqs),
+                ggml_row_size(gdn_out->type, attn_score_elems));
+
+        // recurrent cache view [D, n_seqs, n_written]
+        ggml_tensor * cache = ggml_new_tensor_3d(ctx, type, D, n_seqs, n_written);
+        ggml_set_name(cache, "cache");
+        ggml_tensor * dst = ggml_view_3d(ctx, cache,
+                D, n_seqs, n_written,
+                ggml_row_size(cache->type, D),
+                ggml_row_size(cache->type, D * n_seqs), 0);
+
+        ggml_tensor * cpy = ggml_cpy(ctx, src, dst);
+        ggml_set_name(cpy, "gdn_cache_cpy");
+        cpy_node = cpy;
+
+        // read the cpy output (not the plain dst view, which would not pull the cpy into the graph)
+        // so that neither the gdn nor the cpy is the graph output
+        ggml_tensor * out = ggml_sum(ctx, cpy);
+        return out;
+    }
+
+    std::string op_desc(ggml_tensor * t) override {
+        GGML_UNUSED(t);
+        return "GATED_DELTA_NET_CACHE_FUSION";
+    }
+
+    bool run_whole_graph() override { return true; }
+    std::vector<ggml_tensor *> fusion_test_nodes() override { return { cpy_node }; }
+
+    uint64_t op_flops(ggml_tensor * t) override {
+        GGML_UNUSED(t);
+        const uint64_t S_v = head_size;
+        const uint64_t H_v = head_count;
+        const uint64_t T   = n_seq_tokens;
+        const uint64_t B   = n_seqs;
+        return (4ull*S_v + 2ull*S_v*S_v) * H_v * T * B;
+    }
+
+    void initialize_tensors(ggml_context * ctx) override {
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            if (ggml_is_view_op(t->op)) { continue; }
+            if (strcmp(t->name, "g") == 0) {
+                init_tensor_uniform(t, -20.0f, -1e-4f);
+            } else if (strcmp(t->name, "beta") == 0) {
+                init_tensor_uniform(t, 0.0f, 1.0f);
+            } else if (strcmp(t->name, "v") == 0) {
+                init_tensor_uniform(t, -0.3f, 5.0f);
+            } else if (strcmp(t->name, "cache") == 0) {
+                init_tensor_uniform(t, 0.0f, 0.0f);
+            } else {
+                init_tensor_uniform(t);
+            }
+        }
+    }
+};
+
 // GGML_OP_GATED_LINEAR_ATTN
 struct test_gla : public test_case {
     const ggml_type type;
@@ -4884,25 +5008,31 @@ struct test_mul_mat_hadamard : public test_mul_mat {
     }
 };
 
-static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) {
+static void init_mul_mat_id_ids(ggml_context * ctx, int n_mats) {
     std::random_device rd;
     std::default_random_engine rng(rd());
     for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
-        if (t->type == GGML_TYPE_I32) {
-            if (ggml_is_view_op(t->op)) { continue; }
-            // ids
-            for (int64_t r = 0; r < ggml_nrows(t); r++) {
-                std::vector<int32_t> data(t->ne[0]);
-                for (int i = 0; i < t->ne[0]; i++) {
-                    data[i] = i % n_mats;
-                }
-                std::shuffle(data.begin(), data.end(), rng);
-                ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t));
+        if (t->type != GGML_TYPE_I32 || ggml_is_view_op(t->op)) {
+            continue;
+        }
+        for (int64_t r = 0; r < ggml_nrows(t); r++) {
+            std::vector<int32_t> data(t->ne[0]);
+            for (int i = 0; i < t->ne[0]; i++) {
+                data[i] = i % n_mats;
             }
-        } else {
+            std::shuffle(data.begin(), data.end(), rng);
+            ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t));
+        }
+    }
+}
+
+static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) {
+    for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+        if (t->type != GGML_TYPE_I32) {
             init_tensor_uniform(t);
         }
     }
+    init_mul_mat_id_ids(ctx, n_mats);
 }
 
 // GGML_OP_MUL_MAT_ID
@@ -4969,6 +5099,10 @@ struct test_mul_mat_id : public test_case {
     void initialize_tensors(ggml_context * ctx) override {
         init_mul_mat_id_tensors(ctx, n_mats);
     }
+
+    void reinit_perf_iter(ggml_context * ctx) override {
+        init_mul_mat_id_ids(ctx, n_mats);
+    }
 };
 
 // GGML_OP_MUL_MAT_ID + GGML_OP_ADD or GGML_OP_MUL
@@ -9774,6 +9908,19 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, 256, {2, 3}, {1, 1}, {0, 1, 3, 2}));
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, 256, {2, 3}, {1, 1}, {0, 3, 2, 1}));
 
+    // token-tile boundary coverage. With n_used == n_mats every token routes to every expert, so
+    // each expert receives exactly n rows, with no dependence on the random draw. mul_mm_id is used
+    // from 32 tokens up: n = 32, 33, 47, 48, 49 reach it, leaving a last tile of 32, 1, 15, 16 and
+    // 17 rows - 16 and 17 straddle the point where the upper half stops being skipped. The smaller
+    // n cover the same row counts on the mat-vec path.
+    for (ggml_type type_a : {GGML_TYPE_Q4_K, GGML_TYPE_IQ2_XS, GGML_TYPE_F16}) {
+        for (int n : {1, 15, 16, 17, 31, 32, 33, 47, 48, 49}) {
+            test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 4, 4, false, 512, n, 256));
+        }
+        // experts that receive no rows at all
+        test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 8, 1, false, 512, 1, 256));
+    }
+
     for (ggml_type type_a : other_types) {
         for (ggml_type type_b : {GGML_TYPE_F32}) {
             if (ggml_blck_size(type_a) != 256) {
@@ -10574,6 +10721,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
                                                         GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
     }
 
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1},   512, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1},  4096,  64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1},  4096,  16, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {2, 1},  4096, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {4, 1},  4096, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, 4096, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+
     // dense-allocated (non-view) quant K/V at batch >= 64, in cache and native layouts
     test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
     test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {4, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
@@ -10734,6 +10888,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32,   8, 1, 1, false, false, /*K=*/3));
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64,  16, 2, 1, false, false, /*K=*/4));
 
+    // gdn + cache cpy fusion (K > 1)
+    test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 32,   2, 1, 2));
+    test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 64,   4, 1, 2));
+    test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 32,   4, 1, 4));
+    test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 8, 32,   4, 2, 4));
+    test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 32,   8, 1, 4));
+
 #if 0
     // these tests are disabled to save execution time, sbut they can be handy for debugging
     test_cases.emplace_back(new test_llama(2, true));
@@ -11031,6 +11192,15 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
     test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
     test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
 
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, 4096, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, 16384, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, 16384, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, 65536, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, 65536, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, 131072, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, 131072, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+
     for (int kv : { 4096, 8192, 16384,32768, 65536, }) {
         for (int hs : { 64, 128, 256, 576, }) {
             const int  hsv    = hs == 576 ? 512 : hs;

tests/test-fusion.cpp

@@ -0,0 +1,565 @@
+// test-fusion: verify the backend fusion logic against a per-device baseline.
+//
+// for every dummy model generated by test-llama-archs, the tool runs the model on a single
+// device with fusion enabled and disabled, and reports:
+//   - the per-fusion-type counters for each mode (prefill / decode, merged into "any" when the
+//     per-graph counts match)
+//   - the NMSE between the fused and unfused logits
+//   - the NMSE between the device and a CPU reference
+//
+// the per-fusion-type counters are compared against a per-device baseline file (CSV) so a
+// fusion pattern that silently stops matching (or fires when it should not) is caught as a
+// regression.
+//
+// usage:
+//   test-fusion --models DIR --device MTL0 --record baseline.csv   # generate a baseline
+//   test-fusion --models DIR --device MTL0 --check  baseline.csv   # validate against it
+//   test-fusion --model FILE --device MTL0 --check  baseline.csv   # validate a single model
+
+#include "common.h"
+#include "log.h"
+#include "llama-cpp.h"
+
+#include "ggml.h"
+#include "gguf.h"
+
+#include <algorithm>
+#include <array>
+#include <cstring>
+#include <filesystem>
+#include <fstream>
+#include <iomanip>
+#include <iostream>
+#include <map>
+#include <random>
+#include <string>
+#include <vector>
+
+// generic fusion debugging API, resolved through the ad-hoc get_proc_address mechanism
+// (not part of the official ggml backend interface yet). a backend that adopts fusion debugging
+// exports these exact names.
+typedef void * ggml_backend_fusion_t;
+
+typedef ggml_backend_fusion_t ( * fusion_get_t)       (ggml_backend_dev_t);
+typedef void ( * fusion_stats_init_t)  (ggml_backend_fusion_t);
+typedef void ( * fusion_stats_reset_t) (ggml_backend_fusion_t);
+typedef int  ( * fusion_stats_get_t)   (ggml_backend_fusion_t, const char **, uint64_t *, int);
+typedef void ( * fusion_set_enabled_t) (ggml_backend_fusion_t, bool);
+
+static bool silent_model_load_progress(float, void *) {
+    return true;
+}
+
+struct gguf_context_ptr {
+    gguf_context * ctx;
+    gguf_context_ptr(gguf_context * c) : ctx(c) {}
+    ~gguf_context_ptr() { if (ctx) { gguf_free(ctx); } }
+    gguf_context * get() const { return ctx; }
+    gguf_context_ptr(const gguf_context_ptr &) = delete;
+    gguf_context_ptr & operator=(const gguf_context_ptr &) = delete;
+};
+
+// NMSE between two vectors (same as tests/test-llama-archs.cpp)
+static double nmse(const std::vector<float> & a, const std::vector<float> & b) {
+    GGML_ASSERT(a.size() == b.size());
+    double mse_a_b = 0.0;
+    double mse_a_0 = 0.0;
+
+    for (size_t i = 0; i < a.size(); i++) {
+        const float a_i = a[i];
+        const float b_i = b[i];
+
+        mse_a_b += (a_i - b_i) * (a_i - b_i);
+        mse_a_0 += a_i * a_i;
+    }
+
+    return mse_a_b / mse_a_0;
+}
+
+// deterministic token sequence
+static std::vector<llama_token> get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed) {
+    std::mt19937 gen(seed);
+    std::uniform_int_distribution<> dis(0, n_vocab - 1);
+    std::vector<llama_token> ret;
+    ret.reserve(n_tokens);
+    for (uint32_t i = 0; i < n_tokens; i++) {
+        ret.push_back(dis(gen));
+    }
+    return ret;
+}
+
+// trim leading/trailing whitespace (used when parsing padded CSV columns)
+static std::string trim(const std::string & s) {
+    const size_t b = s.find_first_not_of(" \t\r\n");
+    if (b == std::string::npos) {
+        return "";
+    }
+    const size_t e = s.find_last_not_of(" \t\r\n");
+    return s.substr(b, e - b + 1);
+}
+
+static std::string get_arch(const std::string & path) {
+    gguf_init_params params = { /*no_alloc=*/true, /*ctx=*/nullptr };
+    gguf_context_ptr ctx(gguf_init_from_file(path.c_str(), params));
+    if (!ctx.get()) {
+        throw std::runtime_error("failed to read gguf: " + path);
+    }
+    const int idx = gguf_find_key(ctx.get(), "general.architecture");
+    if (idx < 0) {
+        return "unknown";
+    }
+    const char * val = gguf_get_val_str(ctx.get(), idx);
+    return val ? val : "unknown";
+}
+
+static llama_model_ptr load_model(const std::string & path, ggml_backend_dev_t dev) {
+    llama_model_params model_params = llama_model_default_params();
+    model_params.progress_callback = silent_model_load_progress;
+    std::vector<ggml_backend_dev_t> devs = { dev, nullptr };
+    model_params.devices = devs.data();
+    model_params.split_mode = LLAMA_SPLIT_MODE_LAYER;
+
+    llama_model_ptr model(llama_model_load_from_file(path.c_str(), model_params));
+    if (!model) {
+        throw std::runtime_error("failed to load model: " + path);
+    }
+    return model;
+}
+
+// a fresh context (fresh state) from an already-loaded model
+static llama_context_ptr create_ctx(llama_model * model, int n_ubatch) {
+    llama_context_params ctx_params = llama_context_default_params();
+    ctx_params.n_ctx = 0;
+    ctx_params.n_threads = 4;
+    ctx_params.n_threads_batch = 4;
+    ctx_params.n_ubatch = n_ubatch;
+    ctx_params.n_batch = n_ubatch;
+
+    llama_context_ptr lctx(llama_init_from_model(model, ctx_params));
+    if (!lctx) {
+        throw std::runtime_error("failed to init context");
+    }
+    return lctx;
+}
+
+// decode all tokens in one batch; returns the logits of every token
+static std::vector<float> decode_prefill(llama_model * model, llama_context * lctx, const std::vector<llama_token> & tokens) {
+    const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
+    llama_batch batch = llama_batch_init(tokens.size(), 0, 1);
+    for (size_t i = 0; i < tokens.size(); i++) {
+        common_batch_add(batch, tokens[i], i, { 0 }, true);
+    }
+    batch.n_tokens = tokens.size();
+    if (llama_decode(lctx, batch)) {
+        llama_batch_free(batch);
+        throw std::runtime_error("prefill decode failed");
+    }
+
+    std::vector<float> ret;
+    ret.reserve(tokens.size() * n_vocab);
+    for (size_t i = 0; i < tokens.size(); i++) {
+        const float * logits_ith = llama_get_logits_ith(lctx, i);
+        for (uint32_t j = 0; j < n_vocab; j++) {
+            ret.push_back(logits_ith[j]);
+        }
+    }
+    llama_batch_free(batch);
+    return ret;
+}
+
+// decode one token at a time; returns the logits of the last token of each step
+static std::vector<float> decode_gen(llama_model * model, llama_context * lctx, const std::vector<llama_token> & tokens) {
+    const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
+    llama_batch batch = llama_batch_init(1, 0, 1);
+    std::vector<float> ret;
+    for (size_t i = 0; i < tokens.size(); i++) {
+        common_batch_clear(batch);
+        common_batch_add(batch, tokens[i], i, { 0 }, true);
+        if (llama_decode(lctx, batch)) {
+            llama_batch_free(batch);
+            throw std::runtime_error("decode failed");
+        }
+        const float * logits = llama_get_logits_ith(lctx, 0);
+        for (uint32_t j = 0; j < n_vocab; j++) {
+            ret.push_back(logits[j]);
+        }
+    }
+    llama_batch_free(batch);
+    return ret;
+}
+
+static void read_counts(fusion_stats_get_t api_stats_get, ggml_backend_fusion_t finfo,
+                        std::vector<const char *> & labels, std::vector<uint64_t> & counts) {
+    const int n = api_stats_get(finfo, nullptr, nullptr, 0);
+    labels.assign(n, nullptr);
+    counts.assign(n, 0);
+    api_stats_get(finfo, labels.data(), counts.data(), n);
+}
+
+// one row of the per-label report
+struct fusion_row {
+    std::string  arch;
+    bool         moe;
+    std::string  mode;
+    std::string  label;
+    uint64_t     count_fused;
+    uint64_t     count_unfused;
+    uint64_t     expected;
+    double       nmse_fus;
+    double       nmse_dev;
+    bool         ok_count; // counts match the baseline
+    bool         ok_nmse;  // nmse within epsilon
+};
+
+static void usage(const char * argv0) {
+    printf("%s: verify fusion counts on a device against a per-device baseline\n\n", argv0);
+    printf("usage: %s [options]\n\n", argv0);
+    printf("options:\n");
+    printf("  --models DIR   run over all .gguf models in a directory\n");
+    printf("  --model FILE   run over a single model file (mutually exclusive with --models)\n");
+    printf("  --device NAME  device to run on (e.g. MTL0, CPU)\n");
+    printf("  --record CSV   write the golden baseline\n");
+    printf("  --check  CSV   validate the counters against a baseline (default)\n");
+    printf("  -h, --help     show this message and exit\n");
+}
+
+int main(int argc, char ** argv) {
+    std::string models_dir;
+    std::string model_file;
+    std::string device_name;
+    std::string record_path;
+    std::string check_path;
+
+    for (int i = 1; i < argc; i++) {
+        const std::string arg = argv[i];
+        const auto next = [&](const char * name) -> std::string {
+            if (i + 1 >= argc) {
+                LOG_ERR("%s: %s requires an argument\n", __func__, name);
+                exit(1);
+            }
+            return argv[++i];
+        };
+        if (arg == "-h" || arg == "--help") {
+            usage(argv[0]);
+            exit(0);
+        }
+        if (arg == "--models")     { models_dir  = next("--models"); }
+        else if (arg == "--model") { model_file  = next("--model"); }
+        else if (arg == "--device"){ device_name = next("--device"); }
+        else if (arg == "--record"){ record_path = next("--record"); }
+        else if (arg == "--check") { check_path  = next("--check"); }
+        else {
+            LOG_ERR("%s: unknown argument: %s\n", __func__, arg.c_str());
+            return 1;
+        }
+    }
+
+    if (device_name.empty()) {
+        LOG_ERR("%s: --device NAME is required\n", __func__);
+        return 1;
+    }
+    if (models_dir.empty() && model_file.empty()) {
+        LOG_ERR("%s: --models DIR or --model FILE is required\n", __func__);
+        return 1;
+    }
+    if (!models_dir.empty() && !model_file.empty()) {
+        LOG_ERR("%s: --models DIR and --model FILE are mutually exclusive\n", __func__);
+        return 1;
+    }
+    if (!record_path.empty() && !check_path.empty()) {
+        LOG_ERR("%s: --record and --check are mutually exclusive\n", __func__);
+        return 1;
+    }
+
+    std::vector<std::string> models;
+    if (!model_file.empty()) {
+        if (!std::filesystem::is_regular_file(model_file)) {
+            LOG_ERR("%s: model file '%s' does not exist\n", __func__, model_file.c_str());
+            return 1;
+        }
+        models.push_back(model_file);
+    } else {
+        if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) {
+            LOG_ERR("%s: models directory '%s' does not exist\n", __func__, models_dir.c_str());
+            return 1;
+        }
+        for (const auto & entry : std::filesystem::directory_iterator(models_dir)) {
+            if (entry.is_regular_file() && entry.path().extension() == ".gguf") {
+                models.push_back(entry.path().string());
+            }
+        }
+        std::sort(models.begin(), models.end());
+
+        if (models.empty()) {
+            LOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str());
+            return 1;
+        }
+    }
+
+    common_init();
+    ggml_backend_load_all();
+
+    ggml_backend_dev_t dev = ggml_backend_dev_by_name(device_name.c_str());
+    if (!dev) {
+        LOG_WRN("%s: device '%s' not found - skipping (baseline is device-specific)\n",
+                __func__, device_name.c_str());
+        return 0;
+    }
+
+    // resolve the generic fusion debugging functions through the ad-hoc get_proc_address
+    // mechanism; a backend that does not adopt fusion debugging exports none of them
+    auto * reg = ggml_backend_dev_backend_reg(dev);
+
+    // output naming uses the backend base name (e.g. "MTL") rather than the specific device
+    // name (e.g. "MTL0") the test was invoked with
+    const std::string base_name = ggml_backend_reg_name(reg);
+
+    auto api_get         = (fusion_get_t)         ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_get");
+    auto api_stats_init  = (fusion_stats_init_t)  ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_init");
+    auto api_stats_reset = (fusion_stats_reset_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_reset");
+    auto api_stats_get   = (fusion_stats_get_t)   ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_get");
+    auto api_set_enabled = (fusion_set_enabled_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_set_enabled");
+
+    if (!api_get || !api_stats_init || !api_set_enabled || !api_stats_reset || !api_stats_get) {
+        LOG_ERR("%s: device '%s' does not export the generic fusion debugging API "
+                "(ggml_backend_fusion_*) - cannot run the fusion regression test\n",
+                __func__, device_name.c_str());
+        return 1;
+    }
+
+    ggml_backend_fusion_t finfo = api_get(dev);
+
+    // enable fusions stats
+    api_stats_init(finfo);
+
+    const bool has_counts = true;
+
+    // load the baseline (if any): key arch|moe|mode|label -> expected count
+    std::map<std::string, uint64_t> baseline;
+    if (!check_path.empty()) {
+        std::ifstream in(check_path);
+        if (!in) {
+            LOG_ERR("%s: cannot open baseline '%s'\n", __func__, check_path.c_str());
+            return 1;
+        }
+        std::string line;
+        while (std::getline(in, line)) {
+            if (line.empty() || line[0] == '#') {
+                continue;
+            }
+            std::vector<std::string> cols;
+            size_t pos = 0;
+            while ((pos = line.find(',')) != std::string::npos) {
+                cols.push_back(trim(line.substr(0, pos)));
+                line.erase(0, pos + 1);
+            }
+            cols.push_back(trim(line));
+            if (cols.size() != 5) {
+                continue;
+            }
+            baseline[cols[0] + "|" + cols[1] + "|" + cols[2] + "|" + cols[3]] = std::stoull(cols[4]);
+        }
+    }
+
+    std::vector<fusion_row> rows;
+
+    LOG_INF("%s: running fusion test over %zu models on '%s'\n", __func__, models.size(), base_name.c_str());
+
+    const size_t seed = 1;
+
+    for (const auto & model_path : models) {
+        const std::string arch = get_arch(model_path);
+        const bool moe = arch.find("moe") != std::string::npos;
+
+        llama_model_ptr model;
+        llama_model_ptr model_cpu;
+        uint32_t n_vocab = 0;
+        try {
+            model = load_model(model_path, dev);
+            model_cpu = load_model(model_path, ggml_backend_dev_by_name("CPU"));
+            n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model.get()));
+        } catch (const std::exception & e) {
+            LOG_ERR("%s: %s: %s\n", __func__, model_path.c_str(), e.what());
+            continue;
+        }
+
+        struct mode_cfg {
+            std::string name;
+            std::vector<float> (*decode)(llama_model *, llama_context *, const std::vector<llama_token> &);
+            int n_tokens;
+            int n_graphs;   // graph runs per mode (prefill=1, decode=16)
+        };
+        const mode_cfg modes[] = {
+            { "prefill", decode_prefill, 32, 1  },
+            { "decode",  decode_gen,     16, 16 },
+        };
+
+        // per-label, per-mode data for this model; prefill and decode are merged into a single
+        // "any" row when their per-graph counts match
+        struct mode_data {
+            bool     present;
+            uint64_t count_fused;   // per graph
+            uint64_t count_unfused; // per graph
+            double   nmse_fus;
+            double   nmse_dev;
+            bool     ok_nmse;
+        };
+        std::map<std::string, std::array<mode_data, 2>> mdata;
+
+        for (int mi = 0; mi < 2; mi++) {
+            const mode_cfg & mode = modes[mi];
+            const auto tokens = get_tokens(mode.n_tokens, n_vocab, seed);
+
+            // CPU reference for this mode (fresh context, fresh state)
+            std::vector<float> logits_cpu;
+            try {
+                llama_context_ptr ctx = create_ctx(model_cpu.get(), 32);
+                logits_cpu = mode.decode(model_cpu.get(), ctx.get(), tokens);
+            } catch (const std::exception & e) {
+                LOG_WRN("%s: %s: cpu reference: %s\n", __func__, model_path.c_str(), e.what());
+            }
+
+            // fused run on a fresh context (fresh state)
+            std::vector<float> logits_fused;
+            std::vector<const char *> labels;
+            std::vector<uint64_t> counts_fused;
+            {
+                llama_context_ptr ctx = create_ctx(model.get(), 32);
+                if (has_counts) {
+                    api_set_enabled(finfo, true);
+                    api_stats_reset(finfo);
+                }
+                logits_fused = mode.decode(model.get(), ctx.get(), tokens);
+                if (has_counts) {
+                    read_counts(api_stats_get, finfo, labels, counts_fused);
+                }
+            }
+
+            // unfused run on another fresh context (fresh state)
+            std::vector<float> logits_unfused;
+            std::vector<uint64_t> counts_unfused;
+            {
+                llama_context_ptr ctx = create_ctx(model.get(), 32);
+                if (has_counts) {
+                    api_set_enabled(finfo, false);
+                    api_stats_reset(finfo);
+                }
+                logits_unfused = mode.decode(model.get(), ctx.get(), tokens);
+                if (has_counts) {
+                    read_counts(api_stats_get, finfo, labels, counts_unfused);
+                }
+            }
+
+            const double nmse_fus = nmse(logits_fused, logits_unfused);
+            const double nmse_dev = logits_cpu.empty() ? 0.0 : nmse(logits_fused, logits_cpu);
+
+            if (has_counts) {
+                for (int i = 0; i < (int) labels.size(); i++) {
+                    const uint64_t fused   = counts_fused[i]   / mode.n_graphs;
+                    const uint64_t unfused = counts_unfused[i] / mode.n_graphs;
+                    if (fused == 0 && unfused == 0) {
+                        continue;
+                    }
+                    auto & d = mdata[labels[i]][mi];
+                    d.present       = true;
+                    d.count_fused   = fused;
+                    d.count_unfused = unfused;
+                    d.nmse_fus      = nmse_fus;
+                    d.nmse_dev      = nmse_dev;
+                    d.ok_nmse       = nmse_fus <= 1e-4;
+                }
+            } else {
+                rows.push_back({ arch, moe, mode.name, "?", 0, 0, 0, nmse_fus, nmse_dev, true, nmse_fus <= 1e-4 });
+            }
+        }
+
+        // build the per-label rows, merging prefill and decode into "any" when the per-graph
+        // counts match (they always do for the deterministic fusion table)
+        if (has_counts) {
+            for (auto & kv : mdata) {
+                const std::string & label = kv.first;
+                const auto & d = kv.second;
+                const bool both = d[0].present && d[1].present;
+                const bool match = both && d[0].count_fused == d[1].count_fused;
+
+                if (match) {
+                    // one "any" row; use the worst NMSE across the two modes
+                    const std::string any_key = arch + "|" + (moe ? "1" : "0") + "|any|" + label;
+                    const uint64_t expected = baseline.count(any_key) ? baseline.at(any_key) : 0;
+                    const bool ok_count = check_path.empty() || d[0].count_fused == expected;
+                    const bool ok_nmse   = d[0].ok_nmse && d[1].ok_nmse;
+                    const double nmse_fus = std::max(d[0].nmse_fus, d[1].nmse_fus);
+                    const double nmse_dev = std::max(d[0].nmse_dev, d[1].nmse_dev);
+                    rows.push_back({ arch, moe, "any", label, d[0].count_fused, d[0].count_unfused,
+                                     expected, nmse_fus, nmse_dev, ok_count, ok_nmse });
+                } else {
+                    // counts differ - keep a separate row per mode
+                    for (int mi = 0; mi < 2; mi++) {
+                        if (!d[mi].present) {
+                            continue;
+                        }
+                        const mode_data & a = d[mi];
+                        const std::string mode_key = arch + "|" + (moe ? "1" : "0") + "|" + modes[mi].name + "|" + label;
+                        const uint64_t expected = baseline.count(mode_key) ? baseline.at(mode_key) : 0;
+                        const bool ok_count = check_path.empty() || a.count_fused == expected;
+                        rows.push_back({ arch, moe, modes[mi].name, label, a.count_fused, a.count_unfused,
+                                         expected, a.nmse_fus, a.nmse_dev, ok_count, a.ok_nmse });
+                    }
+                }
+            }
+        }
+
+        LOG_INF("%s: %-20s (%s) done\n", __func__, arch.c_str(), model_path.c_str());
+    }
+
+    // print the report
+    {
+        std::ofstream out(record_path);
+        std::ostream & os = record_path.empty() ? std::cout : out;
+        if (!record_path.empty()) {
+            os << "# test-fusion baseline for device " << base_name << "\n";
+            os << "# " << std::left
+               << std::setw(18) << "arch"  << ','
+               << std::setw(4)  << "moe"   << ','
+               << std::setw(8)  << "mode"  << ','
+               << std::setw(28) << "label" << ','
+               << std::right << std::setw(7) << "count" << '\n';
+        }
+
+        LOG_INF("%-20s %-4s %-8s %-22s %7s %7s %7s %10s %10s %s\n",
+                "arch", "moe", "mode", "label", "fused", "unfused", "expected", "nmse_fus", "nmse_dev", "status");
+        int n_ok = 0;
+        int n_bad = 0;
+        for (const auto & r : rows) {
+            const bool ok = r.ok_count && r.ok_nmse;
+            const char * status = ok ? "ok" : "FAIL";
+            if (ok) { n_ok++; } else { n_bad++; }
+            LOG_INF("%-20s %-4s %-8s %-22s %7llu %7llu %7llu %10.2e %10.2e %s\n",
+                    r.arch.c_str(), r.moe ? "moe" : "dense", r.mode.c_str(), r.label.c_str(),
+                    (unsigned long long) r.count_fused, (unsigned long long) r.count_unfused,
+                    (unsigned long long) r.expected, r.nmse_fus, r.nmse_dev, status);
+            if (!record_path.empty()) {
+                os << std::left
+                   << std::setw(20) << r.arch << ','
+                   << std::setw(4)  << (r.moe ? "1" : "0") << ','
+                   << std::setw(8)  << r.mode << ','
+                   << std::setw(28) << r.label << ','
+                   << std::right << std::setw(7) << r.count_fused << '\n';
+            }
+        }
+        LOG_INF("summary: %d ok, %d failed\n", n_ok, n_bad);
+        if (!record_path.empty()) {
+            LOG_INF("%s: baseline written to '%s'\n", __func__, record_path.c_str());
+        }
+
+        if (n_bad && !models_dir.empty() && !check_path.empty()) {
+            LOG_WRN("%s: if the fusion counts are expected to change, run with --record to update the baseline:\n"
+                    "\n"
+                    "./bin/test-llama-archs -o %s\n"
+                    "%s --device %s --models %s --record %s\n",
+                    __func__, models_dir.c_str(), argv[0], device_name.c_str(), models_dir.c_str(), check_path.c_str());
+        }
+
+        return n_bad;
+    }
+}

tests/test-llama-archs.cpp

@@ -128,7 +128,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
     } else if (arch == LLM_ARCH_CHAMELEON) {
         n_vocab = 10240;
     } else if (arch == LLM_ARCH_QWEN3TTS) {
-        n_vocab = 4096; // must be >= the hard-coded codec head size (3072)
+        //n_vocab = 4096; // must be >= the hard-coded codec head size (3072)
+        n_vocab = 3072; // TODO: should be 4096, but user code cannot get `n_vocab_out` yet [TAG_LLAMA_N_VOCAB_OUT]
     }
 
     uint32_t n_head_kv = n_head;

tests/test-save-load-state.cpp

@@ -109,7 +109,7 @@ static bool test_seq_rm_isolated(
     for (llama_seq_id seq_id = 0; seq_id < 2; ++seq_id) {
         llama_batch_ptr batch(n_tokens, 0, 1);
         for (size_t i = 0; i < n_tokens; ++i) {
-            common_batch_add(batch.get(), tokens[i], i, { seq_id }, false);
+            common_batch_add(batch.get(), tokens[i], i, { seq_id }, i == n_tokens - 1);
         }
 
         if (llama_decode(ctx.get(), batch.get())) {
@@ -373,7 +373,7 @@ static bool test_seq_cp_scatter(struct llama_model * model, const struct common_
 
     auto decode_one = [&](llama_token tok, int pos, llama_seq_id seq) {
         llama_batch_ptr batch(1, 0, 1);
-        common_batch_add(batch.get(), tok, pos, { seq }, false);
+        common_batch_add(batch.get(), tok, pos, { seq }, true);
         return llama_decode(ctx.get(), batch.get()) == 0;
     };
 

tools/mtmd/CMakeLists.txt

@@ -84,6 +84,13 @@ target_link_libraries     (mtmd PUBLIC ggml llama)
 target_link_libraries     (mtmd PRIVATE Threads::Threads vendor::hash vendor::miniaudio vendor::stb vendor::sheredom)
 target_include_directories(mtmd PUBLIC  .)
 target_compile_features   (mtmd PRIVATE cxx_std_17)
+target_precompile_headers (mtmd PRIVATE models/models.h)
+
+set_source_files_properties(
+    mtmd-helper.cpp
+    mtmd-helper-gen.cpp
+    PROPERTIES SKIP_PRECOMPILE_HEADERS ON
+)
 
 if (MTMD_VIDEO)
     target_compile_definitions(mtmd PRIVATE MTMD_VIDEO)

tools/server/CMakeLists.txt

@@ -32,6 +32,7 @@ endif()
 target_include_directories(${TARGET} PRIVATE ../mtmd)
 target_include_directories(${TARGET} PRIVATE ${CMAKE_SOURCE_DIR})
 target_link_libraries(${TARGET} PUBLIC llama-common mtmd ${CMAKE_THREAD_LIBS_INIT})
+target_precompile_headers(${TARGET} PRIVATE ${CMAKE_SOURCE_DIR}/common/common.h)
 
 # llama-server-impl: server logic, reusable by app
 
@@ -49,6 +50,7 @@ set_target_properties(${TARGET} PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS ON)
 target_include_directories(${TARGET} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 target_include_directories(${TARGET} PRIVATE ../mtmd ${CMAKE_SOURCE_DIR})
 target_link_libraries(${TARGET} PUBLIC server-context llama-ui cpp-httplib ${CMAKE_THREAD_LIBS_INIT})
+target_precompile_headers(${TARGET} PRIVATE ${CMAKE_SOURCE_DIR}/common/common.h)
 
 add_dependencies(${TARGET} llama-ui-assets)
 

tools/server/server-context.cpp

@@ -3028,7 +3028,7 @@ private:
                         common_speculative_get_draft_params(spec.get(), slot.id) = {
                             /* .drafting = */ true,
                             /* .n_max    = */ n_draft_max,
-                            /* .n_past   = */ slot.prompt.n_tokens(),
+                            /* .pos0     = */ slot.prompt.tokens.pos_next(),
                             /* .id_last  = */ slot.sampled,
                             /* .prompt   = */ &slot.spec_prompt,
                             /* .result   = */ &slot.spec_draft,

tools/server/tests/unit/test_completion.py

@@ -394,7 +394,12 @@ def test_completion_unified(n_ctx, n_slots, n_predict_vals, expected_success):
     results = parallel_function_calls(tasks)
     for res, n_predict, expect_ok in zip(results, n_predict_vals, expected_success):
         if expect_ok:
-            assert res.status_code == 200
+            # the pool is aborted as a whole, so a request that fits on its own
+            # is still dropped when the slots overlap, and it says so explicitly
+            assert res.status_code == 200 or (
+                res.status_code == 500
+                and "context size has been exceeded" in res.body["error"]["message"].lower()
+            )
 
         # note: https://github.com/ggml-org/llama.cpp/pull/18700#issuecomment-3728695581
         if res.status_code == 200: