HIP: mmq: add use_typical_moe_ncols to mmq-config-pascal-older

Rows for M2 Max (30 GPU cores) collected with 'ggml-metal-tuning fa-vec
--dtype f16,q8_0', pasted into fa_vec_tuned_table.

ref: https://github.com/ggml-org/llama.cpp/discussions/27668#discussioncomment-18205786

Assisted-by: pi:llama.cpp/Qwen3.8-27B

.github/workflows/build-android.yml

@@ -110,7 +110,7 @@ jobs:
       #   cache on:  https://github.com/ggerganov/tmp2/actions/runs/26534713799/job/78224189394
       #
       #- name: ccache
-      #  uses: ggml-org/ccache-action@v1.2.21
+      #  uses: ggml-org/ccache-action@v1.2.24
       #  with:
       #    key: android-ubuntu-arm64
       #    evict-old-files: 1d

.github/workflows/build-apple.yml

@@ -47,7 +47,7 @@ jobs:
         uses: actions/checkout@v6
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: apple-arm64
           evict-old-files: 1d
@@ -93,7 +93,7 @@ jobs:
         uses: actions/checkout@v6
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: apple-x64
           evict-old-files: 1d

.github/workflows/build-cpu.yml

@@ -62,7 +62,7 @@ jobs:
         uses: actions/checkout@v6
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: cpu-${{ matrix.os }}
           evict-old-files: 1d
@@ -156,7 +156,7 @@ jobs:
         uses: actions/checkout@v6
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: cpu-windows-2025-${{ matrix.build }}
           variant: ccache

.github/workflows/build-cuda-ubuntu.yml

@@ -53,7 +53,7 @@ jobs:
           apt install -y cmake build-essential ninja-build libgomp1 git libssl-dev jq python3 python3-venv python3-pip
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: cuda-ubuntu-24.04-cuda
           save: false
@@ -108,7 +108,7 @@ jobs:
           sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev libssl-dev rocwmma-dev jq python3-venv
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: cuda-ubuntu-22.04-hip
           save: false
@@ -159,7 +159,7 @@ jobs:
           apt-get install -y build-essential git cmake libssl-dev jq
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: cuda-ubuntu-22.04-musa
           save: false

.github/workflows/build-cuda-windows.yml

@@ -47,7 +47,7 @@ jobs:
         uses: actions/checkout@v6
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-windows-2022-${{ matrix.arch }}-cuda-${{ matrix.cuda }}
 
@@ -152,7 +152,7 @@ jobs:
           & "${env:HIP_PATH}\lib\llvm\bin\clang.exe" --version
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           # TODO: this build does not match the build in release.yml, so we use a different cache key
           #       ideally, the builds should match, similar to the CUDA build above so that we would be able

.github/workflows/build-msys.yml

@@ -35,7 +35,7 @@ jobs:
         uses: actions/checkout@v6
 
       #- name: ccache
-      #  uses: ggml-org/ccache-action@v1.2.16
+      #  uses: ggml-org/ccache-action@v1.2.24
       #  with:
       #    key: msys-windows-2025-x64
       #    variant: ccache

.github/workflows/build-opencl.yml

@@ -44,7 +44,7 @@ jobs:
         uses: actions/checkout@v6
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: opencl-windows-2025-x64
           variant: ccache

.github/workflows/build-openvino.yml

@@ -105,7 +105,7 @@ jobs:
         uses: actions/checkout@v6
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: openvino-windows-2022
           variant: ccache

.github/workflows/build-riscv.yml

@@ -67,7 +67,7 @@ jobs:
 
       # note: sparing some ccache since these jobs run on dedicated runners that are not part of the organitzation
       #- name: ccache
-      #  uses: ggml-org/ccache-action@afde29e5b5422e5da23cb1f639e8baecadeadfc3 # https://github.com/ggml-org/ccache-action/pull/1
+      #  uses: ggml-org/ccache-action@v1.2.24
       #  with:
       #    key: riscv-ubuntu-native
       #    evict-old-files: 1d
@@ -137,7 +137,7 @@ jobs:
 
       # note: sparing some ccache since these jobs run on dedicated runners that are not part of the organitzation
       #- name: ccache
-      #  uses: ggml-org/ccache-action@afde29e5b5422e5da23cb1f639e8baecadeadfc3 # https://github.com/ggml-org/ccache-action/pull/1
+      #  uses: ggml-org/ccache-action@v1.2.24
       #  with:
       #    key: riscv-ubuntu-native-sanitizer-${{ matrix.sanitizer }}-${{ matrix.build_type }}
       #    evict-old-files: 1d

.github/workflows/build-sanitize.yml

@@ -55,7 +55,7 @@ jobs:
         uses: actions/checkout@v6
 
       # - name: ccache
-      #   uses: ggml-org/ccache-action@v1.2.21
+      #   uses: ggml-org/ccache-action@v1.2.24
       #   if: ${{ matrix.sanitizer != 'UNDEFINED' }}
       #   with:
       #     key: ctest-${{ matrix.sanitizer }}-ubuntu-24.04

.github/workflows/build-sycl.yml

@@ -75,7 +75,7 @@ jobs:
           sudo apt-get install -y ./level-zero.deb ./level-zero-devel.deb
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: sycl-ubuntu-24-${{ matrix.build }}
           evict-old-files: 1d
@@ -137,7 +137,7 @@ jobs:
           "LEVEL_ZERO_V1_SDK_PATH=C:/level-zero-sdk" | Out-File -FilePath $env:GITHUB_ENV -Append
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: sycl-windows-latest
           variant: ccache

.github/workflows/build-vulkan.yml

@@ -53,7 +53,7 @@ jobs:
           echo "CXX=g++-14" >> "$GITHUB_ENV"
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: vulkan-ubuntu-24.04-arm
           variant: ccache
@@ -112,7 +112,7 @@ jobs:
           strip: 1
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: vulkan-ubuntu-24.04-llvmpipe
           evict-old-files: 1d
@@ -160,7 +160,7 @@ jobs:
         uses: actions/checkout@v6
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: cpu-windows-2025-x64-vulkan
           variant: ccache

.github/workflows/build-wasm.yml

@@ -54,7 +54,7 @@ jobs:
         uses: actions/checkout@v6
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: webgpu-ubuntu-24.04-arm-wasm
           evict-old-files: 1d

.github/workflows/build-webgpu.yml

@@ -69,7 +69,7 @@ jobs:
         uses: actions/checkout@v6
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: webgpu-macos-latest
           evict-old-files: 1d
@@ -120,7 +120,7 @@ jobs:
         uses: actions/checkout@v6
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: webgpu-ubuntu-24.04
           evict-old-files: 1d

.github/workflows/copilot-setup-steps.yml

@@ -29,7 +29,7 @@ jobs:
         uses: actions/checkout@v6
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: copilot-setup-steps
           evict-old-files: 1d

.github/workflows/hip-quality-check.yml

@@ -52,7 +52,7 @@ jobs:
           sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev libssl-dev python3
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: hip-quality-check-ubuntu-22.04
           evict-old-files: 1d

.github/workflows/release.yml

@@ -103,7 +103,7 @@ jobs:
           path: tools/ui/dist
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-${{ matrix.os }}-${{ matrix.arch }}
 
@@ -187,7 +187,7 @@ jobs:
 
       - name: ccache
         if: ${{ matrix.build != 's390x' }}
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-${{ matrix.os }}-cpu
 
@@ -272,7 +272,7 @@ jobs:
           fi
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-${{ matrix.os }}-vulkan
 
@@ -358,7 +358,7 @@ jobs:
       #   cache on:  https://github.com/ggerganov/tmp2/actions/runs/26534713799/job/78224189394
       #
       #- name: ccache
-      #  uses: ggml-org/ccache-action@v1.2.21
+      #  uses: ggml-org/ccache-action@v1.2.24
       #  with:
       #    key: release-android-arm64
 
@@ -436,7 +436,7 @@ jobs:
           path: tools/ui/dist
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-ubuntu-24.04-openvino-release-no-preset-v1
 
@@ -551,7 +551,7 @@ jobs:
           path: tools/ui/dist
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-windows-2022-openvino
           variant: ccache
@@ -679,7 +679,7 @@ jobs:
           choco install ninja
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-windows-2025-vs2026-${{ matrix.arch }}-cpu
 
@@ -741,7 +741,7 @@ jobs:
           choco install ninja
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
           evict-old-files: 1d
@@ -923,7 +923,7 @@ jobs:
 
       # TODO: these jobs need to use llvm toolchain in order to utilize the ccache
       #- name: ccache
-      #  uses: ggml-org/ccache-action@v1.2.21
+      #  uses: ggml-org/ccache-action@v1.2.24
       #  with:
       #    key: release-windows-2025-${{ matrix.arch }}-${{ matrix.backend }}
 
@@ -1011,7 +1011,7 @@ jobs:
           choco install ninja
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-windows-2022-${{ matrix.arch }}-cuda-${{ matrix.cuda }}
 
@@ -1107,7 +1107,7 @@ jobs:
           "LEVEL_ZERO_V1_SDK_PATH=C:/level-zero-sdk" | Out-File -FilePath $env:GITHUB_ENV -Append
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-windows-2022-x64-sycl
 
@@ -1225,7 +1225,7 @@ jobs:
           path: tools/ui/dist
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-ubuntu-24.04-sycl-${{ matrix.build }}
 
@@ -1302,7 +1302,7 @@ jobs:
           tool-cache: true
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: release-ubuntu-24.04-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
           evict-old-files: 1d

.github/workflows/server.yml

@@ -80,7 +80,7 @@ jobs:
           ref: ${{ github.event.inputs.sha || github.event.pull_request.head.sha || github.sha || github.head_ref || github.ref_name }}
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: server-ubuntu-24.04-arm
           evict-old-files: 1d
@@ -150,7 +150,7 @@ jobs:
           ref: ${{ github.event.inputs.sha || github.event.pull_request.head.sha || github.sha || github.head_ref || github.ref_name }}
 
       - name: ccache
-        uses: ggml-org/ccache-action@v1.2.21
+        uses: ggml-org/ccache-action@v1.2.24
         with:
           key: server-windows-2025-x64
           evict-old-files: 1d

CODEOWNERS

@@ -57,6 +57,7 @@
 /ggml/src/ggml-cann/                    @ggml-org/ggml-cann
 /ggml/src/ggml-common.h                 @ggerganov
 /ggml/src/ggml-cpu/                     @ggerganov
+/ggml/src/ggml-cpu/iqp.*                @bartowski1182
 /ggml/src/ggml-cpu/spacemit/            @alex-spacemit
 /ggml/src/ggml-cuda/                    @ggml-org/ggml-cuda
 /ggml/src/ggml-cuda/vendors/hip.h       @IMbackK

ci/run.sh

@@ -732,6 +732,11 @@ function gg_check_build_requirements {
         gg_printf 'ctest not found, please install\n'
         exit 1
     fi
+
+    if ! command -v unzip &> /dev/null; then
+        gg_printf 'unzip not found, please install\n'
+        exit 1
+    fi
 }
 
 function gg_run_test_backend_ops_cpu {

common/log.cpp

@@ -438,7 +438,7 @@ void common_log_flush(struct common_log * log) {
     log->resume();
 }
 
-static int common_get_verbosity(enum ggml_log_level level) {
+int common_log_get_verbosity(enum ggml_log_level level) {
     switch (level) {
         case GGML_LOG_LEVEL_DEBUG: return LOG_LEVEL_DEBUG;
         case GGML_LOG_LEVEL_INFO:  return LOG_LEVEL_TRACE;
@@ -452,7 +452,7 @@ static int common_get_verbosity(enum ggml_log_level level) {
 }
 
 void common_log_default_callback(enum ggml_log_level level, const char * text, void * /*user_data*/) {
-    auto verbosity = common_get_verbosity(level);
+    auto verbosity = common_log_get_verbosity(level);
     if (verbosity <= common_log_verbosity_thold) {
         common_log_add(common_log_main(), level, "%s", text);
     }

common/log.h

@@ -43,6 +43,8 @@ int  common_log_get_verbosity_thold(void);
 
 void common_log_set_verbosity_thold(int verbosity); // not thread-safe
 
+int common_log_get_verbosity(enum ggml_log_level level);
+
 void common_log_default_callback(enum ggml_log_level level, const char * text, void * user_data);
 
 // the common_log uses an internal worker thread to print/write log messages

common/speculative.cpp

@@ -941,9 +941,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
     const int32_t * target_layer_ids   = nullptr; // model_dft's extract layer indices
     uint32_t        target_layer_ids_n = 0;
 
-    // scratch buffer for concatenated target features [n_tokens, n_embd_enc]
-    std::vector<float> features_buf;
-
     common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq,
             common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)
         : common_speculative_impl(type, n_seq, params.draft.n_max)
@@ -1011,7 +1008,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
         this->n_max = this->params.n_max;
 
         batch        = llama_batch_init(llama_n_batch(ctx_dft), 0,          n_seq);
-        batch_inject = llama_batch_init(llama_n_batch(ctx_dft), n_embd_dec, n_seq);
+        batch_inject = llama_batch_init(llama_n_ubatch(ctx_dft), n_embd_enc, n_seq);
 
         // embd batches on an M-RoPE draft need 4 position rows per token
         is_mrope = llama_model_rope_type(model_dft) == LLAMA_ROPE_TYPE_MROPE;
@@ -1137,58 +1134,21 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
             for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) {
                 const int32_t n_chunk = std::min(n_ubatch, n_rows - offset);
 
-                // gather this chunk's target features, interleaved by extract layer
-                features_buf.resize((size_t) n_chunk * n_embd_enc);
+                // gather target features per extract layer; the fused decode encodes and
+                // injects them into the K/V cache at the target positions
+                batch_inject.n_tokens = n_chunk;
                 for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
                     const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
                     if (!layer) {
                         GGML_ABORT("DFlash: target layer %d input not extracted.", target_layer_ids[k]);
                     }
                     for (int32_t i = 0; i < n_chunk; ++i) {
-                        float       * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
+                        float       * dst = batch_inject.embd + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
                         const float * src = layer + (size_t) (i_batch_beg[seq_id] + offset + i) * n_embd_tgt;
                         std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));
                     }
                 }
 
-                // fuse extracted features through DFlash encoder
-                // M-RoPE drafts read 4 position rows per token from embd batches, so pass them explicitly
-                std::vector<llama_pos> enc_pos;
-                if (is_mrope) {
-                    enc_pos.resize((size_t) 4 * n_chunk);
-                    for (int32_t i = 0; i < n_chunk; ++i) {
-                        const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i];
-                        enc_pos[0 * n_chunk + i] = p;
-                        enc_pos[1 * n_chunk + i] = p;
-                        enc_pos[2 * n_chunk + i] = p;
-                        enc_pos[3 * n_chunk + i] = 0;
-                    }
-                }
-
-                llama_batch enc_batch = {
-                    /*.n_tokens =*/ n_chunk,
-                    /*.token    =*/ nullptr,
-                    /*.embd     =*/ features_buf.data(),
-                    /*.pos      =*/ is_mrope ? enc_pos.data() : nullptr,
-                    /*.n_seq_id =*/ nullptr,
-                    /*.seq_id   =*/ nullptr,
-                    /*.logits   =*/ nullptr,
-                };
-
-                int32_t rc = llama_encode(ctx_dft, enc_batch);
-                if (rc != 0) {
-                    LOG_ERR("%s: llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
-                            __func__, rc, (int) n_chunk, (int) offset);
-                    return false;
-                }
-
-                const float * inp_g = llama_get_embeddings_nextn(ctx_dft);
-                GGML_ASSERT(inp_g && "DFlash encoder produced no output.");
-
-                // inject the DFlash decoder K/V cache at the tokens' target positions
-                batch_inject.n_tokens = n_chunk;
-                std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float));
-
                 for (int32_t i = 0; i < n_chunk; ++i) {
                     const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i];
                     batch_inject.pos[i] = p;
@@ -1201,7 +1161,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
                     batch_inject.seq_id[i][0] = seq_id;
                     batch_inject.logits[i]    = false;
                 }
-                rc = llama_decode(ctx_dft, batch_inject);
+                const int32_t rc = llama_decode(ctx_dft, batch_inject);
                 if (rc != 0) {
                     LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
                             __func__, rc, (int) n_chunk, (int) offset);

docs/android.md

@@ -53,7 +53,7 @@ To see what it might look like visually, here's an old demo of an interactive se
 https://user-images.githubusercontent.com/271616/225014776-1d567049-ad71-4ef2-b050-55b0b3b9274c.mp4
 
 ## Cross-compile CLI using Android NDK
-It's possible to build `llama.cpp` for Android on your host system via CMake and the Android NDK. If you are interested in this path, ensure you already have an environment prepared to cross-compile programs for Android (i.e., install the Android SDK). Note that, unlike desktop environments, the Android environment ships with a limited set of native libraries, and so only those libraries are available to CMake when building with the Android NDK (see: https://developer.android.com/ndk/guides/stable_apis.)
+It's possible to build `llama.cpp` for Android on your host system via CMake and the Android NDK. If you are interested in this path, ensure you already have an environment prepared to cross-compile programs for Android (i.e., install the Android SDK/NDK and set `ANDROID_NDK` to the NDK root). Note that, unlike desktop environments, the Android environment ships with a limited set of native libraries, and so only those libraries are available to CMake when building with the Android NDK (see: https://developer.android.com/ndk/guides/stable_apis.)
 
 Once you're ready and have cloned `llama.cpp`, invoke the following in the project directory:
 
@@ -62,18 +62,22 @@ $ cmake \
   -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \
   -DANDROID_ABI=arm64-v8a \
   -DANDROID_PLATFORM=android-28 \
-  -DCMAKE_C_FLAGS="-march=armv8.7a" \
-  -DCMAKE_CXX_FLAGS="-march=armv8.7a" \
+  -DGGML_NATIVE=OFF \
   -DGGML_OPENMP=OFF \
   -DGGML_LLAMAFILE=OFF \
+  -DLLAMA_OPENSSL=OFF \
   -B build-android
 ```
 
 Notes:
+  - `GGML_NATIVE=OFF` is required for cross-compilation because the host CPU is not the Android target CPU
   - While later versions of Android NDK ship with OpenMP, it must still be installed by CMake as a dependency, which is not supported at this time
   - `llamafile` does not appear to support Android devices (see: https://github.com/Mozilla-Ocho/llamafile/issues/325)
+  - `LLAMA_OPENSSL=OFF` avoids depending on OpenSSL, which is not part of the Android NDK stable native API set
 
-The above command should configure `llama.cpp` with the most performant options for modern devices. Even if your device is not running `armv8.7a`, `llama.cpp` includes runtime checks for available CPU features it can use.
+The above command configures a portable Android `arm64-v8a` build. Do not add a global `-march` flag unless you intentionally want to raise the baseline instruction set for every compiled source.
+
+For optional KleidiAI acceleration on Android `arm64-v8a`, see the [Arm KleidiAI section in build.md](./build.md#arm-kleidiai).
 
 Feel free to adjust the Android ABI for your target. Once the project is configured:
 

docs/backend/SYCL.md

@@ -795,7 +795,9 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_ENABLE_FLASH_ATTN | 1 (default) or 0| Enable Flash-Attention. It can reduce memory usage. The performance impact depends on the LLM.|
 | GGML_SYCL_ENABLE_OPT | 0 or 1 (default)| Enable optimize features for Intel GPUs. (Recommended to 0 for Intel devices older than Gen 10) |
 | GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. |
-| GGML_SYCL_ENABLE_HOST_PINNED_MEM | 0 or 1 (default) | Enable host pinned memory to speed up copy data from host to device. When disable it, host memory will common malloc() on CPU.|
+| GGML_SYCL_ENABLE_HOST_PINNED_MEM | 0 or 1 (default) | Enable host pinned memory to speed up copy data from host to device. When disable it, host memory will common malloc() on CPU. Disable it when use `--load-model mlock`.|
+| GGML_SYCL_HOST_PINNED_MEM_2G | 0 (default) or 1 | Limit the max memory allocation to be no more than 2GB when enable host pinned memory. USM allocations above 2 GiB take the relaxed/large-allocation path, which serializes H2D copies with compute and prevents copy/compute overlap. It will impact the startup time. Need more test. Depend on `GGML_SYCL_ENABLE_HOST_PINNED_MEM=1`.|
+| GGML_SYCL_GET_MEM_API | 0 (default) or 1  | Set to get memory info (free, total) by Level Zero or SYCL API:<br>0 - Level Zero API: support more GPUs, only run on Level Zero running time. When there is an error, fallback to call SYCL API. Depend on GGML_SYCL_SUPPORT_LEVEL_ZERO_API.<br>1 - SYCL API: legacy, support more running time, it can't get the free size of some GPUs (like Arc770). In such case, return total size for free size.|
 | GGML_SYCL_USE_LEVEL_ZERO_API | 1 (default) or 0 | Use Level Zero API for device memory allocation instead of SYCL. Reduces system RAM usage on Intel dGPUs by avoiding DMA-buf/TTM host memory staging. Requires GGML_SYCL_SUPPORT_LEVEL_ZERO_API=ON at build time. SYCL backend always runs on Level Zero running time even if it's set as OFF (The SYCL api will be usage for memory allocation).|
 | GGML_SYCL_ENABLE_DNN | 0 or 1 (default)| Enable running computations through oneDNN and always use oneMKL. |
 | GGML_SYCL_FA_ONEDNN | 1 (default) or 0 | Enable the oneDNN fused SDPA (flash-attention) path on supported GPUs. Set to 0 to always use the native SYCL flash-attention kernel. |

docs/build.md

@@ -614,30 +614,100 @@ You can test with:
 For detailed information about hardware support, setup instructions, and performance optimization, refer to [llama.cpp for ZenDNN](./backend/ZenDNN.md).
 
 ## Arm® KleidiAI™
-KleidiAI is a library of optimized microkernels for AI workloads, specifically designed for Arm CPUs. These microkernels enhance performance and can be enabled for use by the CPU backend.
+KleidiAI provides optimized Arm CPU microkernels used by the ggml CPU backend. Enabling it at build time makes those kernels available; it does not force every operation to use KleidiAI. At runtime, llama.cpp selects the best compatible CPU kernel from the detected CPU features, tensor type, operation shape, and active backend priority.
+
+Supported targets:
+
+| Platform | Supported ABI / architecture | Notes |
+| --- | --- | --- |
+| Linux | AArch64 / arm64 | Runtime CPU feature detection is automatic. |
+| Android | `arm64-v8a` | Use the Android NDK command below for a portable build. |
+| Apple | arm64 | Runtime CPU feature detection is automatic. Non-streaming SVE vector length is treated as unavailable. |
+| Windows | arm64 | Runtime CPU feature detection is automatic. SMCU count is treated as unknown until a detection path is verified. |
+
+`GGML_CPU_KLEIDIAI=ON` is valid only for AArch64/arm64 builds. Do not enable it for x86, 32-bit Arm, or Android ABIs other than `arm64-v8a`.
+
+### Native AArch64/arm64 build
+
+From the llama.cpp source directory:
 
-To enable KleidiAI, go to the llama.cpp directory and build using CMake
 ```bash
-cmake -B build -DGGML_CPU_KLEIDIAI=ON
+cmake -S . -B build -DGGML_CPU_KLEIDIAI=ON
 cmake --build build --config Release
 ```
-You can verify that KleidiAI is being used by running
+
+### Android arm64-v8a NDK build
+
+Set `ANDROID_NDK` to the Android NDK root, then run the following from the llama.cpp source directory. This command configures a portable Android `arm64-v8a` build with KleidiAI enabled and avoids Android dependencies that are not part of the NDK stable native API set.
+
+```bash
+cmake -S . -B build-android \
+  -DCMAKE_BUILD_TYPE=Release \
+  -DCMAKE_TOOLCHAIN_FILE="$ANDROID_NDK/build/cmake/android.toolchain.cmake" \
+  -DANDROID_ABI=arm64-v8a \
+  -DANDROID_PLATFORM=android-28 \
+  -DGGML_CPU_KLEIDIAI=ON \
+  -DGGML_NATIVE=OFF \
+  -DGGML_OPENMP=OFF \
+  -DGGML_LLAMAFILE=OFF \
+  -DLLAMA_OPENSSL=OFF
+cmake --build build-android --config Release --parallel
+cmake --install build-android --prefix {install-dir} --config Release
+```
+
+Important Android options:
+
+- `GGML_CPU_KLEIDIAI=ON` enables KleidiAI for Android `arm64-v8a`.
+- `GGML_NATIVE=OFF` is required for cross-compilation because the build host CPU is not the Android target CPU.
+- `GGML_OPENMP=OFF` avoids adding an OpenMP runtime dependency to this NDK command-line build.
+- `GGML_LLAMAFILE=OFF` avoids the llamafile backend, which is not supported on Android.
+- `LLAMA_OPENSSL=OFF` avoids depending on OpenSSL, which is not part of the Android NDK stable native API set.
+
+The Android Studio project under `examples/llama.android` enables KleidiAI automatically for `arm64-v8a`. For Android command-line CMake builds on `arm64-v8a`, pass `-DGGML_CPU_KLEIDIAI=ON` explicitly.
+
+Global -march flags such as `-march=armv8.7a` flag are not required for a portable Android `arm64-v8a` build. Global `-march` flags raise the baseline instruction set for generic code. No manual architecture-specific source selection is required; llama.cpp selects compatible KleidiAI kernels at runtime. The KleidiAI libraries internal CMake handles the -march flags for each particular kernel.
+
+### Verifying the build
+
+Run an installed or in-tree binary:
+
 ```bash
 ./build/bin/llama-cli -m PATH_TO_MODEL -p "What is a car?"
 ```
-If KleidiAI is enabled, the output will contain a line similar to:
+
+If KleidiAI is enabled, the output contains a line similar to:
+
 ```
 load_tensors: CPU_KLEIDIAI model buffer size =  3474.00 MiB
 ```
-KleidiAI’s microkernels implement optimized tensor operations using Arm CPU features such as dotprod, int8mm, SVE, and SME. Llama.cpp selects the most efficient kernels at runtime based on detected CPU capabilities.
-On CPUs that support SME, SME microkernels are enabled automatically using runtime detection.
-The environment variable GGML_KLEIDIAI_SME can be used to control SME behavior:
-- Not set: enable SME automatically if supported and detected.
-- 0: disable SME.
-- <n> > 0: enable SME and assume <n> available SME units (override auto detection).
-If SME is not supported by the CPU, SME microkernels are always disabled.
 
-Depending on your build target, other higher priority backends may be enabled by default. To ensure the CPU backend is used, you must disable the higher priority backends either at compile time, e.g. -DGGML_METAL=OFF, or during run-time using the command line option `--device none`.
+This confirms that the model has tensors allocated through the KleidiAI CPU buffer. It does not prove that every operation, or any specific SME-family operation, used a KleidiAI microkernel. Runtime CPU features, tensor type, operation shape, and backend priority still control dispatch.
+
+Depending on the build target, another backend may have higher priority than the CPU backend. To force CPU execution for a run, disable higher priority backends at build time, for example `-DGGML_METAL=OFF`, or use a runtime device option such as `--device none` where supported.
+
+### Runtime dispatch
+
+KleidiAI microkernels use Arm CPU features such as dotprod, i8mm, SVE, and SME/SME2. Build-time configuration makes the kernels available. Runtime dispatch selects a compatible kernel for the detected CPU and operation. Older or lower-feature CPUs fall back automatically to compatible kernels.
+
+KleidiAI accelerates selected `GGML_OP_MUL_MAT` paths for F32 and common quantized formats. Exact coverage depends on the bundled KleidiAI version and the llama.cpp runtime selector, so unsupported tensor types, unsupported operation shapes, or higher priority backends may bypass KleidiAI even when the CPU supports the required Arm feature. This is also why a model may not use SME-family kernels on SME-capable hardware.
+
+The current llama.cpp KleidiAI SVE selector only enables SVE kernels when the runtime SVE vector length is known to be QK8_0 bytes, currently 32 bytes. Linux and Android query this at runtime. Apple reports SVE capability separately from userspace non-streaming SVE availability, so llama.cpp treats the SVE vector length as unknown there. Windows exposes SVE feature presence but not the runtime SVE vector length used by this selector, so that value is also treated as unknown. Windows arm64 also treats SMCU count as unknown until a detection mechanism is verified.
+
+The set of available SME-family kernels depends on the bundled KleidiAI version and the detected CPU capabilities. Production configuration does not require any KleidiAI runtime environment variables.
+
+### Diagnostics and debug overrides
+
+KleidiAI runtime environment variables are diagnostics/debug overrides, not production configuration. Leave them unset for normal use.
+
+`GGML_KLEIDIAI_SME` controls SME-family kernel selection and overrides the maximum number of threads assigned to selected quantized SME-family kernels:
+
+- Not set: use automatic runtime detection.
+- `0`: disable SME-family kernels.
+- `<n> > 0`: enable compatible SME-family kernels and allow up to `<n>` threads for quantized SME-family kernels.
+
+On Windows arm64, use `GGML_KLEIDIAI_SME=<n>` as the temporary diagnostics/debug override for SME thread-cap calibration until automatic SMCU count detection is verified.
+
+If the CPU does not support the required SME-family capability for a bundled kernel, that kernel is disabled regardless of the environment variable.
 
 ## OpenCL
 

ggml/src/ggml-backend.cpp

@@ -2115,6 +2115,7 @@ ggml_backend_t ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched,
 bool ggml_backend_op_alloc_size_may_expand(enum ggml_op op) {
     switch (op) {
         case GGML_OP_FLASH_ATTN_EXT:
+        case GGML_OP_MUL_MAT:
         case GGML_OP_MUL_MAT_ID:
         case GGML_OP_CUMSUM:
         case GGML_OP_ARGSORT:

ggml/src/ggml-common.h

@@ -1131,7 +1131,7 @@ GGML_TABLE_END()
 #define NGRID_IQ1S 2048
 #define IQ1S_DELTA 0.125f
 #define IQ1M_DELTA 0.125f
-#if defined(GGML_COMMON_IMPL_C)
+#if defined(GGML_COMMON_IMPL_C) || defined(GGML_COMMON_IMPL_CPP)
 GGML_TABLE_BEGIN(uint64_t, iq1s_grid, NGRID_IQ1S)
     0xffffffffffffffff, 0xffffffffffffff01, 0xffffffffffff0000, 0xffffffffffff01ff,
     0xffffffffffff0101, 0xffffffffff00ff00, 0xffffffffff000000, 0xffffffffff01ffff,

ggml/src/ggml-cpu/CMakeLists.txt

@@ -31,6 +31,8 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
         ggml-cpu/ggml-cpu.cpp
         ggml-cpu/repack.cpp
         ggml-cpu/repack.h
+        ggml-cpu/iqp.cpp
+        ggml-cpu/iqp.h
         ggml-cpu/hbm.cpp
         ggml-cpu/hbm.h
         ggml-cpu/quants.c

ggml/src/ggml-cpu/ggml-cpu.c

@@ -4,6 +4,7 @@
 #include "ggml-backend-impl.h"
 #include "ggml-backend.h"
 #include "traits.h"
+#include "iqp.h"
 #include "ggml-cpu-impl.h"
 #include "ggml-impl.h"
 #include "quants.h"
@@ -1363,6 +1364,13 @@ UseGgmlGemm1:;
 
     ggml_barrier(params->threadpool);
 
+    // IQ panel gemm (see iqp.h) - must come after the barrier above, it consumes the q8_K rows
+    // of src1 from the work buffer
+    if (ggml_cpu_iqp_supports_mul_mat(dst) && !params->use_ref) {
+        ggml_compute_forward_mul_mat_iqp(params, dst);
+        return;
+    }
+
 #if GGML_USE_LLAMAFILE
     if (src1->type != vec_dot_type) {
         const void* wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata;
@@ -1580,6 +1588,16 @@ static void ggml_compute_forward_mul_mat_id(
     char (*atomic_current_chunk)[CACHE_LINE_SIZE] = // [n_as]
         incr_ptr_aligned(&wdata_cur, CACHE_LINE_SIZE * n_as, CACHE_LINE_SIZE);
 
+    // IQ panel gemm (see iqp.h); per expert eligibility is decided below, but the work buffer is
+    // reserved for the whole node (ggml_graph_plan sizes it without params, use_ref only skips the dispatch)
+    const bool iqp = ggml_cpu_iqp_supports_mul_mat_id(dst) && !params->use_ref;
+
+    char * iqp_panels = NULL;
+
+    if (iqp) {
+        iqp_panels = incr_ptr_aligned(&wdata_cur, nth * ggml_cpu_iqp_scratch_size(dst), 64);
+    }
+
     GGML_ASSERT(params->wsize >= (size_t)((char *) wdata_cur - (char *) params->wdata));
 
     if (src1->type != vec_dot_type) {
@@ -1651,6 +1669,13 @@ static void ggml_compute_forward_mul_mat_id(
             continue;
         }
 
+        if (iqp && ggml_cpu_iqp_mul_mat_id_min_batch(cne1)) {
+            ggml_compute_forward_mul_mat_id_iqp(params, dst, cur_a, cne1, (const int32_t *) &MMID_MATRIX_ROW(cur_a, 0),
+                                                iqp_panels);
+
+            continue;
+        }
+
         const char * src0_cur = (const char *) src0->data + cur_a * nb02;
         const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata;
         const size_t row_size = ggml_row_size(vec_dot_type, ne10);
@@ -2858,6 +2883,11 @@ struct ggml_cplan ggml_graph_plan(
                         if (node->src[1]->type != vec_dot_type) {
                             cur = ggml_row_size(vec_dot_type, ggml_nelements(node->src[1]));
                         }
+
+                        // the IQ panel path needs one scratch panel per thread past the q8_K rows
+                        if (ggml_cpu_iqp_supports_mul_mat(node)) {
+                            cur = GGML_PAD(cur, 64) + n_tasks * ggml_cpu_iqp_scratch_size(node);
+                        }
                     } break;
                 case GGML_OP_MUL_MAT_ID:
                     {
@@ -2877,6 +2907,10 @@ struct ggml_cplan ggml_graph_plan(
                         cur += n_as*ids->ne[0]*ids->ne[1]*sizeof(struct mmid_row_mapping) + sizeof(int64_t);
                         // atomic_current_chunk
                         cur += CACHE_LINE_SIZE*n_as + CACHE_LINE_SIZE;
+                        // the IQ panel path needs one scratch panel per thread on top of that
+                        if (ggml_cpu_iqp_supports_mul_mat_id(node)) {
+                            cur += n_tasks * ggml_cpu_iqp_scratch_size(node) + 64;
+                        }
                     } break;
                 case GGML_OP_OUT_PROD:
                     {

ggml/src/ggml-cpu/iqp.h

@@ -0,0 +1,39 @@
+#pragma once
+
+#include "ggml-cpu-impl.h"
+#include "ggml.h"
+
+// GGML internal header
+
+// batched mul_mat path for the grid based IQ types: decode 8 src0 rows at a time into per thread scratch
+// (block_iqp_x8, see iqp.cpp) and run an integer gemm over them against all src1 columns
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// whether cne1 rows of src1 are enough for the decode to pay for itself, per expert, for MUL_MAT_ID
+bool ggml_cpu_iqp_mul_mat_id_min_batch(int64_t cne1);
+
+bool ggml_cpu_iqp_supports_mul_mat(const struct ggml_tensor * dst);
+
+// node level test only - per expert eligibility is decided with ggml_cpu_iqp_mul_mat_id_min_batch
+bool ggml_cpu_iqp_supports_mul_mat_id(const struct ggml_tensor * dst);
+
+// per thread panel scratch bytes, padded
+size_t ggml_cpu_iqp_scratch_size(const struct ggml_tensor * dst);
+
+// must be called after src1 has been converted to q8_K into params->wdata and the threads have synchronized on it
+void ggml_compute_forward_mul_mat_iqp(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+
+// one expert: expert_rows points at its row of the matrix_rows table of (i1, i2) int32 pairs, panels at the base of the per thread panel scratches
+void ggml_compute_forward_mul_mat_id_iqp(const struct ggml_compute_params * params,
+                                         struct ggml_tensor *               dst,
+                                         int64_t                            cur_a,
+                                         int64_t                            cne1,
+                                         const int32_t *                    expert_rows,
+                                         void *                             panels);
+
+#ifdef __cplusplus
+}
+#endif

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

@@ -2,6 +2,7 @@
 #include "cp-async.cuh"
 #include "mma.cuh"
 #include "fattn-common.cuh"
+#include "fattn-swizzle.cuh"
 
 using namespace ggml_cuda_mma;
 
@@ -66,7 +67,7 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 32, 128, 2,  32,  96,  64,  64, 2, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 64, 128, 2,  32,  96,  64,  64, 2, true);
 
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8,  64, 4,  64, 128, 128, 128, 2, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8, 128, 2,  64, 128, 128, 128, 2, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16,  64, 4,  32, 128, 128, 128, 2, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 128, 2,  32, 128, 128, 128, 2, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2,  32, 128, 128, 128, 2, true);
@@ -360,7 +361,7 @@ static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, c
 
 // ------------------------------------------------------------------------------------------------------------------
 
-template<int stride_tile, int nwarps, int nbatch_fa, bool use_cp_async, bool oob_check>
+template<int stride_tile, bool swz, int nwarps, int nbatch_fa, bool use_cp_async, bool oob_check>
 static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
         const half2 * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int D2, const int stride_KV, const int i_sup) {
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
@@ -397,7 +398,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
                 for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) {
                     const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k);
 
-                    cp_async_cg_16<preload>(tile_KV_32 + i*(stride_tile*sizeof(half2)) + k*16, KV + i*stride_KV + k*h2_per_chunk);
+                    if constexpr (swz) {
+                        const int smem_offs_b = ggml_cuda_fattn_smem_swizzle::bytes_rc<stride_tile>(i, k*h2_per_chunk);
+                        cp_async_cg_16<preload>(tile_KV_32 + smem_offs_b, KV + i*stride_KV + k*h2_per_chunk);
+                    } else {
+                        cp_async_cg_16<preload>(tile_KV_32 + i*(stride_tile*sizeof(half2)) + k*16, KV + i*stride_KV + k*h2_per_chunk);
+                    }
                 }
             }
         };
@@ -432,8 +438,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
                 for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) {
                     const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k);
 
-                    ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4,
-                        !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero);
+                    if constexpr (swz) {
+                        ggml_cuda_memcpy_1<16>((char *) tile_KV + ggml_cuda_fattn_smem_swizzle::bytes_rc<stride_tile>(i, k*h2_per_chunk),
+                            !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero);
+                    } else {
+                        ggml_cuda_memcpy_1<16>(tile_KV + i*stride_tile + k*4,
+                            !oob_check || i < i_sup ? KV + i*stride_KV + k*h2_per_chunk : zero);
+                    }
                 }
             }
         };
@@ -568,9 +579,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
     constexpr bool Q_in_reg        = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols);
     constexpr int  nstages         = ggml_cuda_fattn_mma_get_nstages  (DKQ, DV, ncols1, ncols2);
 
-    constexpr int stride_tile_K = nbatch_K2 + 4;
-
-    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : nbatch_V2 + 4;
+    // swizzle the tile stride for K and V based on the batch size.
+    constexpr int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2);
+    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2);
+    constexpr bool swz_K = ggml_cuda_fattn_smem_swizzle::enabled(nbatch_K2);
+    constexpr bool swz_V = V_is_K_view ? swz_K : ggml_cuda_fattn_smem_swizzle::enabled(nbatch_V2);
 
     const int k_VKQ_0 = kb0 * nbatch_fa;
 #if defined(TURING_MMA_AVAILABLE)
@@ -588,7 +601,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
         constexpr bool use_cp_async = true;
         cp_async_wait_all();
         __syncthreads();
-        flash_attn_ext_f16_load_tile<stride_tile_V, nwarps, nbatch_fa, use_cp_async, oob_check>
+        flash_attn_ext_f16_load_tile<stride_tile_V, swz_V, nwarps, nbatch_fa, use_cp_async, oob_check>
             (V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, nbatch_V2, stride_V, k_VKQ_sup);
     } else {
         constexpr bool use_cp_async = nstages == 1;
@@ -607,7 +620,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
         if constexpr (nstages <= 1) {
             const int k0_diff = k0_stop - k0_start;
             constexpr bool use_cp_async = nstages == 1;
-            flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check>
+            flash_attn_ext_f16_load_tile<stride_tile_K, swz_K, nwarps, nbatch_fa, use_cp_async, oob_check>
                 (K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup);
             if (use_cp_async) {
                 cp_async_wait_all();
@@ -623,7 +636,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
 #pragma unroll
                 for (int k_KQ_0 = k0_start; k_KQ_0 < k0_stop; k_KQ_0 += T_A_KQ::J) {
                     T_A_KQ K_A;
-                    load_ldmatrix(K_A, tile_K + i_KQ_0*stride_tile_K + (k_KQ_0 - k0_start), stride_tile_K);
+                    ggml_cuda_fattn_smem_swizzle::load_ldmatrix<stride_tile_K, swz_K>(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start);
                     if constexpr (cols_per_warp == 8) {
                         mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[k_KQ_0/T_A_KQ::J]);
                     } else {
@@ -649,7 +662,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
                     const int i_KQ_0 = i_KQ_00 + (threadIdx.y % np)*T_A_KQ::I;
 
                     T_A_KQ K_A;
-                    load_ldmatrix(K_A, tile_K + i_KQ_0*stride_tile_K + (k_KQ_0 - k0_start), stride_tile_K);
+                    ggml_cuda_fattn_smem_swizzle::load_ldmatrix<stride_tile_K, swz_K>(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start);
 
                     if constexpr (cols_per_warp == 8) {
                         mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[0]);
@@ -943,7 +956,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
                 flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, use_cp_async, oob_check>
                     (mask_h + k_VKQ_0 + nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01);
             }
-            flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check>
+            flash_attn_ext_f16_load_tile<stride_tile_K, swz_K, nwarps, nbatch_fa, use_cp_async, oob_check>
                 (K_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup);
         }
     }
@@ -959,7 +972,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
             const int i0_diff = i0_stop - i0_start;
             if (!V_is_K_view || i0_stop > 2*nbatch_K2) {
                 constexpr bool use_cp_async = nstages == 1;
-                flash_attn_ext_f16_load_tile<stride_tile_V, nwarps, nbatch_fa, use_cp_async, oob_check>
+                flash_attn_ext_f16_load_tile<stride_tile_V, swz_V, nwarps, nbatch_fa, use_cp_async, oob_check>
                     (V_h2 + int64_t(k_VKQ_0)*stride_V + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_sup);
                 if (use_cp_async) {
                     cp_async_wait_all();
@@ -978,7 +991,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
                 const int k0 = k00 + (threadIdx.y % np)*T_A_VKQ::J;
 
                 T_A_VKQ A; // Transposed in SRAM but not in registers, gets transposed on load.
-                load_ldmatrix_trans(A, tile_V_i + 2*k0*stride_tile_V + (i_VKQ_0 - i0_start)/2, stride_tile_V);
+                ggml_cuda_fattn_smem_swizzle::load_ldmatrix_trans<stride_tile_V, swz_V>(A, tile_V, (int)(tile_V_i - tile_V) + 2*k0*stride_tile_V + (i_VKQ_0 - i0_start)/2);
                 if constexpr (T_B_KQ::I == 8) {
                     mma(VKQ_C[i_VKQ_0/T_A_VKQ::I], A, B[k00/(np*T_A_VKQ::J)]);
                 } else {
@@ -1004,7 +1017,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
                 const int k0 = k00 + (threadIdx.y % np)*T_A_VKQ::I;
 
                 T_A_VKQ A; // Transposed in both SRAM and registers, load normally.
-                load_ldmatrix(A, tile_V_i + k0*stride_tile_V + (i_VKQ_0 - i0_start)/2, stride_tile_V);
+                ggml_cuda_fattn_smem_swizzle::load_ldmatrix<stride_tile_V, swz_V>(A, tile_V, (int)(tile_V_i - tile_V) + k0*stride_tile_V + (i_VKQ_0 - i0_start)/2);
                 mma(VKQ_C[i_VKQ_0/i0_stride], B[k00/(np*T_A_VKQ::I)], A);
             }
         }
@@ -1168,10 +1181,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
     static_assert(nwarps * (cols_per_warp/ncols2) % ncols1 == 0, "bad nwarps");
 
     constexpr int stride_tile_Q = DKQ/2     + 4;
-    constexpr int stride_tile_K = nbatch_K2 + 4;
-
-    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : nbatch_V2 + 4;
+    // swizzle the tile stride for K and V based on the batch size.
+    constexpr int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2);
+    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2);
     constexpr int stride_tile_KV_max = stride_tile_K > stride_tile_V ? stride_tile_K : stride_tile_V;
+    constexpr bool swz_K = ggml_cuda_fattn_smem_swizzle::enabled(nbatch_K2);
+    constexpr bool swz_V = V_is_K_view ? swz_K : ggml_cuda_fattn_smem_swizzle::enabled(nbatch_V2);
 
     extern __shared__ half2 tile_Q[];
     half2 * tile_K    = Q_in_reg              ? tile_Q                             : tile_Q + ncols     * stride_tile_Q;
@@ -1265,7 +1280,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
             flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, use_cp_async, oob_check>
                 (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01);
         }
-        flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check>
+        flash_attn_ext_f16_load_tile<stride_tile_K, swz_K, nwarps, nbatch_fa, use_cp_async, oob_check>
             (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup);
     }
 
@@ -1430,11 +1445,17 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
     constexpr int tile_stride = nbatch_combine + 4;
     static_assert((DV/2) % nbatch_combine == 0, "bad nbatch_combine");
 
+    constexpr bool combine_needs_sync = swz_K || swz_V;
+
     if constexpr (cols_per_warp == 8) {
         const int jc_cwmo = (threadIdx.x % (2*T_C_VKQ::J)) / T_C_VKQ::J; // jc combine write meta offset
         const int jc_cwm = threadIdx.y*(2*T_C_VKQ::J) + 2*T_C_VKQ::get_j(-1) + jc_cwmo; // jc combine write meta
         const float2 KQ_cmr = make_float2(KQ_max[jc_cwmo], KQ_rowsum[jc_cwmo]); // KQ combine max rowsum
 
+        if constexpr (combine_needs_sync) {
+            __syncthreads();
+        }
+
         if (((!needs_fixup && !is_fixup) || np > 1) && threadIdx.x < 2*T_C_VKQ::J) {
             // Use the 16 bytes of padding in each row to store the meta data: KQ max, KQ rowsum, KQ max scale.
             ((float2 *) tile_Q)[jc_cwm*(tile_stride/2) + nbatch_combine/2] = KQ_cmr;
@@ -1471,6 +1492,10 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
         const bool thread_should_write = T_C_KQ::J == 8 || T_C_KQ::get_j(threadIdx.x & 2) < 8;
 #endif // defined(TURING_MMA_AVAILABLE)
 
+        if constexpr (combine_needs_sync) {
+            __syncthreads();
+        }
+
         if (((!needs_fixup && !is_fixup) || np > 1) && thread_should_write) {
             ((float2 *) tile_Q)[jc_cwm*(tile_stride/2) + nbatch_combine/2] = KQ_cmr;
         }
@@ -1914,8 +1939,11 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
 
     constexpr bool V_is_K_view = DKQ == 576; // Guaranteed by the kernel selection logic in fattn.cu
 
-    const size_t nbytes_shared_KV_1stage = nbatch_fa            * std::max(nbatch_K2 + 4,  nbatch_V2 + 4) * sizeof(half2);
-    const size_t nbytes_shared_KV_2stage = nbatch_fa            *         (nbatch_K2 + 4 + nbatch_V2 + 4) * sizeof(half2);
+    // KV tile strides must match flash_attn_ext_f16_iter / _process_tile.
+    const int stride_tile_K = ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_K2, cc);
+    const int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_smem_swizzle::tile_stride(nbatch_V2, cc);
+    const size_t nbytes_shared_KV_1stage = nbatch_fa            * std::max(stride_tile_K,  stride_tile_V) * sizeof(half2);
+    const size_t nbytes_shared_KV_2stage = nbatch_fa            *         (stride_tile_K + stride_tile_V) * sizeof(half2);
     const size_t nbytes_shared_Q         = ncols                * (DKQ/2 + 4)                             * sizeof(half2);
     const size_t nbytes_shared_mask      = ncols1               * (nbatch_fa/2 + 4)                       * sizeof(half2);
     const size_t nbytes_shared_combine   = nwarps*cols_per_warp * (nbatch_combine + 4)                    * sizeof(half2);

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

@@ -0,0 +1,126 @@
+#pragma once
+
+#include "common.cuh"
+#include "mma.cuh"
+
+// XOR swizzle for K/V SMEM tiles to avoid bank conflicts without row padding (Turing+ only).
+// Stride must be a multiple of 32 half2 columns, otherwise we keep +4 row padding.
+
+namespace ggml_cuda_fattn_smem_swizzle {
+
+static __host__ __device__ constexpr bool bank_aligned(const int nbatch_2) {
+    return nbatch_2 >= 32 && nbatch_2 % 32 == 0;
+}
+
+static __device__ constexpr bool enabled(const int nbatch_2) {
+#if defined(TURING_MMA_AVAILABLE)
+    return bank_aligned(nbatch_2);
+#else
+    GGML_UNUSED(nbatch_2);
+    return false;
+#endif // defined(TURING_MMA_AVAILABLE)
+}
+
+static __host__ bool enabled(const int nbatch_2, const int cc) {
+#ifdef GGML_USE_HIP
+    GGML_UNUSED(nbatch_2);
+    GGML_UNUSED(cc);
+    return false;
+#else
+    return turing_mma_available(cc) && bank_aligned(nbatch_2);
+#endif // GGML_USE_HIP
+}
+
+static __device__ constexpr int tile_stride(const int nbatch_2) {
+    return enabled(nbatch_2) ? nbatch_2 : nbatch_2 + 4;
+}
+
+static __host__ int tile_stride(const int nbatch_2, const int cc) {
+    return enabled(nbatch_2, cc) ? nbatch_2 : nbatch_2 + 4;
+}
+
+// Swizzled byte offset for tile element (row, col_h2), same map used for writes and reads.
+template<int stride_h2>
+static __device__ __forceinline__ int bytes_rc(const int row, const int col_h2) {
+    static_assert(bank_aligned(stride_h2), "swizzled tile needs a stride that is a multiple of 32");
+    return ((row * stride_h2 + col_h2) * (int) sizeof(half2)) ^ ((row & 7) << 4);
+}
+
+// ldmatrix.x4 via 64-bit generic pointer.
+static __device__ __forceinline__ void ldmatrix_x4(int * xi, const half2 * addr) {
+#if defined(TURING_MMA_AVAILABLE)
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3])
+        : "l"(addr));
+#else
+    GGML_UNUSED_VARS(xi, addr);
+    NO_DEVICE_CODE;
+#endif // defined(TURING_MMA_AVAILABLE)
+}
+
+static __device__ __forceinline__ void ldmatrix_x4_trans(int * xi, const half2 * addr) {
+#if defined(TURING_MMA_AVAILABLE)
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];"
+        : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3])
+        : "l"(addr));
+#else
+    GGML_UNUSED_VARS(xi, addr);
+    NO_DEVICE_CODE;
+#endif // defined(TURING_MMA_AVAILABLE)
+}
+
+// Per-lane swizzled address for one tile<16, 8, half2> ldmatrix: 16 rows, 4 half2 columns per lane.
+template<int stride_h2>
+static __device__ __forceinline__ const half2 * lane_addr(
+        const half2 * tile_base, const int base_row, const int base_col_h2, const int I, const int J) {
+    static_assert(bank_aligned(stride_h2), "swizzled tile needs a stride that is a multiple of 32");
+    const int lane_row = threadIdx.x % I;
+    const int lane_col = (threadIdx.x / I) * (J / 2);
+    uint32_t byte_off = (uint32_t) ((base_row + lane_row)*stride_h2 + base_col_h2 + lane_col) * (uint32_t) sizeof(half2);
+    byte_off ^= (uint32_t) (((base_row + lane_row) & 7) << 4);
+    return (const half2 *) ((const char *) tile_base + byte_off);
+}
+
+template<int stride_h2, bool swz, typename TileT>
+static __device__ __forceinline__ void load_ldmatrix(
+        TileT & t, const half2 * tile_base, const int base_row, const int base_col_h2) {
+    if constexpr (swz) {
+        static_assert(std::is_same_v<TileT, ggml_cuda_mma::tile<16, 8, half2>>,
+            "the swizzled layout is only supported for tile<16, 8, half2>");
+        ldmatrix_x4((int *) t.x, lane_addr<stride_h2>(tile_base, base_row, base_col_h2, TileT::I, TileT::J));
+    } else {
+        ggml_cuda_mma::load_ldmatrix(t, tile_base + base_row*stride_h2 + base_col_h2, stride_h2);
+    }
+}
+
+template<int stride_h2, bool swz, typename TileT>
+static __device__ __forceinline__ void load_ldmatrix(TileT & t, const half2 * tile_base, const int off_h2) {
+    if constexpr (swz) {
+        load_ldmatrix<stride_h2, swz>(t, tile_base, off_h2 / stride_h2, off_h2 % stride_h2);
+    } else {
+        ggml_cuda_mma::load_ldmatrix(t, tile_base + off_h2, stride_h2);
+    }
+}
+
+template<int stride_h2, bool swz, typename TileT>
+static __device__ __forceinline__ void load_ldmatrix_trans(
+        TileT & t, const half2 * tile_base, const int base_row, const int base_col_h2) {
+    if constexpr (swz) {
+        static_assert(std::is_same_v<TileT, ggml_cuda_mma::tile<16, 8, half2>>,
+            "the swizzled layout is only supported for tile<16, 8, half2>");
+        ldmatrix_x4_trans((int *) t.x, lane_addr<stride_h2>(tile_base, base_row, base_col_h2, TileT::I, TileT::J));
+    } else {
+        ggml_cuda_mma::load_ldmatrix_trans(t, tile_base + base_row*stride_h2 + base_col_h2, stride_h2);
+    }
+}
+
+template<int stride_h2, bool swz, typename TileT>
+static __device__ __forceinline__ void load_ldmatrix_trans(TileT & t, const half2 * tile_base, const int off_h2) {
+    if constexpr (swz) {
+        load_ldmatrix_trans<stride_h2, swz>(t, tile_base, off_h2 / stride_h2, off_h2 % stride_h2);
+    } else {
+        ggml_cuda_mma::load_ldmatrix_trans(t, tile_base + off_h2, stride_h2);
+    }
+}
+
+} // namespace ggml_cuda_fattn_smem_swizzle

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -1807,7 +1807,7 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
         return false;
     }
 
-    if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] != 1) {
+    if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] > get_mmvq_mmid_max_batch(src0->type, cc)) {
         return false;
     }
 
@@ -2983,9 +2983,10 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph,
     };
 
     bool is_ok = true;
-    // exception for topk-moe, as each row is read entirely before writing
-    if (ggml_nrows(cgraph->nodes[node_idx]) == 1 && is_topk_moe) {
-        return true;
+    // one block reads all logits before it writes, so logits may alias the out nodes
+    const ggml_tensor * logits_may_alias = nullptr;
+    if (is_topk_moe && ggml_nrows(cgraph->nodes[node_idx]) <= TOPK_MOE_ROWS_PER_BLOCK) {
+        logits_may_alias = cgraph->nodes[node_idx]->src[0];
     }
 
     for (int i = 0; i < out_count; ++i) {
@@ -2999,7 +3000,7 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph,
             for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
                 const ggml_tensor * src = cgraph->nodes[j]->src[src_idx];
 
-                if (!src || src->op == GGML_OP_NONE) {
+                if (!src || src->op == GGML_OP_NONE || src == logits_may_alias) {
                     continue;
                 }
 
@@ -5272,6 +5273,11 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         case GGML_OP_SUM:
             return ggml_is_contiguous_rows(op->src[0]);
         case GGML_OP_TOP_K:
+#if defined(GGML_USE_HIP) || defined(GGML_CUDA_USE_CUB)
+            return true;
+#else
+            return op->src[0]->ne[0] <= 1024;
+#endif // defined(GGML_USE_HIP) || defined(GGML_CUDA_USE_CUB)
         case GGML_OP_ARGSORT:
 #ifndef GGML_CUDA_USE_CUB
             return op->src[0]->ne[0] <= 1024;

ggml/src/ggml-cuda/mmq-config-ampere.cuh

@@ -1,4 +1,5 @@
 static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_ampere(ggml_type type, int J, bool fallback) {
+    constexpr bool use_typical_moe_ncols = false;
     CASE(GGML_TYPE_Q1_0, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
     CASE(GGML_TYPE_Q1_0, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
     CASE(GGML_TYPE_Q1_0, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
@@ -379,5 +380,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
     CASE(GGML_TYPE_NVFP4, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
     CASE(GGML_TYPE_NVFP4, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
 
-    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
 }

ggml/src/ggml-cuda/mmq-config-blackwell.cuh

@@ -1,4 +1,5 @@
 static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_blackwell(ggml_type type, int J, bool fallback) {
+    constexpr bool use_typical_moe_ncols = false;
     CASE(GGML_TYPE_MXFP4, 256, 1, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
     CASE(GGML_TYPE_MXFP4, 256, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);
     CASE(GGML_TYPE_MXFP4, 256, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_FP4, MMQ_ITER_K_FP4, true, true);

ggml/src/ggml-cuda/mmq-config-cdna.cuh

@@ -1,4 +1,5 @@
 static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_cdna(ggml_type type, int J, bool fallback) {
+    constexpr bool use_typical_moe_ncols = false;
     CASE(GGML_TYPE_Q1_0, 512, 1, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
     CASE(GGML_TYPE_Q1_0, 512, 1, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
     CASE(GGML_TYPE_Q1_0, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
@@ -181,5 +182,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
     CASE(GGML_TYPE_NVFP4, 512, 1, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
     CASE(GGML_TYPE_NVFP4, 512, 1, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, true, false);
 
-    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 512, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
 }

ggml/src/ggml-cuda/mmq-config-pascal-dp4a.cuh

@@ -1,4 +1,5 @@
 static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal_dp4a(ggml_type type, int J, bool fallback) {
+    constexpr bool use_typical_moe_ncols = false;
     CASE(GGML_TYPE_Q1_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
@@ -269,5 +270,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
     CASE(GGML_TYPE_NVFP4, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_NVFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
 
-    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
 }

ggml/src/ggml-cuda/mmq-config-pascal-older.cuh

@@ -1,4 +1,5 @@
 static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal_older(ggml_type type, int J, bool fallback) {
+    constexpr bool use_typical_moe_ncols = false;
     CASE(GGML_TYPE_Q1_0, 256, 2, 64,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 256, 2, 64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 256, 2, 64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
@@ -269,5 +270,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
     CASE(GGML_TYPE_NVFP4, 256, 2, 64,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_NVFP4, 256, 2, 64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
 
-    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, false, true);
 }

ggml/src/ggml-cuda/mmq-config-rdna2.cuh

@@ -1,4 +1,5 @@
 static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna2(ggml_type type, int J, bool fallback) {
+    constexpr bool use_typical_moe_ncols = false;
     CASE(GGML_TYPE_Q1_0, 256, 2, 128,   8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 256, 2, 128,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 256, 2, 128,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
@@ -269,5 +270,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
     CASE(GGML_TYPE_NVFP4, 256, 2, 128,  48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_NVFP4, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
 
-    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
 }

ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh

@@ -1,4 +1,5 @@
 static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna3_5(ggml_type type, int J, bool fallback) {
+    constexpr bool use_typical_moe_ncols = false;
     CASE(GGML_TYPE_Q1_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 256, 2, 128,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
@@ -286,5 +287,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
     CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
 
-    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
 }

ggml/src/ggml-cuda/mmq-config-rdna3.cuh

@@ -1,4 +1,5 @@
 static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna3(ggml_type type, int J, bool fallback) {
+    constexpr bool use_typical_moe_ncols = true;
     CASE(GGML_TYPE_Q1_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
@@ -270,5 +271,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
     CASE(GGML_TYPE_NVFP4, 256, 2, 128,  96, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
 
-    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
 }

ggml/src/ggml-cuda/mmq-config-rdna4.cuh

@@ -1,4 +1,5 @@
 static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_rdna4(ggml_type type, int J, bool fallback) {
+    constexpr bool use_typical_moe_ncols = true;
     CASE(GGML_TYPE_Q1_0, 128, 2,  64,  16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 128, 2,  64,  32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
     CASE(GGML_TYPE_Q1_0, 128, 2,  64,  64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true);
@@ -286,5 +287,5 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
     CASE(GGML_TYPE_NVFP4, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
     CASE(GGML_TYPE_NVFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false);
 
-    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true);
+    return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, use_typical_moe_ncols, false, true);
 }

ggml/src/ggml-cuda/mmq.cuh

@@ -170,12 +170,13 @@ struct ggml_cuda_mmq_config {
     int                       J;           // SRAM tile width in src1->ne[1]/dst->ne[1] direction.
     ggml_cuda_mmq_sram_layout sram_layout; // SRAM tile length in src0->ne[0]/src1->ne[0] direction (physical 32 bit elements).
     int                       K_vram;      // VRAM tile length in src0->ne[0]/src1->ne[0] direction (logical elements).
+    bool                      use_typical_moe_ncols;
     bool                      stream_k;    // Whether or not to use stream-k decomposition.
     bool                      fallback;    // Whether a fallback for out-of-bounds check in src0->ne[1] direction is needed.
 
     constexpr __host__ __device__ ggml_cuda_mmq_config(
-            ggml_type type, int nthreads, int occupancy, int I, int J, ggml_cuda_mmq_sram_layout sram_layout, int K_vram, bool stream_k, bool fallback) :
-        type(type), nthreads(nthreads), occupancy(occupancy), I(I), J(J), sram_layout(sram_layout), K_vram(K_vram), stream_k(stream_k), fallback(fallback) {}
+            ggml_type type, int nthreads, int occupancy, int I, int J, ggml_cuda_mmq_sram_layout sram_layout, int K_vram, bool use_typical_moe_ncols, bool stream_k, bool fallback) :
+        type(type), nthreads(nthreads), occupancy(occupancy), I(I), J(J), sram_layout(sram_layout), K_vram(K_vram), use_typical_moe_ncols(use_typical_moe_ncols), stream_k(stream_k), fallback(fallback) {}
 
     constexpr __device__ int rows_per_warp() const {
 #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
@@ -210,7 +211,7 @@ struct ggml_cuda_mmq_config {
         static_assert((I_)        %  32 == 0,                             "bad I");                                                       \
         static_assert((J_)        %   8 == 0,                             "bad J");                                                       \
         static_assert((K_vram_)   % 256 == 0,                             "bad K_vram");                                                  \
-        return ggml_cuda_mmq_config((type_), (nthreads_), (occupancy_), (I_), (J_), (sram_layout_), (K_vram_), (stream_k_), (fallback_)); \
+        return ggml_cuda_mmq_config((type_), (nthreads_), (occupancy_), (I_), (J_), (sram_layout_), (K_vram_), use_typical_moe_ncols, (stream_k_), (fallback_)); \
     }                                                                                                                                     \
 
 #include "mmq-config-pascal-older.cuh"
@@ -1478,6 +1479,20 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
     const int    cc    = ggml_cuda_info().devices[id].cc;
     const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
 
+    int64_t ncols_picker = args.ncols_max;
+    if (args.expert_bounds != nullptr && args.nchannels_x > 0) {
+        const int J_max = ggml_cuda_mmq_get_J_max(type, fallback, cc, 128);
+        const ggml_cuda_mmq_config config_max = ggml_cuda_mmq_get_config(type, J_max, fallback, cc);
+        if (config_max.use_typical_moe_ncols) {
+            // Use the typical expert width only for tile selection.
+            // The launch grid still uses args.ncols_max.
+            const int64_t ncols_typical = (args.ncols_dst + args.nchannels_x - 1) / args.nchannels_x;
+            if (ncols_typical >= 1 && ncols_typical < J_max && ncols_typical < ncols_picker) {
+                ncols_picker = ncols_typical;
+            }
+        }
+    }
+
     int J_best        = 0;
     int ntiles_J_best = INT_MAX;
 
@@ -1491,7 +1506,7 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
             continue;
         }
 
-        const int ntiles_x = (args.ncols_max + config.J - 1) / config.J;
+        const int ntiles_x = (ncols_picker + config.J - 1) / config.J;
 
         if (ntiles_x < ntiles_J_best) {
             J_best = J;

ggml/src/ggml-cuda/mmvq.cu

@@ -773,10 +773,10 @@ static __global__ void mul_mat_vec_q(
 // Grid: (ceil(nrows_x / c_rows_per_block), nchannels_dst)
 // Block: (warp_size, ncols_dst) - each warp handles one token independently.
 // No shared memory reduction needed since each warp works alone.
-template <ggml_type type, int c_rows_per_block>
+template <ggml_type type, int c_rows_per_block, bool has_fusion = false>
 __launch_bounds__(get_mmvq_mmid_max_batch_for_device<type>()*ggml_cuda_get_physical_warp_size(), 1)
 static __global__ void mul_mat_vec_q_moe(
-        const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr,
+        const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion,
         float * dst_ptr,
         const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x,
         const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst,
@@ -794,6 +794,29 @@ static __global__ void mul_mat_vec_q_moe(
 
     constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type);
 
+    // fuse gate, bias, scales, and glu_op into the up projection
+    bool use_gate = false;
+    const void  * vgate      = nullptr;
+    const float * x_bias     = nullptr;
+    const float * gate_bias  = nullptr;
+    const float * x_scale    = nullptr;
+    const float * gate_scale = nullptr;
+    ggml_glu_op   active_glu = GGML_GLU_OP_SWIGLU;
+    float         glu_limit  = 0.0f;
+
+    if constexpr (has_fusion) {
+        use_gate   = fusion.gate != nullptr;
+        vgate      = fusion.gate;
+        x_bias     = (const float *) fusion.x_bias;
+        gate_bias  = (const float *) fusion.gate_bias;
+        active_glu = fusion.glu_op;
+        glu_limit  = fusion.glu_limit;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            x_scale    = (const float *) fusion.x_scale;
+            gate_scale = (const float *) fusion.gate_scale;
+        }
+    }
+
     const uint32_t token_idx   = threadIdx.y;
     const int      row0        = c_rows_per_block*blockIdx.x;
     const int      blocks_per_row_x = ncols_x / qk;
@@ -814,6 +837,7 @@ static __global__ void mul_mat_vec_q_moe(
 
     // partial sum for each thread
     float tmp[c_rows_per_block] = {0.0f};
+    float tmp_gate[c_rows_per_block] = {0.0f};
 
     for (int kbx = threadIdx.x / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) {
         const int kby = kbx * (qk/QK8_1);
@@ -822,6 +846,11 @@ static __global__ void mul_mat_vec_q_moe(
 #pragma unroll
         for (int i = 0; i < c_rows_per_block; ++i) {
             tmp[i] += vec_dot_q_cuda(vx, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs);
+            if constexpr (has_fusion) {
+                if (use_gate) {
+                    tmp_gate[i] += vec_dot_q_cuda(vgate, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs);
+                }
+            }
         }
     }
 
@@ -831,11 +860,63 @@ static __global__ void mul_mat_vec_q_moe(
 #pragma unroll
     for (int i = 0; i < c_rows_per_block; ++i) {
         tmp[i] = warp_reduce_sum<warp_size>(tmp[i]);
+        if constexpr (has_fusion) {
+            if (use_gate) {
+                tmp_gate[i] = warp_reduce_sum<warp_size>(tmp_gate[i]);
+            }
+        }
     }
 
     // Write results
     if (threadIdx.x < c_rows_per_block && (c_rows_per_block == 1 || uint32_t(row0 + threadIdx.x) < nrows_x)) {
-        dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = tmp[threadIdx.x];
+        float result = tmp[threadIdx.x];
+        if constexpr (has_fusion) {
+            const uint32_t bias_idx = channel_x*stride_channel_dst + row0 + threadIdx.x;
+
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                if (x_scale) {
+                    result *= x_scale[channel_x];
+                }
+            }
+            if (x_bias) {
+                result += x_bias[bias_idx];
+            }
+            if (use_gate) {
+                float gate_value = tmp_gate[threadIdx.x];
+                if constexpr (type == GGML_TYPE_NVFP4) {
+                    if (gate_scale) {
+                        gate_value *= gate_scale[channel_x];
+                    }
+                }
+                if (gate_bias) {
+                    gate_value += gate_bias[bias_idx];
+                }
+                switch (active_glu) {
+                    case GGML_GLU_OP_SWIGLU:
+                        result *= ggml_cuda_op_silu_single(gate_value);
+                        break;
+                    case GGML_GLU_OP_GEGLU:
+                        result *= ggml_cuda_op_gelu_single(gate_value);
+                        break;
+                    case GGML_GLU_OP_SWIGLU_OAI:
+                        result = ggml_cuda_op_swiglu_oai_single(gate_value, result);
+                        break;
+                    case GGML_GLU_OP_SWIGLU_CLAMP:
+                        result = ggml_cuda_op_swiglu_clamp_single(gate_value, result, glu_limit);
+                        break;
+                    default:
+                        result = result * gate_value;
+                        break;
+                }
+            }
+        }
+        dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = result;
+    }
+
+    if constexpr (!has_fusion) {
+        GGML_UNUSED_VARS(use_gate, tmp_gate, vgate, x_bias, gate_bias, active_glu, glu_limit, x_scale, gate_scale);
+    } else if constexpr (type != GGML_TYPE_NVFP4) {
+        GGML_UNUSED_VARS(x_scale, gate_scale);
     }
 }
 
@@ -885,7 +966,7 @@ static void mul_mat_vec_q_switch_fusion(
 
 template <ggml_type type>
 static void mul_mat_vec_q_moe_launch(
-        const void * vx, const void * vy, const int32_t * ids, float * dst,
+        const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
         const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x,
         const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst,
         const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
@@ -898,11 +979,22 @@ static void mul_mat_vec_q_moe_launch(
     const dim3 block_dims(warp_size, ncols_dst);
     const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
 
-    ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block>, launch_params,
-        vx, vy, ids, dst, ncols_x, nchannels_y, nrows_x,
-        stride_row_x, stride_col_y, stride_col_dst,
-        stride_channel_x, stride_channel_y, stride_channel_dst,
-        ncols_dst, ids_stride);
+    const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr ||
+                            fusion.x_scale != nullptr || fusion.gate_scale != nullptr;
+
+    if (has_fusion) {
+        ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block, true>, launch_params,
+            vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x,
+            stride_row_x, stride_col_y, stride_col_dst,
+            stride_channel_x, stride_channel_y, stride_channel_dst,
+            ncols_dst, ids_stride);
+    } else {
+        ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block, false>, launch_params,
+            vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x,
+            stride_row_x, stride_col_y, stride_col_dst,
+            stride_channel_x, stride_channel_y, stride_channel_dst,
+            ncols_dst, ids_stride);
+    }
 }
 
 template <ggml_type type>
@@ -998,7 +1090,7 @@ static void mul_mat_vec_q_switch_ncols_dst(
     if (has_ids && ncols_dst > 1) {
         // Multi-token MUL_MAT_ID path - dedicated MoE kernel
         mul_mat_vec_q_moe_launch<type>(
-            vx, vy, ids, dst, ncols_x, nchannels_y_fd, nrows_x,
+            vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, nrows_x,
             stride_row_x, stride_col_y, stride_col_dst,
             stride_channel_x, stride_channel_y, stride_channel_dst,
             ncols_dst, ids_stride, warp_size, nchannels_dst, stream);
@@ -1280,7 +1372,8 @@ void ggml_cuda_mul_mat_vec_q(
     ggml_cuda_mm_fusion_args_device fusion_local{};
 
     if (fusion) {
-        GGML_ASSERT( !ids || dst->ne[2] == 1);
+        const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+        GGML_ASSERT( !ids || dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc));
         GGML_ASSERT(  ids || dst->ne[1] == 1);
         // Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is
         // non-negligible for some models such as gpt-oss-20b

ggml/src/ggml-cuda/top-k.cu

@@ -48,6 +48,168 @@ static int next_power_of_2(int x) {
 
 #endif                            // CUB_TOP_K_AVAILABLE
 
+#if !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP)
+
+static __device__ __forceinline__ uint32_t top_k_float_to_ordered(float value) {
+    const uint32_t bits = __float_as_uint(value);
+    const uint32_t mask = (uint32_t) (-(int32_t) (bits >> 31)) | 0x80000000U;
+    return bits ^ mask;
+}
+
+struct top_k_radix_state {
+    uint32_t prefix;
+    uint32_t prefix_mask;
+    int rank;
+    int greater_count;
+    int equal_count;
+};
+
+static __global__ void top_k_radix_init(top_k_radix_state * states, int nrows, int k) {
+    const int row = blockIdx.x * blockDim.x + threadIdx.x;
+    if (row < nrows) {
+        states[row] = {0, 0, k, 0, 0};
+    }
+}
+
+template<int BLOCK_SIZE, int RADIX_BITS>
+static __global__ void top_k_radix_histogram(
+        const float * __restrict__ src,
+        const top_k_radix_state * __restrict__ states,
+        int * __restrict__ block_histograms,
+        int ncols,
+        int blocks_per_row,
+        int shift) {
+    constexpr int NBINS = 1 << RADIX_BITS;
+
+    const int row = blockIdx.x / blocks_per_row;
+    const int row_block = blockIdx.x % blocks_per_row;
+    const int tid = threadIdx.x;
+    const float * row_src = src + (size_t) row * ncols;
+    __shared__ int histogram[NBINS];
+
+    histogram[tid] = 0;
+    __syncthreads();
+
+    const top_k_radix_state state = states[row];
+    for (int col = row_block * BLOCK_SIZE + tid;
+         col < ncols;
+         col += blocks_per_row * BLOCK_SIZE) {
+        const uint32_t key = top_k_float_to_ordered(row_src[col]);
+        if ((key & state.prefix_mask) == state.prefix) {
+            atomicAdd(&histogram[(key >> shift) & (NBINS - 1)], 1);
+        }
+    }
+    __syncthreads();
+
+    const size_t histogram_offset =
+        ((size_t) row * blocks_per_row + row_block) * NBINS;
+    block_histograms[histogram_offset + tid] = histogram[tid];
+}
+
+template<int BLOCK_SIZE, int RADIX_BITS>
+static __global__ void top_k_radix_select(
+        const int * __restrict__ block_histograms,
+        top_k_radix_state * __restrict__ states,
+        int blocks_per_row,
+        int shift) {
+    constexpr int NBINS = 1 << RADIX_BITS;
+
+    const int row = blockIdx.x;
+    const int tid = threadIdx.x;
+    __shared__ int histogram[NBINS];
+
+    int count = 0;
+    for (int row_block = 0; row_block < blocks_per_row; ++row_block) {
+        const size_t offset = ((size_t) row * blocks_per_row + row_block) * NBINS;
+        count += block_histograms[offset + tid];
+    }
+    histogram[tid] = count;
+    __syncthreads();
+
+    if (tid == 0) {
+        top_k_radix_state state = states[row];
+        int bin = NBINS - 1;
+        while (bin > 0 && histogram[bin] < state.rank) {
+            state.rank -= histogram[bin--];
+        }
+        state.prefix |= (uint32_t) bin << shift;
+        state.prefix_mask |= (uint32_t) (NBINS - 1) << shift;
+        states[row] = state;
+    }
+}
+
+static __global__ void top_k_radix_reset_counters(top_k_radix_state * states, int nrows) {
+    const int row = blockIdx.x * blockDim.x + threadIdx.x;
+    if (row < nrows) {
+        states[row].greater_count = 0;
+        states[row].equal_count = 0;
+    }
+}
+
+template<int BLOCK_SIZE>
+static __global__ void top_k_radix_gather(
+        const float * __restrict__ src,
+        int * __restrict__ dst,
+        top_k_radix_state * __restrict__ states,
+        int ncols,
+        int k,
+        int blocks_per_row) {
+    const int row = blockIdx.x / blocks_per_row;
+    const int row_block = blockIdx.x % blocks_per_row;
+    const int tid = threadIdx.x;
+    const float * row_src = src + (size_t) row * ncols;
+    int * row_dst = dst + (size_t) row * k;
+    top_k_radix_state * state = &states[row];
+
+    for (int col = row_block * BLOCK_SIZE + tid;
+         col < ncols;
+         col += blocks_per_row * BLOCK_SIZE) {
+        const uint32_t key = top_k_float_to_ordered(row_src[col]);
+        if (key > state->prefix) {
+            const int pos = atomicAdd(&state->greater_count, 1);
+            row_dst[pos] = col;
+        } else if (key == state->prefix) {
+            const int pos = atomicAdd(&state->equal_count, 1);
+            if (pos < state->rank) {
+                row_dst[k - state->rank + pos] = col;
+            }
+        }
+    }
+}
+
+static void top_k_radix_cuda(
+        ggml_cuda_pool & pool,
+        const float * src, int * dst, int ncols, int nrows, int k, cudaStream_t stream) {
+    constexpr int BLOCK_SIZE = 256;
+    constexpr int RADIX_BITS = 8;
+    constexpr int NBINS = 1 << RADIX_BITS;
+    const int blocks_per_row = std::min((ncols + 1023) / 1024, 64);
+
+    ggml_cuda_pool_alloc<top_k_radix_state> states_alloc(pool, nrows);
+    ggml_cuda_pool_alloc<int> histograms_alloc(pool, (size_t) nrows * blocks_per_row * NBINS);
+    top_k_radix_state * states = states_alloc.get();
+    int * histograms = histograms_alloc.get();
+
+    top_k_radix_init<<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows, k);
+
+    const dim3 row_grid(blocks_per_row * nrows);
+    for (int shift = 32 - RADIX_BITS; shift >= 0; shift -= RADIX_BITS) {
+        top_k_radix_histogram<BLOCK_SIZE, RADIX_BITS>
+            <<<row_grid, BLOCK_SIZE, 0, stream>>>(
+                src, states, histograms, ncols, blocks_per_row, shift);
+        top_k_radix_select<BLOCK_SIZE, RADIX_BITS>
+            <<<nrows, BLOCK_SIZE, 0, stream>>>(histograms, states, blocks_per_row, shift);
+    }
+
+    top_k_radix_reset_counters
+        <<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows);
+    top_k_radix_gather<BLOCK_SIZE>
+        <<<row_grid, BLOCK_SIZE, 0, stream>>>(
+            src, dst, states, ncols, k, blocks_per_row);
+}
+
+#endif // !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP)
+
 void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const ggml_tensor * src0   = dst->src[0];
     const float *       src0_d = (const float *) src0->data;
@@ -96,10 +258,18 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
         dst_d  += k     * iter_nrows;
     }
 #else                             // GGML_CUDA_USE_CUB
-    ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows);
-    int *                     tmp_dst = temp_dst_alloc.get();
-    argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream);
-    CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows,
-                                 cudaMemcpyDeviceToDevice, stream));
+#if defined(GGML_USE_HIP)
+    if (ncols > 1024) {
+        top_k_radix_cuda(pool, src0_d, dst_d, ncols, nrows, k, stream);
+    } else {
+#endif // defined(GGML_USE_HIP)
+        ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows);
+        int *                     tmp_dst = temp_dst_alloc.get();
+        argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream);
+        CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows,
+                                     cudaMemcpyDeviceToDevice, stream));
+#if defined(GGML_USE_HIP)
+    }
+#endif // defined(GGML_USE_HIP)
 #endif
 }

ggml/src/ggml-cuda/topk-moe.cu

@@ -88,15 +88,16 @@ __device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], co
     It is intended as fusion of softmax->top-k->get_rows pipeline for MoE models
 */
 template <int n_experts, bool has_bias>
-__launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float *         logits,
-                                                                  float *               weights,
-                                                                  int32_t *             ids,
-                                                                  float *               bias,
-                                                                  const int             n_rows,
-                                                                  const int             n_expert_used,
-                                                                  const float           clamp_val,
-                                                                  const float           scale_val,
-                                                                  const topk_moe_config config) {
+__launch_bounds__(TOPK_MOE_ROWS_PER_BLOCK * WARP_SIZE, 1)
+__global__ void topk_moe_cuda(const float *         logits,
+                              float *               weights,
+                              int32_t *             ids,
+                              float *               bias,
+                              const int             n_rows,
+                              const int             n_expert_used,
+                              const float           clamp_val,
+                              const float           scale_val,
+                              const topk_moe_config config) {
     const int row = blockIdx.x * blockDim.y + threadIdx.y;
     if (row >= n_rows) {
         return;
@@ -123,6 +124,9 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float *
         wt[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? logits[expert] : -INFINITY;
     }
 
+    // Weights and IDs can alias logits, so wait until every row in the block reads its logits.
+    __syncthreads();
+
     if (!config.delayed_softmax) {
         if (config.use_sigmoid) {
            sigmoid_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
@@ -282,7 +286,7 @@ static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx,
                                  const topk_moe_config       config) {
     GGML_ASSERT(!(config.with_norm && config.delayed_softmax) &&
                 "delayed softmax is not supported with weight normalization");
-    const int    rows_per_block = 4;
+    const int    rows_per_block = TOPK_MOE_ROWS_PER_BLOCK;
     dim3         grid_dims((n_rows + rows_per_block - 1) / rows_per_block, 1, 1);
     dim3         block_dims(WARP_SIZE, rows_per_block, 1);
     cudaStream_t stream = ctx.stream();

ggml/src/ggml-cuda/topk-moe.cuh

@@ -3,6 +3,9 @@
 
 #include <initializer_list>
 
+// Rows that one CUDA block handles.
+#define TOPK_MOE_ROWS_PER_BLOCK 8
+
 struct ggml_cuda_topk_moe_args {
     bool sigmoid{};
     bool sqrt_softplus{};

ggml/src/ggml-hexagon/htp/cpy-ops.c

@@ -330,7 +330,7 @@ int op_cpy(struct htp_ops_context * octx) {
     }
 
     const struct htp_tensor *sync = octx->src[1];
-    if (sync) {
+    if (sync && (sync->flags & HTP_TENSOR_FENCE)) {
         if (!use_dma) {
             // htp_tensor_flush_all(octx->ctx, octx->dsts, 1);
             qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);

ggml/src/ggml-metal/CMakeLists.txt

@@ -163,19 +163,37 @@ else()
         )
     endforeach()
 
+    # the tensor API kernels go in a separate metallib, loaded only where supported
+    set(AIR_MM_TENSOR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/mul_mm_tensor.air")
+    add_custom_command(
+        OUTPUT ${AIR_MM_TENSOR}
+        COMMAND xcrun -sdk macosx metal ${XC_FLAGS} -DGGML_METAL_HAS_TENSOR -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/mul_mm.metal -o ${AIR_MM_TENSOR}
+        DEPENDS kernels/mul_mm.metal kernels/common.h kernels/dequantize.h ${METALLIB_COMMON} ggml-metal-impl.h
+        COMMENT "Compiling kernels/mul_mm.metal (tensor API)"
+        VERBATIM
+    )
+
+    add_custom_command(
+        OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-tensor.metallib
+        COMMAND xcrun -sdk macosx metallib ${AIR_MM_TENSOR} -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-tensor.metallib
+        DEPENDS ${AIR_MM_TENSOR}
+        COMMENT "Linking tensor API Metal kernels into ggml-tensor.metallib"
+    )
+
     add_custom_command(
         OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
         COMMAND xcrun -sdk macosx metallib ${AIR_FILES} -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
         COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h
         COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal-impl.h
         COMMAND rm -rf ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels
-        DEPENDS ${AIR_FILES}
+        DEPENDS ${AIR_FILES} ${AIR_MM_TENSOR}
         COMMENT "Linking Metal kernels into default.metallib"
     )
 
     add_custom_target(
         ggml-metal-lib ALL
         DEPENDS ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
+                ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-tensor.metallib
     )
 endif() # GGML_METAL_EMBED_LIBRARY
 
@@ -188,6 +206,7 @@ if (NOT GGML_METAL_EMBED_LIBRARY)
 
     install(
         FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
+              ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-tensor.metallib
         DESTINATION ${CMAKE_INSTALL_BINDIR}
     )
 endif()

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

@@ -1336,7 +1336,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht(ggml_metal_
     return res;
 }
 
-// note: reuse the argsort kernel for top_k
+// note: reuse the argsort kernel for the bitonic top_k fallback
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal_library_t lib, const ggml_tensor * op) {
     assert(op->op == GGML_OP_TOP_K);
 
@@ -1364,6 +1364,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix(ggml_metal_library_t lib, const ggml_tensor * op) {
+    assert(op->op == GGML_OP_TOP_K);
+
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_top_k_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->type));
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+    }
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge(ggml_metal_library_t lib, const ggml_tensor * op) {
     assert(op->op == GGML_OP_TOP_K);
 

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

@@ -145,6 +145,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge     (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht              (ggml_metal_library_t lib, int n);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k             (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix       (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge       (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin               (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse );
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin_one           (ggml_metal_library_t lib, enum ggml_op op);

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

@@ -27,6 +27,9 @@
 static const NSInteger MTLGPUFamilyMetal3_GGML = 5001;
 static const NSInteger MTLGPUFamilyMetal4_GGML = 5002;
 
+// MTLLanguageVersion4_0 is not present in older SDKs
+static const NSUInteger MTLLanguageVersion4_0_GGML = 4 << 16;
+
 #if !GGML_METAL_EMBED_LIBRARY
 // Here to assist with NSBundle Path Hack
 @interface GGMLMetalClass : NSObject
@@ -154,6 +157,9 @@ struct ggml_metal_library {
     // nil in single_library mode (everything resolves to objs[0]).
     NSMutableDictionary<NSString *, NSNumber *> * fn_to_lib;
 
+    // kernels from a second metallib, resolved ahead of the combined library
+    NSSet<NSString *> * override_fns;
+
     ggml_metal_device_t dev;
     ggml_metal_pipelines_t pipelines; // cache of compiled pipelines
 
@@ -174,6 +180,18 @@ static void ggml_metal_library_build_index(ggml_metal_library_t lib) {
     }
 }
 
+// note: defined below, after struct ggml_metal_device
+static void ggml_metal_device_disable_tensor(ggml_metal_device_t dev);
+
+// the tensor API headers are exposed to the shader compiler only at Metal language version 4.0
+static void ggml_metal_compile_options_set_lang(MTLCompileOptions * options, bool has_tensor) {
+    if (!has_tensor) {
+        return;
+    }
+
+    options.languageVersion = (MTLLanguageVersion) MTLLanguageVersion4_0_GGML;
+}
+
 // Parse a `#include "name"` line. Returns the quoted name in *include_name on
 // success. Whitespace-tolerant; ignores `#include <...>` (system headers).
 static bool ggml_metal_library_parse_quoted_include(NSString * line, NSString ** include_name) {
@@ -313,6 +331,7 @@ static bool ggml_metal_library_compile_all(
             @autoreleasepool {
                 MTLCompileOptions * options = [MTLCompileOptions new];
                 options.preprocessorMacros = prep;
+                ggml_metal_compile_options_set_lang(options, ggml_metal_device_get_props(res->dev)->has_tensor);
 
                 lib = [device newLibraryWithSource:src options:options error:&error];
 
@@ -369,6 +388,46 @@ static bool ggml_metal_library_compile_all(
     return ok;
 }
 
+// look for <name>.metallib as a bundle resource, then next to the running binary
+static NSString * ggml_metal_find_metallib(NSBundle * bundle, NSString * name) {
+    NSError * error = nil;
+
+    NSString * path_lib = [bundle pathForResource:name ofType:@"metallib"];
+    if (path_lib == nil) {
+        // Try to find the resource in the directory where the current binary located.
+        NSString * bin_cur = [[NSProcessInfo processInfo] arguments][0];
+        NSString * bin_dir = [bin_cur stringByDeletingLastPathComponent];
+
+        NSString * path_lib_default = [NSString pathWithComponents:@[bin_dir, [name stringByAppendingPathExtension:@"metallib"]]];
+        if ([[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
+            GGML_LOG_INFO("%s: found '%s'\n", __func__, [path_lib_default UTF8String]);
+
+            NSDictionary * atts = [[NSFileManager defaultManager] attributesOfItemAtPath:path_lib_default error:&error];
+            if (atts && atts[NSFileType] == NSFileTypeSymbolicLink) {
+                // Optionally, if this is a symlink, try to resolve it.
+                path_lib_default = [[NSFileManager defaultManager] destinationOfSymbolicLinkAtPath:path_lib_default error:&error];
+                if (path_lib_default && [path_lib_default length] > 0 && ![[path_lib_default substringToIndex:1] isEqualToString:@"/"]) {
+                    // It is a relative path, adding the binary directory as directory prefix.
+                    path_lib_default = [NSString pathWithComponents:@[bin_dir, path_lib_default]];
+                }
+                if (!path_lib_default || ![[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
+                    // Link to the resource could not be resolved.
+                    path_lib_default = nil;
+                } else {
+                    GGML_LOG_INFO("%s: symlink resolved '%s'\n", __func__, [path_lib_default UTF8String]);
+                }
+            }
+        } else {
+            // The resource couldn't be found in the binary's directory.
+            path_lib_default = nil;
+        }
+
+        path_lib = path_lib_default;
+    }
+
+    return path_lib;
+}
+
 ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
     id<MTLDevice> device = ggml_metal_device_get_obj(dev);
 
@@ -432,38 +491,7 @@ ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
     const int64_t t_start = ggml_time_us();
 
     NSError * error = nil;
-    NSString * path_lib = [bundle pathForResource:@"default" ofType:@"metallib"];
-    if (path_lib == nil) {
-        // Try to find the resource in the directory where the current binary located.
-        NSString * bin_cur = [[NSProcessInfo processInfo] arguments][0];
-        NSString * bin_dir = [bin_cur stringByDeletingLastPathComponent];
-
-        NSString * path_lib_default = [NSString pathWithComponents:@[bin_dir, @"default.metallib"]];
-        if ([[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
-            GGML_LOG_INFO("%s: found '%s'\n", __func__, [path_lib_default UTF8String]);
-
-            NSDictionary * atts = [[NSFileManager defaultManager] attributesOfItemAtPath:path_lib_default error:&error];
-            if (atts && atts[NSFileType] == NSFileTypeSymbolicLink) {
-                // Optionally, if this is a symlink, try to resolve it.
-                path_lib_default = [[NSFileManager defaultManager] destinationOfSymbolicLinkAtPath:path_lib_default error:&error];
-                if (path_lib_default && [path_lib_default length] > 0 && ![[path_lib_default substringToIndex:1] isEqualToString:@"/"]) {
-                    // It is a relative path, adding the binary directory as directory prefix.
-                    path_lib_default = [NSString pathWithComponents:@[bin_dir, path_lib_default]];
-                }
-                if (!path_lib_default || ![[NSFileManager defaultManager] isReadableFileAtPath:path_lib_default]) {
-                    // Link to the resource could not be resolved.
-                    path_lib_default = nil;
-                } else {
-                    GGML_LOG_INFO("%s: symlink resolved '%s'\n", __func__, [path_lib_default UTF8String]);
-                }
-            }
-        } else {
-            // The resource couldn't be found in the binary's directory.
-            path_lib_default = nil;
-        }
-
-        path_lib = path_lib_default;
-    }
+    NSString * path_lib = ggml_metal_find_metallib(bundle, @"default");
 
     if (path_lib != nil) {
         // pre-compiled library found: a single combined default.metallib
@@ -478,6 +506,30 @@ ggml_metal_library_t ggml_metal_library_init(ggml_metal_device_t dev) {
             return NULL;
         }
 
+        // the tensor API kernels are built into a separate metallib
+        if (ggml_metal_device_get_props(dev)->has_tensor) {
+            NSString * path_mm = ggml_metal_find_metallib(bundle, @"ggml-tensor");
+
+            id<MTLLibrary> lib_mm = nil;
+            if (path_mm != nil) {
+                lib_mm = [device newLibraryWithURL:[NSURL fileURLWithPath:path_mm] error:&error];
+                if (!lib_mm && error) {
+                    GGML_LOG_ERROR("%s: %s\n", __func__, [[error description] UTF8String]);
+                }
+            }
+
+            if (lib_mm) {
+                GGML_LOG_INFO("%s: loaded '%s'\n", __func__, [path_mm UTF8String]);
+
+                res->objs[GGML_METAL_LIB_MUL_MM] = [lib_mm retain];
+                res->override_fns                = [[NSSet setWithArray:[lib_mm functionNames]] retain];
+            } else {
+                GGML_LOG_INFO("%s: ggml-tensor.metallib not found - disabling the tensor API\n", __func__);
+
+                ggml_metal_device_disable_tensor(dev);
+            }
+        }
+
         GGML_LOG_INFO("%s: loaded in %.3f sec\n", __func__, (ggml_time_us() - t_start) / 1e6);
         return res;
     }
@@ -557,6 +609,7 @@ ggml_metal_library_t ggml_metal_library_init_from_source(ggml_metal_device_t dev
 
         MTLCompileOptions * options = [MTLCompileOptions new];
         options.preprocessorMacros = prep;
+        ggml_metal_compile_options_set_lang(options, ggml_metal_device_get_props(dev)->has_tensor);
 
         library = [device newLibraryWithSource:src options:options error:&error];
         if (error) {
@@ -615,6 +668,10 @@ void ggml_metal_library_free(ggml_metal_library_t lib) {
         [lib->fn_to_lib release];
     }
 
+    if (lib->override_fns) {
+        [lib->override_fns release];
+    }
+
     ggml_metal_pipelines_free(lib->pipelines);
 
     [lib->lock release];
@@ -676,7 +733,9 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_compile_pipeline(ggml_
         // route to the library that actually defines this kernel; fn_to_lib is
         // built from -[MTLLibrary functionNames] so it's always in sync
         int lib_idx = 0;
-        if (!lib->single_library) {
+        if (lib->override_fns && [lib->override_fns containsObject:base_func]) {
+            lib_idx = GGML_METAL_LIB_MUL_MM;
+        } else if (!lib->single_library) {
             NSNumber * idx = lib->fn_to_lib[base_func];
             if (!idx) {
                 [lib->lock unlock];
@@ -1538,6 +1597,12 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                         return true;
                     case GGML_TYPE_BF16:
                         return has_bfloat;
+                    case GGML_TYPE_Q4_0:
+                    case GGML_TYPE_Q4_1:
+                    case GGML_TYPE_Q5_0:
+                    case GGML_TYPE_Q5_1:
+                    case GGML_TYPE_Q8_0:
+                        return true;
                     default:
                         return false;
                 }
@@ -1856,6 +1921,10 @@ const struct ggml_metal_device_props * ggml_metal_device_get_props(ggml_metal_de
     return &dev->props;
 }
 
+static void ggml_metal_device_disable_tensor(ggml_metal_device_t dev) {
+    dev->props.has_tensor = false;
+}
+
 //
 // device buffers
 //

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

@@ -1189,6 +1189,17 @@ typedef struct {
     int32_t  len;
 } ggml_metal_kargs_argsort_merge;
 
+typedef struct {
+    int32_t  ne00;   // number of columns (elements per row)
+    int32_t  ne01;   // rows
+    int32_t  ne02;
+    int32_t  ne03;
+    uint64_t nb01;   // row stride in src0
+    uint64_t nb02;
+    uint64_t nb03;
+    int32_t  top_k;  // k
+} ggml_metal_kargs_top_k;
+
 typedef struct {
     int32_t nrows;
 } ggml_metal_kargs_fwht;

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

@@ -552,8 +552,24 @@ int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) {
 
     const int32_t dim = ((const int32_t *) op->op_params)[0];
 
+    const bool is_q = ggml_is_quantized(op->type);
+
+    // for quantized types, concat is done at the block level (nb0 == type_size == block size)
+    int32_t ne00_arg = ne00;
+    int32_t ne10_arg = ne10;
+    int32_t ne0_arg  = ne0;
+    if (is_q) {
+        const int32_t blck = ggml_blck_size(op->type);
+        GGML_ASSERT(ne00 % blck == 0);
+        GGML_ASSERT(ne10 % blck == 0);
+        GGML_ASSERT(ne0  % blck == 0);
+        ne00_arg = ne00/blck;
+        ne10_arg = ne10/blck;
+        ne0_arg  = ne0/blck;
+    }
+
     ggml_metal_kargs_concat args = {
-        /*.ne00 =*/ ne00,
+        /*.ne00 =*/ ne00_arg,
         /*.ne01 =*/ ne01,
         /*.ne02 =*/ ne02,
         /*.ne03 =*/ ne03,
@@ -561,7 +577,7 @@ int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) {
         /*.nb01 =*/ nb01,
         /*.nb02 =*/ nb02,
         /*.nb03 =*/ nb03,
-        /*.ne10 =*/ ne10,
+        /*.ne10 =*/ ne10_arg,
         /*.ne11 =*/ ne11,
         /*.ne12 =*/ ne12,
         /*.ne13 =*/ ne13,
@@ -569,7 +585,7 @@ int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) {
         /*.nb11 =*/ nb11,
         /*.nb12 =*/ nb12,
         /*.nb13 =*/ nb13,
-        /*.ne0  =*/ ne0,
+        /*.ne0  =*/ ne0_arg,
         /*.ne1  =*/ ne1,
         /*.ne2  =*/ ne2,
         /*.ne3  =*/ ne3,
@@ -588,7 +604,7 @@ int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) {
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[1]), 2);
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         3);
 
-    int nth = std::min(256, ne0);
+    int nth = std::min(256, ne0_arg);
 
     // when rows are small, we can batch them together in a single threadgroup
     int nrptg = 1;
@@ -5091,7 +5107,9 @@ int ggml_metal_op_argsort(ggml_metal_op_t ctx, int idx) {
     return 1;
 }
 
-int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) {
+// bitonic-sort + merge fallback: efficient when k is small and there are few rows,
+// where the single-workgroup-per-row radix-select cannot reach enough parallelism
+static void ggml_metal_op_top_k_bitonic(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 
     ggml_metal_library_t lib = ctx->lib;
@@ -5199,6 +5217,74 @@ int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) {
 
         len <<= 1;
     }
+}
+
+// radix-select: one workgroup per row. Maps each float to an order-preserving unsigned
+// key, finds the k-th largest via 4 radix-8 histogram passes, then compacts the top-k
+// indices. Fast for large k and/or many rows.
+static void ggml_metal_op_top_k_radix(ggml_metal_op_t ctx, int idx) {
+    ggml_tensor * op = ctx->node(idx);
+
+    ggml_metal_library_t lib = ctx->lib;
+    ggml_metal_encoder_t enc = ctx->enc;
+
+    GGML_ASSERT(ggml_is_contiguous_rows(op->src[0]));
+
+    GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
+    GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
+
+    auto pipeline = ggml_metal_library_get_pipeline_top_k_radix(lib, op);
+
+    // one workgroup per row; radix-select the k-th largest value
+    const int nth = std::min(1024, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
+
+    ggml_metal_kargs_top_k args = {
+        /*.ne00  =*/ ne00,
+        /*.ne01  =*/ ne01,
+        /*.ne02  =*/ ne02,
+        /*.ne03  =*/ ne03,
+        /*.nb01  =*/ nb01,
+        /*.nb02  =*/ nb02,
+        /*.nb03  =*/ nb03,
+        /*.top_k =*/ (int32_t) op->ne[0],
+    };
+
+    // shared memory: 256-entry histogram + bucket/above scalars + output counter
+    const size_t smem_histo  = GGML_PAD(256*sizeof(uint32_t), 16);
+    const size_t smem_bucket = GGML_PAD(    sizeof(uint32_t), 16);
+    const size_t smem_above  = GGML_PAD(    sizeof(uint32_t), 16);
+    const size_t smem_out    = GGML_PAD(    sizeof(uint32_t), 16);
+
+    ggml_metal_encoder_set_pipeline(enc, pipeline);
+    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         2);
+
+    ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_histo,  0);
+    ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_bucket, 1);
+    ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_above,  2);
+    ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_out,    3);
+
+    ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1);
+}
+
+int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) {
+    ggml_tensor * op = ctx->node(idx);
+
+    // radix-select has a fixed single-workgroup-per-row cost (~50-60us) that is only
+    // amortized for long rows, many rows, or a large k; otherwise the bitonic path wins
+    const int ncols = op->src[0]->ne[0];
+    const int k     = op->ne[0];
+    const int nrows = ggml_nrows(op->src[0]);
+
+    const bool use_radix =
+        ncols > 2048 && (k > 64 || (nrows > 4 && ncols >= 8192));
+
+    if (use_radix) {
+        ggml_metal_op_top_k_radix(ctx, idx);
+    } else {
+        ggml_metal_op_top_k_bitonic(ctx, idx);
+    }
 
     return 1;
 }

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

@@ -68,6 +68,196 @@ fa_vec_cfg_t fa_vec_baseline_cfg(int dk, int dv) {
 // sweep and paste its output. See ggml-metal-tuning.h for the row/lookup semantics.
 // ref: https://github.com/ggml-org/llama.cpp/pull/27824
 constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 1, 1 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 1, 3 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 3 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, 1, 1 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, 1, 3 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 192, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 192, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 1, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 2, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 3, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 96, 96, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 96, 96, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 96, 96, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 192, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 192, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, 1, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, 2, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, 3, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 512, 512, 3, 3 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 1, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 512, 512, 2, 2 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 512, 512, 3, 2 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
+
     { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } },
     { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } },
@@ -467,6 +657,183 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
     { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
 
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 128, 128, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 128, 128, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 128, 128, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 192, 128, 1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 192, 128, 1, 3 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_F16, 320, 256, 1, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 32, 32, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 32, 32, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 32, 32, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 64, 64, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 96, 96, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 96, 96, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 96, 96, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 96, 96, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 96, 96, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 32, 32, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 32, 32, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 32, 32, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 96, 96, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 96, 96, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 128, 128, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 128, 128, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 128, 128, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 256, 256, 3, 2 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 32, 32, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 32, 32, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 192, 192, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 192, 192, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 192, 192, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 192, 192, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 256, 256, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 320, 256, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 320, 256, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_0, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 32, 32, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 32, 32, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 32, 32, 2, 4 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 32, 32, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 32, 32, 3, 4 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 96, 96, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 96, 96, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 256, 256, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 320, 256, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 320, 256, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 320, 256, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 320, 256, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 320, 256, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 320, 256, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q5_1, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 32, 32, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 32, 32, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 32, 32, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 64, 64, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 64, 64, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 96, 96, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 96, 96, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 96, 96, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 320, 256, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 320, 256, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 320, 256, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 320, 256, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 320, 256, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 320, 256, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M1_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
+
     { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 32, 32, 1, 3 }, { 2, 4 } },
     { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } },
@@ -516,6 +883,119 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
     { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 512, 512, 3, 3 }, { 2, 2 } },
     { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 32, 32, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 32, 32, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 32, 32, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 32, 32, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 96, 96, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 96, 96, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 96, 96, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 96, 96, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 96, 96, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 192, 192, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 192, 192, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 192, 192, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 192, 192, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 256, 256, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 576, 512, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 576, 512, 1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 576, 512, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 576, 512, 1, 3 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_1, 576, 512, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 96, 96, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 128, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 128, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 128, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 192, 192, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 192, 192, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 192, 128, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 192, 128, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 192, 128, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 192, 128, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 192, 128, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_0, 576, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 96, 96, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 96, 96, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 96, 96, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 128, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 192, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 192, 1, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 192, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 192, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 128, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 128, 1, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 128, 2, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 192, 128, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 320, 256, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q5_1, 576, 512, -1, 1 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } },
     { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 1, 4 } },
@@ -549,6 +1029,83 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
     { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
 
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 32, 32, 1, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 64, 64, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 64, 64, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 64, 64, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 64, 64, 3, 2 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 64, 64, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 64, 64, 3, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 96, 96, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 128, 128, 3, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 128, 128, 3, 3 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 192, 192, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 192, 128, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 192, 128, 1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 192, 128, 2, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 192, 128, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 192, 128, 3, 1 }, { 2, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 320, 256, 2, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 320, 256, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 320, 256, 2, 3 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 320, 256, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 512, 512, 2, 0 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 512, 512, 2, 1 }, { 4, 1 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 512, 512, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 64, 64, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 64, 64, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 96, 96, 1, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 4 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 128, 128, 2, 3 }, { 4, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 128, 128, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 128, 128, 3, 3 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 192, 128, 1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 192, 128, 3, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 192, 128, 3, 3 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 192, 128, 3, 4 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 256, 256, 2, 2 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 256, 256, 2, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 320, 256, 1, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 320, 256, 1, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 320, 256, 2, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 320, 256, 2, 2 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 320, 256, 2, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 320, 256, 2, 4 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 320, 256, 3, 1 }, { 2, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 320, 256, 3, 3 }, { 4, 2 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 512, 512, 3, 1 }, { 1, 1 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
+    { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
+
     { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } },
     { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 128, 128, 2, 1 }, { 1, 4 } },

ggml/src/ggml-metal/kernels/argsort.metal

@@ -230,3 +230,108 @@ kernel void kernel_argsort_merge_f32_i32(
 
 template [[host_name("kernel_argsort_merge_f32_i32_asc")]]  kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_ASC>;
 template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_DESC>;
+
+static inline uint ggml_top_k_f2ui(float x) {
+    uint y = as_type<uint>(x);
+    if ((y & 0x80000000u) != 0u) {
+        y ^= 0xFFFFFFFFu; // negative floats: flip all bits
+    } else {
+        y |= 0x80000000u; // positive floats: set the sign bit
+    }
+    return y;
+}
+
+kernel void kernel_top_k_f32_i32(
+        constant   ggml_metal_kargs_top_k & args,
+        device   const char * src0,
+        device      int32_t * dst,
+        threadgroup atomic_uint * histo     [[threadgroup(0)]],
+        threadgroup        uint * sh_bucket [[threadgroup(1)]],
+        threadgroup        uint * sh_above  [[threadgroup(2)]],
+        threadgroup atomic_uint * out_count [[threadgroup(3)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+
+    const uint ncols = args.ne00;
+    const uint top_k = args.top_k;
+    const uint i01   = tgpig[0];
+    const uint i02   = tgpig[1];
+    const uint i03   = tgpig[2];
+
+    device const float * src0_row = (device const float *) (src0 + args.nb01*i01 + args.nb02*i02 + args.nb03*i03);
+
+    device int32_t * dst_row = dst + top_k*(i01 + args.ne01*i02 + args.ne01*args.ne02*i03);
+
+    const uint tid = tpitg.x;
+    const uint ntg_x = ntg.x;
+
+    uint prefix  = 0;     // fixed high bits of the threshold key
+    uint desired = top_k; // count still needed from the candidate range
+
+    for (int shift = 24; shift >= 0; shift -= 8) {
+        for (uint i = tid; i < 256; i += ntg_x) {
+            atomic_store_explicit(&histo[i], 0u, memory_order_relaxed);
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        const uint hi_mask   = (shift + 8 >= 32) ? 0u : (0xFFFFFFFFu << uint(shift + 8));
+        const uint prefix_hi = prefix & hi_mask;
+
+        for (uint i = tid; i < ncols; i += ntg_x) {
+            const uint key = ggml_top_k_f2ui(src0_row[i]);
+            if ((key & hi_mask) == prefix_hi) {
+                atomic_fetch_add_explicit(&histo[(key >> uint(shift)) & 0xFFu], 1u, memory_order_relaxed);
+            }
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        // top-down scan for the bucket holding the k-th value
+        if (tid == 0) {
+            uint acc = 0;
+            uint b   = 0;
+            for (int bb = 255; bb >= 0; --bb) {
+                const uint c = atomic_load_explicit(&histo[bb], memory_order_relaxed);
+                if (acc + c >= desired) {
+                    b = uint(bb);
+                    break;
+                }
+                acc += c;
+            }
+            *sh_bucket = b;
+            *sh_above  = acc;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        prefix  |= *sh_bucket << uint(shift);
+        desired -= *sh_above;
+
+        // ensure every thread has consumed sh_bucket/sh_above before the next pass
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    if (tid == 0) {
+        atomic_store_explicit(out_count, 0u, memory_order_relaxed);
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // emit everything above the threshold, then fill the rest from ties
+    const uint threshold = prefix;
+
+    for (uint i = tid; i < ncols; i += ntg_x) {
+        if (ggml_top_k_f2ui(src0_row[i]) > threshold) {
+            const uint pos = atomic_fetch_add_explicit(out_count, 1u, memory_order_relaxed);
+            dst_row[pos] = (int32_t) i;
+        }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    for (uint i = tid; i < ncols; i += ntg_x) {
+        if (ggml_top_k_f2ui(src0_row[i]) == threshold) {
+            const uint pos = atomic_fetch_add_explicit(out_count, 1u, memory_order_relaxed);
+            if (pos < top_k) {
+                dst_row[pos] = (int32_t) i;
+            }
+        }
+    }
+}

ggml/src/ggml-metal/kernels/quantize.metal

@@ -207,6 +207,51 @@ template [[host_name("kernel_concat_i16")]]  kernel kernel_concat_t kernel_conca
 template [[host_name("kernel_concat_i32")]]  kernel kernel_concat_t kernel_concat<int>;
 template [[host_name("kernel_concat_i64")]]  kernel kernel_concat_t kernel_concat<long>;
 
+template<typename block_q>
+kernel void kernel_concat_q(
+        constant ggml_metal_kargs_concat & args,
+        device  const char * src0,
+        device  const char * src1,
+        device        char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
+
+    // note: for quantized types, the args are in units of blocks (nb0 == type_size)
+    const int i3 = tgpig.z;
+    const int i2 = tgpig.y;
+    const int i1 = ntg.y == 1 ? tgpig.x : tgpig.x*ntg.y + tpitg.y;
+
+    if (i1 >= args.ne1) {
+        return;
+    }
+
+    int o[4] = {0, 0, 0, 0};
+    o[args.dim] = args.dim == 0 ? args.ne00 : (args.dim == 1 ? args.ne01 : (args.dim == 2 ? args.ne02 : args.ne03));
+
+    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
+        device const block_q * x;
+
+        if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
+            x = (device const block_q *)(src0 + (i3       )*args.nb03 + (i2       )*args.nb02 + (i1       )*args.nb01 + (i0       )*args.nb00);
+        } else {
+            x = (device const block_q *)(src1 + (i3 - o[3])*args.nb13 + (i2 - o[2])*args.nb12 + (i1 - o[1])*args.nb11 + (i0 - o[0])*args.nb10);
+        }
+
+        device block_q * y = (device block_q *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
+
+        *y = *x;
+    }
+}
+
+typedef decltype(kernel_concat_q<block_q4_0>) kernel_concat_q_t;
+
+template [[host_name("kernel_concat_q4_0")]] kernel kernel_concat_q_t kernel_concat_q<block_q4_0>;
+template [[host_name("kernel_concat_q4_1")]] kernel kernel_concat_q_t kernel_concat_q<block_q4_1>;
+template [[host_name("kernel_concat_q5_0")]] kernel kernel_concat_q_t kernel_concat_q<block_q5_0>;
+template [[host_name("kernel_concat_q5_1")]] kernel kernel_concat_q_t kernel_concat_q<block_q5_1>;
+template [[host_name("kernel_concat_q8_0")]] kernel kernel_concat_q_t kernel_concat_q<block_q8_0>;
+
 template<typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread float4x4 &)>
 kernel void kernel_get_rows_q(
         constant ggml_metal_kargs_get_rows & args,

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

@@ -20002,7 +20002,8 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
                 }
 
                 kernel = backend_ctx->kernel_mul_mm_q4_k_f32_l4_lm;
-                nth0 = 128; // calculated as (BM*BN)/(TM*TN)
+                // (BM*BN)/(TM*TN): Intel uses an 8x8 microtile (WG=64), others 4x8 (WG=128)
+                nth0 = (backend_ctx->gpu_family == INTEL) ? 64 : 128;
 
                 int batch_stride_a = ne00*ne01;
                 int batch_stride_b = ne10*ne11;
@@ -20046,7 +20047,7 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
                 }
 
                 kernel = backend_ctx->kernel_mul_mm_q5_k_f32_l4_lm;
-                nth0 = 128; // calculated as (BM*BN)/(TM*TN)
+                nth0 = (backend_ctx->gpu_family == INTEL) ? 64 : 128; // Intel 8x8 microtile
 
                 int batch_stride_a = ne00*ne01;
                 int batch_stride_b = ne10*ne11;
@@ -20860,7 +20861,7 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
             if (backend_ctx->gpu_family == INTEL) {
                 nth0 = 16;
                 nth1 = 1;
-                ndst = 4;
+                ndst = 16; // 8->16 rows per subgroup — matches N_DST in mul_mv_q4_k_f32_flat.cl (32 spills)
             } else if (backend_ctx->gpu_family == ADRENO) {
                 nth0 = 64;
                 nth1 = 2;
@@ -20934,7 +20935,7 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
             if (backend_ctx->gpu_family == INTEL) {
                 nth0 = 16;
                 nth1 = 1;
-                ndst = 4;
+                ndst = 8; // 4->8 rows per subgroup (2x activation reuse)
             } else if (backend_ctx->gpu_family == ADRENO) {
                 nth0 = 64;
                 nth1 = 2;

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

@@ -1,13 +1,23 @@
 #pragma OPENCL EXTENSION cl_khr_fp16 : enable
 
+#ifdef cl_intel_required_subgroup_size
+#define INTEL_GPU 1
+#endif
+
 #define LOAD_VEC_A 4
 #define LOAD_VEC_B 4
 
 #define BM 64
 #define BN 64
 #define BK 32
+#ifdef INTEL_GPU
+// Intel Xe iGPU: 8x8 microtile (WG = BM*BN/(TM*TN) = 64) — ~+12% pp512 vs 4x8
+#define TM 8
+#define TN 8
+#else
 #define TM 4
 #define TN 8
+#endif
 
 kernel void kernel_mul_mm_q4_k_f32_l4_lm(
     global uchar4 * src0_q,

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

@@ -1,13 +1,23 @@
 #pragma OPENCL EXTENSION cl_khr_fp16 : enable
 
+#ifdef cl_intel_required_subgroup_size
+#define INTEL_GPU 1
+#endif
+
 #define LOAD_VEC_A 4
 #define LOAD_VEC_B 4
 
 #define BM 64
 #define BN 64
 #define BK 32
+#ifdef INTEL_GPU
+// Intel Xe iGPU: 8x8 microtile (WG=64)
+#define TM 8
+#define TN 8
+#else
 #define TM 4
 #define TN 8
+#endif
 
 kernel void kernel_mul_mm_q5_k_f32_l4_lm(
     global uchar4 * src0_q,

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

@@ -40,7 +40,7 @@ typedef struct {
 #undef N_SIMDWIDTH
 
 #ifdef INTEL_GPU
-#define N_DST 4 // number of rows each SIMD group works on
+#define N_DST 16 // number of rows each SIMD group works on (Intel: 8->16, 2x further activation reuse; 32 spills registers)
 #define N_SIMDGROUP 1 // number of SIMD groups in a thread group
 #define N_SIMDWIDTH 16 // SIMD group size
 #elif defined (ADRENO_GPU)

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

@@ -38,7 +38,7 @@ typedef struct {
 #undef N_SIMDWIDTH
 
 #ifdef INTEL_GPU
-#define N_DST       4
+#define N_DST       8 // Intel: 4->8 for 2x activation reuse (see mul_mv_q4_k_f32_flat.cl)
 #define N_SIMDGROUP 1
 #define N_SIMDWIDTH 16
 #elif defined(ADRENO_GPU)

ggml/src/ggml-rpc/ggml-rpc.cpp

@@ -625,7 +625,7 @@ static bool ggml_backend_buffer_is_rpc(ggml_backend_buffer_t buffer) {
     return buffer->iface.free_buffer == ggml_backend_rpc_buffer_free_buffer;
 }
 
-static rpc_tensor serialize_tensor(const ggml_tensor * tensor) {
+static rpc_tensor serialize_tensor(const ggml_tensor * tensor, const std::shared_ptr<rpc_dispatcher> & dispatcher = nullptr) {
     rpc_tensor result;
     if (!tensor) {
         memset(&result, 0, sizeof(result));
@@ -637,8 +637,14 @@ static rpc_tensor serialize_tensor(const ggml_tensor * tensor) {
     if (tensor->buffer && ggml_backend_buffer_is_rpc(tensor->buffer)) {
         ggml_backend_buffer_t buffer = tensor->buffer;
         ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
-        result.buffer = ctx != nullptr ? ctx->remote_ptr : 0;
-        result.data = reinterpret_cast<uint64_t>(tensor->data);
+        // ref: https://github.com/ggml-org/llama.cpp/pull/26500
+        if (ctx != nullptr && (dispatcher == nullptr || ctx->dispatcher == dispatcher)) {
+            result.buffer = ctx->remote_ptr;
+            result.data = reinterpret_cast<uint64_t>(tensor->data);
+        } else {
+            result.buffer = 0;
+            result.data = 0;
+        }
     } else {
         result.buffer = 0;
         result.data   = 0;
@@ -958,7 +964,7 @@ static void ggml_backend_rpc_synchronize(ggml_backend_t backend) {
     rpc_ctx->dispatcher->synchronize();
 }
 
-static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::vector<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) {
+static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, const std::shared_ptr<rpc_dispatcher> & dispatcher, std::vector<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) {
     if (tensor == nullptr) {
         return;
     }
@@ -967,10 +973,10 @@ static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::ve
     }
     visited.insert(tensor);
     for (int i = 0; i < GGML_MAX_SRC; i++) {
-        add_tensor(tensor->src[i], cgraph, tensors, visited);
+        add_tensor(tensor->src[i], cgraph, dispatcher, tensors, visited);
     }
-    add_tensor(tensor->view_src, cgraph, tensors, visited);
-    rpc_tensor result = serialize_tensor(tensor);
+    add_tensor(tensor->view_src, cgraph, dispatcher, tensors, visited);
+    rpc_tensor result = serialize_tensor(tensor, dispatcher);
     const size_t hash_pos = ggml_hash_find(&cgraph->visited_hash_set, tensor);
     if (hash_pos != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, hash_pos)) {
         result.use_count = cgraph->use_counts[hash_pos];
@@ -978,12 +984,12 @@ static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::ve
     tensors.push_back(result);
 }
 
-static uint8_t * serialize_graph(uint32_t device, const ggml_cgraph * cgraph, size_t * output_size) {
+static uint8_t * serialize_graph(uint32_t device, const ggml_cgraph * cgraph, const std::shared_ptr<rpc_dispatcher> & dispatcher, size_t * output_size) {
     uint32_t n_nodes = cgraph->n_nodes;
     std::vector<rpc_tensor> tensors;
     std::unordered_set<ggml_tensor*> visited;
     for (uint32_t i = 0; i < n_nodes; i++) {
-        add_tensor(cgraph->nodes[i], cgraph, tensors, visited);
+        add_tensor(cgraph->nodes[i], cgraph, dispatcher, tensors, visited);
     }
     // serialization format:
     // | device (4 bytes) | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) |
@@ -1020,7 +1026,7 @@ static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, g
     } else {
         rpc_dev_ctx->last_graph_uid = cgraph->uid;
         size_t input_size = 0;
-        uint8_t * input = serialize_graph(rpc_ctx->device, cgraph, &input_size);
+        uint8_t * input = serialize_graph(rpc_ctx->device, cgraph, rpc_ctx->dispatcher, &input_size);
         std::shared_ptr<uint8_t> input_ptr(input, std::default_delete<uint8_t[]>());
         rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_COMPUTE, input_ptr, input_size);
     }

ggml/src/ggml-sycl/base.hpp

@@ -0,0 +1,36 @@
+#ifndef GGML_SYCL_BASE_HPP
+#define GGML_SYCL_BASE_HPP
+
+/**
+ * Module: base
+ *
+ * Description:
+ * Provides zero-dependency, foundational primitives, core abstractions,
+ * and low-level system interfaces. This module acts as the lowest layer
+ * of the architecture and is consumed globally across all subsystems.
+ *
+ * Constraints:
+ * - STRICTLY zero upstream dependencies (leaf module).
+ * - High stability and backward compatibility required.
+ */
+
+#include <cstdio>
+
+extern int g_ggml_sycl_debug;
+
+#if defined(__clang__) && __has_builtin(__builtin_expect)
+// Hint the optimizer to pipeline the more likely following instruction in branches
+#    define LIKELY(expr)   __builtin_expect(expr, true)
+#    define UNLIKELY(expr) __builtin_expect(expr, false)
+#else
+#    define LIKELY(expr)   (expr)
+#    define UNLIKELY(expr) (expr)
+#endif
+
+#define GGML_SYCL_DEBUG(...)              \
+    do {                                  \
+        if (UNLIKELY(g_ggml_sycl_debug))  \
+            fprintf(stderr, __VA_ARGS__); \
+    } while (0)
+
+#endif  // GGML_SYCL_BASE_HPP

ggml/src/ggml-sycl/common.hpp

@@ -18,6 +18,7 @@
 #include <iostream>
 #include <string>
 
+#include "base.hpp"
 #include "dpct/helper.hpp"
 #include "ggml.h"
 #include "ggml-impl.h"
@@ -69,21 +70,6 @@ extern int g_ggml_sycl_fa_onednn;
 extern int g_ggml_sycl_fa_onednn_max_kv;
 
 
-#if defined(__clang__) && __has_builtin(__builtin_expect)
-// Hint the optimizer to pipeline the more likely following instruction in branches
-#    define LIKELY(expr)   __builtin_expect(expr, true)
-#    define UNLIKELY(expr) __builtin_expect(expr, false)
-#else
-#    define LIKELY(expr)   (expr)
-#    define UNLIKELY(expr) (expr)
-#endif
-
-#define GGML_SYCL_DEBUG(...)              \
-    do {                                  \
-        if (UNLIKELY(g_ggml_sycl_debug))  \
-            fprintf(stderr, __VA_ARGS__); \
-    } while (0)
-
 #define CHECK_TRY_ERROR(expr)                                            \
   [&]() {                                                                \
     try {                                                                \

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -35,6 +35,7 @@
 #include <sycl/backend.hpp>
 #ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API
 #include <level_zero/ze_api.h>
+#include <level_zero/zes_api.h>
 #endif
 #if defined(GGML_SYCL_GRAPH) && SYCL_EXT_ONEAPI_ASYNC_MEMORY_ALLOC
 #    include <sycl/ext/oneapi/experimental/async_alloc/async_alloc.hpp>
@@ -61,6 +62,7 @@
 #include "ggml-sycl/fwht.hpp"
 #include "ggml-sycl/gemm.hpp"
 #include "ggml-sycl/getrows.hpp"
+#include "ggml-sycl/mem.hpp"
 #include "ggml-sycl/norm.hpp"
 #include "ggml-sycl/presets.hpp"
 #include "ggml-sycl/quantize.hpp"
@@ -105,6 +107,9 @@ int g_ggml_sycl_enable_flash_attention = 1;
 int g_ggml_sycl_dev2dev_memcpy = DEV2DEV_MEMCPY_SYCL;
 int g_ggml_sycl_usm_system = 0;
 int g_ggml_sycl_enable_host_pinned_mem = 1;
+int g_ggml_sycl_host_pinned_mem_2g = 0;
+int g_ggml_sycl_get_mem_api = MEMORY_API_TYPE_LEVEL_ZERO;
+
 
 static ggml_sycl_device_info ggml_sycl_init() {
     ggml_sycl_device_info info = {};
@@ -301,10 +306,27 @@ static const char* dev2dev_int2str(int dev2dev) {
     }
 }
 
+/*
+* There are several entry APIs to be called as first function in SYCL backend in different cases.
+* It's the first internal function to be called by them in SYCL backend.
+* This function is used to do initialize work for the SYCL backend and set the global variables.
+*/
+void initialize_sycl_begining() {
+#ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API
+    ze_result_t zes_init = zesInit(0);
+    if (zes_init != ZE_RESULT_SUCCESS) {
+        std::cerr << "Warning: zesInit failed [ggml_check_sycl] with code " << static_cast<int>(zes_init)
+                            << ". Sysman free-memory query may be unavailable.\n";
+    }
+#endif
+}
+
 static void ggml_check_sycl() try {
     static bool initialized = false;
 
     if (!initialized) {
+        initialize_sycl_begining();
+
         g_ggml_sycl_debug = ggml_sycl_get_env("GGML_SYCL_DEBUG", 0);
         g_ggml_sycl_enable_optimize = ggml_sycl_get_env("GGML_SYCL_ENABLE_OPT", 1);
         g_ggml_sycl_enable_graph = ggml_sycl_get_env("GGML_SYCL_ENABLE_GRAPH", 0);
@@ -317,8 +339,11 @@ static void ggml_check_sycl() try {
         g_ggml_sycl_prioritize_dmmv = ggml_sycl_get_env("GGML_SYCL_PRIORITIZE_DMMV", 0);
 
         g_ggml_sycl_dev2dev_memcpy = ggml_sycl_get_env("GGML_SYCL_DEV2DEV_MEMCPY", DEV2DEV_MEMCPY_SYCL);
+        g_ggml_sycl_get_mem_api = ggml_sycl_get_env("GGML_SYCL_GET_MEM_API", MEMORY_API_TYPE_LEVEL_ZERO);
+
         if (g_ggml_sycl_use_level_zero_api == 0) {
             g_ggml_sycl_dev2dev_memcpy = DEV2DEV_MEMCPY_SYCL;
+            g_ggml_sycl_get_mem_api = MEMORY_API_TYPE_SYCL;
         }
 
 #ifdef SYCL_FLASH_ATTN
@@ -331,6 +356,9 @@ static void ggml_check_sycl() try {
         g_ggml_sycl_enable_host_pinned_mem =
             ggml_sycl_get_env("GGML_SYCL_ENABLE_HOST_PINNED_MEM", 1);
 
+        g_ggml_sycl_host_pinned_mem_2g =
+            ggml_sycl_get_env("GGML_SYCL_HOST_PINNED_MEM_2G", 0) & g_ggml_sycl_enable_host_pinned_mem;
+
         GGML_SYCL_DEBUG("[SYCL] call ggml_check_sycl\n");
 
         GGML_LOG_INFO("Build with Macros:\n");
@@ -374,9 +402,12 @@ static void ggml_check_sycl() try {
 
 #ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API
         GGML_LOG_INFO("  GGML_SYCL_DEV2DEV_MEMCPY: %d (%s)\n", g_ggml_sycl_dev2dev_memcpy, dev2dev_int2str(g_ggml_sycl_dev2dev_memcpy));
+        GGML_LOG_INFO("  GGML_SYCL_GET_MEM_API: %d (%s)\n", g_ggml_sycl_get_mem_api, mem_api_int2str(g_ggml_sycl_get_mem_api));
 #else
         GGML_LOG_INFO("  GGML_SYCL_DEV2DEV_MEMCPY: %d (%s), enable to SYCL API since missing GGML_SYCL_SUPPORT_LEVEL_ZERO_API\n",
                       g_ggml_sycl_dev2dev_memcpy, dev2dev_int2str(g_ggml_sycl_dev2dev_memcpy));
+        GGML_LOG_INFO("  GGML_SYCL_GET_MEM_API: %d (%s), enable to SYCL API since missing GGML_SYCL_SUPPORT_LEVEL_ZERO_API\n",
+                      g_ggml_sycl_get_mem_api, mem_api_int2str(g_ggml_sycl_get_mem_api));
 #endif
 
 #if defined(GGML_SYCL_DNNL)
@@ -429,6 +460,7 @@ static void ggml_check_sycl() try {
 
         GGML_LOG_INFO("  GGML_SYCL_USM_SYSTEM: %d\n", g_ggml_sycl_usm_system);
         GGML_LOG_INFO("  GGML_SYCL_ENABLE_HOST_PINNED_MEM: %d\n", g_ggml_sycl_enable_host_pinned_mem);
+        GGML_LOG_INFO("  GGML_SYCL_HOST_PINNED_MEM_2G: %d\n", g_ggml_sycl_host_pinned_mem_2g);
 
 /* NOT REMOVE, keep it for next optimize for XMX.
 #if defined(SYCL_USE_XMX)
@@ -949,8 +981,12 @@ static size_t ggml_backend_sycl_buffer_type_get_alignment(ggml_backend_buffer_ty
 }
 
 static size_t ggml_backend_sycl_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) {
-    return dpct::get_current_device().get_max_mem_alloc_size();
-
+    size_t max_alloc_size = dpct::get_current_device().get_max_mem_alloc_size();
+    if (g_ggml_sycl_host_pinned_mem_2g) {
+        return std::min(max_alloc_size, (size_t) 2LL*1024*1024*1024);
+    } else {
+        return max_alloc_size;
+    }
     GGML_UNUSED(buft);
 }
 
@@ -1523,7 +1559,12 @@ static size_t ggml_backend_sycl_host_buffer_type_get_max_size(ggml_backend_buffe
 
     if (g_ggml_sycl_enable_host_pinned_mem) {
         ggml_backend_sycl_device_context * dev_ctx = (ggml_backend_sycl_device_context *) buft->device->context;
-        return dpct::dev_mgr::instance().get_device(dev_ctx->device).get_max_mem_alloc_size();
+        size_t max_alloc_size = dpct::dev_mgr::instance().get_device(dev_ctx->device).get_max_mem_alloc_size();
+        if (g_ggml_sycl_host_pinned_mem_2g) {
+            return std::min(max_alloc_size, (size_t) 2LL*1024*1024*1024);
+        } else {
+            return max_alloc_size;
+        }
     } else {
         return SIZE_MAX;
     }
@@ -5208,6 +5249,7 @@ catch (sycl::exception const &exc) {
 
 static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct ggml_tensor * dst) try {
     if (!g_sycl_loaded) return false;
+    initialize_sycl_begining();
 
     if (dst->src[0] != nullptr && ggml_backend_buffer_is_sycl_split(dst->src[0]->buffer)) {
         ggml_sycl_set_peer_access(dst->src[1]->ne[1], ctx.device);
@@ -5590,18 +5632,16 @@ catch (sycl::exception const &exc) {
   std::exit(1);
 }
 
-void ggml_backend_sycl_get_device_memory(int device, size_t *free,
-                                                   size_t *total) try {
+void ggml_backend_sycl_get_device_memory(int device, size_t * free, size_t * total) try {
     GGML_SYCL_DEBUG("[SYCL] call ggml_backend_sycl_get_device_memory\n");
-    ggml_sycl_set_device(device);
-
-    SYCL_CHECK(CHECK_TRY_ERROR(
-        dpct::dev_mgr::instance().get_device(device).get_memory_info(*free, *total)));
-}
-catch (sycl::exception const &exc) {
-  std::cerr << exc.what() << "Exception caught at file:" << __FILE__
-            << ", line:" << __LINE__ << std::endl;
-  std::exit(1);
+    bool res = get_memory_size(dpct::dev_mgr::instance().get_device(device), *free, *total,
+                               (MemoryAPIType) g_ggml_sycl_get_mem_api);
+    if (!res) {
+        GGML_ABORT("[%s] failed to get device memory size", __func__);
+    }
+} catch (const sycl::exception & exc) {
+    std::cerr << exc.what() << "Exception caught at file:" << __FILE__ << ", line:" << __LINE__ << std::endl;
+    std::exit(1);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -6020,10 +6060,12 @@ static const char * ggml_backend_sycl_device_get_description(ggml_backend_dev_t
 }
 
 static void ggml_backend_sycl_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) {
-    ggml_backend_sycl_device_context * ctx = (ggml_backend_sycl_device_context *)dev->context;
-    ggml_sycl_set_device(ctx->device);
-    SYCL_CHECK(CHECK_TRY_ERROR(
-    dpct::dev_mgr::instance().get_device(ctx->device).get_memory_info(*free, *total)));
+    ggml_backend_sycl_device_context * ctx = (ggml_backend_sycl_device_context *) dev->context;
+    bool res = get_memory_size(dpct::dev_mgr::instance().get_device(ctx->device), *free, *total,
+                               (MemoryAPIType) g_ggml_sycl_get_mem_api);
+    if (!res) {
+        GGML_ABORT("[%s] failed to get device memory size", __func__);
+    }
 }
 
 static enum ggml_backend_dev_type ggml_backend_sycl_device_get_type(ggml_backend_dev_t dev) {
@@ -6906,6 +6948,7 @@ ggml_backend_reg_t ggml_backend_sycl_reg() {
         static std::mutex mutex;
         std::lock_guard<std::mutex> lock(mutex);
         if (!initialized) {
+            initialize_sycl_begining();
             ggml_backend_sycl_reg_context * ctx = new ggml_backend_sycl_reg_context;
             const int min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32;
 

ggml/src/ggml-sycl/mem.cpp

@@ -0,0 +1,162 @@
+#include <sycl/ext/oneapi/backend/level_zero.hpp>
+#include <sycl/sycl.hpp>
+
+#ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API
+#include <level_zero/ze_api.h>
+#include <level_zero/zes_api.h>
+#endif
+
+#include "base.hpp"
+#include "mem.hpp"
+
+#include <cstdint>
+#include <iostream>
+#include <vector>
+
+const char * mem_api_int2str(int mem_api) {
+    if (mem_api == MEMORY_API_TYPE_SYCL) {
+        return "SYCL API";
+    } else if (mem_api == MEMORY_API_TYPE_LEVEL_ZERO) {
+        return "Level Zero API";
+    } else {
+        return "Unknown";
+    }
+}
+
+#ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API
+bool query_free_memory_by_ze(sycl::device dev, size_t & free_bytes, size_t & total_bytes) {
+    free_bytes  = 0;
+    total_bytes = 0;
+
+    uint32_t module_count = 0;
+
+#if defined(SYCL_EXT_ONEAPI_BACKEND_LEVEL_ZERO)
+    constexpr sycl::backend kL0Backend = sycl::backend::ext_oneapi_level_zero;
+#else
+    constexpr sycl::backend kL0Backend = sycl::backend::level_zero;
+#endif
+
+    try {
+        ze_result_t zes_init = zesInit(0);
+        if (zes_init != ZE_RESULT_SUCCESS) {
+            std::cerr << "Warning: zesInit failed with code " << static_cast<int>(zes_init)
+                      << ". Sysman free-memory query may be unavailable.\n";
+        }
+
+        if (dev.get_platform().get_backend() != kL0Backend) {
+            GGML_SYCL_DEBUG("Device backend is not Level Zero; falling back to SYCL memory query.\n");
+            total_bytes = dev.get_info<sycl::info::device::global_mem_size>();
+            free_bytes  = total_bytes;
+            return false;
+        }
+
+        ze_device_handle_t ze_dev = sycl::get_native<kL0Backend>(dev);
+        if (ze_dev == nullptr) {
+            GGML_SYCL_DEBUG("Level Zero device handle is null; falling back to SYCL memory query.\n");
+            total_bytes = dev.get_info<sycl::info::device::global_mem_size>();
+            free_bytes  = total_bytes;
+            return false;
+        }
+
+        ze_result_t r = zesDeviceEnumMemoryModules(ze_dev, &module_count, nullptr);
+        if (r != ZE_RESULT_SUCCESS || module_count == 0) {
+            GGML_SYCL_DEBUG("Failed to enumerate Level Zero memory modules. Falling back to SYCL memory query.\n");
+            total_bytes = dev.get_info<sycl::info::device::global_mem_size>();
+            free_bytes  = total_bytes;
+            return false;
+        }
+
+        std::vector<zes_mem_handle_t> modules(module_count);
+        r = zesDeviceEnumMemoryModules(ze_dev, &module_count, modules.data());
+        if (r != ZE_RESULT_SUCCESS || module_count == 0) {
+            GGML_SYCL_DEBUG("Failed to enumerate Level Zero memory modules. Falling back to SYCL memory query.\n");
+            total_bytes = dev.get_info<sycl::info::device::global_mem_size>();
+            free_bytes  = total_bytes;
+            return false;
+        }
+
+        for (uint32_t i = 0; i < module_count; ++i) {
+            zes_mem_state_t state = {};
+            state.stype           = ZES_STRUCTURE_TYPE_MEM_STATE;
+            state.pNext           = nullptr;
+
+            r = zesMemoryGetState(modules[i], &state);
+            if (r != ZE_RESULT_SUCCESS) {
+                continue;
+            }
+
+            free_bytes += state.free;
+            total_bytes += state.size;
+        }
+
+        if (total_bytes == 0) {
+            GGML_SYCL_DEBUG("Level Zero memory query returned zero total bytes. Falling back to SYCL memory query.\n");
+            total_bytes = dev.get_info<sycl::info::device::global_mem_size>();
+            free_bytes  = total_bytes;
+            return false;
+        }
+        return true;
+    } catch (const sycl::exception & e) {
+        GGML_SYCL_DEBUG("Level Zero memory query failed: %s\n", e.what());
+        total_bytes = dev.get_info<sycl::info::device::global_mem_size>();
+        free_bytes  = total_bytes;
+        return false;
+    }
+}
+#endif
+
+bool get_memory_size_by_sycl_api(sycl::device dev, size_t & free_bytes, size_t & total_bytes) {
+    GGML_SYCL_DEBUG("[%s]Querying free memory using SYCL API.\n", __func__);
+    total_bytes = dev.get_info<sycl::info::device::global_mem_size>();
+
+#if (defined(__SYCL_COMPILER_VERSION) && __SYCL_COMPILER_VERSION >= 20221105)
+    if (dev.has(sycl::aspect::ext_intel_free_memory)) {
+        try {
+            GGML_SYCL_DEBUG("Querying free memory using SYCL aspect::ext_intel_free_memory.");
+            free_bytes = dev.get_info<sycl::ext::intel::info::device::free_memory>();
+            return true;
+        } catch (const sycl::exception &) {
+            GGML_SYCL_DEBUG(
+                "Failed to query free memory using SYCL aspect::ext_intel_free_memory. Using total memory as free "
+                "memory.");
+            free_bytes = total_bytes;
+            return false;
+        }
+    } else {
+        GGML_SYCL_DEBUG(
+            "Device does not support SYCL aspect::ext_intel_free_memory. Using total memory as free memory.");
+        free_bytes = total_bytes;
+    }
+#else
+    GGML_SYCL_DEBUG("SYCL Compiler version is older than 20221105. Using total memory as free memory.");
+    free_bytes = total_bytes;
+#endif
+    return true;
+}
+
+bool get_memory_size(sycl::device dev, size_t & free_bytes, size_t & total_bytes, MemoryAPIType api_type) {
+    const auto name       = dev.get_info<sycl::info::device::name>();
+    const auto vendor     = dev.get_info<sycl::info::device::vendor>();
+    const auto global_mem = dev.get_info<sycl::info::device::global_mem_size>();
+
+    GGML_SYCL_DEBUG("[%s]GPU Name:          %s\n", __func__, name.c_str());
+    GGML_SYCL_DEBUG("[%s]GPU Vendor:        %s\n", __func__, vendor.c_str());
+    GGML_SYCL_DEBUG("[%s]GPU Global Memory: %zu bytes\n", __func__, static_cast<size_t>(global_mem));
+
+    if (api_type == MEMORY_API_TYPE_LEVEL_ZERO) {
+#ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API
+        GGML_SYCL_DEBUG("[%s]Querying free memory using Level Zero API.\n", __func__);
+        if (!query_free_memory_by_ze(dev, free_bytes, total_bytes)) {
+            //fallback to SYCL API if Level Zero API fails
+            GGML_SYCL_DEBUG("[%s]Falling back to SYCL API for memory query.\n", __func__);
+            return get_memory_size_by_sycl_api(dev, free_bytes, total_bytes);
+        }
+        return true;
+#else
+        GGML_SYCL_DEBUG("[%s]Level Zero API support is not enabled. Please enable it to use this feature.\n", __func__);
+        return false;
+#endif
+    } else {  //MEMORY_API_TYPE_SYCL
+        return get_memory_size_by_sycl_api(dev, free_bytes, total_bytes);
+    }
+}

ggml/src/ggml-sycl/mem.hpp

@@ -0,0 +1,16 @@
+#ifndef GGML_SYCL_MEM_HPP
+#define GGML_SYCL_MEM_HPP
+
+#include <sycl/sycl.hpp>
+
+enum MemoryAPIType {
+    MEMORY_API_TYPE_LEVEL_ZERO = 0,
+    MEMORY_API_TYPE_SYCL = 1,
+};
+
+const char* mem_api_int2str(int mem_api);
+
+bool get_memory_size(sycl::device dev, size_t & free_bytes, size_t & total_bytes,
+    MemoryAPIType api_type);
+
+#endif  // GGML_SYCL_MEM_HPP

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -657,6 +657,21 @@ static constexpr std::initializer_list<ggml_op> snake_pattern              { GGM
                                                                              GGML_OP_SQR,      GGML_OP_MUL,
                                                                              GGML_OP_ADD };
 
+// qwen4 QSA indexer: gather per-block scores to cells + add f16 mask (cast+reshape) + top-k,
+// fused into one radix-select. The cast/reshape are elided; the raw f16 mask is read in-shader.
+static constexpr std::initializer_list<ggml_op> topk_qsa_pattern { GGML_OP_GET_ROWS, GGML_OP_PERMUTE,
+                                                                   GGML_OP_CONT,     GGML_OP_CPY,
+                                                                   GGML_OP_RESHAPE,  GGML_OP_ADD,
+                                                                   GGML_OP_TOP_K };
+static constexpr std::initializer_list<std::array<int, 3>> topk_qsa_edges {
+    { 1, 0, 0 }, // permute->src[0] == get_rows
+    { 2, 0, 1 }, // cont->src[0]    == permute
+    { 4, 0, 3 }, // reshape->src[0] == cpy (mask cast)
+    { 5, 0, 2 }, // add->src[0]     == cont
+    { 5, 1, 4 }, // add->src[1]     == reshape
+    { 6, 0, 5 }, // top_k->src[0]   == add
+};
+
 //node #978 (  SOFT_MAX):     ffn_moe_probs-15 (   0K) [Vulka         ] use=2:    ffn_moe_logits-15 (   0K) [Vulka         ]
 //node #979 (   RESHAPE): ffn_moe_probs-15 (re (   0K) [Vulka         ] use=1:     ffn_moe_probs-15 (   0K) [Vulka         ]
 //node #980 (   ARGSORT):   ffn_moe_argsort-15 (   0K) [Vulka         ] use=1:     ffn_moe_probs-15 (   0K) [Vulka         ]
@@ -1057,6 +1072,8 @@ struct vk_device_struct {
     vk_pipeline pipeline_argsort_f32[num_argsort_pipelines];
     vk_pipeline pipeline_argsort_large_f32[num_argsort_pipelines];
     vk_pipeline pipeline_topk_f32[num_topk_pipelines];
+    vk_pipeline pipeline_topk_radix_f32;
+    vk_pipeline pipeline_topk_radix_qsa; // qwen4 QSA indexer fusion (f16 mask)
     vk_pipeline pipeline_sum_rows_f32;
     vk_pipeline pipeline_cross_entropy_loss_f32, pipeline_cross_entropy_loss_f32_wg512;
     vk_pipeline pipeline_cross_entropy_loss_back_f32, pipeline_cross_entropy_loss_back_f32_wg512;
@@ -1749,6 +1766,15 @@ struct vk_op_topk_push_constants {
     uint32_t last_pass;
 };
 
+struct vk_op_topk_radix_push_constants {
+    uint32_t ncols;
+    uint32_t k;
+    uint32_t nrows;
+    uint32_t n_tps;    // QSA only
+    uint32_t n_blocks; // QSA only
+    uint32_t n_stream; // QSA only
+};
+
 struct vk_op_im2col_push_constants {
     uint64_t dst_addr;
     uint32_t batch_offset; uint32_t offset_delta;
@@ -2439,6 +2465,8 @@ struct ggml_backend_vk_context {
     int fused_ops_write_mask {};
     topk_moe_mode fused_topk_moe_mode {};
     bool fused_topk_moe_scale {};
+    // QSA indexer gather+add+top_k fused into one radix-select
+    bool fused_topk_qsa {};
 
     // for GGML_VK_PERF_LOGGER
     std::unique_ptr<vk_perf_logger> perf_logger;
@@ -5261,6 +5289,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         rm_stdq = 2;
         rm_stdq_int = 2;
     }
+    // RDNA3: above four columns, static 4 rows for all types bench faster than the default
+    const bool is_rdna3 = device->vendor_id == VK_VENDOR_ID_AMD && device->architecture == AMD_RDNA3;
+    auto const &rm_int_n = [&](uint32_t rows, uint32_t i) { return (is_rdna3 && i >= 4) ? 4u : rows; };
+    // RDNA3: Static 4 rows for all types bench faster than the default
+    auto const &rm_id = [&](uint32_t rows) { return is_rdna3 ? 4u : rows; };
     uint32_t rm_iq = 2 * rm_kq;
 
     const bool use_subgroups = device->subgroup_arithmetic;
@@ -5357,20 +5390,20 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
                 const uint32_t subgroup_size_int = (device->vendor_id == VK_VENDOR_ID_INTEL && device->subgroup_size_control) ? device->subgroup_min_size : device->subgroup_size;
                 const uint32_t wg_size_subgroup_int = (w == DMMV_WG_SIZE_SUBGROUP) ? subgroup_size_int : (subgroup_size_int * 4);
 
-                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q2_0][i], "mul_mat_vec_q2_0_q8_1_f32", arr_dmmv_q2_0_q8_1_f32_len[reduc], arr_dmmv_q2_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 2*rm_kq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
-                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_q8_1_f32", arr_dmmv_q4_0_q8_1_f32_len[reduc], arr_dmmv_q4_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
-                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_q8_1_f32", arr_dmmv_q4_1_q8_1_f32_len[reduc], arr_dmmv_q4_1_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
-                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_q8_1_f32", arr_dmmv_q5_0_q8_1_f32_len[reduc], arr_dmmv_q5_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
-                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_q8_1_f32", arr_dmmv_q5_1_q8_1_f32_len[reduc], arr_dmmv_q5_1_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
-                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_q8_1_f32", arr_dmmv_q8_0_q8_1_f32_len[reduc], arr_dmmv_q8_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q2_0][i], "mul_mat_vec_q2_0_q8_1_f32", arr_dmmv_q2_0_q8_1_f32_len[reduc], arr_dmmv_q2_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(2*rm_kq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(2*rm_kq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int);
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_q8_1_f32", arr_dmmv_q4_0_q8_1_f32_len[reduc], arr_dmmv_q4_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(1*rm_stdq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(1*rm_stdq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int);
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_q8_1_f32", arr_dmmv_q4_1_q8_1_f32_len[reduc], arr_dmmv_q4_1_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(1*rm_stdq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(1*rm_stdq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int);
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_q8_1_f32", arr_dmmv_q5_0_q8_1_f32_len[reduc], arr_dmmv_q5_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(1*rm_stdq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(1*rm_stdq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int);
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_q8_1_f32", arr_dmmv_q5_1_q8_1_f32_len[reduc], arr_dmmv_q5_1_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(1*rm_stdq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(1*rm_stdq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int);
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_q8_1_f32", arr_dmmv_q8_0_q8_1_f32_len[reduc], arr_dmmv_q8_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(1*rm_stdq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(1*rm_stdq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int);
 
-                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_MXFP4][i], "mul_mat_vec_mxfp4_q8_1_f32", arr_dmmv_mxfp4_q8_1_f32_len[reduc], arr_dmmv_mxfp4_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 2*rm_stdq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_MXFP4][i], "mul_mat_vec_mxfp4_q8_1_f32", arr_dmmv_mxfp4_q8_1_f32_len[reduc], arr_dmmv_mxfp4_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(2*rm_stdq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(2*rm_stdq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int);
 
-                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_q8_1_f32", arr_dmmv_q2_k_q8_1_f32_len[reduc], arr_dmmv_q2_k_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 2*rm_kq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
-                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_q8_1_f32", arr_dmmv_q3_k_q8_1_f32_len[reduc], arr_dmmv_q3_k_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_kq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
-                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_q8_1_f32", arr_dmmv_q4_k_q8_1_f32_len[reduc], arr_dmmv_q4_k_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_kq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
-                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_q8_1_f32", arr_dmmv_q5_k_q8_1_f32_len[reduc], arr_dmmv_q5_k_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_kq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
-                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_q8_1_f32", arr_dmmv_q6_k_q8_1_f32_len[reduc], arr_dmmv_q6_k_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_kq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_q8_1_f32", arr_dmmv_q2_k_q8_1_f32_len[reduc], arr_dmmv_q2_k_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(2*rm_kq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(2*rm_kq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int);
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_q8_1_f32", arr_dmmv_q3_k_q8_1_f32_len[reduc], arr_dmmv_q3_k_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(1*rm_kq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(1*rm_kq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int);
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_q8_1_f32", arr_dmmv_q4_k_q8_1_f32_len[reduc], arr_dmmv_q4_k_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(1*rm_kq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(1*rm_kq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int);
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_q8_1_f32", arr_dmmv_q5_k_q8_1_f32_len[reduc], arr_dmmv_q5_k_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(1*rm_kq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(1*rm_kq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int);
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_q8_1_f32", arr_dmmv_q6_k_q8_1_f32_len[reduc], arr_dmmv_q6_k_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(1*rm_kq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(1*rm_kq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int);
 
                 ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_IQ1_S][i], "mul_mat_vec_iq1_s_q8_1_f32", arr_dmmv_iq1_s_q8_1_f32_len[reduc], arr_dmmv_iq1_s_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_iq_int(i), 1, 1}, {wg_size_subgroup_int, 1*rm_iq_int(i), i+1}, 1, true, use_subgroups, subgroup_size_int);
                 ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_IQ1_M][i], "mul_mat_vec_iq1_m_q8_1_f32", arr_dmmv_iq1_m_q8_1_f32_len[reduc], arr_dmmv_iq1_m_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_iq_int(i), 1, 1}, {wg_size_subgroup_int, 1*rm_iq_int(i), i+1}, 1, true, use_subgroups, subgroup_size_int);
@@ -5412,20 +5445,20 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
             const uint32_t subgroup_size_int = (device->vendor_id == VK_VENDOR_ID_INTEL && device->subgroup_size_control) ? device->subgroup_min_size : device->subgroup_size;
             const uint32_t wg_size_subgroup_int = (w == DMMV_WG_SIZE_SUBGROUP) ? subgroup_size_int : (subgroup_size_int * 4);
 
-            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q2_0], "mul_mat_vec_id_q2_0_q8_1_f32", arr_dmmv_id_q2_0_q8_1_f32_len[reduc], arr_dmmv_id_q2_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 2*rm_kq_int}, 1, true, use_subgroups, subgroup_size_int);
-            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q4_0], "mul_mat_vec_id_q4_0_q8_1_f32", arr_dmmv_id_q4_0_q8_1_f32_len[reduc], arr_dmmv_id_q4_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int}, 1, true, use_subgroups, subgroup_size_int);
-            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q4_1], "mul_mat_vec_id_q4_1_q8_1_f32", arr_dmmv_id_q4_1_q8_1_f32_len[reduc], arr_dmmv_id_q4_1_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int}, 1, true, use_subgroups, subgroup_size_int);
-            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q5_0], "mul_mat_vec_id_q5_0_q8_1_f32", arr_dmmv_id_q5_0_q8_1_f32_len[reduc], arr_dmmv_id_q5_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int}, 1, true, use_subgroups, subgroup_size_int);
-            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q5_1], "mul_mat_vec_id_q5_1_q8_1_f32", arr_dmmv_id_q5_1_q8_1_f32_len[reduc], arr_dmmv_id_q5_1_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int}, 1, true, use_subgroups, subgroup_size_int);
-            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q8_0], "mul_mat_vec_id_q8_0_q8_1_f32", arr_dmmv_id_q8_0_q8_1_f32_len[reduc], arr_dmmv_id_q8_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int}, 1, true, use_subgroups, subgroup_size_int);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q2_0], "mul_mat_vec_id_q2_0_q8_1_f32", arr_dmmv_id_q2_0_q8_1_f32_len[reduc], arr_dmmv_id_q2_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(2*rm_kq_int), 1, 1}, {wg_size_subgroup_int, rm_id(2*rm_kq_int)}, 1, true, use_subgroups, subgroup_size_int);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q4_0], "mul_mat_vec_id_q4_0_q8_1_f32", arr_dmmv_id_q4_0_q8_1_f32_len[reduc], arr_dmmv_id_q4_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(1*rm_stdq_int), 1, 1}, {wg_size_subgroup_int, rm_id(1*rm_stdq_int)}, 1, true, use_subgroups, subgroup_size_int);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q4_1], "mul_mat_vec_id_q4_1_q8_1_f32", arr_dmmv_id_q4_1_q8_1_f32_len[reduc], arr_dmmv_id_q4_1_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(1*rm_stdq_int), 1, 1}, {wg_size_subgroup_int, rm_id(1*rm_stdq_int)}, 1, true, use_subgroups, subgroup_size_int);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q5_0], "mul_mat_vec_id_q5_0_q8_1_f32", arr_dmmv_id_q5_0_q8_1_f32_len[reduc], arr_dmmv_id_q5_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(1*rm_stdq_int), 1, 1}, {wg_size_subgroup_int, rm_id(1*rm_stdq_int)}, 1, true, use_subgroups, subgroup_size_int);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q5_1], "mul_mat_vec_id_q5_1_q8_1_f32", arr_dmmv_id_q5_1_q8_1_f32_len[reduc], arr_dmmv_id_q5_1_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(1*rm_stdq_int), 1, 1}, {wg_size_subgroup_int, rm_id(1*rm_stdq_int)}, 1, true, use_subgroups, subgroup_size_int);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q8_0], "mul_mat_vec_id_q8_0_q8_1_f32", arr_dmmv_id_q8_0_q8_1_f32_len[reduc], arr_dmmv_id_q8_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(1*rm_stdq_int), 1, 1}, {wg_size_subgroup_int, rm_id(1*rm_stdq_int)}, 1, true, use_subgroups, subgroup_size_int);
 
-            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_MXFP4], "mul_mat_vec_id_mxfp4_q8_1_f32", arr_dmmv_id_mxfp4_q8_1_f32_len[reduc], arr_dmmv_id_mxfp4_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 2*rm_stdq_int}, 1, true, use_subgroups, subgroup_size_int);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_MXFP4], "mul_mat_vec_id_mxfp4_q8_1_f32", arr_dmmv_id_mxfp4_q8_1_f32_len[reduc], arr_dmmv_id_mxfp4_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(2*rm_stdq_int), 1, 1}, {wg_size_subgroup_int, rm_id(2*rm_stdq_int)}, 1, true, use_subgroups, subgroup_size_int);
 
-            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q2_K], "mul_mat_vec_id_q2_k_q8_1_f32", arr_dmmv_id_q2_k_q8_1_f32_len[reduc], arr_dmmv_id_q2_k_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 2*rm_kq_int}, 1, true, use_subgroups, subgroup_size_int);
-            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q3_K], "mul_mat_vec_id_q3_k_q8_1_f32", arr_dmmv_id_q3_k_q8_1_f32_len[reduc], arr_dmmv_id_q3_k_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_kq_int}, 1, true, use_subgroups, subgroup_size_int);
-            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q4_K], "mul_mat_vec_id_q4_k_q8_1_f32", arr_dmmv_id_q4_k_q8_1_f32_len[reduc], arr_dmmv_id_q4_k_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_kq_int}, 1, true, use_subgroups, subgroup_size_int);
-            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q5_K], "mul_mat_vec_id_q5_k_q8_1_f32", arr_dmmv_id_q5_k_q8_1_f32_len[reduc], arr_dmmv_id_q5_k_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_kq_int}, 1, true, use_subgroups, subgroup_size_int);
-            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q6_K], "mul_mat_vec_id_q6_k_q8_1_f32", arr_dmmv_id_q6_k_q8_1_f32_len[reduc], arr_dmmv_id_q6_k_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_kq_int}, 1, true, use_subgroups, subgroup_size_int);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q2_K], "mul_mat_vec_id_q2_k_q8_1_f32", arr_dmmv_id_q2_k_q8_1_f32_len[reduc], arr_dmmv_id_q2_k_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(2*rm_kq_int), 1, 1}, {wg_size_subgroup_int, rm_id(2*rm_kq_int)}, 1, true, use_subgroups, subgroup_size_int);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q3_K], "mul_mat_vec_id_q3_k_q8_1_f32", arr_dmmv_id_q3_k_q8_1_f32_len[reduc], arr_dmmv_id_q3_k_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(1*rm_kq_int), 1, 1}, {wg_size_subgroup_int, rm_id(1*rm_kq_int)}, 1, true, use_subgroups, subgroup_size_int);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q4_K], "mul_mat_vec_id_q4_k_q8_1_f32", arr_dmmv_id_q4_k_q8_1_f32_len[reduc], arr_dmmv_id_q4_k_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(1*rm_kq_int), 1, 1}, {wg_size_subgroup_int, rm_id(1*rm_kq_int)}, 1, true, use_subgroups, subgroup_size_int);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q5_K], "mul_mat_vec_id_q5_k_q8_1_f32", arr_dmmv_id_q5_k_q8_1_f32_len[reduc], arr_dmmv_id_q5_k_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(1*rm_kq_int), 1, 1}, {wg_size_subgroup_int, rm_id(1*rm_kq_int)}, 1, true, use_subgroups, subgroup_size_int);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q6_K], "mul_mat_vec_id_q6_k_q8_1_f32", arr_dmmv_id_q6_k_q8_1_f32_len[reduc], arr_dmmv_id_q6_k_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(1*rm_kq_int), 1, 1}, {wg_size_subgroup_int, rm_id(1*rm_kq_int)}, 1, true, use_subgroups, subgroup_size_int);
 
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_IQ1_S], "mul_mat_vec_id_iq1_s_q8_1_f32", arr_dmmv_id_iq1_s_q8_1_f32_len[reduc], arr_dmmv_id_iq1_s_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_iq_int(0), 1, 1}, {wg_size_subgroup_int, 1*rm_iq_int(0)}, 1, true, use_subgroups, subgroup_size_int);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_IQ1_M], "mul_mat_vec_id_iq1_m_q8_1_f32", arr_dmmv_id_iq1_m_q8_1_f32_len[reduc], arr_dmmv_id_iq1_m_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_iq_int(0), 1, 1}, {wg_size_subgroup_int, 1*rm_iq_int(0)}, 1, true, use_subgroups, subgroup_size_int);
@@ -5439,6 +5472,9 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 #if !defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
     GGML_UNUSED(rm_stdq_int);
     GGML_UNUSED(rm_kq_int);
+    GGML_UNUSED(is_rdna3);
+    GGML_UNUSED(rm_int_n);
+    GGML_UNUSED(rm_id);
     GGML_UNUSED(rm_iq_int);
 #endif
 
@@ -5814,6 +5850,14 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         }
     }
 
+    // large-k fallback: one workgroup per row, radix-select instead of a full sort. The QSA
+    // variant (spec constant 1) additionally gathers the qwen4 indexer input on the fly.
+    {
+        const uint32_t BLOCK_SIZE = 1u << std::min(10u, device->max_workgroup_size_log2);
+        ggml_vk_create_pipeline2(device, device->pipeline_topk_radix_f32, "topk_radix_f32", topk_radix_select_f32_len, topk_radix_select_f32_data, "main", 5, sizeof(vk_op_topk_radix_push_constants), {BLOCK_SIZE, 1, 1}, {BLOCK_SIZE, 0}, 1, true);
+        ggml_vk_create_pipeline2(device, device->pipeline_topk_radix_qsa, "topk_radix_qsa", topk_radix_select_f32_len, topk_radix_select_f32_data, "main", 5, sizeof(vk_op_topk_radix_push_constants), {BLOCK_SIZE, 1, 1}, {BLOCK_SIZE, 1}, 1, true);
+    }
+
     ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_sum_rows_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
@@ -13940,6 +13984,31 @@ static void ggml_vk_topk(ggml_backend_vk_context * ctx, vk_context& subctx, cons
     uint32_t nrows = ggml_nrows(src0);
     uint32_t k = dst->ne[0];
 
+    // tournament path is faster where it fits; use radix-select only past its k limit
+    const uint32_t k_min_pipeline = std::max((uint32_t) log2f(float(k)) + 1, ctx->device->subgroup_size_log2);
+    if (k_min_pipeline >= num_topk_pipelines || ctx->device->pipeline_topk_f32[k_min_pipeline] == nullptr) {
+        vk_pipeline pipeline = ctx->device->pipeline_topk_radix_f32;
+        GGML_ASSERT(pipeline != nullptr);
+
+        if (ctx->prealloc_x_need_sync) {
+            ggml_vk_sync_buffers(ctx, subctx);
+        }
+
+        vk_op_topk_radix_push_constants pc { ncols, k, nrows, 0, 0, 0 };
+        std::array<uint32_t, 3> elements {
+            pipeline->wg_denoms[0],
+            std::min(nrows, ctx->device->properties.limits.maxComputeWorkGroupCount[1]),
+            1,
+        };
+        // the non-QSA path only uses bindings 0/1; bind valid buffers for the unused QSA slots
+        vk_subbuffer src0_buf = ggml_vk_tensor_subbuffer(ctx, src0);
+        vk_subbuffer dst_buf  = ggml_vk_tensor_subbuffer(ctx, dst);
+        ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+        ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
+            { src0_buf, dst_buf, src0_buf, src0_buf, src0_buf }, pc, elements);
+        return;
+    }
+
     vk_op_topk_push_constants pc { ncols, ncols, ncols, k, nrows, 0, 0 };
 
     if (ctx->prealloc_x_need_sync) {
@@ -14043,6 +14112,55 @@ static void ggml_vk_topk(ggml_backend_vk_context * ctx, vk_context& subctx, cons
     ctx->prealloc_x_need_sync = true;
 }
 
+static void ggml_vk_topk_qsa(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_cgraph * cgraph, int node_idx) {
+    const ggml_tensor * get_rows = cgraph->nodes[node_idx + 0];
+    const ggml_tensor * add      = cgraph->nodes[node_idx + ctx->num_additional_fused_ops - 1];
+    ggml_tensor *       top_k    = cgraph->nodes[node_idx + ctx->num_additional_fused_ops];
+
+    const ggml_tensor * scores   = get_rows->src[0]; // [n_tps, n_blocks, n_stream]
+    const ggml_tensor * cell_blk = get_rows->src[1]; // [n_kv, n_stream]
+
+    // raw f16 mask: follow the reshape/cpy chain back to the materialized input
+    const ggml_tensor * mask = add->src[1];
+    while (mask->op == GGML_OP_RESHAPE || mask->op == GGML_OP_CPY) {
+        mask = mask->src[0];
+    }
+
+    const uint32_t n_tps    = scores->ne[0];
+    const uint32_t n_blocks = scores->ne[1];
+    const uint32_t n_stream = scores->ne[2];
+    const uint32_t n_kv     = cell_blk->ne[0];
+    const uint32_t width    = top_k->ne[0];
+    const uint32_t nrows    = n_tps * n_stream;
+
+    vk_pipeline pipeline = ctx->device->pipeline_topk_radix_qsa;
+    GGML_ASSERT(pipeline != nullptr);
+
+    // scratch holds the gathered+masked input, materialized once and reused across passes
+    const size_t scratch_size = size_t{ n_kv } * nrows * sizeof(float);
+    if (ctx->prealloc_size_x < scratch_size) {
+        ctx->prealloc_size_x = scratch_size;
+        ggml_vk_preallocate_buffers(ctx, subctx);
+    }
+    if (ctx->prealloc_x_need_sync) {
+        ggml_vk_sync_buffers(ctx, subctx);
+    }
+
+    vk_op_topk_radix_push_constants pc { n_kv, width, nrows, n_tps, n_blocks, n_stream };
+    std::array<uint32_t, 3> elements {
+        pipeline->wg_denoms[0],
+        std::min(nrows, ctx->device->properties.limits.maxComputeWorkGroupCount[1]),
+        1,
+    };
+    vk_subbuffer scratch_buf { ctx->prealloc_x, 0, ctx->prealloc_x->size };
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+    ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
+        { ggml_vk_tensor_subbuffer(ctx, scores), ggml_vk_tensor_subbuffer(ctx, top_k),
+          ggml_vk_tensor_subbuffer(ctx, cell_blk), ggml_vk_tensor_subbuffer(ctx, mask),
+          scratch_buf }, pc, elements);
+    ctx->prealloc_x_need_sync = true;
+}
+
 static void ggml_vk_sum(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
     vk_op_sum_rows_push_constants p = vk_op_sum_rows_push_constants_init(src0, dst, ggml_nelements(src0));
     ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_SUM, p);
@@ -15704,7 +15822,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
 
         break;
     case GGML_OP_GET_ROWS:
-        ggml_vk_get_rows(ctx, compute_ctx, src0, src1, node);
+        if (ctx->fused_topk_qsa) {
+            ggml_vk_topk_qsa(ctx, compute_ctx, cgraph, node_idx);
+        } else {
+            ggml_vk_get_rows(ctx, compute_ctx, src0, src1, node);
+        }
 
         break;
     case GGML_OP_GET_ROWS_BACK:
@@ -17116,6 +17238,92 @@ static bool ggml_vk_can_fuse_topk_moe(ggml_backend_vk_context * ctx, const struc
     return true;
 }
 
+// Manual op-sequence match (ggml_can_fuse_subgraph rejects the mask's external reshape/cpy).
+static bool ggml_vk_match_ops(const struct ggml_cgraph * cgraph, int node_idx,
+                              const std::initializer_list<ggml_op> & ops) {
+    if (node_idx + (int) ops.size() > cgraph->n_nodes) {
+        return false;
+    }
+    for (size_t j = 0; j < ops.size(); ++j) {
+        const ggml_tensor * node = cgraph->nodes[node_idx + j];
+        if (node->op != ops.begin()[j] ||
+            (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0 ||
+            (node->flags & GGML_TENSOR_FLAG_OUTPUT) != 0) {
+            return false;
+        }
+    }
+    return true;
+}
+
+// True if the qwen4 QSA indexer top-k can be fused at node_idx (the get_rows).
+static bool ggml_vk_can_fuse_topk_qsa(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx) {
+    if (ctx->device->disable_fusion || !ctx->device->pipeline_topk_radix_qsa) {
+        return false;
+    }
+
+    const int n_ops = topk_qsa_pattern.size();
+    if (!ggml_vk_match_ops(cgraph, node_idx, topk_qsa_pattern) ||
+        !ggml_check_edges(cgraph, node_idx, topk_qsa_edges)) {
+        return false;
+    }
+
+    // elided nodes must be single-use (cpy counts its own src[1] self-reference)
+    for (int j = 0; j < n_ops - 1; ++j) {
+        const ggml_tensor * node = cgraph->nodes[node_idx + j];
+        const int32_t want = node->op == GGML_OP_CPY ? 2 : 1;
+        if (ggml_node_get_use_count(cgraph, node_idx + j) != want) {
+            return false;
+        }
+    }
+
+    const ggml_tensor * get_rows = cgraph->nodes[node_idx + 0];
+    const ggml_tensor * add      = cgraph->nodes[node_idx + n_ops - 2];
+    const ggml_tensor * top_k    = cgraph->nodes[node_idx + n_ops - 1];
+
+    const ggml_tensor * scores   = get_rows->src[0]; // [n_tps, n_blocks, n_stream]
+    const ggml_tensor * cell_blk = get_rows->src[1]; // [n_kv, n_stream]
+    const ggml_tensor * expanded = add->src[0];      // [n_kv, n_tps, n_stream]
+
+    // raw mask: follow the reshape/cpy chain back to the materialized f16 input
+    const ggml_tensor * mask = add->src[1];
+    while (mask && (mask->op == GGML_OP_RESHAPE || mask->op == GGML_OP_CPY)) {
+        mask = mask->src[0];
+    }
+    if (!mask || mask->type != GGML_TYPE_F16) {
+        return false;
+    }
+
+    if (scores->type != GGML_TYPE_F32 || cell_blk->type != GGML_TYPE_I32 || top_k->type != GGML_TYPE_I32) {
+        return false;
+    }
+    if (!ggml_is_contiguous(scores) || !ggml_is_contiguous(cell_blk) || !ggml_is_contiguous(mask) ||
+        !ggml_is_contiguous(expanded) || !ggml_is_contiguous(top_k)) {
+        return false;
+    }
+
+    const int64_t n_tps    = scores->ne[0];
+    const int64_t n_blocks = scores->ne[1];
+    const int64_t n_stream = scores->ne[2];
+    const int64_t n_kv     = cell_blk->ne[0];
+    const int64_t width    = top_k->ne[0];
+
+    // pin the indexer layout the shader's addressing assumes
+    if (scores->ne[3] != 1 || cell_blk->ne[1] != n_stream || ggml_nrows(cell_blk) != n_stream ||
+        ggml_nelements(mask) != n_kv * n_tps * n_stream ||
+        expanded->ne[0] != n_kv || expanded->ne[1] != n_tps || expanded->ne[2] != n_stream ||
+        top_k->ne[1] != n_tps || top_k->ne[2] != n_stream || top_k->ne[3] != 1 ||
+        n_blocks <= 0 || n_kv <= 0 || width <= 0 || width > n_kv) {
+        return false;
+    }
+
+    // only worth it in the radix regime; small k uses the faster tournament unfused
+    const uint32_t k_min_pipeline = std::max((uint32_t) log2f(float(width)) + 1, ctx->device->subgroup_size_log2);
+    if (k_min_pipeline < num_topk_pipelines && ctx->device->pipeline_topk_f32[k_min_pipeline]) {
+        return false;
+    }
+    return true;
+}
+
 static bool ggml_vk_can_fuse_rope_set_rows(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph,
                                            int node_idx) {
     GGML_UNUSED(ctx);
@@ -17495,6 +17703,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
 
         ctx->fused_topk_moe_mode = TOPK_MOE_COUNT;
         ctx->fused_topk_moe_scale = false;
+        ctx->fused_topk_qsa = false;
         const char *fusion_string {};
         if (!ctx->device->disable_fusion) {
             uint32_t num_adds = ggml_vk_fuse_multi_add(ctx, cgraph, i);
@@ -17584,6 +17793,11 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
                 // with a data dependency on that register. The overlap check still
                 // rejects partial overlaps (different base or size).
                 std::fill_n(op_srcs_fused_elementwise, 5, true);
+            } else if (ggml_vk_can_fuse_topk_qsa(ctx, cgraph, i)) {
+                ctx->num_additional_fused_ops = topk_qsa_pattern.size() - 1;
+                ctx->fused_topk_qsa = true;
+                fusion_string = "TOPK_QSA";
+                std::fill_n(op_srcs_fused_elementwise, ctx->num_additional_fused_ops + 1, false);
             } else if (ggml_can_fuse_subgraph(cgraph, i, topk_moe_early_softmax_norm, { i + 3, i + 9 }) &&
                        ggml_check_edges(cgraph, i, topk_moe_early_softmax_norm_edges) &&
                        ggml_vk_can_fuse_topk_moe(ctx, cgraph, i, TOPK_MOE_EARLY_SOFTMAX_NORM)) {
@@ -17700,6 +17914,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
                 ctx->fused_ops_write_mask = 1;
                 ctx->fused_topk_moe_mode = TOPK_MOE_COUNT;
                 ctx->fused_topk_moe_scale = false;
+                ctx->fused_topk_qsa = false;
             }
         }
 
@@ -17896,6 +18111,9 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
         if (keep_pattern(snake_pattern)) {
             continue;
         }
+        if (keep_pattern(topk_qsa_pattern)) {
+            continue;
+        }
 
         // First, grab the next unused node.
         current_set.push_back(first_unused);
@@ -17914,13 +18132,23 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
             if (is_empty(graph->nodes[j])) {
                 continue;
             }
-            // Don't pull forward nodes from fusion patterns
+            // Protect every interior QSA node (not just the start): the mask branch is
+            // independent, so it gets pulled out and breaks keep_pattern otherwise.
+            auto const &in_qsa_pattern = [&](int n) -> bool {
+                for (int o = 0; o < (int) topk_qsa_pattern.size(); ++o) {
+                    if (n - o >= 0 && match_pattern(topk_qsa_pattern, n - o)) {
+                        return true;
+                    }
+                }
+                return false;
+            };
             if (match_pattern(topk_moe_early_softmax_norm, j) ||
                 match_pattern(topk_moe_sigmoid_norm_bias, j) ||
                 match_pattern(topk_moe_sqrt_softplus_norm_bias, j) ||
                 match_pattern(topk_moe_early_softmax, j) ||
                 match_pattern(topk_moe_late_softmax, j) ||
-                match_pattern(snake_pattern, j)) {
+                match_pattern(snake_pattern, j) ||
+                in_qsa_pattern(j)) {
                 continue;
             }
             bool ok = true;
@@ -18723,15 +18951,14 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
                 if (!ggml_is_contiguous(op) || !ggml_is_contiguous(op->src[0])) {
                     return false;
                 }
-                // We could potentially support larger, using argsort to sort the
-                // whole thing. Not clear if this is needed.
-                uint32_t min_pipeline = (uint32_t)log2f(float(op->ne[0])) + 1;
-                if (min_pipeline >= num_topk_pipelines ||
-                    !device->pipeline_topk_f32[min_pipeline]) {
-                    return false;
+                // large k falls back to radix-select
+                const uint32_t min_pipeline =
+                    std::max((uint32_t) log2f(float(op->ne[0])) + 1, device->subgroup_size_log2);
+                if (min_pipeline < num_topk_pipelines && device->pipeline_topk_f32[min_pipeline]) {
+                    return true;
                 }
+                return device->pipeline_topk_radix_f32 != nullptr;
             }
-            return true;
         case GGML_OP_UPSCALE:
             if (op->op_params[0] & GGML_SCALE_FLAG_ANTIALIAS) {
                 if ((op->op_params[0] & 0xFF) != GGML_SCALE_MODE_BILINEAR) {

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

@@ -0,0 +1,144 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : enable
+#extension GL_EXT_shader_16bit_storage : require
+
+#include "types.glsl"
+
+layout(constant_id = 0) const int BLOCK_SIZE = 1024;
+layout(constant_id = 1) const int QSA = 0; // 1: fuse the qwen4 QSA indexer gather + f16 mask
+
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A  {float     data_a[];};   // input values, or QSA block scores [n_tps, n_blocks, n_stream]
+layout (binding = 1) writeonly buffer D {int       data_d[];};   // [k, ...]
+layout (binding = 2) readonly buffer CB {int       cell_blk[];}; // QSA: cell->block map [n_kv, n_stream]
+layout (binding = 3) readonly buffer M  {float16_t mask[];};     // QSA: raw f16 kq_mask [n_kv, n_tps, n_stream]
+layout (binding = 4) buffer S           {float     scratch[];};  // QSA: [nrows, n_kv] gathered inputs
+
+layout (push_constant) uniform parameter {
+    uint ncols;
+    uint k;
+    uint nrows;
+    uint n_tps;    // QSA only
+    uint n_blocks; // QSA only
+    uint n_stream; // QSA only
+} p;
+
+#define RADIX_BITS 8
+#define RADIX_SIZE (1 << RADIX_BITS)
+
+shared uint histo[RADIX_SIZE];
+shared uint sh_bucket;
+shared uint sh_above;
+shared uint out_count;
+
+// order-preserving float -> uint mapping
+uint f2ui(float x) {
+    uint y = floatBitsToUint(x);
+    if ((y & 0x80000000u) != 0u) {
+        y ^= 0xFFFFFFFFu;
+    } else {
+        y |= 0x80000000u;
+    }
+    return y;
+}
+
+// QSA element i of row (t,s): score[cell_blk[i,s], t, s] + mask[i,t,s]
+float gather(uint row, uint i) {
+    const uint t = row % p.n_tps;
+    const uint s = row / p.n_tps;
+    const uint block = uint(cell_blk[s * p.ncols + i]);
+    const float a = data_a[(s * p.n_blocks + block) * p.n_tps + t];
+    const float m = float(mask[(s * p.n_tps + t) * p.ncols + i]);
+    return a + m;
+}
+
+float load(uint row, uint i, bool first) {
+    if (QSA == 0) {
+        return data_a[row * p.ncols + i];
+    }
+    // materialize the scattered gather on the first pass and reuse it after; each
+    // invocation only touches its own scratch entries, so no barrier is needed
+    const uint off = row * p.ncols + i;
+    if (first) {
+        const float v = gather(row, i);
+        scratch[off] = v;
+        return v;
+    }
+    return scratch[off];
+}
+
+// one workgroup per row: radix-select the K-th largest, then compact it plus enough ties
+void topk(const uint row) {
+    const uint tid = gl_LocalInvocationID.x;
+    const uint ncols = p.ncols;
+    const uint row_out = row * p.k;
+
+    uint prefix  = 0;     // fixed high bits of the threshold key
+    uint desired = p.k;   // count still needed from the candidate range
+
+    [[unroll]] for (int shift = 32 - RADIX_BITS; shift >= 0; shift -= RADIX_BITS) {
+        for (uint i = tid; i < RADIX_SIZE; i += BLOCK_SIZE) {
+            histo[i] = 0;
+        }
+        barrier();
+
+        const bool first     = (shift == 32 - RADIX_BITS);
+        const uint hi_mask   = (shift + RADIX_BITS >= 32) ? 0u : (0xFFFFFFFFu << uint(shift + RADIX_BITS));
+        const uint prefix_hi = prefix & hi_mask;
+        for (uint i = tid; i < ncols; i += BLOCK_SIZE) {
+            const uint key = f2ui(load(row, i, first));
+            if ((key & hi_mask) == prefix_hi) {
+                atomicAdd(histo[(key >> uint(shift)) & (RADIX_SIZE - 1)], 1u);
+            }
+        }
+        barrier();
+
+        // top-down scan for the bucket holding the K-th value
+        if (tid == 0) {
+            uint acc = 0;
+            uint b   = 0;
+            for (int bb = RADIX_SIZE - 1; bb >= 0; --bb) {
+                const uint c = histo[bb];
+                if (acc + c >= desired) { b = uint(bb); break; }
+                acc += c;
+            }
+            sh_bucket = b;
+            sh_above  = acc;
+        }
+        barrier();
+
+        prefix  |= sh_bucket << uint(shift);
+        desired -= sh_above;
+        barrier();
+    }
+
+    if (tid == 0) {
+        out_count = 0;
+    }
+    barrier();
+
+    // emit everything above the threshold, then fill the rest from ties
+    const uint threshold = prefix;
+    for (uint i = tid; i < ncols; i += BLOCK_SIZE) {
+        if (f2ui(load(row, i, false)) > threshold) {
+            data_d[row_out + atomicAdd(out_count, 1u)] = int(i);
+        }
+    }
+    barrier();
+    for (uint i = tid; i < ncols; i += BLOCK_SIZE) {
+        if (f2ui(load(row, i, false)) == threshold) {
+            const uint pos = atomicAdd(out_count, 1u);
+            if (pos < p.k) {
+                data_d[row_out + pos] = int(i);
+            }
+        }
+    }
+}
+
+void main() {
+    for (uint row = gl_WorkGroupID.y; row < p.nrows; row += gl_NumWorkGroups.y) {
+        topk(row);
+    }
+}

ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp

@@ -1028,6 +1028,7 @@ void process_shaders() {
 
     string_to_spv("topk_argsort_f32", "topk_argsort.comp", {{"A_TYPE", "float"}});
     string_to_spv("topk_nary_search_f32", "topk_nary_search.comp", {{"A_TYPE", "float"}});
+    string_to_spv("topk_radix_select_f32", "topk_radix_select.comp", {{"A_TYPE", "float"}});
 
     string_to_spv("argmax_f32", "argmax.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "int"}}));
     string_to_spv("sum_rows_f32", "sum_rows.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));

ggml/src/ggml-webgpu/ggml-webgpu.cpp

@@ -3713,11 +3713,18 @@ static void ggml_backend_webgpu_buffer_get_tensor(ggml_backend_buffer_t buffer,
 
     size_t total_offset = ggml_webgpu_tensor_offset(tensor) + offset;
 
-    size_t final_size = size;
-    if (size % 4 != 0) {
+    size_t local_offset = total_offset % 4;
+    if (local_offset != 0) {
+        // If offset is not a multiple of 4, we need to round it down to the previous
+        // multiple of 4
+        total_offset -= local_offset;
+    }
+
+    size_t final_size = size + local_offset;
+    if (final_size % 4 != 0) {
         // If size is not a multiple of 4, we need to round it up to the next
         // multiple of 4
-        final_size = size + (4 - (size % 4));
+        final_size += 4 - (final_size % 4);
     }
 
     std::lock_guard<std::recursive_mutex> lock(buf_ctx->global_ctx->mutex);
@@ -3748,7 +3755,7 @@ static void ggml_backend_webgpu_buffer_get_tensor(ggml_backend_buffer_t buffer,
     const void * mapped_range = buf_ctx->global_ctx->get_tensor_staging_buf.GetConstMappedRange(0, final_size);
 
     // Copy the data from the mapped range to the output buffer
-    std::memcpy(data, mapped_range, size);
+    std::memcpy(data, (const void *) ((const char *) mapped_range + local_offset), size);
     buf_ctx->global_ctx->get_tensor_staging_buf.Unmap();
     WEBGPU_CPU_PROFILE_TOTAL_END(get_tensor, buf_ctx->global_ctx);
 }

src/llama-arch.cpp

@@ -1100,6 +1100,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
     switch (arch) {
         case LLM_ARCH_QWEN35:
         case LLM_ARCH_QWEN35MOE:
+        case LLM_ARCH_QWEN4EXP:
         case LLM_ARCH_DEEPSEEK4:
         case LLM_ARCH_NEMOTRON_H:
         case LLM_ARCH_NEMOTRON_H_MOE:
@@ -1142,6 +1143,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
         case LLM_ARCH_BAILINGMOE3:
         case LLM_ARCH_KIMI_K3:
         case LLM_ARCH_QWEN3TTS:
+        case LLM_ARCH_QWEN4EXP:   // TODO: fix test-llama-archs
             return false;
         default:
             return true;

src/llama-context.cpp

@@ -1660,7 +1660,9 @@ int llama_context::decode(const llama_batch & batch_inp) {
 
     const int64_t n_vocab = vocab.n_tokens();
     const bool    mtp_embd = cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && batch_inp.embd;
-    const int64_t n_embd  = mtp_embd ? hparams.n_embd_out() : hparams.n_embd_inp();
+    // DFlash embd batches carry the fused target features at the encoder input width
+    const bool    dflash_embd = model.arch == LLM_ARCH_DFLASH && batch_inp.embd;
+    const int64_t n_embd  = mtp_embd ? hparams.n_embd_out() : dflash_embd ? hparams.n_embd_inp_enc() : hparams.n_embd_inp();
 
     // when computing embeddings, all tokens are output
     const bool output_all   = cparams.embeddings;
@@ -2905,17 +2907,28 @@ public:
             }
 
             if (mbuf_cur.n_tensors == mbuf.n_tensors) {
-                // same chunking: copy 1:1 by index
+                // an equal tensor count does not imply the same chunking, e.g. save ranges [2,1] vs restore runs [1,2]
+                bool same_chunking = true;
                 for (size_t i = 0; i < mbuf_cur.org.size(); ++i) {
-                    GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == ggml_nbytes(mbuf.org[i]));
-                    ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]);
+                    if (ggml_nbytes(mbuf_cur.cpy[i]) != ggml_nbytes(mbuf.org[i])) {
+                        same_chunking = false;
+                        break;
+                    }
+                }
+
+                if (same_chunking) {
+                    // same chunking: copy 1:1 by index
+                    for (size_t i = 0; i < mbuf_cur.org.size(); ++i) {
+                        ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]);
+                    }
+                    continue;
                 }
-                continue;
             }
 
             // different chunking: copy the write-side data (mbuf_cur.cpy) into the read-side targets (mbuf.org)
             // with a byte cursor. Write and read enumerate the same logical data in the same order but may chunk
-            // it differently, so copy across tensor boundaries rather than 1:1 by index.
+            // it differently (even with an equal number of tensors), so copy across tensor boundaries rather than
+            // 1:1 by index.
             const size_t total = mbuf_cur.total_size;
 
             ggml_init_params params_scratch = {

src/llama-kv-cache.cpp

@@ -2387,6 +2387,12 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
 
         ubatch.seq_id_unq[0] = dest_seq_id;
 
+        // the ext as it was saved, to put back after apply_ubatch()
+        std::vector<llama_kv_cell_ext> exts;
+        if (has_cell_ext()) {
+            exts.resize(cell_count);
+        }
+
         for (uint32_t i = 0; i < cell_count; ++i) {
             llama_pos pos;
             uint32_t n_seq_id;
@@ -2410,6 +2416,8 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
 
                 // apply_ubatch() below restores ext.tok from the ubatch tokens
                 ubatch.token[i] = ext.tok;
+
+                exts[i] = ext;
             }
 
             // read the sequence id, but directly discard it - we will use dest_seq_id instead
@@ -2461,6 +2469,14 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
         //       see: https://github.com/ggml-org/llama.cpp/pull/16825#issuecomment-3460868350
         apply_ubatch(sinfo, ubatch);
 
+        // apply_ubatch() takes the 2D position from the ubatch, and that ubatch is built with this
+        // cache's own n_pos_per_embd. a cache that does not use M-RoPE itself but mirrors one that
+        // does (the qwen4exp QSA indexer) would drop x and y. put the saved ext back instead, which
+        // is what the whole-context path below already does.
+        for (uint32_t i = 0; i < (uint32_t) exts.size(); ++i) {
+            cells.ext_set(sinfo.idxs[0][i], exts[i]);
+        }
+
         LLAMA_LOG_DEBUG("%s: cell_count = %d, dest_seq_id = %d\n", __func__, cell_count, dest_seq_id);
 
         // DEBUG CHECK: verify that all cells were allocated and have correct seq_id and pos values
@@ -2533,6 +2549,24 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
 bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, const slot_info & sinfo) {
     auto & cells = v_cells[strm];
 
+    // batch the scatter reads per contiguous run of destination indices
+    // from inclusive, to exclusive - same convention as cell_ranges_t
+    // contiguous cells yield a single run covering the whole block
+    struct cell_run { uint32_t from; uint32_t to; };
+    std::vector<cell_run> runs;
+    if (cell_count > 0) {
+        const auto & idxs = sinfo.idxs[0];
+        uint32_t i0 = 0;
+        while (i0 < cell_count) {
+            uint32_t i1 = i0 + 1;
+            while (i1 < cell_count && idxs[i1] == idxs[i1 - 1] + 1) {
+                ++i1;
+            }
+            runs.push_back({idxs[i0], idxs[i1 - 1] + 1});
+            i0 = i1;
+        }
+    }
+
     uint32_t v_trans;
     uint32_t n_layer;
 
@@ -2580,17 +2614,8 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
             return false;
         }
 
-        if (cell_count) {
-            if (sinfo.is_contiguous()) {
-                // Fast path: contiguous cells, single memcpy
-                io.read_tensor(k, sinfo.head() * k_size_row, cell_count * k_size_row);
-            } else {
-                // Slow path: scatter to non-contiguous positions
-                for (uint32_t i = 0; i < cell_count; ++i) {
-                    const size_t dst_offset = sinfo.idxs[0][i] * k_size_row;
-                    io.read_tensor(k, dst_offset, k_size_row);
-                }
-            }
+        for (const auto & r : runs) {
+            io.read_tensor(k, (size_t) r.from * k_size_row, (size_t) (r.to - r.from) * k_size_row);
         }
     }
 
@@ -2623,17 +2648,8 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
                 return false;
             }
 
-            if (cell_count) {
-                if (sinfo.is_contiguous()) {
-                    // Fast path: contiguous cells, single memcpy
-                    io.read_tensor(v, sinfo.head() * v_size_row, cell_count * v_size_row);
-                } else {
-                    // Slow path: scatter to non-contiguous positions
-                    for (uint32_t i = 0; i < cell_count; ++i) {
-                        const size_t dst_offset = sinfo.idxs[0][i] * v_size_row;
-                        io.read_tensor(v, dst_offset, v_size_row);
-                    }
-                }
+            for (const auto & r : runs) {
+                io.read_tensor(v, (size_t) r.from * v_size_row, (size_t) (r.to - r.from) * v_size_row);
             }
         }
     } else {
@@ -2674,22 +2690,10 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
                 return false;
             }
 
-            if (cell_count) {
-                if (sinfo.is_contiguous()) {
-                    // Fast path: contiguous cells
-                    const uint32_t h = sinfo.head();
-                    for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
-                        const size_t dst_offset = (h + j * cells.size()) * v_size_el;
-                        io.read_tensor(v, dst_offset, cell_count * v_size_el);
-                    }
-                } else {
-                    // Slow path: scatter to non-contiguous positions
-                    for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
-                        for (uint32_t i = 0; i < cell_count; ++i) {
-                            const size_t dst_offset = (sinfo.idxs[0][i] + j * cells.size()) * v_size_el;
-                            io.read_tensor(v, dst_offset, v_size_el);
-                        }
-                    }
+            for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
+                for (const auto & r : runs) {
+                    const size_t dst_offset = ((size_t) r.from + j * cells.size()) * v_size_el;
+                    io.read_tensor(v, dst_offset, (size_t) (r.to - r.from) * v_size_el);
                 }
             }
         }

src/llama-kv-cells.h

@@ -35,6 +35,8 @@ struct llama_kv_cell_ext {
 // TODO: add unit tests
 class llama_kv_cells {
 public:
+    using seq_set_t = std::bitset<LLAMA_MAX_SEQ>;
+
     void reset() {
         for (uint32_t i = 0; i < pos.size(); ++i) {
             pos[i]   = -1;
@@ -301,6 +303,13 @@ public:
         return seq[i].count();
     }
 
+    // the full set of sequences this cell is visible to
+    const seq_set_t & seq_get_all(uint32_t i) const {
+        assert(i < pos.size());
+
+        return seq[i];
+    }
+
     // check if the cell contains seq_id
     bool seq_has(uint32_t i, llama_seq_id seq_id) const {
         assert(i < pos.size());
@@ -511,8 +520,6 @@ private:
     //
     std::vector<llama_pos> shift;
 
-    using seq_set_t = std::bitset<LLAMA_MAX_SEQ>;
-
     // the bitset seq[i] tells us which sequences are currently occupying the i-th cell
     std::vector<seq_set_t> seq;
 

src/llama-memory-hybrid-idx.cpp

@@ -5,6 +5,7 @@
 #include "llama-io.h"
 #include "llama-model.h"
 
+
 #include <algorithm>
 #include <cassert>
 #include <cmath>
@@ -50,6 +51,10 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx(
         std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1);
         hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size;
 
+        // the cached indexer keys are raw, rotation happens after pooling at read time, so a
+        // K-shift must not rotate them while the stream copies in the same update still apply
+        hparams_idx.rope_type = LLAMA_ROPE_TYPE_NONE;
+
         LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);
 
         return new llama_kv_cache(
@@ -261,6 +266,324 @@ llama_kv_cache * llama_memory_hybrid_idx::get_mem_idx() const {
     return mem_idx.get();
 }
 
+void llama_memory_hybrid_idx::set_input_qsa(
+        ggml_tensor * cell_blk,
+        ggml_tensor * blk_cells,
+        ggml_tensor * blk_pos,
+        ggml_tensor * bias,
+        const llama_ubatch * ubatch,
+        uint32_t ratio,
+        bool blk_bias) const {
+    GGML_ASSERT(ratio > 0);
+    GGML_ASSERT(get_mem_idx() != nullptr);
+
+    GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer));
+
+    const int64_t n_kv     = cell_blk->ne[0];
+    const int64_t n_ns     = cell_blk->ne[1];        // streams in this ubatch
+    const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns);
+    const int64_t n_tokens = ubatch->n_tokens;
+    const int64_t r        = ratio;
+
+    GGML_ASSERT(n_tokens % n_ns == 0);
+    const int64_t n_tps = n_tokens/n_ns;             // tokens per stream
+
+    int32_t * dst_cell_blk  = (int32_t *) cell_blk->data;
+    int32_t * dst_blk_cells = (int32_t *) blk_cells->data;
+    int32_t * dst_blk_pos   = (int32_t *) blk_pos->data;
+    float   * dst_bias      = (float   *) bias->data;
+
+    // a block is keyed on (sequence set, index bucket): a unified cache counts every sequence
+    // from zero, so the bucket alone would pool two sequences into one block
+    GGML_ASSERT(r <= 64);
+    const uint64_t slots_full = r == 64 ? ~uint64_t(0) : ((uint64_t(1) << r) - 1);
+
+    // TODO: this runs per ubatch and is O(n_kv) per stream, about 865 us at 33k context. the cost
+    //       is the per-cell scan rather than these allocations, so hoisting them buys nothing
+    std::vector<int32_t>  blk_of(n_kv);
+    std::vector<int32_t>  cell_grp(n_kv);
+    std::vector<int32_t>  grp_head(n_blocks);
+    std::vector<int32_t>  grp_next;
+    std::vector<int32_t>  grp_first;
+    std::vector<int32_t>  grp_slot0;
+    std::vector<uint64_t> grp_slots;
+    std::vector<int32_t>  grp_bid;
+    std::vector<int32_t>  bid_idx;
+    std::vector<int32_t>  bid_cell;
+    std::vector<int32_t>  bid_slot0;
+
+    std::vector<int32_t> order;
+    std::vector<int32_t> rank;
+
+    std::fill(dst_blk_pos, dst_blk_pos + 4*n_blocks*n_ns, 0);
+
+    for (int64_t s = 0; s < n_ns; ++s) {
+        // ubatch index s*n_tps belongs to this stream; ask which cells array it uses
+        const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0];
+        const auto & cells = get_mem_idx()->get_cells(seq_of_stream);
+
+        int32_t * cur_cell_blk  = dst_cell_blk  + s*n_kv;
+        int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks);
+
+        std::fill(cur_blk_cells, cur_blk_cells + r*n_blocks, 0);
+
+        bid_idx  .clear();
+        bid_cell .clear();
+        bid_slot0.clear();
+
+        int n_seq_present = 0;
+
+        for (int sq = 0; sq < LLAMA_MAX_SEQ && n_seq_present < 2; ++sq) {
+            if (cells.seq_pos_min(sq) >= 0) {
+                n_seq_present++;
+            }
+        }
+
+        const bool one_seq = n_seq_present <= 1;
+
+        // a cell no block covers needs its own -inf, which a per-block bias cannot carry
+        // every cache path keeps the position below the cell window, so this stays false
+        bool oor = false;
+
+        bool dup = false;
+
+        bool ranked = false;
+
+        auto group_cells = [&]() {
+            // -1 means no usable block: an incomplete or short group cannot be pooled
+            std::fill(blk_of.begin(),   blk_of.end(),   -1);
+            std::fill(cell_grp.begin(), cell_grp.end(), -1);
+            std::fill(grp_head.begin(), grp_head.end(), -1);
+
+            grp_next .clear();
+            grp_first.clear();
+            grp_slot0.clear();
+            grp_slots.clear();
+            grp_bid  .clear();
+
+            oor = false;
+            dup = false;
+
+            for (int64_t j = 0; j < n_kv; ++j) {
+                if (cells.is_empty(j)) {
+                    continue;
+                }
+
+                const int64_t idx = ranked ? rank[j] : cells.pos_get(j);
+                const int64_t pb  = idx/r;
+
+                if (pb >= n_blocks) {
+                    oor = true;
+                    continue;
+                }
+
+                int32_t g = -1;
+
+                for (int32_t c = grp_head[pb]; c >= 0; c = grp_next[c]) {
+                    if (one_seq || cells.seq_get_all((uint32_t) grp_first[c]) == cells.seq_get_all((uint32_t) j)) {
+                        g = c;
+                        break;
+                    }
+                }
+
+                if (g < 0) {
+                    g = (int32_t) grp_first.size();
+
+                    grp_next .push_back(grp_head[pb]);
+                    grp_first.push_back((int32_t) j);
+                    grp_slot0.push_back(-1);
+                    grp_slots.push_back(0);
+                    grp_bid  .push_back(-1);
+
+                    grp_head[pb] = g;
+                }
+
+                const uint64_t bit = uint64_t(1) << (idx%r);
+
+                dup |= (grp_slots[g] & bit) != 0;
+
+                cell_grp[j]   = g;
+                grp_slots[g] |= bit;
+
+                if (idx%r == 0) {
+                    grp_slot0[g] = (int32_t) j;
+                }
+            }
+        };
+
+        group_cells();
+
+        // mrope repeats one position across an image, so rank cells instead of using the position
+        if (dup && ubatch->is_pos_2d() && one_seq) {
+            order.clear();
+            order.reserve(n_kv);
+
+            for (int64_t j = 0; j < n_kv; ++j) {
+                if (!cells.is_empty(j)) {
+                    order.push_back((int32_t) j);
+                }
+            }
+
+            // same total order the mrope causal mask uses: pos, then ext.y, then ext.x
+            std::sort(order.begin(), order.end(), [&cells](int32_t a, int32_t b) {
+                const llama_pos pa = cells.pos_get(a);
+                const llama_pos pb = cells.pos_get(b);
+
+                if (pa != pb) {
+                    return pa < pb;
+                }
+
+                const auto & ea = cells.ext_get(a);
+
+                return cells.ext_get(b).is_2d_gt(ea.x, ea.y);
+            });
+
+            rank.assign(n_kv, -1);
+
+            for (int64_t k = 0; k < (int64_t) order.size(); ++k) {
+                rank[order[k]] = (int32_t) k;
+            }
+
+            ranked = true;
+
+            group_cells();
+        }
+
+        GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window");
+
+        int32_t n_bid = 0;
+
+        for (int64_t pb = 0; pb < n_blocks; ++pb) {
+            for (int32_t g = grp_head[pb]; g >= 0; g = grp_next[g]) {
+                if (grp_slots[g] != slots_full) {
+                    continue;
+                }
+
+                grp_bid[g] = n_bid++;
+
+                bid_idx  .push_back((int32_t) (pb*r));
+                bid_cell .push_back(grp_first[g]);
+                bid_slot0.push_back(grp_slot0[g]);
+            }
+        }
+
+        GGML_ASSERT(n_bid <= n_blocks);
+
+        for (int32_t b = 0; b < n_bid; ++b) {
+            int32_t sec_pos[4] = { bid_idx[b], bid_idx[b], bid_idx[b], bid_idx[b] };
+
+            if (ranked) {
+                const int32_t   c = bid_slot0[b];
+                const llama_pos p = cells.pos_get(c);
+                const auto &    e = cells.ext_get(c);
+
+                sec_pos[0] = p;
+                sec_pos[1] = e.y;
+                sec_pos[2] = e.x;
+                sec_pos[3] = p;
+            }
+
+            for (int64_t sec = 0; sec < 4; ++sec) {
+                dst_blk_pos[sec*(n_blocks*n_ns) + s*n_blocks + b] = sec_pos[sec];
+            }
+        }
+
+        // unpooled cells all point at one spare block. a spare block exists only when some
+        // cell is unpooled: n_bid == n_blocks means every cell sits in a full block.
+        const bool     have_dead = n_bid < n_blocks;
+        const int32_t  dead_bid  = have_dead ? n_bid : n_blocks - 1;
+
+        for (int64_t j = 0; j < n_kv; ++j) {
+            const int32_t g = cell_grp[j];
+
+            blk_of[j] = g < 0 ? -1 : grp_bid[g];
+
+            if (blk_of[j] >= 0) {
+                const int64_t idx = ranked ? rank[j] : cells.pos_get(j);
+
+                cur_blk_cells[blk_of[j]*r + (idx%r)] = (int32_t) j;
+            }
+
+            cur_cell_blk[j] = blk_of[j] < 0 ? dead_bid : blk_of[j];
+        }
+
+        for (int64_t ii = 0; ii < n_tps; ++ii) {
+            const int64_t      i      = s*n_tps + ii;
+            const llama_seq_id seq_id = ubatch->seq_id[i][0];
+
+            int64_t q = ubatch->pos[i];
+
+            if (ranked) {
+                const llama_pos qt = ubatch->pos[i];
+                const llama_pos qy = ubatch->pos[i + n_tokens];
+                const llama_pos qx = ubatch->pos[i + n_tokens*2];
+
+                int64_t lo = 0;
+                int64_t hi = (int64_t) order.size();
+
+                while (lo < hi) {
+                    const int64_t   mid = (lo + hi)/2;
+                    const int32_t   c   = order[mid];
+                    const llama_pos pc  = cells.pos_get(c);
+
+                    if (pc < qt || (pc == qt && !cells.ext_get(c).is_2d_gt(qx, qy))) {
+                        lo = mid + 1;
+                    } else {
+                        hi = mid;
+                    }
+                }
+
+                q = lo - 1;
+            }
+
+            // the tail is an incomplete block and is always visible, as in the reference
+            const int64_t tail_start = (q + 1)/r*r;
+
+            if (blk_bias) {
+                // a block sits wholly inside or outside the tail, so one value covers it
+                // the caller adds the attention mask, which drops empty, foreign and future cells
+                float * cur_blk_bias = dst_bias + i*n_blocks;
+
+                for (int64_t b = 0; b < n_blocks; ++b) {
+                    if (b >= n_bid || !cells.seq_has((uint32_t) bid_cell[b], seq_id)) {
+                        cur_blk_bias[b] = -INFINITY;
+                        continue;
+                    }
+
+                    // finite, so it can never meet a -inf and produce a nan
+                    cur_blk_bias[b] = bid_idx[b] >= tail_start ? 1e9f : 0.0f;
+                }
+
+                // the spare block holds the unpooled cells, which are the incomplete tail, so
+                // it gets the tail value. it must stay finite: a sequence with fewer than
+                // `ratio` cells owns no full block, and a row of -inf only gives a nan.
+                if (have_dead) {
+                    cur_blk_bias[dead_bid] = 1e9f;
+                }
+
+                continue;
+            }
+
+            float * cur_bias = dst_bias + i*n_kv;
+
+            for (int64_t j = 0; j < n_kv; ++j) {
+                float v = -INFINITY;
+
+                if (!cells.is_empty(j) && cells.seq_has(j, seq_id)) {
+                    const int64_t idx = ranked ? rank[j] : cells.pos_get(j);
+
+                    if (idx <= q) {
+                        // finite, so it can never meet a -inf and produce a nan
+                        v = idx >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f);
+                    }
+                }
+
+                cur_bias[j] = v;
+            }
+        }
+    }
+}
+
 //
 // llama_memory_hybrid_idx_context
 //
@@ -295,7 +618,10 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(
                   llama_context * lctx,
                            bool   optimize) :
     llama_memory_hybrid_context(mem, lctx, optimize),
-    mem(mem) {}
+    mem(mem),
+    // update() applies a pending cross-stream seq_cp, else the copy keeps stale indexer keys
+    ctx_idx(mem->get_mem_idx() == nullptr ? nullptr :
+        mem->get_mem_idx()->init_update(lctx, optimize)) {}
 
 llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(
         llama_memory_hybrid_idx * mem,
@@ -347,119 +673,7 @@ void llama_memory_hybrid_idx_context::set_input_qsa(
         const llama_ubatch * ubatch,
         uint32_t ratio,
         bool blk_bias) const {
-    GGML_ASSERT(ratio > 0);
-    GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr);
+    GGML_ASSERT(mem != nullptr);
 
-    GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer));
-
-    const int64_t n_kv     = cell_blk->ne[0];
-    const int64_t n_ns     = cell_blk->ne[1];        // streams in this ubatch
-    const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns);
-    const int64_t n_tokens = ubatch->n_tokens;
-    const int64_t r        = ratio;
-
-    GGML_ASSERT(n_tokens % n_ns == 0);
-    const int64_t n_tps = n_tokens/n_ns;             // tokens per stream
-
-    int32_t * dst_cell_blk  = (int32_t *) cell_blk->data;
-    int32_t * dst_blk_cells = (int32_t *) blk_cells->data;
-    int32_t * dst_blk_pos   = (int32_t *) blk_pos->data;
-    float   * dst_bias      = (float   *) bias->data;
-
-    // block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio
-    // all mrope sections carry it: exact for text, approximate for images
-    for (int64_t sec = 0; sec < 4; ++sec) {
-        for (int64_t s = 0; s < n_ns; ++s) {
-            for (int64_t b = 0; b < n_blocks; ++b) {
-                dst_blk_pos[sec*(n_blocks*n_ns) + s*n_blocks + b] = (int32_t) (b*r);
-            }
-        }
-    }
-
-    // one pass per stream: cell j is a different token in each, so no mapping is shared
-    std::vector<int32_t> blk_of(n_kv);
-    std::vector<int32_t> filled(n_blocks);
-
-    for (int64_t s = 0; s < n_ns; ++s) {
-        // ubatch index s*n_tps belongs to this stream; ask which cells array it uses
-        const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0];
-        const auto & cells = mem->get_mem_idx()->get_cells(seq_of_stream);
-
-        int32_t * cur_cell_blk  = dst_cell_blk  + s*n_kv;
-        int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks);
-
-        // an incomplete block cannot be pooled; the bias below forces those tail cells in
-        // -1 means no usable block, and block 0 only keeps the gather in range
-        std::fill(blk_of.begin(),  blk_of.end(),  -1);
-        std::fill(filled.begin(),  filled.end(),   0);
-        std::fill(cur_blk_cells, cur_blk_cells + r*n_blocks, 0);
-
-        // a cell no block covers needs its own -inf, which a per-block bias cannot carry
-        // every cache path keeps the position below the cell window, so this stays false
-        bool oor = false;
-
-        for (int64_t j = 0; j < n_kv; ++j) {
-            if (cells.is_empty(j)) {
-                continue;
-            }
-
-            const llama_pos p = cells.pos_get(j);
-            const int64_t   b = p/r;
-
-            if (b >= n_blocks) {
-                oor = true;
-                continue;
-            }
-
-            blk_of[j] = (int32_t) b;
-            cur_blk_cells[b*r + (p%r)] = (int32_t) j;
-            filled[b]++;
-        }
-
-        GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window");
-
-        // per-block mode keeps an unpooled cell's real block, so the block's own -inf reaches it
-        // per-cell mode carries that -inf itself and only needs the gather in range
-        for (int64_t j = 0; j < n_kv; ++j) {
-            if (blk_of[j] >= 0 && filled[blk_of[j]] < r && !blk_bias) {
-                blk_of[j] = -1;
-            }
-            cur_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j];
-        }
-
-        for (int64_t ii = 0; ii < n_tps; ++ii) {
-            const int64_t      i      = s*n_tps + ii;
-            const llama_seq_id seq_id = ubatch->seq_id[i][0];
-            const llama_pos    q      = ubatch->pos[i];
-
-            // the tail is an incomplete block and is always visible, as in the reference
-            const llama_pos tail_start = (q + 1)/r*r;
-
-            if (blk_bias) {
-                // a block sits wholly inside or outside the tail, so one value covers it
-                // the caller adds the attention mask, which drops empty, foreign and future cells
-                float * cur_blk_bias = dst_bias + i*n_blocks;
-
-                for (int64_t b = 0; b < n_blocks; ++b) {
-                    // finite, so it can never meet a -inf and produce a nan
-                    cur_blk_bias[b] = b*r >= tail_start ? 1e9f : (filled[b] < r ? -INFINITY : 0.0f);
-                }
-
-                continue;
-            }
-
-            float * cur_bias = dst_bias + i*n_kv;
-
-            for (int64_t j = 0; j < n_kv; ++j) {
-                float v = -INFINITY;
-
-                if (!cells.is_empty(j) && cells.seq_has(j, seq_id) && cells.pos_get(j) <= q) {
-                    // finite, so it can never meet a -inf and produce a nan
-                    v = cells.pos_get(j) >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f);
-                }
-
-                cur_bias[j] = v;
-            }
-        }
-    }
+    mem->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias);
 }

src/llama-memory-hybrid-idx.h

@@ -75,6 +75,18 @@ public:
 
     llama_kv_cache * get_mem_idx() const;   // nullptr when the model carries no indexer
 
+    // block-compressed sparse attention (qwen4exp QSA) over the cells of the indexer cache.
+    // Blocks cut the position line, not the cell array, so no caller assumes a contiguous layout:
+    //   cell_blk  I32 [n_kv, ns]           block each cell belongs to
+    //   blk_cells I32 [ratio*n_blocks, ns] cells making up each block
+    //   blk_pos   I32 [4*n_blocks*ns]      mrope position rows of each block's first token
+    //   bias      F32 [n_kv, n_tokens/ns, ns] -inf where invisible, large where always visible
+    // blk_bias asks for the bias per block instead: [n_blocks, n_tokens/ns, ns]
+    // the caller then adds the attention mask, the only part of the bias that varies within a block
+    void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos,
+                       ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio,
+                       bool blk_bias) const;
+
 private:
     // forget seq_id (all of it if seq_id < 0) in every cache at once, so a failed restore cannot leave the caches out of step
     // seq_id < 0 drops the whole context, as the caches themselves do on a failed restore
@@ -123,20 +135,12 @@ public:
     // llama_memory_hybrid_idx_context specific API
     //
 
-    // nullptr with no indexer, and for the update context, which builds no sparse graph
+    // nullptr with no indexer
     const llama_kv_cache_context * get_idx() const;
 
     // streams in the current slot info, the `ns` of get_k/get_v; 1 if unified
     uint32_t get_n_stream() const;
 
-    // block-compressed sparse attention (qwen4exp QSA) over the cells of the indexer cache.
-    // Blocks cut the position line, not the cell array, so no caller assumes a contiguous layout:
-    //   cell_blk  I32 [n_kv, ns]           block each cell belongs to
-    //   blk_cells I32 [ratio*n_blocks, ns] cells making up each block
-    //   blk_pos   I32 [4*n_blocks*ns]      mrope position rows of each block's first token
-    //   bias      F32 [n_kv, n_tokens/ns, ns] -inf where invisible, large where always visible
-    // blk_bias asks for the bias per block instead: [n_blocks, n_tokens/ns, ns]
-    // the caller then adds the attention mask, the only part of the bias that varies within a block
     void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos,
                        ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio,
                        bool blk_bias) const;
@@ -148,7 +152,7 @@ private:
     // declared first, so it is initialised while sinfos_idx is still intact
     const std::vector<uint32_t> ns_ubatch;
 
-    // null unless the model has an indexer and this is a batch or full context
+    // null unless the model has an indexer
     const llama_memory_context_ptr ctx_idx;
 
     // mirrors the base class's ubatch cursor, which is private there

src/llama-quant.cpp

@@ -742,12 +742,28 @@ static ggml_type llama_tensor_get_type(quantize_state_impl & qs, const llama_mod
 // quantization implementation
 //
 
-static size_t llama_tensor_quantize_impl(enum ggml_type new_type, const float * f32_data, void * new_data, const int64_t chunk_size, int64_t nrows, int64_t n_per_row, const float * imatrix, std::vector<std::thread> & workers, const int nthread) {
+// quantize rows [first_row, first_row + nrows), indexed globally across all expert matrices
+// note: chunks never cross an expert boundary since each expert has its own imatrix slice
+static size_t llama_tensor_quantize_impl(enum ggml_type new_type, const float * f32_data, void * new_data, const int64_t chunk_size, int64_t first_row, int64_t nrows, int64_t nrows_per_expert, int64_t n_per_row, const float * imatrix, std::vector<std::thread> & workers, const int nthread) {
+    const size_t row_size = ggml_row_size(new_type, n_per_row);
+
+    auto imatrix_for_row = [=](int64_t row_global) {
+        return imatrix ? imatrix + (row_global / nrows_per_expert) * n_per_row : nullptr;
+    };
+
     if (nthread < 2) {
         // single-thread
-        size_t new_size = ggml_quantize_chunk(new_type, f32_data, new_data, 0, nrows, n_per_row, imatrix);
-        if (!ggml_validate_row_data(new_type, new_data, new_size)) {
-            throw std::runtime_error("quantized data validation failed");
+        size_t new_size = 0;
+        for (int64_t row = 0; row < nrows;) {
+            const int64_t row_global = first_row + row;
+            const int64_t this_nrow  = std::min(nrows - row, nrows_per_expert - row_global % nrows_per_expert);
+            void * this_data = (char *) new_data + row * row_size;
+            size_t this_size = ggml_quantize_chunk(new_type, f32_data + row * n_per_row, this_data, 0, this_nrow, n_per_row, imatrix_for_row(row_global));
+            if (!ggml_validate_row_data(new_type, this_data, this_size)) {
+                throw std::runtime_error("quantized data validation failed");
+            }
+            new_size += this_size;
+            row += this_nrow;
         }
         return new_size;
     }
@@ -757,26 +773,29 @@ static size_t llama_tensor_quantize_impl(enum ggml_type new_type, const float *
     size_t new_size = 0;
     bool valid = true;
     auto compute = [&mutex, &counter, &new_size, &valid, new_type, f32_data, new_data, chunk_size,
-            nrows, n_per_row, imatrix]() {
+            first_row, nrows, nrows_per_expert, n_per_row, row_size, imatrix_for_row]() {
         const int64_t nrows_per_chunk = chunk_size / n_per_row;
         size_t local_size = 0;
         while (true) {
             std::unique_lock<std::mutex> lock(mutex);
-            int64_t first_row = counter; counter += nrows_per_chunk;
-            if (first_row >= nrows) {
+            if (counter >= nrows) {
                 if (local_size > 0) {
                     new_size += local_size;
                 }
                 break;
             }
+            const int64_t row        = counter;
+            const int64_t row_global = first_row + row;
+            // stop at the expert boundary
+            const int64_t this_nrow  = std::min(std::min(nrows - row, nrows_per_chunk), nrows_per_expert - row_global % nrows_per_expert);
+            counter += this_nrow;
             lock.unlock();
-            const int64_t this_nrow = std::min(nrows - first_row, nrows_per_chunk);
-            size_t this_size = ggml_quantize_chunk(new_type, f32_data, new_data, first_row * n_per_row, this_nrow, n_per_row, imatrix);
+
+            void * this_data = (char *) new_data + row * row_size;
+            size_t this_size = ggml_quantize_chunk(new_type, f32_data + row * n_per_row, this_data, 0, this_nrow, n_per_row, imatrix_for_row(row_global));
             local_size += this_size;
 
             // validate the quantized data
-            const size_t row_size  = ggml_row_size(new_type, n_per_row);
-            void * this_data = (char *) new_data + first_row * row_size;
             if (!ggml_validate_row_data(new_type, this_data, this_size)) {
                 std::unique_lock<std::mutex> lock(mutex);
                 valid = false;
@@ -1258,52 +1277,49 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
                 fflush(stdout);
 
                 const int64_t n_per_row = tensor->ne[0];
-                const int64_t nrows = tensor->ne[1];
+                const int64_t nrows_per_expert = tensor->ne[1];
+                const int64_t nrows_total = tensor->ne[1] * tensor->ne[2];
 
                 const size_t row_size_src = ggml_row_size(tensor->type, n_per_row);
                 const size_t row_size_dst = ggml_row_size(new_type,     n_per_row);
 
                 // process the rows in slabs, so that the buffers stay below max_buf_size
                 const size_t bytes_per_row = row_size_src + row_size_dst + (tensor->type == GGML_TYPE_F32 ? 0 : n_per_row*sizeof(float));
-                const int64_t nrows_slab = std::max<int64_t>(1, std::min<int64_t>(nrows, max_buf_size/bytes_per_row));
+                const int64_t nrows_slab = std::max<int64_t>(1, std::min<int64_t>(nrows_total, max_buf_size/bytes_per_row));
 
                 static const int64_t min_chunk_size = 32 * 512;
                 const int64_t chunk_size = (n_per_row >= min_chunk_size ? n_per_row : n_per_row * ((min_chunk_size + n_per_row - 1)/n_per_row));
 
-                // quantize each expert separately since they have different importance matrices
+                // process rows across all experts in one pass to keep all threads busy
                 new_size = 0;
-                for (int64_t i03 = 0; i03 < tensor->ne[2]; ++i03) {
-                    const float * imatrix_03 = imatrix ? imatrix + i03 * n_per_row : nullptr;
+                for (int64_t ir = 0; ir < nrows_total; ir += nrows_slab) {
+                    const int64_t nrows_cur = std::min(nrows_slab, nrows_total - ir);
+                    const int64_t nelements_cur = nrows_cur * n_per_row;
 
-                    for (int64_t ir = 0; ir < nrows; ir += nrows_slab) {
-                        const int64_t nrows_cur = std::min(nrows_slab, nrows - ir);
-                        const int64_t nelements_cur = nrows_cur * n_per_row;
+                    const void * src = load_range(ir*row_size_src, nrows_cur*row_size_src);
 
-                        const void * src = load_range((i03*nrows + ir)*row_size_src, nrows_cur*row_size_src);
-
-                        const float * f32_data;
-                        if (tensor->type == GGML_TYPE_F32) {
-                            f32_data = (const float *) src;
-                        } else {
-                            if (f32_conv_buf.size() < (size_t) nelements_cur) {
-                                f32_conv_buf.resize(nelements_cur);
-                            }
-                            llama_tensor_dequantize_impl(tensor->type, src, (float *) f32_conv_buf.data(), workers, nelements_cur, nthread);
-                            f32_data = (const float *) f32_conv_buf.data();
+                    const float * f32_data;
+                    if (tensor->type == GGML_TYPE_F32) {
+                        f32_data = (const float *) src;
+                    } else {
+                        if (f32_conv_buf.size() < (size_t) nelements_cur) {
+                            f32_conv_buf.resize(nelements_cur);
                         }
-
-                        if (work.size() < nrows_cur*row_size_dst) {
-                            work.resize(nrows_cur*row_size_dst);
-                        }
-
-                        const int64_t nchunk = (nelements_cur + chunk_size - 1)/chunk_size;
-                        const int64_t nthread_use = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, nchunk)) : 1;
-
-                        const size_t size_cur = llama_tensor_quantize_impl(new_type, f32_data, work.data(), chunk_size, nrows_cur, n_per_row, imatrix_03, workers, nthread_use);
-
-                        fout.write((const char *) work.data(), size_cur);
-                        new_size += size_cur;
+                        llama_tensor_dequantize_impl(tensor->type, src, (float *) f32_conv_buf.data(), workers, nelements_cur, nthread);
+                        f32_data = (const float *) f32_conv_buf.data();
                     }
+
+                    if (work.size() < nrows_cur*row_size_dst) {
+                        work.resize(nrows_cur*row_size_dst);
+                    }
+
+                    const int64_t nchunk = (nelements_cur + chunk_size - 1)/chunk_size;
+                    const int64_t nthread_use = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, nchunk)) : 1;
+
+                    const size_t size_cur = llama_tensor_quantize_impl(new_type, f32_data, work.data(), chunk_size, ir, nrows_cur, nrows_per_expert, n_per_row, imatrix, workers, nthread_use);
+
+                    fout.write((const char *) work.data(), size_cur);
+                    new_size += size_cur;
                 }
                 LLAMA_LOG_INFO("size = %8.2f MiB -> %8.2f MiB\n", tensor_size/1024.0/1024.0, new_size/1024.0/1024.0);
             }

src/models/dflash.cpp

@@ -257,9 +257,10 @@ std::unique_ptr<llm_graph_context> llama_model_dflash::build_arch_graph(const ll
 
 template <>
 ggml_tensor * llama_model_dflash::graph<true>::build_inp_embd_enc() const {
-    auto inp_target = std::make_unique<llm_graph_input_embd>(hparams.n_embd_inp_enc());
+    const int64_t n_embd_inp = hparams.n_embd_inp_enc();
+    auto inp_target = std::make_unique<llm_graph_input_embd>(n_embd_inp);
 
-    inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp_enc(), n_tokens);
+    inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
     ggml_set_input(inp_target->embd);
 
     ggml_tensor * cur = inp_target->embd;
@@ -567,6 +568,7 @@ static void build_dflash2_selector(llm_graph_context & g, const llama_model & mo
 //   * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens
 template <>
 llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
+    const int64_t n_embd_inp = hparams.n_embd_inp_enc();
     const int64_t n_embd_head = hparams.n_embd_head_v();
 
     GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
@@ -602,16 +604,21 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
 
     // KV cache injection
     if (ubatch.embd) {
-        auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
+        auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
 
-        inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
+        inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
         ggml_set_input(inp->embd);
 
-        ggml_tensor * inp_g = inp->embd;
-        cb(inp_g, "inp_g_embeddings", -1);
+        ggml_tensor * inp_target = inp->embd;
+        cb(inp_target, "inp_target_features", -1);
 
         res->add_input(std::move(inp));
 
+        // fuse the target features through the encoder
+        ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);
+        inp_g = build_norm(inp_g, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);
+        cb(inp_g, "inp_g_embeddings", -1);
+
         for (int il = 0; il < n_layer; ++il) {
             const auto & layer = model.layers[il];
 
@@ -823,6 +830,7 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
 //   * token batch -> noise block through 3 full DSV4 stages (hc + MLA + MoE), markov + confidence heads
 llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_graph_params & params) :
     llama_model_deepseek4::graph(params) {
+    const int64_t n_embd_inp       = hparams.n_embd_inp_enc();
     const int64_t n_embd_head      = hparams.n_embd_head_k();
     const int64_t n_embd_head_rope = hparams.n_rot();
     const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
@@ -833,16 +841,21 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
 
     // KV cache injection: fused target features from the encoder
     if (ubatch.embd) {
-        auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
+        auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
 
-        inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
+        inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
         ggml_set_input(inp->embd);
 
-        ggml_tensor * inp_g = inp->embd;
-        cb(inp_g, "inp_g_embeddings", -1);
+        ggml_tensor * inp_target = inp->embd;
+        cb(inp_target, "inp_target_features", -1);
 
         res->add_input(std::move(inp));
 
+        // fuse the target features through the encoder
+        ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);
+        inp_g = build_norm(inp_g, model.output_norm_enc, nullptr, LLM_NORM_RMS, -1);
+        cb(inp_g, "inp_g_embeddings", -1);
+
         for (int il = 0; il < n_layer; ++il) {
             const auto & layer = model.layers[il];
 

src/models/qwen4exp.cpp

@@ -6,6 +6,23 @@
 #include <algorithm>
 #include <cinttypes>
 
+// bad metadata must be catchable: GGML_ASSERT aborts the whole process
+static void qwen4exp_require_nonzero(const llama_model_loader & ml, llm_kv kid, uint32_t value) {
+    if (value == 0) {
+        throw std::runtime_error(format("%s must be greater than zero, got %u", ml.llm_kv(kid).c_str(), value));
+    }
+}
+
+// get_arr() copies a short array as-is, leaving a zero tail the n-gram hash silently drops
+static void qwen4exp_require_arr_len(llama_model_loader & ml, llm_kv kid, uint32_t n_min) {
+    uint32_t n_arr = 0;
+    ml.get_arr_n(kid, n_arr, true);
+    if (n_arr < n_min) {
+        throw std::runtime_error(format("%s has %u entries, but at least %u are required",
+                                        ml.llm_kv(kid).c_str(), n_arr, n_min));
+    }
+}
+
 void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,        hparams.n_ff_exp, false);
     ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
@@ -18,21 +35,30 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key(LLM_KV_SSM_STATE_SIZE,     hparams.ssm_d_state);
     ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
     ml.get_key(LLM_KV_SSM_GROUP_COUNT,    hparams.ssm_n_group);
-    GGML_ASSERT(hparams.ssm_d_conv  > 0 && hparams.ssm_d_inner > 0 && hparams.ssm_d_state > 0 &&
-                hparams.ssm_dt_rank > 0 && hparams.ssm_n_group > 0);
+    qwen4exp_require_nonzero(ml, LLM_KV_SSM_CONV_KERNEL,    hparams.ssm_d_conv);
+    qwen4exp_require_nonzero(ml, LLM_KV_SSM_INNER_SIZE,     hparams.ssm_d_inner);
+    qwen4exp_require_nonzero(ml, LLM_KV_SSM_STATE_SIZE,     hparams.ssm_d_state);
+    qwen4exp_require_nonzero(ml, LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
+    qwen4exp_require_nonzero(ml, LLM_KV_SSM_GROUP_COUNT,    hparams.ssm_n_group);
 
     // HC; low_rank is qwen4exp-specific, DeepSeek-V4 leaves it absent (full rank)
     ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT,    hparams.dsv4_hc_mult);
     ml.get_key(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank);
-    GGML_ASSERT(hparams.dsv4_hc_mult > 0 && hparams.hc_low_rank > 0);
+    // a count of 1 has nothing to mix: transformers configuration_qwen4_exp.py:196, vLLM
+    // config.py:49 and SGLang configs/qwen4_exp.py:38 all raise on hc_count <= 1
+    if (hparams.dsv4_hc_mult <= 1) {
+        throw std::runtime_error(format("%s must be greater than one, got %u",
+                                        ml.llm_kv(LLM_KV_HYPER_CONNECTION_COUNT).c_str(), hparams.dsv4_hc_mult));
+    }
+    qwen4exp_require_nonzero(ml, LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank);
     hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd;
 
     ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);
     ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
     ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K,      hparams.indexer_top_k);
-    GGML_ASSERT(hparams.indexer_n_head > 0
-             && hparams.indexer_head_size > 0
-             && hparams.indexer_top_k > 0);
+    qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);
+    qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
+    qwen4exp_require_nonzero(ml, LLM_KV_ATTENTION_INDEXER_TOP_K,      hparams.indexer_top_k);
     ml.get_key_or_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, hparams.n_layer_all, false);
 
     // PLE n-gram hash embeddings; if the key group is absent every field stays zero
@@ -44,7 +70,11 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
     if (n_ple > 0) {
         std::vector<uint32_t> ple_layers;
         ml.get_arr(LLM_KV_PLE_LAYERS, ple_layers);
-        GGML_ASSERT(n_ple == 1 && "qwen4exp supports only one PLE layer");
+        if (n_ple != 1) {
+            // hparams holds one set of hash constants, so several PLE modules cannot be represented
+            throw std::runtime_error(format("%s lists %u layers, but only one PLE layer is supported",
+                                            ml.llm_kv(LLM_KV_PLE_LAYERS).c_str(), n_ple));
+        }
         for (uint32_t il : ple_layers) {
             if (il >= hparams.n_layer_all) {
                 throw std::runtime_error(format("PLE layer %u is out of range", il));
@@ -59,7 +89,8 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
         // optional: files written before this key fall back to the EOS token
         ml.get_key(LLM_KV_PLE_IMAGE_TOKEN_ID,  hparams.ple_image_token_id, false);
         ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer);
-        GGML_ASSERT(hparams.ple_conv_kernel > 0 && hparams.n_embd_per_layer > 0);
+        qwen4exp_require_nonzero(ml, LLM_KV_PLE_CONV_KERNEL,             hparams.ple_conv_kernel);
+        qwen4exp_require_nonzero(ml, LLM_KV_EMBEDDING_LENGTH_PER_LAYER,  hparams.n_embd_per_layer);
 
         hparams.ple_n_heads  = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram;
         hparams.ple_head_dim = hparams.n_embd_per_layer;
@@ -70,6 +101,10 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
             throw std::runtime_error(format("PLE head count %u is out of range", hparams.ple_n_heads));
         }
 
+        qwen4exp_require_arr_len(ml, LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_ngram_size);
+        qwen4exp_require_arr_len(ml, LLM_KV_PLE_HEAD_OFFSETS,      hparams.ple_n_heads);
+        qwen4exp_require_arr_len(ml, LLM_KV_PLE_HEAD_VOCAB_SIZES,  hparams.ple_n_heads);
+
         ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_layer_multipliers);
 
         // the file stores the head ranges as uint64, so read at that width and narrow to the int32 the gather uses
@@ -93,12 +128,19 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) {
     if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {
         uint32_t full_attn_interval = 4;
         ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
-        GGML_ASSERT(full_attn_interval > 0);
+        qwen4exp_require_nonzero(ml, LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval);
         for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
             hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);
         }
     }
 
+    // the PLE conv history is a row of the recurrent cache, which linear layers alone have
+    for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
+        if (hparams.is_ple(i) && !hparams.is_recr(i)) {
+            throw std::runtime_error(format("PLE layer %u is not a linear attention layer", i));
+        }
+    }
+
     switch (hparams.n_layer()) {
         case 48: type = LLM_TYPE_A3B; break;
         default: type = LLM_TYPE_UNKNOWN;
@@ -124,18 +166,24 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
         output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);
     }
 
-    // flat [ple_head_dim, n_rows] gather target; n_rows is padded, so read it back
+    // flat [ple_head_dim, n_rows] gather target
     if (hparams.ple_n_heads > 0) {
-        const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str();
-        const auto & ple_w = ml.require_weight(ple_name.c_str());
-        const int64_t ple_rows = ple_w.tensor->ne[1];
-
-        // sanity check
+        // the head ranges are what the gather indexes, so they set the minimum row count
+        int64_t ple_rows = 0;
         for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) {
-            if ((int64_t) hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h] > ple_rows) {
-                throw std::runtime_error(format("PLE head %u range exceeds the %" PRId64 " table rows", h, ple_rows));
-            }
+            ple_rows = std::max(ple_rows, (int64_t) hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h]);
         }
+
+        // the converter pads the table; a model synthesised from metadata has no tensor to ask
+        const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str();
+        if (const auto * ple_w = ml.get_weight(ple_name.c_str())) {
+            if (ple_w->tensor->ne[1] < ple_rows) {
+                throw std::runtime_error(format("%s has %" PRId64 " rows, too few for the PLE head ranges (%" PRId64 ")",
+                                                ple_name.c_str(), ple_w->tensor->ne[1], ple_rows));
+            }
+            ple_rows = ple_w->tensor->ne[1];
+        }
+
         per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"),
                                            { hparams.ple_head_dim, ple_rows }, TENSOR_READ_LAZY);
     }
@@ -556,9 +604,12 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
     pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r);
     cb(pooled, "indexer_k_pooled", il);
 
+    // count blocks along ne1: rms_norm launches gridDim.y = ne2, capped at 65535, and 262144/4 = 65536
+    pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks*n_stream, 1);
+    pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il);
+
     // rope wants [n_dims, n_head, n_tokens]: lay every stream's blocks flat, split after.
     pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks*n_stream);
-    pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il);
     pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr,
             n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
             ext_factor, attn_factor, beta_fast, beta_slow);
@@ -576,12 +627,19 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k(
     // rectify each head dot product before the sum, as in the DeepSeek lightning indexer
     // mul_mat matches ne[2], so the queries of stream s only meet the blocks of stream s
     ggml_tensor * score = ggml_mul_mat(ctx0, pooled,
-            ggml_reshape_3d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tps, n_stream));
+            ggml_reshape_3d(ctx0, q, idx_dim, n_idx_h*n_tps, n_stream));
     score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream);
     score = ggml_relu(ctx0, score);
-    score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3));
-    score = ggml_sum_rows(ctx0, score);
-    score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream);
+
+    // the heads sit side by side on ne[1] and there are only a few of them
+    ggml_tensor * summed = nullptr;
+    for (int64_t h = 0; h < n_idx_h; ++h) {
+        ggml_tensor * slice = ggml_view_3d(ctx0, score, n_blocks, n_tps, n_stream,
+                score->nb[2], score->nb[3], h*score->nb[1]);
+        summed = summed ? ggml_add(ctx0, summed, slice) : ggml_cont(ctx0, slice);
+    }
+
+    score = summed;
     cb(score, "indexer_score", il);
 
     // one value per block, so it is cheaper to bias here than after the cells are expanded
@@ -1080,20 +1138,28 @@ ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at(
 
     ggml_tensor * conv_input = ggml_concat(ctx0, state, ggml_transpose(ctx0, x), 0);
 
-    // keep the last state_cols columns for the next ubatch
+    // [TAG_RECURRENT_ROLLBACK_SPLITS] keep the last state_cols columns once per rollback slot,
+    // slot s ending s tokens earlier so a rollback of s tokens reads a history that never saw them
     const size_t row_size = ggml_row_size(conv_states_all->type, row_total);
+    const uint32_t mem_size = mctx_cur->get_size();
 
-    ggml_tensor * tail = ggml_view_3d(ctx0, conv_input,
-            state_cols, channels, n_seqs,
-            conv_input->nb[1], conv_input->nb[2],
-            ggml_row_size(conv_input->type, conv_input->ne[0] - state_cols));
+    const int64_t n_slots = (int64_t) cparams.n_rs_seq + 1;
 
-    ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all,
-            state_cols * channels, n_seqs,
-            conv_states_all->nb[1],
-            kv_head * row_size);
+    for (int64_t slot = 0; slot < n_slots; ++slot) {
+        const int64_t s_idx = std::max<int64_t>(0, conv_input->ne[0] - state_cols - slot);
 
-    ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_cont(ctx0, tail), dst));
+        ggml_tensor * tail = ggml_view_3d(ctx0, conv_input,
+                state_cols, channels, n_seqs,
+                conv_input->nb[1], conv_input->nb[2],
+                ggml_row_size(conv_input->type, s_idx));
+
+        ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all,
+                state_cols * channels, n_seqs,
+                conv_states_all->nb[1],
+                (slot * mem_size + kv_head) * row_size);
+
+        ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_cont(ctx0, tail), dst));
+    }
 
     return conv_input;
 }

tests/CMakeLists.txt

@@ -299,6 +299,14 @@ if (NOT LLAMA_SANITIZE_ADDRESS AND NOT GGML_SCHED_NO_REALLOC)
 endif()
 llama_build_and_test(test-backend-ops.cpp)
 
+# the tensor API kernels come from a separate metallib - check they produce correct results
+# ref: https://github.com/ggml-org/llama.cpp/issues/27473
+if (GGML_METAL AND NOT GGML_METAL_EMBED_LIBRARY)
+    llama_test(test-backend-ops NAME test-backend-ops-metallib-tensor
+               ARGS test -b MTL0 -o MUL_MAT -p type_a=q6_K)
+    set_tests_properties(test-backend-ops-metallib-tensor PROPERTIES ENVIRONMENT GGML_METAL_TENSOR_ENABLE=1)
+endif()
+
 llama_build_and_test(test-model-load-cancel.cpp LABEL "model")
 llama_build_and_test(test-autorelease.cpp       LABEL "model")
 llama_build_and_test(test-backend-sampler.cpp   LABEL "model")

tests/test-backend-ops.cpp

@@ -6287,6 +6287,87 @@ struct test_top_k : public test_case {
     }
 };
 
+// qwen4exp QSA indexer top-k fusion: expand per-block scores to cells, add the f16 mask, top-k.
+struct test_topk_qsa : public test_case {
+    const int64_t n_blocks;
+    const int64_t n_kv;
+    const int64_t n_tps;
+    const int64_t n_stream;
+    const int     width;
+    ggml_tensor * out {};
+
+    std::string op_desc(ggml_tensor * t) override {
+        GGML_UNUSED(t);
+        return "TOPK_QSA";
+    }
+
+    std::string vars() override {
+        return VARS_TO_STR5(n_blocks, n_kv, n_tps, n_stream, width);
+    }
+
+    test_topk_qsa(int64_t n_blocks = 512, int64_t n_kv = 2048, int64_t n_tps = 2, int64_t n_stream = 1, int width = 1500)
+        : n_blocks(n_blocks), n_kv(n_kv), n_tps(n_tps), n_stream(n_stream), width(width) {}
+
+    double max_err() override { return 0.0; }
+    bool run_whole_graph() override { return true; }
+
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_blocks, n_tps, n_stream);
+        ggml_set_name(score, "score");
+        ggml_tensor * cell_blk = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_kv, n_stream);
+        ggml_set_name(cell_blk, "cell_blk");
+        ggml_tensor * kq_mask = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, n_kv, n_tps, n_stream);
+        ggml_set_name(kq_mask, "kq_mask");
+
+        ggml_tensor * a = ggml_cont(ctx, ggml_permute(ctx, score, 1, 0, 2, 3));
+        ggml_tensor * e = ggml_get_rows(ctx, a, cell_blk);
+        e = ggml_cont(ctx, ggml_permute(ctx, e, 1, 0, 2, 3));
+        ggml_tensor * m = ggml_cast(ctx, kq_mask, GGML_TYPE_F32);
+        e = ggml_add(ctx, e, ggml_reshape_3d(ctx, m, n_kv, n_tps, n_stream));
+        out = ggml_top_k(ctx, e, width);
+        ggml_set_name(out, "out");
+        return out;
+    }
+
+    std::vector<ggml_tensor *> fusion_test_nodes() override { return { out }; }
+
+    // distinct mask ramp + small scores keep every cell value unique, so no top-k ties
+    void initialize_tensors(ggml_context * ctx) override {
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->op != GGML_OP_NONE) {
+                continue;
+            }
+            if (t->type == GGML_TYPE_I32) {
+                std::vector<int32_t> data(ggml_nelements(t));
+                for (auto & v : data) { v = rand() % n_blocks; }
+                ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(int32_t));
+            } else if (t->type == GGML_TYPE_F16) {
+                std::vector<ggml_fp16_t> data(ggml_nelements(t));
+                for (int64_t r = 0; r < ggml_nrows(t); r++) {
+                    for (int64_t i = 0; i < n_kv; i++) {
+                        data[r * n_kv + i] = ggml_fp32_to_fp16((float) i);
+                    }
+                }
+                ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(ggml_fp16_t));
+            } else {
+                init_tensor_uniform(t, 0.0f, 0.5f);
+            }
+        }
+    }
+
+    // top-k output order is unspecified; compare as a set of indices
+    double err(const float * a, const float * b, size_t n) override {
+        std::vector<int32_t> ia(n), ib(n);
+        double diff = 0.0;
+        for (size_t i = 0; i < n; i++) {
+            ia[i] = (int32_t) a[i];
+            ib[i] = (int32_t) b[i];
+            diff += std::fabs(a[i] - ia[i]) + std::fabs(b[i] - ib[i]);
+        }
+        return diff + jdst(ia.data(), ib.data(), n);
+    }
+};
+
 enum MoeGatingFunc {
     GATING_FUNC_SOFTMAX,
     GATING_FUNC_SIGMOID,
@@ -9339,6 +9420,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
             test_cases.emplace_back(new test_mul_mat(type_a, type_b, 16, 1, 256, {1,  1}, {1, 1}));
         }
     }
+
+    // Test IQP panel path for all grid IQ types
+    for (ggml_type type_a : {GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS, GGML_TYPE_IQ2_S, GGML_TYPE_IQ3_XXS,
+                             GGML_TYPE_IQ3_S, GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M, GGML_TYPE_IQ4_XS}) {
+        test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 16, 10, 256, {1, 1}, {1, 1}));
+    }
 #else
     // m = a rows
     // n = b rows
@@ -9450,6 +9537,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 4, 2, false, 64, 16, 3*ggml_blck_size(type_a)));
     }
 
+    // Test IQP panel path for all grid IQ types
+    for (ggml_type type_a : {GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS, GGML_TYPE_IQ2_S, GGML_TYPE_IQ3_XXS,
+                             GGML_TYPE_IQ3_S, GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M, GGML_TYPE_IQ4_XS}) {
+        test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 4, 4, false, 16, 10, 256));
+    }
+
     for (ggml_type type_a : base_types) {
         for (ggml_type type_b : {GGML_TYPE_F32 /*, GGML_TYPE_F16 */}) {
             for (int n_mats : {4, 8}) {
@@ -9813,6 +9906,22 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2049, 2, 1, 3}, k));
     }
 
+    // Large-k, including multi-row and ties (qwen4exp)
+    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, { 1024,  1, 1, 1 }, 1024));
+    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, { 2048,  2, 1, 1 }, 1024));
+    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, { 4096,  1, 1, 1 }, 2048));
+    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, { 8192,  2, 1, 1 }, 2051));
+    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, { 33024, 1, 1, 1 }, 2051));
+    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, { 33024, 4, 1, 1 }, 2051));
+    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, { 8192,  2, 1, 1 }, 2051, true));
+    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, { 33024, 4, 1, 1 }, 2051, true));
+
+    // qwen4exp QSA indexer top-k fusion (get_rows + f16 mask + top_k)
+    test_cases.emplace_back(new test_topk_qsa(512,  2048,  1, 1, 1500));
+    test_cases.emplace_back(new test_topk_qsa(512,  2048,  2, 1, 1500));
+    test_cases.emplace_back(new test_topk_qsa(256,  2048,  4, 2, 2000));
+    test_cases.emplace_back(new test_topk_qsa(64,   256,   2, 1, 200));  // small k: unfused fallback
+
     // exhaustive top_k tests
     //for (int i = 1; i < 9999; ++i) {
     //    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {i, 2, 1, 3}, rand() % i + 1));
@@ -10074,6 +10183,16 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 1024, 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, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, false));
 
+    // FLASH_ATTN_EXT MMA: non-pow2 head size and MLA K/V view.
+    test_cases.emplace_back(new test_flash_attn_ext(192, 128, 8, {8, 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(576, 512, 1, {20, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true));
+
+    // FLASH_ATTN_EXT MMA, swizzled K/V tiles, power-of-two stride: nbatch_K2 = 32, 64, 128, 256.
+    test_cases.emplace_back(new test_flash_attn_ext( 64,  64, 8, {8, 1}, 4096,  4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(128, 128, 8, {4, 1}, 4096,  8, true,  true, 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}, 1024, 32, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(512, 512, 4, {2, 1}, 1024,  4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+
     test_cases.emplace_back(new test_cross_entropy_loss     (GGML_TYPE_F32, {   10, 5, 4, 3}));
     test_cases.emplace_back(new test_cross_entropy_loss     (GGML_TYPE_F32, {30000, 1, 1, 1}));
     test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, {   10, 5, 4, 3}));
@@ -10108,12 +10227,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
                                     use_id, 16, 8, b, with_bias, with_gate, with_lane_scale));
                                 test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 1, 32, 256,
                                     use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}));
-                                if (!use_id && with_gate && !with_bias && glu_op != GGML_GLU_OP_SWIGLU_CLAMP) {
-                                    // small multi-token batches (speculative decoding / MTP verify)
-                                    for (int64_t m_batch : { 2, 4, 8 }) {
-                                        test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, m_batch, 32, 256,
-                                            use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}));
-                                    }
+                                // multi-token batches (spec decoding)
+                                for (int64_t m_batch : { 2, 4, 8 }) {
+                                    test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, m_batch, 32, 256,
+                                        use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}));
                                 }
                             }
                         }
@@ -10142,6 +10259,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
                     test_cases.emplace_back(new test_topk_moe({160, 4, 1, 1}, 160, with_norm, bias_probs, gate, scale_w));
                     test_cases.emplace_back(new test_topk_moe({256, 22, 1, 1}, 6, with_norm, bias_probs, gate, scale_w)); // Used by DeepSeek-V4
                     test_cases.emplace_back(new test_topk_moe({288, 22, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); // Used by StepFun 3.7
+                    // rows at and just past the limit where one block still covers all rows
+                    test_cases.emplace_back(new test_topk_moe({32, 8, 1, 1}, 4, with_norm, bias_probs, gate, scale_w));
+                    test_cases.emplace_back(new test_topk_moe({32, 8, 1, 1}, 8, with_norm, bias_probs, gate, scale_w));
+                    test_cases.emplace_back(new test_topk_moe({32, 9, 1, 1}, 8, with_norm, bias_probs, gate, scale_w));
                 }
             }
         }
@@ -10470,10 +10591,14 @@ 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));
 
-    for (int kv : { 4096, 8192, 16384, }) {
-        for (int hs : { 64, 128, }) {
-            for (int nr : { 1, 4, }) {
-                test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 8, {nr, 1}, kv, 1, 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;
+            const bool v_view = hs == 576;
+            for (int nr : { 1, 4, 8, }) {
+                for (int nb : { 1, 4096, }) {
+                    test_cases.emplace_back(new test_flash_attn_ext(hs, hsv, 8, {nr, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, v_view));
+                }
             }
         }
     }

tests/test-llama-archs.cpp

@@ -65,7 +65,7 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) {
 }
 
 static void usage(char ** argv) {
-    printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v/--verbose] [-h/--help]\n", argv[0]);
+    printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help]\n", argv[0]);
 }
 
 static std::vector<llama_token> get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){
@@ -254,6 +254,30 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
         ms.add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, uint32_t(8));
         // without this the QSA layers fall back to dense and go uncovered
         ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector<uint32_t>(n_layer, 4));
+
+        // has_cell_ext() needs ple_n_heads here: the indexer cache serializes no ext without it
+        const uint32_t ple_ngram_size      = 3;
+        const uint32_t ple_heads_per_ngram = 2;
+        const uint32_t ple_n_heads         = (ple_ngram_size - 1)*ple_heads_per_ngram;
+        GGML_ASSERT(n_embd % ple_n_heads == 0);
+        const uint32_t ple_head_dim = n_embd/ple_n_heads;
+
+        std::vector<uint64_t> ple_head_offsets(ple_n_heads);
+        std::vector<uint64_t> ple_head_vocab_sizes(ple_n_heads, n_vocab);
+        for (uint32_t h = 0; h < ple_n_heads; h++) {
+            ple_head_offsets[h] = uint64_t(h)*n_vocab;
+        }
+
+        // the PLE history lives in the recurrent cache, so it must sit on a linear attention layer
+        ms.add_kv(LLM_KV_PLE_LAYERS,                  std::vector<uint32_t>({ 0 }));
+        ms.add_kv(LLM_KV_PLE_NGRAM_SIZE,              ple_ngram_size);
+        ms.add_kv(LLM_KV_PLE_HEADS_PER_NGRAM,         ple_heads_per_ngram);
+        ms.add_kv(LLM_KV_PLE_CONV_KERNEL,             uint32_t(4));
+        ms.add_kv(LLM_KV_PLE_EOS_TOKEN_ID,            uint32_t(0));
+        ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER,  ple_head_dim);
+        ms.add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS,       std::vector<uint64_t>({ 1, 3, 5 }));
+        ms.add_kv(LLM_KV_PLE_HEAD_OFFSETS,            ple_head_offsets);
+        ms.add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES,        ple_head_vocab_sizes);
     }
 
     // minimax-m3 keeps one indexer head per GQA head; the rest use a fixed 64 to match the fused
@@ -535,22 +559,27 @@ static bool arch_supported(const llm_arch arch) {
     return true;
 }
 
-static int save_models(const llm_arch target_arch, const size_t seed, const ggml_log_level log_level, const std::string & dir) {
+static int save_models(const llm_arch target_arch, const size_t seed, const int verbosity, const std::string & dir) {
     struct user_data_t {
         struct {
             ggml_log_callback callback;
             void * user_data;
-        } original_logger;
-        ggml_log_level min_level; // prints below this log level go to debug log
+        } log_old;
+
+        int verbosity;
+
+        user_data_t(int verbosity) : verbosity(verbosity) {
+            llama_log_get(&log_old.callback, &log_old.user_data);
+        }
     };
-    user_data_t ud;
-    llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data);
-    ud.min_level = log_level;
+    user_data_t ud(verbosity);
 
     llama_log_set([](ggml_log_level level, const char * text, void * user_data) {
         const user_data_t * ud = (const user_data_t *) user_data;
-        const ggml_log_level level_eff = level >= ud->min_level ? level : GGML_LOG_LEVEL_DEBUG;
-        ud->original_logger.callback(level_eff, text, ud->original_logger.user_data);
+        int verbosity = common_log_get_verbosity(level);
+        if (verbosity <= ud->verbosity) {
+            ud->log_old.callback(level, text, ud->log_old.user_data);
+        }
     }, &ud);
 
     for (const llm_arch & arch : llm_arch_all()) {
@@ -584,26 +613,31 @@ static int save_models(const llm_arch target_arch, const size_t seed, const ggml
             llama_model_save_to_file(model_and_ctx.first.get(), path.c_str());
         }
     }
-    llama_log_set(ud.original_logger.callback, ud.original_logger.user_data);
+    llama_log_set(ud.log_old.callback, ud.log_old.user_data);
     return 0;
 }
 
-static int test_backends(const llm_arch target_arch, const size_t seed, const ggml_log_level log_level) {
+static int test_backends(const llm_arch target_arch, const size_t seed, const int verbosity) {
     struct user_data_t {
         struct {
             ggml_log_callback callback;
             void * user_data;
-        } original_logger;
-        ggml_log_level min_level; // prints below this log level go to debug log
+        } log_old;
+
+        int verbosity;
+
+        user_data_t(int verbosity) : verbosity(verbosity) {
+            llama_log_get(&log_old.callback, &log_old.user_data);
+        }
     };
-    user_data_t ud;
-    llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data);
-    ud.min_level = log_level;
+    user_data_t ud(verbosity);
 
     llama_log_set([](ggml_log_level level, const char * text, void * user_data) {
         const user_data_t * ud = (const user_data_t *) user_data;
-        const ggml_log_level level_eff = level >= ud->min_level ? level : GGML_LOG_LEVEL_DEBUG;
-        ud->original_logger.callback(level_eff, text, ud->original_logger.user_data);
+        int verbosity = common_log_get_verbosity(level);
+        if (verbosity <= ud->verbosity) {
+            ud->log_old.callback(level, text, ud->log_old.user_data);
+        }
     }, &ud);
 
     const std::vector<llama_token> tokens = get_tokens(128, 128, seed);
@@ -699,6 +733,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg
                 std::string status_nmse      = "\033[1;33mSKIP\033[0m";
                 std::string status_roundtrip = "\033[1;33mSKIP\033[0m";
                 char nmse_str[12] = {0};
+
                 bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty());
                 if (!skip) {
                     if (logits_cpu.empty()) {
@@ -749,20 +784,23 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg
             }
         }
     }
-    llama_log_set(ud.original_logger.callback, ud.original_logger.user_data);
+    llama_log_set(ud.log_old.callback, ud.log_old.user_data);
     return all_ok ? 0 : 1;
 }
 
 int main(int argc, char ** argv) {
-    // FIXME these tests are disabled in the CI for macOS-latest-cmake-arm64 because they are segfaulting
+    // init the logger at max verbosity. filter with a custom callback respecting the user-configure verbosity
+    common_log_set_verbosity_thold(LOG_LEVEL_DEBUG);
     common_init();
+
     std::random_device rd;
 
     llm_arch arch = LLM_ARCH_UNKNOWN;
     size_t seed = rd();
-    ggml_log_level log_level = GGML_LOG_LEVEL_ERROR;
     std::string out;
 
+    int verbosity = LOG_LEVEL_ERROR;
+
     for (int i = 1; i < argc; i++) {
         if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) {
             usage(argv);
@@ -789,9 +827,13 @@ int main(int argc, char ** argv) {
                 return 1;
             }
         }
-        if (strcmp(argv[i], "-v") == 0 || strcmp(argv[i], "--verbose") == 0) {
-            log_level = GGML_LOG_LEVEL_INFO;
-            continue;
+        if (strcmp(argv[i], "-v") == 0) {
+            if (i + 1 < argc) {
+                verbosity = std::stoull(argv[++i]);
+            } else {
+                usage(argv);
+                return 1;
+            }
         }
         if (strcmp(argv[i], "-o") == 0 || strcmp(argv[i], "--out") == 0) {
             if (i + 1 < argc) {
@@ -806,9 +848,9 @@ int main(int argc, char ** argv) {
 
     try {
         if (!out.empty()) {
-            return save_models(arch, seed, log_level, out);
+            return save_models(arch, seed, verbosity, out);
         }
-        return test_backends(arch, seed, log_level);
+        return test_backends(arch, seed, verbosity);
     } catch (const std::exception & err) {
         fprintf(stderr, "encountered runtime error: %s\n", err.what());
         return -1;

tests/test-rpc-multi-server.cpp

@@ -0,0 +1,47 @@
+#include "ggml-alloc.h"
+#include "ggml-backend.h"
+#include "ggml-impl.h"
+#include "ggml-rpc.h"
+#include "ggml.h"
+
+int main(int argc, char ** argv) {
+    GGML_ASSERT(argc == 3);
+    ggml_backend_load_all();
+
+    const char * endpoint_a = argv[1];
+    const char * endpoint_b = argv[2];
+
+    ggml_backend_t backend_a = ggml_backend_rpc_init(endpoint_a, 0);
+    ggml_backend_t backend_b = ggml_backend_rpc_init(endpoint_b, 0);
+    GGML_ASSERT(backend_a != nullptr);
+    GGML_ASSERT(backend_b != nullptr);
+
+    ggml_init_params params = {
+        /* .mem_size   = */ ggml_tensor_overhead() + ggml_graph_overhead_custom(1, false),
+        /* .mem_buffer = */ nullptr,
+        /* .no_alloc   = */ true,
+    };
+    ggml_context * ctx = ggml_init(params);
+    GGML_ASSERT(ctx != nullptr);
+
+    ggml_tensor * tensor = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+    ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend_a);
+    GGML_ASSERT(buffer != nullptr);
+
+    // A remote pointer allocated by server A is not meaningful to server B.
+    ggml_cgraph * graph = ggml_new_graph_custom(ctx, 1, false);
+    graph->nodes[0] = tensor;
+    graph->n_nodes = 1;
+
+    GGML_ASSERT(ggml_backend_graph_compute(backend_b, graph) == GGML_STATUS_SUCCESS);
+    // Wait for server B to finish the graph before the script checks its log.
+    size_t free_mem;
+    size_t total_mem;
+    ggml_backend_rpc_get_device_memory(endpoint_b, 0, &free_mem, &total_mem);
+    GGML_ASSERT(total_mem > 0);
+    ggml_backend_buffer_free(buffer);
+    ggml_free(ctx);
+    ggml_backend_free(backend_b);
+    ggml_backend_free(backend_a);
+    return 0;
+}

tests/test-rpc-multi-server.sh

@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+server=$1
+client=$2
+port_a=$((40000 + $$ % 10000))
+port_b=$((port_a + 1))
+endpoint_a="127.0.0.1:${port_a}"
+endpoint_b="127.0.0.1:${port_b}"
+test_dir=$(mktemp -d)
+
+cleanup() {
+    kill "${pid_a:-}" "${pid_b:-}" 2>/dev/null || true
+    rm -rf "$test_dir"
+}
+trap cleanup EXIT
+
+wait_for_port() {
+    local port=$1
+    for _ in {1..600}; do
+        if (exec 3<>"/dev/tcp/127.0.0.1/$port") 2>/dev/null; then
+            exec 3>&-
+            exec 3<&-
+            return 0
+        fi
+        sleep 0.05
+    done
+    return 1
+}
+
+"$server" --device CPU --host 127.0.0.1 --port "$port_a" >"$test_dir/server-a.log" 2>&1 &
+pid_a=$!
+"$server" --device CPU --host 127.0.0.1 --port "$port_b" >"$test_dir/server-b.log" 2>&1 &
+pid_b=$!
+wait_for_port "$port_a"
+wait_for_port "$port_b"
+
+"$client" "$endpoint_a" "$endpoint_b"
+
+if grep -q "invalid data ptr" "$test_dir/server-b.log"; then
+    cat "$test_dir/server-b.log"
+    exit 1
+fi

tests/test-save-load-state.cpp

@@ -355,7 +355,160 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p
 }
 
 
-// Run the full save/load test suite (tests 1-5) for a single model.
+// Test 6/7: seq copy (scatter)
+// - decode the same prefix on two sequences, interleaving seq 0 cells between the seq 1 cells
+// - save the seq 1 state, free the interleaved seq 0 cells, and restore via the given io path
+// - the restore destination is non-contiguous: scatter reads are batched per contiguous run
+// - save again on the host and compare the two blobs byte for byte
+static bool test_seq_cp_scatter(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, int test_num, bool on_device) {
+    auto params_ctx = common_context_params_to_llama(params);
+    params_ctx.n_ctx      = 256;
+    params_ctx.n_seq_max  = 2;
+    params_ctx.kv_unified = true;
+    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
+
+    LOG("\n=== Test %d: seq copy (%s, scatter) ===\n", test_num, on_device ? "device" : "host");
+
+    const uint32_t flags = on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : LLAMA_STATE_SEQ_FLAGS_NONE;
+
+    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);
+        return llama_decode(ctx.get(), batch.get()) == 0;
+    };
+
+    // seq 0 cells 0,1,4 interleave the seq 1 cells 2,3,5
+    if (!decode_one(tokens[0], 0, 0) ||
+        !decode_one(tokens[1], 1, 0) ||
+        !decode_one(tokens[0], 0, 1) ||
+        !decode_one(tokens[1], 1, 1) ||
+        !decode_one(tokens[2], 2, 0) ||
+        !decode_one(tokens[2], 2, 1)) {
+        LOG_ERR("%s: failed to build interleaved state\n", __func__);
+        return false;
+    }
+
+    const auto get_seq_state = [&](llama_seq_id seq_id, uint32_t fl, std::vector<uint8_t> & state) {
+        const size_t state_size = llama_state_seq_get_size_ext(ctx.get(), seq_id, fl);
+        if (state_size == 0) {
+            LOG_ERR("%s: sequence state is empty\n", __func__);
+            return false;
+        }
+
+        state.resize(state_size);
+        const size_t ncopy = llama_state_seq_get_data_ext(ctx.get(), state.data(), state.size(), seq_id, fl);
+        if (ncopy != state.size()) {
+            LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n",
+                    __func__, ncopy, state.size());
+            return false;
+        }
+
+        return true;
+    };
+
+    // host blob: contains the KV data, used for the byte-for-byte comparison
+    std::vector<uint8_t> state_before;
+    if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) {
+        return false;
+    }
+
+    // save via the io path under test
+    std::vector<uint8_t> state_save;
+    if (!get_seq_state(1, flags, state_save)) {
+        return false;
+    }
+    LOG_TRC("%s: seq 1 saved via %s, %zu bytes\n", __func__, on_device ? "device" : "host", state_save.size());
+
+    // free seq 0's cells so the ring is fragmented: the restore destination (seq 1's interleaved cells) stays non-contiguous
+    if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
+        LOG_ERR("%s: failed to remove sequence 0\n", __func__);
+        return false;
+    }
+
+    // restore via the io path under test
+    const size_t nset = llama_state_seq_set_data_ext(ctx.get(), state_save.data(), state_save.size(), 1, flags);
+    if (nset != state_save.size()) {
+        LOG_ERR("%s: seq set data length %zu does not match expected length %zu\n", __func__, nset, state_save.size());
+        return false;
+    }
+    LOG_TRC("%s: seq 1 restored via %s, %zu bytes\n", __func__, on_device ? "device" : "host", nset);
+
+    std::vector<uint8_t> state_after;
+    if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_after)) {
+        return false;
+    }
+
+    // the blob is serialized in sequence cell order, so identical bytes iff the restore wrote the same KV
+    if (state_before.size() != state_after.size() || memcmp(state_before.data(), state_after.data(), state_before.size()) != 0) {
+        LOG_ERR("\n%s: error: restored KV state is not byte-identical to the saved state\n", __func__);
+        return false;
+    }
+
+    LOG("\nPASS\n");
+    return true;
+}
+
+
+// Test 8: state blob round-trip
+// compares blobs rather than generated text: a partially restored cell still decodes to plausible tokens
+static bool test_state_roundtrip(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) {
+    auto params_ctx = common_context_params_to_llama(params);
+    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
+
+    LOG("\n=== Test 8: state blob round-trip ===\n");
+
+    if (llama_decode(ctx.get(), llama_batch_get_one(const_cast<llama_token *>(tokens.data()), (int32_t) tokens.size()))) {
+        LOG_ERR("\n%s: failed to decode prompt\n", __func__);
+        return false;
+    }
+
+    std::vector<uint8_t> blob_a(llama_state_seq_get_size(ctx.get(), 0));
+    const size_t n_a = llama_state_seq_get_data(ctx.get(), blob_a.data(), blob_a.size(), 0);
+    if (n_a != blob_a.size()) {
+        LOG_ERR("\n%s: saved %zu bytes, expected %zu\n", __func__, n_a, blob_a.size());
+        return false;
+    }
+
+    if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
+        LOG_ERR("\n%s: failed to erase seq 0\n", __func__);
+        return false;
+    }
+
+    if (llama_state_seq_set_data(ctx.get(), blob_a.data(), blob_a.size(), 0) != blob_a.size()) {
+        LOG_ERR("\n%s: failed to restore seq 0\n", __func__);
+        return false;
+    }
+
+    std::vector<uint8_t> blob_b(llama_state_seq_get_size(ctx.get(), 0));
+    const size_t n_b = llama_state_seq_get_data(ctx.get(), blob_b.data(), blob_b.size(), 0);
+    if (n_b != n_a) {
+        LOG_ERR("\n%s: re-saved %zu bytes, expected %zu\n", __func__, n_b, n_a);
+        return false;
+    }
+
+    size_t n_diff = 0;
+    size_t i_diff = 0;
+    for (size_t i = 0; i < n_a; i++) {
+        if (blob_a[i] != blob_b[i]) {
+            if (n_diff == 0) {
+                i_diff = i;
+            }
+            n_diff++;
+        }
+    }
+
+    if (n_diff > 0) {
+        LOG_ERR("\n%s: state changed across a restore: %zu of %zu bytes differ, first at offset %zu\n",
+                __func__, n_diff, n_a, i_diff);
+        return false;
+    }
+
+    LOG("\nPASS\n");
+    return true;
+}
+
+
+// Run the full save/load test suite (tests 1-8) for a single model.
 // Returns true if all tests pass, false otherwise.
 static bool run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) {
     struct common_params params = base_params;
@@ -422,6 +575,21 @@ static bool run_save_load_tests_for_model(const std::string & model_path, const
         return false;
     }
 
+    // Test 6: seq copy (host, scatter)
+    if (!test_seq_cp_scatter(model, params, tokens, 6, false)) {
+        return false;
+    }
+
+    // Test 7: seq copy (device, scatter)
+    if (!test_seq_cp_scatter(model, params, tokens, 7, true)) {
+        return false;
+    }
+
+    // Test 8: state blob round-trip
+    if (!test_state_roundtrip(model, params, tokens)) {
+        return false;
+    }
+
     LOG("\nAll tests passed.\n");
 
     return true;

tools/rpc/CMakeLists.txt

@@ -3,6 +3,18 @@ add_executable(${TARGET} rpc-server.cpp)
 target_link_libraries(${TARGET} PRIVATE ggml)
 target_compile_features(${TARGET} PRIVATE cxx_std_17)
 
+if (LLAMA_BUILD_TESTS AND UNIX AND NOT GGML_BACKEND_DL)
+    add_executable(test-rpc-multi-server ${PROJECT_SOURCE_DIR}/tests/test-rpc-multi-server.cpp)
+    target_link_libraries(test-rpc-multi-server PRIVATE ggml ggml-rpc)
+    target_include_directories(test-rpc-multi-server PRIVATE ${PROJECT_SOURCE_DIR}/ggml/src)
+    add_test(
+        NAME test-rpc-multi-server
+        COMMAND bash ${PROJECT_SOURCE_DIR}/tests/test-rpc-multi-server.sh
+            $<TARGET_FILE:ggml-rpc-server>
+            $<TARGET_FILE:test-rpc-multi-server>)
+    set_property(TEST test-rpc-multi-server PROPERTY LABELS main)
+endif()
+
 if(LLAMA_TOOLS_INSTALL)
     install(TARGETS ${TARGET} RUNTIME)
 endif()

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte

@@ -7,7 +7,12 @@
 		ChatMessageSystem,
 		ChatMessageUser
 	} from '$lib/components/app/chat';
-	import { REASONING_TAGS, ROUTES, SYSTEM_MESSAGE_PLACEHOLDER } from '$lib/constants';
+	import {
+		AGENTIC_TEXT_COPY_SEPARATOR,
+		REASONING_TAGS,
+		ROUTES,
+		SYSTEM_MESSAGE_PLACEHOLDER
+	} from '$lib/constants';
 	import { setChatMessageActionsContext, setChatMessageEditContext } from '$lib/contexts';
 	import { AgenticSectionType, AttachmentType, MessageRole } from '$lib/enums';
 	import { DatabaseService } from '$lib/services/database.service';
@@ -237,6 +242,24 @@
 	}
 
 	function handleCopy() {
+		// Agentic sessions render as a single entry anchored on the first assistant
+		// turn, whose own content is typically just the first tool call. Copy the
+		// text sections of the whole session so the clipboard matches the visible
+		// response instead of the anchor turn.
+		if (message.role === MessageRole.ASSISTANT) {
+			const sections = deriveAgenticSections(message, toolMessages, [], false);
+			const text = sections
+				.filter((section) => section.type === AgenticSectionType.TEXT)
+				.map((section) => section.content)
+				.join(AGENTIC_TEXT_COPY_SEPARATOR);
+
+			if (text) {
+				chatActions.copy(message, text);
+
+				return;
+			}
+		}
+
 		chatActions.copy(message);
 	}
 

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte

@@ -29,10 +29,10 @@
 			refreshAllMessages();
 		},
 
-		copy: async (message: DatabaseMessage) => {
+		copy: async (message: DatabaseMessage, contentOverride?: string) => {
 			const asPlainText = Boolean(currentConfig.copyTextAttachmentsAsPlainText);
 			const clipboardContent = formatMessageForClipboard(
-				message.content,
+				contentOverride ?? message.content,
 				message.extra,
 				asPlainText
 			);

tools/ui/src/lib/constants/agentic.constants.ts

@@ -20,6 +20,10 @@ export const SEARCH_SUMMARY = {
 // wraps mid-paragraph.
 export const RESULT_STAT_SEPARATOR = ' - ';
 
+// Separator between the assistant text sections of a grouped agentic
+// session when they are joined for the clipboard.
+export const AGENTIC_TEXT_COPY_SEPARATOR = '\n\n';
+
 export const DEFAULT_AGENTIC_CONFIG: AgenticConfig = {
 	enabled: true,
 	maxTurns: 100

tools/ui/src/lib/types/chat.d.ts

@@ -249,7 +249,7 @@ export interface ChatMessageDeletionInfo {
  * refresh + user-action notification), passed to each ChatMessage as a prop.
  */
 export interface ChatMessageActions {
-	copy: (message: DatabaseMessage) => void;
+	copy: (message: DatabaseMessage, contentOverride?: string) => void;
 	delete: (message: DatabaseMessage) => void;
 	navigateToSibling: (siblingId: string) => void;
 	editWithBranching: (