rpc: bump major version

* llama: re-pool each shared k-pool rep once

With shared cells every pool is re-pooled, and since the pooled keys
are always scattered, the pools a seq_cp shares between sequences
wrote the same rep row from several scatter entries, a data race on
the CPU backend. Mark each rep once: the sharing sequences read the
same row through pool_cells.

* llama: assert whole-sequence seq_cp in the hybrid idx memory

The recurrent state is always copied whole whatever the range, and a
k-pool cell shared by a partial copy could carry two pool groupings
with a single pooled row. Every caller copies whole sequences, so
reject partial ranges instead of supporting them.

* llama: drop the k-pool cache_safe mode

With whole-sequence seq_cp, sequences sharing cells share their pools
too, so the pooled row of a shared rep is valid for all of them. Mark
each rep once in every ubatch instead of re-pooling everything while
cells are shared, which removes the sharing scan and the stale-all
workarounds in seq_rm, state_read and state_drop. seq_cp now only
stales the destination.

.github/workflows/build-apple.yml

@@ -103,8 +103,9 @@ jobs:
         id: cmake_test
         run: |
           cd build
+          # Metal Paravirtual devices are difficult to support -> disable
           # ref: https://github.com/ggml-org/llama.cpp/pull/19802#issuecomment-4013704023
-          ctest -L main -E "test-llama-archs|test-save-load-state" --verbose --timeout 900
+          ctest -L main -E "test-llama-archs|test-save-load-state|test-recurrent-state-rollback" --verbose --timeout 900
 
   macos-latest-x64:
     runs-on: macos-15-intel

.github/workflows/docker.yml

@@ -23,7 +23,6 @@ on:
 cache-mode: none
 permissions:
   contents: read
-  packages: write
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
@@ -62,6 +61,9 @@ jobs:
   build_ui:
     name: Build UI
     needs: create_tag
+    permissions:
+      actions: write
+      contents: read
     uses: ./.github/workflows/ui-build.yml
     with:
       ui_version: ${{ needs.create_tag.outputs.source_tag }}
@@ -146,6 +148,11 @@ jobs:
     needs: [prepare_matrices, create_tag, build_ui]
 
     runs-on: ${{ matrix.config.runs_on }}
+    # cache-mode: write # for QEMU
+    permissions:
+      actions: write
+      contents: read
+      packages: write
     strategy:
       fail-fast: false
       matrix:
@@ -165,11 +172,11 @@ jobs:
           name: llama-ui.zip
           path: tools/ui/dist
 
-      - name: Set up QEMU
-        if: ${{ contains(matrix.config.platforms, 'linux/amd64') }}
-        uses: docker/setup-qemu-action@ce360397dd3f832beb865e1373c09c0e9f86d70a # v4
-        with:
-          image: tonistiigi/binfmt:qemu-v10.2.1
+      # - name: Set up QEMU
+      #   if: ${{ contains(matrix.config.platforms, 'linux/amd64') }}
+      #   uses: docker/setup-qemu-action@ce360397dd3f832beb865e1373c09c0e9f86d70a # v4
+      #   with:
+      #     image: tonistiigi/binfmt:qemu-v10.2.1
 
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@4d04d5d9486b7bd6fa91e7baf45bbb4f8b9deedd # v4

.github/workflows/make-release.yml

@@ -18,6 +18,11 @@ on:
         required: false
         type: boolean
         default: false
+      require_docker:
+        description: 'Require the Docker workflow to have completed successfully'
+        required: true
+        type: boolean
+        default: true
       apiabi_compare_tag:
         description: 'Tag to compare against for API/ABI check (default: latest release)'
         required: false
@@ -56,6 +61,7 @@ jobs:
           RELEASE_BRANCH: ${{ github.ref_name }}
           SKIP_APIABI_CHECK: ${{ github.event.inputs.skip_apiabi_check }}
           APIABI_COMPARE_TAG: ${{ github.event.inputs.apiabi_compare_tag }}
+          REQUIRE_DOCKER: ${{ github.event.inputs.require_docker }}
 
       - name: Create release tag
         if: ${{ github.event.inputs.dry_run == 'false' }}
@@ -132,7 +138,7 @@ jobs:
             });
 
       - name: Re-tag container images with release version
-        if: ${{ github.event.inputs.dry_run == 'false' && steps.desc.outputs.nightly_tag != '' }}
+        if: ${{ github.event.inputs.dry_run == 'false' && github.event.inputs.require_docker != 'false' && steps.desc.outputs.nightly_tag != '' }}
         env:
           GITHUB_REPOSITORY_OWNER: ${{ github.repository_owner }}
         run: |
@@ -145,14 +151,23 @@ jobs:
 
           VARIANTS=("" "-cuda" "-cuda13" "-vulkan" "-rocm" "-intel" "-musa" "-openvino")
           TYPES=("full" "light" "server")
+          # the release is already created at this point, so keep going on a
+          # missing image and report all of them at the end
+          MISSING=()
           for type in "${TYPES[@]}"; do
             for variant in "${VARIANTS[@]}"; do
               src="${IMAGE_REPO}:${type}${variant}-${NIGHTLY_TAG}"
               dst="${IMAGE_REPO}:${type}${variant}-${VERSION}"
               echo "Tagging ${src} -> ${dst}"
-              docker buildx imagetools create --tag "${dst}" "${src}"
+              if ! docker buildx imagetools create --tag "${dst}" "${src}"; then
+                MISSING+=("${type}${variant}")
+              fi
             done
           done
+          if [[ ${#MISSING[@]} -gt 0 ]]; then
+            echo "::error::failed to re-tag container images for ${NIGHTLY_TAG}:${MISSING[*]}"
+            exit 1
+          fi
 
       - name: Dry run summary
         if: ${{ github.event.inputs.dry_run == 'true' }}

.github/workflows/models-check.yml

@@ -159,7 +159,7 @@ jobs:
           GGML_METAL_DEVICES=4 ./build/bin/test-llama-archs -s 1
 
   rocm:
-    runs-on: [self-hosted, Linux, gfx1201]
+    runs-on: [self-hosted, Linux, gfx1201, 1accel]
     container: "rocm/dev-ubuntu-24.04:7.2.4-complete"
 
     steps:
@@ -299,7 +299,7 @@ jobs:
           ./build/bin/test-llama-archs -s 1
 
   vulkan-amd:
-    runs-on: [self-hosted, Linux, gfx1201]
+    runs-on: [self-hosted, Linux, gfx1201, 1accel]
     container: "ubuntu:26.04"
 
     steps:

.github/workflows/winget.yml

@@ -39,13 +39,14 @@ jobs:
             const { browser_download_url: asset_url_arm64 } = assets.find(asset => asset.name.includes('win-vulkan-arm64'));
             console.log("Latest release:", version);
             core.setOutput('VERSION', version);
-            core.setOutput('ASSETURL', `${asset_url_x64} ${asset_url_arm64}`);
+            core.setOutput('ASSETURL_X64', asset_url_x64);
+            core.setOutput('ASSETURL_ARM64', asset_url_arm64);
 
       - name: Update manifest
         run: |
           echo "Updating manifest..."
           komac update --version ${{ steps.find_latest_release.outputs.VERSION }} \
-            --urls "${{ steps.find_latest_release.outputs.ASSETURL }}" \
+            --urls "${{ steps.find_latest_release.outputs.ASSETURL_X64 }}" "${{ steps.find_latest_release.outputs.ASSETURL_ARM64 }}" \
             --token ${{ secrets.WINGET_GITHUB_TOKEN }} \
             --submit \
             ggml.llamacpp

CMakeLists.txt

@@ -4,7 +4,7 @@ include(CheckIncludeFileCXX)
 
 ### llama.cpp version
 set(LLAMA_VERSION_MAJOR 0)
-set(LLAMA_VERSION_MINOR 5)
+set(LLAMA_VERSION_MINOR 6)
 set(LLAMA_VERSION_PATCH 0)
 set(LLAMA_VERSION_BASE "${LLAMA_VERSION_MAJOR}.${LLAMA_VERSION_MINOR}.${LLAMA_VERSION_PATCH}")
 

common/arg.cpp

@@ -682,7 +682,10 @@ void common_models_handler_apply(common_models_handler & handler, common_params
             // if HF repo is a preset repo, we simply run server in router mode with the preset.ini file
             params.models_preset_hf = params.model.hf_repo; // only for showing a warning
             params.models_preset    = hf_cache::finalize_file(plan.preset);
-            params.model = common_params_model{}; // make sure to clear model, so server starts in router mode
+            // clear the model so the server starts in router mode
+            params.model.path.clear();
+            params.model.hf_repo.clear();
+            params.model.docker_repo.clear();
         });
     }
 

common/common.cpp

@@ -1,4 +1,5 @@
 #include "ggml.h"
+#include "ggml-cpp.h"
 #include "gguf.h"
 
 #include "build-info.h"
@@ -1065,6 +1066,20 @@ bool tty_can_use_colors() {
     return common_is_tty(stdout) || common_is_tty(stderr);
 }
 
+bool tty_enable_ansi() {
+#if defined(_WIN32)
+    // a Windows console renders ANSI sequences only in virtual terminal mode, pipes and files take them as is
+    for (DWORD id : { STD_OUTPUT_HANDLE, STD_ERROR_HANDLE }) {
+        HANDLE h    = GetStdHandle(id);
+        DWORD  mode = 0;
+        if (GetConsoleMode(h, &mode) && !SetConsoleMode(h, mode | ENABLE_VIRTUAL_TERMINAL_PROCESSING)) {
+            return false;
+        }
+    }
+#endif
+    return true;
+}
+
 //
 // Model utils
 //
@@ -1177,6 +1192,41 @@ common_decision_type common_get_decision_type(const struct llama_model * model)
     return common_decision_type_from_string(buf);
 }
 
+common_decision_type common_get_decision_type(const std::string & fname) {
+    struct gguf_init_params gguf_params = {
+        /* .no_alloc = */ true,
+        /* .ctx      = */ nullptr,
+    };
+
+    gguf_context_ptr gguf_ctx(gguf_init_from_file(fname.c_str(), gguf_params));
+    if (!gguf_ctx) {
+        return COMMON_DECISION_TYPE_UNKNOWN; // missing or unreadable file
+    }
+
+    std::string arch;
+    const int64_t arch_id = gguf_find_key(gguf_ctx.get(), "general.architecture");
+    if (arch_id < 0) {
+        return COMMON_DECISION_TYPE_UNKNOWN; // no architecture in the metadata
+    }
+    if (gguf_get_kv_type(gguf_ctx.get(), arch_id) != GGUF_TYPE_STRING) {
+        return COMMON_DECISION_TYPE_UNKNOWN; // malformed metadata
+    }
+    arch = gguf_get_val_str(gguf_ctx.get(), arch_id);
+    if (arch.empty()) {
+        return COMMON_DECISION_TYPE_UNKNOWN;
+    }
+
+    const std::string key = arch + ".decision.type";
+    const int64_t type_id = gguf_find_key(gguf_ctx.get(), key.c_str());
+    if (type_id < 0) {
+        return COMMON_DECISION_TYPE_NONE;
+    }
+    if (gguf_get_kv_type(gguf_ctx.get(), type_id) != GGUF_TYPE_STRING) {
+        return COMMON_DECISION_TYPE_UNKNOWN; // malformed metadata
+    }
+    return common_decision_type_from_string(gguf_get_val_str(gguf_ctx.get(), type_id));
+}
+
 common_init_result::common_init_result(common_params & params, bool model_only) :
     pimpl(new impl{}) {
     auto mparams = common_model_params_to_llama(params);

common/common.h

@@ -940,6 +940,7 @@ void fs_write_atomic(const std::filesystem::path & path, const std::string & dat
 
 // Auto-detect if colors can be enabled based on terminal and environment
 bool tty_can_use_colors();
+bool tty_enable_ansi(); // false when stdout or stderr is a console that cannot render ANSI sequences
 
 // Check if the given file is attached to a terminal
 bool common_is_tty(FILE * file);
@@ -964,6 +965,10 @@ enum common_decision_type {
 
 common_decision_type common_get_decision_type(const struct llama_model * model);
 
+// same as above, but reads a GGUF file; it does not load the model
+// returns COMMON_DECISION_TYPE_UNKNOWN if the file is missing, unreadable, or invalid
+common_decision_type common_get_decision_type(const std::string & fname);
+
 // note: defines the model, context, samplers, ets. lifetimes
 struct common_init_result {
     common_init_result(common_params & params, bool model_only = false);

common/log.cpp

@@ -144,12 +144,16 @@ struct common_log_entry {
             }
         }
 
-        fprintf(fcur, "%s", msg.data());
+        // the reset goes before the trailing newlines, so that every line carries its own colors
+        const bool reset = level == GGML_LOG_LEVEL_WARN || level == GGML_LOG_LEVEL_ERROR || level == GGML_LOG_LEVEL_DEBUG;
 
-        if (level == GGML_LOG_LEVEL_WARN || level == GGML_LOG_LEVEL_ERROR || level == GGML_LOG_LEVEL_DEBUG) {
-            fprintf(fcur, "%s", g_col[COMMON_LOG_COL_DEFAULT]);
+        size_t end = strlen(msg.data());
+        while (end > 0 && msg[end - 1] == '\n') {
+            end--;
         }
 
+        fprintf(fcur, "%.*s%s%s", (int) end, msg.data(), reset ? g_col[COMMON_LOG_COL_DEFAULT] : "", msg.data() + end);
+
         fflush(fcur);
     }
 };
@@ -158,6 +162,7 @@ struct common_log {
     // default capacity
     common_log(size_t capacity = 512) {
         file       = nullptr;
+        colors     = false;
         prefix     = false;
         timestamps = false;
         running    = false;
@@ -185,6 +190,7 @@ private:
 
     FILE * file;
 
+    bool colors;
     bool prefix;
     bool timestamps;
     bool running;
@@ -394,10 +400,16 @@ public:
         resume();
     }
 
+    bool get_colors() const {
+        return colors;
+    }
+
     void set_colors(bool colors) {
         pause();
 
-        if (colors) {
+        this->colors = colors && tty_enable_ansi();
+
+        if (this->colors) {
             g_col[COMMON_LOG_COL_DEFAULT] = LOG_COL_DEFAULT;
             g_col[COMMON_LOG_COL_BOLD]    = LOG_COL_BOLD;
             g_col[COMMON_LOG_COL_RED]     = LOG_COL_RED;
@@ -500,6 +512,10 @@ void common_log_set_colors(struct common_log * log, log_colors colors) {
     log->set_colors(true);
 }
 
+bool common_log_get_colors(struct common_log * log) {
+    return log->get_colors();
+}
+
 void common_log_set_prefix(struct common_log * log, bool prefix) {
     log->set_prefix(prefix);
 }

common/log.h

@@ -93,6 +93,7 @@ void common_log_add(struct common_log * log, enum ggml_log_level level, const ch
 
 void common_log_set_file      (struct common_log * log, const char * file); // not thread-safe
 void common_log_set_colors    (struct common_log * log, log_colors colors); // not thread-safe
+bool common_log_get_colors    (struct common_log * log);                    // whether colors are enabled
 void common_log_set_prefix    (struct common_log * log, bool prefix);       // whether to output prefix to each log
 void common_log_set_timestamps(struct common_log * log, bool timestamps);   // whether to output timestamps in the prefix
 void common_log_flush         (struct common_log * log);                    // flush all pending log messages

conversion/clef.py

@@ -11,8 +11,9 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import MmprojModel, ModelBase, gguf, logger
+from .base import ModelBase, gguf, logger
 from .qwen import Qwen3_5TextModel
+from .qwen3vl import Qwen3VLVisionModel
 
 
 def _is_clef_checkpoint(dir_model: Path) -> bool:
@@ -65,6 +66,7 @@ class ClefModel(Qwen3_5TextModel):
 
         # the pieces of the prompt are tokenized one by one, the server gives the text that separates them (sep)
         # and the text that starts the span of a question or of an option (mark_question, mark_option)
+        # images is one media marker per image, the vision start and end tokens are added by the server
         # the keys of JSON objects are given in sorted order
         option = (
             "{% set d = o.description %}"
@@ -75,6 +77,7 @@ class ClefModel(Qwen3_5TextModel):
         )
         return (
             text(f"<|im_start|>system\n{cls._SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\nSTATE:\n")
+            + "{% if images %}{{ sep }}{% for image in images %}{{ image }}{% endfor %}" + text("\n") + "{% endif %}"
             + "{{ sep }}" + render("state")
             + "{{ sep }}" + text("\n\nSCHEMA FIELDS:\n")
             + "{% for q in questions %}"
@@ -143,8 +146,5 @@ class ClefModel(Qwen3_5TextModel):
 
 
 @ModelBase.register("ClefModel")
-class ClefVisionModel(MmprojModel):
-    def __init__(self, *args, **kwargs):
-        del args, kwargs
-        raise NotImplementedError(
-            "multimodal input is not supported yet for Clef, requires https://github.com/ggml-org/llama.cpp/pull/29622 to be merged first")
+class ClefVisionModel(Qwen3VLVisionModel):
+    pass

docs/build-riscv64-spacemit.md

@@ -63,7 +63,7 @@ ${QEMU_ROOT_PATH}/bin/qemu-riscv64 -L ${RISCV_ROOT_PATH_IME1}/sysroot -cpu max,v
 | Q5_1 |  | :heavy_check_mark: |
 | Q5_K |  | :heavy_check_mark: |
 | Q6_K |  | :heavy_check_mark: |
-| Q8_0 |  | :heavy_check_mark: |
+| Q8_0 | :heavy_check_mark: | :heavy_check_mark: |
 
 
 ## Performance

ggml/CMakeLists.txt

@@ -4,8 +4,8 @@ project("ggml" C CXX ASM)
 
 ### GGML Version
 set(GGML_VERSION_MAJOR 0)
-set(GGML_VERSION_MINOR 25)
-set(GGML_VERSION_PATCH 3)
+set(GGML_VERSION_MINOR 26)
+set(GGML_VERSION_PATCH 0)
 set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")
 
 list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/")

ggml/include/ggml-rpc.h

@@ -6,7 +6,7 @@
 extern "C" {
 #endif
 
-#define RPC_PROTO_MAJOR_VERSION    7
+#define RPC_PROTO_MAJOR_VERSION    8
 #define RPC_PROTO_MINOR_VERSION    0
 #define RPC_PROTO_PATCH_VERSION    0
 

ggml/src/ggml-backend-meta.cpp

@@ -869,7 +869,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
         ggml_backend_meta_split_state split_state;
         switch (tensor->op) {
             case GGML_OP_NONE: {
-                split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+                if (tensor->view_src != nullptr) {
+                    // full-tensor view created with ggml_view_tensor, transparent for the split state
+                    split_state = ggml_backend_meta_get_split_state(stc, tensor->view_src, assume_sync);
+                } else {
+                    split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+                }
             } break;
             case GGML_OP_DUP: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);

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

@@ -1305,7 +1305,7 @@ void ggml_compute_forward_mul_mat(
 
     const bool src1_cont = ggml_is_contiguous(src1);
 
-    if (src1_cont) {
+    if (!params->use_ref && src1_cont) {
         for (int64_t i13 = 0; i13 < ne13; i13++)
             for (int64_t i12 = 0; i12 < ne12; i12++)
                 if (!llamafile_sgemm(params,
@@ -1384,7 +1384,7 @@ UseGgmlGemm1:;
     ggml_barrier(params->threadpool);
 
 #if GGML_USE_LLAMAFILE
-    if (src1->type != vec_dot_type) {
+    if (!params->use_ref && src1->type != vec_dot_type) {
         const void* wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata;
         const size_t row_size = ggml_row_size(vec_dot_type, ne10);
 

ggml/src/ggml-cpu/llamafile/sgemm.cpp

@@ -384,6 +384,80 @@ template <> inline __m256bh load(const float *p) {
 }
 #endif
 
+#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
+template <typename T, typename U> T load_partial(const U *, int);
+template <typename T> T load_partial_u16(const void *, int);
+
+template <> inline __m128i load_partial_u16(const void *p, int n) {
+#if defined(__AVX512BW__) && defined(__AVX512VL__)
+    return _mm_maskz_loadu_epi16((1u << n) - 1, p);
+#else
+    const __m128i index = _mm_setr_epi32(0, 1, 2, 3);
+    const __m128i pairs = _mm_set1_epi32(n / 2);
+    __m128i v = _mm_castps_si128(_mm_maskload_ps((const float *)p, _mm_cmpgt_epi32(pairs, index)));
+    if (n & 1) {
+        uint16_t last;
+        memcpy(&last, (const char *)p + 2*(n - 1), sizeof(last));
+        v = _mm_or_si128(v, _mm_and_si128(_mm_cmpeq_epi32(pairs, index), _mm_set1_epi32(last)));
+    }
+    return v;
+#endif
+}
+
+template <> inline __m256 load_partial(const float *p, int n) {
+    const __m256 index = _mm256_setr_ps(0, 1, 2, 3, 4, 5, 6, 7);
+    return _mm256_maskload_ps(p, _mm256_castps_si256(_mm256_cmp_ps(index, _mm256_set1_ps(n), _CMP_LT_OQ)));
+}
+
+#if defined(__F16C__)
+template <> inline __m256 load_partial(const ggml_fp16_t *p, int n) {
+    return _mm256_cvtph_ps(load_partial_u16<__m128i>(p, n));
+}
+#endif
+
+#if defined(__AVX2__) || defined(__AVX512F__)
+template <> inline __m256 load_partial(const ggml_bf16_t *p, int n) {
+    return _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(load_partial_u16<__m128i>(p, n)), 16));
+}
+#endif
+
+#if defined(__AVX512F__)
+template <> inline __m256i load_partial_u16(const void *p, int n) {
+#if defined(__AVX512BW__) && defined(__AVX512VL__)
+    return _mm256_maskz_loadu_epi16((1u << n) - 1, p);
+#else
+    const __m256i index = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+    const __m256i pairs = _mm256_set1_epi32(n / 2);
+    __m256i v = _mm256_maskload_epi32((const int *)p, _mm256_cmpgt_epi32(pairs, index));
+    if (n & 1) {
+        uint16_t last;
+        memcpy(&last, (const char *)p + 2*(n - 1), sizeof(last));
+        v = _mm256_or_si256(v, _mm256_and_si256(_mm256_cmpeq_epi32(pairs, index), _mm256_set1_epi32(last)));
+    }
+    return v;
+#endif
+}
+
+template <> inline __m512 load_partial(const float *p, int n) {
+    return _mm512_maskz_loadu_ps((1u << n) - 1, p);
+}
+
+template <> inline __m512 load_partial(const ggml_fp16_t *p, int n) {
+    return _mm512_cvtph_ps(load_partial_u16<__m256i>(p, n));
+}
+
+template <> inline __m512 load_partial(const ggml_bf16_t *p, int n) {
+    return _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(load_partial_u16<__m256i>(p, n)), 16));
+}
+#endif
+
+#if defined(__AVX512BF16__)
+template <> inline __m512bh load_partial(const ggml_bf16_t *p, int n) {
+    return (__m512bh) _mm512_maskz_loadu_epi16((uint64_t(1) << n) - 1, p);
+}
+#endif
+#endif
+
 #if defined(__riscv_v_intrinsic)
 template <> inline vfloat32m1_t load(const float *p) {
     return __riscv_vle32_v_f32m1(p, __riscv_vsetvlmax_e32m1());
@@ -492,8 +566,10 @@ class tinyBLAS {
     }
 
     bool matmul(int64_t m, int64_t n) {
+#if !defined(__AVX__) && !defined(__AVX2__) && !defined(__AVX512F__)
         if (k % KN != 0)
             return false;
+#endif
         // compute RM for only need tile with size RM&RM-1
 #if VECTOR_REGISTERS == 32
         if (m % 16 == 0 && (m/16 >= params->nth)) {
@@ -548,7 +624,7 @@ class tinyBLAS {
     template <int RM, int RN>
     inline void gemm_bloc(int64_t ii, int64_t jj) {
         D Cv[RN][RM] = {};
-        for (int64_t l = 0; l < k; l += KN) {
+        for (int64_t l = 0; l + KN <= k; l += KN) {
             // help compiler for op order.
             if constexpr (RM <= RN) {
                 V Av[RM];
@@ -574,6 +650,21 @@ class tinyBLAS {
                 }
             }
         }
+#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
+        const int64_t rem = k % KN;
+        if (rem != 0) {
+            V Av[RM];
+            for (int64_t i = 0; i < RM; ++i) {
+                Av[i] = load_partial<V>(A + lda * (ii + i) + k - rem, rem);
+            }
+            for (int64_t j = 0; j < RN; ++j) {
+                V Bv = load_partial<V>(B + ldb * (jj + j) + k - rem, rem);
+                for (int64_t i = 0; i < RM; ++i) {
+                    Cv[j][i] = madd(Av[i], Bv, Cv[j][i]);
+                }
+            }
+        }
+#endif
         for (int64_t j = 0; j < RN; ++j)
             for (int64_t i = 0; i < RM; ++i)
                 C[ldc * (jj + j) + (ii + i)] = hsum(Cv[j][i]);

ggml/src/ggml-cpu/spacemit/ime.cpp

@@ -321,6 +321,9 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
                           std::is_same_v<BLOC_TYPE, block_q4_K>) {
                 gemm_kernel     = spacemit_kernels::ime1::gemm_kernel_i8i4;
                 set_kernel_impl = true;
+            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q8_0>) {
+                gemm_kernel     = spacemit_kernels::ime1::gemm_kernel_i8i8;
+                set_kernel_impl = true;
             }
         }
 #endif
@@ -624,6 +627,9 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
                           std::is_same_v<BLOC_TYPE, block_q4_K>) {
                 gemm_kernel     = spacemit_kernels::ime1::gemm_kernel_i8i4;
                 set_kernel_impl = true;
+            } else if constexpr (std::is_same_v<BLOC_TYPE, block_q8_0>) {
+                gemm_kernel     = spacemit_kernels::ime1::gemm_kernel_i8i8;
+                set_kernel_impl = true;
             }
         }
 #endif
@@ -1237,6 +1243,7 @@ class tensor_traits_common : public tensor_traits_base {
 static const tensor_traits<block_q4_0, 32, 16>  q4_0_16x32_q8_0;
 static const tensor_traits<block_q4_1, 32, 16>  q4_1_16x32_q8_0;
 static const tensor_traits<block_q4_K, 32, 16>  q4_k_16x32_q8_0;
+static const tensor_traits<block_q8_0, 32, 16>  q8_0_16x32_q8_0;
 // Impl By IME2
 static const tensor_traits<block_q2_K, 256, 32> q2_k_32x256_q8_0;
 static const tensor_traits<block_q3_K, 256, 32> q3_k_32x256_q8_0;
@@ -1348,6 +1355,12 @@ static const ggml::cpu::tensor_traits * ggml_riscv64_spacemit_get_optimal_repack
                     return &ggml::cpu::riscv64_spacemit::q8_0_32x32_q8_0;
                 }
 #endif
+
+#if defined(RISCV64_SPACEMIT_IME1)
+                if (cur->ne[1] % 16 == 0 && (ggml::cpu::riscv64_spacemit::global_spine_env_info.use_ime1)) {
+                    return &ggml::cpu::riscv64_spacemit::q8_0_16x32_q8_0;
+                }
+#endif
             }
             break;
         case GGML_TYPE_MXFP4:

ggml/src/ggml-cpu/spacemit/ime1_kernels.cpp

@@ -991,6 +991,224 @@ void SQ4BitGemmM1Kernel_CompInt8_ScaleFp16_Impl(size_t          BlkLen,
         }
     }
 }
+
+
+// ---- Q8_0 IME1 int8xint8 M4 kernel ----------------------------------------
+// Handles 4 rows at once using the 4-row interleaved A produced by quantize_a_4row_i8:
+// per K-block (stride 144B) the first 16B are the four fp32 row scales and the following
+// 128B are four 32B chunks, each holding 8 K-values for each of the 4 rows. That maps
+// directly onto vmadot's 4-lane row dimension, so the A side is identical to Q4_0.
+// B is block_q8_0x16 as in the M1 kernel: 32B of fp16 scales then 512B of interleaved
+// int8, loaded with 8 plain vle8 into v2..v9 (no nibble unpack).
+static void SQ8BitGemmM4Kernel_CompInt8_ScaleFp16_Impl(size_t          BlkLen,
+                                                       const uint8_t * QuantA,
+                                                       const uint8_t * QuantBData,
+                                                       float *         C,
+                                                       size_t          CountN,
+                                                       size_t          BlockCountK,
+                                                       const size_t    ldc) {
+    // Same invariant as the Q4_0 M4 kernel: SAVE_RESULT_4x16 stores a full 4x16 tile with no tail
+    // handling. Q8_0 is only admitted to the IME1 path when ne[1] % 16 == 0 (ime.cpp),
+    // and the n-tiling step is NB_COLS == 16, so a partial tile never reaches here.
+    GGML_ASSERT(CountN % 16 == 0);
+
+    const size_t INNER = BlkLen / 16;
+    const size_t LDC   = ldc * sizeof(float);
+
+    for (size_t n = 0; n < CountN; n += 16) {
+        uint8_t * QuantBDataPtr =
+            (uint8_t *) QuantBData + (n / 16) * BlockCountK * (16 * sizeof(_Float16) + 512);
+        float * CPtr = C + n;
+
+        __asm__ volatile(
+            "vsetvli            t0, zero, e32, m8           \n\t"
+            "vxor.vv            v24, v24, v24               \n\t"
+            "addi               t3, %[BlockCountK], 0       \n\t"
+            "addi               a1, %[A], 0                 \n\t"
+            "addi               s1, %[B], 0                 \n\t"
+
+            "BLOCK_COUNTK_LOOP%=:                           \n\t"
+            "addi               s5, s1, 0                   \n\t"
+            "addi               s1, s5, 32                  \n\t"
+            "vsetvli            t0, zero, e32, m8           \n\t"
+            "vxor.vv            v16, v16, v16               \n\t"
+            "flw                f1, (a1)                    \n\t"
+            "flw                f2, 4(a1)                   \n\t"
+            "flw                f3, 8(a1)                   \n\t"
+            "flw                f4, 12(a1)                  \n\t"
+            "addi               a1, a1, 16                  \n\t"
+            "addi               t2, %[INNER], 0             \n\t"
+
+            "BLOCK_INNER_LOOP%=:                            \n\t"
+            "vsetvli            t0, zero, e8, m1            \n\t"
+            "vle8.v             v2, (s1)                    \n\t"
+            "addi               s1, s1, 32                  \n\t"
+            "vle8.v             v3, (s1)                    \n\t"
+            "addi               s1, s1, 32                  \n\t"
+            "vle8.v             v4, (s1)                    \n\t"
+            "addi               s1, s1, 32                  \n\t"
+            "vle8.v             v5, (s1)                    \n\t"
+            "addi               s1, s1, 32                  \n\t"
+            "vle8.v             v6, (s1)                    \n\t"
+            "addi               s1, s1, 32                  \n\t"
+            "vle8.v             v7, (s1)                    \n\t"
+            "addi               s1, s1, 32                  \n\t"
+            "vle8.v             v8, (s1)                    \n\t"
+            "addi               s1, s1, 32                  \n\t"
+            "vle8.v             v9, (s1)                    \n\t"
+            "addi               s1, s1, 32                  \n\t"
+            "vle8.v             v10, (a1)                   \n\t"
+            "addi               a1, a1, 32                  \n\t"
+            "vle8.v             v11, (a1)                   \n\t"
+            "addi               a1, a1, 32                  \n\t"
+
+            SQ4BIT_KERNEL_COMP_4x16x16
+
+            "addi               t2, t2, -1                  \n\t"
+            "bnez               t2, BLOCK_INNER_LOOP%=      \n\t"
+
+            LOAD_SCALE_4x16_FP16
+
+            "vsetvli            t0, zero, e32, m8           \n\t"
+            "vfcvt.f.x.v        v16, v16                    \n\t"
+            "vfmacc.vv          v24, v16, v8                \n\t"
+            "addi               t3, t3, -1                  \n\t"
+            "bnez               t3, BLOCK_COUNTK_LOOP%=     \n\t"
+
+            "RESULT_SAVE%=:                                 \n\t"
+
+            SAVE_RESULT_4x16
+
+            :
+            : [INNER] "r"(INNER), [A] "r"(QuantA), [B] "r"(QuantBDataPtr), [LDC] "r"(LDC),
+              [BlockCountK] "r"(BlockCountK), [C] "r"(CPtr)
+            : "cc", "t0", "t1", "t2", "t3", "a1", "a2", "a3", "a4", "f1", "f2", "f3", "f4",
+              "s1", "s2", "s3", "s4", "s5", "s6");
+    }
+}
+
+// ---- Q8_0 IME1 int8xint8 M1 kernel ----------------------------------------
+// B comes from block_q8_0x16 (make_block_q8_0x16): 16 fp16 scales (32B) then 512B of interleaved
+// int8 weights laid out as [INNER0: reg0..reg7][INNER1: reg0..reg7], each reg 32B = 4 columns x
+// (even=K-first-half / odd=K-second-half). A comes from quantize_a_row_i8 (same as Q4_0 path).
+// Reuses the vmadot COMP macro and the ACC_F16 dequant tail from the Q4_0 kernel.
+static void SQ8BitGemmM1Kernel_CompInt8_ScaleFp16_Impl(size_t          BlkLen,
+                                                       const uint8_t * QuantA,
+                                                       const uint8_t * QuantBData,
+                                                       float *         C,
+                                                       size_t          CountN,
+                                                       size_t          BlockCountK,
+                                                       const size_t    ldc) {
+    GGML_UNUSED(ldc);
+    const size_t INNER = BlkLen / 16;  // = 2 for QK8_0=32
+
+    for (size_t n = 0; n < CountN; n += 16) {
+        size_t    nblks = (CountN - n) > 16 ? 16 : CountN - n;
+        // Each x16 K-block is {16 fp16 scales (32B), 512B interleaved int8}; stride = 544B.
+        uint8_t * QuantBDataPtr = (uint8_t *) QuantBData + (n / 16) * BlockCountK * (16 * sizeof(_Float16) + 512);
+        float *   CPtr = C + n;
+        size_t    cnt  = BlockCountK;
+
+        __asm__ volatile(
+            "vsetvli      t0, zero, e32, m4       \n\t"
+            "vxor.vv      v28, v28, v28           \n\t"
+            // s7 = per-K-block base (scale@+0, data@+32, block stride 544)
+            "addi         s7, %[B], 0             \n\t"
+            "addi         s5, %[A], 0             \n\t"   // A scale (fp32)
+            "addi         s6, %[A], 12            \n\t"   // A data (int8), offset like Q4_0 M1
+            "LOOP_K%=:                            \n\t"
+            "addi         s1, s7, 32              \n\t"   // data base for this K-block
+            // B scales: d[0..15] fp16 at block start. Load in 4 groups of 4 (d[0-3]/[4-7]/[8-11]/[12-15])
+            // matching the 4 accumulators (each covers columns [g*4 .. g*4+3]).
+            "addi         s2, s7, 8               \n\t"
+            "addi         s3, s7, 16              \n\t"
+            "addi         s4, s7, 24              \n\t"
+            "vsetvli      t0, zero, e16, mf4      \n\t"
+            "vle16.v      v4, (s7)                \n\t"
+            "vle16.v      v5, (s2)                \n\t"
+            "vle16.v      v6, (s3)                \n\t"
+            "vle16.v      v7, (s4)                \n\t"
+            "addi         s7, s7, 544             \n\t"   // advance to next K-block (32 scale + 512 data)
+            "flw          f1, (s5)                \n\t"
+            "addi         s5, s5, 4               \n\t"
+            "vfwcvt.f.f.v v8, v4                  \n\t"
+            "vfwcvt.f.f.v v9, v5                  \n\t"
+            "vfwcvt.f.f.v v10, v6                 \n\t"
+            "vfwcvt.f.f.v v11, v7                 \n\t"
+            "vsetvli      t0, zero, e32, mf2      \n\t"
+            "addi         t5, %[INNER], 0         \n\t"
+            "vxor.vv      v16, v16, v16           \n\t"
+            "vxor.vv      v18, v18, v18           \n\t"
+            "vxor.vv      v20, v20, v20           \n\t"
+            "vxor.vv      v22, v22, v22           \n\t"
+            // combined scale (A_scale * B_scale) -> v24..v27 (one per accumulator)
+            "vfmul.vf     v24, v8, f1             \n\t"
+            "vfmul.vf     v25, v9, f1             \n\t"
+            "vfmul.vf     v26, v10, f1            \n\t"
+            "vfmul.vf     v27, v11, f1            \n\t"
+            "addi         %[CNT], %[CNT], -1      \n\t"
+            "vsetvli      t0, zero, e8, m1        \n\t"
+            "LOOP_INNER%=:                        \n\t"
+            // load 8 B data regs (v0..v7) directly (int8, no nibble unpack)
+            "vle8.v       v0, (s1)                \n\t"
+            "addi         s1, s1, 32              \n\t"
+            "vle8.v       v1, (s1)                \n\t"
+            "addi         s1, s1, 32              \n\t"
+            "vle8.v       v2, (s1)                \n\t"
+            "addi         s1, s1, 32              \n\t"
+            "vle8.v       v3, (s1)                \n\t"
+            "addi         s1, s1, 32              \n\t"
+            "vle8.v       v4, (s1)                \n\t"
+            "addi         s1, s1, 32              \n\t"
+            "vle8.v       v5, (s1)                \n\t"
+            "addi         s1, s1, 32              \n\t"
+            "vle8.v       v6, (s1)                \n\t"
+            "addi         s1, s1, 32              \n\t"
+            "vle8.v       v7, (s1)                \n\t"
+            "addi         s1, s1, 32              \n\t"
+            // load A (2 halves): v14 from s5, v15 from s6 (matches Q4_0 M1 A packing)
+            "vsetvli      t0, zero, e8, mf4       \n\t"
+            "vle8.v       v14, (s5)               \n\t"
+            "addi         s5, s5, 16              \n\t"
+            "vle8.v       v15, (s6)               \n\t"
+            "addi         s6, s6, 16              \n\t"
+            "addi         t5, t5, -1              \n\t"
+            "vsetvli      t0, zero, e8, m1        \n\t"
+            SQ4BIT_KERNEL_COMP_1x8x2_4X8X4
+            "bnez         t5, LOOP_INNER%=        \n\t"
+            "vsetvli      t0, zero, e32, mf2      \n\t"
+            SQ4BIT_KERNEL_ACC_F16_1X4X4
+            "bnez         %[CNT], LOOP_K%=        \n\t"
+            "addi         t3, zero, 16            \n\t"
+            "addi         s1, %[C], 16            \n\t"
+            "addi         s2, %[C], 32            \n\t"
+            "addi         s3, %[C], 48            \n\t"
+            "blt          %[NBLKS], t3, ST_TAIL%= \n\t"
+            "vse32.v      v28, (%[C])             \n\t"
+            "vse32.v      v29, (s1)               \n\t"
+            "vse32.v      v30, (s2)               \n\t"
+            "vse32.v      v31, (s3)               \n\t"
+            "jal          x0, END%=               \n\t"
+            "ST_TAIL%=:                           \n\t"
+            "vsetvli      t0, %[NBLKS], e32, mf2  \n\t"
+            "sub          %[NBLKS], %[NBLKS], t0  \n\t"
+            "vse32.v      v28, (%[C])             \n\t"
+            "vsetvli      t0, %[NBLKS], e32, mf2  \n\t"
+            "sub          %[NBLKS], %[NBLKS], t0  \n\t"
+            "vse32.v      v29, (s1)               \n\t"
+            "vsetvli      t0, %[NBLKS], e32, mf2  \n\t"
+            "sub          %[NBLKS], %[NBLKS], t0  \n\t"
+            "vse32.v      v30, (s2)               \n\t"
+            "vsetvli      t0, %[NBLKS], e32, mf2  \n\t"
+            "sub          %[NBLKS], %[NBLKS], t0  \n\t"
+            "vse32.v      v31, (s3)               \n\t"
+            "END%=:                               \n\t"
+            : [CNT] "+r"(cnt), [NBLKS] "+r"(nblks)
+            : [INNER] "r"(INNER), [A] "r"(QuantA), [B] "r"(QuantBDataPtr), [C] "r"(CPtr)
+            : "cc", "t0", "t3", "t5", "f1", "s1", "s2", "s3", "s4", "s5", "s6", "s7");
+    }
+}
+
 }  // namespace
 
 namespace ime1 {
@@ -1023,5 +1241,24 @@ size_t gemm_kernel_i8i4(size_t          blk_len,
         return 1;
     }
 }
+
+size_t gemm_kernel_i8i8(size_t          blk_len,
+                        const uint8_t * quant_a_ptr,
+                        const uint8_t * quant_b_data,
+                        const uint8_t * quant_b_zp,
+                        float *         c_ptr,
+                        size_t          count_m,
+                        size_t          count_n,
+                        size_t          k_blks,
+                        size_t          ldc) {
+    GGML_UNUSED(quant_b_zp);
+    if (count_m >= 4) {
+        SQ8BitGemmM4Kernel_CompInt8_ScaleFp16_Impl(blk_len, quant_a_ptr, quant_b_data, c_ptr, count_n, k_blks,
+                                                   ldc);
+        return 4;
+    }
+    SQ8BitGemmM1Kernel_CompInt8_ScaleFp16_Impl(blk_len, quant_a_ptr, quant_b_data, c_ptr, count_n, k_blks, ldc);
+    return 1;
+}
 }  // namespace ime1
 }  // namespace spacemit_kernels

ggml/src/ggml-cpu/spacemit/ime_kernels.h

@@ -79,6 +79,16 @@ size_t gemm_kernel_i8i4(size_t          blk_len,
                         size_t          k_blks,
                         size_t          ldc);
 
+size_t gemm_kernel_i8i8(size_t          blk_len,
+                        const uint8_t * quant_a_ptr,
+                        const uint8_t * quant_b_data,
+                        const uint8_t * quant_b_zp,
+                        float *         c_ptr,
+                        size_t          count_m,
+                        size_t          count_n,
+                        size_t          k_blks,
+                        size_t          ldc);
+
 void quantize_a_row_i8(size_t blk_len, const float * a_ptr, size_t count_k, uint8_t * quant_a_ptr);
 
 void quantize_a_4row_i8(size_t blk_len, const float * a_ptr, size_t count_k, uint8_t * quant_a_ptr);

ggml/src/ggml-cpu/spacemit/repack.cpp

@@ -370,6 +370,74 @@ static block_q8_0x32 make_block_q8_0x32(block_q8_0 * in, unsigned int blck_size_
     return out;
 }
 
+// IME1: interleave 16 q8_0 rows so a plain vle8 sequence in the i8i8 kernel lands weights in the
+// vmadot group/parity layout. Mirrors make_block_q4_0x16 but stores full int8 (no nibble packing).
+// qs (512B) = [INNER step 0: reg0..reg7][INNER step 1: reg0..reg7], each reg 32B holding 4 columns
+// x (even=K-first-half / odd=K-second-half). Column col -> acc=col/4, cgrp=col%4.
+static block_q8_0x16 make_block_q8_0x16(block_q8_0 * in, unsigned int blck_size_interleave) {
+    block_q8_0x16 out;
+    GGML_ASSERT(QK8_0 / blck_size_interleave == 2);
+    GGML_UNUSED(blck_size_interleave);
+
+    for (int i = 0; i < 16; i++) {
+        out.d[i] = in[i].d;
+    }
+
+    memset(out.qs, 0, sizeof(out.qs));
+    for (int col = 0; col < 16; col++) {
+        const int      acc  = col / 4;
+        const int      cgrp = col % 4;
+        const int8_t * q    = in[col].qs;
+        for (int s = 0; s < 2; s++) {
+            const int base   = s * 16;
+            uint8_t * reg_lo = out.qs + (s * 8 + acc)     * 32;
+            uint8_t * reg_hi = out.qs + (s * 8 + acc + 4) * 32;
+            for (int i = 0; i < 4; i++) {
+                reg_lo[(2 * cgrp)     * 4 + i] = (uint8_t) q[base + 0  + i];
+                reg_lo[(2 * cgrp + 1) * 4 + i] = (uint8_t) q[base + 4  + i];
+                reg_hi[(2 * cgrp)     * 4 + i] = (uint8_t) q[base + 8  + i];
+                reg_hi[(2 * cgrp + 1) * 4 + i] = (uint8_t) q[base + 12 + i];
+            }
+        }
+    }
+    return out;
+}
+
+static int repack_q8_0_to_q8_0_16_bl_ref(ggml_tensor *              t,
+                                         int                        interleave_block,
+                                         const void * GGML_RESTRICT data,
+                                         size_t                     data_size) {
+    GGML_ASSERT(t->type == GGML_TYPE_Q8_0);
+    GGML_ASSERT(interleave_block == 16);
+
+    constexpr int nrows_interleaved = 16;
+
+    block_q8_0x16 *    dst = (block_q8_0x16 *) t->data;
+    const block_q8_0 * src = (const block_q8_0 *) data;
+    block_q8_0         dst_tmp[16];
+    int                nrow    = ggml_nrows(t);
+    int                nblocks = t->ne[0] / QK8_0;
+
+    GGML_ASSERT(data_size == nrow * nblocks * sizeof(block_q8_0));
+
+    if (t->ne[1] % nrows_interleaved != 0 || t->ne[0] % QK8_0 != 0) {
+        return -1;
+    }
+
+    for (int b = 0; b < nrow; b += nrows_interleaved) {
+        for (int64_t x = 0; x < nblocks; x++) {
+            for (int i = 0; i < nrows_interleaved; i++) {
+                dst_tmp[i] = src[x + i * nblocks];
+            }
+            *dst++ = make_block_q8_0x16(dst_tmp, interleave_block);
+        }
+        src += nrows_interleaved * nblocks;
+    }
+    return 0;
+
+    GGML_UNUSED(data_size);
+}
+
 static int repack_q2_k_to_q2_k_32_bl(ggml_tensor *              t,
                                      int                        interleave_block,
                                      const void * GGML_RESTRICT data,
@@ -1768,6 +1836,10 @@ template <> int repack<block_q6_K, 32, 32>(ggml_tensor * t, const void * data, s
 #endif
 }
 
+template <> int repack<block_q8_0, 32, 16>(ggml_tensor * t, const void * data, size_t data_size) {
+    return repack_q8_0_to_q8_0_16_bl_ref(t, 16, data, data_size);
+}
+
 template <> int repack<block_q8_0, 32, 32>(ggml_tensor * t, const void * data, size_t data_size) {
 #if 1
     return repack_q8_0_to_q8_0_32_bl_ref(t, 32, data, data_size);

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

@@ -981,7 +981,7 @@ template <int DV, int ncols1, int ncols2>
 void launch_fattn(
     ggml_backend_cuda_context & ctx, ggml_tensor * dst, fattn_kernel_t fattn_kernel, const int nwarps, const size_t nbytes_shared,
     const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const bool use_sparse,
-    const int warp_size = WARP_SIZE
+    const int warp_size = WARP_SIZE, const bool async_kv_preload = false
 ) {
     constexpr int ncols = ncols1 * ncols2;
 
@@ -1114,7 +1114,8 @@ void launch_fattn(
     // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped.
     // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or
     //     multiple sequences of possibly different lengths.
-    if (!use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) {
+    const bool scan_mask = !use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1);
+    if (scan_mask) {
         const int64_t s31 = mask->nb[1] / sizeof(half2);
         const int64_t s33 = mask->nb[3] / sizeof(half2);
 
@@ -1142,10 +1143,16 @@ void launch_fattn(
 
     dim3 blocks_num;
     if (stream_k) {
-        auto should_use_stream_k = [](const int cc, const int ntiles_dst, const int max_blocks, const int DKQ) {
+        // Stream-K splits the work before the mask scan is applied, so skipped KV tiles make the blocks uneven.
+        const bool prefer_whole_tiles = GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_DGX_SPARK && async_kv_preload && scan_mask;
+
+        auto should_use_stream_k = [prefer_whole_tiles](const int cc, const int ntiles_dst, const int max_blocks, const int DKQ) {
             const int tiles_nwaves             = (ntiles_dst + max_blocks - 1) / max_blocks;
             const int tiles_efficiency_percent = 100 * ntiles_dst / (max_blocks*tiles_nwaves);
 
+            if (prefer_whole_tiles && tiles_efficiency_percent >= 75) {
+                return false;
+            }
             if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_ADA_LOVELACE) {
                 return true;
             }

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

@@ -329,32 +329,6 @@ static constexpr __device__ bool ggml_cuda_fattn_mma_get_Q_in_reg(const int DKQ,
     return ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols).Q_in_reg;
 }
 
-// Swizzling needs a tile stride that is a multiple of 32 half2 columns.
-static constexpr __host__ __device__ bool ggml_cuda_fattn_mma_bank_aligned(const int nbatch_2) {
-    return nbatch_2 >= 32 && nbatch_2 % 32 == 0;
-}
-
-// Swizzling needs ldmatrix, on other hardware the tiles keep the row padding.
-static __host__ bool ggml_cuda_fattn_mma_get_swizzled(const int DKQ, const int DV, const int ncols1, const int ncols2, const int cc) {
-    const fattn_mma_config cfg = ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols1*ncols2, cc);
-    return turing_mma_available(cc) && ggml_cuda_fattn_mma_bank_aligned(cfg.nbatch_K2) && ggml_cuda_fattn_mma_bank_aligned(cfg.nbatch_V2);
-}
-
-static constexpr __device__ bool ggml_cuda_fattn_mma_get_swizzled(const int DKQ, const int DV, const int ncols1, const int ncols2) {
-#if defined(TURING_MMA_AVAILABLE)
-    const fattn_mma_config cfg = ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols1*ncols2);
-    return ggml_cuda_fattn_mma_bank_aligned(cfg.nbatch_K2) && ggml_cuda_fattn_mma_bank_aligned(cfg.nbatch_V2);
-#else
-    GGML_UNUSED_VARS(DKQ, DV, ncols1, ncols2);
-    return false;
-#endif // defined(TURING_MMA_AVAILABLE)
-}
-
-// Row padding is only needed if the tile is not swizzled.
-static constexpr __host__ __device__ int ggml_cuda_fattn_mma_get_stride_tile(const int nbatch_2, const bool swizzled) {
-    return swizzled ? nbatch_2 : nbatch_2 + 4;
-}
-
 static constexpr __device__ int get_cols_per_thread() {
 #if defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
     return 1; // AMD has a single column per thread.
@@ -372,6 +346,20 @@ static __host__ int get_cols_per_warp(const int cc) {
     }
 }
 
+static __host__ bool ggml_cuda_fattn_mma_get_swizzled(const int DKQ, const int DV, const int ncols, const int cc) {
+    return turing_mma_available(cc) &&
+        ggml_cuda_fattn_mma_get_nbatch_K2(DKQ, DV, ncols, cc) % 32 == 0 && ggml_cuda_fattn_mma_get_nbatch_V2(DKQ, DV, ncols, cc) % 32 == 0;
+}
+
+static constexpr __device__ bool ggml_cuda_fattn_mma_get_swizzled(const int DKQ, const int DV, const int ncols) {
+#ifdef TURING_MMA_AVAILABLE
+    return ggml_cuda_fattn_mma_get_nbatch_K2(DKQ, DV, ncols) % 32 == 0 && ggml_cuda_fattn_mma_get_nbatch_V2(DKQ, DV, ncols) % 32 == 0;
+#else
+    GGML_UNUSED_VARS(DKQ, DV, ncols);
+    return false;
+#endif // TURING_MMA_AVAILABLE
+}
+
 // ------------------------------------------------------------------------------------------------------------------
 
 static __host__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, const int DV, const int ncols1, const int ncols2, const int cc) {
@@ -392,14 +380,15 @@ static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(
 
 // ------------------------------------------------------------------------------------------------------------------
 
-template<int stride_tile, bool swz, int nwarps, int nbatch_fa, bool use_cp_async, bool oob_check, bool use_sparse>
+template<int stride_tile, int nwarps, int nbatch_fa, bool use_cp_async, bool oob_check, bool use_sparse>
 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 k_VKQ_0, const int i_sup, const int32_t * const __restrict__ indices) {
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
     // K/V data is loaded with decreasing granularity for D for better memory bandwidth.
     // The minimum granularity is 16 bytes.
-    constexpr int h2_per_chunk = 16/sizeof(half2);
+    constexpr int chunk_size = 16;
+    constexpr int h2_per_chunk = chunk_size / sizeof(half2);
     const int chunks_per_row = D2 / h2_per_chunk;
     if constexpr (use_cp_async) {
         static_assert(warp_size == 32, "bad warp_size");
@@ -439,7 +428,7 @@ 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 + swizzle_bytes<swz, half2>(i, k*h2_per_chunk, stride_tile), KV + i_KV*stride_KV + k*h2_per_chunk);
+                    cp_async_cg_16<preload>(tile_KV_32 + swizzle<stride_tile*sizeof(half2), char>(i*stride_tile*sizeof(half2) + k*chunk_size, i), KV + i_KV*stride_KV + k*h2_per_chunk);
                 }
             }
         };
@@ -481,7 +470,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
                     } else {
                         src = !oob_check || i < i_sup ? KV + int64_t(k_VKQ_0 + i)*stride_KV + k*h2_per_chunk : zero;
                     }
-                    ggml_cuda_memcpy_1<16>((char *) tile_KV + swizzle_bytes<swz, half2>(i, k*h2_per_chunk, stride_tile), src);
+                    ggml_cuda_memcpy_1<16>(swizzle<stride_tile>(tile_KV, i*stride_tile + k*h2_per_chunk, i), src);
                 }
             }
         };
@@ -624,9 +613,9 @@ 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, use_sparse);
 
-    constexpr bool swz = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols1, ncols2);
-    constexpr int stride_tile_K = ggml_cuda_fattn_mma_get_stride_tile(nbatch_K2, swz);
-    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_mma_get_stride_tile(nbatch_V2, swz);
+    constexpr bool swz = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols);
+    constexpr int stride_tile_K = swz ? nbatch_K2 : nbatch_K2 + 4;
+    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : (swz ? nbatch_V2 : nbatch_V2 + 4);
 
     const int k_VKQ_0 = kb0 * nbatch_fa;
 #if defined(TURING_MMA_AVAILABLE)
@@ -644,7 +633,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, swz, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+        flash_attn_ext_f16_load_tile<stride_tile_V, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
             (V_h2, tile_V, nbatch_V2, stride_V, k_VKQ_0, k_VKQ_sup, nullptr);
     } else {
         // the sparse mask values are gathered per element, always load them synchronously
@@ -664,7 +653,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, swz, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+            flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
                 (K_h2 + k0_start, tile_K, k0_diff, stride_K, k_VKQ_0, k_VKQ_sup, indices);
             if (use_cp_async) {
                 cp_async_wait_all();
@@ -680,7 +669,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<swz>(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start, stride_tile_K);
+                    load_ldmatrix_swizzled<stride_tile_K>(K_A, tile_K, i_KQ_0*stride_tile_K + 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 {
@@ -706,7 +695,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<swz>(K_A, tile_K, i_KQ_0, k_KQ_0 - k0_start, stride_tile_K);
+                    load_ldmatrix_swizzled<stride_tile_K>(K_A, tile_K, i_KQ_0*stride_tile_K + 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]);
@@ -1001,7 +990,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, use_sparse>
                     (mask_h, tile_mask, stride_mask, k_VKQ_0 + nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr);
             }
-            flash_attn_ext_f16_load_tile<stride_tile_K, swz, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+            flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
                 (K_h2, tile_K, nbatch_K2, stride_K, k_VKQ_0 + nbatch_fa, k_VKQ_sup, nullptr);
         }
     }
@@ -1017,7 +1006,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, swz, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+                flash_attn_ext_f16_load_tile<stride_tile_V, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
                     (V_h2 + i0_start/2, tile_V, i0_diff/2, stride_V, k_VKQ_0, k_VKQ_sup, indices);
                 if (use_cp_async) {
                     cp_async_wait_all();
@@ -1025,7 +1014,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
                 __syncthreads();
             }
         }
-        const half2 * tile_V_i = !V_is_K_view || i0_stop > 2*nbatch_K2 ? tile_V : tile_V + i0_start/2;
+        const int tile_V_offset_i = !V_is_K_view || i0_stop > 2*nbatch_K2 ? 0 : i0_start/2;
 
 #if defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
 #pragma unroll
@@ -1036,7 +1025,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<swz>(A, tile_V, 2*k0, (int)(tile_V_i - tile_V) + (i_VKQ_0 - i0_start)/2, stride_tile_V);
+                load_ldmatrix_trans_swizzled<stride_tile_V>(A, tile_V, tile_V_offset_i + 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 {
@@ -1062,8 +1051,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.
-                static_assert(!swz, "Volta has no ldmatrix");
-                load_ldmatrix(A, tile_V_i + k0*stride_tile_V + (i_VKQ_0 - i0_start)/2, stride_tile_V);
+                load_ldmatrix_swizzled<stride_tile_V>(A, tile_V, tile_V_offset_i + 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);
             }
         }
@@ -1253,10 +1241,10 @@ 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 bool swz = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols1, ncols2);
-    constexpr int stride_tile_K = ggml_cuda_fattn_mma_get_stride_tile(nbatch_K2, swz);
-    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_mma_get_stride_tile(nbatch_V2, swz);
+    constexpr bool swz = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols);
+    constexpr int stride_tile_Q = DKQ/2 + 4;
+    constexpr int stride_tile_K = swz ? nbatch_K2 : nbatch_K2 + 4;
+    constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : (swz ? nbatch_V2 : nbatch_V2 + 4);
     constexpr int stride_tile_KV_max = stride_tile_K > stride_tile_V ? stride_tile_K : stride_tile_V;
 
     extern __shared__ half2 tile_Q[];
@@ -1354,7 +1342,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, use_sparse>
                 (mask_h, tile_mask, stride_mask, kb0*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr);
         }
-        flash_attn_ext_f16_load_tile<stride_tile_K, swz, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
+        flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
             (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, nullptr);
     }
 
@@ -2039,9 +2027,9 @@ 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
 
     // KV tile strides must match flash_attn_ext_f16_iter / _process_tile.
-    const bool swizzled     = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols1, ncols2, cc);
-    const int stride_tile_K = ggml_cuda_fattn_mma_get_stride_tile(nbatch_K2, swizzled);
-    const int stride_tile_V = V_is_K_view ? stride_tile_K : ggml_cuda_fattn_mma_get_stride_tile(nbatch_V2, swizzled);
+    const bool swz = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols, cc);
+    const int stride_tile_K = swz ? nbatch_K2 : nbatch_K2 + 4;
+    const int stride_tile_V = V_is_K_view ? stride_tile_K : (swz ? nbatch_V2 : nbatch_V2 + 4);
     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);
@@ -2112,8 +2100,9 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
         }
     }
 
+    const bool async_kv_preload = nstages == 2 && !use_sparse;
     launch_fattn<DV, ncols1, ncols2>
-        (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host);
+        (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host, async_kv_preload);
 }
 
 

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

@@ -1778,8 +1778,9 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) {
         (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) &&
         src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32;
 
-    const int cc      = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
-    use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, is_mul_mat_id ? src1->ne[2] : src1->ne[1]);
+    const int cc        = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+    const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
+    use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, warp_size, src0->ne, src0->nb, is_mul_mat_id ? src1->ne[2] : src1->ne[1]);
 
     //we only support fusion for ncols_dst = 1
     if (tensor->op == GGML_OP_MUL_MAT && dst->ne[1] != 1) {
@@ -1845,11 +1846,14 @@ static bool ggml_cuda_match_shared_expert(const ggml_cgraph * graph, int routed_
         return (a == graph->nodes[idx] && b == graph->nodes[idx + 1]) ||
                (b == graph->nodes[idx] && a == graph->nodes[idx + 1]);
     };
+    // only batch-size independent checks here: graph_optimize must produce the same graph topology for every ubatch
+    // size, otherwise ggml-alloc has to re-reserve (and the scheduler to synchronize) at runtime.
+    // the MMVQ batch size check is done in ggml_cuda_try_fuse
     if (!is_pair(gate, up, routed_idx) || !is_pair(shared_gate, shared_up, shared_idx) ||
             !ggml_cuda_should_fuse_mul_mat(up, gate, routed) ||
             !ggml_cuda_should_fuse_mul_mat(shared_up, shared_gate, shared) ||
             !up->src[0]->buffer ||
-            !ggml_cuda_should_fuse_mul_mat_vec_q(up)) {
+            !ggml_is_quantized(up->src[0]->type)) {
         return false;
     }
     const ggml_tensor * input = up->src[1];
@@ -1892,7 +1896,7 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor
     const int cc        = ggml_cuda_info().devices[ctx.device].cc;
     const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size;
 
-    if (ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, ne11)) {
+    if (ggml_cuda_should_use_mmvf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11)) {
         // The custom F16 vector kernel can be used over batched cuBLAS GEMM.
         // But this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention)
         ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst);
@@ -1902,7 +1906,7 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor
     if (ne01 == 1 && ne11 > MMVF_MAX_BATCH_SIZE && ne2 == 1 && ne3 == 1
             && src0->type == GGML_TYPE_F32
             && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst)
-            && ggml_cuda_should_use_mmvf(src1->type, cc, src1->ne, src1->nb, /*ne11 =*/ 1)) {
+            && ggml_cuda_should_use_mmvf(src1->type, cc, warp_size, src1->ne, src1->nb, /*ne11 =*/ 1)) {
         ggml_tensor dst_vec = *dst;
         dst_vec.ne[0] = ne11;
         dst_vec.ne[1] = 1;
@@ -3163,7 +3167,7 @@ static bool ggml_cuda_match_moe_weighted_reduction(
 
     const int     n_expert_used = (int) weighted->ne[1];
     const int64_t n_tokens      = weighted->ne[2] * weighted->ne[3];
-    if (n_expert_used < 2 || n_expert_used > MOE_WEIGHTED_REDUCTION_MAX_EXPERTS || n_tokens <= 0) {
+    if (n_expert_used < 2 || n_expert_used > MOE_WEIGHTED_REDUCTION_MAX_EXPERTS) {
         return false;
     }
 
@@ -3509,7 +3513,8 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
     ggml_tensor * node = cgraph->nodes[i];
 
     if (node->op == GGML_OP_MUL_MAT_ID && cuda_ctx->stream_context().concurrent_events.empty() &&
-            ggml_cuda_match_shared_expert(cgraph, i, i + 3)) {
+            ggml_cuda_match_shared_expert(cgraph, i, i + 3) &&
+            ggml_cuda_should_fuse_mul_mat_vec_q(cgraph->nodes[i + 2]->src[1])) {
         const int outputs[] = { i + 2, i + 5 };
         if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, 6, outputs, 2)) {
             ggml_tensor * routed = cgraph->nodes[i + 2];

ggml/src/ggml-cuda/lightning-indexer.cu

@@ -236,6 +236,9 @@ static __global__ void lightning_indexer_kernel_wmma(
 #endif // defined(TURING_MMA_AVAILABLE)
 #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
 
+// tokens scored per block by the tile kernel
+#define LIGHTNING_INDEXER_TILE_TOKENS 8
+
 // TODO there is one ugly assumption used in this kernel - that WARP_SIZE is equal to 32
 // thanks to that one warp operating on float4 processes whole indexer K/Q vectors
 // 32 * 4 = 128 (N_EMBD)
@@ -382,6 +385,152 @@ static __global__ void lightning_indexer_kernel_vec(
     }
 }
 
+// one block scores a tile of K_VECS_PER_BLOCK keys against TOKENS_PER_BLOCK tokens: the keys are
+// staged in half precision and the queries of every head in float, each thread owns KEYS_PER_THREAD
+// keys for one token, a warp shares its token so the query reads are broadcasts, and every key
+// element is widened once for all heads, so no dot product needs a cross thread reduction
+template <int WARPS_PER_BLOCK, int K_VECS_PER_BLOCK, int64_t N_EMBD, int64_t N_HEAD, ggml_type TYPE_K>
+static __global__ void lightning_indexer_kernel_tile(
+        const float * Q, const char * K, const float * W, const half * M, float * dst,
+        int64_t n_stream, int64_t n_batch, int64_t n_kv,
+        size_t nb1, size_t nb2, size_t nb3,
+        size_t nbq1, size_t nbq2, size_t nbq3,
+        size_t nbk1, size_t nbk2, size_t nbk3,
+        size_t nbw1, size_t nbw2, size_t nbw3,
+        size_t nbm1, size_t nbm2, size_t nbm3,
+        int64_t nem3
+    ) {
+
+    constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WARP_SIZE;
+    constexpr int TOKENS_PER_BLOCK  = LIGHTNING_INDEXER_TILE_TOKENS;
+    constexpr int KEY_LANES         = THREADS_PER_BLOCK / TOKENS_PER_BLOCK;
+    constexpr int KEYS_PER_THREAD   = K_VECS_PER_BLOCK / KEY_LANES;
+    constexpr int N_EMBD_H2         = N_EMBD / 2;
+
+    static_assert(THREADS_PER_BLOCK % TOKENS_PER_BLOCK == 0, "threads must cover the token tile");
+    static_assert(K_VECS_PER_BLOCK % KEY_LANES == 0, "key lanes must cover the key tile");
+
+    const int tid         = threadIdx.y * WARP_SIZE + threadIdx.x;
+    const int start_kv    = blockIdx.x * K_VECS_PER_BLOCK;
+    const int start_batch = blockIdx.y * TOKENS_PER_BLOCK;
+    const int i_stream    = blockIdx.z;
+
+    // the row padding keeps the keys of consecutive threads in distinct banks
+    __shared__ half2 k_shared[K_VECS_PER_BLOCK][N_EMBD_H2 + 1];
+    __shared__ float2 q_shared[N_HEAD][TOKENS_PER_BLOCK][N_EMBD_H2];
+    __shared__ float w_shared[N_HEAD][TOKENS_PER_BLOCK];
+
+    // phase 1 - stage the key tile four elements at a time, rows past n_kv are zero
+
+#pragma unroll
+    for (int i = tid; i < K_VECS_PER_BLOCK * (N_EMBD / 4); i += THREADS_PER_BLOCK) {
+        const int r  = i / (N_EMBD / 4);
+        const int c4 = i % (N_EMBD / 4);
+
+        half2 lo = make_half2(0.0f, 0.0f);
+        half2 hi = lo;
+        if (start_kv + r < n_kv) {
+            const char * k_row = K + (start_kv + r)*nbk2 + i_stream*nbk3;
+            if constexpr (TYPE_K == GGML_TYPE_F16) {
+                lo = ((const half2 *) k_row)[2*c4 + 0];
+                hi = ((const half2 *) k_row)[2*c4 + 1];
+            } else {
+                float4 v;
+                if constexpr (TYPE_K == GGML_TYPE_F32) {
+                    v = ((const float4 *) k_row)[c4];
+                } else {
+                    constexpr dequantize_V_t dequantize_k = get_dequantize_V<TYPE_K, float, 4>();
+                    dequantize_k(k_row, &v, c4 * 4);
+                }
+                lo = make_half2(v.x, v.y);
+                hi = make_half2(v.z, v.w);
+            }
+        }
+
+        k_shared[r][2*c4 + 0] = lo;
+        k_shared[r][2*c4 + 1] = hi;
+    }
+
+    // phase 2 - stage the queries and weights of every head, tokens past n_batch are zero
+
+#pragma unroll
+    for (int i = tid; i < N_HEAD * TOKENS_PER_BLOCK * (N_EMBD / 4); i += THREADS_PER_BLOCK) {
+        const int h  = i / (TOKENS_PER_BLOCK * (N_EMBD / 4));
+        const int r  = i / (N_EMBD / 4) % TOKENS_PER_BLOCK;
+        const int c4 = i % (N_EMBD / 4);
+
+        float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+        if (start_batch + r < n_batch) {
+            v = *(const float4 *) ((const char *) Q + h*nbq1 + (start_batch + r)*nbq2 + i_stream*nbq3 + c4*sizeof(float4));
+        }
+
+        q_shared[h][r][2*c4 + 0] = make_float2(v.x, v.y);
+        q_shared[h][r][2*c4 + 1] = make_float2(v.z, v.w);
+    }
+
+    if (tid < N_HEAD * TOKENS_PER_BLOCK) {
+        const int h = tid / TOKENS_PER_BLOCK;
+        const int r = tid % TOKENS_PER_BLOCK;
+        w_shared[h][r] = start_batch + r < n_batch ?
+            ((const float *) ((const char *) W + (start_batch + r)*nbw1 + i_stream*nbw3))[h] : 0.0f;
+    }
+
+    __syncthreads();
+
+    // phase 3 - float products of the widened keys for every head, ReLU, weight
+
+    const int kl = tid % KEY_LANES;
+    const int tl = tid / KEY_LANES;
+
+    float qk[N_HEAD][KEYS_PER_THREAD] = { { 0.0f } };
+
+#pragma unroll 8
+    for (int c = 0; c < N_EMBD_H2; ++c) {
+        float2 k_val[KEYS_PER_THREAD];
+#pragma unroll
+        for (int j = 0; j < KEYS_PER_THREAD; ++j) {
+            k_val[j] = __half22float2(k_shared[kl + j*KEY_LANES][c]);
+        }
+#pragma unroll
+        for (int h = 0; h < N_HEAD; ++h) {
+            const float2 q_val = q_shared[h][tl][c];
+#pragma unroll
+            for (int j = 0; j < KEYS_PER_THREAD; ++j) {
+                qk[h][j] = fmaf(k_val[j].x, q_val.x, qk[h][j]);
+                qk[h][j] = fmaf(k_val[j].y, q_val.y, qk[h][j]);
+            }
+        }
+    }
+
+    float score[KEYS_PER_THREAD] = { 0.0f };
+
+#pragma unroll
+    for (int h = 0; h < N_HEAD; ++h) {
+#pragma unroll
+        for (int j = 0; j < KEYS_PER_THREAD; ++j) {
+            score[j] += fmaxf(qk[h][j], 0.0f) * w_shared[h][tl];
+        }
+    }
+
+    // phase 4 - add the mask and write, consecutive threads write consecutive keys
+
+    const int i_batch = start_batch + tl;
+    if (i_batch >= n_batch) {
+        return;
+    }
+
+    const half * m_base = (const half *) ((const char *) M + i_batch*nbm1 + (i_stream%nem3)*nbm3);
+    float * dst_base = (float *) ((char *) dst + i_batch*nb1 + i_stream*nb3);
+
+#pragma unroll
+    for (int j = 0; j < KEYS_PER_THREAD; ++j) {
+        const int i_kv = start_kv + kl + j*KEY_LANES;
+        if (i_kv < n_kv) {
+            dst_base[i_kv] = score[j] + __half2float(m_base[i_kv]);
+        }
+    }
+}
+
 #define LIGHTNING_INDEXER_CASE(lightning_indexer_kernel, n_embd, n_head, K, type_K)         \
     if (K->type == (type_K)) {                                                              \
         lightning_indexer_kernel<WARPS_PER_BLOCK, K_VECS_PER_BLOCK, n_embd, n_head, type_K> \
@@ -528,8 +677,30 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor *
             LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_F32)
             GGML_ABORT("fatal error");
         }
+#ifndef GGML_USE_MUSA
+    // MUSA archs 21 and 22 cap static shared memory at 28 KB, below what the tile kernel stages
+    } else if (n_embd == 128 && n_head == 4 && n_batch >= LIGHTNING_INDEXER_TILE_TOKENS) {
+        // too few heads for a wmma tile, the tile kernel shares the keys across the tokens
+        constexpr int WARPS_PER_BLOCK = 8;
+        constexpr int K_VECS_PER_BLOCK = 64;
+
+        dim3 block(32, WARPS_PER_BLOCK);
+        int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK);
+        int num_batch_blocks = (n_batch + LIGHTNING_INDEXER_TILE_TOKENS - 1) / LIGHTNING_INDEXER_TILE_TOKENS;
+        dim3 grid(num_kv_blocks, num_batch_blocks, n_stream);
+
+        LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_F16)
+        LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_Q4_0)
+        LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_Q4_1)
+        LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_Q5_0)
+        LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_Q5_1)
+        LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_Q8_0)
+        LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_BF16)
+        LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_F32)
+        GGML_ABORT("fatal error");
+#endif // GGML_USE_MUSA
     } else if (n_embd == 128 && n_head == 4) {
-        // too few heads for a wmma tile, use vector kernel
+        // a batch smaller than a token tile, or MUSA, use vector kernel
         constexpr int K_VECS_PER_WARP = 8;
         constexpr int WARPS_PER_BLOCK = 8;
         constexpr int K_VECS_PER_BLOCK = K_VECS_PER_WARP * WARPS_PER_BLOCK;

ggml/src/ggml-cuda/mma.cuh

@@ -782,18 +782,27 @@ namespace ggml_cuda_mma {
         }
     }
 
-    // Byte offset of tile element (i, j). If swz, XOR swizzle it to avoid bank conflicts without row padding.
-    template <bool swz, typename T>
-    static __device__ __forceinline__ int swizzle_bytes(const int i, const int j, const int stride) {
-        static_assert(!swz || sizeof(T) == 4, "swizzled tiles need 32 bit elements");
-        const int off = (i*stride + j) * (int) sizeof(T);
-        return swz ? off ^ ((i & 7) << 4) : off;
+    template <int stride, typename T>
+    static __device__ __forceinline__ uint32_t swizzle(const uint32_t offset, const uint32_t i) {
+        static_assert(sizeof(T) <= 4, "unsupported type size");
+        constexpr int stride_bytes = stride*sizeof(T);
+        static_assert(stride_bytes % 16 == 0, "bad stride");
+        constexpr uint32_t shift = sizeof(T) == 1 ? 4 : (sizeof(T) == 2 ? 3 : 2);
+        if (stride_bytes % 32 != 0) {
+            return offset; // Equivalent to padding with 16 bytes.
+        }
+        if (stride_bytes % 64 != 0) {
+            return offset ^ (((i / 4) % 2) << shift);
+        }
+        if (stride_bytes % 128 != 0) {
+            return offset ^ (((i / 2) % 4) << shift);
+        }
+        return offset ^ ((i % 8) << shift);
     }
 
-    template <bool swz, typename T>
-    static __device__ __forceinline__ const T * swizzle(
-            const T * __restrict__ tile_base, const int i, const int j, const int stride) {
-        return (const T *) ((const char *) tile_base + swizzle_bytes<swz, T>(i, j, stride));
+    template <int stride, typename T>
+    static __device__ __forceinline__ T * swizzle(T * ptr, const uint32_t offset, const uint32_t i) {
+        return ptr + swizzle<stride, T>(offset, i);
     }
 
     template <typename T>
@@ -872,29 +881,6 @@ namespace ggml_cuda_mma {
 #endif // TURING_MMA_AVAILABLE
     }
 
-    // Load from tile element (i0, j0), swz tells if the tile is stored swizzled.
-    template <bool swz, int I, int J, typename T, data_layout dl>
-    static __device__ __forceinline__ void load_ldmatrix(
-            tile<I, J, T, dl> & t, const T * __restrict__ tile_base, const int i0, const int j0, const int stride) {
-        if constexpr (!swz) {
-            load_ldmatrix(t, tile_base + i0*stride + j0, stride);
-            return;
-        }
-#if defined(TURING_MMA_AVAILABLE)
-        static_assert(I == 16, "bad tile width");
-        static_assert(J ==  8, "bad tile height");
-        const int i = i0 + threadIdx.x % t.I;
-        const int j = j0 + (threadIdx.x / t.I) * (t.J / 2);
-        int * xi = (int *) t.x;
-        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"(swizzle<true>(tile_base, i, j, stride)));
-#else
-        GGML_UNUSED_VARS(t, tile_base, i0, j0, stride);
-        NO_DEVICE_CODE;
-#endif // defined(TURING_MMA_AVAILABLE)
-    }
-
     static __device__ __forceinline__ void load_ldmatrix(
             tile<8, 4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int stride) {
         ggml_cuda_memcpy_1<4*sizeof(half2)>(t.x, xs0 + t.get_i(0)*stride);
@@ -902,10 +888,15 @@ namespace ggml_cuda_mma {
 
     static __device__ __forceinline__ void load_ldmatrix(
             tile<8, 4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int stride) {
+#ifdef VOLTA_MMA_AVAILABLE
 #pragma unroll
         for (int l0 = 0; l0 < t.ne; l0 += 2) {
             ggml_cuda_memcpy_1<2*sizeof(half2)>(t.x + l0, xs0 + t.get_i(l0)*stride + t.get_j(l0));
         }
+#else
+        GGML_UNUSED_VARS(t, xs0, stride);
+        NO_DEVICE_CODE;
+#endif // VOLTA_MMA_AVAILABLE
     }
 
     static __device__ __forceinline__ void load_ldmatrix(
@@ -954,25 +945,112 @@ namespace ggml_cuda_mma {
 #endif // TURING_MMA_AVAILABLE
     }
 
-    // Load from tile element (i0, j0), swz tells if the tile is stored swizzled.
-    template <bool swz, int I, typename T, data_layout dl>
-    static __device__ __forceinline__ void load_ldmatrix_trans(
-            tile<I, 8, T, dl> & t, const T * __restrict__ tile_base, const int i0, const int j0, const int stride) {
-        if constexpr (!swz) {
-            load_ldmatrix_trans(t, tile_base + i0*stride + j0, stride);
-            return;
+    template <int stride, int I, int J, typename T, data_layout dl>
+    static __device__ __forceinline__ void load_ldmatrix_swizzled(
+            tile<I, J, T, dl> & t, const T * __restrict__ xs0, const int offset) {
+#if defined(TURING_MMA_AVAILABLE)
+        static_assert(I == 16, "bad tile width");
+        static_assert(J ==  8, "bad tile height");
+        const int i = threadIdx.x % t.I;
+        const int j = (threadIdx.x / t.I) * (t.J / 2);
+        int offset_ij = offset + i * stride + j;
+        offset_ij = swizzle<stride, T>(offset_ij, i);
+        int * xi = (int *) t.x;
+        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"(xs0 + offset_ij));
+#elif defined(VOLTA_MMA_AVAILABLE)
+#pragma unroll
+        for (int o = 0; o < t.ne; o += 4) {
+            const int offset_ij = offset + t.get_i(o) * stride + o;
+            ggml_cuda_memcpy_1<4*sizeof(T)>(t.x + o, swizzle<stride>(xs0, offset_ij, t.get_i(o)));
         }
+#elif defined(AMD_WMMA_AVAILABLE)
+#ifdef RDNA3
+        static_assert(dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout");
+        static_assert(sizeof(t.x) == 32, "bad ne");
+        static_assert(I == 16, "bad tile width");
+        static_assert(J ==  8, "bad tile height");
+#pragma unroll
+        for (int o = 0; o < 8; o += 4) {
+            const int offset_ij = offset + t.get_i(0) * stride + o;
+            ggml_cuda_memcpy_1<16>(t.x + o, swizzle<stride>(xs0, offset_ij, t.get_i(0)));
+        }
+#else
+        static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout");
+        static_assert(sizeof(t.x) == 16, "bad ne");
+        const int offset_ij = offset + t.get_i(0)*stride + t.get_j(0);
+        ggml_cuda_memcpy_1<16>(t.x, swizzle<stride>(xs0, offset_ij, t.get_i(0)));
+#endif // RDNA3
+#elif defined(AMD_MFMA_AVAILABLE)
+        static_assert(sizeof(t.x) == 8, "bad ne");
+        const int offset_ij = offset + t.get_i(0)*stride + t.get_j(0);
+        ggml_cuda_memcpy_1<8>(t.x, swizzle<stride>(xs0, offset_ij, t.get_i(0)));
+#else
+        GGML_UNUSED_VARS(t, xs0, offset);
+        NO_DEVICE_CODE;
+#endif // defined(TURING_MMA_AVAILABLE)
+    }
+
+    template <int stride>
+    static __device__ __forceinline__ void load_ldmatrix_swizzled(
+            tile<8, 4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int offset) {
+#ifdef VOLTA_MMA_AVAILABLE
+#pragma unroll
+        for (int l0 = 0; l0 < t.ne; l0 += 2) {
+            const int offset_ij = offset + t.get_i(l0)*stride + t.get_j(l0);
+            ggml_cuda_memcpy_1<2*sizeof(half2)>(t.x + l0, swizzle<stride>(xs0, offset_ij, t.get_i(l0)));
+        }
+#else
+        GGML_UNUSED_VARS(t, xs0, offset);
+        NO_DEVICE_CODE;
+#endif // VOLTA_MMA_AVAILABLE
+    }
+
+    template <int stride, int I, typename T, data_layout dl>
+    static __device__ __forceinline__ void load_ldmatrix_trans_swizzled(
+            tile<I, 8, T, dl> & t, const T * __restrict__ xs0, const int offset) {
 #if defined(TURING_MMA_AVAILABLE)
         static_assert(I == 16, "bad tile width");
         static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout");
-        const int i = i0 + threadIdx.x % t.I;
-        const int j = j0 + (threadIdx.x / t.I) * (t.J / 2);
+        const int i = threadIdx.x % t.I;
+        const int j = (threadIdx.x / t.I) * (t.J / 2);
+        int offset_ij = offset + i * stride + j;
+        offset_ij = swizzle<stride, T>(offset_ij, i);
         int * xi = (int *) t.x;
         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"(swizzle<true>(tile_base, i, j, stride)));
+            : "l"(xs0 + offset_ij));
+#elif defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+        static_assert(dl == DATA_LAYOUT_I_MAJOR || dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout");
+        if constexpr (I == 32) {
+#pragma unroll
+            for (int l0 = 0; l0 < t.ne/2; ++l0) {
+                half2 tmp[2];
+#pragma unroll
+                for (int o = 0; o < 2; ++o) {
+                    const int j = 2*t.get_j(l0) + o;
+                    int offset_ij = offset + j*stride + t.get_i(l0)/2;
+                    offset_ij = swizzle<stride, T>(offset_ij, j);
+                    tmp[o] = xs0[offset_ij];
+                }
+
+                t.x[l0]          =  __lows2half2(tmp[0], tmp[1]);
+                t.x[l0 + t.ne/2] = __highs2half2(tmp[0], tmp[1]);
+            }
+        } else {
+            half * xh = (half *) t.x;
+#pragma unroll
+            for (int l = 0; l < t.ne; ++l) {
+#pragma unroll
+                for (int o = 0; o < 2; ++o) {
+                    const int j = 2*t.get_j(l) + o;
+                    xh[2*l + o] = ((const half *) xs0)[swizzle<2*stride, half>(2*offset + j*(2*stride) + t.get_i(l), j)];
+                }
+            }
+        }
 #else
-        GGML_UNUSED_VARS(t, tile_base, i0, j0, stride);
+        GGML_UNUSED_VARS(t, xs0, offset);
         NO_DEVICE_CODE;
 #endif // defined(TURING_MMA_AVAILABLE)
     }

ggml/src/ggml-cuda/mmid.cu

@@ -37,9 +37,6 @@ static __global__ void mm_ids_helper(
     const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
     const int expert = blockIdx.x;
 
-    // token slots per warp lane group, padded to a power of 2 so a warp divides evenly
-    constexpr int neu_padded = mm_ids_pow2<n_expert_used_template>::value;
-
     extern __shared__ char data_mm_ids_helper[];
     mm_ids_helper_store * store = (mm_ids_helper_store *) data_mm_ids_helper;
 
@@ -69,6 +66,7 @@ static __global__ void mm_ids_helper(
     } else {
         // Implementation optimized for specific numbers of experts used:
         // a warp holds a whole number of token slots, so the slot count is padded to a power of 2
+        constexpr int neu_padded = mm_ids_pow2<n_expert_used_template>::value;
         static_assert(neu_padded <= warp_size && warp_size % neu_padded == 0, "bad n_expert_used");
         for (int it0 = 0; it0 < n_tokens; it0 += warp_size/neu_padded) {
             const int it = it0 + threadIdx.x / neu_padded;

ggml/src/ggml-cuda/mmq-vec-dot.cuh

@@ -1225,14 +1225,11 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
 
 #pragma unroll
         for (int n = 0; n < ntx; ++n) {
+            // accumulate in place into the output sum array
+            tile_C & C = *reinterpret_cast<tile_C *>(sum + (j0 / tile_C::J + n) * tile_C::ne);
 #pragma unroll
             for (int frag = 0; frag < nfrags; ++frag) {
-                tile_C C = {};
                 mma_block_scaled_fp4<type>(C, A[n][frag], B[frag], scaleA[n][frag], scaleB[frag]);
-#pragma unroll
-                for (int l = 0; l < tile_C::ne; ++l) {
-                    sum[(j0 / tile_C::J + n) * tile_C::ne + l] += C.x[l];
-                }
             }
         }
     }

ggml/src/ggml-cuda/mmvf.cu

@@ -2,6 +2,7 @@
 #include "common.cuh"
 #include "unary.cuh"
 #include "mmvf.cuh"
+#include "mmf.cuh"
 #include "convert.cuh"
 
 template <typename T, typename type_acc, int ncols_dst, int block_size, bool has_fusion = false, bool is_multi_token_id = false>
@@ -789,7 +790,7 @@ void ggml_cuda_op_mul_mat_vec_f(
     GGML_UNUSED_VARS(ctx, src1, dst, src1_ddq_i, src1_ncols, src1_padded_row_size);
 }
 
-bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0_ne, const size_t * src0_nb, int64_t ne11) {
+bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, int warp_size, const int64_t * src0_ne, const size_t * src0_nb, int64_t ne11) {
     if (src0_ne[0] % 2 != 0) {
         return false;
     }
@@ -820,12 +821,16 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0
                 if (fp32_mma_hardware_available(cc)) {
                     return ne11 <= 3;
                 }
-                return ne11 <= 8;
+                return ne11 <= MMVF_MAX_BATCH_SIZE;
             }
-            return ne11 <= 8;
+            return ne11 <= MMVF_MAX_BATCH_SIZE;
         case GGML_TYPE_F16:
             if (GGML_CUDA_CC_IS_NVIDIA(cc)) {
                 const bool src0_small = (src0_ne[1] <= 512 || src0_ne[2]*src0_ne[3] == 1);
+                // MMF needs full row tiles, for other row counts MMVF still beats cuBLAS at small batch size
+                if (src0_small && !ggml_cuda_should_use_mmf(type, cc, warp_size, src0_ne, src0_nb, ne11, /*mul_mat_id =*/ false)) {
+                    return ne11 <= MMVF_MAX_BATCH_SIZE;
+                }
                 if (ampere_mma_available(cc)) {
                     return src0_small && ne11 == 1;
                 }
@@ -835,8 +840,11 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0
                 if (fp16_mma_hardware_available(cc)) {
                     return src0_small && ne11 <= 3;
                 }
-                return ne11 <= 8;
+                return ne11 <= MMVF_MAX_BATCH_SIZE;
             } else if (GGML_CUDA_CC_IS_AMD(cc)) {
+                if (GGML_CUDA_CC_IS_RDNA(cc) && !ggml_cuda_should_use_mmf(type, cc, warp_size, src0_ne, src0_nb, ne11, /*mul_mat_id =*/ false)) {
+                    return ne11 <= MMVF_MAX_BATCH_SIZE;
+                }
                 if (fp16_mma_hardware_available(cc)) {
                     if (GGML_CUDA_CC_IS_RDNA3(cc)) {
                         return ne11 <= 3;
@@ -846,12 +854,16 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0
                     }
                     return ne11 <= 2;
                 }
-                return ne11 <= 8;
+                return ne11 <= MMVF_MAX_BATCH_SIZE;
             }
-            return ne11 <= 8;
+            return ne11 <= MMVF_MAX_BATCH_SIZE;
         case GGML_TYPE_BF16:
             if (GGML_CUDA_CC_IS_NVIDIA(cc)) {
                 const bool src0_small = (src0_ne[1] <= 512 || src0_ne[2]*src0_ne[3] == 1);
+                // MMF needs full row tiles, for other row counts MMVF still beats cuBLAS at small batch size
+                if (src0_small && !ggml_cuda_should_use_mmf(type, cc, warp_size, src0_ne, src0_nb, ne11, /*mul_mat_id =*/ false)) {
+                    return ne11 <= MMVF_MAX_BATCH_SIZE;
+                }
                 if (ampere_mma_available(cc)) {
                     return src0_small && ne11 == 1;
                 }
@@ -861,14 +873,17 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0
                 if (bf16_mma_hardware_available(cc)) {
                     return src0_small && ne11 <= 3;
                 }
-                return ne11 <= 8;
+                return ne11 <= MMVF_MAX_BATCH_SIZE;
             } else if (GGML_CUDA_CC_IS_AMD(cc)) {
+                if (GGML_CUDA_CC_IS_RDNA(cc) && !ggml_cuda_should_use_mmf(type, cc, warp_size, src0_ne, src0_nb, ne11, /*mul_mat_id =*/ false)) {
+                    return ne11 <= MMVF_MAX_BATCH_SIZE;
+                }
                 if (bf16_mma_hardware_available(cc)) {
                     return ne11 <= 3;
                 }
-                return ne11 <= 8;
+                return ne11 <= MMVF_MAX_BATCH_SIZE;
             }
-            return ne11 <= 8;
+            return ne11 <= MMVF_MAX_BATCH_SIZE;
         default:
             return false;
     }

ggml/src/ggml-cuda/mmvf.cuh

@@ -11,4 +11,4 @@ void ggml_cuda_op_mul_mat_vec_f(
     const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols,
     const int64_t src1_padded_row_size, cudaStream_t stream);
 
-bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0_ne, const size_t * src0_nb, int64_t ne11);
+bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, int warp_size, const int64_t * src0_ne, const size_t * src0_nb, int64_t ne11);

ggml/src/ggml-hexagon/ggml-hexagon.cpp

@@ -1,3 +1,5 @@
+#define _USE_MATH_DEFINES
+
 #include <assert.h>
 #include <inttypes.h>
 #include <stdio.h>
@@ -24,6 +26,10 @@
 #include <cmath>
 #include <initializer_list>
 
+#ifndef M_LOG2E
+#    define M_LOG2E 1.44269504088896340736
+#endif
+
 #ifdef _WIN32
 #    define WIN32_LEAN_AND_MEAN
 #    ifndef NOMINMAX
@@ -60,6 +66,7 @@
 #include "htp/get-rows-ops.h"
 #include "htp/set-rows-ops.h"
 #include "htp/softmax-ops.h"
+#include "htp/pool-ops.h"
 #include "htp/rope-ops.h"
 #include "htp/ssm-conv.h"
 #include "htp/gated-delta-net-ops.h"
@@ -101,6 +108,7 @@ static bool   opt_dma64   = false;
 
 static int    opt_mm_select  = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
 static int    opt_fa_select  = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
+static int    opt_fa_head_split = 1; // 1 = partition flash_attn by KV heads in multicore (default on), 0 = token-based (original)
 static int    opt_gdn_select = 2; // 2 = HMX -> HVX, 1 = HVX, 0 = CPU (unsupported)
 static int    opt_ar_select  = 2; // 2 = fused ALLREDUCE+ADD (default), 1 = unfused ALLREDUCE, 0 = fallback to CPY+FENCE
 static int    opt_ar_scatter = 1; // 1 = reduce-scatter the fused ALLREDUCE+ADD (default), 0 = full reduction
@@ -384,10 +392,18 @@ static void ggml_hexagon_precompute_sort_params(
     struct htp_sort_kernel_params * kparams
 );
 
+static void ggml_hexagon_precompute_pool_2d_params(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * dst,
+    struct htp_pool_2d_kernel_params * kparams,
+    bool is_pool_1d
+);
+
 static void ggml_hexagon_precompute_fused_mmnx_params(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * src0,
-    const struct ggml_tensor * src1,
+    const struct ggml_tensor * act,
     int32_t n_weights,
     struct htp_mm_kernel_params * kparams
 );
@@ -395,7 +411,7 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
 static void ggml_hexagon_precompute_fused_mmidnx_params(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * src0,
-    const struct ggml_tensor * src1,
+    const struct ggml_tensor * act,
     const struct ggml_tensor * dst,
     int32_t n_weights,
     struct htp_mm_kernel_params * kparams
@@ -412,6 +428,10 @@ static bool ggml_hexagon_precompute_allreduce_params(
     struct htp_allreduce_kernel_params * kparams
 );
 
+static bool ggml_hexagon_rows_stride(const int64_t * ne, const size_t * nb, size_t * stride);
+static bool ggml_hexagon_matmul_can_collapse(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst);
+static ggml_tensor ggml_hexagon_tensor_collapse_rows(const struct ggml_tensor * t);
+
 static bool mm_is_hmx_eligible(const ggml_tensor * t);
 static htp_op_code op_remap_to_htp(const ggml_tensor * t);
 static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams);
@@ -444,6 +464,13 @@ static inline bool ggml_hexagon_tensors_overlap(const struct ggml_tensor * a, co
     return a0 < b1 && b0 < a1;
 }
 
+static inline bool ggml_hexagon_can_row_partition(const struct ggml_tensor * t) {
+    if (t->ne[1] > 1 && (t->nb[1] & 127) != 0) return false;
+    if (t->ne[2] > 1 && (t->nb[2] & 127) != 0) return false;
+    if (t->ne[3] > 1 && (t->nb[3] & 127) != 0) return false;
+    return true;
+}
+
 struct htp_opnode;
 
 struct ggml_hexagon_opbatch;
@@ -3410,6 +3437,10 @@ struct ggml_hexagon_opbatch {
             return false;
         }
 
+        if (orig_kparams->collapse != kparams.collapse) {
+            return false;
+        }
+
         const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3];
         const bool can_fuse = (kparams.n_hmx > 0) || (src1_nrows == 1);
         if (!can_fuse) return false;
@@ -3493,8 +3524,20 @@ struct ggml_hexagon_opbatch {
                 return false;
             }
 
+            const struct htp_mm_kernel_params * orig_kparams = (const struct htp_mm_kernel_params *) last_node.kernel_params;
+            const bool collapse = orig_kparams->collapse && ggml_hexagon_matmul_can_collapse(w_in, x, d_in);
+            if (orig_kparams->collapse && !collapse) {
+                return false;
+            }
+
             struct htp_mm_kernel_params kparams;
-            ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, curr_n + 1, &kparams);
+            if (collapse) {
+                const ggml_tensor x_collapsed = ggml_hexagon_tensor_collapse_rows(x);
+                ggml_hexagon_precompute_fused_mmnx_params(sess, w0, &x_collapsed, curr_n + 1, &kparams);
+                kparams.collapse = 1;
+            } else {
+                ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, curr_n + 1, &kparams);
+            }
             if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) {
                 return false;
             }
@@ -3543,9 +3586,23 @@ struct ggml_hexagon_opbatch {
             const ggml_tensor * w0 = last_node.src0();
             const ggml_tensor * x  = last_node.src1();
             const ggml_tensor * w1 = node.src0();
+            const ggml_tensor * dst_0 = last_node.dst();
+            const ggml_tensor * dst_1 = node.dst();
+
+            const struct htp_mm_kernel_params * orig_kparams = (const struct htp_mm_kernel_params *) last_node.kernel_params;
+            const bool collapse = orig_kparams->collapse && ggml_hexagon_matmul_can_collapse(w1, x, dst_1);
+            if (orig_kparams->collapse && !collapse) {
+                return false;
+            }
 
             struct htp_mm_kernel_params kparams;
-            ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, 2, &kparams);
+            if (collapse) {
+                const ggml_tensor x_collapsed = ggml_hexagon_tensor_collapse_rows(x);
+                ggml_hexagon_precompute_fused_mmnx_params(sess, w0, &x_collapsed, 2, &kparams);
+                kparams.collapse = 1;
+            } else {
+                ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, 2, &kparams);
+            }
             if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) {
                 return false;
             }
@@ -3559,9 +3616,6 @@ struct ggml_hexagon_opbatch {
                 return false;
             }
 
-            const ggml_tensor * dst_0 = last_node.dst();
-            const ggml_tensor * dst_1 = node.dst();
-
             last_node.opcode = HTP_OP_MUL_MAT_NX;
             last_node.name   = "MUL_MAT_NX";
             last_node.inputs.clear();
@@ -4831,15 +4885,48 @@ static bool ggml_hexagon_flash_attn_is_hmx_eligible(
         return false;
     }
 
-    // Fall back to HVX for small token counts if head dimension is small (DK <= 128)
-    const uint32_t neq1 = q->ne[1];
-    if (DK <= 128 && neq1 < 5) {
-        return false;
+    GGML_UNUSED(sinks);
+
+    // Explicit force mode
+    if (opt_fa_select > 2) {
+        return true;
     }
 
-    return true;
+    const uint32_t M = q->ne[1];
 
-    GGML_UNUSED(sinks);
+    // Prefill or batched decode
+    if (M > 1) {
+        return true;
+    }
+
+    // Compute-bound head dim
+    if (DK >= 256) {
+        return true;
+    }
+
+    const uint32_t n_head = q->ne[2];
+    const uint32_t n_kv_heads = k->ne[2];
+    const uint32_t G = n_kv_heads > 0 ? n_head / n_kv_heads : 1;
+    const uint32_t S = k->ne[1];
+
+    // Tile alignment for 32-row HMX tiles
+    const bool is_tile_aligned = (G > 0 && (32 % G == 0));
+    if (!is_tile_aligned) {
+        if (sess->n_threads >= 6) {
+            return false;
+        }
+        return S >= 1024;
+    }
+
+    // Context depth crossover
+    uint32_t s_cross = 512;
+    if (DK <= 64) {
+        s_cross = (sess->n_threads >= 8) ? 2048 : ((sess->n_threads >= 6) ? 768 : 512);
+    } else {
+        s_cross = (sess->n_threads >= 8) ? 1024 : 512;
+    }
+
+    return S >= s_cross;
 }
 
 static bool ggml_hexagon_precompute_flash_attn_params(
@@ -4857,7 +4944,6 @@ static bool ggml_hexagon_precompute_flash_attn_params(
     const struct ggml_tensor * k    = op->src[1];
     const struct ggml_tensor * v    = op->src[2];
     const struct ggml_tensor * mask = op->src[3];
-    const struct ggml_tensor * dst  = op;
 
     const uint32_t neq0 = q->ne[0];  // head_dim (DK)
     const uint32_t neq1 = q->ne[1];  // n_tokens
@@ -4888,8 +4974,8 @@ static bool ggml_hexagon_precompute_flash_attn_params(
     kparams->max_bias = max_bias;
     kparams->logit_softcap = logit_softcap;
 
-    kparams->is_q_fp32 = (q->type == GGML_TYPE_F32) ? 1 : 0;
-    kparams->is_dst_fp32 = (dst->type == GGML_TYPE_F32) ? 1 : 0;
+    kparams->head_split = (opt_fa_head_split != 0) ? 1 : 0;
+    kparams->flags      = 0;
     kparams->G = G;
 
     const uint32_t n_head = q->ne[2];
@@ -4906,9 +4992,16 @@ static bool ggml_hexagon_precompute_flash_attn_params(
         const uint32_t DK_pad = hex_round_up(DK, 64);
         const uint32_t DV_pad = hex_round_up(DV, 64);
         size_t Br = 0, Bc = 0;
-        int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, sess->vtcm_size, sess->n_threads, kparams->is_q_fp32 != 0, sinks != nullptr, n_head);
+        int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, sess->vtcm_size, sess->n_threads, (q->type == GGML_TYPE_F32), sinks != nullptr, n_head);
         if (ret == 0) {
             kparams->kernel_type = HTP_FA_KERNEL_HMX;
+            if (logit_softcap == 0.0f) {
+                kparams->scale = scale * (float) M_LOG2E;
+                kparams->logit_softcap = 0.0f;
+            } else {
+                kparams->scale = scale;
+                kparams->logit_softcap = logit_softcap * (float) M_LOG2E;
+            }
             kparams->Br = Br;
             kparams->Bc = Bc;
             kparams->n_kv_blocks = (nek1 + Bc - 1) / Bc;
@@ -4916,7 +5009,7 @@ static bool ggml_hexagon_precompute_flash_attn_params(
 
             kparams->u.hmx.g_br = hex_align_up(G * Br, 32);
             kparams->u.hmx.pipeline = (kparams->n_kv_blocks >= 3 && sess->n_threads >= 2) ? 1 : 0;
-            kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK_pad, DV_pad, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0, kparams->is_q_fp32 != 0, sinks != nullptr, n_head);
+            kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK_pad, DV_pad, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0, (q->type == GGML_TYPE_F32), sinks != nullptr, n_head);
 
             const size_t row_vec_bytes = hex_align_up(Bc * sizeof(uint16_t), 256);
             kparams->u.hmx.row_buf_stride = row_vec_bytes / 128; // HVX vector is 128 bytes
@@ -4943,11 +5036,11 @@ static bool ggml_hexagon_precompute_flash_attn_params(
     kparams->n_kv_blocks = (k->ne[1] + 64 - 1) / 64;
     kparams->n_threads = sess->n_threads;
 
-    const size_t size_q_row_padded = hex_round_up(q->ne[0] * (kparams->is_q_fp32 ? 4 : 2), 128);
+    const size_t size_q_row_padded = hex_round_up(q->ne[0] * ((q->type == GGML_TYPE_F32) ? 4 : 2), 128);
     const size_t size_k_row_padded = hex_round_up(k->ne[0] * 2, 128);
     const size_t size_v_row_padded = hex_round_up(v->ne[0] * 2, 128);
 
-    kparams->vtcm_size = hvx_fa_compute_vtcm_usage(DK, DV, kparams->is_q_fp32 != 0, mask != nullptr, sinks != nullptr, n_head, sess->n_threads);
+    kparams->vtcm_size = hvx_fa_compute_vtcm_usage(DK, DV, (q->type == GGML_TYPE_F32), mask != nullptr, sinks != nullptr, n_head, sess->n_threads);
 
     kparams->u.hvx.size_q_row_padded = size_q_row_padded;
     kparams->u.hvx.size_k_row_padded = size_k_row_padded;
@@ -5102,7 +5195,7 @@ static bool ggml_hexagon_matmul_is_hmx_eligible(
     bool is_matmul_id,
     bool is_batched
 ) {
-    if (src1->type != GGML_TYPE_F32) {
+    if (src1->type != GGML_TYPE_F32 && (src1->type != GGML_TYPE_F16 || is_matmul_id)) {
         return false;
     }
 
@@ -5111,8 +5204,8 @@ static bool ggml_hexagon_matmul_is_hmx_eligible(
     const int ne12  = src1->ne[2];
     const int wtype = src0->type;
 
-    // HMX weight tile requires N to be 32-aligned.
-    if (ne01_padded % 32 != 0) {
+    // HMX weight tiles accept non-32-aligned N for non-matmul_id.
+    if (ne01_padded % 32 != 0 && is_matmul_id) {
         return false;
     }
 
@@ -5151,7 +5244,7 @@ static bool ggml_hexagon_matmul_is_hmx_eligible(
 static bool ggml_hexagon_precompute_hmx_mm_params(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * src0,
-    const struct ggml_tensor * src1,
+    const struct ggml_tensor * act,
     const struct ggml_tensor * dst,
     int wtype,
     int ne00_padded,
@@ -5167,9 +5260,22 @@ static bool ggml_hexagon_precompute_hmx_mm_params(
     struct htp_mm_kernel_params * kparams
 ) {
     const int aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
-    const bool pipeline = is_matmul_id ? false : htp_mm_hmx_pipeline(ne11);
     const int n_threads = (int)sess->n_threads;
-    const int ne10 = src1->ne[0];
+    const int ne10 = act->ne[0];
+
+    int m_for_solver = ne11;
+    int m_for_solver_padded = ne11_padded;
+    // matmul_id partitions by expert; regular matmul partitions M rows (ne11) across devices
+    if (!is_matmul_id && sess->mdev.count > 1 && ((uint32_t) ne11 >= sess->mdev.count)) {
+        // when dst is null, padded dims are used for estimate which are 128-byte aligned
+        const bool dst_row_split = dst ? ggml_hexagon_can_row_partition(dst) : true;
+        const bool act_row_split = ggml_hexagon_can_row_partition(act);
+        if (dst_row_split && act_row_split) {
+            m_for_solver = (ne11 + (int) sess->mdev.count - 1) / (int) sess->mdev.count;
+            m_for_solver_padded = hex_round_up(std::max(m_for_solver, 32), 32);
+        }
+    }
+    const bool pipeline = is_matmul_id ? false : htp_mm_hmx_pipeline(m_for_solver);
 
     const bool is_batched_val = is_matmul_id ? false : is_batched;
     const int group_size = (ne02 > 0 ? ne12 / ne02 : 1);
@@ -5182,15 +5288,25 @@ static bool ggml_hexagon_precompute_hmx_mm_params(
 
     if (is_batched_val && wtype == GGML_TYPE_F16 && group_size > 1) {
         // Try grouped path first
-        if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, ne11, group_size, n_threads, pipeline, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
+        if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, m_for_solver, group_size, n_threads, pipeline, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
             use_grouped = true;
         }
     }
 
     if (!use_grouped) {
         // Fallback to simple 2D path (group_size = 1)
-        const int m_id_rows = (dst && is_matmul_id) ? (int) ((size_t) dst->ne[1] * dst->ne[2]) : 0;
-        if (!htp_mm_hmx_solve_2d_params(wtype, ne00_padded, m_id_rows, ne01_padded, ne11_padded, ne11, n_threads, pipeline, is_matmul_id, aligned_tile_size, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
+        int m_id_rows = 0;
+        if (dst && is_matmul_id) {
+            const int n_experts = ne02 > 0 ? ne02 : 1;
+            const size_t total_expert_rows = (size_t) dst->ne[1] * dst->ne[2];
+            int m_per_expert = (int) ((total_expert_rows + n_experts - 1) / n_experts);
+            if (sess->mdev.count > 1 && ggml_hexagon_can_row_partition(dst)) {
+                m_per_expert = (m_per_expert + (int) sess->mdev.count - 1) / (int) sess->mdev.count;
+            }
+            m_id_rows = hex_round_up(std::max(m_per_expert, 32), 32);
+        }
+        const uint32_t cost_m = is_matmul_id ? (uint32_t) m_id_rows : (uint32_t) m_for_solver;
+        if (!htp_mm_hmx_solve_2d_params(wtype, ne00_padded, m_id_rows, ne01_padded, m_for_solver_padded, cost_m, n_threads, pipeline, is_matmul_id, aligned_tile_size, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
             return false;
         }
     }
@@ -5203,14 +5319,14 @@ static bool ggml_hexagon_precompute_hmx_mm_params(
     kparams->n_act_threads = act_threads_selected;
     kparams->tile_size = htp_mm_get_weight_tile_size(wtype);
     kparams->aligned_tile_size = aligned_tile_size;
-    kparams->src1_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
+    kparams->act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
     kparams->vtcm_size = vtcm_size;
     kparams->vtcm_src0_size = 0;
     kparams->div_n_act_threads = init_fastdiv_values(act_threads_selected);
     kparams->div_ne00_padded   = init_fastdiv_values(ne00_padded);
-    kparams->vtcm_src1_size = 0;
-    kparams->vtcm_src2_size = (int32_t) src2_size;
-    kparams->vtcm_dst_size = 0;
+    kparams->vtcm_act_size     = 0;
+    kparams->vtcm_bias_size    = (int32_t) src2_size;
+    kparams->vtcm_dst_size     = 0;
 
     if (is_batched && !is_matmul_id) {
         kparams->kernel_type = HTP_MM_KERNEL_HMX_F16_BATCHED;
@@ -5247,6 +5363,9 @@ static void ggml_hexagon_precompute_hvx_mm_params(
     kparams->n_hmx = 0;
     kparams->n_threads = sess->n_threads;
 
+    GGML_UNUSED(ne02);
+    GGML_UNUSED(ne03);
+
     const bool is_quant = (wtype != GGML_TYPE_F16 && wtype != GGML_TYPE_F32);
     const int src1_nrows = ne11 * ne12 * ne13;
 
@@ -5259,7 +5378,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
 
         if (is_matmul_id) {
             kparams->kernel_type   = (src1_nrows < (int) sess->n_threads) ? HTP_MM_KERNEL_HVX_QUANT_BLOCK : HTP_MM_KERNEL_HVX_QUANT_ROW;
-            kparams->src1_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
+            kparams->act_row_size  = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
 
             struct htp_mm_hvx_vtcm_layout L;
             uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
@@ -5267,7 +5386,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
             for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
                 htp_mm_hvx_vtcm_layout_build(
                     &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                    0, src0->nb[1], kparams->src1_row_size, 0, d, true, false
+                    0, src0->nb[1], kparams->act_row_size, 0, d, true, false
                 );
                 if (L.total_bytes <= vtcm_budget) {
                     best_n_prefetch = d;
@@ -5281,13 +5400,14 @@ static void ggml_hexagon_precompute_hvx_mm_params(
             kparams->n_prefetch     = best_n_prefetch;
             kparams->vtcm_size      = L.total_bytes;
             kparams->vtcm_src0_size = L.src0_bytes;
-            kparams->vtcm_src1_size = L.src1_bytes;
+            kparams->vtcm_act_size  = L.act_bytes;
+            kparams->vtcm_bias_size = 0;
             kparams->vtcm_dst_size  = L.dst_bytes;
             goto done_quant;
         } else {
             bool try_tiled = (k_align && opt_mm_select >= 1);
             if (try_tiled) {
-                kparams->src1_row_size = htp_mm_weight_has_offset(wtype)
+                kparams->act_row_size = htp_mm_weight_has_offset(wtype)
                                        ? htp_mm_q8_1_tiled_row_size(ne10)
                                        : htp_mm_q8_0_tiled_row_size(ne10);
                 if (src1_nrows < (int) sess->n_threads) {
@@ -5319,8 +5439,8 @@ static void ggml_hexagon_precompute_hvx_mm_params(
                     kparams->m_chunk        = (m_chunk < (uint32_t) src1_nrows) ? m_chunk : 0;
                     kparams->vtcm_size      = L.total_bytes;
                     kparams->vtcm_src0_size = L.src0_bytes;
-                    kparams->vtcm_src1_size = L.src1_bytes;
-                    kparams->vtcm_src2_size = L.src2_bytes;
+                    kparams->vtcm_act_size  = L.act_bytes;
+                    kparams->vtcm_bias_size = L.bias_bytes;
                     kparams->vtcm_dst_size  = L.dst_bytes;
                     goto done_quant;
                 }
@@ -5341,11 +5461,11 @@ static void ggml_hexagon_precompute_hvx_mm_params(
                 &L, &m_chunk)) {
             kparams->kernel_type = HTP_MM_KERNEL_HVX_F16_F16_VTCM;
             kparams->m_chunk = (m_chunk < (uint32_t) src1_nrows) ? m_chunk : 0;
-            kparams->src1_row_size = hex_round_up(ne10 * 2, 128);
+            kparams->act_row_size = hex_round_up(ne10 * 2, 128);
             kparams->vtcm_size = L.total_bytes;
             kparams->vtcm_src0_size = L.src0_bytes;
-            kparams->vtcm_src1_size = L.src1_bytes;
-            kparams->vtcm_src2_size = L.src2_bytes;
+            kparams->vtcm_act_size  = L.act_bytes;
+            kparams->vtcm_bias_size = L.bias_bytes;
             kparams->vtcm_dst_size = L.dst_bytes;
             kparams->n_prefetch = 16;
             return;
@@ -5363,11 +5483,11 @@ static void ggml_hexagon_precompute_hvx_mm_params(
                 &L, &m_chunk)) {
             kparams->kernel_type = HTP_MM_KERNEL_HVX_F32_F32_VTCM;
             kparams->m_chunk = (m_chunk < (uint32_t) src1_nrows) ? m_chunk : 0;
-            kparams->src1_row_size = hex_round_up(ne10 * 4, 128);
+            kparams->act_row_size = hex_round_up(ne10 * 4, 128);
             kparams->vtcm_size = L.total_bytes;
             kparams->vtcm_src0_size = L.src0_bytes;
-            kparams->vtcm_src1_size = L.src1_bytes;
-            kparams->vtcm_src2_size = L.src2_bytes;
+            kparams->vtcm_act_size  = L.act_bytes;
+            kparams->vtcm_bias_size = L.bias_bytes;
             kparams->vtcm_dst_size = L.dst_bytes;
             kparams->n_prefetch = 16;
             return;
@@ -5404,6 +5524,9 @@ static void ggml_hexagon_precompute_matmul_params_impl(
     const int ne00_padded = is_repack ? hex_round_up(ne00, 32) : ne00;
     const int ne01_padded = is_repack ? hex_round_up(ne01, 32) : ne01;
     const int ne11_padded = hex_round_up(ne11, 32);
+    // VTCM has to hold whole 32-row weight tiles, so size for the rounded-up N
+    // even when the tensor itself is ragged.
+    const int  ne01_tiled  = hex_round_up(ne01_padded, 32);
 
     const bool is_matmul_id = (dst->op == GGML_OP_MUL_MAT_ID);
     const bool is_batched   = (ne02 * ne03 > 1 || ne12 * ne13 > 1);
@@ -5413,7 +5536,7 @@ static void ggml_hexagon_precompute_matmul_params_impl(
     // Check HMX eligibility and try precomputing HMX parameters
     bool hmx_enabled = (sess->n_hmx > 0) && (opt_mm_select >= 2);
     if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, src1, dst, ne01_padded, is_matmul_id, is_batched)) {
-        if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, dst, wtype, ne00_padded, ne01_padded, ne02, ne11, ne12, ne11_padded, is_matmul_id, is_batched, src2_size, vtcm_budget, kparams)) {
+        if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, dst, wtype, ne00_padded, ne01_tiled, ne02, ne11, ne12, ne11_padded, is_matmul_id, is_batched, src2_size, vtcm_budget, kparams)) {
             goto finalize;
         }
     }
@@ -5429,6 +5552,66 @@ finalize:
     kparams->div_ne12     = init_fastdiv_values(ne12);
 }
 
+// The rows of dims 1..3 can be walked with one stride (size-1 dims skipped); returns that stride
+static bool ggml_hexagon_rows_stride(const int64_t * ne, const size_t * nb, size_t * stride) {
+    size_t s = 0, next = 0;
+    for (int i = 1; i < GGML_MAX_DIMS; i++) {
+        if (ne[i] == 1) continue;
+        if (s == 0) { s = nb[i]; next = s * ne[i]; continue; }
+        if (nb[i] != next) return false;
+        next *= ne[i];
+    }
+    *stride = s ? s : nb[1];
+    return true;
+}
+
+// A 2D weight applied to a batched activation whose rows are evenly strided is the same matmul over ne11 * ne12 * ne13 rows
+static bool ggml_hexagon_matmul_can_collapse(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
+    size_t s1, sd;
+    return (dst->op == GGML_OP_MUL_MAT || dst->op == GGML_OP_ADD) &&
+           src0->ne[2] == 1 && src0->ne[3] == 1 && src1->ne[2] * src1->ne[3] > 1 &&
+           src1->nb[0] == ggml_type_size(src1->type) && ggml_hexagon_rows_stride(src1->ne, src1->nb, &s1) &&
+           dst->nb[0] == ggml_type_size(dst->type) && ggml_hexagon_rows_stride(dst->ne, dst->nb, &sd);
+}
+
+static bool ggml_hexagon_matmul_add_can_collapse(
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * src2,
+    const struct ggml_tensor * dst
+) {
+    if (!ggml_hexagon_matmul_can_collapse(src0, src1, dst)) {
+        return false;
+    }
+    if (!src2) {
+        return true;
+    }
+    const int64_t src2_nrows = src2->ne[1] * src2->ne[2] * src2->ne[3];
+    if (src2_nrows == 1) {
+        return true;
+    }
+    size_t s2;
+    return src2->nb[0] == ggml_type_size(src2->type) &&
+           src2->ne[0] == dst->ne[0] &&
+           src2->ne[1] == dst->ne[1] &&
+           src2->ne[2] == dst->ne[2] &&
+           src2->ne[3] == dst->ne[3] &&
+           ggml_hexagon_rows_stride(src2->ne, src2->nb, &s2);
+}
+
+static ggml_tensor ggml_hexagon_tensor_collapse_rows(const struct ggml_tensor * t) {
+    size_t stride = 0;
+    ggml_hexagon_rows_stride(t->ne, t->nb, &stride);
+    ggml_tensor c = *t;
+    c.ne[1] = t->ne[1] * t->ne[2] * t->ne[3];
+    c.ne[2] = 1;
+    c.ne[3] = 1;
+    c.nb[1] = stride;
+    c.nb[2] = c.nb[1] * c.ne[1];
+    c.nb[3] = c.nb[2];
+    return c;
+}
+
 static void ggml_hexagon_precompute_matmul_params(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * src0,
@@ -5436,6 +5619,13 @@ static void ggml_hexagon_precompute_matmul_params(
     const struct ggml_tensor * dst,
     struct htp_mm_kernel_params * kparams
 ) {
+    if (ggml_hexagon_matmul_can_collapse(src0, src1, dst)) {
+        const ggml_tensor src1_collapsed = ggml_hexagon_tensor_collapse_rows(src1);
+        const ggml_tensor dst_collapsed  = ggml_hexagon_tensor_collapse_rows(dst);
+        ggml_hexagon_precompute_matmul_params_impl(sess, src0, &src1_collapsed, &dst_collapsed, 0, 0, kparams);
+        kparams->collapse = 1;
+        return;
+    }
     ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, 0, kparams);
 }
 
@@ -5447,6 +5637,18 @@ static void ggml_hexagon_precompute_fused_matmul_add_params(
     const struct ggml_tensor * dst,
     struct htp_mm_kernel_params * kparams
 ) {
+    if (ggml_hexagon_matmul_add_can_collapse(src0, src1, src2, dst)) {
+        const ggml_tensor src1_collapsed = ggml_hexagon_tensor_collapse_rows(src1);
+        const ggml_tensor dst_collapsed  = ggml_hexagon_tensor_collapse_rows(dst);
+        const ggml_tensor src2_collapsed = (src2 && (src2->ne[1] * src2->ne[2] * src2->ne[3] > 1))
+                                           ? ggml_hexagon_tensor_collapse_rows(src2)
+                                           : (src2 ? *src2 : ggml_tensor{});
+        const struct ggml_tensor * p_src2 = src2 ? &src2_collapsed : nullptr;
+        const size_t src2_size = p_src2 ? hex_round_up(ggml_nbytes(p_src2), 128) : 0;
+        ggml_hexagon_precompute_matmul_params_impl(sess, src0, &src1_collapsed, &dst_collapsed, p_src2 ? p_src2->nb[1] : 0, src2_size, kparams);
+        kparams->collapse = 1;
+        return;
+    }
     const size_t src2_size = src2 ? hex_round_up(ggml_nbytes(src2), 128) : 0;
     ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, src2 ? src2->nb[1] : 0, src2_size, kparams);
 }
@@ -5878,68 +6080,63 @@ static void ggml_hexagon_precompute_ssm_conv_params(
 
     const uint32_t raw_rpt = (d_inner + n_threads - 1) / n_threads;
     const uint32_t d_inner_per_thread = hex_round_up(raw_rpt, 32);
-    kparams->d_inner_per_thread = d_inner_per_thread;
 
-    kparams->src0_row_size_aligned = hex_round_up(ncs * sizeof(float), 128);
-    kparams->src1_row_size_aligned = hex_round_up(d_conv * sizeof(float), 128);
-    kparams->dst_row_size_aligned  = hex_round_up(d_inner * sizeof(float), 128);
+    const uint32_t src1_raw_bytes = hex_round_up(d_inner_per_thread * d_conv * sizeof(float), 128) + 128;
+    const uint32_t src1_T_bytes   = hex_round_up(d_conv * d_inner_per_thread * sizeof(float), 128);
+    const uint32_t vtcm_src1_per_thread = src1_raw_bytes + src1_T_bytes;
+
+    uint32_t vtcm_src0_per_thread = 0;
+    uint32_t vtcm_dst_per_thread  = 0;
 
     if (n_t == 1) {
         kparams->d_inner_tile = d_inner_per_thread;
 
-        const uint32_t src1_raw_bytes = hex_round_up(d_inner_per_thread * d_conv * sizeof(float), 128) + 128;
-        const uint32_t src1_T_bytes   = hex_round_up(d_conv * d_inner_per_thread * sizeof(float), 128);
-        const uint32_t vtcm_src1_per_thread = src1_raw_bytes + src1_T_bytes;
+        const uint32_t src0_tile_raw_bytes = hex_round_up(d_inner_per_thread * d_conv * sizeof(float), 128);
+        const uint32_t src0_T_bytes        = hex_round_up(d_conv * d_inner_per_thread * sizeof(float), 128);
+        vtcm_src0_per_thread = 2 * src0_tile_raw_bytes + src0_T_bytes;
 
-        const uint32_t src0_raw_bytes = hex_round_up(d_inner_per_thread * d_conv * sizeof(float), 128) + 128;
-        const uint32_t src0_T_bytes   = hex_round_up(d_conv * d_inner_per_thread * sizeof(float), 128);
-        const uint32_t vtcm_src0_per_thread = src0_raw_bytes + src0_T_bytes;
-
-        const uint32_t vtcm_dst_per_thread = hex_round_up(d_inner_per_thread * sizeof(float), 128);
-
-        kparams->vtcm_src0_size_per_thread = vtcm_src0_per_thread;
-        kparams->vtcm_src1_size_per_thread = vtcm_src1_per_thread;
-        kparams->vtcm_dst_size_per_thread  = vtcm_dst_per_thread;
-
-        kparams->vtcm_src0_size = vtcm_src0_per_thread * n_threads;
-        kparams->vtcm_src1_size = vtcm_src1_per_thread * n_threads;
-        kparams->vtcm_dst_size  = vtcm_dst_per_thread  * n_threads;
-        kparams->vtcm_size      = kparams->vtcm_src0_size + kparams->vtcm_src1_size + kparams->vtcm_dst_size;
+        const uint32_t dst_tile_bytes = hex_round_up(d_inner_per_thread * sizeof(float), 128);
+        vtcm_dst_per_thread = 2 * dst_tile_bytes;
     } else {
-        const uint32_t src1_raw_bytes = hex_round_up(d_inner_per_thread * d_conv * sizeof(float), 128) + 128;
-        const uint32_t src1_T_bytes   = hex_round_up(d_conv * d_inner_per_thread * sizeof(float), 128);
-        const uint32_t vtcm_src1_per_thread = src1_raw_bytes + src1_T_bytes;
-
         const size_t vtcm_budget = (sess->vtcm_size > 0 ? sess->vtcm_size / n_threads : (1024 * 1024));
-        const size_t avail_for_src0 = vtcm_budget > vtcm_src1_per_thread ? vtcm_budget - vtcm_src1_per_thread : (128 * 1024);
 
-        uint32_t d_inner_tile = (uint32_t)((avail_for_src0 / 2) / (ncs * sizeof(float) + n_t * sizeof(float) + 1));
+        // the kernel double-buffers the raw src0 tile and the dst tile, and transposes
+        // one 32-channel block at a time
+        const uint32_t src0_block_T    = hex_round_up(ncs * 32 * sizeof(float), 128);
+        const size_t   fixed_bytes     = vtcm_src1_per_thread + src0_block_T;
+        const size_t   avail_for_tiles = vtcm_budget > fixed_bytes ? vtcm_budget - fixed_bytes : (128 * 1024);
+
+        uint32_t target_max_tile = hex_round_up((d_inner_per_thread + 3) / 4, 32);
+        target_max_tile = (std::max)(target_max_tile, 32u);
+        target_max_tile = (std::min)(target_max_tile, 128u);
+
+        uint32_t d_inner_tile = (uint32_t)(avail_for_tiles / (2 * (ncs + n_t) * sizeof(float)));
         d_inner_tile = (d_inner_tile / 32) * 32;
         if (d_inner_tile == 0) {
             d_inner_tile = 32;
         }
+        if (d_inner_tile > target_max_tile) {
+            d_inner_tile = target_max_tile;
+        }
         if (d_inner_tile > d_inner_per_thread) {
             d_inner_tile = d_inner_per_thread;
         }
         kparams->d_inner_tile = d_inner_tile;
 
-        const uint32_t src0_tile_raw = hex_round_up(d_inner_tile * ncs * sizeof(float), 128) + 128;
-        const uint32_t src0_tile_T   = hex_round_up(ncs * d_inner_tile * sizeof(float), 128);
-        const uint32_t vtcm_src0_per_thread = src0_tile_raw + src0_tile_T;
+        const uint32_t src0_tile_raw = hex_round_up(d_inner_tile * ncs * sizeof(float), 128);
+        vtcm_src0_per_thread = 2 * src0_tile_raw + src0_block_T;
 
-        const uint32_t vtcm_dst_per_thread = hex_round_up(d_inner_tile * n_t * sizeof(float), 128);
-
-        kparams->vtcm_src0_size_per_thread = vtcm_src0_per_thread;
-        kparams->vtcm_src1_size_per_thread = vtcm_src1_per_thread;
-        kparams->vtcm_dst_size_per_thread  = vtcm_dst_per_thread;
-
-        kparams->vtcm_src0_size = vtcm_src0_per_thread * n_threads;
-        kparams->vtcm_src1_size = vtcm_src1_per_thread * n_threads;
-        kparams->vtcm_dst_size  = vtcm_dst_per_thread  * n_threads;
-        kparams->vtcm_size      = kparams->vtcm_src0_size + kparams->vtcm_src1_size + kparams->vtcm_dst_size;
+        vtcm_dst_per_thread = 2 * hex_round_up(d_inner_tile * n_t * sizeof(float), 128);
     }
 
-    kparams->div_n_threads = init_fastdiv_values(n_threads);
+    kparams->vtcm_src0_size_per_thread = vtcm_src0_per_thread;
+    kparams->vtcm_src1_size_per_thread = vtcm_src1_per_thread;
+    kparams->vtcm_dst_size_per_thread  = vtcm_dst_per_thread;
+
+    kparams->vtcm_src0_size = vtcm_src0_per_thread * n_threads;
+    kparams->vtcm_src1_size = vtcm_src1_per_thread * n_threads;
+    kparams->vtcm_dst_size  = vtcm_dst_per_thread  * n_threads;
+    kparams->vtcm_size      = kparams->vtcm_src0_size + kparams->vtcm_src1_size + kparams->vtcm_dst_size;
 }
 
 static void ggml_hexagon_precompute_gated_delta_net_params(
@@ -6063,10 +6260,100 @@ static void ggml_hexagon_precompute_sort_params(
     kparams->n_slots           = (int32_t) layout.n_slots;
 }
 
+static void ggml_hexagon_pool_interior_range(
+    uint32_t src_x, uint32_t dst_x, uint32_t kernel_x, uint32_t stride_x, int32_t pad_x,
+    uint32_t * ox_lo, uint32_t * ox_hi) {
+    const uint32_t lo_raw = ((uint32_t) pad_x + stride_x - 1) / stride_x;
+    *ox_lo = (lo_raw < dst_x) ? lo_raw : dst_x;
+
+    const int32_t numer_hi = (int32_t) src_x - (int32_t) kernel_x + pad_x;
+    if (numer_hi < 0) {
+        *ox_hi = 0;
+    } else {
+        const uint32_t hi_raw = (uint32_t) numer_hi / stride_x + 1;
+        *ox_hi = (hi_raw < dst_x) ? hi_raw : dst_x;
+    }
+    if (*ox_hi < *ox_lo) {
+        *ox_hi = *ox_lo;
+    }
+}
+
+static bool ggml_hexagon_pool_shape_fits(
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * dst
+) {
+    return ggml_nbytes(src0) <= UINT32_MAX - 256 && ggml_nbytes(dst) <= UINT32_MAX - 256;
+}
+
+static void ggml_hexagon_precompute_pool_2d_params(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * dst,
+    struct htp_pool_2d_kernel_params * kparams,
+    bool is_pool_1d
+) {
+    memset(kparams, 0, sizeof(*kparams));
+
+    const uint32_t src_x = (uint32_t) src0->ne[0];
+    const uint32_t src_y = is_pool_1d ? 1 : (uint32_t) src0->ne[1];
+    const uint32_t dst_x = (uint32_t) dst->ne[0];
+    const uint32_t dst_y = is_pool_1d ? 1 : (uint32_t) dst->ne[1];
+    const uint32_t src_plane_bytes = src_x * src_y * sizeof(float);
+    const uint32_t dst_plane_bytes = dst_x * dst_y * sizeof(float);
+    const uint32_t planes = (uint32_t) (is_pool_1d
+        ? (uint64_t) src0->ne[1] * (uint64_t) src0->ne[2] * (uint64_t) src0->ne[3]
+        : (uint64_t) src0->ne[2] * (uint64_t) src0->ne[3]);
+    const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, planes > 0 ? planes : 1);
+
+    struct htp_pool_vtcm_layout layout;
+    const bool ok = htp_pool_solve_layout(&layout, src_x, src_y, dst_x, dst_y, n_threads, sess->vtcm_size);
+    GGML_ASSERT(ok);
+
+    kparams->src_x = src_x;
+    kparams->src_y = src_y;
+    kparams->dst_x = dst_x;
+    kparams->dst_y = dst_y;
+    kparams->kernel_x = (uint32_t) ggml_get_op_params_i32(dst, 1);
+    kparams->kernel_y = is_pool_1d ? 1 : (uint32_t) ggml_get_op_params_i32(dst, 2);
+    kparams->stride_x = (uint32_t) ggml_get_op_params_i32(dst, is_pool_1d ? 2 : 3);
+    kparams->stride_y = is_pool_1d ? 1 : (uint32_t) ggml_get_op_params_i32(dst, 4);
+    kparams->pad_x = ggml_get_op_params_i32(dst, is_pool_1d ? 3 : 5);
+    kparams->pad_y = is_pool_1d ? 0 : ggml_get_op_params_i32(dst, 6);
+    kparams->src_plane_bytes = src_plane_bytes;
+    kparams->dst_plane_bytes = dst_plane_bytes;
+    kparams->src_plane_bytes_aligned = (uint32_t) layout.src_spad_half_size;
+    kparams->dst_plane_bytes_aligned = (uint32_t) layout.dst_spad_half_size;
+    kparams->n_threads = n_threads;
+    kparams->planes = planes;
+    kparams->pool_op = (uint32_t) ggml_get_op_params_i32(dst, 0);
+    // Fast HVX path requires exact tiling and a supported kernel width.
+    kparams->fast_path = (kparams->pad_x == 0 && kparams->pad_y == 0 &&
+                           kparams->stride_x == kparams->kernel_x && kparams->stride_y == kparams->kernel_y &&
+                           (kparams->kernel_x == 1 || kparams->kernel_x == 2)) ? 1 : 0;
+    kparams->global_path = (kparams->pad_x == 0 && kparams->pad_y == 0 &&
+                            kparams->stride_x == kparams->kernel_x &&
+                            kparams->stride_y == kparams->kernel_y &&
+                            kparams->kernel_x == kparams->src_x &&
+                            kparams->kernel_y == kparams->src_y &&
+                            kparams->dst_x == 1 && kparams->dst_y == 1) ? 1 : 0;
+    kparams->block_path = (kparams->pad_x == 0 && kparams->pad_y == 0 &&
+                           kparams->kernel_y == 1 && kparams->stride_y == 1 &&
+                           kparams->stride_x == kparams->kernel_x &&
+                           kparams->kernel_x >= 32) ? 1 : 0;
+    kparams->avg_divide_count = (is_pool_1d && kparams->pool_op == GGML_OP_POOL_AVG) ? 1 : 0;
+    ggml_hexagon_pool_interior_range(kparams->src_x, kparams->dst_x, kparams->kernel_x,
+                                     kparams->stride_x, kparams->pad_x,
+                                     &kparams->ox_lo, &kparams->ox_hi);
+
+    const bool narrow_ok = (uint64_t) kparams->dst_x * kparams->kernel_x <= 32;
+    kparams->narrow_path = (kparams->fast_path && narrow_ok) ? 1 : 0;
+    kparams->inv_kernel_area = 1.0f / (float) (kparams->kernel_x * kparams->kernel_y);
+}
+
 static void ggml_hexagon_precompute_fused_mmnx_params(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * src0, // W0
-    const struct ggml_tensor * src1, // x
+    const struct ggml_tensor * act,  // x
     int32_t n_weights,
     struct htp_mm_kernel_params * kparams
 ) {
@@ -6078,23 +6365,24 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
     const int ne02 = src0->ne[2];
     const int ne03 = src0->ne[3];
 
-    const int ne10 = src1->ne[0];
-    const int ne11 = src1->ne[1];
-    const int ne12 = src1->ne[2];
-    const int ne13 = src1->ne[3];
+    const int ne10 = act->ne[0];
+    const int ne11 = act->ne[1];
+    const int ne12 = act->ne[2];
+    const int ne13 = act->ne[3];
 
     const int wtype = src0->type;
     const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype);
     const int ne00_padded = is_repack ? hex_round_up(ne00, 32) : ne00;
     const int ne01_padded = is_repack ? hex_round_up(ne01, 32) : ne01;
     const int ne11_padded = hex_round_up(ne11, 32);
+    const int ne01_tiled  = hex_round_up(ne01_padded, 32);
 
     const size_t vtcm_budget = sess->vtcm_size;
     const bool is_batched = (ne02 * ne03 > 1 || ne12 * ne13 > 1);
 
     bool hmx_enabled = (sess->n_hmx > 0) && (opt_mm_select >= 2);
-    if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, src1, nullptr, ne01_padded, false, is_batched)) {
-        if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, nullptr, wtype, ne00_padded, ne01_padded, ne02, ne11, ne12, ne11_padded, false, is_batched, 0, vtcm_budget, kparams)) {
+    if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, act, nullptr, ne01_padded, false, is_batched)) {
+        if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, act, nullptr, wtype, ne00_padded, ne01_tiled, ne02, ne11, ne12, ne11_padded, false, is_batched, 0, vtcm_budget, kparams)) {
             kparams->n_weights = n_weights;
             goto finalize;
         }
@@ -6106,20 +6394,20 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
     }
 
     {
-        const int src1_nrows = ne11 * ne12 * ne13;
-        const size_t src1_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
+        const int act_nrows = ne11 * ne12 * ne13;
+        const size_t act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
         const size_t src0_row_size = src0->nb[1];
 
         uint32_t best_n_prefetch = 16;
 
         if (is_repack) {
-            const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
+            const uint32_t max_prefetch = (act_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
             best_n_prefetch = 2;
             for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
                 struct htp_mm_hvx_vtcm_layout L;
                 htp_mm_hvx_vtcm_layout_build(
-                    &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
-                    0, src0_row_size, src1_row_size, 0, d, false, true
+                    &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, act_nrows, sess->n_threads,
+                    0, src0_row_size, act_row_size, 0, d, false, true
                 );
                 if (L.total_bytes <= sess->vtcm_size) {
                     best_n_prefetch = d;
@@ -6133,14 +6421,16 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
 
         // Test tiled first
         htp_mm_hvx_vtcm_layout_build(
-            &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
-            0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true
+            &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, act_nrows, sess->n_threads,
+            0, src0_row_size, act_row_size, 0, best_n_prefetch, false, true
         );
 
         if (try_tiled && L.total_bytes <= sess->vtcm_size) {
-            kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
+            kparams->kernel_type    = HTP_MM_KERNEL_HVX_QUANT_ROW;
+            kparams->act_row_size   = act_row_size;
             kparams->vtcm_src0_size = L.src0_bytes;
-            kparams->vtcm_src1_size = L.src1_bytes;
+            kparams->vtcm_act_size  = L.act_bytes;
+            kparams->vtcm_bias_size = 0;
             kparams->vtcm_dst_size  = L.dst_bytes;
             kparams->vtcm_size      = L.total_bytes;
             kparams->n_prefetch     = best_n_prefetch;
@@ -6162,12 +6452,12 @@ finalize:
 static void ggml_hexagon_precompute_fused_mmidnx_params(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * src0, // W0
-    const struct ggml_tensor * src1, // x
+    const struct ggml_tensor * act,  // x
     const struct ggml_tensor * dst,  // dst0
     int32_t n_weights,
     struct htp_mm_kernel_params * kparams
 ) {
-    ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, 0, kparams);
+    ggml_hexagon_precompute_matmul_params_impl(sess, src0, act, dst, 0, 0, kparams);
     kparams->n_weights = n_weights;
 }
 
@@ -6505,6 +6795,76 @@ static bool ggml_hexagon_supported_argmax(const struct ggml_hexagon_session * se
     GGML_UNUSED(sess);
 }
 
+static bool ggml_hexagon_supported_pool_2d(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
+    const struct ggml_tensor * src0 = op->src[0];
+    const int32_t * params = op->op_params;
+
+    if (params[0] != GGML_OP_POOL_AVG && params[0] != GGML_OP_POOL_MAX) {
+        return false;
+    }
+    if (src0->type != GGML_TYPE_F32 || op->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(op)) {
+        return false;
+    }
+    if (!ggml_hexagon_pool_shape_fits(src0, op)) {
+        return false;
+    }
+
+    const int32_t kernel_x = params[1];
+    const int32_t kernel_y = params[2];
+    const int32_t stride_x = params[3];
+    const int32_t stride_y = params[4];
+    const int32_t pad_x    = params[5];
+    const int32_t pad_y    = params[6];
+
+    // Keep invalid parameters on CPU, avoiding unsafe unsigned values in HTP.
+    if (kernel_x <= 0 || kernel_y <= 0 || stride_x <= 0 || stride_y <= 0 ||
+        pad_x < 0 || pad_y < 0) {
+        return false;
+    }
+
+    const uint32_t planes = (uint32_t) (src0->ne[2] * src0->ne[3]);
+    const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, planes > 0 ? planes : 1);
+
+    struct htp_pool_vtcm_layout layout;
+    return htp_pool_solve_layout(&layout, (uint32_t) src0->ne[0], (uint32_t) src0->ne[1],
+                                 (uint32_t) op->ne[0], (uint32_t) op->ne[1],
+                                 n_threads, sess->vtcm_size);
+}
+
+static bool ggml_hexagon_supported_pool_1d(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
+    const struct ggml_tensor * src0 = op->src[0];
+    const int32_t * params = op->op_params;
+
+    if (params[0] != GGML_OP_POOL_AVG && params[0] != GGML_OP_POOL_MAX) {
+        return false;
+    }
+    if (src0->type != GGML_TYPE_F32 || op->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(op)) {
+        return false;
+    }
+    if (!ggml_hexagon_pool_shape_fits(src0, op)) {
+        return false;
+    }
+
+    const int32_t kernel = params[1];
+    const int32_t stride = params[2];
+    const int32_t pad    = params[3];
+    if (kernel <= 0 || stride <= 0 || pad < 0) {
+        return false;
+    }
+
+    const uint32_t planes = (uint32_t) (src0->ne[1] * src0->ne[2] * src0->ne[3]);
+    const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, planes > 0 ? planes : 1);
+
+    struct htp_pool_vtcm_layout layout;
+    return htp_pool_solve_layout(&layout, (uint32_t) src0->ne[0], 1, (uint32_t) op->ne[0], 1, n_threads, sess->vtcm_size);
+}
+
 static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
     const struct ggml_tensor * src0 = op->src[0];
     const struct ggml_tensor * src1 = op->src[1];
@@ -7062,6 +7422,8 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
         case GGML_OP_SET_ROWS:        return HTP_OP_SET_ROWS;
         case GGML_OP_SUM:             return HTP_OP_SUM;
         case GGML_OP_SUM_ROWS:        return HTP_OP_SUM_ROWS;
+        case GGML_OP_POOL_2D:         return HTP_OP_POOL_2D;
+        case GGML_OP_POOL_1D:         return HTP_OP_POOL_1D;
         case GGML_OP_ARGSORT:         return HTP_OP_ARGSORT;
         case GGML_OP_TOP_K:           return HTP_OP_TOP_K;
         case GGML_OP_ARGMAX:          return HTP_OP_ARGMAX;
@@ -7137,6 +7499,14 @@ static bool mm_is_hmx_eligible(const ggml_tensor * t) {
     const ggml_tensor * src0 = t->src[0];
     const ggml_tensor * src1 = t->src[1];
 
+    if (ggml_hexagon_matmul_can_collapse(src0, src1, t)) {
+        const ggml_tensor src1_c = ggml_hexagon_tensor_collapse_rows(src1);
+        const int wtype = src0->type;
+        const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype);
+        const int ne01_padded = is_repack ? hex_round_up(src0->ne[1], 32) : src0->ne[1];
+        return ggml_hexagon_matmul_is_hmx_eligible(src0, &src1_c, t, ne01_padded, false, false);
+    }
+
     const int wtype = src0->type;
     const bool is_repack    = ggml_hexagon_is_repack_type((ggml_type) wtype);
     const bool is_matmul_id = (t->op == GGML_OP_MUL_MAT_ID);
@@ -7176,13 +7546,15 @@ static bool is_mergeable_mul_mat(const ggml_tensor * t) {
 
     const ggml_tensor * src0 = t->src[0];
     const ggml_tensor * src1 = t->src[1];
-    if (src1->type != GGML_TYPE_F32) return false;
     if (src0->ne[2] != 1 || src0->ne[3] != 1) return false;
 
     if (mm_is_hmx_eligible(t)) {
         return ggml_hexagon_is_hmx_weight_type(src0->type);
     }
 
+    // HVX path requires F32 activations and repacked weights (except Q6_K)
+    if (src1->type != GGML_TYPE_F32) return false;
+
     return ggml_hexagon_is_repack_type(src0->type) && src0->type != GGML_TYPE_Q6_K;
 }
 
@@ -7340,6 +7712,12 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg
                     node.opcode == HTP_OP_TOP_K,
                     (struct htp_sort_kernel_params *) node.kernel_params
                 );
+            } else if (node.opcode == HTP_OP_POOL_2D || node.opcode == HTP_OP_POOL_1D) {
+                ggml_hexagon_precompute_pool_2d_params(
+                    sess, node.node->src[0], node.dst(),
+                    (struct htp_pool_2d_kernel_params *)node.kernel_params,
+                    node.opcode == HTP_OP_POOL_1D
+                );
             }
             computed_nodes.push_back(std::move(node));
         }
@@ -8157,6 +8535,14 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
             supp = ggml_hexagon_supported_argmax(sess, op);
             break;
 
+        case GGML_OP_POOL_2D:
+            supp = ggml_hexagon_supported_pool_2d(sess, op);
+            break;
+
+        case GGML_OP_POOL_1D:
+            supp = ggml_hexagon_supported_pool_1d(sess, op);
+            break;
+
         case GGML_OP_SOFT_MAX:
             supp = ggml_hexagon_supported_softmax(sess, op);
             break;
@@ -8683,6 +9069,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
     const char * str_nhmx     = getenv("GGML_HEXAGON_NHMX");
     const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT");
     const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT");
+    const char * str_fa_head_split = getenv("GGML_HEXAGON_FA_HEAD_SPLIT");
     const char * str_gdn_select = getenv("GGML_HEXAGON_GDN_SELECT");
     const char * str_ar_select = getenv("GGML_HEXAGON_AR_SELECT");
     const char * str_ar_scatter = getenv("GGML_HEXAGON_AR_SCATTER");
@@ -8737,6 +9124,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
     opt_nhmx      = str_nhmx     ? atoi(str_nhmx)                         : opt_nhmx;
     opt_mm_select = str_mm_select ? atoi(str_mm_select)                   : opt_mm_select;
     opt_fa_select = str_fa_select ? atoi(str_fa_select)                   : opt_fa_select;
+    opt_fa_head_split = str_fa_head_split ? atoi(str_fa_head_split)        : opt_fa_head_split;
     opt_gdn_select = str_gdn_select ? atoi(str_gdn_select)                 : opt_gdn_select;
     opt_ar_select = str_ar_select ? atoi(str_ar_select)                   : opt_ar_select;
     opt_ar_scatter = str_ar_scatter ? atoi(str_ar_scatter)                : opt_ar_scatter;

ggml/src/ggml-hexagon/htp/CMakeLists.txt

@@ -44,6 +44,7 @@ add_library(${HTP_LIB} SHARED
     argsort-ops.c
     im2col-ops.c
     roll-ops.c
+    pool-ops.c
     allreduce-ops.c
 )
 

ggml/src/ggml-hexagon/htp/dma-queue.c

@@ -161,9 +161,9 @@ bool dma_queue_push_fallback_contig(dma_queue * q, dma_data ddata, size_t total)
     while (rem_bytes > 0) {
         const uint32_t cur_bytes = MIN(rem_bytes, DMA_SAFE_CHUNK_SIZE);
         dma_data cur_data = dma_make_data(cur_dst, cur_src);
-        if (!dma_ring_push_single_1d(r1, cur_data, cur_bytes)) {
+        if (!dma_ring_push_single_contig(r1, cur_data, cur_bytes)) {
             dma_ring_flush(r1);
-            dma_ring_push_single_1d(r1, cur_data, cur_bytes);
+            dma_ring_push_single_contig(r1, cur_data, cur_bytes);
         }
         cur_dst   += cur_bytes;
         cur_src   += cur_bytes;

ggml/src/ggml-hexagon/htp/dma-queue.h

@@ -107,6 +107,14 @@ typedef struct {
 #define DMA_MAX_STRIDE_24B     0x00FFFFFFu    // 24-bit HW descriptor limit for strides (16MB - 1)
 #define DMA_SAFE_CHUNK_SIZE    0x00F00000u    // ~15MB safe contiguous chunk size
 
+#if __HVX_ARCH__ < 75
+#define DMA_MAX_2D_ROW_SIZE    DMA_MAX_SIZE_16B
+#define DMA_MAX_2D_STRIDE      DMA_MAX_STRIDE_16B
+#else
+#define DMA_MAX_2D_ROW_SIZE    DMA_MAX_SIZE_24B
+#define DMA_MAX_2D_STRIDE      DMA_MAX_STRIDE_24B
+#endif
+
 #define DMA_FALLBACK_CAPACITY  16u            // descriptors in secondary fallback ring
 
 typedef struct dma_ring_s dma_ring;
@@ -216,13 +224,10 @@ static inline bool dma_ring_push_single_1d(dma_ring * r, dma_data ddata, size_t
 
 static inline bool dma_ring_push_single_2d(dma_ring * r, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
 #if __HVX_ARCH__ > 79
+    assert(!((ddata.src | ddata.dst) >> 40) || nrows == 0);
     const uint32_t src_hi = (uint32_t) (ddata.src >> 32);
     const uint32_t dst_hi = (uint32_t) (ddata.dst >> 32);
     const bool is_ext     = (src_hi | dst_hi) != 0;
-
-    if (is_ext && ((ddata.src >> 40) || (ddata.dst >> 40))) {
-        return false;
-    }
 #endif
 
     if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
@@ -284,6 +289,16 @@ static inline bool dma_ring_push_single_2d(dma_ring * r, dma_data ddata, size_t
     return true;
 }
 
+#if __HVX_ARCH__ < 75
+static inline bool dma_ring_push_single_contig(dma_ring * r, dma_data ddata, size_t size) {
+    return dma_ring_push_single_1d(r, ddata, size);
+}
+#else
+static inline bool dma_ring_push_single_contig(dma_ring * r, dma_data ddata, size_t size) {
+    return dma_ring_push_single_2d(r, ddata, size, size, size, 1);
+}
+#endif
+
 static inline dma_data dma_ring_pop(dma_ring * r) {
     dma_data ddata = { 0 };
 
@@ -374,57 +389,36 @@ static inline uint32_t dma_queue_capacity(dma_queue * q) {
     return dma_ring_capacity(q->ring0);
 }
 
-#if __HVX_ARCH__ < 75
-
-static inline bool dma_queue_push(dma_queue *q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
-    // Fast path: everything fits in 16 bits
-    if (nrows == 0 || __builtin_expect(
-            nrows      <= DMA_MAX_NROWS &&
-            row_size   <= DMA_MAX_SIZE_16B &&
-            src_stride <= DMA_MAX_STRIDE_16B &&
-            dst_stride <= DMA_MAX_STRIDE_16B, 1)) {
-        return dma_ring_push_single_2d(q->ring0, ddata, dst_stride, src_stride, row_size, nrows);
+static inline bool dma_queue_push(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
+    if (__builtin_expect(nrows == 0, 0)) {
+        return dma_ring_push_single_1d(q->ring0, ddata, 0);
     }
 
-    // Contiguous block: 1D DMA mode supports up to 24-bit size (16MB)
-    if (nrows == 1 || (row_size == src_stride && row_size == dst_stride)) {
-        size_t total = row_size * nrows;
-        if (total <= DMA_MAX_SIZE_24B) {
-            return dma_ring_push_single_1d(q->ring0, ddata, total);
+    // 1. Hot path: Contiguous or single-row (80-90% of calls)
+    if (nrows == 1 || !((row_size ^ src_stride) | (row_size ^ dst_stride))) {
+        const size_t total = row_size * nrows;
+        if (__builtin_expect(total <= DMA_MAX_SIZE_24B, 1)) {
+            return dma_ring_push_single_contig(q->ring0, ddata, total);
         }
         return dma_queue_push_fallback_contig(q, ddata, total);
     }
 
-    // Row count overflow with 16-bit strides: chunk 2D descriptors via fallback ring
-    if (row_size <= DMA_MAX_SIZE_16B && src_stride <= DMA_MAX_STRIDE_16B && dst_stride <= DMA_MAX_STRIDE_16B) {
-        return dma_queue_push_fallback_2d(q, ddata, dst_stride, src_stride, row_size, nrows);
-    }
-
-    // Stride or row_size overflow: row-by-row 1D via fallback ring
-    return dma_queue_push_fallback_1d(q, ddata, dst_stride, src_stride, row_size, nrows);
-}
-
-#else // HVX_ARCH >= 75
-
-static inline bool dma_queue_push(dma_queue *q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
-    if (nrows == 0 || __builtin_expect(
-            nrows      <= DMA_MAX_NROWS &&
-            row_size   <= DMA_MAX_SIZE_24B &&
-            src_stride <= DMA_MAX_STRIDE_24B &&
-            dst_stride <= DMA_MAX_STRIDE_24B, 1)) {
+    // 2. Hot path: Standard strided 2D (10-20% of calls)
+    if (__builtin_expect(nrows <= DMA_MAX_NROWS &&
+                         (row_size | src_stride | dst_stride) <= DMA_MAX_2D_ROW_SIZE, 1)) {
         return dma_ring_push_single_2d(q->ring0, ddata, dst_stride, src_stride, row_size, nrows);
     }
 
-    // Contiguous block exceeding 24 bits
-    if (nrows == 1 || (row_size == src_stride && row_size == dst_stride)) {
-        size_t total = row_size * nrows;
-        return dma_queue_push_fallback_contig(q, ddata, total);
+    // 3. Cold path: Descriptor chunking fallbacks (< 0.1%)
+#if __HVX_ARCH__ < 75
+    if (row_size <= DMA_MAX_SIZE_16B && (src_stride | dst_stride) <= DMA_MAX_STRIDE_16B) {
+        return dma_queue_push_fallback_2d(q, ddata, dst_stride, src_stride, row_size, nrows);
     }
-
+    return dma_queue_push_fallback_1d(q, ddata, dst_stride, src_stride, row_size, nrows);
+#else
     return dma_queue_push_fallback_2d(q, ddata, dst_stride, src_stride, row_size, nrows);
-}
-
 #endif
+}
 
 static inline void dma_sync_read(dma_queue * dma_q, void * dst, dma_addr_t src, size_t bytes) {
     const uint32_t b = (uint32_t) bytes;

ggml/src/ggml-hexagon/htp/flash-attn-ops.c

@@ -1869,8 +1869,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     factx.Bc             = kparams->Bc;
     factx.g_br           = kparams->u.hmx.g_br;
     factx.n_kv_blocks    = kparams->n_kv_blocks;
-    factx.is_q_fp32      = (kparams->is_q_fp32 != 0);
-    factx.is_dst_fp32    = (kparams->is_dst_fp32 != 0);
+    factx.is_q_fp32      = (q->type == HTP_TYPE_F32);
+    factx.is_dst_fp32    = (dst->type == HTP_TYPE_F32);
     factx.pipeline       = (kparams->u.hmx.pipeline != 0);
     factx.mask_broadcast = (kparams->u.hmx.mask_broadcast != 0);
     if (mask) {
@@ -1879,13 +1879,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     }
 
     factx.has_softcap   = (kparams->logit_softcap != 0.0f);
-    if (!factx.has_softcap) {
-        factx.scale = (__fp16) (kparams->scale * EXP_LOG2E_F);  // log2(e)
-    } else {
-        factx.scale = (__fp16) kparams->scale;
-    }
+    factx.scale         = (__fp16) kparams->scale;
     factx.max_bias      = kparams->max_bias;
-    factx.logit_softcap = factx.has_softcap ? (__fp16) (kparams->logit_softcap * EXP_LOG2E_F) : 0;
+    factx.logit_softcap = 0;
+    if (factx.has_softcap) {
+        factx.logit_softcap = (__fp16) kparams->logit_softcap;
+    }
 
     factx.n_head_log2 = kparams->n_head_log2;
     factx.m0          = kparams->m0;
@@ -1898,22 +1897,36 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     const uint32_t n_threads = factx.n_threads;
     const uint32_t G = factx.G;
 
-    // Multi-device: split Q blocks across devices
+    // Multi-device: prefer head-parallel partitioning (each core owns a disjoint head
+    // shard), falling back to Q-block (token) split when heads don't divide evenly.
     const uint32_t n_q_blocks = (neq1 + Br - 1) / Br;
-    uint32_t q_start_min = 0;
-    uint32_t q_start_max = neq1;
+    uint32_t q_start_min  = 0;
+    uint32_t q_start_max  = neq1;
+    uint32_t kv_head_min  = 0;
+    uint32_t kv_head_max  = n_kv_heads;
 
     if (octx->ctx->mdev.count > 1) {
-        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(n_q_blocks, htp_tensor_mdev_data_aligned(dst) ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
-        const uint32_t block_start = range.start;
-        const uint32_t block_end   = range.start + range.count;
+        const uint32_t mdev_count = octx->ctx->mdev.count;
+        const uint32_t mdev_idx   = octx->ctx->mdev.idx;
+        const uint32_t dst_e_size = (dst->type == HTP_TYPE_F32) ? sizeof(float) : sizeof(__fp16);
+        const bool can_split      = htp_tensor_can_row_partition(dst, dst_e_size);
 
-        if (block_start >= block_end) {
-            return HTP_STATUS_OK;
+        if (kparams->head_split && can_split && n_kv_heads >= mdev_count && n_kv_heads % mdev_count == 0) {
+            const uint32_t kv_per_core = n_kv_heads / mdev_count;
+            kv_head_min = mdev_idx * kv_per_core;
+            kv_head_max = kv_head_min + kv_per_core;
+        } else {
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(n_q_blocks, can_split ? 1 : 0, mdev_idx, mdev_count, &octx->ctx->mdev.count_div);
+            const uint32_t block_start = range.start;
+            const uint32_t block_end   = range.start + range.count;
+
+            if (block_start >= block_end) {
+                return HTP_STATUS_OK;
+            }
+
+            q_start_min = block_start * Br;
+            q_start_max = MIN(block_end * Br, neq1);
         }
-
-        q_start_min = block_start * Br;
-        q_start_max = MIN(block_end * Br, neq1);
     }
 
     // ======== VTCM allocation (GQA-aware) ========
@@ -2032,7 +2045,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
             const size_t   g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
             const size_t   n_row_tiles = g_br_actual / HMX_FP16_TILE_N_ROWS;
 
-            for (uint32_t kv_head = 0; kv_head < n_kv_heads; ++kv_head) {
+            for (uint32_t kv_head = kv_head_min; kv_head < kv_head_max; ++kv_head) {
                 const uint32_t ik2 = kv_head;
                 const uint32_t ik3 = fastdiv(ib3, &kparams->broadcast_rk3);
                 const uint32_t iv2 = kv_head;
@@ -2040,7 +2053,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
 
                 // 1. Push Q and KV DMAs for the very first iteration.
                 // Subsequent iterations are enqueued early at the end of the previous iteration.
-                if (ib3 == 0 && q_start == q_start_min && kv_head == 0) {
+                if (ib3 == 0 && q_start == q_start_min && kv_head == kv_head_min) {
                     const dma_addr_t q_ptr = q->data + q_start * q->nb[1] +
                                             (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3];
                     const size_t q_row_bytes = q_transposed ? n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
@@ -2358,8 +2371,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                 uint32_t next_kv_head = kv_head + 1;
                 uint32_t next_q_start = q_start;
                 uint32_t next_ib3     = ib3;
-                if (next_kv_head >= n_kv_heads) {
-                    next_kv_head = 0;
+                if (next_kv_head >= kv_head_max) {
+                    next_kv_head = kv_head_min;
                     next_q_start = q_start + Br;
                     if (next_q_start >= q_start_max) {
                         next_q_start = q_start_min;
@@ -2478,7 +2491,7 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
         factx.src3_div3 = kparams->src3_div3;
     }
 
-    factx.is_q_fp32 = (kparams->is_q_fp32 != 0);
+    factx.is_q_fp32 = (q->type == HTP_TYPE_F32);
     factx.size_q_row_padded = kparams->u.hvx.size_q_row_padded;
     factx.size_k_row_padded = kparams->u.hvx.size_k_row_padded;
     factx.size_v_row_padded = kparams->u.hvx.size_v_row_padded;
@@ -2488,7 +2501,10 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     factx.scale = kparams->scale;
     factx.max_bias = kparams->max_bias;
     factx.has_softcap = (kparams->logit_softcap != 0.0f);
-    factx.logit_softcap = factx.has_softcap ? (__fp16) kparams->logit_softcap : 0;
+    factx.logit_softcap = 0;
+    if (factx.has_softcap) {
+        factx.logit_softcap = (__fp16) kparams->logit_softcap;
+    }
 
     factx.n_head_log2 = kparams->n_head_log2;
     factx.m0          = kparams->m0;
@@ -2512,10 +2528,25 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     uint32_t qrows      = total_qrows;
 
     if (octx->ctx->mdev.count > 1) {
-        const bool can_split = htp_tensor_mdev_data_aligned(dst) && ((dst->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0);
-        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_qrows, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
-        qrow_start = range.start;
-        qrows      = range.count;
+        const uint32_t mdev_count = octx->ctx->mdev.count;
+        const uint32_t mdev_idx   = octx->ctx->mdev.idx;
+        const uint32_t n_kv_heads = k->ne[2];
+        const uint32_t dst_e_size = (dst->type == HTP_TYPE_F32) ? sizeof(float) : sizeof(__fp16);
+        const bool can_split      = htp_tensor_can_row_partition(dst, dst_e_size);
+
+        // head range is contiguous in flat row space only when neq3 == 1
+        if (kparams->head_split && can_split && neq3 == 1 && n_kv_heads >= mdev_count && n_kv_heads % mdev_count == 0) {
+            const uint32_t G              = kparams->G;
+            const uint32_t kv_per_core    = n_kv_heads / mdev_count;
+            const uint32_t heads_per_core = kv_per_core * G;
+            const uint32_t head_start     = mdev_idx * heads_per_core;
+            qrow_start = head_start * neq1;
+            qrows      = heads_per_core * neq1;
+        } else {
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_qrows, can_split ? 1 : 0, mdev_idx, mdev_count, &octx->ctx->mdev.count_div);
+            qrow_start = range.start;
+            qrows      = range.count;
+        }
     }
 
     if (qrows == 0) {

ggml/src/ggml-hexagon/htp/flash-attn-ops.h

@@ -34,8 +34,8 @@ enum htp_fa_kernel_type {
 
 struct htp_fa_kernel_params {
     uint8_t  kernel_type;        // enum htp_fa_kernel_type
-    uint8_t  is_q_fp32;          // 1 = Q type is F32, 0 = F16
-    uint8_t  is_dst_fp32;        // 1 = dst type is F32, 0 = F16
+    uint8_t  head_split;         // 1 = partition by KV heads in multicore, 0 = token partition
+    uint8_t  flags;              // reserved
     uint8_t  n_threads;          // Number of threads to run
 
     // Common parameters

ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h

@@ -745,7 +745,7 @@ void convert_f16_weight_to_fp16_tiles_task(
                 const uint8_t *r0 = state->src + row0 * state->row_stride;
                 const uint8_t *r1 = state->src + row1 * state->row_stride;
 
-                HVX_Vector v0 = hvx_vmemu((const __fp16 *)(r0 + byte_off));
+                HVX_Vector v0 = (row0 < state->n_cols) ? hvx_vmemu((const __fp16 *)(r0 + byte_off)) : Q6_V_vzero();
                 HVX_Vector v1 = (row1 < state->n_cols) ? hvx_vmemu((const __fp16 *)(r1 + byte_off)) : Q6_V_vzero();
 
                 Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v0);
@@ -788,7 +788,7 @@ void quantize_f32_weight_to_fp16_tiles_task(
                 const uint8_t *r0 = state->src + row0 * state->row_stride;
                 const uint8_t *r1 = state->src + row1 * state->row_stride;
 
-                HVX_Vector v0_f32 = hvx_vmem((const float *)(r0 + byte_off));
+                HVX_Vector v0_f32 = (row0 < state->n_cols) ? hvx_vmem((const float *)(r0 + byte_off)) : Q6_V_vzero();
                 HVX_Vector v1_f32 = (row1 < state->n_cols) ? hvx_vmem((const float *)(r1 + byte_off)) : Q6_V_vzero();
 
                 HVX_Vector v_out = hvx_vec_f32_to_f16(v0_f32, v1_f32);
@@ -988,9 +988,7 @@ static void transfer_output_chunk_fp16_to_fp32_col_chunk(
     uint32_t src2_stride,
     uint32_t dst_cols
 ) {
-    assert(c_len % HTP_MM_HMX_TILE_N_COLS == 0);
-    assert(total_n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
-    const size_t tile_row_stride = (total_n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;
+    const size_t tile_row_stride = hmx_ceil_div(total_n_cols, HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;
 
     const HVX_Vector one = hvx_vec_splat_f16(1.0);
 
@@ -1137,6 +1135,73 @@ static void transfer_activation_row_pair_fp32_to_fp16(
     }
 }
 
+// F16-input variant of transfer_activation_row_pair_fp32_to_fp16, for F16 activation
+// (src1). Same shape as the F16 Q-prep in hmx-fa-kernels.h: one 128-byte load carries 64 f16
+// = two tile columns, and Q6_W_vshuff_VVR interleaves the two rows straight into the HMX tile
+// layout, so no F32 round-trip is needed. Rows are only 64-byte aligned when k_block is an odd
+// multiple of the tile width, hence the unaligned load type.
+static void transfer_activation_row_pair_f16_to_f16(__fp16 * restrict vtcm_dst,
+                                                    const __fp16 * restrict row0,
+                                                    const __fp16 * restrict row1,
+                                                    uint32_t r,
+                                                    uint32_t k_block,
+                                                    uint32_t k_valid,
+                                                    bool     row0_valid,
+                                                    bool     row1_valid) {
+    uint32_t r0 = r / HTP_MM_HMX_TILE_N_ROWS;  // tile row index
+    uint32_t r1 = r % HTP_MM_HMX_TILE_N_ROWS;  // intra-tile row idx
+
+    const uint32_t n_tile_cols = k_block / HTP_MM_HMX_TILE_N_COLS;
+    __fp16 * restrict tile_row = vtcm_dst + (size_t) r0 * n_tile_cols * HTP_MM_HMX_TILE_N_ELMS;
+
+    const HVX_UVector * pv0 = row0_valid ? (const HVX_UVector *) row0 : NULL;
+    const HVX_UVector * pv1 = row1_valid ? (const HVX_UVector *) row1 : NULL;
+
+    uint32_t c = 0;
+    for (; c + 64 <= k_valid; c += 64) {
+        HVX_Vector     v0 = pv0 ? pv0[c / 64] : Q6_V_vzero();
+        HVX_Vector     v1 = pv1 ? pv1[c / 64] : Q6_V_vzero();
+        HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+        uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS;
+
+        HVX_Vector * tile0 = (HVX_Vector *) (tile_row + (size_t) c0 * HTP_MM_HMX_TILE_N_ELMS);
+        HVX_Vector * tile1 = (HVX_Vector *) (tile_row + (size_t) (c0 + 1) * HTP_MM_HMX_TILE_N_ELMS);
+
+        tile0[r1 / 2] = Q6_V_lo_W(vp);
+        tile1[r1 / 2] = Q6_V_hi_W(vp);
+    }
+    // Tail: fewer than 64 valid columns left, plus the k_valid..k_block padding that HMX will
+    // still multiply, so it has to be written as zeros.
+    for (; c < k_block; c += 64) {
+        HVX_Vector v0 = Q6_V_vzero();
+        HVX_Vector v1 = Q6_V_vzero();
+
+        if (c < k_valid) {
+            uint32_t       rem  = k_valid - c;  // 1..63 valid f16 lanes
+            HVX_VectorPred mask = Q6_Q_vsetq2_R(rem * sizeof(__fp16));
+            if (pv0) {
+                v0 = Q6_V_vmux_QVV(mask, pv0[c / 64], Q6_V_vzero());
+            }
+            if (pv1) {
+                v1 = Q6_V_vmux_QVV(mask, pv1[c / 64], Q6_V_vzero());
+            }
+        }
+
+        HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+        uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS;
+
+        HVX_Vector * tile0 = (HVX_Vector *) (tile_row + (size_t) c0 * HTP_MM_HMX_TILE_N_ELMS);
+        tile0[r1 / 2]      = Q6_V_lo_W(vp);
+
+        if (c0 + 1 < n_tile_cols) {
+            HVX_Vector * tile1 = (HVX_Vector *) (tile_row + (size_t) (c0 + 1) * HTP_MM_HMX_TILE_N_ELMS);
+            tile1[r1 / 2]      = Q6_V_hi_W(vp);
+        }
+    }
+}
+
 static void transfer_activation_row_pair_fp32_to_fp16_col_chunk(
         __fp16 *restrict vtcm_dst,
         const float *restrict row0, // offset by c_first

ggml/src/ggml-hexagon/htp/hmx-utils.h

@@ -73,19 +73,20 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
         for (uint32_t r = start_row; r < end_row; r += 2) {
             const uint32_t   ct             = r / HMX_FP16_TILE_N_ROWS;
             const uint32_t   local_r        = r % HMX_FP16_TILE_N_ROWS;
+            const bool       row0_valid     = r < n_cols;
             const bool       next_row_valid = (r + 1) < end_row && (r + 1) < n_cols;
             const HVX_Vector v_off0         = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4));
             const HVX_Vector v_off1         = Q6_Vw_vadd_VwVw(v_off0, v_scat_step);
 
             __fp16 * tile_base = vtcm_dst + (size_t) ct * n_k_tiles * HMX_FP16_TILE_N_ELMS;
-            const uint8_t * p0 = (const uint8_t *) (vtcm_src + r * src_stride);
+            const uint8_t * p0 = row0_valid ? (const uint8_t *) (vtcm_src + r * src_stride) : NULL;
             const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL;
 
-            assert(hex_is_aligned(p0, 128));
-            assert(hex_is_aligned(p1, 128));
+            assert(!p0 || hex_is_aligned(p0, 128));
+            assert(!p1 || hex_is_aligned(p1, 128));
             assert(c_byte_step % 128 == 0);
 
-            if (p1) {
+            if (p0 && p1) {
                 for (uint32_t i = 0; i < n_c_iters; ++i) {
                     HVX_Vector v0 = hvx_vmem(p0); p0 += c_byte_step;
                     HVX_Vector v1 = hvx_vmem(p1); p1 += c_byte_step;
@@ -96,9 +97,12 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
             } else {
                 const HVX_Vector vzero = Q6_V_vzero();
                 for (uint32_t i = 0; i < n_c_iters; ++i) {
-                    HVX_Vector v0 = hvx_vmem(p0); p0 += c_byte_step;
+                    HVX_Vector v0 = p0 ? hvx_vmem(p0) : vzero;
+                    if (p0) p0 += c_byte_step;
+                    HVX_Vector v1 = p1 ? hvx_vmem(p1) : vzero;
+                    if (p1) p1 += c_byte_step;
                     Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0);
-                    Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, vzero);
+                    Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, v1);
                     tile_base += dst_step;
                 }
             }
@@ -113,15 +117,16 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
         for (uint32_t r = start_row; r < end_row; r += 2) {
             const uint32_t   ct             = r / HMX_FP16_TILE_N_ROWS;
             const uint32_t   local_r        = r % HMX_FP16_TILE_N_ROWS;
+            const bool       row0_valid     = r < n_cols;
             const bool       next_row_valid = (r + 1) < end_row && (r + 1) < n_cols;
             const HVX_Vector v_off0         = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4));
             const HVX_Vector v_off1         = Q6_Vw_vadd_VwVw(v_off0, v_scat_step);
 
             __fp16 * tile_base = vtcm_dst + (size_t) ct * n_k_tiles * HMX_FP16_TILE_N_ELMS;
-            const uint8_t * p0 = (const uint8_t *) (vtcm_src + r * src_stride);
+            const uint8_t * p0 = row0_valid ? (const uint8_t *) (vtcm_src + r * src_stride) : NULL;
             const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL;
 
-            if (p1) {
+            if (p0 && p1) {
                 for (uint32_t i = 0; i < n_c_iters; ++i) {
                     HVX_Vector v0 = hvx_vmemu(p0); p0 += c_byte_step;
                     HVX_Vector v1 = hvx_vmemu(p1); p1 += c_byte_step;
@@ -132,9 +137,12 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
             } else {
                 const HVX_Vector vzero = Q6_V_vzero();
                 for (uint32_t i = 0; i < n_c_iters; ++i) {
-                    HVX_Vector v0 = hvx_vmemu(p0); p0 += c_byte_step;
+                    HVX_Vector v0 = p0 ? hvx_vmemu(p0) : vzero;
+                    if (p0) p0 += c_byte_step;
+                    HVX_Vector v1 = p1 ? hvx_vmemu(p1) : vzero;
+                    if (p1) p1 += c_byte_step;
                     Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0);
-                    Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, vzero);
+                    Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, v1);
                     tile_base += dst_step;
                 }
             }

ggml/src/ggml-hexagon/htp/htp-ctx.h

@@ -173,6 +173,8 @@ int op_solve_tri(struct htp_ops_context * octx);
 int op_gated_delta_net(struct htp_ops_context * octx);
 int op_pad(struct htp_ops_context * octx);
 int op_im2col(struct htp_ops_context * octx);
+int op_pool_2d(struct htp_ops_context * octx);
+int op_pool_1d(struct htp_ops_context * octx);
 int op_allreduce(struct htp_ops_context * octx);
 int op_roll(struct htp_ops_context * octx);
 

ggml/src/ggml-hexagon/htp/htp-ops.h

@@ -115,6 +115,8 @@ enum htp_op_code {
     HTP_OP_ARGMAX,
     HTP_OP_UNARY_GELU_ERF,
     HTP_OP_GLU_GEGLU_ERF,
+    HTP_OP_POOL_2D,
+    HTP_OP_POOL_1D,
 
     HTP_OP_INVALID
 };

ggml/src/ggml-hexagon/htp/main.c

@@ -931,6 +931,12 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_ROLL:
             return op_roll(octx);
 
+        case HTP_OP_POOL_2D:
+            return op_pool_2d(octx);
+
+        case HTP_OP_POOL_1D:
+            return op_pool_1d(octx);
+
         case HTP_OP_CONCAT:
             return op_concat(octx);
 

ggml/src/ggml-hexagon/htp/matmul-ops.h

@@ -87,24 +87,26 @@ enum htp_mm_kernel_type {
 
 // Op-specific struct for precomputed matmul params
 struct htp_mm_kernel_params {
-    int32_t  kernel_type;        // enum htp_mm_kernel_type
-    int32_t  pipeline;           // 1 = pipelined execution, 0 = standard
+    uint8_t  kernel_type;        // enum htp_mm_kernel_type
+    uint8_t  pipeline;           // 1 = pipelined execution, 0 = standard
+    uint8_t  collapse;           // 1 = collapse outer dims into 2D, 0 = standard
+    uint8_t  n_hmx;              // 1 = use HMX, 0 = use HVX
+
+    uint8_t  n_threads;          // Number of threads to spawn
+    uint8_t  n_act_threads;      // Number of threads for activation preparation
+    uint8_t  n_prefetch;         // Prefetch lookahead buffers/rows in VTCM
+    uint8_t  n_weights;          // Number of weights for fused NX
+
     int32_t  m_chunk;            // Row chunk size (M chunk)
     int32_t  n_chunk;            // Col chunk size (N chunk)
-    int32_t  n_threads;          // Number of threads to spawn
-    int32_t  n_act_threads;      // Number of threads for activation preparation
-    int32_t  n_hmx;              // 1 = use HMX, 0 = use HVX
-    int32_t  n_prefetch;         // Prefetch lookahead buffers/rows in VTCM
     int32_t  tile_size;          // Weight tile size
     int32_t  aligned_tile_size;  // Aligned weight tile size (padded to 128)
-    int32_t  src1_row_size;      // Row size for quantized activation
+    int32_t  act_row_size;       // Row size for activation scratchpad
     int32_t  vtcm_size;          // Total required scratchpad size in VTCM
     int32_t  vtcm_src0_size;     // src0 scratchpad size in VTCM
-    int32_t  vtcm_src1_size;     // src1 scratchpad size in VTCM
-    int32_t  vtcm_src2_size;     // src2 scratchpad size in VTCM (fused only)
-    int32_t  vtcm_src3_size;     // src3 scratchpad size in VTCM (fused only)
+    int32_t  vtcm_act_size;      // activation scratchpad size in VTCM
+    int32_t  vtcm_bias_size;     // bias scratchpad size in VTCM (fused only)
     int32_t  vtcm_dst_size;      // dst scratchpad size in VTCM
-    int32_t  n_weights;          // Number of weights for fused NX
 
     // Precomputed division values
     struct fastdiv_values div_ne12_ne1;
@@ -147,6 +149,7 @@ static inline int htp_mm_hmx_compute_chunks(size_t   vtcm_total,
     const size_t usable = vtcm_total - overhead;
 
     size_t best_cost = SIZE_MAX;
+    size_t best_tail_waste = SIZE_MAX;
     size_t best_mn   = 0;
     size_t best_m = 0, best_n = 0;
 
@@ -173,12 +176,17 @@ static inline int htp_mm_hmx_compute_chunks(size_t   vtcm_total,
             size_t mblocks = ((size_t) m + mc - 1) / mc;
             size_t nblocks = ((size_t) n + nc - 1) / nc;
             size_t cost    = mblocks * m_block_cost + nblocks * n_block_cost;
+            size_t rem     = n % nc;
+            size_t tail_waste = (rem == 0) ? 0 : (nc - rem);
             size_t mn      = mc * nc;
-            if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
-                best_cost = cost;
-                best_mn   = mn;
-                best_m    = mc;
-                best_n    = nc;
+            if (cost < best_cost ||
+                (cost == best_cost && tail_waste < best_tail_waste) ||
+                (cost == best_cost && tail_waste == best_tail_waste && mn > best_mn)) {
+                best_cost       = cost;
+                best_tail_waste = tail_waste;
+                best_mn         = mn;
+                best_m          = mc;
+                best_n          = nc;
             }
         }
 
@@ -349,7 +357,7 @@ struct htp_mm_hmx_vtcm_layout {
     size_t off_dst[2];        // [1] is only used when pipelined
     size_t off_scratch[2];    // dequantization scratch pads
     size_t off_scales;        // HMX scales (256 bytes)
-    size_t off_src2;          // src2 bias in VTCM
+    size_t off_bias;          // bias in VTCM
 
     // Cached sizes of regions for HMX kernel use
     size_t weight_area_bytes;
@@ -358,25 +366,23 @@ struct htp_mm_hmx_vtcm_layout {
     size_t output_area_bytes;
     size_t scratch_bytes[2];
     size_t act_head_stride;
-    size_t src2_bytes;
+    size_t bias_bytes;
 
     size_t total_bytes;
 };
 
 struct htp_mm_hvx_vtcm_layout {
     // Byte offsets from vtcm_base for each region
-    size_t off_src1;          // vtcm_src1 (activation)
+    size_t off_act;           // vtcm_act (activation)
     size_t off_src0;          // vtcm_src0 (weight/Wk)
-    size_t off_src2;          // vtcm_src2 (Wq / fused only)
-    size_t off_src3;          // vtcm_src3 (Wv / fused only)
+    size_t off_bias;          // vtcm_bias (bias / fused add only)
     size_t off_dst;           // vtcm_dst (output scratch)
     size_t off_act_raw;       // vtcm_act_raw (raw activation DMA staging)
 
     // Cached sizes
     size_t src0_bytes;
-    size_t src1_bytes;
-    size_t src2_bytes;
-    size_t src3_bytes;
+    size_t act_bytes;
+    size_t bias_bytes;
     size_t dst_bytes;
     size_t act_raw_bytes;
 
@@ -394,7 +400,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
     bool pipeline,
     uint32_t act_threads,
     uint32_t aligned_tile_size,
-    size_t src2_size
+    size_t bias_size
 ) {
     size_t off = 0;
 
@@ -411,7 +417,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
         size_t off_group_a = 0;
         VTCM_LAYOUT_ALLOC(off_group_a, off_act, activation_area_size);
         VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data
-        VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src2, hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE), src2_size > 0);
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_bias, hex_align_up(bias_size, HTP_MM_HMX_TILE_SIZE), bias_size > 0);
 
         // Group B: Compute-only buffers (starts at off_group_a)
         size_t off_group_b = off_group_a;
@@ -439,7 +445,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
         L->scratch_bytes[0]  = scratch_area_size;
         L->scratch_bytes[1]  = scratch_area_size;
         L->act_head_stride   = act_head_stride;
-        L->src2_bytes        = src2_size;
+        L->bias_bytes        = bias_size;
 
         off = off_group_a + hex_smax(group_b_size, group_c_size);
     } else {
@@ -463,7 +469,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
         size_t off_group_a = 0;
         VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data
         VTCM_LAYOUT_ALLOC(off_group_a, off_act, act_area_size);
-        VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src2, hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE), src2_size > 0);
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_bias, hex_align_up(bias_size, HTP_MM_HMX_TILE_SIZE), bias_size > 0);
 
         // Group B: Compute-only buffers (starts at off_group_a)
         size_t off_group_b = off_group_a;
@@ -491,7 +497,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
         L->scratch_bytes[0]  = scratch0_size;
         L->scratch_bytes[1]  = scratch1_size;
         L->act_head_stride   = 0;
-        L->src2_bytes        = src2_size;
+        L->bias_bytes        = bias_size;
 
         off = off_group_a + hex_smax(group_b_size, group_c_size);
     }
@@ -504,21 +510,20 @@ static inline void htp_mm_hvx_vtcm_layout_build(
     int kernel_type,
     int wtype,
     uint32_t ne10,       // k
-    uint32_t src1_nrows, // m_total
+    uint32_t act_nrows,  // m_total
     uint32_t n_threads,
     size_t dst_row_size,
     size_t src0_row_size,
-    size_t src1_row_size,
-    size_t src2_row_size,
+    size_t act_row_size,
+    size_t bias_row_size,
     uint32_t n_prefetch,
     bool is_matmul_id,
     bool is_fused_nx
 ) {
-    (void)src1_row_size;
+    (void)act_row_size;
     size_t src0_sz    = 0;
-    size_t src1_sz    = 0;
-    size_t src2_sz    = src2_row_size > 0 ? htp_mm_round_up(src2_row_size, 128) : 0;
-    size_t src3_sz    = 0;
+    size_t act_sz     = 0;
+    size_t bias_sz    = bias_row_size > 0 ? htp_mm_round_up(bias_row_size, 128) : 0;
     size_t dst_sz     = 0;
     size_t act_raw_sz = 0;
 
@@ -544,22 +549,21 @@ static inline void htp_mm_hvx_vtcm_layout_build(
         }
 
         size_t tiled_act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
-        size_t act_sz = hex_round_up(tiled_act_row_size * src1_nrows, 128);
+        size_t q_act_sz = hex_round_up(tiled_act_row_size * act_nrows, 128);
         size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
 
         src0_sz    = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer
-        src1_sz    = act_sz;                           // quantized activation buffer
-        src2_sz    = 0;
-        src3_sz    = 0;
+        act_sz     = q_act_sz;                         // quantized activation buffer
+        bias_sz    = 0;
         dst_sz     = 0;
-        act_raw_sz = hex_round_up(raw_row_size * src1_nrows, 128);
+        act_raw_sz = hex_round_up(raw_row_size * act_nrows, 128);
     } else if (is_matmul_id) {
         const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
-        const size_t src1_row_size_tiled = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10)
-                                                                                               : htp_mm_q8_0_tiled_row_size(ne10);
+        const size_t act_row_size_tiled = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10)
+                                                                          : htp_mm_q8_0_tiled_row_size(ne10);
 
         size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
-        src1_sz                   = htp_mm_round_up(src1_row_size_tiled * src1_nrows, 256);
+        act_sz                    = htp_mm_round_up(act_row_size_tiled * act_nrows, 256);
 
         if (is_repack) {
             const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
@@ -573,25 +577,24 @@ static inline void htp_mm_hvx_vtcm_layout_build(
 
         src0_sz    = src0_sz_per_thread * n_threads;
         dst_sz     = 0;
-        src2_sz    = 0;
-        src3_sz    = 0;
-        act_raw_sz = hex_round_up(raw_row_size * src1_nrows, 128);
+        bias_sz    = 0;
+        act_raw_sz = hex_round_up(raw_row_size * act_nrows, 128);
     } else {
         const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
-        const size_t dst_nrows = (src1_nrows > 1) ? 0 : 1;
+        const size_t dst_nrows = (act_nrows > 1) ? 0 : 1;
 
         switch (kernel_type) {
             case HTP_MM_KERNEL_HVX_F16_F16_VTCM: {
-                size_t f16_src1_row_size = htp_mm_round_up(ne10 * 2, 128);
-                src1_sz    = htp_mm_round_up(f16_src1_row_size * src1_nrows, 256);
+                size_t f16_act_row_size = htp_mm_round_up(ne10 * 2, 128);
+                act_sz     = htp_mm_round_up(f16_act_row_size * act_nrows, 256);
                 src0_sz    = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
                 dst_sz     = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
-                act_raw_sz = hex_round_up(hex_round_up(ne10 * sizeof(float), 128) * src1_nrows, 128);
+                act_raw_sz = hex_round_up(hex_round_up(ne10 * sizeof(float), 128) * act_nrows, 128);
                 break;
             }
             case HTP_MM_KERNEL_HVX_F32_F32_VTCM: {
-                size_t f32_src1_row_size = htp_mm_round_up(ne10 * 4, 128);
-                src1_sz    = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256);
+                size_t f32_act_row_size = htp_mm_round_up(ne10 * 4, 128);
+                act_sz     = htp_mm_round_up(f32_act_row_size * act_nrows, 256);
                 src0_sz    = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
                 dst_sz     = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
                 act_raw_sz = 0;
@@ -599,10 +602,10 @@ static inline void htp_mm_hvx_vtcm_layout_build(
             }
             case HTP_MM_KERNEL_HVX_QUANT_BLOCK:
             case HTP_MM_KERNEL_HVX_QUANT_ROW: {
-                size_t q_src1_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
+                size_t q_act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
 
                 src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
-                src1_sz = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
+                act_sz  = htp_mm_round_up(q_act_row_size * act_nrows, 256);
 
                 src0_sz = src0_sz * n_threads;
 
@@ -614,10 +617,10 @@ static inline void htp_mm_hvx_vtcm_layout_build(
                     src0_sz = repacked_vtcm_size * n_threads;
                 }
 
-                size_t dst_slice_per_thread = (dst_nrows > 0 && src1_nrows == 1) ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0;
+                size_t dst_slice_per_thread = (dst_nrows > 0 && act_nrows == 1) ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0;
                 dst_sz = dst_slice_per_thread * n_threads;
                 size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
-                act_raw_sz = hex_round_up(raw_row_size * src1_nrows, 128);
+                act_raw_sz = hex_round_up(raw_row_size * act_nrows, 128);
                 break;
             }
             default:
@@ -627,9 +630,8 @@ static inline void htp_mm_hvx_vtcm_layout_build(
 
     // Group A: Persistent buffers across chunk compute
     size_t off_group_a = 0;
-    VTCM_LAYOUT_ALLOC(off_group_a, off_src1, src1_sz);
-    VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src2, src2_sz, src2_sz > 0);
-    VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src3, src3_sz, src3_sz > 0);
+    VTCM_LAYOUT_ALLOC(off_group_a, off_act, act_sz);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_bias, bias_sz, bias_sz > 0);
 
     // Group B: Compute-only buffers (starts at off_group_a)
     size_t off_group_b = off_group_a;
@@ -643,9 +645,8 @@ static inline void htp_mm_hvx_vtcm_layout_build(
     const size_t group_c_size = off_group_c - off_group_a;
 
     L->src0_bytes    = src0_sz;
-    L->src1_bytes    = src1_sz;
-    L->src2_bytes    = src2_sz;
-    L->src3_bytes    = src3_sz;
+    L->act_bytes     = act_sz;
+    L->bias_bytes    = bias_sz;
     L->dst_bytes     = dst_sz;
     L->act_raw_bytes = act_raw_sz;
     L->total_bytes   = off_group_a + hex_smax(group_b_size, group_c_size);
@@ -655,12 +656,12 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
     int kernel_type,
     int wtype,
     uint32_t ne10,
-    uint32_t src1_nrows,
+    uint32_t act_nrows,
     uint32_t n_threads,
     size_t dst_row_size,
     size_t src0_row_size,
-    size_t src1_row_size,
-    size_t src2_row_size,
+    size_t act_row_size,
+    size_t bias_row_size,
     uint32_t n_prefetch,
     size_t vtcm_budget,
     struct htp_mm_hvx_vtcm_layout * L_out,
@@ -668,17 +669,17 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
 ) {
     struct htp_mm_hvx_vtcm_layout L;
     htp_mm_hvx_vtcm_layout_build(
-        &L, kernel_type, wtype, ne10, src1_nrows, n_threads,
-        dst_row_size, src0_row_size, src1_row_size, src2_row_size, n_prefetch, false, false
+        &L, kernel_type, wtype, ne10, act_nrows, n_threads,
+        dst_row_size, src0_row_size, act_row_size, bias_row_size, n_prefetch, false, false
     );
 
     if (L.total_bytes <= vtcm_budget) {
         *L_out = L;
-        *m_chunk_out = src1_nrows;
+        *m_chunk_out = act_nrows;
         return true;
     }
 
-    const size_t fixed_bytes = L.src0_bytes + L.src2_bytes + L.dst_bytes;
+    const size_t fixed_bytes = L.src0_bytes + L.bias_bytes + L.dst_bytes;
     if (vtcm_budget <= fixed_bytes) {
         return false;
     }
@@ -707,8 +708,8 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
     if (m_chunk > 1) {
         m_chunk &= ~1U;
     }
-    if (m_chunk > src1_nrows) {
-        m_chunk = src1_nrows;
+    if (m_chunk > act_nrows) {
+        m_chunk = act_nrows;
     }
     if (m_chunk < 1) {
         return false;
@@ -716,14 +717,14 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
 
     htp_mm_hvx_vtcm_layout_build(
         &L, kernel_type, wtype, ne10, m_chunk, n_threads,
-        dst_row_size, src0_row_size, src1_row_size, src2_row_size, n_prefetch, false, false
+        dst_row_size, src0_row_size, act_row_size, bias_row_size, n_prefetch, false, false
     );
 
     while (m_chunk > 2 && L.total_bytes > vtcm_budget) {
         m_chunk -= 2;
         htp_mm_hvx_vtcm_layout_build(
             &L, kernel_type, wtype, ne10, m_chunk, n_threads,
-            dst_row_size, src0_row_size, src1_row_size, src2_row_size, n_prefetch, false, false
+            dst_row_size, src0_row_size, act_row_size, bias_row_size, n_prefetch, false, false
         );
     }
 
@@ -737,18 +738,18 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
 }
 
 static inline size_t htp_mm_hmx_get_2d_vtcm_size(
-    int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size, size_t src2_size
+    int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size, size_t bias_size
 ) {
     struct htp_mm_hmx_vtcm_layout L;
-    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, pipeline, act_threads, aligned_tile_size, src2_size);
+    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, pipeline, act_threads, aligned_tile_size, bias_size);
     return L.total_bytes;
 }
 
 static inline size_t htp_mm_hmx_get_batched_vtcm_size(
-    int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool pipeline, uint32_t act_threads, size_t src2_size) {
+    int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool pipeline, uint32_t act_threads, size_t bias_size) {
     (void)pipeline;
     struct htp_mm_hmx_vtcm_layout L;
-    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, false, act_threads, 0, src2_size);
+    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, false, act_threads, 0, bias_size);
     return L.total_bytes;
 }
 
@@ -760,7 +761,7 @@ static inline bool htp_mm_hmx_solve_batched_params(
     uint32_t group_size,
     int n_threads,
     bool pipeline,
-    size_t src2_size,
+    size_t bias_size,
     size_t vtcm_budget,
     size_t * m_chunk_out,
     size_t * n_chunk_out,
@@ -775,7 +776,7 @@ static inline bool htp_mm_hmx_solve_batched_params(
 
     int act_threads = n_threads;
     while (act_threads >= 1) {
-        size_t group_overhead = htp_mm_hmx_get_batched_overhead() + (src2_size > 0 ? hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE) : 0);
+        size_t group_overhead = htp_mm_hmx_get_batched_overhead() + (bias_size > 0 ? hex_align_up(bias_size, HTP_MM_HMX_TILE_SIZE) : 0);
         size_t group_size_per_n, group_size_per_m, group_size_per_mn;
         htp_mm_hmx_get_batched_chunk_costs(k, group_size, &group_size_per_n, &group_size_per_m, &group_size_per_mn);
 
@@ -785,8 +786,8 @@ static inline bool htp_mm_hmx_solve_batched_params(
 
         if (htp_mm_hmx_compute_chunks(vtcm_budget, group_overhead, group_size_per_n, group_size_per_m, group_size_per_mn, hex_align_up(ne11, 32), ne01_padded,
                                (size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) ne11 * HTP_MM_HMX_COST_A_CONVERT,
-                               &m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
-            size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, pipeline, act_threads, src2_size);
+                                &m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
+            size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, pipeline, act_threads, bias_size);
             if (exact_size <= vtcm_budget) {
                 size_t mblocks = ((size_t) ne11 + m_chunk_candidate - 1) / m_chunk_candidate;
                 if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) {
@@ -826,7 +827,7 @@ static inline bool htp_mm_hmx_solve_2d_params(
     bool pipeline,
     bool is_matmul_id,
     uint32_t aligned_tile_size,
-    size_t src2_size,
+    size_t bias_size,
     size_t vtcm_budget,
     size_t * m_chunk_out,
     size_t * n_chunk_out,
@@ -843,7 +844,7 @@ static inline bool htp_mm_hmx_solve_2d_params(
 
     int act_threads = n_threads;
     while (act_threads >= 1) {
-        size_t simple_2d_overhead = htp_mm_hmx_get_2d_overhead(pipeline, is_matmul_id) + (src2_size > 0 ? hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE) : 0);
+        size_t simple_2d_overhead = htp_mm_hmx_get_2d_overhead(pipeline, is_matmul_id) + (bias_size > 0 ? hex_align_up(bias_size, HTP_MM_HMX_TILE_SIZE) : 0);
         size_t simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn;
         htp_mm_hmx_get_2d_chunk_costs(wtype, k, pipeline, aligned_tile_size, &simple_2d_size_per_n, &simple_2d_size_per_m, &simple_2d_size_per_mn);
 
@@ -854,7 +855,7 @@ static inline bool htp_mm_hmx_solve_2d_params(
         if (htp_mm_hmx_compute_chunks(vtcm_budget, simple_2d_overhead, simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn, m_for_chunks, ne01_padded,
                                (size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) m_for_cost * HTP_MM_HMX_COST_A_CONVERT,
                                &m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
-            size_t exact_size = htp_mm_hmx_get_2d_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, pipeline, is_matmul_id ? 0 : act_threads, aligned_tile_size, src2_size);
+            size_t exact_size = htp_mm_hmx_get_2d_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, pipeline, is_matmul_id ? 0 : act_threads, aligned_tile_size, bias_size);
             if (exact_size <= vtcm_budget) {
                 size_t mblocks = ((size_t) m_for_cost + m_chunk_candidate - 1) / m_chunk_candidate;
                 if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) {

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

@@ -0,0 +1,399 @@
+#pragma clang diagnostic ignored "-Wunused-variable"
+
+#include <float.h>
+#include <HAP_farf.h>
+
+#include "hex-common.h"
+#include "dma-queue.h"
+#include "hex-profile.h"
+#include "htp-ctx.h"
+#include "htp-ops.h"
+#include "htp-tensor.h"
+#include "hvx-inverse.h"
+#include "hvx-types.h"
+#include "hvx-utils.h"
+#include "pool-ops.h"
+
+#define HTP_POOL_MAX 0
+#define HTP_POOL_AVG 1
+
+// Fast path: exact non-overlapping tiling (stride == kernel, no padding), kernel_x in {1,2}.
+// Every window is guaranteed fully in-bounds, so this never needs boundary clamping.
+static void pool_plane_hvx(
+    const float * src, float * dst, const struct htp_pool_2d_kernel_params * p) {
+    const bool is_max = (p->pool_op == HTP_POOL_MAX);
+    const HVX_Vector scale = hvx_vec_splat_f32(p->inv_kernel_area);
+    const HVX_Vector seed  = is_max ? hvx_vec_splat_f32(-FLT_MAX) : Q6_V_vsplat_R(0);
+
+    for (uint32_t oy = 0; oy < p->dst_y; ++oy) {
+        const uint32_t sy = oy * p->kernel_y;
+        for (uint32_t ox = 0; ox < p->dst_x; ox += VLEN_FP32) {
+            const uint32_t rem = p->dst_x - ox;
+            const uint32_t nbytes = (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN;
+            HVX_Vector acc = seed;
+            for (uint32_t ky = 0; ky < p->kernel_y; ++ky) {
+                const float * row = src + (sy + ky) * p->src_x;
+                if (p->kernel_x == 2) {
+                    const HVX_Vector v0 = *(const HVX_UVector *) (row + ox * 2);
+                    const HVX_Vector v1 = *(const HVX_UVector *) (row + ox * 2 + VLEN_FP32);
+                    const HVX_VectorPair deinterleaved = Q6_W_vdeal_VVR(v1, v0, -4);
+                    const HVX_Vector lo = Q6_V_lo_W(deinterleaved);
+                    const HVX_Vector hi = Q6_V_hi_W(deinterleaved);
+                    acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, Q6_Vsf_vmax_VsfVsf(lo, hi))
+                                 : hvx_vec_add_f32_f32(acc, hvx_vec_add_f32_f32(lo, hi));
+                } else if (p->kernel_x == 1) {
+                    const HVX_Vector v = *(const HVX_UVector *) (row + ox);
+                    acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, v) : hvx_vec_add_f32_f32(acc, v);
+                }
+            }
+            hvx_vec_store_u(dst + oy * p->dst_x + ox, nbytes, is_max ? acc : hvx_vec_mul_f32_f32(acc, scale));
+        }
+    }
+}
+
+// Narrow exact-tiling path. The input is staged in VTCM with one vector of
+// guard space, so full-width loads are safe even when src_x is below 32.
+static void pool_plane_hvx_narrow(
+    const float * src, float * dst, const struct htp_pool_2d_kernel_params * p) {
+    const bool is_max = (p->pool_op == HTP_POOL_MAX);
+    const HVX_Vector scale = hvx_vec_splat_f32(p->inv_kernel_area);
+    const HVX_Vector seed  = is_max ? hvx_vec_splat_f32(-FLT_MAX) : Q6_V_vsplat_R(0);
+
+    for (uint32_t oy = 0; oy < p->dst_y; ++oy) {
+        const uint32_t sy = oy * p->kernel_y;
+        HVX_Vector acc = seed;
+        for (uint32_t ky = 0; ky < p->kernel_y; ++ky) {
+            const float * row = src + (sy + ky) * p->src_x;
+            const HVX_Vector v0 = *(const HVX_UVector *) row;
+            HVX_Vector v = v0;
+            if (p->kernel_x == 2) {
+                // The second vector is zero because a narrow row has fewer
+                // than 32 input elements. The low lanes still contain the
+                // complete even/odd pairs needed by the output.
+                const HVX_Vector zero = Q6_V_vsplat_R(0);
+                const HVX_VectorPair deinterleaved = Q6_W_vdeal_VVR(zero, v0, -4);
+                const HVX_Vector even = Q6_V_lo_W(deinterleaved);
+                const HVX_Vector odd  = Q6_V_hi_W(deinterleaved);
+                v = is_max ? Q6_Vsf_vmax_VsfVsf(even, odd)
+                           : hvx_vec_add_f32_f32(even, odd);
+            }
+            acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, v)
+                         : hvx_vec_add_f32_f32(acc, v);
+        }
+        hvx_vec_store_u(dst + oy * p->dst_x, p->dst_x * sizeof(float),
+                        is_max ? acc : hvx_vec_mul_f32_f32(acc, scale));
+    }
+}
+
+static void pool_plane_global(
+    const float * src, float * dst, const struct htp_pool_2d_kernel_params * p) {
+    const uint32_t n = p->src_x * p->src_y;
+    float val;
+    if (p->pool_op == HTP_POOL_MAX) {
+        val = hvx_reduce_max_f32((const uint8_t *) src, n);
+    } else {
+        val = hvx_reduce_sum_f32((const uint8_t *) src, n) * p->inv_kernel_area;
+    }
+    hvx_vec_store_u(dst, sizeof(float), hvx_vec_splat_f32(val));
+}
+
+static void pool_plane_block(
+    const float * src, float * dst, const struct htp_pool_2d_kernel_params * p) {
+    for (uint32_t oy = 0; oy < p->dst_y; ++oy) {
+        const float * row = src + oy * p->src_x;
+        float * dst_row = dst + oy * p->dst_x;
+        for (uint32_t ox = 0; ox < p->dst_x; ++ox) {
+            const uint8_t * block = (const uint8_t *) (row + ox * p->kernel_x);
+            float val;
+            if (p->pool_op == HTP_POOL_MAX) {
+                val = hvx_reduce_max_f32(block, p->kernel_x);
+            } else {
+                val = hvx_reduce_sum_f32(block, p->kernel_x) * p->inv_kernel_area;
+            }
+            hvx_vec_store_u(dst_row + ox, sizeof(float), hvx_vec_splat_f32(val));
+        }
+    }
+}
+
+// General path: arbitrary kernel/stride/padding
+
+// Vertical padding is uniform across a whole output row (iy0 depends only on oy, not ox),
+// so it collapses to one valid-ky range per row instead of a per-element check.
+static inline void pool_row_bounds_y(
+    const struct htp_pool_2d_kernel_params * p, uint32_t oy, int32_t * iy0, uint32_t * ky_lo, uint32_t * ky_hi) {
+    *iy0 = (int32_t) (oy * p->stride_y) - p->pad_y;
+    const int32_t lo = -(*iy0);
+    const int32_t hi = (int32_t) p->src_y - *iy0;
+    *ky_lo = (uint32_t) MAX(0, lo);
+    *ky_hi = (uint32_t) MAX(0, MIN((int32_t) p->kernel_y, hi));
+}
+
+static inline void pool_pixel_boundary_vec(
+    const float * src, float * dst_row, const struct htp_pool_2d_kernel_params * p,
+    uint32_t ox, int32_t iy0, uint32_t ky_lo, uint32_t ky_hi, bool is_max) {
+    const int32_t ix0 = (int32_t) (ox * p->stride_x) - p->pad_x;
+    const int32_t kx_lo = MAX(0, -ix0);
+    const int32_t kx_hi = MIN((int32_t) p->kernel_x, (int32_t) p->src_x - ix0);
+
+    if (kx_lo >= kx_hi || ky_lo >= ky_hi) {
+        HVX_Vector empty_val = is_max ? hvx_vec_splat_f32(-FLT_MAX) : Q6_V_vsplat_R(0);
+        hvx_vec_store_u(dst_row + ox, sizeof(float), empty_val);
+        return;
+    }
+
+    const uint32_t valid_kx = (uint32_t) (kx_hi - kx_lo);
+    const HVX_Vector mask_identity = is_max ? hvx_vec_splat_f32(-FLT_MAX) : Q6_V_vsplat_R(0);
+
+    HVX_Vector acc = mask_identity;
+    for (uint32_t ky = ky_lo; ky < ky_hi; ++ky) {
+        const float * row = src + (uint32_t) (iy0 + (int32_t) ky) * p->src_x;
+        for (uint32_t k = 0; k < valid_kx; k += VLEN_FP32) {
+            const uint32_t k_rem = valid_kx - k;
+            const uint32_t n = (k_rem < VLEN_FP32) ? k_rem : VLEN_FP32;
+            const HVX_VectorPred q = Q6_Q_vsetq_R(n * sizeof(float));
+            const HVX_Vector raw = *(const HVX_UVector *) (row + ix0 + kx_lo + k);
+            const HVX_Vector v = Q6_V_vmux_QVV(q, raw, mask_identity);
+            acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, v) : hvx_vec_add_f32_f32(acc, v);
+        }
+    }
+
+    HVX_Vector reduced = is_max ? hvx_vec_reduce_max_f32(acc) : hvx_vec_reduce_sum_f32(acc);
+    if (!is_max) {
+        HVX_Vector scale_vec;
+        if (p->avg_divide_count) {
+            const uint32_t count = (ky_hi - ky_lo) * valid_kx;
+            scale_vec = hvx_vec_inverse_f32(hvx_vec_splat_f32((float) count));
+        } else {
+            scale_vec = hvx_vec_splat_f32(p->inv_kernel_area);
+        }
+        reduced = hvx_vec_mul_f32_f32(reduced, scale_vec);
+    }
+    hvx_vec_store_u(dst_row + ox, sizeof(float), reduced);
+}
+
+static inline void pool_row_general_boundary_vec(
+    const float * src, float * dst_row, const struct htp_pool_2d_kernel_params * p,
+    uint32_t ox_start, uint32_t ox_end, int32_t iy0, uint32_t ky_lo, uint32_t ky_hi, bool is_max) {
+    for (uint32_t ox = ox_start; ox < ox_end; ++ox) {
+        pool_pixel_boundary_vec(src, dst_row, p, ox, iy0, ky_lo, ky_hi, is_max);
+    }
+}
+
+// Vectorized interior loop.
+static inline void pool_row_general_vec(
+    const float * src, float * dst_row, const struct htp_pool_2d_kernel_params * p,
+    uint32_t ox_start, uint32_t ox_end, int32_t iy0, uint32_t ky_lo, uint32_t ky_hi, bool is_max) {
+    if (ox_start >= ox_end) {
+        return;
+    }
+
+    if (p->stride_x != 1 && p->stride_x != 2) {
+        pool_row_general_boundary_vec(src, dst_row, p, ox_start, ox_end, iy0, ky_lo, ky_hi, is_max);
+        return;
+    }
+
+    const HVX_Vector scale = hvx_vec_splat_f32(p->inv_kernel_area);
+    const HVX_Vector seed  = is_max ? hvx_vec_splat_f32(-FLT_MAX) : Q6_V_vsplat_R(0);
+
+    for (uint32_t ox = ox_start; ox < ox_end; ox += VLEN_FP32) {
+        const uint32_t rem = ox_end - ox;
+        const uint32_t nbytes = (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN;
+        const int32_t ix0 = (int32_t) (ox * p->stride_x) - p->pad_x;
+        HVX_Vector acc = seed;
+
+        for (uint32_t ky = ky_lo; ky < ky_hi; ++ky) {
+            const float * row = src + (uint32_t) (iy0 + (int32_t) ky) * p->src_x;
+            if (p->stride_x == 1) {
+                for (uint32_t kx = 0; kx < p->kernel_x; ++kx) {
+                    const HVX_Vector v = *(const HVX_UVector *) (row + ix0 + (int32_t) kx);
+                    acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, v) : hvx_vec_add_f32_f32(acc, v);
+                }
+            } else {
+                for (uint32_t kx = 0; kx < p->kernel_x; kx += 2) {
+                    const HVX_Vector v0 = *(const HVX_UVector *) (row + ix0 + (int32_t) kx);
+                    const HVX_Vector v1 = *(const HVX_UVector *) (row + ix0 + (int32_t) kx + VLEN_FP32);
+                    const HVX_VectorPair deinterleaved = Q6_W_vdeal_VVR(v1, v0, -4);
+                    const HVX_Vector lo = Q6_V_lo_W(deinterleaved);
+                    acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, lo) : hvx_vec_add_f32_f32(acc, lo);
+                    if (kx + 1 < p->kernel_x) {
+                        const HVX_Vector hi = Q6_V_hi_W(deinterleaved);
+                        acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, hi) : hvx_vec_add_f32_f32(acc, hi);
+                    }
+                }
+            }
+        }
+        hvx_vec_store_u(dst_row + ox, nbytes, is_max ? acc : hvx_vec_mul_f32_f32(acc, scale));
+    }
+}
+
+static void pool_plane_general(
+    const float * src, float * dst, const struct htp_pool_2d_kernel_params * p) {
+    const bool is_max = (p->pool_op == HTP_POOL_MAX);
+
+    for (uint32_t oy = 0; oy < p->dst_y; ++oy) {
+        int32_t iy0;
+        uint32_t ky_lo, ky_hi;
+        pool_row_bounds_y(p, oy, &iy0, &ky_lo, &ky_hi);
+        float * dst_row = dst + oy * p->dst_x;
+
+        pool_row_general_boundary_vec(src, dst_row, p, 0, p->ox_lo, iy0, ky_lo, ky_hi, is_max);
+        pool_row_general_vec(src, dst_row, p, p->ox_lo, p->ox_hi, iy0, ky_lo, ky_hi, is_max);
+        pool_row_general_boundary_vec(src, dst_row, p, p->ox_hi, p->dst_x, iy0, ky_lo, ky_hi, is_max);
+    }
+}
+
+typedef void (*pool_plane_fn_t)(const float * src, float * dst, const struct htp_pool_2d_kernel_params * p);
+
+struct pool_2d_context {
+    struct htp_ops_context * octx;
+    const struct htp_pool_2d_kernel_params * kparams;
+    pool_plane_fn_t pool_plane;
+    uint32_t n_threads;
+    uint32_t plane_start;
+    uint32_t plane_count;
+    uint32_t planes_per_thread;
+};
+
+static void pool_2d_thread(unsigned int nth, unsigned int ith, void * data) {
+    struct pool_2d_context * ctx = (struct pool_2d_context *) data;
+    const struct htp_pool_2d_kernel_params * p = ctx->kparams;
+    const struct htp_tensor * src0 = ctx->octx->src[0];
+    const struct htp_tensor * dst = ctx->octx->dst;
+    pool_plane_fn_t pool_plane = ctx->pool_plane;
+    const uint32_t planes_per_thread = ctx->planes_per_thread;
+    const uint32_t first = ctx->plane_start + ith * planes_per_thread;
+    const uint32_t last = MIN(first + planes_per_thread, ctx->plane_start + ctx->plane_count);
+
+    if (first >= last) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &ctx->octx->ctx->trace[ith];
+    dma_queue * dma_queue = ctx->octx->ctx->dma[ith];
+
+    const uint32_t src_spad_half = p->src_plane_bytes_aligned;
+    const uint32_t dst_spad_half = p->dst_plane_bytes_aligned;
+    const uint32_t src_bytes_per_thread = 2 * src_spad_half;
+    const uint32_t dst_bytes_per_thread = 2 * dst_spad_half;
+    const size_t off_dst = (size_t) ctx->n_threads * src_bytes_per_thread;
+
+    uint8_t * vtcm_base = (uint8_t *) ctx->octx->ctx->vtcm_base;
+    uint8_t * src_spad  = vtcm_base + ith * src_bytes_per_thread;
+    uint8_t * dst_spad  = vtcm_base + off_dst + ith * dst_bytes_per_thread;
+
+    float * srcb2[2] = { (float *) src_spad, (float *) (src_spad + src_spad_half) };
+    float * dstb2[2] = { (float *) dst_spad, (float *) (dst_spad + dst_spad_half) };
+
+    const uint32_t total = last - first;
+
+    // Warm up the pipeline: push up to 2 initial (dummy dst, src) transfer pairs.
+    for (uint32_t i = 0; i < total && i < 2; ++i) {
+        dma_queue_push(dma_queue,
+                       dma_make_data(dst->data, dstb2[i]),
+                       p->dst_plane_bytes, dst_spad_half,
+                       p->dst_plane_bytes, 0);
+
+        const dma_addr_t src_addr = src0->data + (first + i) * p->src_plane_bytes;
+        dma_queue_push(dma_queue,
+                       dma_make_data(srcb2[i], src_addr),
+                       src_spad_half, p->src_plane_bytes,
+                       p->src_plane_bytes, 1);
+    }
+
+    for (uint32_t i = 0; i < total; ++i) {
+        const uint32_t plane = first + i;
+        const uint32_t buf   = i & 1u;
+        float * srcb = srcb2[buf];
+        float * dstb = dstb2[buf];
+
+        dma_queue_pop(dma_queue); // dst writeback from plane i - 2 (or dummy on iter 0, 1)
+        dma_queue_pop(dma_queue); // input for plane i
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) plane);
+        pool_plane(srcb, dstb, p);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) plane);
+
+        const dma_addr_t dst_addr = dst->data + plane * p->dst_plane_bytes;
+        dma_queue_push(dma_queue,
+                       dma_make_data(dst_addr, dstb),
+                       p->dst_plane_bytes, dst_spad_half,
+                       p->dst_plane_bytes, 1);
+
+        if (i + 2 < total) {
+            const dma_addr_t next_src_addr = src0->data + (plane + 2) * p->src_plane_bytes;
+            dma_queue_push(dma_queue,
+                           dma_make_data(srcb, next_src_addr),
+                           src_spad_half, p->src_plane_bytes,
+                           p->src_plane_bytes, 1);
+        }
+    }
+
+    dma_queue_flush(dma_queue);
+
+    FARF(HIGH, "pool2d-f32-dma %d/%d: %ux%ux%ux%u -> %ux%ux%ux%u (%u:%u)\n",
+         ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
+         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], first, last);
+    (void) nth;
+}
+
+int op_pool_2d(struct htp_ops_context * octx) {
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst = octx->dst;
+    const struct htp_pool_2d_kernel_params * p =
+        (const struct htp_pool_2d_kernel_params *) octx->kernel_params;
+
+    if (src0->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_F32 ||
+        (p->pool_op != HTP_POOL_AVG && p->pool_op != HTP_POOL_MAX)) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    uint32_t plane_start = 0;
+    uint32_t plane_count = p->planes;
+    if (octx->ctx->mdev.count > 1) {
+        const uint32_t planes_per_chunk = (p->dst_plane_bytes > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(p->dst_plane_bytes, HEX_L2_LINE_SIZE)) : 1;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+            plane_count, htp_tensor_mdev_data_aligned(dst) ? planes_per_chunk : 0,
+            octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        plane_start = range.start;
+        plane_count = range.count;
+    }
+    if (plane_count == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = MIN(p->n_threads, plane_count);
+    if (!htp_ops_context_set_n_threads(octx, n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    const uint32_t planes_per_thread = fastdiv(plane_count + n_threads - 1, &octx->n_threads_div);
+
+    pool_plane_fn_t pool_plane;
+    if (p->global_path) {
+        pool_plane = pool_plane_global;
+    } else if (p->block_path) {
+        pool_plane = pool_plane_block;
+    } else if (p->narrow_path) {
+        pool_plane = pool_plane_hvx_narrow;
+    } else if (p->fast_path) {
+        pool_plane = pool_plane_hvx;
+    } else {
+        pool_plane = pool_plane_general;
+    }
+
+    struct pool_2d_context ctx = {
+        .octx = octx,
+        .kparams = p,
+        .pool_plane = pool_plane,
+        .n_threads = n_threads,
+        .plane_start = plane_start,
+        .plane_count = plane_count,
+        .planes_per_thread = planes_per_thread,
+    };
+    work_queue_run(octx->ctx->work_queue, pool_2d_thread, &ctx, n_threads);
+    return HTP_STATUS_OK;
+}
+
+int op_pool_1d(struct htp_ops_context * octx) {
+    return op_pool_2d(octx);
+}

ggml/src/ggml-hexagon/htp/pool-ops.h

@@ -0,0 +1,85 @@
+#ifndef HTP_POOL_OPS_H
+#define HTP_POOL_OPS_H
+
+#include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
+#include <string.h>
+
+#include "hex-common.h"
+
+struct htp_pool_2d_kernel_params {
+    uint32_t src_x;
+    uint32_t src_y;
+    uint32_t dst_x;
+    uint32_t dst_y;
+    uint32_t kernel_x;
+    uint32_t kernel_y;
+    uint32_t stride_x;
+    uint32_t stride_y;
+    int32_t  pad_x;
+    int32_t  pad_y;
+    uint32_t src_plane_bytes;
+    uint32_t dst_plane_bytes;
+    uint32_t src_plane_bytes_aligned;
+    uint32_t dst_plane_bytes_aligned;
+    uint32_t n_threads;
+    uint32_t planes;
+    uint32_t pool_op;
+    uint32_t fast_path;
+    uint32_t narrow_path;
+    uint32_t global_path;
+    uint32_t block_path;
+    uint32_t avg_divide_count;
+    uint32_t ox_lo;
+    uint32_t ox_hi;
+    float    inv_kernel_area;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_pool_2d_kernel_params) <= 128, "htp_pool_2d_kernel_params is too large");
+#else
+_Static_assert(sizeof(struct htp_pool_2d_kernel_params) <= 128, "htp_pool_2d_kernel_params is too large");
+#endif
+
+struct htp_pool_vtcm_layout {
+    size_t   total_bytes;
+    size_t   off_src;
+    size_t   off_dst;
+    size_t   src_bytes_per_thread;
+    size_t   dst_bytes_per_thread;
+    size_t   src_spad_half_size;
+    size_t   dst_spad_half_size;
+};
+
+static inline bool htp_pool_solve_layout(
+    struct htp_pool_vtcm_layout * layout,
+    uint32_t src_x,
+    uint32_t src_y,
+    uint32_t dst_x,
+    uint32_t dst_y,
+    uint32_t n_threads,
+    size_t   vtcm_budget
+) {
+    // Full-plane double buffering (256 bytes guard space for vector loads)
+    const size_t src_plane_bytes   = (size_t) src_x * src_y * sizeof(float);
+    const size_t dst_plane_bytes   = (size_t) dst_x * dst_y * sizeof(float);
+    const size_t src_plane_aligned = hex_round_up((uint32_t) src_plane_bytes + 256, 128);
+    const size_t dst_plane_aligned = hex_round_up((uint32_t) dst_plane_bytes, 128);
+
+    const size_t spad_per_thread = 2 * (src_plane_aligned + dst_plane_aligned);
+    if (spad_per_thread * n_threads > vtcm_budget) {
+        return false;
+    }
+
+    layout->src_spad_half_size   = src_plane_aligned;
+    layout->dst_spad_half_size   = dst_plane_aligned;
+    layout->src_bytes_per_thread = 2 * src_plane_aligned;
+    layout->dst_bytes_per_thread = 2 * dst_plane_aligned;
+    layout->off_src              = 0;
+    layout->off_dst              = layout->src_bytes_per_thread * n_threads;
+    layout->total_bytes          = layout->off_dst + layout->dst_bytes_per_thread * n_threads;
+    return true;
+}
+
+#endif // HTP_POOL_OPS_H

ggml/src/ggml-hexagon/htp/ssm-conv.c

@@ -135,48 +135,49 @@ static inline void hvx_ssm_conv_unpack_to_T(const float * raw, float * T, uint32
     }
 }
 
-// HVX 32x32 src0 transpose for prefill: src0 {tile_n, ncs} (VTCM) -> src0_T {ncs, d_inner_tile} (VTCM)
-static inline void transpose_src0_block(const float * src0_block,
-                                        uint32_t      ncs,
-                                        uint32_t      cb_n,
-                                        uint32_t      d_inner_tile,
-                                        float *       src0_T_block_dst,
-                                        uint32_t      cb) {
-    const uint32_t T_TILE = VLEN_FP32;
+// Decode dot product specialization for d_conv == 4: multiply in the raw channel-major layout,
+// then deinterleave the products so each vector holds one tap of 32 channels, and sum.
+// Keeps both operands in DMA layout - no transpose, no scratch.
+static inline void hvx_ssm_conv_decode_4(const float * x, const float * w, float * out, uint32_t n_ch) {
+    for (uint32_t cb = 0; cb < n_ch; cb += VLEN_FP32) {
+        const float * xp = x + cb * 4;
+        const float * wp = w + cb * 4;
 
-    HVX_Vector __attribute__((aligned(VLEN))) sub[32];
+        HVX_Vector p0 = Q6_Vqf32_vmpy_VsfVsf(*(const HVX_Vector *)(xp +  0), *(const HVX_Vector *)(wp +  0));
+        HVX_Vector p1 = Q6_Vqf32_vmpy_VsfVsf(*(const HVX_Vector *)(xp + 32), *(const HVX_Vector *)(wp + 32));
+        HVX_Vector p2 = Q6_Vqf32_vmpy_VsfVsf(*(const HVX_Vector *)(xp + 64), *(const HVX_Vector *)(wp + 64));
+        HVX_Vector p3 = Q6_Vqf32_vmpy_VsfVsf(*(const HVX_Vector *)(xp + 96), *(const HVX_Vector *)(wp + 96));
 
-    for (uint32_t t0 = 0; t0 < ncs; t0 += T_TILE) {
-        const uint32_t t_n = MIN(T_TILE, ncs - t0);
+        HVX_VectorPair p01 = Q6_W_vdeal_VVR(p1, p0, -4);
+        HVX_VectorPair p23 = Q6_W_vdeal_VVR(p3, p2, -4);
 
-        uint32_t __attribute__((aligned(VLEN))) mask_buf[VLEN_FP32] = { 0 };
-        for (uint32_t k = 0; k < t_n; ++k) {
-            mask_buf[k] = 0xFFFFFFFF;
-        }
-        const HVX_Vector mask = *(const HVX_Vector *) mask_buf;
+        HVX_VectorPair q02 = Q6_W_vdeal_VVR(Q6_V_lo_W(p23), Q6_V_lo_W(p01), -4);
+        HVX_VectorPair q13 = Q6_W_vdeal_VVR(Q6_V_hi_W(p23), Q6_V_hi_W(p01), -4);
 
-        for (uint32_t r = 0; r < cb_n; ++r) {
-            const float * src_row = src0_block + r * ncs + t0;
-            sub[r] = (t_n == T_TILE) ? *(const HVX_UVector *) src_row : Q6_V_vand_VV(*(const HVX_UVector *) src_row, mask);
-        }
-        for (uint32_t r = cb_n; r < T_TILE; ++r) {
-            sub[r] = hvx_vec_splat_f32(0.0f);
-        }
+        HVX_Vector a = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(q02), Q6_V_lo_W(q13));
+        HVX_Vector b = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_hi_W(q02), Q6_V_hi_W(q13));
 
-        hvx_transpose_32x32_f32(sub);
-
-        for (uint32_t r = 0; r < t_n; ++r) {
-            float * dst = src0_T_block_dst + (t0 + r) * d_inner_tile + cb;
-            if (cb_n == T_TILE) {
-                *(HVX_UVector *) dst = sub[r];
-            } else {
-                hvx_vec_store_u(dst, cb_n * sizeof(float), sub[r]);
-            }
-        }
+        *(HVX_Vector *)(out + cb) = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(a, b));
     }
 }
 
-// Single-row decode worker (n_t == 1)
+// Transpose src0 for prefill: one 32-channel block {32, ncs} (VTCM) -> T {ncs, 32} (VTCM).
+// One VTCM gather per output row: lane r picks channel cb+r, the region bound drops
+// the lanes past the channel tail.
+static inline void hvx_ssm_conv_transpose_block(const float * raw_block,
+                                                float *       T,
+                                                uint32_t      ncs,
+                                                uint32_t      cb_n,
+                                                HVX_Vector    vv) {
+    const size_t   base = (size_t) raw_block;
+    const uint32_t mu   = cb_n * ncs * sizeof(float) - 1;
+
+    for (uint32_t t = 0; t < ncs; ++t) {
+        Q6_vgather_ARMVw((HVX_Vector *) (T + (size_t) t * VLEN_FP32), base + t * sizeof(float), mu, vv);
+    }
+}
+
+// Single-token decode worker (n_t == 1)
 static void ssm_conv_thread_f32_decode(unsigned int nth, unsigned int ith, void * data) {
     struct htp_ssm_conv_context *             scctx   = (struct htp_ssm_conv_context *) data;
     struct htp_ops_context *                  octx    = scctx->octx;
@@ -202,9 +203,11 @@ static void ssm_conv_thread_f32_decode(unsigned int nth, unsigned int ith, void
 
     const uint32_t d_inner_per_thread = ir1 - ir0;
     const uint32_t d_inner_stride     = hex_round_up(d_inner_per_thread, VLEN_FP32);
+    const uint32_t d_inner_tile       = scctx->d_inner_tile;
 
-    const size_t src0_stride_seq_bytes = src0->nb[2];
-    const size_t dst_stride_seq_bytes  = dst->nb[2];
+    const size_t src0_stride_inner_bytes = src0->nb[1];
+    const size_t src0_stride_seq_bytes   = src0->nb[2];
+    const size_t dst_stride_seq_bytes    = dst->nb[2];
 
     uint8_t * src1_spad_base = octx->src1_spad.data + ith * octx->src1_spad.size_per_thread;
     uint8_t * src0_spad_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread;
@@ -216,57 +219,121 @@ static void ssm_conv_thread_f32_decode(unsigned int nth, unsigned int ith, void
     float * src1_raw = (float *) src1_spad_base;
     float * src1_T   = (float *) (src1_spad_base + weight_raw_size);
 
-    float * src0_raw = (float *) src0_spad_base;
-    float * src0_T   = (float *) (src0_spad_base + weight_raw_size);
+    const size_t src0_tile_raw_bytes = hex_round_up(d_inner_tile * d_conv * sizeof(float), 128);
+    const size_t dst_tile_bytes      = hex_round_up(d_inner_tile * sizeof(float), 128);
 
-    float * dst_spad = (float *) dst_spad_base;
+    float * src0_tile_raw[2] = { (float *) src0_spad_base, (float *) (src0_spad_base + src0_tile_raw_bytes) };
+    float * src0_T           = (float *) (src0_spad_base + 2 * src0_tile_raw_bytes);
+
+    float * dst_tile[2] = { (float *) dst_spad_base, (float *) (dst_spad_base + dst_tile_bytes) };
 
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
-    // 1. Fetch weights src1 from DDR into VTCM via DMA (DMA64-safe)
+    // raw_dot keeps both operands in the DMA layout, so src1 needs no prep pass
+    const bool raw_dot = (d_conv == 4) && (d_inner_per_thread % VLEN_FP32 == 0);
+
+    const uint32_t n_tiles   = (d_inner_per_thread + d_inner_tile - 1) / d_inner_tile;
+    const uint32_t n_chunks  = n_s * n_tiles;
+    const size_t   row_bytes = d_conv * sizeof(float);
+
+    uint32_t s_fetch        = 0;
+    uint32_t tile_off_fetch = 0;
+
+    #define SSM_CONV_DECODE_PUSH_FETCH(c)                                                         \
+        do {                                                                                      \
+            const uint32_t   cur_tile_n = MIN(d_inner_tile, d_inner_per_thread - tile_off_fetch); \
+            const dma_addr_t fetch_ddr  = src0->data + s_fetch * src0_stride_seq_bytes +          \
+                                          (ir0 + tile_off_fetch) * src0_stride_inner_bytes;       \
+            dma_queue_push(dma_q,                                                                 \
+                           dma_make_data((uint8_t *) src0_tile_raw[(c) & 1], fetch_ddr),          \
+                           row_bytes, src0_stride_inner_bytes, row_bytes, cur_tile_n);            \
+            tile_off_fetch += d_inner_tile;                                                       \
+            if (tile_off_fetch >= d_inner_per_thread) {                                           \
+                tile_off_fetch = 0;                                                               \
+                s_fetch++;                                                                        \
+            }                                                                                     \
+        } while (0)
+
+    // Queue weights and initial input tiles together so DDR reads overlap
     const dma_addr_t src1_ddr = src1->data + ir0 * d_conv * sizeof(float);
     dma_queue_push(dma_q, dma_make_data((uint8_t *) src1_raw, src1_ddr), weight_bytes, weight_bytes, weight_bytes, 1);
-    dma_queue_pop(dma_q);
 
-    // 2. Unpack/transpose src1_raw into src1_T {d_conv, d_inner_stride}
-    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
-    hvx_ssm_conv_unpack_to_T(src1_raw, src1_T, d_inner_per_thread, d_inner_stride, d_conv);
-    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
+    SSM_CONV_DECODE_PUSH_FETCH(0);
+    if (n_chunks > 1) {
+        SSM_CONV_DECODE_PUSH_FETCH(1);
+    }
 
-    const size_t input_bytes  = (size_t) d_inner_per_thread * d_conv * sizeof(float);
-    const size_t output_bytes = (size_t) d_inner_per_thread * sizeof(float);
+    dma_queue_pop(dma_q);  // weights
 
-    // 3. Process each sequence
-    for (uint32_t s = 0; s < n_s; ++s) {
-        const dma_addr_t src0_ddr = src0->data + s * src0_stride_seq_bytes + ir0 * d_conv * sizeof(float);
-        dma_queue_push(dma_q, dma_make_data((uint8_t *) src0_raw, src0_ddr), input_bytes, input_bytes, input_bytes, 1);
-        dma_queue_pop(dma_q);
+    if (!raw_dot) {
+        // Unpack/transpose src1_raw into src1_T {d_conv, d_inner_stride}
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_W_PREP, (uint16_t) ir0);
+        hvx_ssm_conv_unpack_to_T(src1_raw, src1_T, d_inner_per_thread, d_inner_stride, d_conv);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_W_PREP, (uint16_t) ir0);
+    }
 
-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) s);
-        hvx_ssm_conv_unpack_to_T(src0_raw, src0_T, d_inner_per_thread, d_inner_stride, d_conv);
+    uint32_t i3       = 0;
+    uint32_t tile_off = 0;
 
-        for (uint32_t cb = 0; cb < d_inner_per_thread; cb += VLEN_FP32) {
-            const uint32_t cb_n = MIN(VLEN_FP32, d_inner_per_thread - cb);
-            HVX_Vector acc = hvx_vec_splat_f32(0.0f);
-            for (uint32_t j = 0; j < d_conv; ++j) {
-                HVX_Vector x = *(const HVX_Vector *)(src0_T + j * d_inner_stride + cb);
-                HVX_Vector w = *(const HVX_Vector *)(src1_T + j * d_inner_stride + cb);
-                acc          = Q6_Vqf32_vadd_Vqf32Vqf32(acc, Q6_Vqf32_vmpy_VsfVsf(x, w));
-            }
-            HVX_Vector y = Q6_Vsf_equals_Vqf32(acc);
-            if (cb_n == VLEN_FP32) {
-                *(HVX_Vector *)(dst_spad + cb) = y;
-            } else {
-                hvx_vec_store_u(dst_spad + cb, cb_n * sizeof(float), y);
+    for (uint32_t c = 0; c < n_chunks; ++c) {
+        const uint32_t tile_n = MIN(d_inner_tile, d_inner_per_thread - tile_off);
+
+        if (c >= 2) {
+            dma_queue_pop(dma_q);  // writeback of chunk c-2, frees dst_tile[c & 1]
+        }
+        dma_queue_pop(dma_q);      // fetch chunk c
+
+        float * restrict out = dst_tile[c & 1];
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i3);
+        if (raw_dot) {
+            const float * xp = src0_tile_raw[c & 1];
+            const float * wp = src1_raw + tile_off * 4;
+            hvx_ssm_conv_decode_4(xp, wp, out, tile_n);
+        } else {
+            const uint32_t tile_stride = hex_round_up(tile_n, VLEN_FP32);
+            hvx_ssm_conv_unpack_to_T(src0_tile_raw[c & 1], src0_T, tile_n, tile_stride, d_conv);
+
+            for (uint32_t cb = 0; cb < tile_n; cb += VLEN_FP32) {
+                const uint32_t cb_n = MIN(VLEN_FP32, tile_n - cb);
+                HVX_Vector acc = hvx_vec_splat_f32(0.0f);
+                for (uint32_t j = 0; j < d_conv; ++j) {
+                    HVX_Vector x = *(const HVX_Vector *)(src0_T + j * tile_stride + cb);
+                    HVX_Vector w = *(const HVX_Vector *)(src1_T + j * d_inner_stride + tile_off + cb);
+                    acc          = Q6_Vqf32_vadd_Vqf32Vqf32(acc, Q6_Vqf32_vmpy_VsfVsf(x, w));
+                }
+                HVX_Vector y = Q6_Vsf_equals_Vqf32(acc);
+                if (cb_n == VLEN_FP32) {
+                    *(HVX_Vector *)(out + cb) = y;
+                } else {
+                    hvx_vec_store_u(out + cb, cb_n * sizeof(float), y);
+                }
             }
         }
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) s);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i3);
 
-        const dma_addr_t dst_ddr = dst->data + s * dst_stride_seq_bytes + ir0 * sizeof(float);
-        dma_queue_push(dma_q, dma_make_data(dst_ddr, (uint8_t *) dst_spad), output_bytes, output_bytes, output_bytes, 1);
-        dma_queue_pop(dma_q);
+        const dma_addr_t dst_ddr        = dst->data + i3 * dst_stride_seq_bytes + (ir0 + tile_off) * sizeof(float);
+        const size_t     tile_out_bytes = (size_t) tile_n * sizeof(float);
+        dma_queue_push(dma_q, dma_make_data(dst_ddr, (uint8_t *) out),
+                       tile_out_bytes, tile_out_bytes, tile_out_bytes, 1);
+
+        if (c + 2 < n_chunks) {
+            SSM_CONV_DECODE_PUSH_FETCH(c + 2);
+        }
+
+        tile_off += d_inner_tile;
+        if (tile_off >= d_inner_per_thread) {
+            tile_off = 0;
+            i3++;
+        }
     }
 
+    for (uint32_t k = MIN(n_chunks, 2); k > 0; --k) {
+        dma_queue_pop(dma_q);  // drain the last writebacks
+    }
+
+    #undef SSM_CONV_DECODE_PUSH_FETCH
+
     FARF(HIGH, "ssm-conv-f32-decode %d/%d: %ux%ux%ux%u (%u:%u) * %ux%ux%ux%u -> %ux%ux%ux%u\n",
          ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1,
          src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0], dst->ne[1],
@@ -318,82 +385,157 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
     float * src1_raw = (float *) src1_spad_base;
     float * src1_T   = (float *) (src1_spad_base + weight_raw_size);
 
+    // src0 spad holds two raw tiles (fetch of tile n+1 overlaps compute of tile n) plus
+    // the transposed block. dst spad holds two tiles so a writeback can stay in flight.
     const size_t src0_tile_raw_bytes = hex_round_up(d_inner_tile * ncs * sizeof(float), 128);
-    float * src0_tile_raw = (float *) src0_spad_base;
-    float * src0_T        = (float *) (src0_spad_base + src0_tile_raw_bytes);
+    const size_t dst_tile_bytes      = hex_round_up(d_inner_tile * n_t * sizeof(float), 128);
 
-    float * dst_tile = (float *) dst_spad_base;
+    float * src0_tile_raw[2] = { (float *) src0_spad_base, (float *) (src0_spad_base + src0_tile_raw_bytes) };
+    float * src0_T           = (float *) (src0_spad_base + 2 * src0_tile_raw_bytes);
+
+    float * dst_tile[2] = { (float *) dst_spad_base, (float *) (dst_spad_base + dst_tile_bytes) };
 
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
-    // 1. Fetch weights src1 from DDR into VTCM via DMA (DMA64-safe)
+    const uint32_t n_tiles   = (d_inner_per_thread + d_inner_tile - 1) / d_inner_tile;
+    const uint32_t n_chunks  = n_s * n_tiles;
+    const size_t   row_bytes = ncs * sizeof(float);
+
+    uint32_t s_fetch        = 0;
+    uint32_t tile_off_fetch = 0;
+
+    // Chunk c fetches into src0_tile_raw[c & 1] and writes back from dst_tile[c & 1].
+    // Two fetches run ahead, so the queue order is F0 F1 W0 F2 W1 ... and pops follow it.
+    #define SSM_CONV_PUSH_FETCH(c)                                                                \
+        do {                                                                                      \
+            const uint32_t   cur_tile_n = MIN(d_inner_tile, d_inner_per_thread - tile_off_fetch); \
+            const dma_addr_t fetch_ddr  = src0->data + s_fetch * src0_stride_seq_bytes +          \
+                                          (ir0 + tile_off_fetch) * src0_stride_inner_bytes;       \
+            dma_queue_push(dma_q,                                                                 \
+                           dma_make_data((uint8_t *) src0_tile_raw[(c) & 1], fetch_ddr),          \
+                           row_bytes, src0_stride_inner_bytes, row_bytes, cur_tile_n);            \
+            tile_off_fetch += d_inner_tile;                                                       \
+            if (tile_off_fetch >= d_inner_per_thread) {                                           \
+                tile_off_fetch = 0;                                                               \
+                s_fetch++;                                                                        \
+            }                                                                                     \
+        } while (0)
+
+    // Queue weights and initial input tiles together so DDR reads overlap
     const dma_addr_t src1_ddr = src1->data + ir0 * d_conv * sizeof(float);
     dma_queue_push(dma_q, dma_make_data((uint8_t *) src1_raw, src1_ddr), weight_bytes, weight_bytes, weight_bytes, 1);
-    dma_queue_pop(dma_q);
 
-    // 2. Unpack/transpose src1_raw into src1_T {d_conv, d_inner_stride}
-    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
+    SSM_CONV_PUSH_FETCH(0);
+    if (n_chunks > 1) {
+        SSM_CONV_PUSH_FETCH(1);
+    }
+
+    dma_queue_pop(dma_q);  // weights
+
+    // Unpack/transpose src1_raw into src1_T {d_conv, d_inner_stride}
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_W_PREP, (uint16_t) ir0);
     hvx_ssm_conv_unpack_to_T(src1_raw, src1_T, d_inner_per_thread, d_inner_stride, d_conv);
-    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_W_PREP, (uint16_t) ir0);
 
     const uint32_t C_TILE = VLEN_FP32;
 
-    for (uint32_t i3 = 0; i3 < n_s; ++i3) {
-        for (uint32_t tile_off = 0; tile_off < d_inner_per_thread; tile_off += d_inner_tile) {
-            const uint32_t tile_n = MIN(d_inner_tile, d_inner_per_thread - tile_off);
+    // gather offsets: lane r reads channel r of the raw tile block
+    uint32_t __attribute__((aligned(VLEN))) gather_off[VLEN_FP32];
+    for (uint32_t r = 0; r < VLEN_FP32; ++r) {
+        gather_off[r] = r * ncs * sizeof(float);
+    }
+    const HVX_Vector vv = *(const HVX_Vector *) gather_off;
 
-            // Fetch src0 chunk from DDR to VTCM via 2D DMA
-            const dma_addr_t src0_tile_ddr = src0->data +
-                i3 * src0_stride_seq_bytes +
-                (ir0 + tile_off) * src0_stride_inner_bytes;
-            const size_t row_bytes = ncs * sizeof(float);
+    uint32_t i3       = 0;
+    uint32_t tile_off = 0;
 
-            dma_queue_push(dma_q, dma_make_data((uint8_t *) src0_tile_raw, src0_tile_ddr),
-                           row_bytes, src0_stride_inner_bytes, row_bytes, tile_n);
-            dma_queue_pop(dma_q);
+    for (uint32_t c = 0; c < n_chunks; ++c) {
+        const uint32_t tile_n = MIN(d_inner_tile, d_inner_per_thread - tile_off);
+
+        const float * restrict raw = src0_tile_raw[c & 1];
+        float * restrict       out = dst_tile[c & 1];
+
+        if (c >= 2) {
+            dma_queue_pop(dma_q);  // writeback of chunk c-2, frees dst_tile[c & 1]
+        }
+        dma_queue_pop(dma_q);      // fetch chunk c
+
+        // Channel block outer, token inner: the taps and the sliding window stay in
+        // registers, so a new output row costs one src0_T load.
+        const uint32_t dst_tile_stride = hex_round_up(tile_n, C_TILE);
+
+        for (uint32_t cb = 0; cb < tile_n; cb += C_TILE) {
+            const uint32_t cb_n = MIN(C_TILE, tile_n - cb);
+
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, (uint16_t) tile_off);
+            hvx_ssm_conv_transpose_block(raw + (size_t) cb * ncs, src0_T, ncs, cb_n, vv);
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, (uint16_t) tile_off);
+
+            const float * restrict wp = src1_T + tile_off + cb;
+            float * restrict       op = out + cb;
 
-            // Transpose src0 chunk in VTCM into {d_inner_tile, ncs} layout
             htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) tile_off);
-            for (uint32_t cb = 0; cb < tile_n; cb += C_TILE) {
-                const uint32_t cb_n = MIN(C_TILE, tile_n - cb);
-                transpose_src0_block(src0_tile_raw + cb * ncs, ncs, cb_n, d_inner_tile, src0_T, cb);
-            }
+            if (d_conv == 4) {
+                const HVX_Vector w0 = *(const HVX_Vector *) (wp);
+                const HVX_Vector w1 = *(const HVX_Vector *) (wp + d_inner_stride);
+                const HVX_Vector w2 = *(const HVX_Vector *) (wp + 2 * d_inner_stride);
+                const HVX_Vector w3 = *(const HVX_Vector *) (wp + 3 * d_inner_stride);
 
-            // Compute convolution
-            for (uint32_t t = 0; t < n_t; ++t) {
-                for (uint32_t cb = 0; cb < tile_n; cb += C_TILE) {
-                    const uint32_t cb_n = MIN(C_TILE, tile_n - cb);
+                HVX_Vector x0 = *(const HVX_Vector *) (src0_T);
+                HVX_Vector x1 = *(const HVX_Vector *) (src0_T + C_TILE);
+                HVX_Vector x2 = *(const HVX_Vector *) (src0_T + 2 * C_TILE);
 
+                for (uint32_t t = 0; t < n_t; ++t) {
+                    const HVX_Vector x3 = *(const HVX_Vector *) (src0_T + (t + 3) * C_TILE);
+
+                    HVX_Vector a = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_Vqf32_vmpy_VsfVsf(x0, w0),
+                                                            Q6_Vqf32_vmpy_VsfVsf(x1, w1));
+                    HVX_Vector b = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_Vqf32_vmpy_VsfVsf(x2, w2),
+                                                            Q6_Vqf32_vmpy_VsfVsf(x3, w3));
+
+                    *(HVX_Vector *) (op + t * dst_tile_stride) = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(a, b));
+
+                    x0 = x1;
+                    x1 = x2;
+                    x2 = x3;
+                }
+            } else {
+                for (uint32_t t = 0; t < n_t; ++t) {
                     HVX_Vector acc = hvx_vec_splat_f32(0.0f);
                     for (uint32_t j = 0; j < d_conv; ++j) {
-                        HVX_Vector x = *(const HVX_Vector *) (src0_T + (t + j) * d_inner_tile + cb);
-                        HVX_Vector w = *(const HVX_Vector *) (src1_T + j * d_inner_stride + tile_off + cb);
+                        HVX_Vector x = *(const HVX_Vector *) (src0_T + (t + j) * C_TILE);
+                        HVX_Vector w = *(const HVX_Vector *) (wp + j * d_inner_stride);
                         acc          = Q6_Vqf32_vadd_Vqf32Vqf32(acc, Q6_Vqf32_vmpy_VsfVsf(x, w));
                     }
-
-                    HVX_Vector y = Q6_Vsf_equals_Vqf32(acc);
-                    float * dst_tile_ptr = dst_tile + t * tile_n + cb;
-                    if (cb_n == C_TILE) {
-                        *(HVX_Vector *) dst_tile_ptr = y;
-                    } else {
-                        hvx_vec_store_u(dst_tile_ptr, cb_n * sizeof(float), y);
-                    }
+                    *(HVX_Vector *) (op + t * dst_tile_stride) = Q6_Vsf_equals_Vqf32(acc);
                 }
             }
             htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) tile_off);
+        }
 
-            // Writeback dst_tile from VTCM to DDR via 2D DMA
-            const dma_addr_t dst_tile_ddr = dst->data +
-                i3 * dst_stride_seq_bytes +
-                (ir0 + tile_off) * sizeof(float);
-            const size_t dst_row_bytes = tile_n * sizeof(float);
+        // Writeback dst_tile VTCM -> DDR via 2D DMA
+        dma_queue_push(dma_q,
+                       dma_make_data(dst->data + i3 * dst_stride_seq_bytes + (ir0 + tile_off) * sizeof(float),
+                                     (uint8_t *) out),
+                       dst_stride_token_bytes, dst_tile_stride * sizeof(float), tile_n * sizeof(float), n_t);
 
-            dma_queue_push(dma_q, dma_make_data(dst_tile_ddr, (uint8_t *) dst_tile),
-                           dst_stride_token_bytes, dst_row_bytes, dst_row_bytes, n_t);
-            dma_queue_pop(dma_q);
+        if (c + 2 < n_chunks) {
+            SSM_CONV_PUSH_FETCH(c + 2);
+        }
+
+        tile_off += d_inner_tile;
+        if (tile_off >= d_inner_per_thread) {
+            tile_off = 0;
+            i3++;
         }
     }
 
+    for (uint32_t k = MIN(n_chunks, 2); k > 0; --k) {
+        dma_queue_pop(dma_q);  // drain the last writebacks
+    }
+
+    #undef SSM_CONV_PUSH_FETCH
+
     FARF(HIGH, "ssm-conv-f32-prefill %d/%d: %ux%ux%ux%u (%u:%u) * %ux%ux%ux%u -> %ux%ux%ux%u\n",
          ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1,
          src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0], dst->ne[1],

ggml/src/ggml-hexagon/htp/ssm-conv.h

@@ -12,13 +12,8 @@ struct htp_ssm_conv_kernel_params {
     uint32_t d_inner;
     uint32_t n_t;
     uint32_t n_s;
-    uint32_t d_inner_per_thread;
     uint32_t d_inner_tile;
 
-    uint32_t src0_row_size_aligned;
-    uint32_t src1_row_size_aligned;
-    uint32_t dst_row_size_aligned;
-
     uint32_t vtcm_src0_size_per_thread;
     uint32_t vtcm_src1_size_per_thread;
     uint32_t vtcm_dst_size_per_thread;
@@ -27,8 +22,6 @@ struct htp_ssm_conv_kernel_params {
     uint32_t vtcm_src1_size;
     uint32_t vtcm_dst_size;
     uint32_t vtcm_size;
-
-    struct fastdiv_values div_n_threads;
 };
 
 #if defined(__cplusplus)

ggml/src/ggml-metal/CMakeLists.txt

@@ -53,6 +53,7 @@ set(METALLIB_KERNEL_SOURCES
     kernels/fa_vec_q5_1.metal
     kernels/fa_vec_q8_0.metal
     kernels/mul_mv.metal
+    kernels/mul_mv_mma.metal
     kernels/mul_mm.metal
     kernels/quantize.metal
     kernels/softmax.metal

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

@@ -48,6 +48,90 @@ bool ggml_metal_op_mul_mat_id_use_mm(const struct ggml_tensor * op, bool has_sim
     return has_simdgroup_mm && ne00 >= 64 && ne21 >= 32;
 }
 
+// the most src1 rows of the few-row MMA kernels
+static constexpr int64_t GGML_METAL_MMA_ROWS_MAX = 16;
+
+// src1 rows per 8x8 simdgroup matrix tile of the few-row MMA kernels
+static constexpr int64_t GGML_METAL_MMA_TILE_ROWS = 8;
+
+// weights per K step of the q5_K and generic few-row MMA kernels
+static constexpr int64_t GGML_METAL_MMA_K_CHUNK = 64;
+
+enum ggml_metal_mma_kind ggml_metal_mul_mv_mma_kind(enum ggml_type type, int rt) {
+    if (type == GGML_TYPE_Q4_0 || (type == GGML_TYPE_Q8_0 && rt == 1)) {
+        return GGML_METAL_MMA_KIND_BLK;
+    }
+    return type == GGML_TYPE_Q5_K ? GGML_METAL_MMA_KIND_Q5_K : GGML_METAL_MMA_KIND_GEN;
+}
+
+int ggml_metal_mul_mv_mma_rt(const struct ggml_tensor * op) {
+    return op->src[1]->ne[1] > GGML_METAL_MMA_TILE_ROWS ? 2 : 1;
+}
+
+static bool ggml_metal_mul_mv_mma_type_supported(enum ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_F32:
+        case GGML_TYPE_F16:
+        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:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
+            return true;
+        default:
+            return false;
+    }
+}
+
+int64_t ggml_metal_mul_mv_mma_k_step(enum ggml_type type, int rt) {
+    if (!ggml_metal_mul_mv_mma_type_supported(type)) {
+        return 0;
+    }
+    return ggml_metal_mul_mv_mma_kind(type, rt) == GGML_METAL_MMA_KIND_BLK ? ggml_blck_size(type) : GGML_METAL_MMA_K_CHUNK;
+}
+
+static bool ggml_metal_mul_mat_mma_type_ok(const struct ggml_tensor * op) {
+    const ggml_tensor * src0 = op->src[0];
+    const int64_t step = ggml_metal_mul_mv_mma_k_step(src0->type, ggml_metal_mul_mv_mma_rt(op));
+
+    return step > 0 && src0->ne[0] % step == 0 && src0->nb[0] == ggml_type_size(src0->type);
+}
+
+// the fewest src1 rows at which the few-row MMA kernels beat the mat-vec kernels (measured on an M3 Ultra)
+static int64_t ggml_metal_mul_mv_mma_rows_min(enum ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_F32:
+            return 6;
+        case GGML_TYPE_F16:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q5_1:
+            return 3;
+        default:
+            return 2;
+    }
+}
+
+bool ggml_metal_op_mul_mat_use_mma(const struct ggml_tensor * op) {
+    const ggml_tensor * src0 = op->src[0];
+    const ggml_tensor * src1 = op->src[1];
+
+    // the batch shape goes into int16 function constants
+    const bool batch_ok = src1->ne[2] <= INT16_MAX && src1->ne[2]/src0->ne[2] <= INT16_MAX && src1->ne[3]/src0->ne[3] <= INT16_MAX;
+
+    return ggml_metal_mul_mat_mma_type_ok(op) && batch_ok &&
+        src1->type == GGML_TYPE_F32 && src1->ne[1] >= ggml_metal_mul_mv_mma_rows_min(src0->type) && src1->ne[1] <= GGML_METAL_MMA_ROWS_MAX &&
+        !ggml_is_transposed(src0) && !ggml_is_transposed(src1) &&
+        src1->nb[0] == sizeof(float) && src1->nb[1] % 16 == 0 && src1->nb[2] % 16 == 0 && src1->nb[3] % 16 == 0;
+}
+
+bool ggml_metal_op_mul_mat_may_use_mma(const struct ggml_tensor * op) {
+    return ggml_metal_mul_mv_mma_type_supported(op->src[0]->type) && op->src[1]->type == GGML_TYPE_F32;
+}
+
 // represents a memory range (i.e. an interval from a starting address p0 to an ending address p1 in a given buffer pb)
 // the type indicates whether it is a source range (i.e. ops read data from it) or a destination range (i.e. ops write data to it)
 struct ggml_mem_range {

ggml/src/ggml-metal/ggml-metal-common.h

@@ -2,6 +2,8 @@
 
 #pragma once
 
+#include "ggml.h"
+
 #include <stdbool.h>
 #include <stddef.h>
 
@@ -42,17 +44,30 @@ bool ggml_mem_ranges_add(ggml_mem_ranges_t mrs, const struct ggml_tensor * tenso
 // - new dst range overlaps with any existing range (src or dst)
 bool ggml_mem_ranges_check(ggml_mem_ranges_t mrs, const struct ggml_tensor * tensor);
 
-// reorder the nodes in the graph to improve concurrency, while respecting fusion
+// reorder the nodes in the graph to improve concurrency, while respecting the fusions of a device with props
 //
 // note: this implementation is generic and not specific to metal
 //       if it proves to work well, we can start using it for other backends in the future
 void ggml_graph_optimize(struct ggml_cgraph * gf);
 
 // mat-mat vs mat-vec dispatch; used by both supports_op and ggml_metal_op_mul_mat*
-bool ggml_metal_op_mul_mat_use_fwht(const struct ggml_tensor * op, size_t max_tg_mem);
+bool ggml_metal_op_mul_mat_use_fwht (const struct ggml_tensor * op, size_t max_tg_mem);
 bool ggml_metal_op_mul_mat_use_mm   (const struct ggml_tensor * op, bool has_simdgroup_mm);
 bool ggml_metal_op_mul_mat_id_use_mm(const struct ggml_tensor * op, bool has_simdgroup_mm);
 
+bool ggml_metal_op_mul_mat_use_mma    (const struct ggml_tensor * op);
+bool ggml_metal_op_mul_mat_may_use_mma(const struct ggml_tensor * op); // graph structure only
+
+// the few-row MMA kernel for a src0 type and rt src1 tiles: per 32-weight block (q4_0, q8_0 with one tile), q5_K, or the generic 64-weight chunk kernel
+enum ggml_metal_mma_kind { GGML_METAL_MMA_KIND_BLK, GGML_METAL_MMA_KIND_Q5_K, GGML_METAL_MMA_KIND_GEN };
+enum ggml_metal_mma_kind ggml_metal_mul_mv_mma_kind(enum ggml_type type, int rt);
+
+// the src1 tiles of the few-row MMA kernels for mat-mul op: one 8-row tile, or two above 8 rows
+int ggml_metal_mul_mv_mma_rt(const struct ggml_tensor * op);
+
+// the weights of K per simdgroup step of the few-row MMA kernel for a src0 type and rt src1 tiles, 0 if none takes the type
+int64_t ggml_metal_mul_mv_mma_k_step(enum ggml_type type, int rt);
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -1,4 +1,5 @@
 #include "ggml-metal-device.h"
+#include "ggml-metal-common.h"
 
 #include "ggml-metal-impl.h"
 #include "ggml-metal-tuning.h"
@@ -801,6 +802,125 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext(ggml_
     return res;
 }
 
+// threadgroups needed to fill the GPU
+// TODO: dedup
+static constexpr int64_t GGML_METAL_MIN_THREADGROUPS = 128;
+
+// few-row MMA mat-mul: 8x8 simdgroup matrices for the src1 rows of ggml_metal_mul_mat_use_mma
+static constexpr int     GGML_METAL_MMA_TILE          = 8;
+static constexpr int     GGML_METAL_MMA_NT_MAX        = 4;
+static constexpr int     GGML_METAL_MMA_SMEM_MAX      = 16384;
+
+// simdgroups per threadgroup for at most FEW_ROWS, at most MANY_ROWS, and more src0 rows
+static constexpr int64_t GGML_METAL_MMA_FEW_ROWS      = 64;
+static constexpr int64_t GGML_METAL_MMA_MANY_ROWS     = 6144;
+static constexpr int     GGML_METAL_MMA_NSG_FEW_ROWS  = 32;
+static constexpr int     GGML_METAL_MMA_NSG_MID_ROWS  = 16;
+static constexpr int     GGML_METAL_MMA_NSG_MANY_ROWS = 8;
+
+struct ggml_metal_mma_tiling {
+    int nsg; // simdgroups per threadgroup, each over a slice of K
+    int nt;  // 8-row src0 tiles per threadgroup
+    int rt;  // 8-row src1 tiles per threadgroup
+};
+
+// halves n while it is above limit
+static int ggml_metal_halve_to_limit(int n, int64_t limit) {
+    while (n > 1 && n > limit) {
+        n /= 2;
+    }
+    return n;
+}
+
+static size_t ggml_metal_mul_mv_mma_smem(int nsg, int nt, int rt) {
+    // one 8x8 float simdgroup matrix per output tile
+    constexpr size_t tile_bytes = 8*8*sizeof(float);
+    return (size_t) nsg*nt*rt*tile_bytes;
+}
+
+// fewer src0 rows need more simdgroups per threadgroup (a finer K split) to fill the GPU.
+// the tiles become narrower when the K-slice reduction buffer is too large or when too few threadgroups fill the GPU.
+static ggml_metal_mma_tiling ggml_metal_op_mul_mat_mma_tiling(const ggml_tensor * op) {
+    const ggml_type type = op->src[0]->type;
+    const int64_t   ne00 = op->src[0]->ne[0];
+    const int64_t   ne01 = op->src[0]->ne[1];
+
+    ggml_metal_mma_tiling res;
+    res.rt = ggml_metal_mul_mv_mma_rt(op);
+
+    const int64_t n_steps = ne00/ggml_metal_mul_mv_mma_k_step(type, res.rt);
+
+    int nsg = GGML_METAL_MMA_NSG_MANY_ROWS;
+    if (ne01 <= GGML_METAL_MMA_FEW_ROWS) {
+        nsg = GGML_METAL_MMA_NSG_FEW_ROWS;
+    } else if (ne01 <= GGML_METAL_MMA_MANY_ROWS) {
+        nsg = GGML_METAL_MMA_NSG_MID_ROWS;
+    }
+    res.nsg = ggml_metal_halve_to_limit(nsg, n_steps);
+
+    const int64_t nt_smem   = GGML_METAL_MMA_SMEM_MAX/(int64_t) ggml_metal_mul_mv_mma_smem(res.nsg, 1, res.rt);
+    const int64_t nt_rows   = ne01/(GGML_METAL_MIN_THREADGROUPS*GGML_METAL_MMA_TILE);
+    res.nt = ggml_metal_halve_to_limit(GGML_METAL_MMA_NT_MAX, std::min(nt_smem, nt_rows));
+
+    return res;
+}
+
+// the pipeline for the tiling, with fewer simdgroups when the device cannot run that many threads
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_mma_auto(ggml_metal_library_t lib, const ggml_tensor * op, bool add) {
+    ggml_metal_mma_tiling tiling = ggml_metal_op_mul_mat_mma_tiling(op);
+
+    auto pipeline = ggml_metal_library_get_pipeline_mul_mv_mma(lib, op, tiling.nsg, tiling.nt, tiling.rt, add);
+    while (tiling.nsg > 1 && ggml_metal_pipeline_max_theads_per_threadgroup(pipeline) < tiling.nsg*32) {
+        tiling.nsg /= 2;
+        pipeline = ggml_metal_library_get_pipeline_mul_mv_mma(lib, op, tiling.nsg, tiling.nt, tiling.rt, add);
+    }
+
+    return pipeline;
+}
+
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_mma(ggml_metal_library_t lib, const ggml_tensor * op, int nsg, int nt, int rt, bool add) {
+    char base[256];
+    char name[256];
+
+    const ggml_type tsrc0 = op->src[0]->type;
+    const ggml_type tsrc1 = op->src[1]->type;
+    const int       ne12  = op->src[1]->ne[2];
+    const int       r2    = ne12 / op->src[0]->ne[2];
+    const int       r3    = op->src[1]->ne[3] / op->src[0]->ne[3];
+
+    GGML_ASSERT(ne12 <= INT16_MAX && r2 <= INT16_MAX && r3 <= INT16_MAX);
+
+    // the specialized kernels unroll over a compile-time row length
+    const int ne00 = ggml_metal_mul_mv_mma_kind(tsrc0, rt) != GGML_METAL_MMA_KIND_GEN ? op->src[0]->ne[0] : 0;
+
+    snprintf(base, 256, "kernel_mul_mv_mma_%s_%s_nt%d_rt%d", ggml_type_name(tsrc0), ggml_type_name(tsrc1), nt, rt);
+    snprintf(name, 256, "%s_nsg=%d_ne12=%d_r2=%d_r3=%d_ne00=%d_add=%d", base, nsg, ne12, r2, r3, ne00, add);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        ggml_metal_cv_t cv = ggml_metal_cv_init();
+
+        ggml_metal_cv_set_int32(cv, ne00,           FC_MUL_MV_MMA + 4);
+        ggml_metal_cv_set_int16(cv, nsg,            FC_MUL_MV_MMA + 0);
+        ggml_metal_cv_set_int16(cv, (int16_t) ne12, FC_MUL_MV_MMA + 1);
+        ggml_metal_cv_set_int16(cv, (int16_t) r2,   FC_MUL_MV_MMA + 2);
+        ggml_metal_cv_set_int16(cv, (int16_t) r3,   FC_MUL_MV_MMA + 3);
+        ggml_metal_cv_set_bool (cv, add,            FC_MUL_MV_MMA + 5);
+
+        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
+
+        ggml_metal_cv_free(cv);
+    }
+
+    res.nsg  = nsg;
+    res.smem = ggml_metal_mul_mv_mma_smem(nsg, nt, rt);
+
+    res.nr0 = GGML_METAL_MMA_TILE*nt;
+    res.nr1 = GGML_METAL_MMA_TILE*rt;
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_metal_library_t lib, const ggml_tensor * op) {
     char base[256];
     char name[256];

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

@@ -135,6 +135,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net   (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext        (ggml_metal_library_t lib, const struct ggml_tensor * op, int nsg, int nxpsg, int r1ptg);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_mma        (ggml_metal_library_t lib, const struct ggml_tensor * op, int nsg, int nt, int rt, bool add);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_mma_auto   (ggml_metal_library_t lib, const struct ggml_tensor * op, bool add);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm            (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv            (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0    (ggml_metal_library_t lib, int ne02, int ne20);
@@ -302,8 +304,6 @@ struct ggml_metal_device_props {
     bool use_residency_sets;
     bool use_shared_buffers;
 
-    bool supports_gpu_family_apple7;
-
     enum ggml_metal_device_id device_id;
     int gpu_family;
 

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

@@ -126,6 +126,7 @@ int ggml_metal_pipeline_max_theads_per_threadgroup(struct ggml_metal_pipeline_wi
     X(FA_VEC_Q5_1,     fa_vec_q5_1)    \
     X(FA_VEC_Q8_0,     fa_vec_q8_0)    \
     X(MUL_MV,          mul_mv)         \
+    X(MUL_MV_MMA,      mul_mv_mma)     \
     X(MUL_MM,          mul_mm)         \
     X(QUANTIZE,        quantize)       \
     X(SOFTMAX,         softmax)        \
@@ -1278,8 +1279,6 @@ ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) {
                     dev->props.use_shared_buffers = true;
                 }
 
-                dev->props.supports_gpu_family_apple7 = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7];
-
                 dev->props.device_id = ggml_metal_device_id_parse([[dev->mtl_device name] UTF8String]);
 
                 dev->props.op_offload_min_batch_size  = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32;

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

@@ -1,8 +1,10 @@
 #include "ggml-metal-fusion.h"
 
-#include "ggml-backend-impl.h"
+#include "ggml-metal-common.h"
 #include "ggml-metal-device.h"
 
+#include "ggml-backend-impl.h"
+
 #include <algorithm>
 #include <cstddef>
 #include <cstring>
@@ -85,17 +87,34 @@ static bool ggml_metal_fusion_same_buffer(const ggml_tensor * a, const ggml_tens
     return ggml_metal_buffer_get_id(ca, a).metal == ggml_metal_buffer_get_id(cb, b).metal;
 }
 
+// true if the memory of two tensors overlaps in the same Metal buffer
+static bool ggml_metal_fusion_overlap(const ggml_tensor * a, const ggml_tensor * b) {
+    ggml_backend_buffer_t ba = a->view_src ? a->view_src->buffer : a->buffer;
+    ggml_backend_buffer_t bb = b->view_src ? b->view_src->buffer : b->buffer;
+
+    const ggml_metal_buffer_id bid_a = ggml_metal_buffer_get_id((ggml_metal_buffer_t) ba->context, a);
+    const ggml_metal_buffer_id bid_b = ggml_metal_buffer_get_id((ggml_metal_buffer_t) bb->context, b);
+
+    if (bid_a.metal == nullptr || bid_a.metal != bid_b.metal) {
+        return false;
+    }
+
+    return bid_a.offs <= bid_b.offs
+        ? bid_b.offs - bid_a.offs < ggml_nbytes(a)
+        : bid_a.offs - bid_b.offs < ggml_nbytes(b);
+}
+
 // ---- pattern checks ------------------------------------------------------
 
 // NORM/RMS_NORM + MUL + ADD: the weight/bias of each fused step must match the norm input
 // width, be contiguous rows, and the fused outputs must stay F32
 static bool ggml_metal_fusion_check_norm(
-        const ggml_metal_fusion      * fusion,
-        const ggml_tensor * const    * nodes,
-        const ggml_cgraph            * gf,
-        const int                    * node_idxs,
-              int                      idx,
-              ggml_metal_fusion_mode   mode) {
+        const ggml_metal_fusion       * fusion,
+        const ggml_tensor * const     * nodes,
+        const ggml_cgraph             * gf,
+        const int                     * node_idxs,
+              int                       idx,
+              ggml_metal_fusion_mode    mode) {
     GGML_UNUSED(mode);
     GGML_UNUSED(gf);
     GGML_UNUSED(node_idxs);
@@ -140,12 +159,12 @@ static bool ggml_metal_fusion_check_norm(
 
 // SSM_CONV + UNARY (silu)
 static bool ggml_metal_fusion_check_ssm_conv_silu(
-        const ggml_metal_fusion      * fusion,
-        const ggml_tensor * const    * nodes,
-        const ggml_cgraph            * gf,
-        const int                    * node_idxs,
-              int                      idx,
-              ggml_metal_fusion_mode   mode) {
+        const ggml_metal_fusion       * fusion,
+        const ggml_tensor * const     * nodes,
+        const ggml_cgraph             * gf,
+        const int                     * node_idxs,
+              int                       idx,
+              ggml_metal_fusion_mode    mode) {
     GGML_UNUSED(fusion);
     GGML_UNUSED(gf);
     GGML_UNUSED(node_idxs);
@@ -173,12 +192,12 @@ static bool ggml_metal_fusion_check_ssm_conv_silu(
 // ADD x N: each ADD reads the previous ADD as src0, and all addends must share layout
 // (and, in FULL mode, live in the same Metal buffer)
 static bool ggml_metal_fusion_check_add_chain(
-        const ggml_metal_fusion      * fusion,
-        const ggml_tensor * const    * nodes,
-        const ggml_cgraph            * gf,
-        const int                    * node_idxs,
-              int                      idx,
-              ggml_metal_fusion_mode   mode) {
+        const ggml_metal_fusion       * fusion,
+        const ggml_tensor * const     * nodes,
+        const ggml_cgraph             * gf,
+        const int                     * node_idxs,
+              int                       idx,
+              ggml_metal_fusion_mode    mode) {
     GGML_UNUSED(gf);
     GGML_UNUSED(node_idxs);
     GGML_UNUSED(idx);
@@ -209,11 +228,11 @@ static bool ggml_metal_fusion_check_add_chain(
 // attn scores view), so unlike the other patterns this is not an elision chain: the structural
 // checks live entirely in this callback (unsafe = true).
 static bool ggml_metal_fusion_check_gdn_cache(
-        const ggml_metal_fusion      * fusion,
-        const ggml_tensor * const    * nodes,
-        const ggml_cgraph            * gf,
-        const int                    * node_idxs,
-              int                      idx,
+        const ggml_metal_fusion       * fusion,
+        const ggml_tensor * const     * nodes,
+        const ggml_cgraph             * gf,
+        const int                     * node_idxs,
+              int                       idx,
               ggml_metal_fusion_mode    mode) {
     GGML_UNUSED(fusion);
     GGML_UNUSED(gf);
@@ -271,12 +290,12 @@ static bool ggml_metal_fusion_check_gdn_cache(
 
 // MUL + SIN + SQR + MUL + ADD (snake activation)
 static bool ggml_metal_fusion_check_snake(
-        const ggml_metal_fusion      * fusion,
-        const ggml_tensor * const    * nodes,
-        const ggml_cgraph            * gf,
-        const int                    * node_idxs,
-              int                      idx,
-              ggml_metal_fusion_mode   mode) {
+        const ggml_metal_fusion       * fusion,
+        const ggml_tensor * const     * nodes,
+        const ggml_cgraph             * gf,
+        const int                     * node_idxs,
+              int                       idx,
+              ggml_metal_fusion_mode    mode) {
     GGML_UNUSED(fusion);
     GGML_UNUSED(mode);
     GGML_UNUSED(gf);
@@ -345,12 +364,12 @@ static const std::vector<ggml_op> ops_topk_moe_norm_scale = {
 };
 
 static bool ggml_metal_fusion_check_topk_moe(
-        const ggml_metal_fusion      * fusion,
-        const ggml_tensor * const    * nodes,
-        const ggml_cgraph            * gf,
-        const int                    * node_idxs,
-              int                      idx,
-              ggml_metal_fusion_mode   mode) {
+        const ggml_metal_fusion       * fusion,
+        const ggml_tensor * const     * nodes,
+        const ggml_cgraph             * gf,
+        const int                     * node_idxs,
+              int                       idx,
+              ggml_metal_fusion_mode    mode) {
     GGML_ASSERT(fusion->ops.size() >= 3);
     GGML_UNUSED(nodes);
 
@@ -591,12 +610,12 @@ static bool ggml_metal_fusion_match_moe_reduce(
 }
 
 static bool ggml_metal_fusion_check_moe_reduce(
-        const ggml_metal_fusion      * fusion,
-        const ggml_tensor * const    * nodes,
-        const ggml_cgraph            * gf,
-        const int                    * node_idxs,
-              int                      idx,
-              ggml_metal_fusion_mode   mode) {
+        const ggml_metal_fusion       * fusion,
+        const ggml_tensor * const     * nodes,
+        const ggml_cgraph             * gf,
+        const int                     * node_idxs,
+              int                       idx,
+              ggml_metal_fusion_mode    mode) {
     GGML_UNUSED(nodes);
 
     ggml_metal_moe_reduce_match match;
@@ -622,6 +641,71 @@ static bool ggml_metal_fusion_check_moe_reduce(
     return true;
 }
 
+// true if t is or views a tensor in a buffer marked as weights, such as a bias; the model loader marks its buffers before
+// any graph is optimized, and tensors in unmarked or not yet allocated buffers count as non-weights in both phases
+static bool ggml_metal_tensor_is_weight(const struct ggml_tensor * t) {
+    const ggml_tensor * base = t->view_src != NULL ? t->view_src : t;
+
+    return base->buffer != NULL && ggml_backend_buffer_get_usage(base->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS;
+}
+
+static const struct ggml_tensor * ggml_metal_mul_mat_add_operand(const struct ggml_tensor * mm, const struct ggml_tensor * add) {
+    if (add->op != GGML_OP_ADD || (add->src[0] == mm) == (add->src[1] == mm)) {
+        return NULL;
+    }
+
+    const ggml_tensor * other = add->src[0] == mm ? add->src[1] : add->src[0];
+
+    const bool ok = other->type == GGML_TYPE_F32 && add->type == GGML_TYPE_F32 && !ggml_metal_tensor_is_weight(other);
+
+    return ok ? other : NULL;
+}
+
+static const struct ggml_tensor * ggml_metal_mul_mat_add_residual(const struct ggml_tensor * mm, const struct ggml_tensor * add) {
+    const ggml_tensor * res = ggml_metal_mul_mat_add_operand(mm, add);
+
+    const bool ok = res != NULL && ggml_are_same_shape(res, mm) &&
+        ggml_is_contiguous(res) && ggml_is_contiguous(mm) && ggml_is_contiguous(add);
+
+    return ok ? res : NULL;
+}
+
+// MUL_MAT + ADD of an f32 non-weight: the reorder packs it without reading row counts, so ubatch sizes share one order;
+// the encoder fuses only a same-shape residual in the few-row MMA store, which the sum may overlap only in place
+static bool ggml_metal_fusion_check_mul_mat_add(
+        const ggml_metal_fusion       * fusion,
+        const ggml_tensor * const     * nodes,
+        const ggml_cgraph             * gf,
+        const int                     * node_idxs,
+              int                       idx,
+              ggml_metal_fusion_mode    mode) {
+    GGML_UNUSED(gf);
+    GGML_UNUSED(node_idxs);
+    GGML_UNUSED(idx);
+    GGML_UNUSED(fusion);
+
+    const ggml_tensor * mm  = nodes[0];
+    const ggml_tensor * add = nodes[1];
+
+    if (ggml_metal_mul_mat_add_operand(mm, add) == nullptr ||
+        !ggml_metal_op_mul_mat_may_use_mma(mm)) {
+        return false;
+    }
+
+    if (mode == GGML_METAL_FUSION_STRUCTURAL) {
+        return true;
+    }
+
+    const ggml_tensor * res = ggml_metal_mul_mat_add_residual(mm, add);
+
+    if (res == nullptr || !ggml_metal_op_mul_mat_use_mma(mm)) {
+        return false;
+    }
+
+    return !ggml_metal_fusion_overlap(add, mm->src[0]) && !ggml_metal_fusion_overlap(add, mm->src[1]) &&
+        (add->data == res->data || !ggml_metal_fusion_overlap(add, res));
+}
+
 // ---- patterns ------------------------------------------------------------
 
 static const std::vector<ggml_op> ops_norm_mul         = { GGML_OP_NORM, GGML_OP_MUL };
@@ -671,6 +755,8 @@ static const std::vector<ggml_op> ops_moe_reduce_8 = {
     GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD
 };
 
+static const std::vector<ggml_op> ops_mul_mat_add = { GGML_OP_MUL_MAT, GGML_OP_ADD };
+
 static const std::vector<ggml_metal_fusion> ggml_metal_fusions = {
     { GGML_METAL_FUSION_NORM_MUL,       ops_norm_mul,               {},     false, ggml_metal_fusion_check_norm },
     { GGML_METAL_FUSION_NORM_MUL_ADD,   ops_norm_mul_add,           {},     false, ggml_metal_fusion_check_norm },
@@ -698,6 +784,7 @@ static const std::vector<ggml_metal_fusion> ggml_metal_fusions = {
     { GGML_METAL_FUSION_MOE_REDUCE,     ops_moe_reduce_7,           {},     true,  ggml_metal_fusion_check_moe_reduce },
     { GGML_METAL_FUSION_MOE_REDUCE,     ops_moe_reduce_8,           {},     true,  ggml_metal_fusion_check_moe_reduce },
     { GGML_METAL_FUSION_SSM_CONV_SILU,  ops_ssm_conv_silu,          {},     false, ggml_metal_fusion_check_ssm_conv_silu },
+    { GGML_METAL_FUSION_MUL_MAT_ADD,    ops_mul_mat_add,            {},     false, ggml_metal_fusion_check_mul_mat_add },
 };
 
 // ---- alloc deps -----------------------------------------------------------

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

@@ -2,8 +2,9 @@
 //
 // every fusable subgraph is declared exactly once as a ggml_metal_fusion entry in
 // the table in ggml-metal-fusion.cpp. both the graph optimizer (ggml_metal_fusion_max)
-// and the op encoders (ggml_metal_fusion_next) consult this same table, so the two
-// phases can never disagree about what can be fused.
+// and the op encoders (ggml_metal_fusion_next) consult this same table with the same device
+// properties. a check that passes in FULL mode also passes in STRUCTURAL mode, so the encoders
+// fuse only groups that the optimizer may pack.
 
 #pragma once
 
@@ -15,13 +16,15 @@
 extern "C" {
 #endif
 
+struct ggml_metal_device_props;
+
 // the maximum number of nodes that can be fused in a single kernel
 // (also the maximum length of a packed fusion group during graph optimization)
 #define GGML_METAL_FUSION_MAX 16
 
 typedef enum ggml_metal_fusion_mode {
-    // structural checks only; used by the graph optimizer, at which point the graph
-    // tensors are not allocated yet, so buffer placement cannot be verified
+    // structural checks for the graph optimizer, which runs before the graph tensors are allocated (weights
+    // already are); a check may skip conditions that differ between batch sizes, and so accept more than FULL
     GGML_METAL_FUSION_STRUCTURAL = 0,
     // full checks, including buffer placement; used by the op encoders
     GGML_METAL_FUSION_FULL,
@@ -39,6 +42,7 @@ typedef enum ggml_metal_fusion_id {
     GGML_METAL_FUSION_TOPK_MOE,     // SOFT_MAX + ARGSORT + GET_ROWS + norm/scale (MoE routing)
     GGML_METAL_FUSION_MOE_REDUCE,   // MUL + expert VIEWs + ADD chain (MoE output reduction)
     GGML_METAL_FUSION_SSM_CONV_SILU, // SSM_CONV + UNARY (silu)
+    GGML_METAL_FUSION_MUL_MAT_ADD,  // MUL_MAT + ADD (residual added in the few-row MMA store)
 } ggml_metal_fusion_id;
 
 struct ggml_metal_fusion; // defined in ggml-metal-fusion.cpp
@@ -79,8 +83,8 @@ void ggml_metal_fusion_info_stats_reset(      struct ggml_metal_fusion_info * fi
 int  ggml_metal_fusion_info_stats_get  (const struct ggml_metal_fusion_info * finfo, const char ** labels, uint64_t * counts, int n);
 void ggml_metal_fusion_info_labels_init(      struct ggml_metal_fusion_info * finfo);
 
-// compute phase: longest fusion starting at idx (a position in node_idxs) that matches in `mode`.
-// returns the matching pattern (nullptr if no fusion) and sets *n_out to the number of nodes consumed.
+// compute phase: longest fusion starting at idx (a position in node_idxs) that matches in `mode` on a device with
+// props. returns the matching pattern (nullptr if no fusion) and sets *n_out to the number of nodes consumed.
 const ggml_metal_fusion * ggml_metal_fusion_next(
         const struct ggml_cgraph * gf,
         const int * node_idxs,
@@ -90,7 +94,7 @@ const ggml_metal_fusion * ggml_metal_fusion_next(
         int * n_out);
 
 // optimize phase: maximum number of nodes starting at idx (a raw sequential graph index) that
-// could be fused, chaining patterns back-to-back. returns at least 1.
+// could be fused on a device with props, chaining patterns back-to-back. returns at least 1.
 int ggml_metal_fusion_max(const struct ggml_cgraph * gf, int idx);
 
 #ifdef __cplusplus

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

@@ -123,6 +123,7 @@
 #define FC_PAD                         2100
 #define FC_FLASH_ATTN_EXT_TENSOR       2200
 #define FC_LIGHTNING_INDEXER           2200
+#define FC_MUL_MV_MMA                  2300
 
 // op-specific constants
 #define OP_FLASH_ATTN_EXT_NQPSG 8
@@ -215,6 +216,7 @@ typedef struct {
     uint64_t nb2;
     uint64_t nb3;
     int32_t  dim;
+    int32_t  nc0;
 } ggml_metal_kargs_concat;
 
 typedef struct {

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

@@ -15,6 +15,10 @@
 #include <limits>
 #include <cmath>
 
+// threadgroups needed to fill the GPU
+// TODO: dedup
+static constexpr int64_t GGML_METAL_MIN_THREADGROUPS = 128;
+
 static ggml_metal_buffer_id ggml_metal_get_buffer_id(const ggml_tensor * t) {
     if (!t) {
         return { nullptr, 0 };
@@ -584,6 +588,27 @@ int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) {
         ne0_arg  = ne0/blck;
     }
 
+    int nth = std::min(256, ne0_arg);
+
+    // when rows are small, we can batch them together in a single threadgroup
+    int nrptg = 1;
+    if (nth < 256) {
+        nrptg = std::min((256 + nth - 1) / nth, ne1);
+        if (nrptg * nth > 256) {
+            nrptg = 256 / nth;
+        }
+    }
+
+    const int nw0 = (ne1 + nrptg - 1) / nrptg;
+
+    // split long rows across threadgroups when there are too few rows to fill the GPU
+    const int64_t n_rows = (int64_t) nw0*ne2*ne3;
+
+    int nc0 = 1;
+    if (nrptg == 1) {
+        nc0 = (int) std::max<int64_t>(1, std::min<int64_t>((ne0_arg + nth - 1)/nth, (GGML_METAL_MIN_THREADGROUPS + n_rows - 1)/n_rows));
+    }
+
     ggml_metal_kargs_concat args = {
         /*.ne00 =*/ ne00_arg,
         /*.ne01 =*/ ne01,
@@ -610,6 +635,7 @@ int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) {
         /*.nb2  =*/ nb2,
         /*.nb3  =*/ nb3,
         /*.dim  =*/ dim,
+        /*.nc0  =*/ nc0,
     };
 
     auto pipeline = ggml_metal_library_get_pipeline_concat(lib, op->type);
@@ -620,20 +646,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_arg);
-
-    // when rows are small, we can batch them together in a single threadgroup
-    int nrptg = 1;
-    if (nth < 256) {
-        nrptg = std::min((256 + nth - 1) / nth, ne1);
-        if (nrptg * nth > 256) {
-            nrptg = 256 / nth;
-        }
-    }
-
-    const int nw0 = (ne1 + nrptg - 1) / nrptg;
-
-    ggml_metal_encoder_dispatch_threadgroups(enc, nw0, ne2, ne3, nth, nrptg, 1);
+    ggml_metal_encoder_dispatch_threadgroups(enc, nw0*nc0, ne2, ne3, nth, nrptg, 1);
 
     return 1;
 }
@@ -2421,6 +2434,157 @@ int ggml_metal_op_pool_2d(ggml_metal_op_t ctx, int idx) {
     return 1;
 }
 
+// the number of nodes from idx on that the fusion table fuses as pattern id into one dispatch, 1 if it fuses none
+static int ggml_metal_op_try_fusion(ggml_metal_op_t ctx, int idx, ggml_metal_fusion_id id) {
+    if (!ctx->use_fusion()) {
+        return 1;
+    }
+
+    int n = 1;
+    const ggml_metal_fusion * fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n);
+    if (fusion == nullptr || ggml_metal_fusion_get_id(fusion) != id) {
+        return 1;
+    }
+
+    ctx->count_fusions(fusion);
+
+    if (ggml_metal_fusion_info_debug(ctx->finfo) > 1) {
+        GGML_LOG_DEBUG("%s: fuse: %s to %s, %d nodes\n", __func__, ggml_op_name(ctx->node(idx)->op), ggml_op_name(ctx->node(idx + n - 1)->op), n);
+    }
+
+    return n;
+}
+
+static int ggml_metal_op_mul_mat_mma(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_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
+    GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
+    GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne);
+    GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb);
+    GGML_TENSOR_LOCALS( int32_t, ne,  op,         ne);
+
+    // the MMA store adds the residual when the table fuses the ADD after this mat-mul
+    const int  n_fuse   = ggml_metal_op_try_fusion(ctx, idx, GGML_METAL_FUSION_MUL_MAT_ADD);
+    const bool fuse_add = n_fuse > 1;
+
+    const ggml_tensor * dst = op;
+    const ggml_tensor * res = dst;
+
+    if (fuse_add) {
+        dst = ctx->node(idx + n_fuse - 1);
+        res = dst->src[0]->op == GGML_OP_MUL_MAT ? dst->src[1] : dst->src[0];
+    }
+
+    auto pipeline = ggml_metal_library_get_pipeline_mul_mv_mma_auto(lib, op, fuse_add);
+
+    ggml_metal_kargs_mul_mv_ext args = {
+        /*.ne00  =*/ ne00,
+        /*.ne01  =*/ ne01,
+        /*.ne02  =*/ ne02,
+        /*.nb00  =*/ nb00,
+        /*.nb01  =*/ nb01,
+        /*.nb02  =*/ nb02,
+        /*.nb03  =*/ nb03,
+        /*.ne10  =*/ ne10,
+        /*.ne11  =*/ ne11,
+        /*.ne12  =*/ ne12,
+        /*.nb10  =*/ nb10,
+        /*.nb11  =*/ nb11,
+        /*.nb12  =*/ nb12,
+        /*.nb13  =*/ nb13,
+        /*.ne0   =*/ ne0,
+        /*.ne1   =*/ ne1,
+        /*.r2    =*/ (int16_t) (ne12/ne02),
+        /*.r3    =*/ (int16_t) (ne13/ne03),
+    };
+
+    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->src[1]), 2);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(dst),        3);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(res),        4);
+
+    ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0);
+
+    const int rows0 = pipeline.nr0;
+    const int rows1 = pipeline.nr1;
+
+    ggml_metal_encoder_dispatch_threadgroups(enc, (ne01 + rows0 - 1)/rows0, (ne11 + rows1 - 1)/rows1, ne12*ne13, 32, pipeline.nsg, 1);
+
+    return n_fuse;
+}
+
+// the generic mat-vec kernel, or its 2-row Q4_0 variant if nc
+static int ggml_metal_op_mul_mat_mv(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_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
+    GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
+    GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne);
+    GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb);
+    GGML_TENSOR_LOCALS( int32_t, ne,  op,         ne);
+
+    const int16_t r2 = ne12/ne02;
+    const int16_t r3 = ne13/ne03;
+
+    auto pipeline = ggml_metal_library_get_pipeline_mul_mv(lib, op);
+
+    const int nr0 = pipeline.nr0;
+    const int nr1 = pipeline.nr1;
+    const int nsg = pipeline.nsg;
+
+    const size_t smem = pipeline.smem;
+
+    ggml_metal_kargs_mul_mv args = {
+        /*.ne00 =*/ ne00,
+        /*.ne01 =*/ ne01,
+        /*.ne02 =*/ ne02,
+        /*.nb00 =*/ nb00,
+        /*.nb01 =*/ nb01,
+        /*.nb02 =*/ nb02,
+        /*.nb03 =*/ nb03,
+        /*.ne10 =*/ ne10,
+        /*.ne11 =*/ ne11,
+        /*.ne12 =*/ ne12,
+        /*.nb10 =*/ nb10,
+        /*.nb11 =*/ nb11,
+        /*.nb12 =*/ nb12,
+        /*.nb13 =*/ nb13,
+        /*.ne0  =*/ ne0,
+        /*.ne1  =*/ ne1,
+        /*.nr0  =*/ nr0,
+        /*.r2   =*/ r2,
+        /*.r3   =*/ r3,
+    };
+
+    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->src[1]), 2);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         3);
+
+    ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
+
+    if (op->src[0]->type == GGML_TYPE_F32 ||
+        op->src[0]->type == GGML_TYPE_F16 ||
+        op->src[0]->type == GGML_TYPE_BF16 ||
+        op->src[0]->type == GGML_TYPE_Q8_0) {
+        ggml_metal_encoder_dispatch_threadgroups(enc, ((ne01 + nr0 - 1)/(nr0)), ((ne11 + nr1 - 1)/nr1), ne12*ne13, 32, nsg, 1);
+    } else {
+        ggml_metal_encoder_dispatch_threadgroups(enc, ((ne01 + nr0*nsg - 1)/(nr0*nsg)), ((ne11 + nr1 - 1)/nr1), ne12*ne13, 32, nsg, 1);
+    }
+
+    return 1;
+}
+
 int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 
@@ -2433,6 +2597,10 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
         return ggml_metal_op_fwht(ctx, idx);
     }
 
+    if (props_dev->has_simdgroup_mm && ggml_metal_op_mul_mat_use_mma(op)) {
+        return ggml_metal_op_mul_mat_mma(ctx, idx);
+    }
+
     GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
     GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
     GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne);
@@ -2597,52 +2765,7 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
 
         ggml_metal_encoder_dispatch_threadgroups(enc, ((ne11 + nr1 - 1) / nr1), ((ne01 + nr0 - 1) / nr0), ne12 * ne13, 32, nsg, 1);
     } else {
-        auto pipeline = ggml_metal_library_get_pipeline_mul_mv(lib, op);
-
-        const int nr0 = pipeline.nr0;
-        const int nr1 = pipeline.nr1;
-        const int nsg = pipeline.nsg;
-
-        const size_t smem = pipeline.smem;
-
-        ggml_metal_kargs_mul_mv args = {
-            /*.ne00 =*/ ne00,
-            /*.ne01 =*/ ne01,
-            /*.ne02 =*/ ne02,
-            /*.nb00 =*/ nb00,
-            /*.nb01 =*/ nb01,
-            /*.nb02 =*/ nb02,
-            /*.nb03 =*/ nb03,
-            /*.ne10 =*/ ne10,
-            /*.ne11 =*/ ne11,
-            /*.ne12 =*/ ne12,
-            /*.nb10 =*/ nb10,
-            /*.nb11 =*/ nb11,
-            /*.nb12 =*/ nb12,
-            /*.nb13 =*/ nb13,
-            /*.ne0  =*/ ne0,
-            /*.ne1  =*/ ne1,
-            /*.nr0  =*/ nr0,
-            /*.r2   =*/ r2,
-            /*.r3   =*/ r3,
-        };
-
-        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->src[1]), 2);
-        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         3);
-
-        ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
-
-        if (op->src[0]->type == GGML_TYPE_F32 ||
-            op->src[0]->type == GGML_TYPE_F16 ||
-            op->src[0]->type == GGML_TYPE_BF16 ||
-            op->src[0]->type == GGML_TYPE_Q8_0) {
-            ggml_metal_encoder_dispatch_threadgroups(enc, ((ne01 + nr0 - 1)/(nr0)), ((ne11 + nr1 - 1)/nr1), ne12*ne13, 32, nsg, 1);
-        } else {
-            ggml_metal_encoder_dispatch_threadgroups(enc, ((ne01 + nr0*nsg - 1)/(nr0*nsg)), ((ne11 + nr1 - 1)/nr1), ne12*ne13, 32, nsg, 1);
-        }
+        return ggml_metal_op_mul_mat_mv(ctx, idx);
     }
 
     return 1;
@@ -3968,7 +4091,9 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
     int n_fuse = 1;
     const ggml_metal_fusion * fusion = nullptr;
 
-    if (ctx->use_fusion()) {
+    const bool use_fusion = ctx->use_fusion();
+
+    if (use_fusion) {
         int n = 1;
         fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n);
         n_fuse = n;
@@ -3991,8 +4116,6 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
     ggml_metal_library_t lib = ctx->lib;
     ggml_metal_encoder_t enc = ctx->enc;
 
-    const bool use_fusion = ctx->use_fusion();
-
     const int debug_fusion = ggml_metal_fusion_info_debug(ctx->finfo);
 
     GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);

ggml/src/ggml-metal/kernels/mul_mv_mma.metal

@@ -0,0 +1,608 @@
+#include "common.h"
+#include "dequantize.h"
+
+constant short FC_mul_mv_mma_nsg  [[function_constant(FC_MUL_MV_MMA + 0)]];
+constant short FC_mul_mv_mma_ne12 [[function_constant(FC_MUL_MV_MMA + 1)]];
+constant short FC_mul_mv_mma_r2   [[function_constant(FC_MUL_MV_MMA + 2)]];
+constant short FC_mul_mv_mma_r3   [[function_constant(FC_MUL_MV_MMA + 3)]];
+constant int   FC_mul_mv_mma_ne00 [[function_constant(FC_MUL_MV_MMA + 4)]];
+constant bool  FC_mul_mv_mma_add  [[function_constant(FC_MUL_MV_MMA + 5)]];
+
+// a lane of a few-row MMA tile: A fragment row fm, B fragment columns fn and fn + 1
+struct mul_mv_mma_tile {
+    short    fm;
+    short    fn;
+    int      i01;
+    int      i11;
+    int      i1m;
+    uint64_t offset0;
+    uint64_t offset1;
+};
+
+// the A fragment row and the first B fragment column that lane l holds in an 8x8 simdgroup matrix
+inline short mul_mv_mma_lane_fm(ushort l) { return ((l/4) & 4) + ((l/2) % 4); }
+inline short mul_mv_mma_lane_fn(ushort l) { return ((l/4) & 2)*2 + (l % 2)*2; }
+
+template<short NT, short RT>
+inline mul_mv_mma_tile mul_mv_mma_tile_init(constant ggml_metal_kargs_mul_mv_ext & args, uint3 tgpig, ushort tiisg) {
+    mul_mv_mma_tile tile;
+    tile.fm  = mul_mv_mma_lane_fm(tiisg);
+    tile.fn  = mul_mv_mma_lane_fn(tiisg);
+    tile.i01 = tgpig.x*(8*NT);
+    tile.i11 = tgpig.y*(8*RT);
+    tile.i1m = tgpig.z;
+
+    const int i12 = tile.i1m%FC_mul_mv_mma_ne12;
+    const int i13 = tile.i1m/FC_mul_mv_mma_ne12;
+
+    tile.offset0 = (i12/FC_mul_mv_mma_r2)*args.nb02 + (i13/FC_mul_mv_mma_r3)*args.nb03;
+    tile.offset1 = i12*args.nb12 + i13*args.nb13;
+
+    return tile;
+}
+
+// the src0 row of A fragment row fm in 8-row tile t, clamped to the last row
+inline device const char * mul_mv_mma_src0_row(
+        thread const mul_mv_mma_tile & tile, constant ggml_metal_kargs_mul_mv_ext & args, device const char * src0, short t) {
+    const int r = min(tile.i01 + 8*t + tile.fm, args.ne01 - 1);
+    return src0 + tile.offset0 + (uint64_t) r*args.nb01;
+}
+
+// the src1 row of B fragment column fn + e in 8-row tile rt, clamped to the last row
+inline device const float * mul_mv_mma_src1_row(
+        thread const mul_mv_mma_tile & tile, constant ggml_metal_kargs_mul_mv_ext & args, device const char * src1, short rt, short e) {
+    const int r = min(tile.i11 + 8*rt + tile.fn + e, args.ne11 - 1);
+    return (device const float *) (src1 + tile.offset1 + (uint64_t) r*args.nb11);
+}
+
+constexpr constant static ushort mma_f16_1024_bits = 0x6400;
+constexpr constant static half   mma_f16_1024      = 1024.0h;
+
+// the halves 1024 + q for integers q < 1024: exact normal values, unlike the subnormal q*2^-24 that Metal may flush to zero
+inline half2 mul_mv_mma_1024_plus(ushort2 q) {
+    return as_type<half2>(q | mma_f16_1024_bits);
+}
+
+// adds up the K slices of the NSG simdgroups for an 8*NT x 8*RT output tile and writes it, plus the residual in src2 if FC_mul_mv_mma_add
+template<short NT, short RT>
+inline void mul_mv_mma_store(
+        thread float (&acc)[RT][NT][2],
+        constant ggml_metal_kargs_mul_mv_ext & args,
+        device const char * src2,
+        device char * dst,
+        threadgroup char * shmem,
+        thread const mul_mv_mma_tile & tile, ushort tiisg, ushort sgitg) {
+    const short NSG = FC_mul_mv_mma_nsg;
+
+    threadgroup float * red = (threadgroup float *) shmem;
+
+    FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+        FOR_UNROLL (short t = 0; t < NT; ++t) {
+            red[((sgitg*RT + rt)*NT + t)*64 + 2*tiisg + 0] = acc[rt][t][0];
+            red[((sgitg*RT + rt)*NT + t)*64 + 2*tiisg + 1] = acc[rt][t][1];
+        }
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    device       float * dst_f32 = (device       float *) dst  + (uint64_t) tile.i1m*args.ne0*args.ne1;
+    device const float * res_f32 = (device const float *) src2 + (uint64_t) tile.i1m*args.ne0*args.ne1;
+
+    for (short idx = sgitg*32 + tiisg; idx < RT*NT*64; idx += NSG*32) {
+        float sum = 0.0f;
+        for (short sg = 0; sg < NSG; ++sg) {
+            sum += red[sg*(RT*NT*64) + idx];
+        }
+
+        const short rt = idx/(NT*64);
+        const short t  = (idx/64) % NT;
+        const short l  = (idx % 64)/2;
+        const short e  = idx % 2;
+
+        const int r0 = tile.i01 + 8*t  + mul_mv_mma_lane_fm(l);
+        const int r1 = tile.i11 + 8*rt + mul_mv_mma_lane_fn(l) + e;
+
+        if (r0 < args.ne01 && r1 < args.ne11) {
+            const uint64_t i = (uint64_t) r1*args.ne0 + r0;
+            dst_f32[i] = FC_mul_mv_mma_add ? sum + res_f32[i] : sum;
+        }
+    }
+}
+
+// the per-type parts of kernel_mul_mv_mma_blk for 32-weight blocks. a src1 block splits into halves b0 and b1,
+// and MMA step s uses half b1 when b1_step(s) and the .y value of a pair when y_step(s)
+// q4_0: A lane (m, j) holds qs ushorts j and j + 1 (j even); a high nibble stays in place as 16*q, so b1 is divided by 16.
+// B lane k = fm holds src1 values 2*k, 2*k + 1 (b0, low nibbles) and 2*k + 16, 2*k + 17 (b1, high nibbles) of a block
+struct mul_mv_mma_q4_0 {
+    typedef block_q4_0 block;
+    typedef ushort2    quants;
+
+    // weights per block, and the float2 offset of b1 in a src1 block
+    enum { qk = QK4_0, b1 = QK4_0/4 };
+
+    static short a_off(short fn) { return 1 + fn; }
+    static short b_off(short fm) { return 2*fm; }
+
+    static quants load(device const ushort * qs) { return ushort2(qs[0], qs[1]); }
+    static quants prep(quants q) { return q; }
+    static float2 prep_b1(float2 v) {
+        constexpr float hi_scale = 1.0f/16;
+        return v*hi_scale;
+    }
+
+    static bool b1_step(short s) { return s % 2 != 0; }
+    static bool y_step (short s) { return s >= 2; }
+
+    static half2 frag(quants q, short s) {
+        constexpr ushort lo_mask = 0x000F;
+        constexpr ushort hi_mask = 0x00F0;
+        constexpr half   lo_zero = 8.0h;
+        constexpr half   hi_zero = 16*lo_zero;
+
+        const ushort2 qq = s < 2 ? q : q >> 8;
+        return s % 2 == 0 ? mul_mv_mma_1024_plus(qq & lo_mask) - (mma_f16_1024 + lo_zero) : mul_mv_mma_1024_plus(qq & hi_mask) - (mma_f16_1024 + hi_zero);
+    }
+};
+
+// q8_0: A lane (m, j) holds qs bytes 4*j .. 4*j + 7 (j even); flipping the sign bit of a quant byte gives the unsigned q + 128.
+// B lane k = fm holds src1 values b, b + 1 (b0) and b + 4, b + 5 (b1) of a block, b = 8*(k/2) + 2*(k%2)
+struct mul_mv_mma_q8_0 {
+    typedef block_q8_0 block;
+    typedef ushort4    quants;
+
+    // weights per block, and the float2 offset of b1 in a src1 block
+    enum { qk = QK8_0, b1 = 2 };
+
+    static short a_off(short fn) { return 1 + 2*fn; }
+    static short b_off(short fm) { return 8*(fm/2) + 2*(fm%2); }
+
+    static quants load(device const ushort * qs) { return ushort4(qs[0], qs[1], qs[2], qs[3]); }
+    static quants prep(quants q) {
+        constexpr ushort sign_bits = 0x8080;
+        return q ^ ushort4(sign_bits);
+    }
+    static float2 prep_b1(float2 v) { return v; }
+
+    static bool b1_step(short s) { return s >= 2; }
+    static bool y_step (short s) { return s % 2 != 0; }
+
+    static half2 frag(quants q, short s) {
+        constexpr ushort byte_mask = 0x00FF;
+        constexpr half   q_bias    = mma_f16_1024 + 128.0h;
+
+        const ushort2 w  = s < 2 ? q.xy : q.zw;
+        const ushort2 qq = s % 2 == 0 ? w : w >> 8;
+        return mul_mv_mma_1024_plus(qq & byte_mask) - q_bias;
+    }
+};
+
+template<typename Q, short NT>
+inline void load_mma_blk_a(device const ushort * const x[NT], int off, short fn, thread typename Q::quants * q, thread float * d) {
+    FOR_UNROLL (short t = 0; t < NT; ++t) {
+        device const ushort * qs = x[t] + off;
+        q[t] = Q::load(qs);
+        d[t] = as_type<half>(*(qs - Q::a_off(fn)));
+    }
+}
+
+template<typename Q, short RT>
+inline void load_mma_blk_b(device const float2 * const y[RT][2], int ib, thread float2 (*b0)[2], thread float2 (*b1)[2]) {
+    FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+        FOR_UNROLL (short e = 0; e < 2; ++e) {
+            b0[rt][e] = y[rt][e][ib*(Q::qk/2)];
+            b1[rt][e] = y[rt][e][ib*(Q::qk/2) + Q::b1];
+        }
+    }
+}
+
+// few-row mat-mat (2..16 src1 rows) on 8x8 simdgroup matrices for 32-weight block types: a threadgroup reads each weight once
+// for 8*NT src0 rows x 8*RT src1 rows, and its NSG simdgroups split K
+template<short NT, short RT, typename Q>
+kernel void kernel_mul_mv_mma_blk(
+        constant ggml_metal_kargs_mul_mv_ext & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        device const char * src2,
+        threadgroup  char * shmem [[threadgroup(0)]],
+        uint3  tgpig[[threadgroup_position_in_grid]],
+        ushort tiisg[[thread_index_in_simdgroup]],
+        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+    const short NSG = FC_mul_mv_mma_nsg;
+
+    const mul_mv_mma_tile tile = mul_mv_mma_tile_init<NT, RT>(args, tgpig, tiisg);
+
+    device const ushort * x[NT];
+    FOR_UNROLL (short t = 0; t < NT; ++t) {
+        x[t] = (device const ushort *) mul_mv_mma_src0_row(tile, args, src0, t) + Q::a_off(tile.fn);
+    }
+
+    device const float2 * y[RT][2];
+    FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+        FOR_UNROLL (short e = 0; e < 2; ++e) {
+            y[rt][e] = (device const float2 *) (mul_mv_mma_src1_row(tile, args, src1, rt, e) + Q::b_off(tile.fm));
+        }
+    }
+
+    float acc[RT][NT][2] = {};
+
+    const int nb = FC_mul_mv_mma_ne00/Q::qk;
+
+    // a block is d, then the quants
+    constexpr short us_blk = sizeof(typename Q::block)/2;
+
+    typename Q::quants q[NT];
+    float  d[NT];
+    float2 b0[RT][2];
+    float2 b1[RT][2];
+
+    const int ib0 = min((int) sgitg, nb - 1);
+    load_mma_blk_a<Q, NT>(x, ib0*us_blk, tile.fn, q, d);
+    load_mma_blk_b<Q, RT>(y, ib0, b0, b1);
+
+    for (int ib = sgitg; ib < nb; ib += NSG) {
+        typename Q::quants qc[NT];
+        float  dc[NT];
+        float2 b0c[RT][2];
+        float2 b1c[RT][2];
+
+        FOR_UNROLL (short t = 0; t < NT; ++t) {
+            qc[t] = Q::prep(q[t]);
+            dc[t] = d[t];
+        }
+        FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+            FOR_UNROLL (short e = 0; e < 2; ++e) {
+                b0c[rt][e] = b0[rt][e];
+                b1c[rt][e] = Q::prep_b1(b1[rt][e]);
+            }
+        }
+
+        const int ibn = min(ib + NSG, nb - 1);
+        load_mma_blk_a<Q, NT>(x, ibn*us_blk, tile.fn, q, d);
+        load_mma_blk_b<Q, RT>(y, ibn, b0, b1);
+
+        simdgroup_float8x8 mp[RT][NT];
+        FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+            FOR_UNROLL (short t = 0; t < NT; ++t) {
+                mp[rt][t] = make_filled_simdgroup_matrix<float, 8>(0.0f);
+            }
+        }
+
+        FOR_UNROLL (short s = 0; s < 4; ++s) {
+            simdgroup_float8x8 mb[RT];
+            FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+                const float2 v0 = Q::b1_step(s) ? b1c[rt][0] : b0c[rt][0];
+                const float2 v1 = Q::b1_step(s) ? b1c[rt][1] : b0c[rt][1];
+                mb[rt].thread_elements()[0] = Q::y_step(s) ? v0.y : v0.x;
+                mb[rt].thread_elements()[1] = Q::y_step(s) ? v1.y : v1.x;
+            }
+
+            FOR_UNROLL (short t = 0; t < NT; ++t) {
+                const half2 h = Q::frag(qc[t], s);
+
+                simdgroup_half8x8 ma;
+                ma.thread_elements()[0] = h.x;
+                ma.thread_elements()[1] = h.y;
+
+                FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+                    simdgroup_multiply_accumulate(mp[rt][t], ma, mb[rt], mp[rt][t]);
+                }
+            }
+        }
+
+        FOR_UNROLL (short t = 0; t < NT; ++t) {
+            FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+                acc[rt][t][0] = fma(dc[t], mp[rt][t].thread_elements()[0], acc[rt][t][0]);
+                acc[rt][t][1] = fma(dc[t], mp[rt][t].thread_elements()[1], acc[rt][t][1]);
+            }
+        }
+    }
+
+    mul_mv_mma_store<NT, RT>(acc, args, src2, dst, shmem, tile, tiisg, sgitg);
+}
+
+typedef decltype(kernel_mul_mv_mma_blk<4, 1, mul_mv_mma_q4_0>) mul_mv_mma_t;
+
+template [[host_name("kernel_mul_mv_mma_q4_0_f32_nt1_rt1")]] kernel mul_mv_mma_t kernel_mul_mv_mma_blk<1, 1, mul_mv_mma_q4_0>;
+template [[host_name("kernel_mul_mv_mma_q4_0_f32_nt2_rt1")]] kernel mul_mv_mma_t kernel_mul_mv_mma_blk<2, 1, mul_mv_mma_q4_0>;
+template [[host_name("kernel_mul_mv_mma_q4_0_f32_nt4_rt1")]] kernel mul_mv_mma_t kernel_mul_mv_mma_blk<4, 1, mul_mv_mma_q4_0>;
+template [[host_name("kernel_mul_mv_mma_q4_0_f32_nt1_rt2")]] kernel mul_mv_mma_t kernel_mul_mv_mma_blk<1, 2, mul_mv_mma_q4_0>;
+template [[host_name("kernel_mul_mv_mma_q4_0_f32_nt2_rt2")]] kernel mul_mv_mma_t kernel_mul_mv_mma_blk<2, 2, mul_mv_mma_q4_0>;
+template [[host_name("kernel_mul_mv_mma_q4_0_f32_nt4_rt2")]] kernel mul_mv_mma_t kernel_mul_mv_mma_blk<4, 2, mul_mv_mma_q4_0>;
+
+template [[host_name("kernel_mul_mv_mma_q8_0_f32_nt1_rt1")]] kernel mul_mv_mma_t kernel_mul_mv_mma_blk<1, 1, mul_mv_mma_q8_0>;
+template [[host_name("kernel_mul_mv_mma_q8_0_f32_nt2_rt1")]] kernel mul_mv_mma_t kernel_mul_mv_mma_blk<2, 1, mul_mv_mma_q8_0>;
+template [[host_name("kernel_mul_mv_mma_q8_0_f32_nt4_rt1")]] kernel mul_mv_mma_t kernel_mul_mv_mma_blk<4, 1, mul_mv_mma_q8_0>;
+
+// q5_K scale and min of sub-block j from the 12 packed bytes, held as 3 words
+inline float2 mul_mv_mma_q5_K_scale_min(thread const uint * w, short j) {
+    if (j < 4) {
+        return float2((w[0] >> 8*j) & 63, (w[1] >> 8*j) & 63);
+    }
+    const short k = 8*(j - 4);
+    return float2(((w[2] >> k) & 0xF) | (((w[0] >> (k + 6)) & 3) << 4), ((w[2] >> (k + 4)) & 0xF) | (((w[1] >> (k + 6)) & 3) << 4));
+}
+
+// the A fragments of one qs/qh word: lo holds the 5-bit quants q of the low-nibble sub-block of a pair, hi holds 16*q of the high-nibble sub-block.
+// step e0 takes bytes 0 and 2 of the word, step e1 takes bytes 1 and 3; hs holds the qh bits of the pair.
+inline void mul_mv_mma_q5_K_frags(uint q, uint hs, thread half2 * lo, thread half2 * hi) {
+    const ushort2 qw = as_type<ushort2>(q);
+    const ushort2 hw = as_type<ushort2>(hs);
+
+    lo[0] = mul_mv_mma_1024_plus((qw        & 0x000F) | ((hw << 4) & 0x0010)) - mma_f16_1024;
+    lo[1] = mul_mv_mma_1024_plus(((qw >> 8) & 0x000F) | ((hw >> 4) & 0x0010)) - mma_f16_1024;
+    hi[0] = mul_mv_mma_1024_plus((qw        & 0x00F0) | ((hw << 7) & 0x0100)) - mma_f16_1024;
+    hi[1] = mul_mv_mma_1024_plus(((qw >> 8) & 0x00F0) | ((hw >> 1) & 0x0100)) - mma_f16_1024;
+}
+
+struct mul_mv_mma_q5_K_a {
+    uint2 q;
+    uint2 h;
+    uint  sc[3];
+    uint  dm;
+};
+
+template<short NT>
+inline void load_q5_K_mma_a(device const block_q5_K * const x[NT], int ip, short fn, thread mul_mv_mma_q5_K_a * a) {
+    constexpr short pairs = QK_K/64;
+
+    FOR_UNROLL (short t = 0; t < NT; ++t) {
+        device const block_q5_K * xb = x[t] + ip/pairs;
+        device const uint * sp = (device const uint *) xb->scales;
+
+        a[t].q     = *((device const uint2 *) (xb->qs + 32*(ip%pairs)) + fn/2);
+        a[t].h     = *((device const uint2 *) xb->qh + fn/2);
+        a[t].sc[0] = sp[0];
+        a[t].sc[1] = sp[1];
+        a[t].sc[2] = sp[2];
+        a[t].dm    = *((device const uint *) xb);
+    }
+}
+
+// xs[rt][e][h]: src1 values of sub-block h of the pair, for the steps (b, b + 1) and (b + 4, b + 5)
+template<short RT>
+inline void load_q5_K_mma_b(device const float2 * const y[RT][2], int ip, thread float2 (*xs)[2][2][2]) {
+    FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+        FOR_UNROLL (short e = 0; e < 2; ++e) {
+            FOR_UNROLL (short h = 0; h < 2; ++h) {
+                xs[rt][e][h][0] = y[rt][e][ip*32 + h*16 + 0];
+                xs[rt][e][h][1] = y[rt][e][ip*32 + h*16 + 2];
+            }
+        }
+    }
+}
+
+// few-row mat-mat for q5_K over pairs of 32-weight sub-blocks, laid out like mul_mv_mma_q4_0: A lane (m, j) holds qs and qh bytes 4*j .. 4*j + 7 of a pair (j even).
+// the 1/16 of the in-place high nibbles and the sub-block scales go into the accumulation; the mins are removed with the src1 sums.
+template<short NT, short RT>
+kernel void kernel_mul_mv_mma_q5_K_f32(
+        constant ggml_metal_kargs_mul_mv_ext & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        device const char * src2,
+        threadgroup  char * shmem [[threadgroup(0)]],
+        uint3  tgpig[[threadgroup_position_in_grid]],
+        ushort tiisg[[thread_index_in_simdgroup]],
+        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+    const short NSG = FC_mul_mv_mma_nsg;
+
+    constexpr float hi_scale = 1.0f/16;
+    constexpr short pairs    = QK_K/64;
+
+    const mul_mv_mma_tile tile = mul_mv_mma_tile_init<NT, RT>(args, tgpig, tiisg);
+
+    device const block_q5_K * x[NT];
+    FOR_UNROLL (short t = 0; t < NT; ++t) {
+        x[t] = (device const block_q5_K *) mul_mv_mma_src0_row(tile, args, src0, t);
+    }
+
+    // B lane k = fm holds src1 values b, b + 1, b + 4, b + 5 of each sub-block, b = 8*(k/2) + 2*(k%2)
+    device const float2 * y[RT][2];
+    FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+        FOR_UNROLL (short e = 0; e < 2; ++e) {
+            y[rt][e] = (device const float2 *) (mul_mv_mma_src1_row(tile, args, src1, rt, e) + 8*(tile.fm/2) + 2*(tile.fm%2));
+        }
+    }
+
+    float acc[RT][NT][2] = {};
+
+    const int np = FC_mul_mv_mma_ne00/64;
+
+    mul_mv_mma_q5_K_a an[NT];
+    float2 xn[RT][2][2][2];
+
+    const int ip0 = min((int) sgitg, np - 1);
+    load_q5_K_mma_a<NT>(x, ip0, tile.fn, an);
+    load_q5_K_mma_b<RT>(y, ip0, xn);
+
+    for (int ip = sgitg; ip < np; ip += NSG) {
+        const short p = ip%pairs;
+
+        mul_mv_mma_q5_K_a ac[NT];
+        float2 xs[RT][2][2][2];
+        FOR_UNROLL (short t = 0; t < NT; ++t) {
+            ac[t] = an[t];
+        }
+        FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+            FOR_UNROLL (short e = 0; e < 2; ++e) {
+                FOR_UNROLL (short h = 0; h < 2; ++h) {
+                    xs[rt][e][h][0] = xn[rt][e][h][0];
+                    xs[rt][e][h][1] = xn[rt][e][h][1];
+                }
+            }
+        }
+
+        const int ipn = min(ip + NSG, np - 1);
+        load_q5_K_mma_a<NT>(x, ipn, tile.fn, an);
+        load_q5_K_mma_b<RT>(y, ipn, xn);
+
+        float c[RT][2][2];
+        FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+            FOR_UNROLL (short e = 0; e < 2; ++e) {
+                FOR_UNROLL (short h = 0; h < 2; ++h) {
+                    float u = xs[rt][e][h][0].x + xs[rt][e][h][0].y + xs[rt][e][h][1].x + xs[rt][e][h][1].y;
+                    u += simd_shuffle_xor(u, 2);
+                    u += simd_shuffle_xor(u, 4);
+                    u += simd_shuffle_xor(u, 16);
+                    c[rt][e][h] = u;
+                }
+            }
+        }
+
+        FOR_UNROLL (short t = 0; t < NT; ++t) {
+            const float2 dm  = float2(as_type<half2>(ac[t].dm));
+            const float2 sm0 = mul_mv_mma_q5_K_scale_min(ac[t].sc, 2*p + 0);
+            const float2 sm1 = mul_mv_mma_q5_K_scale_min(ac[t].sc, 2*p + 1);
+
+            half2 a[2][2][2];
+            mul_mv_mma_q5_K_frags(ac[t].q.x, ac[t].h.x >> 2*p, a[0][0], a[1][0]);
+            mul_mv_mma_q5_K_frags(ac[t].q.y, ac[t].h.y >> 2*p, a[0][1], a[1][1]);
+
+            FOR_UNROLL (short hh = 0; hh < 2; ++hh) {
+                simdgroup_float8x8 mp[RT];
+                FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+                    mp[rt] = make_filled_simdgroup_matrix<float, 8>(0.0f);
+                }
+
+                FOR_UNROLL (short s = 0; s < 4; ++s) {
+                    simdgroup_half8x8 ma;
+                    ma.thread_elements()[0] = a[hh][s/2][s%2].x;
+                    ma.thread_elements()[1] = a[hh][s/2][s%2].y;
+
+                    FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+                        simdgroup_float8x8 mb;
+                        mb.thread_elements()[0] = s % 2 == 0 ? xs[rt][0][hh][s/2].x : xs[rt][0][hh][s/2].y;
+                        mb.thread_elements()[1] = s % 2 == 0 ? xs[rt][1][hh][s/2].x : xs[rt][1][hh][s/2].y;
+
+                        simdgroup_multiply_accumulate(mp[rt], ma, mb, mp[rt]);
+                    }
+                }
+
+                const float2 smh = hh == 0 ? sm0 : sm1;
+                const float  dsc = dm.x*smh.x*(hh == 0 ? 1.0f : hi_scale);
+                const float  dmn = dm.y*smh.y;
+
+                FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+                    acc[rt][t][0] = fma(dsc, mp[rt].thread_elements()[0], fma(-dmn, c[rt][0][hh], acc[rt][t][0]));
+                    acc[rt][t][1] = fma(dsc, mp[rt].thread_elements()[1], fma(-dmn, c[rt][1][hh], acc[rt][t][1]));
+                }
+            }
+        }
+    }
+
+    mul_mv_mma_store<NT, RT>(acc, args, src2, dst, shmem, tile, tiisg, sgitg);
+}
+
+template [[host_name("kernel_mul_mv_mma_q5_K_f32_nt1_rt1")]] kernel mul_mv_mma_t kernel_mul_mv_mma_q5_K_f32<1, 1>;
+template [[host_name("kernel_mul_mv_mma_q5_K_f32_nt2_rt1")]] kernel mul_mv_mma_t kernel_mul_mv_mma_q5_K_f32<2, 1>;
+template [[host_name("kernel_mul_mv_mma_q5_K_f32_nt4_rt1")]] kernel mul_mv_mma_t kernel_mul_mv_mma_q5_K_f32<4, 1>;
+template [[host_name("kernel_mul_mv_mma_q5_K_f32_nt1_rt2")]] kernel mul_mv_mma_t kernel_mul_mv_mma_q5_K_f32<1, 2>;
+template [[host_name("kernel_mul_mv_mma_q5_K_f32_nt2_rt2")]] kernel mul_mv_mma_t kernel_mul_mv_mma_q5_K_f32<2, 2>;
+template [[host_name("kernel_mul_mv_mma_q5_K_f32_nt4_rt2")]] kernel mul_mv_mma_t kernel_mul_mv_mma_q5_K_f32<4, 2>;
+
+// few-row mat-mat for any type with a 16-weight dequantizer: a lane dequantizes 16 consecutive weights of a 64-weight chunk once for all src1 rows.
+// MMA step s at MMA-k index j reads chunk weight 16*(j/2) + 8*(j%2) + s.
+template<short NT, short RT, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread float4x4 &)>
+kernel void kernel_mul_mv_mma_gen(
+        constant ggml_metal_kargs_mul_mv_ext & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        device const char * src2,
+        threadgroup  char * shmem [[threadgroup(0)]],
+        uint3  tgpig[[threadgroup_position_in_grid]],
+        ushort tiisg[[thread_index_in_simdgroup]],
+        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+    const short NSG = FC_mul_mv_mma_nsg;
+
+    const mul_mv_mma_tile tile = mul_mv_mma_tile_init<NT, RT>(args, tgpig, tiisg);
+
+    device const block_q * x[NT];
+    FOR_UNROLL (short t = 0; t < NT; ++t) {
+        x[t] = (device const block_q *) mul_mv_mma_src0_row(tile, args, src0, t);
+    }
+
+    // B lane j = fm reads chunk values 16*(fm/2) + 8*(fm%2) .. +7
+    device const float4 * y[RT][2];
+    FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+        FOR_UNROLL (short e = 0; e < 2; ++e) {
+            y[rt][e] = (device const float4 *) mul_mv_mma_src1_row(tile, args, src1, rt, e) + 4*(tile.fm/2) + 2*(tile.fm%2);
+        }
+    }
+
+    simdgroup_float8x8 mc[RT][NT];
+    FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+        FOR_UNROLL (short t = 0; t < NT; ++t) {
+            mc[rt][t] = make_filled_simdgroup_matrix<float, 8>(0.0f);
+        }
+    }
+
+    const int nch = args.ne00/64;
+
+    for (int g = sgitg; g < nch; g += NSG) {
+        simdgroup_float8x8 mb[RT][8];
+        FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+            const float4 a0 = y[rt][0][16*g + 0];
+            const float4 a1 = y[rt][0][16*g + 1];
+            const float4 b0 = y[rt][1][16*g + 0];
+            const float4 b1 = y[rt][1][16*g + 1];
+            FOR_UNROLL (short s = 0; s < 4; ++s) {
+                mb[rt][s    ].thread_elements()[0] = a0[s];
+                mb[rt][s    ].thread_elements()[1] = b0[s];
+                mb[rt][s + 4].thread_elements()[0] = a1[s];
+                mb[rt][s + 4].thread_elements()[1] = b1[s];
+            }
+        }
+
+        const int ci = 4*g + tile.fn/2;
+
+        FOR_UNROLL (short t = 0; t < NT; ++t) {
+            float4x4 w;
+            dequantize_func(x[t] + ci/nl, ci%nl, w);
+
+            FOR_UNROLL (short s = 0; s < 8; ++s) {
+                simdgroup_float8x8 ma;
+                ma.thread_elements()[0] = w[s/4    ][s%4];
+                ma.thread_elements()[1] = w[s/4 + 2][s%4];
+
+                FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+                    simdgroup_multiply_accumulate(mc[rt][t], ma, mb[rt][s], mc[rt][t]);
+                }
+            }
+        }
+    }
+
+    float acc[RT][NT][2];
+    FOR_UNROLL (short rt = 0; rt < RT; ++rt) {
+        FOR_UNROLL (short t = 0; t < NT; ++t) {
+            acc[rt][t][0] = mc[rt][t].thread_elements()[0];
+            acc[rt][t][1] = mc[rt][t].thread_elements()[1];
+        }
+    }
+
+    mul_mv_mma_store<NT, RT>(acc, args, src2, dst, shmem, tile, tiisg, sgitg);
+}
+
+#define MUL_MV_MMA_GEN(tname, bq, nl, deq) \
+template [[host_name("kernel_mul_mv_mma_" tname "_f32_nt1_rt1")]] kernel mul_mv_mma_t kernel_mul_mv_mma_gen<1, 1, bq, nl, deq>; \
+template [[host_name("kernel_mul_mv_mma_" tname "_f32_nt2_rt1")]] kernel mul_mv_mma_t kernel_mul_mv_mma_gen<2, 1, bq, nl, deq>; \
+template [[host_name("kernel_mul_mv_mma_" tname "_f32_nt4_rt1")]] kernel mul_mv_mma_t kernel_mul_mv_mma_gen<4, 1, bq, nl, deq>; \
+template [[host_name("kernel_mul_mv_mma_" tname "_f32_nt1_rt2")]] kernel mul_mv_mma_t kernel_mul_mv_mma_gen<1, 2, bq, nl, deq>; \
+template [[host_name("kernel_mul_mv_mma_" tname "_f32_nt2_rt2")]] kernel mul_mv_mma_t kernel_mul_mv_mma_gen<2, 2, bq, nl, deq>; \
+template [[host_name("kernel_mul_mv_mma_" tname "_f32_nt4_rt2")]] kernel mul_mv_mma_t kernel_mul_mv_mma_gen<4, 2, bq, nl, deq>;
+
+// q8_0 with 9..16 src1 rows: the per-block scaling above is slower than dequantizing to f32
+template [[host_name("kernel_mul_mv_mma_q8_0_f32_nt1_rt2")]] kernel mul_mv_mma_t kernel_mul_mv_mma_gen<1, 2, block_q8_0, 2, dequantize_q8_0>;
+template [[host_name("kernel_mul_mv_mma_q8_0_f32_nt2_rt2")]] kernel mul_mv_mma_t kernel_mul_mv_mma_gen<2, 2, block_q8_0, 2, dequantize_q8_0>;
+template [[host_name("kernel_mul_mv_mma_q8_0_f32_nt4_rt2")]] kernel mul_mv_mma_t kernel_mul_mv_mma_gen<4, 2, block_q8_0, 2, dequantize_q8_0>;
+
+MUL_MV_MMA_GEN("f32",  float4x4,   1,     dequantize_f32)
+MUL_MV_MMA_GEN("f16",  half4x4,    1,     dequantize_f16)
+MUL_MV_MMA_GEN("q4_1", block_q4_1, 2,     dequantize_q4_1)
+MUL_MV_MMA_GEN("q5_0", block_q5_0, 2,     dequantize_q5_0)
+MUL_MV_MMA_GEN("q5_1", block_q5_1, 2,     dequantize_q5_1)
+MUL_MV_MMA_GEN("q4_K", block_q4_K, QK_NL, dequantize_q4_K)
+MUL_MV_MMA_GEN("q6_K", block_q6_K, QK_NL, dequantize_q6_K)
+
+#undef MUL_MV_MMA_GEN

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

@@ -171,7 +171,8 @@ kernel void kernel_concat(
 
     const int i3 = tgpig.z;
     const int i2 = tgpig.y;
-    const int i1 = ntg.y == 1 ? tgpig.x : tgpig.x*ntg.y + tpitg.y;
+    const int i1 = ntg.y == 1 ? tgpig.x/args.nc0 : tgpig.x*ntg.y + tpitg.y;
+    const int ic = ntg.y == 1 ? tgpig.x%args.nc0 : 0;
 
     if (i1 >= args.ne1) {
         return;
@@ -180,7 +181,11 @@ kernel void kernel_concat(
     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) {
+    // chunk ic of nc0 along the row
+    const int n0  = (args.ne0 + args.nc0 - 1)/args.nc0;
+    const int i0e = min(args.ne0, (ic + 1)*n0);
+
+    for (int i0 = ic*n0 + tpitg.x; i0 < i0e; i0 += ntg.x) {
         device const T * x;
 
         if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
@@ -220,7 +225,8 @@ kernel void kernel_concat_q(
     // 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;
+    const int i1 = ntg.y == 1 ? tgpig.x/args.nc0 : tgpig.x*ntg.y + tpitg.y;
+    const int ic = ntg.y == 1 ? tgpig.x%args.nc0 : 0;
 
     if (i1 >= args.ne1) {
         return;
@@ -229,7 +235,11 @@ kernel void kernel_concat_q(
     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) {
+    // chunk ic of nc0 along the row
+    const int n0  = (args.ne0 + args.nc0 - 1)/args.nc0;
+    const int i0e = min(args.ne0, (ic + 1)*n0);
+
+    for (int i0 = ic*n0 + tpitg.x; i0 < i0e; i0 += ntg.x) {
         device const block_q * x;
 
         if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {

ggml/src/ggml-vulkan/ggml-vulkan-common.h

@@ -233,6 +233,18 @@ ggml_backend_reg_t ggml_backend_vk_reg();
 
 // debug
 int64_t ggml_vk_get_op_batch_size(const ggml_tensor * op);
+#ifdef GGML_VULKAN_RUN_TESTS
+template <typename X_TYPE, typename Y_TYPE>
+void ggml_vk_test_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, int split_k, int shader_size);
+void ggml_vk_test_dequant(ggml_backend_vk_context * ctx, size_t ne, ggml_type quant);
+void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, size_t split_k, size_t shader_size, ggml_type quant, bool mmq = false);
+#endif
+#ifdef GGML_VULKAN_CHECK_RESULTS
+extern size_t vk_skip_checks;
+extern size_t vk_output_tensor;
+void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, int tensor_idx);
+void ggml_vk_check_results_1(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, int tensor_idx);
+#endif
 
 // ggml-vulkan.cpp (residual)
 bool ggml_vk_lightning_indexer_k_type_supported(ggml_type type);

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

@@ -47,12 +47,10 @@ void vk_memory_logger::log_deallocation(vk_buffer_ref buf_ref) {
 }
 
 #ifdef GGML_VULKAN_CHECK_RESULTS
-static size_t vk_skip_checks;
-static size_t vk_output_tensor;
+size_t vk_skip_checks;
+size_t vk_output_tensor;
 
 static void ggml_vk_print_tensor(const ggml_tensor * tensor, const char * name);
-static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, int tensor_idx);
-static void ggml_vk_check_results_1(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, int tensor_idx);
 #endif
 
 #ifdef GGML_VULKAN_RUN_TESTS
@@ -90,7 +88,7 @@ static void ggml_vk_print_matrix_area(const void * data, ggml_type type, int ne0
 }
 
 template <typename X_TYPE, typename Y_TYPE>
-static void ggml_vk_test_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, int split_k, int shader_size) {
+void ggml_vk_test_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, int split_k, int shader_size) {
     VK_LOG_DEBUG("ggml_vk_test_matmul(" << m << ", " << n << ", " << k << ", " << batch << ", " << num_it << ", " << split_k << ", " << shader_size << ")");
     const size_t x_ne = m * k * batch;
     const size_t y_ne = k * n * batch;
@@ -297,6 +295,8 @@ static void ggml_vk_test_matmul(ggml_backend_vk_context * ctx, size_t m, size_t
     free(d);
 }
 
+template void ggml_vk_test_matmul<ggml_fp16_t, float>(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, int split_k, int shader_size);
+
 static void ggml_vk_print_tensor_area(const ggml_tensor * tensor, int i0, int i1, int i2, int i3) {
     if (tensor->type != GGML_TYPE_F32 && tensor->type != GGML_TYPE_F16) {
         return;
@@ -348,7 +348,7 @@ static void ggml_vk_dequantize_data(const void * from, float * to, size_t ne, gg
     dequant_fn(from, to, ne);
 }
 
-static void ggml_vk_test_dequant(ggml_backend_vk_context * ctx, size_t ne, ggml_type quant) {
+void ggml_vk_test_dequant(ggml_backend_vk_context * ctx, size_t ne, ggml_type quant) {
     VK_LOG_DEBUG("ggml_vk_test_dequant(" << ne << ")");
     const size_t x_sz = sizeof(float) * ne;
     const size_t x_sz_f16 = sizeof(ggml_fp16_t) * ne;
@@ -540,7 +540,7 @@ static void ggml_vk_test_dequant(ggml_backend_vk_context * ctx, size_t ne, ggml_
 //     free(qx_res);
 // }
 
-static void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, size_t split_k, size_t shader_size, ggml_type quant, bool mmq = false) {
+void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size_t k, size_t batch, size_t num_it, size_t split_k, size_t shader_size, ggml_type quant, bool mmq) {
     VK_LOG_DEBUG("ggml_vk_test_dequant_matmul(" << m << ", " << n << ", " << k << ", " << batch << ", " << num_it << ", " << split_k << ", " << ggml_type_name(quant) << ")");
     const size_t x_ne = m * k * batch;
     const size_t y_ne = k * n * batch;
@@ -871,7 +871,7 @@ void * comp_result;
 size_t comp_size;
 size_t comp_nb[GGML_MAX_DIMS];
 size_t check_counter = 0;
-static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, int tensor_idx) {
+void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, int tensor_idx) {
     ggml_tensor * tensor = cgraph->nodes[tensor_idx + ctx->num_additional_fused_ops];
     if (tensor->op == GGML_OP_TRANSPOSE || tensor->op == GGML_OP_SET_ROWS) {
         return;
@@ -1374,7 +1374,7 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
     VK_LOG_DEBUG("END ggml_vk_check_results_0(" << tensor->name << ")");
 }
 
-static void ggml_vk_check_results_1(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, int tensor_idx) {
+void ggml_vk_check_results_1(ggml_backend_vk_context * ctx, ggml_cgraph * cgraph, int tensor_idx) {
     ggml_tensor * tensor = cgraph->nodes[tensor_idx + ctx->num_additional_fused_ops];
     if (tensor->op == GGML_OP_TRANSPOSE || tensor->op == GGML_OP_SET_ROWS) {
         return;

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

@@ -7428,11 +7428,10 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
 
     GGML_ASSERT(mmp_map != nullptr);
 
-    const uint32_t n_per_expert = (uint32_t)CEIL_DIV(nei0 * nei1, n_as);
-    const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_pipeline_align_map(ctx, *mmp_map, ne01, n_per_expert, true));
-    const bool aligned = !quantize_y && ne10 == kpad && ne01 > 8 && n_per_expert > 8;
+    const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_pipeline_align_map(ctx, *mmp_map, ne01, nei1, true));
+    const bool aligned = !quantize_y && ne10 == kpad && ne01 > 8 && nei1 > 8;
 
-    vk_pipeline pipeline = ggml_vk_guess_matmul_pipeline_map(ctx, *mmp_map, ne01, n_per_expert, aligned, true);
+    vk_pipeline pipeline = ggml_vk_guess_matmul_pipeline_map(ctx, *mmp_map, ne01, nei1, aligned, true);
 
     if (ggml_nbytes(src0) > ctx->device->properties.limits.maxStorageBufferRange) {
         pipeline = ggml_vk_get_64b_indexing_pipeline(ctx, pipeline);
@@ -8145,11 +8144,14 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
     // Sparse mask hint (op_params[4]): compact the <= n_kv_max finite positions and gather only those.
     const int32_t n_kv_max = mask ? ggml_get_op_params_i32(dst, 4) : 0;
     static const bool disable_sparse = getenv("GGML_VK_FA_SPARSE_DISABLE") != nullptr;
+    const bool kv_f16 = k_type_eff == GGML_TYPE_F16 && v_type_eff == GGML_TYPE_F16;
     // cm2 dense is fast, so it needs a larger reduction to win.
-    const int64_t min_ratio = tuning_params.path == FA_COOPMAT2 ? 4 : 2;
+    // With quantized K/V, sparse only breaks even around 16x (measured on RDNA3/RDNA4).
+    const int64_t min_ratio = tuning_params.path == FA_COOPMAT2 ? 4 : (kv_f16 ? 2 : 16);
     const bool use_sparse = !disable_sparse && n_kv_max > 0 && mask &&
                             max_bias == 0.0f && logit_softcap == 0.0f &&
-                            k_type_eff == GGML_TYPE_F16 && v_type_eff == GGML_TYPE_F16 &&
+                            // the cm2 sparse gather only reads f16
+                            (kv_f16 || tuning_params.path != FA_COOPMAT2) &&
                             nem0 == KV &&
                             (int64_t)KV >= std::max<int64_t>(4096, min_ratio * (int64_t)n_kv_max) &&
                             (gqa_ratio > 1 || (tuning_params.path == FA_SCALAR && N == 1));
@@ -8373,6 +8375,11 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
     vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf;
     vk_subbuffer mask_opt_buf = use_mask_opt ? ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0) : q_buf;
     vk_subbuffer sparse_buf = use_sparse ? ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0) : q_buf;
+    if (use_mask_opt || use_sparse) {
+        // the mask opt bits and the sparse index list overwrite a matmul input converted into prealloc_y
+        ctx->prealloc_y_last_pipeline_used = nullptr;
+        ctx->prealloc_y_last_tensor_used = nullptr;
+    }
 
     if (use_dequant_kv) {
         const uint64_t fp = sizeof(ggml_fp16_t);
@@ -11056,6 +11063,10 @@ void ggml_vk_soft_max(ggml_backend_vk_context * ctx, vk_context& subctx, const g
         vk_subbuffer buf_x = { ctx->prealloc_x, 0, tmp_size };
         vk_subbuffer buf_y = { ctx->prealloc_y, 0, tmp_size };
 
+        // the partial results overwrite a matmul input converted into prealloc_y
+        ctx->prealloc_y_last_pipeline_used = nullptr;
+        ctx->prealloc_y_last_tensor_used = nullptr;
+
         std::array<uint32_t, 3> elements = { num_wgs, nrows_x, 1 };
 
         vk_pipeline pipeline1 = src1 && src1->type == GGML_TYPE_F16 ? ctx->device->pipeline_soft_max_large1_f32_f16 : ctx->device->pipeline_soft_max_large1_f32;

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

@@ -130,20 +130,15 @@ void main() {
         const float qd_inv = qd != 0.0f ? 1.0f / qd : 0.0f;
         vals = round(vals * qd_inv);
 
-        Qf[buf_ib].qs[buf_iqs] = pack32(i8vec4(vals));
+        const float thread_sum = vals.x + vals.y + vals.z + vals.w;
+        // Q8_0 K does not use the row sum.
+        const float sum = FaTypeK == GGML_TYPE_Q8_0 ? 0.0f : subgroupClusteredAdd(thread_sum, 8);
 
-        // Q8_0 K only needs (qd, _); the asymmetric Q4_*/Q5_* family also stores
-        // the row-sum scaled by qd, used in k_dot_correction.
-        if (FaTypeK == GGML_TYPE_Q8_0) {
-            if (buf_iqs == 0) {
-                Qf[buf_ib].ds = FLOAT_TYPEV2(qd, 0.0f);
-            }
-        } else {
-            const float thread_sum = vals.x + vals.y + vals.z + vals.w;
-            const float sum = subgroupClusteredAdd(thread_sum, 8);
+        if (r < Br) {
+            Qf[buf_ib].qs[buf_iqs] = pack32(i8vec4(vals));
 
             if (buf_iqs == 0) {
-                Qf[buf_ib].ds = FLOAT_TYPEV2(qd, sum * qd);
+                Qf[buf_ib].ds = FaTypeK == GGML_TYPE_Q8_0 ? FLOAT_TYPEV2(qd, 0.0f) : FLOAT_TYPEV2(qd, sum * qd);
             }
         }
 #endif
@@ -285,8 +280,9 @@ void main() {
                 const uint32_t block = ib % (HSK / 32);
                 if (idx + gl_WorkGroupSize.x <= quant_iters || c < Bc) {
                     const uint buf_ib = c * qf_stride + block;
-                    if (!KV_bounds_check || j * Bc + c < KV) {
-                        const uint global_ib = (j * Bc + c) * k_stride + block;
+                    uint32_t kcol;
+                    if (fa_kv_index(j * Bc + c, kcol)) {
+                        const uint global_ib = kcol * k_stride + block;
                         k_block_to_shmem(buf_ib, global_ib, iqs, k_offset);
                     } else {
                         k_block_to_shmem_zero(buf_ib, iqs);
@@ -363,7 +359,8 @@ void main() {
                                 (hsk4 % 2 == 0) ? 2 : 1;
 
         [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
-            if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) {
+            uint32_t kcol;
+            if (!fa_kv_index(j * Bc + c * cols_per_iter + col_tid, kcol)) {
                 continue;
             }
 
@@ -400,7 +397,7 @@ void main() {
                         }
                     }
                 } else {
-                    const uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d_tid * (HSK_per_thread / 4) + d_block);
+                    const uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * (d_tid * (HSK_per_thread / 4) + d_block);
                     const uint ib = coord / BLOCK_SIZE_K;
                     const uint iqs = (coord % BLOCK_SIZE_K);
 

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

@@ -26,14 +26,22 @@ layout (push_constant) uniform parameter {
 } p;
 
 #ifdef USE_SUBGROUPS
-shared uvec4 ballots_sh[NUM_SUBGROUPS];
+shared uint counts_sh[NUM_SUBGROUPS];
 #else
 shared uint scan[BLOCK_SIZE];
 #endif
 
+bool is_selected(const uint m_idx) {
+    const float v = float(data_m[m_idx]);
+    return !isinf(v) && !isnan(v);
+}
+
 // One workgroup per mask row: compact the finite-mask KV positions into a
-// per-row index list of length n_kv_max, -1 padded. Emitted in ascending KV
-// order so the downstream attention accumulation is deterministic.
+// per-row index list of length n_kv_max, -1 padded, in ascending KV order so
+// the downstream attention accumulation is deterministic.
+// The row is split into contiguous segments, one per subgroup (or per thread
+// without subgroups), so it needs a single workgroup scan instead of one per
+// BLOCK_SIZE chunk.
 void main() {
     const uint i1 = gl_WorkGroupID.x;
     const uint i2 = gl_WorkGroupID.y;
@@ -43,60 +51,75 @@ void main() {
     const uint m_base   = i3 * p.nbm3 + i2 * p.nbm2 + i1 * p.nbm1;
     const uint out_base = ((i3 * p.nem2 + i2) * p.nem1 + i1) * p.n_kv_max;
 
-    uint base = 0;
-    for (uint chunk = 0; chunk < p.KV; chunk += BLOCK_SIZE) {
-        const uint k = chunk + tid;
-        bool selected = false;
-        if (k < p.KV) {
-            const float v = float(data_m[m_base + k]);
-            selected = !isinf(v) && !isnan(v);
-        }
-
 #ifdef USE_SUBGROUPS
+    const uint sg   = gl_SubgroupID;
+    const uint lane = gl_SubgroupInvocationID;
+    const uint seg  = (p.KV + gl_NumSubgroups - 1) / gl_NumSubgroups;
+    const uint seg_begin = min(sg * seg, p.KV);
+    const uint seg_end   = min(seg_begin + seg, p.KV);
+
+    // Lanes read consecutive positions, so each step is one coalesced load.
+    uint count = 0;
+    for (uint k0 = seg_begin; k0 < seg_end; k0 += gl_SubgroupSize) {
+        const uint k = k0 + lane;
+        count += subgroupBallotBitCount(subgroupBallot(k < seg_end && is_selected(m_base + k)));
+    }
+    if (subgroupElect()) {
+        counts_sh[sg] = count;
+    }
+    barrier();
+
+    uint slot = 0;
+    uint total = 0;
+    for (uint s = 0; s < gl_NumSubgroups; ++s) {
+        slot  += s < sg ? counts_sh[s] : 0u;
+        total += counts_sh[s];
+    }
+
+    for (uint k0 = seg_begin; k0 < seg_end && slot < p.n_kv_max; k0 += gl_SubgroupSize) {
+        const uint k = k0 + lane;
+        const bool selected = k < seg_end && is_selected(m_base + k);
         const uvec4 ballot = subgroupBallot(selected);
-        if (subgroupElect()) {
-            ballots_sh[gl_SubgroupID] = ballot;
+        const uint pos = slot + subgroupBallotExclusiveBitCount(ballot);
+        if (selected && pos < p.n_kv_max) {
+            data_i[out_base + pos] = int32_t(k);
         }
-        barrier();
-
-        uint subgroup_base = 0;
-        uint total = 0;
-        [[unroll]] for (uint s = 0; s < gl_NumSubgroups; ++s) {
-            if (s == gl_SubgroupID) {
-                subgroup_base = total;
-            }
-            total += subgroupBallotBitCount(ballots_sh[s]);
-        }
-        barrier();
-
-        const uint slot = base + subgroup_base + subgroupBallotExclusiveBitCount(ballot);
+        slot += subgroupBallotBitCount(ballot);
+    }
 #else
-        // Hillis-Steele inclusive prefix sum over the workgroup.
-        scan[tid] = selected ? 1u : 0u;
+    const uint run   = (p.KV + BLOCK_SIZE - 1) / BLOCK_SIZE;
+    const uint begin = min(tid * run, p.KV);
+    const uint end   = min(begin + run, p.KV);
+
+    uint count = 0;
+    for (uint k = begin; k < end; ++k) {
+        count += is_selected(m_base + k) ? 1u : 0u;
+    }
+
+    // Hillis-Steele inclusive prefix sum of the per-thread counts.
+    scan[tid] = count;
+    barrier();
+    for (uint off = 1; off < BLOCK_SIZE; off <<= 1) {
+        uint add = 0;
+        if (tid >= off) {
+            add = scan[tid - off];
+        }
         barrier();
-        for (uint off = 1; off < BLOCK_SIZE; off <<= 1) {
-            uint add = 0;
-            if (tid >= off) {
-                add = scan[tid - off];
-            }
-            barrier();
-            scan[tid] += add;
-            barrier();
-        }
-
-        const uint inclusive = scan[tid];
-        const uint total = scan[BLOCK_SIZE - 1];
-        const uint slot = base + inclusive - 1u;
-#endif
-
-        if (selected && slot < p.n_kv_max) {
-            data_i[out_base + slot] = int32_t(k);
-        }
-        base += total;
+        scan[tid] += add;
         barrier();
     }
 
-    for (uint s = min(base, p.n_kv_max) + tid; s < p.n_kv_max; s += BLOCK_SIZE) {
+    const uint total = scan[BLOCK_SIZE - 1];
+    uint slot = scan[tid] - count;
+    for (uint k = begin; k < end && slot < p.n_kv_max; ++k) {
+        if (is_selected(m_base + k)) {
+            data_i[out_base + slot] = int32_t(k);
+            ++slot;
+        }
+    }
+#endif
+
+    for (uint s = min(total, p.n_kv_max) + tid; s < p.n_kv_max; s += BLOCK_SIZE) {
         data_i[out_base + s] = int32_t(-1);
     }
 }

ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp

@@ -1181,11 +1181,18 @@ inline bool ggml_webgpu_can_use_mmvq(const ggml_tensor * src0,
             switch (src1->type) {
                 case GGML_TYPE_F32:
                     switch (src0->type) {
+                        case GGML_TYPE_Q1_0:
                         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:
+                        case GGML_TYPE_MXFP4:
                         case GGML_TYPE_Q2_K:
+                        case GGML_TYPE_Q3_K:
                         case GGML_TYPE_Q4_K:
+                        case GGML_TYPE_Q5_K:
+                        case GGML_TYPE_Q6_K:
                             return src0->ne[0] % 4 == 0;
                         default:
                             break;
@@ -2036,17 +2043,23 @@ class ggml_webgpu_shader_lib {
                     defines.push_back("U32_DEQUANT_HELPERS");
                     defines.push_back("SRC0_INNER_TYPE=u32");
                     switch (context.src0->type) {
-                        case GGML_TYPE_Q8_0:
                         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:
                             if (key.use_mmvq) {
-                                defines.push_back("LEGACY_QUANTS");
+                                defines.push_back("LEGACY_QUANTS_HANDLING");
                             }
                             break;
+                        case GGML_TYPE_Q1_0:
                         case GGML_TYPE_Q2_K:
+                        case GGML_TYPE_Q3_K:
                         case GGML_TYPE_Q4_K:
+                        case GGML_TYPE_Q5_K:
+                        case GGML_TYPE_Q6_K:
                             if (key.use_mmvq) {
-                                defines.push_back("K_QUANTS");
+                                defines.push_back("K_QUANTS_HANDLING");
                             }
                             break;
                         case GGML_TYPE_IQ1_S:
@@ -2064,6 +2077,11 @@ class ggml_webgpu_shader_lib {
                             defines.push_back(type_upper + "_TABLES");
                             break;
                         case GGML_TYPE_MXFP4:
+                            defines.push_back(type_upper + "_LUT");
+                            if (key.use_mmvq) {
+                                defines.push_back("LEGACY_QUANTS_HANDLING");
+                            }
+                            break;
                         case GGML_TYPE_NVFP4:
                             defines.push_back(type_upper + "_LUT");
                             break;

ggml/src/ggml-webgpu/wgsl-shaders/mul_mat_vec_q_acc.tmpl

@@ -14,10 +14,13 @@ fn sbyte_of(v: u32, b: u32) -> i32 {
 #define SRC0_TYPE SRC0_INNER_TYPE
 #define SRC1_TYPE SRC1_INNER_TYPE
 
-#ifdef LEGACY_QUANTS
+#ifdef LEGACY_QUANTS_HANDLING
 #define BLOCK_SIZE 32
 #define THREADS_PER_BLOCK 4
-#elif K_QUANTS
+#elif defined(MUL_ACC_Q1_0)
+#define BLOCK_SIZE 128
+#define THREADS_PER_BLOCK 8
+#elif defined(K_QUANTS_HANDLING)
 #define BLOCK_SIZE 256
 #define THREADS_PER_BLOCK 16
 #endif
@@ -25,6 +28,22 @@ fn sbyte_of(v: u32, b: u32) -> i32 {
 #define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
 #define Q8_BLOCK_SIZE 32
 
+#if (defined(LEGACY_QUANTS_HANDLING) || defined(MUL_ACC_MXFP4)) && !defined(MUL_ACC_Q8_0)
+fn repack_b_qs(block:u32, inner_id: u32) -> vec2<u32> {
+    return vec2<u32>(
+            src1q[block].qs[inner_id],
+            src1q[block].qs[inner_id + 4u],
+        );
+}
+#endif
+
+#if defined(MUL_ACC_Q5_0) || defined(MUL_ACC_Q5_1)
+fn qh_bits(qh: u32, shift: u32) -> u32 {
+    // multiply by 0x00204081 moves bit b to bit 8*b
+    return ((((qh >> shift) & 0xFu) * 0x00204081u) & 0x01010101u) << 4u;
+}
+#endif
+
 #ifdef MUL_ACC_Q4_0
 #define BLOCK_SIZE_BYTES 18
 #define B_DS_TYPE vec2<f32>
@@ -36,12 +55,6 @@ fn repack_a(block_byte_base: u32, inner_id: u32) -> vec2<u32> {
         (qs_packed >> 4u) & 0x0F0F0F0Fu
     );
 }
-fn repack_b_qs(block:u32, inner_id: u32) -> vec2<u32> {
-    return vec2<u32>(
-            src1q[block].qs[inner_id],
-            src1q[block].qs[inner_id + 4u],
-        );
-}
 fn repack_b_dm(block: u32) -> B_DS_TYPE {
     return B_DS_TYPE(
         f32(src1q[block].d),
@@ -64,12 +77,6 @@ fn repack_a(block_byte_base: u32, inner_id: u32) -> vec2<u32> {
         (qs_packed >> 4u) & 0x0F0F0F0Fu
     );
 }
-fn repack_b_qs(block:u32, inner_id: u32) -> vec2<u32> {
-    return vec2<u32>(
-            src1q[block].qs[inner_id],
-            src1q[block].qs[inner_id + 4u],
-        );
-}
 fn repack_b_dm(block: u32) -> B_DS_TYPE {
     return B_DS_TYPE(
         f32(src1q[block].d),
@@ -84,6 +91,78 @@ fn get_dm(block_byte_base: u32) -> vec2<f32> {
 }
 #endif // MUL_ACC_Q4_1
 
+#ifdef MUL_ACC_Q5_0
+#define BLOCK_SIZE_BYTES 22
+#define B_DS_TYPE vec2<f32>
+fn repack_a(block_byte_base: u32, inner_id: u32) -> vec2<u32> {
+    let qh        = load_u32_at_src0(block_byte_base + 2u);
+    let qs_packed = load_u32_at_src0(block_byte_base + 6u + 4u * inner_id);
+
+    return vec2<u32>(
+        (qs_packed & 0x0F0F0F0Fu) | qh_bits(qh, 4u * inner_id),
+        ((qs_packed >> 4u) & 0x0F0F0F0Fu) | qh_bits(qh, 16u + 4u * inner_id)
+    );
+}
+fn repack_b_dm(block: u32) -> B_DS_TYPE {
+    return B_DS_TYPE(
+        f32(src1q[block].d),
+        f32(src1q[block].s)
+    );
+}
+fn get_dm(block_byte_base: u32) -> f32 {
+    return f32(load_f16_at_src0(block_byte_base));
+}
+#endif // MUL_ACC_Q5_0
+
+#ifdef MUL_ACC_Q5_1
+#define BLOCK_SIZE_BYTES 24
+#define B_DS_TYPE vec2<f32>
+fn repack_a(block_byte_base: u32, inner_id: u32) -> vec2<u32> {
+    let qh        = load_u32_at_src0(block_byte_base + 4u);
+    let qs_packed = load_u32_at_src0(block_byte_base + 8u + 4u * inner_id);
+
+    return vec2<u32>(
+        (qs_packed & 0x0F0F0F0Fu) | qh_bits(qh, 4u * inner_id),
+        ((qs_packed >> 4u) & 0x0F0F0F0Fu) | qh_bits(qh, 16u + 4u * inner_id)
+    );
+}
+fn repack_b_dm(block: u32) -> B_DS_TYPE {
+    return B_DS_TYPE(
+        f32(src1q[block].d),
+        f32(src1q[block].s)
+    );
+}
+fn get_dm(block_byte_base: u32) -> vec2<f32> {
+    return vec2<f32>(
+        f32(load_f16_at_src0(block_byte_base)),
+        f32(load_f16_at_src0(block_byte_base + 2u))
+    );
+}
+#endif // MUL_ACC_Q5_1
+
+#ifdef MUL_ACC_MXFP4
+#define BLOCK_SIZE_BYTES 17
+#define B_DS_TYPE f32
+fn repack_a(block_byte_base: u32, inner_id: u32) -> vec2<u32> {
+    let qs_packed = load_u32_at_src0(block_byte_base + 1u + 4u * inner_id);
+
+    var lo = 0u;
+    var hi = 0u;
+    for (var b = 0u; b < 4u; b++) {
+        let q_byte = byte_of(qs_packed, b);
+        lo |= (bitcast<u32>(kvalues_mxfp4[q_byte & 0xFu]) & 0xFFu) << (8u * b);
+        hi |= (bitcast<u32>(kvalues_mxfp4[q_byte >> 4u]) & 0xFFu) << (8u * b);
+    }
+    return vec2<u32>(lo, hi);
+}
+fn repack_b_dm(block: u32) -> B_DS_TYPE {
+    return B_DS_TYPE(src1q[block].d);
+}
+fn get_dm(block_byte_base: u32) -> f32 {
+    return ldexp(1.0, i32(byte_of(load_u32_at_src0(block_byte_base), 0u)) - 128);
+}
+#endif // MUL_ACC_MXFP4
+
 #ifdef MUL_ACC_Q8_0
 #define BLOCK_SIZE_BYTES 34
 #define B_DS_TYPE f32
@@ -107,7 +186,7 @@ fn get_dm(block_byte_base: u32) -> f32 {
 }
 #endif // MUL_ACC_Q8_0
 
-#if defined(LEGACY_QUANTS)
+#if defined(LEGACY_QUANTS_HANDLING)
 fn accumulate_vec_q_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1q_idx_base: u32) -> array<array<f32, OUTPUTS_PER_WG>, NUM_COLS> {
     var acc: array<array<f32, OUTPUTS_PER_WG>, NUM_COLS>;
 
@@ -115,6 +194,13 @@ fn accumulate_vec_q_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, s
 
     for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
         let inner_id = thread_id % THREADS_PER_BLOCK;
+        var b_qs_cols: array<vec2<u32>, NUM_COLS>;
+        var b_ds_cols: array<B_DS_TYPE, NUM_COLS>;
+        for (var col = 0u;col < NUM_COLS;col += 1) {
+            let src1q_idx = src1q_idx_base + col * (params.k / Q8_BLOCK_SIZE) + block;
+            b_qs_cols[col] = repack_b_qs(src1q_idx, inner_id);
+            b_ds_cols[col] = repack_b_dm(src1q_idx);
+        }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
             let output_row = row_base + row;
             if (output_row < params.m) {
@@ -122,9 +208,8 @@ fn accumulate_vec_q_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, s
                 let a_repacked = repack_a(block_byte_base, inner_id);
                 let da = get_dm(block_byte_base);
                 for (var col = 0u;col < NUM_COLS;col += 1) {
-                    let src1q_idx = src1q_idx_base + col * (params.k / Q8_BLOCK_SIZE) + block;
-                    let b_repacked = repack_b_qs(src1q_idx, inner_id);
-                    let b_ds = repack_b_dm(src1q_idx);
+                    let b_repacked = b_qs_cols[col];
+                    let b_ds = b_ds_cols[col];
 
                     let row_sum = dot4I8Packed(a_repacked[0], b_repacked[0]) + dot4I8Packed(a_repacked[1], b_repacked[1]);
 
@@ -132,13 +217,17 @@ fn accumulate_vec_q_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, s
                     acc[col][row] += f32(row_sum) * (da * b_ds.x) - 8.0 * da * b_ds.y / THREADS_PER_BLOCK;
 #endif // MUL_ACC_Q4_0
 
-#if defined(MUL_ACC_Q4_1)
-                    acc[col][row] += f32(row_sum) * (da.x * b_ds.x) + da.y * b_ds.y / THREADS_PER_BLOCK;
-#endif // MUL_ACC_Q4_1
+#if defined(MUL_ACC_Q5_0)
+                    acc[col][row] += f32(row_sum) * (da * b_ds.x) - 16.0 * da * b_ds.y / THREADS_PER_BLOCK;
+#endif // MUL_ACC_Q5_0
 
-#if defined(MUL_ACC_Q8_0)
+#if defined(MUL_ACC_Q4_1) || defined(MUL_ACC_Q5_1)
+                    acc[col][row] += f32(row_sum) * (da.x * b_ds.x) + da.y * b_ds.y / THREADS_PER_BLOCK;
+#endif // MUL_ACC_Q4_1 || MUL_ACC_Q5_1
+
+#if defined(MUL_ACC_Q8_0) || defined(MUL_ACC_MXFP4)
                     acc[col][row] += f32(row_sum) * (da * b_ds);
-#endif // MUL_ACC_Q8_0
+#endif // MUL_ACC_Q8_0 || MUL_ACC_MXFP4
                 }
             }
         }
@@ -146,7 +235,49 @@ fn accumulate_vec_q_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, s
 
     return acc;
 }
-#endif // LEGACY_QUANTS
+#endif // LEGACY_QUANTS_HANDLING
+
+// every k-quant thread covers 16 elements
+#if defined(K_QUANTS_HANDLING)
+fn repack_b_qs(q8_block_idx: u32, tid: u32) -> vec4<u32> {
+    let phase = tid % 2u;
+    return vec4<u32>(
+        src1q[q8_block_idx].qs[4u * phase],
+        src1q[q8_block_idx].qs[4u * phase + 1u],
+        src1q[q8_block_idx].qs[4u * phase + 2u],
+        src1q[q8_block_idx].qs[4u * phase + 3u],
+    );
+}
+#endif
+
+#if defined(MUL_ACC_Q1_0) || defined(MUL_ACC_Q3_K) || defined(MUL_ACC_Q6_K)
+// subtract c from every byte of v, giving 4 packed i8; each byte of v must be below 128
+fn sub_packed_bytes(v: u32, c: u32) -> u32 {
+    return (v + (0x80u - c) * 0x01010101u) ^ 0x80808080u;
+}
+#endif
+
+#ifdef MUL_ACC_Q1_0
+#define BLOCK_SIZE_BYTES 18
+#define B_DS_TYPE f32
+fn repack_a(block_byte_base: u32, tid: u32) -> vec4<u32> {
+    let bits = load_u16_at_src0(block_byte_base + 2u + 2u * tid);
+
+    var res: vec4<u32>;
+    for (var i = 0u; i < 4u; i++) {
+        // spread 4 bits to bit 1 of each byte
+        let twice_bits = ((((bits >> (4u * i)) & 0xFu) * 0x00204081u) & 0x01010101u) << 1u;
+        res[i] = sub_packed_bytes(twice_bits, 1u);
+    }
+    return res;
+}
+fn repack_b_dm(q8_block_idx: u32) -> B_DS_TYPE {
+    return B_DS_TYPE(src1q[q8_block_idx].d);
+}
+fn get_dm(block_byte_base: u32) -> f32 {
+    return f32(load_f16_at_src0(block_byte_base));
+}
+#endif // MUL_ACC_Q1_0
 
 #ifdef MUL_ACC_Q2_K
 #define BLOCK_SIZE_BYTES 84
@@ -164,15 +295,6 @@ fn repack_a(block_byte_base: u32, tid: u32) -> vec4<u32> {
         (load_u32_at_src0_aligned(qs_byte_base + 12u) >> qs_shift) & 0x03030303u,
     );
 }
-fn repack_b_qs(q8_block_idx: u32, tid: u32) -> vec4<u32> {
-    let phase = tid % 2u;
-    return vec4<u32>(
-        src1q[q8_block_idx].qs[4u * phase],
-        src1q[q8_block_idx].qs[4u * phase + 1u],
-        src1q[q8_block_idx].qs[4u * phase + 2u],
-        src1q[q8_block_idx].qs[4u * phase + 3u],
-    );
-}
 fn repack_b_dm(q8_block_idx: u32) -> B_DS_TYPE {
     return B_DS_TYPE(src1q[q8_block_idx].d);
 }
@@ -189,31 +311,51 @@ fn get_scale_min(block_byte_base: u32, tid: u32) -> vec2<f32> {
 }
 #endif // MUL_ACC_Q2_K
 
-#ifdef MUL_ACC_Q4_K
-#define BLOCK_SIZE_BYTES 144
-#define B_DS_TYPE vec2<f32>
+#ifdef MUL_ACC_Q3_K
+#define BLOCK_SIZE_BYTES 110
+#define B_DS_TYPE f32
 fn repack_a(block_byte_base: u32, tid: u32) -> vec4<u32> {
-    let iq4 = tid / 4u;
-    let phase = tid % 2u;
-    let nibble = (tid >> 1u) % 2u;
-    let q_qs_byte_base = block_byte_base + 16u + 32u * iq4 + 16u * phase;
-    let qs_shift = 4u * nibble;
-    return vec4<u32>(
-        (load_u32_at_src0_aligned(q_qs_byte_base) >> qs_shift) & 0x0F0F0F0Fu,
-        (load_u32_at_src0_aligned(q_qs_byte_base + 4u) >> qs_shift) & 0x0F0F0F0Fu,
-        (load_u32_at_src0_aligned(q_qs_byte_base + 8u) >> qs_shift) & 0x0F0F0F0Fu,
-        (load_u32_at_src0_aligned(q_qs_byte_base + 12u) >> qs_shift) & 0x0F0F0F0Fu,
-    );
+    let half_blk = tid / 8u;
+    let sub      = (tid % 8u) / 2u;
+    let phase    = tid % 2u;
+    let qs_byte_base = block_byte_base + 32u + 32u * half_blk + 16u * phase;
+    let hm_byte_base = block_byte_base + 16u * phase;
+    let qs_shift = 2u * sub;
+    let hm_shift = 4u * half_blk + sub;
+
+    var res: vec4<u32>;
+    for (var i = 0u; i < 4u; i++) {
+        let qs = (load_u32_at_src0(qs_byte_base + 4u * i) >> qs_shift) & 0x03030303u;
+        let hm = (load_u32_at_src0(hm_byte_base + 4u * i) >> hm_shift) & 0x01010101u;
+        // the high bit is stored inverted: a clear bit means the value is 4 lower
+        res[i] = sub_packed_bytes(qs | (hm << 2u), 4u);
+    }
+    return res;
 }
-fn repack_b_qs(q8_block_idx: u32, tid: u32) -> vec4<u32> {
-    let phase = tid % 2u;
-    return vec4<u32>(
-        src1q[q8_block_idx].qs[4u * phase],
-        src1q[q8_block_idx].qs[4u * phase + 1u],
-        src1q[q8_block_idx].qs[4u * phase + 2u],
-        src1q[q8_block_idx].qs[4u * phase + 3u],
-    );
+fn repack_b_dm(q8_block_idx: u32) -> B_DS_TYPE {
+    return B_DS_TYPE(src1q[q8_block_idx].d);
 }
+fn get_dm(block_byte_base: u32) -> f32 {
+    return f32(load_f16_at_src0(block_byte_base + 108u));
+}
+fn get_scale_min(block_byte_base: u32, tid: u32) -> f32 {
+    let byte_idx = tid & 3u;
+    let group    = tid / 4u;
+
+    let scales_lo = load_u32_at_src0(block_byte_base + 96u + 4u * (group & 1u));
+    let scales_hi = load_u32_at_src0(block_byte_base + 104u);
+
+    let lo_byte = byte_of(scales_lo, byte_idx);
+    let lo      = select(lo_byte >> 4u, lo_byte & 0x0Fu, group < 2u);
+    let hi      = (byte_of(scales_hi, byte_idx) >> (2u * group)) & 3u;
+
+    return f32(i32(lo | (hi << 4u)) - 32);
+}
+#endif // MUL_ACC_Q3_K
+
+// Q4_K and Q5_K share the scale/min layout
+#if defined(MUL_ACC_Q4_K) || defined(MUL_ACC_Q5_K)
+#define B_DS_TYPE vec2<f32>
 fn repack_b_dm(q8_block_idx: u32) -> B_DS_TYPE {
     return B_DS_TYPE(
         f32(src1q[q8_block_idx].d),
@@ -246,26 +388,104 @@ fn get_scale_min(block_byte_base: u32, tid: u32) -> vec2<f32> {
 
     return vec2<f32>(scale, min_val);
 }
+#endif // MUL_ACC_Q4_K || MUL_ACC_Q5_K
+
+#ifdef MUL_ACC_Q4_K
+#define BLOCK_SIZE_BYTES 144
+fn repack_a(block_byte_base: u32, tid: u32) -> vec4<u32> {
+    let iq4 = tid / 4u;
+    let phase = tid % 2u;
+    let nibble = (tid >> 1u) % 2u;
+    let q_qs_byte_base = block_byte_base + 16u + 32u * iq4 + 16u * phase;
+    let qs_shift = 4u * nibble;
+    return vec4<u32>(
+        (load_u32_at_src0_aligned(q_qs_byte_base) >> qs_shift) & 0x0F0F0F0Fu,
+        (load_u32_at_src0_aligned(q_qs_byte_base + 4u) >> qs_shift) & 0x0F0F0F0Fu,
+        (load_u32_at_src0_aligned(q_qs_byte_base + 8u) >> qs_shift) & 0x0F0F0F0Fu,
+        (load_u32_at_src0_aligned(q_qs_byte_base + 12u) >> qs_shift) & 0x0F0F0F0Fu,
+    );
+}
 #endif // MUL_ACC_Q4_K
 
-#ifdef K_QUANTS
+#ifdef MUL_ACC_Q5_K
+#define BLOCK_SIZE_BYTES 176
+fn repack_a(block_byte_base: u32, tid: u32) -> vec4<u32> {
+    let iq4 = tid / 4u;
+    let phase = tid % 2u;
+    let nibble = (tid >> 1u) % 2u;
+    let ql_byte_base = block_byte_base + 48u + 32u * iq4 + 16u * phase;
+    let qh_byte_base = block_byte_base + 16u + 16u * phase;
+    let ql_shift = 4u * nibble;
+    let qh_shift = 2u * iq4 + nibble;
+
+    var res: vec4<u32>;
+    for (var i = 0u; i < 4u; i++) {
+        let ql = (load_u32_at_src0_aligned(ql_byte_base + 4u * i) >> ql_shift) & 0x0F0F0F0Fu;
+        let qh = (load_u32_at_src0_aligned(qh_byte_base + 4u * i) >> qh_shift) & 0x01010101u;
+        res[i] = ql | (qh << 4u);
+    }
+    return res;
+}
+#endif // MUL_ACC_Q5_K
+
+#ifdef MUL_ACC_Q6_K
+#define BLOCK_SIZE_BYTES 210
+#define B_DS_TYPE f32
+fn repack_a(block_byte_base: u32, tid: u32) -> vec4<u32> {
+    let half_blk = tid / 8u;
+    let sub      = (tid % 8u) / 2u;
+    let phase    = tid % 2u;
+    let ql_byte_base = block_byte_base + 64u * half_blk + 32u * (sub & 1u) + 16u * phase;
+    let qh_byte_base = block_byte_base + 128u + 32u * half_blk + 16u * phase;
+    let ql_shift = 4u * (sub >> 1u);
+    let qh_shift = 2u * sub;
+
+    var res: vec4<u32>;
+    for (var i = 0u; i < 4u; i++) {
+        let ql = (load_u32_at_src0(ql_byte_base + 4u * i) >> ql_shift) & 0x0F0F0F0Fu;
+        let qh = (load_u32_at_src0(qh_byte_base + 4u * i) >> qh_shift) & 0x03030303u;
+        res[i] = sub_packed_bytes(ql | (qh << 4u), 32u);
+    }
+    return res;
+}
+fn repack_b_dm(q8_block_idx: u32) -> B_DS_TYPE {
+    return B_DS_TYPE(src1q[q8_block_idx].d);
+}
+fn get_dm(block_byte_base: u32) -> f32 {
+    return f32(load_f16_at_src0(block_byte_base + 208u));
+}
+fn get_scale_min(block_byte_base: u32, tid: u32) -> f32 {
+    let scale_byte = block_byte_base + 192u + tid;
+    return f32(sbyte_of(load_u32_at_src0_aligned(scale_byte), scale_byte & 3u));
+}
+#endif // MUL_ACC_Q6_K
+
+#if defined(K_QUANTS_HANDLING)
 fn accumulate_vec_q_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1q_idx_base: u32) -> array<array<f32, OUTPUTS_PER_WG>, NUM_COLS> {
     var acc: array<array<f32, OUTPUTS_PER_WG>, NUM_COLS>;
 
     let tid = thread_id % THREADS_PER_BLOCK;
 
     for (var block = thread_id / THREADS_PER_BLOCK; block < params.k / BLOCK_SIZE; block += WG_SIZE / THREADS_PER_BLOCK) {
+        var b_qs_cols: array<vec4<u32>, NUM_COLS>;
+        var b_ds_cols: array<B_DS_TYPE, NUM_COLS>;
+        for (var col = 0u;col < NUM_COLS;col += 1) {
+            let src1q_idx = src1q_idx_base + col * (params.k / Q8_BLOCK_SIZE) + (block * BLOCK_SIZE + ELEMS_PER_THREAD * tid) / Q8_BLOCK_SIZE;
+            b_qs_cols[col] = repack_b_qs(src1q_idx, tid);
+            b_ds_cols[col] = repack_b_dm(src1q_idx);
+        }
         for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
             let output_row = row_base + row;
             if (output_row < params.m) {
                 let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
                 let a_repacked = repack_a(block_byte_base, tid);
                 let dm = get_dm(block_byte_base);
+#ifndef MUL_ACC_Q1_0
                 let scale_min = get_scale_min(block_byte_base, tid);
+#endif
                 for (var col = 0u;col < NUM_COLS;col += 1) {
-                    let src1q_idx = src1q_idx_base + col * (params.k / Q8_BLOCK_SIZE) + (block * BLOCK_SIZE + ELEMS_PER_THREAD * tid) / Q8_BLOCK_SIZE;
-                    let b_repacked = repack_b_qs(src1q_idx, tid);
-                    let b_ds = repack_b_dm(src1q_idx);
+                    let b_repacked = b_qs_cols[col];
+                    let b_ds = b_ds_cols[col];
 
 #if defined(MUL_ACC_Q2_K)
                     let scale_q = i32(scale_min.x);
@@ -279,13 +499,27 @@ fn accumulate_vec_q_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, s
                     acc[col][row] += b_ds * (dm.x * f32(row_sum_d) - dm.y * f32(row_sum_m));
 #endif // MUL_ACC_Q2_K
 
-#if defined(MUL_ACC_Q4_K)
+#if defined(MUL_ACC_Q4_K) || defined(MUL_ACC_Q5_K)
                     let row_sum = dot4I8Packed(a_repacked[0], b_repacked[0]) + dot4I8Packed(a_repacked[1], b_repacked[1])
                                     + dot4I8Packed(a_repacked[2], b_repacked[2]) + dot4I8Packed(a_repacked[3], b_repacked[3]);
 
                     // Each thread covers half of the Q8_1 block, so add only b_ds.y/2.
                     acc[col][row] += b_ds.x * dm.x * scale_min.x * f32(row_sum) - dm.y * scale_min.y * (b_ds.y / (Q8_BLOCK_SIZE / ELEMS_PER_THREAD));
-#endif // MUL_ACC_Q4_K
+#endif // MUL_ACC_Q4_K || MUL_ACC_Q5_K
+
+#if defined(MUL_ACC_Q3_K) || defined(MUL_ACC_Q6_K)
+                    let row_sum = dot4I8Packed(a_repacked[0], b_repacked[0]) + dot4I8Packed(a_repacked[1], b_repacked[1])
+                                    + dot4I8Packed(a_repacked[2], b_repacked[2]) + dot4I8Packed(a_repacked[3], b_repacked[3]);
+
+                    acc[col][row] += b_ds * dm * scale_min * f32(row_sum);
+#endif // MUL_ACC_Q3_K || MUL_ACC_Q6_K
+
+#if defined(MUL_ACC_Q1_0)
+                    let row_sum = dot4I8Packed(a_repacked[0], b_repacked[0]) + dot4I8Packed(a_repacked[1], b_repacked[1])
+                                    + dot4I8Packed(a_repacked[2], b_repacked[2]) + dot4I8Packed(a_repacked[3], b_repacked[3]);
+
+                    acc[col][row] += b_ds * dm * f32(row_sum);
+#endif // MUL_ACC_Q1_0
 
                 }
             }
@@ -294,4 +528,4 @@ fn accumulate_vec_q_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, s
 
     return acc;
 }
-#endif // K_QUANTS
+#endif // K_QUANTS_HANDLING

ggml/src/ggml-webgpu/wgsl-shaders/quantize_q8.wgsl

@@ -34,7 +34,7 @@ fn cluster_max_8(v: f32) -> f32 {
     return r;
 }
 
-#if defined(MUL_ACC_Q4_0) || defined(MUL_ACC_Q4_1) || defined(MUL_ACC_Q4_K)
+#if defined(MUL_ACC_Q4_0) || defined(MUL_ACC_Q4_1) || defined(MUL_ACC_Q5_0) || defined(MUL_ACC_Q5_1) || defined(MUL_ACC_Q4_K) || defined(MUL_ACC_Q5_K)
 fn cluster_add_i4x8(v: i32) -> i32 {
     var r= v;
     r += subgroupShuffleXor(r, 1u);
@@ -113,7 +113,7 @@ fn main(
         src1q[src1q_idx].qs[qs_idx] = q4_quants;
     }
 
-#if defined(MUL_ACC_Q4_0) || defined(MUL_ACC_Q4_1) || defined(MUL_ACC_Q4_K)
+#if defined(MUL_ACC_Q4_0) || defined(MUL_ACC_Q4_1) || defined(MUL_ACC_Q5_0) || defined(MUL_ACC_Q5_1) || defined(MUL_ACC_Q4_K) || defined(MUL_ACC_Q5_K)
     let q4_quants_sum = dot4I8Packed(q4_quants, 0x01010101u);
     let s = f16(d * f32(cluster_add_i4x8(q4_quants_sum)));
 
@@ -158,7 +158,7 @@ fn main(
         }
     }
 
-#if defined(MUL_ACC_Q4_0) || defined(MUL_ACC_Q4_1) || defined(MUL_ACC_Q4_K)
+#if defined(MUL_ACC_Q4_0) || defined(MUL_ACC_Q4_1) || defined(MUL_ACC_Q5_0) || defined(MUL_ACC_Q5_1) || defined(MUL_ACC_Q4_K) || defined(MUL_ACC_Q5_K)
 
     partial_sums[cluster_id][qs_idx] = dot4I8Packed(q4_quants, 0x01010101u);
 

include/llama.h

@@ -43,10 +43,10 @@
 #define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq'
 
 #define LLAMA_SESSION_MAGIC   LLAMA_FILE_MAGIC_GGSN
-#define LLAMA_SESSION_VERSION 10
+#define LLAMA_SESSION_VERSION 11
 
 #define LLAMA_STATE_SEQ_MAGIC   LLAMA_FILE_MAGIC_GGSQ
-#define LLAMA_STATE_SEQ_VERSION 3
+#define LLAMA_STATE_SEQ_VERSION 4
 
 #ifdef __cplusplus
 extern "C" {

scripts/make-release-checks.sh

@@ -8,6 +8,8 @@
 #   GH_TOKEN, GITHUB_REPOSITORY, GITHUB_OUTPUT
 #   RELEASE_BRANCH: when set, HEAD must belong to origin/RELEASE_BRANCH and must
 #     not be older than 3 days from the branch HEAD (skipped when unset)
+#   REQUIRE_DOCKER: set to "false" to skip the container image check and the
+#     release re-tagging of the nightly images in make-release.yml (default: true)
 set -euo pipefail
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
@@ -126,6 +128,8 @@ if [[ -z "${NIGHTLY_TAG}" ]]; then
     echo "Warning: no nightly tag points at ${SHA} - skipping container image check"
 elif [[ -z "${GITHUB_REPOSITORY:-}" ]]; then
     echo "Warning: GITHUB_REPOSITORY not set - skipping container image check (local run)"
+elif [[ "${REQUIRE_DOCKER:-true}" != "true" ]]; then
+    echo "Require Docker is disabled - skipping container image check and re-tagging"
 else
     CONTAINER_REPO="${GITHUB_REPOSITORY,,}"  # lower-case owner/repo for ghcr.io
     GHCR_TOKEN="$(curl -fsSL \

scripts/make-release-summary.txt

@@ -44,10 +44,12 @@ Write a summary of the change log in a few sections:
 Guidelines:
 
 - All bullet point in the summary should be concise and rarely exceed a single line of 120 characters (excluding PR links)
+- Do not wrap or break lines within a bullet point or paragraph - each bullet point must fit on a single line (PR links allowed to exceed the line limit)
 - Avoid repeating `ggml`-specific changes - these should already be covered by the `ggml` release links
 - Provide PR link for each bullet point where possible
 - Don't add bullet point to state that there are no API changes in some module
-- Combine related topics (a full list of commits will be appended independently at end of the summary)
+- Group similar/related items into a single bullet with multiple trailing PR links, one bullet per topic
+  - A full list of commits will be appended independently at end of the summary, so it is not a problem to skip some less important changes
 - Skip minor-impact notes (e.g. "fix compile warnings", "refactored code", ...)
 
 Output just the summary in a markdown block, without any extra text. Save it to a local text file called `release-notes-vX.Y.Z.txt`.

scripts/snapdragon/run.py

@@ -31,8 +31,10 @@ MANAGED_ENV_NAMES = (
     "GGML_HEXAGON_MBUF",
     "GGML_HEXAGON_MM_SELECT",
     "GGML_HEXAGON_FA_SELECT",
+    "GGML_HEXAGON_FA_HEAD_SPLIT",
     "GGML_HEXAGON_GDN_SELECT",
     "GGML_HEXAGON_AR_SELECT",
+    "GGML_HEXAGON_AR_SCATTER",
     "GGML_HEXAGON_ETM",
     "GGML_HEXAGON_ARCH",
     "GGML_HEXAGON_OPTRACE",
@@ -167,6 +169,7 @@ def main():
     parser.add_argument("--hex-mbuf", help="Maximum host buffer size limit in MB to allocate (GGML_HEXAGON_MBUF)")
     parser.add_argument("--hex-mm-select", help="Select MUL_MAT and MUL_MAT_ID kernel (GGML_HEXAGON_MM_SELECT) 2:HMX,1:HVX,0:disable")
     parser.add_argument("--hex-fa-select", help="Select Flash Attention kernel (GGML_HEXAGON_FA_SELECT) 2:HMX,1:HVX,0:disable")
+    parser.add_argument("--hex-fa-head-split", help="Enable (1) or disable (0) head-parallel flash_attn partitioning (GGML_HEXAGON_FA_HEAD_SPLIT)")
     parser.add_argument("--hex-gdn-select", help="Select Gated Delta Net kernel (GGML_HEXAGON_GDN_SELECT) 2:HMX,1:HVX,0:disable")
     parser.add_argument("--hex-ar-select", help="Select All-Reduce kernel (GGML_HEXAGON_AR_SELECT) 1:enable,0:disable")
     parser.add_argument("--hex-ar-scatter", help="Enable (1) or disable (0) reduce-scatter for fused ALLREDUCE+ADD (GGML_HEXAGON_AR_SCATTER)")
@@ -309,6 +312,7 @@ def main():
     set_env("GGML_HEXAGON_MBUF", args.hex_mbuf)
     set_env("GGML_HEXAGON_MM_SELECT", args.hex_mm_select)
     set_env("GGML_HEXAGON_FA_SELECT", args.hex_fa_select)
+    set_env("GGML_HEXAGON_FA_HEAD_SPLIT", args.hex_fa_head_split)
     set_env("GGML_HEXAGON_GDN_SELECT", args.hex_gdn_select)
     set_env("GGML_HEXAGON_AR_SELECT", args.hex_ar_select)
     set_env("GGML_HEXAGON_AR_SCATTER", args.hex_ar_scatter)

scripts/sync-ggml.last

@@ -1 +1 @@
-353b63b439f27ab2cc19dac97ab1681ba6d2d084
+ffa4e8b80930029a35991f94e7c8a93cd67730ab

src/llama-arch.cpp

@@ -1178,6 +1178,21 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
     }
 }
 
+// these models pick weights, routing or input meaning per ubatch based on token vs embd input
+bool llm_arch_supports_mixed_batch(const llm_arch & arch) {
+    switch (arch) {
+        case LLM_ARCH_COGVLM:
+        case LLM_ARCH_DEEPSEEK4:
+        case LLM_ARCH_GRANITE_SWITCH:
+        case LLM_ARCH_EAGLE3:
+        case LLM_ARCH_DFLASH:
+        case LLM_ARCH_GEMMA4_ASSISTANT:
+            return false;
+        default:
+            return true;
+    }
+}
+
 bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
     switch (arch) {
         case LLM_ARCH_GROK:

src/llama-arch.h

@@ -825,3 +825,4 @@ bool llm_arch_is_hybrid         (const llm_arch & arch);
 bool llm_arch_is_diffusion      (const llm_arch & arch);
 bool llm_arch_supports_sm_tensor(const llm_arch & arch);
 bool llm_arch_supports_rs_rollback(const llm_arch & arch);
+bool llm_arch_supports_mixed_batch(const llm_arch & arch);

src/llama-batch.cpp

@@ -12,7 +12,7 @@
 #include <algorithm>
 #include <sstream>
 
-llama_batch_allocr::llama_batch_allocr(uint32_t n_pos_per_embd) : n_pos_per_embd(n_pos_per_embd) {
+llama_batch_allocr::llama_batch_allocr(uint32_t n_pos_per_embd, bool allow_mixed) : n_pos_per_embd(n_pos_per_embd), allow_mixed(allow_mixed) {
     const char * LLAMA_BATCH_DEBUG = getenv("LLAMA_BATCH_DEBUG");
     debug = LLAMA_BATCH_DEBUG ? atoi(LLAMA_BATCH_DEBUG) : 0;
 
@@ -49,25 +49,49 @@ bool llama_batch_allocr::init(
     //
     // determine the content types of the batch
     // an entry can carry a token id, a token embedding, or both (e.g. MTP hook batches)
-    // all entries must carry the same combination
+    // all entries must carry the same combination, or be a mix of token and embd entries
     //
 
-    const bool has_token = batch_inp.tokens[0].id != LLAMA_TOKEN_NULL;
-    const bool has_embd  = batch_inp.tokens[0].has_embd;
+    int32_t n_tok_only  = 0;
+    int32_t n_embd_only = 0;
+    int32_t n_both      = 0;
 
-    for (int32_t i = 1; i < n_tok; ++i) {
-        if ((batch_inp.tokens[i].id != LLAMA_TOKEN_NULL) != has_token ||
-             batch_inp.tokens[i].has_embd                != has_embd) {
-            LLAMA_LOG_ERROR("%s: all entries in the batch must have the same content types\n", __func__);
+    for (int32_t i = 0; i < n_tok; ++i) {
+        const bool is_tok = batch_inp.tokens[i].id != LLAMA_TOKEN_NULL;
+        const bool is_emb = batch_inp.tokens[i].has_embd;
+
+        if (!is_tok && !is_emb) {
+            LLAMA_LOG_ERROR("%s: entry %d has neither a token id nor an embedding\n", __func__, i);
             return false;
         }
+
+        n_tok_only  += is_tok && !is_emb;
+        n_embd_only += is_emb && !is_tok;
+        n_both      += is_tok &&  is_emb;
     }
 
-    if (!has_token && !has_embd) {
-        LLAMA_LOG_ERROR("%s: batch has neither token ids nor embeddings\n", __func__);
+    if (n_both > 0 && n_both != n_tok) {
+        LLAMA_LOG_ERROR("%s: entries with both a token id and an embedding cannot be mixed with other entries\n", __func__);
         return false;
     }
 
+    const bool mixed = n_tok_only > 0 && n_embd_only > 0;
+
+    if (mixed && !allow_mixed) {
+        LLAMA_LOG_ERROR("%s: this model or context does not support batches mixing token and embedding entries\n", __func__);
+        return false;
+    }
+
+    const bool has_token = n_tok_only  > 0 || n_both > 0;
+    const bool has_embd  = n_embd_only > 0 || n_both > 0;
+
+    if (mixed) {
+        is_embd_vec.resize(n_tok);
+        for (int32_t i = 0; i < n_tok; ++i) {
+            is_embd_vec[i] = batch_inp.tokens[i].has_embd ? 1 : 0;
+        }
+    }
+
     //
     // build flat token/embd array
     //
@@ -75,6 +99,10 @@ bool llama_batch_allocr::init(
     if (has_token) {
         token_vec.resize(n_tok);
         for (int32_t i = 0; i < n_tok; ++i) {
+            if (mixed && is_embd_vec[i]) {
+                token_vec[i] = 0; // placeholder
+                continue;
+            }
             const llama_token id = batch_inp.tokens[i].id;
             if (id < 0 || id >= batch_inp.n_vocab) {
                 LLAMA_LOG_ERROR("%s: invalid token[%d] = %d\n", __func__, i, id);
@@ -84,29 +112,35 @@ bool llama_batch_allocr::init(
         }
     }
 
-    if (has_embd) {
+    if (mixed) {
+        embd_vec.assign((size_t) n_tok*n_embd, 0.0f);
+        for (int32_t i = 0; i < n_tok; ++i) {
+            if (is_embd_vec[i]) {
+                const float * src = batch_inp.embd.data() + batch_inp.tokens[i].embd_off;
+                std::copy(src, src + n_embd, embd_vec.data() + (size_t) i*n_embd);
+            }
+        }
+    } else if (has_embd) {
         embd_vec = batch_inp.embd;
     }
 
     //
-    // build flat pos array
-    // token batch:     pos[i]            = tokens[i].pos[0]
-    // embedding batch: pos[j*n_tok + i]  = tokens[i].pos[j]  (section-major)
+    // build flat pos array, section-major: pos[j*n_tok + i] = section j of entry i
+    // token entry: [p, p, p, 0] (M-RoPE text position)
+    // embd entry:  tokens[i].pos as-is
     //
 
-    {
-        const int32_t n_pos_total = has_token ? n_tok : n_tok * (int32_t) n_pos_per_embd;
-        pos.resize(n_pos_total);
-        if (has_token) {
-            for (int32_t i = 0; i < n_tok; ++i) {
-                pos[i] = batch_inp.tokens[i].pos[0];
-            }
-        } else {
-            for (int32_t i = 0; i < n_tok; ++i) {
-                for (uint32_t j = 0; j < n_pos_per_embd; ++j) {
-                    pos[(int32_t) j * n_tok + i] = batch_inp.tokens[i].pos[j];
-                }
+    pos.resize((size_t) n_tok*n_pos_per_embd);
+    for (int32_t i = 0; i < n_tok; ++i) {
+        const auto & tok = batch_inp.tokens[i];
+        const bool expand = tok.id != LLAMA_TOKEN_NULL;
+        for (uint32_t j = 0; j < n_pos_per_embd; ++j) {
+            llama_pos p = tok.pos[j];
+            if (expand) {
+                // expand [p] to [p, p, p, 0] for M-RoPE
+                p = j < 3 ? tok.pos[0] : 0;
             }
+            pos[(size_t) j*n_tok + i] = p;
         }
     }
 
@@ -264,6 +298,7 @@ bool llama_batch_allocr::init(
             /*.seq_id_unq   =*/ this->seq_id_unq.data(),
             /*.seq_idx      =*/ this->seq_idx.data(),
             /*.output       =*/ batch.logits,
+            /*.type         =*/ is_embd_vec.empty() ? nullptr : is_embd_vec.data(),
             /*.decision_order =*/ decision_order.empty() ? nullptr : decision_order.data(),
             /*.data         =*/ {},
         };
@@ -294,6 +329,21 @@ bool llama_batch_allocr::init(
     //
 
     if (n_pos_per_embd > 1) {
+        // in a mixed batch, the first entry of each seq picks the rule
+        std::vector<int8_t> seq_first_embd(n_seq_max, batch.token ? 0 : 1);
+        if (mixed) {
+            std::vector<bool> seen(n_seq_max, false);
+            for (int32_t i = 0; i < batch.n_tokens; ++i) {
+                for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {
+                    const llama_seq_id sid = batch.seq_id[i][s];
+                    if (!seen[sid]) {
+                        seen[sid] = true;
+                        seq_first_embd[sid] = is_embd_vec[i];
+                    }
+                }
+            }
+        }
+
         // M-RoPE case: allow position to "jump" forward only (non-continuous positions are allowed)
         for (uint32_t s = 0; s < n_seq_max; ++s) {
             if (seq_pos[s].empty()) {
@@ -302,7 +352,7 @@ bool llama_batch_allocr::init(
 
             const llama_pos p0 = mem ? mem->seq_pos_max(s) : -1;
 
-            if (batch.token) {
+            if (!seq_first_embd[s]) {
                 if (p0 >= 0 && p0 >= seq_pos_min(s)) {
                     LLAMA_LOG_ERROR(
                             "%s: the tokens of sequence %d in the input batch have inconsistent sequence positions:\n"
@@ -471,6 +521,7 @@ llama_ubatch llama_batch_allocr::ubatch_reserve(uint32_t n_seq_tokens, uint32_t
         /*.seq_id_unq   =*/ udata->seq_id_unq.data(),
         /*.seq_idx      =*/ udata->seq_idx.data(),
         /*.output       =*/ udata->output.data(),
+        /*.type         =*/ nullptr,
         /*.decision_order =*/ nullptr,
         /*.data         =*/ std::move(udata),
     };
@@ -769,6 +820,7 @@ void llama_batch_allocr::clear() {
 
     token_vec   .clear();
     embd_vec    .clear();
+    is_embd_vec .clear();
     seq_id_data .clear();
     pos         .clear();
     n_seq_id    .clear();
@@ -799,7 +851,20 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
 
     auto udata = std::make_shared<llama_ubatch::data_t>();
 
-    const int64_t n_embd_all = batch.embd ? (int64_t) n_tokens*n_embd : 0;
+    const bool mixed_batch = !is_embd_vec.empty();
+
+    // a ubatch with a single kind of rows is emitted as a plain token or embd ubatch
+    uint32_t n_embd_rows = 0;
+    if (mixed_batch) {
+        for (int32_t idx : idxs) {
+            n_embd_rows += is_embd_vec[idx];
+        }
+    }
+    const bool mixed     = mixed_batch && n_embd_rows > 0 && n_embd_rows < n_tokens;
+    const bool use_token = batch.token && !(mixed_batch && n_embd_rows == n_tokens);
+    const bool use_embd  = batch.embd  && !(mixed_batch && n_embd_rows == 0);
+
+    const int64_t n_embd_all = use_embd ? (int64_t) n_tokens*n_embd : 0;
     const int64_t n_pos_all  =              (int64_t) n_tokens*n_pos_per_embd;
 
     udata->token     .resize(n_tokens);
@@ -810,6 +875,7 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
     udata->seq_id_unq.resize(0);
     udata->seq_idx   .resize(LLAMA_MAX_SEQ, -1);
     udata->output    .resize(n_tokens);
+    udata->type      .resize(mixed ? n_tokens : 0);
     udata->decision_order.resize(decision_order.empty() ? 0 : n_tokens);
 
     udata->batch_idxs = idxs;
@@ -818,21 +884,20 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
     seq_set_t seq_set_unq;
 
     for (size_t i = 0; i < idxs.size(); ++i) {
-        if (batch.token) {
+        if (use_token) {
             udata->token[i] = batch.token[idxs[i]];
         }
 
-        if (batch.embd) {
+        if (use_embd) {
             memcpy(udata->embd.data() + i*n_embd, batch.embd + (int64_t) idxs[i]*n_embd, n_embd*sizeof(float));
         }
 
+        if (mixed) {
+            udata->type[i] = is_embd_vec[idxs[i]];
+        }
+
         for (size_t j = 0; j < (size_t)n_pos_per_embd; ++j) {
-            // if we are using M-RoPE
-            //     if the current batch is text, we need to broadcast the same position across all RoPE sections
-            //     otherwise, the input batch is image embeddings, we copy the positions as-is
-            // if we are not using M-RoPE, there is only one position per token (this loop runs only once)
-            size_t src_off = batch.token ? 0 : j*batch.n_tokens;
-            udata->pos[j*n_tokens + i] = batch.pos[src_off + idxs[i]];
+            udata->pos[j*n_tokens + i] = batch.pos[j*batch.n_tokens + idxs[i]];
         }
 
         udata->n_seq_id[i] = batch.n_seq_id[idxs[i]];
@@ -875,14 +940,15 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
         /*.n_seqs_unq   =*/ (uint32_t) udata->seq_id_unq.size(),
         /*.n_pos        =*/ n_pos_per_embd,
 
-        /*.token        =*/ batch.token ? udata->token.data() : nullptr,
-        /*.embd         =*/ batch.embd ? udata->embd.data() : nullptr,
+        /*.token        =*/ use_token ? udata->token.data() : nullptr,
+        /*.embd         =*/ use_embd  ? udata->embd.data()  : nullptr,
         /*.pos          =*/ udata->pos.data(),
         /*.n_seq_id     =*/ udata->n_seq_id.data(),
         /*.seq_id       =*/ udata->seq_id.data(),
         /*.seq_id_unq   =*/ udata->seq_id_unq.data(),
         /*.seq_idx      =*/ udata->seq_idx.data(),
         /*.output       =*/ udata->output.data(),
+        /*.type         =*/ mixed ? udata->type.data() : nullptr,
         /*.decision_order =*/ udata->decision_order.empty() ? nullptr : udata->decision_order.data(),
         /*.data         =*/ std::move(udata),
     };
@@ -933,6 +999,7 @@ void llama_batch_allocr::ubatch_print(const llama_ubatch & ubatch, int debug) {
         LLAMA_LOG_DEBUG("%s:   seq_id_unq = %s\n", __func__, ss_seq_id_unq.str().c_str());
         LLAMA_LOG_DEBUG("%s:   seq_idx    = %s\n", __func__, ss_seq_idx.str().c_str());
         LLAMA_LOG_DEBUG("%s:   output     = %p\n", __func__, (void *) ubatch.output);
+        LLAMA_LOG_DEBUG("%s:   type       = %p\n", __func__, (void *) ubatch.type);
         LLAMA_LOG_DEBUG("%s:   n_outputs  = %d\n", __func__, n_outputs);
 
         if (debug > 0) {
@@ -963,7 +1030,7 @@ void llama_batch_allocr::ubatch_print(const llama_ubatch & ubatch, int debug) {
                     }
                 }
 
-                if (ubatch.token) {
+                if (ubatch.token && !(ubatch.is_mixed() && ubatch.type[i])) {
                     LLAMA_LOG_DEBUG("%s:  %4d: id = %6d (%16s), pos = %4d, n_seq_id = %2d, seq_id = [%s], output = %d\n",
                             __func__, i, ubatch.token[i], vocab->token_to_piece(ubatch.token[i]).c_str(),
                             ubatch.pos[i], ubatch.n_seq_id[i], ss.str().c_str(), ubatch.output[i]);

src/llama-batch.h

@@ -29,6 +29,11 @@ struct llama_ubatch {
         return n_pos >= 3;
     }
 
+    // mixed: type picks token or embd per row, pos has n_pos sections for all rows
+    bool is_mixed() const {
+        return type != nullptr;
+    }
+
     uint32_t b_equal_seqs; // note: this is a boolean, but we use an int32_t for alignment
                            //       otherwise address sanitizer complains
     // TODO: whole_seqs for embeddings?
@@ -52,6 +57,7 @@ struct llama_ubatch {
     llama_seq_id *  seq_id_unq; // [n_seqs_unq]       | s   | seq_id
     int32_t      *  seq_idx;    // [LLAMA_MAX_SEQ]    | -   | seq_idx
     int8_t       *  output;     // [n_tokens]         | i   | -
+    int8_t       *  type;       // [n_tokens]         | i   | -     (mixed ubatch only, 0 - token, 1 - embd)
     int32_t      *  decision_order; // [n_tokens], NULL if no entry has one, see llama_batch_ext_set_decision_order()
 
     struct data_t {
@@ -63,6 +69,7 @@ struct llama_ubatch {
         std::vector<llama_seq_id>   seq_id_unq;
         std::vector<int32_t>        seq_idx;
         std::vector<int8_t>         output;
+        std::vector<int8_t>         type;
         std::vector<int32_t>        batch_idxs;  // original batch index for each token
         std::vector<int32_t>        decision_order;
 
@@ -73,6 +80,9 @@ struct llama_ubatch {
     std::shared_ptr<data_t> data;
 };
 
+// crash if a mixed ubatch reaches code that expects only tokens or only embd
+#define ASSERT_EMBD_OR_TOKEN(ubatch) GGML_ASSERT(!(ubatch).is_mixed() && "mixed token/embd ubatch is not supported here")
+
 struct llama_hparams;
 
 // MTP hook batches carry the target model's hidden state (n_embd_out size).
@@ -134,7 +144,7 @@ struct llama_batch_ext {
 // a helper for sanitizing, fulfilling and splitting a batch
 class llama_batch_allocr {
 public:
-    llama_batch_allocr(uint32_t n_pos_per_embd);
+    llama_batch_allocr(uint32_t n_pos_per_embd, bool allow_mixed = false);
 
     // convert a llama_batch_ext to internal llama_batch and sanitize it
     bool init(
@@ -192,12 +202,15 @@ private:
     //       ref: https://github.com/ggml-org/llama.cpp/issues/13694#issuecomment-2983871762
     const uint32_t n_pos_per_embd;
 
+    const bool allow_mixed;
+
     uint32_t n_embd;
     uint32_t n_seq_max;
     uint32_t n_outputs;
 
     std::vector<llama_token>    token_vec;    // owned token IDs built from llama_batch_ext
     std::vector<float>          embd_vec;     // owned embeddings built from llama_batch_ext
+    std::vector<int8_t>         is_embd_vec;  // mixed batch only (= 1 if embd, 0 if text token)
     std::vector<llama_seq_id>   seq_id_data;  // flat storage for seq_id pointers below
 
     std::vector<llama_pos>      pos;

src/llama-context.cpp

@@ -87,7 +87,9 @@ llama_context::llama_context(
     model(model),
     cvec(std::make_unique<llama_adapter_cvec>()),
     loras(std::make_unique<llama_adapter_loras>()),
-    balloc(std::make_unique<llama_batch_allocr>(model.hparams.n_pos_per_embd())) {
+    // MTP uses the embd input for the hidden state
+    balloc(std::make_unique<llama_batch_allocr>(model.hparams.n_pos_per_embd(),
+                llm_arch_supports_mixed_batch(model.arch) && params.ctx_type == LLAMA_CONTEXT_TYPE_DEFAULT)) {
     // TODO warning when creating llama_context with awkward ctx size that is not a power of 2,
     //     may need to be backend-dependent
     LLAMA_LOG_INFO("%s: constructing llama_context\n", __func__);

src/llama-graph.cpp

@@ -73,13 +73,55 @@ void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) {
         ggml_backend_tensor_set(tokens, ubatch->token, 0, n_tokens*ggml_element_size(tokens));
     }
 
-    if (ubatch->embd) {
+    if (ubatch->embd && embd && !ubatch->is_mixed()) {
         GGML_ASSERT(n_embd == embd->ne[0]);
 
         const int64_t n_tokens = ubatch->n_tokens;
 
         ggml_backend_tensor_set(embd, ubatch->embd, 0, n_tokens*n_embd*ggml_element_size(embd));
     }
+
+    if (ubatch->is_mixed() && embd) {
+        GGML_ASSERT(mixed_tokens && mixed_slots && mixed_embd && "mixed token/embd ubatch is not supported here");
+
+        std::vector<int32_t> ids;
+        std::vector<int64_t> slots;
+        for (uint32_t i = 0; i < ubatch->n_tokens; ++i) {
+            if (!ubatch->type[i]) {
+                ids.push_back(ubatch->token[i]);
+                slots.push_back(i);
+            }
+        }
+        GGML_ASSERT((int64_t) ids.size() == mixed_tokens->ne[0]);
+        GGML_ASSERT(n_embd == mixed_embd->ne[0]);
+
+        ggml_backend_tensor_set(mixed_tokens, ids.data(),    0, ggml_nbytes(mixed_tokens));
+        ggml_backend_tensor_set(mixed_slots,  slots.data(),  0, ggml_nbytes(mixed_slots));
+        ggml_backend_tensor_set(mixed_embd,   ubatch->embd,  0, ggml_nbytes(mixed_embd));
+    }
+
+    if (scale_rows) {
+        const int64_t n_tokens = ubatch->n_tokens;
+
+        std::vector<float> data(n_tokens);
+        for (int64_t i = 0; i < n_tokens; ++i) {
+            const bool is_embd = !ubatch->token || (ubatch->is_mixed() && ubatch->type[i]);
+            data[i] = is_embd ? 1.0f : scale_tok;
+        }
+        ggml_backend_tensor_set(scale_rows, data.data(), 0, ggml_nbytes(scale_rows));
+    }
+}
+
+// number of token rows of the mixed path, a non-mixed ubatch is sized for the worst case
+static int64_t llm_graph_n_tok_rows(const llama_ubatch & ubatch) {
+    if (!ubatch.is_mixed()) {
+        return ubatch.n_tokens;
+    }
+    int64_t n = 0;
+    for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
+        n += !ubatch.type[i];
+    }
+    return n;
 }
 
 bool llm_graph_input_embd::can_reuse(const llm_graph_params & params) {
@@ -87,11 +129,16 @@ bool llm_graph_input_embd::can_reuse(const llm_graph_params & params) {
 
     res &= (!params.ubatch.token) || (tokens && tokens->ne[0] == params.ubatch.n_tokens);
     res &= (!params.ubatch.embd)  || (embd   &&   embd->ne[1] == params.ubatch.n_tokens);
+    res &= (!mixed_tokens) || mixed_tokens->ne[0] == llm_graph_n_tok_rows(params.ubatch);
+    res &= (!mixed_embd)   || mixed_embd->ne[1]   == params.ubatch.n_tokens;
+    res &= (!scale_rows) || scale_rows->ne[1] == params.ubatch.n_tokens;
 
     return res;
 }
 
 void llm_graph_input_embd_h::set_input(const llama_ubatch * ubatch) {
+    ASSERT_EMBD_OR_TOKEN(*ubatch);
+
     const int64_t n_tokens = ubatch->n_tokens;
 
     if (ubatch->token) {
@@ -128,21 +175,7 @@ void llm_graph_input_pos::set_input(const llama_ubatch * ubatch) {
     if (ubatch->pos && pos) {
         const int64_t n_tokens = ubatch->n_tokens;
 
-        if (ubatch->token && n_pos_per_embd == 4) {
-            // in case we're using M-RoPE with text tokens, convert the 1D positions to 4D
-            // the 3 first dims are the same, and 4th dim is all 0
-            std::vector<llama_pos> pos_data(n_tokens*n_pos_per_embd);
-            // copy the first dimension
-            for (int i = 0; i < n_tokens; ++i) {
-                pos_data[               i] = ubatch->pos[i];
-                pos_data[    n_tokens + i] = ubatch->pos[i];
-                pos_data[2 * n_tokens + i] = ubatch->pos[i];
-                pos_data[3 * n_tokens + i] = 0; // 4th dim is 0
-            }
-            ggml_backend_tensor_set(pos, pos_data.data(), 0, pos_data.size()*ggml_element_size(pos));
-        } else {
-            ggml_backend_tensor_set(pos, ubatch->pos, 0, n_tokens*n_pos_per_embd*ggml_element_size(pos));
-        }
+        ggml_backend_tensor_set(pos, ubatch->pos, 0, n_tokens*n_pos_per_embd*ggml_element_size(pos));
     }
 }
 
@@ -2377,7 +2410,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
 }
 
 // input embeddings with optional lora
-ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const {
+ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd, float tok_scale) const {
     const int64_t n_embd_inp = hparams.n_embd_inp();
     const int64_t n_embd     = hparams.n_embd;
 
@@ -2394,15 +2427,9 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const {
     cb(inp->embd, "inp_embd", -1);
     ggml_set_input(inp->embd);
 
-    // select one of the 2 inputs, based on the batch contents
-    // ref: https://github.com/ggml-org/llama.cpp/pull/18550
-    std::array<ggml_tensor *, 2> inps;
-
-    // token embeddings path (ubatch.token != nullptr)
-    {
-        auto & cur = inps[0];
-
-        cur = ggml_get_rows(ctx0, tok_embd, inp->tokens);
+    // token embeddings with lora and padding
+    auto build_tok = [&](ggml_tensor * ids) {
+        ggml_tensor * cur = ggml_get_rows(ctx0, tok_embd, ids);
 
         // apply lora for embedding tokens if needed
         for (const auto & lora : *loras) {
@@ -2416,7 +2443,7 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const {
 
             ggml_tensor * inpL_delta = ggml_scale(ctx0, ggml_mul_mat(
                         ctx0, lw->b, // non-transposed lora_b
-                        ggml_get_rows(ctx0, lw->a, inp->tokens)
+                        ggml_get_rows(ctx0, lw->a, ids)
                         ), scale);
 
             cur = ggml_add(ctx0, cur, inpL_delta);
@@ -2425,19 +2452,48 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const {
         if (n_embd_inp != n_embd) {
             cur = ggml_pad(ctx0, cur, hparams.n_embd_inp() - n_embd, 0, 0, 0);
         }
-    }
+
+        return cur;
+    };
+
+    // select one of the 3 inputs, based on the batch contents
+    // ref: https://github.com/ggml-org/llama.cpp/pull/18550
+    std::array<ggml_tensor *, 3> inps = {};
+
+    // token embeddings path (ubatch.token != nullptr)
+    inps[0] = build_tok(inp->tokens);
 
     // vector embeddings path (ubatch.embd != nullptr)
-    {
-        auto & cur = inps[1];
+    inps[1] = inp->embd;
 
-        cur = inp->embd;
+    // mixed path (ubatch.is_mixed()): set_rows the token rows into a copy of the embd rows, with its own inputs as select branches must not share tensors
+    // TODO: use inp->tokens and inp->embd once ggml_build_forward_select allows it
+    const bool has_mixed = llm_arch_supports_mixed_batch(arch) && cparams.ctx_type == LLAMA_CONTEXT_TYPE_DEFAULT;
+    if (has_mixed) {
+        const int64_t n_tok_rows = llm_graph_n_tok_rows(ubatch);
+
+        inp->mixed_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tok_rows);
+        cb(inp->mixed_tokens, "inp_mixed_tokens", -1);
+        ggml_set_input(inp->mixed_tokens);
+
+        inp->mixed_slots = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_tok_rows);
+        cb(inp->mixed_slots, "inp_mixed_slots", -1);
+        ggml_set_input(inp->mixed_slots);
+
+        inp->mixed_embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, ubatch.n_tokens);
+        cb(inp->mixed_embd, "inp_mixed_embd", -1);
+        ggml_set_input(inp->mixed_embd);
+
+        // note: set_rows writes into its destination, so it gets a copy of the input
+        inps[2] = ggml_set_rows(ctx0, ggml_dup(ctx0, inp->mixed_embd), build_tok(inp->mixed_tokens), inp->mixed_slots);
     }
 
     assert(ggml_are_same_shape (inps[0], inps[1]));
     assert(ggml_are_same_stride(inps[0], inps[1]));
 
-    ggml_tensor * cur = ggml_build_forward_select(gf, inps.data(), inps.size(), ubatch.token ? 0 : 1);
+    const int idx = ubatch.is_mixed() ? 2 : ubatch.token ? 0 : 1;
+
+    ggml_tensor * cur = ggml_build_forward_select(gf, inps.data(), has_mixed ? 3 : 2, idx);
 
     if (n_embd_inp != n_embd) {
         cur = ggml_view_2d(ctx0, cur, n_embd, n_tokens, cur->nb[1], 0);
@@ -2445,16 +2501,31 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const {
 
     res->t_inp_embd = cur;
 
-    // For Granite architecture
     // NOTE: For deepstack models, only apply scale to token inputs (ie text-only input).
     //  Raw embeddings are assumed to be multimodal inputs that should not be scaled.
-    if (hparams.f_embedding_scale != 0.0f && (ubatch.token || hparams.n_deepstack_layers == 0)) {
+    const bool scale_tok_only = hparams.f_embedding_scale != 0.0f && hparams.n_deepstack_layers > 0;
+
+    // For Granite architecture
+    if (hparams.f_embedding_scale != 0.0f && !scale_tok_only) {
         if (!ggml_is_contiguous(cur)) {
             cur = ggml_cont(ctx0, cur);
         }
         cur = ggml_scale(ctx0, cur, hparams.f_embedding_scale);
     }
 
+    // scale the token rows only, applied after the select so that the graph is the same for any batch contents
+    inp->scale_tok = tok_scale*(scale_tok_only ? hparams.f_embedding_scale : 1.0f);
+    if (inp->scale_tok != 1.0f) {
+        inp->scale_rows = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, ubatch.n_tokens);
+        cb(inp->scale_rows, "inp_scale_rows", -1);
+        ggml_set_input(inp->scale_rows);
+
+        if (!ggml_is_contiguous(cur)) {
+            cur = ggml_cont(ctx0, cur);
+        }
+        cur = ggml_mul(ctx0, cur, inp->scale_rows);
+    }
+
     cb(cur, "embd", -1);
 
     res->add_input(std::move(inp));

src/llama-graph.h

@@ -134,8 +134,14 @@ public:
 
     bool can_reuse(const llm_graph_params & params) override;
 
-    ggml_tensor * tokens = nullptr; // I32 [n_batch]
-    ggml_tensor * embd   = nullptr; // F32 [n_embd, n_batch]
+    ggml_tensor * tokens       = nullptr; // I32 [n_batch]
+    ggml_tensor * embd         = nullptr; // F32 [n_embd, n_batch]
+    ggml_tensor * mixed_tokens = nullptr; // I32 [n_tok_rows], mixed path: ids of the token rows
+    ggml_tensor * mixed_slots  = nullptr; // I64 [n_tok_rows], mixed path: batch index of the token rows
+    ggml_tensor * mixed_embd   = nullptr; // F32 [n_embd, n_batch], mixed path: embd rows, token rows are overwritten
+    ggml_tensor * scale_rows   = nullptr; // F32 [1, n_batch], per-row scale: scale_tok for token rows, 1 for embd rows
+
+    float scale_tok = 1.0f;
 
     const int64_t n_embd = 0;
 };
@@ -823,6 +829,7 @@ struct llm_graph_params {
             ubatch.n_seq_tokens == other.ubatch.n_seq_tokens &&
             ubatch.n_seqs       == other.ubatch.n_seqs &&
             ubatch.n_seqs_unq   == other.ubatch.n_seqs_unq &&
+            ubatch.is_mixed()   == other.ubatch.is_mixed() &&
             (
                 (!ubatch.token && !other.ubatch.token) ||
                 (!ubatch.embd  && !other.ubatch.embd)  ||
@@ -1166,7 +1173,8 @@ struct llm_graph_context {
     // inputs
     //
 
-    ggml_tensor * build_inp_embd(ggml_tensor * tok_embd) const;
+    // tok_scale: applied to token rows only
+    ggml_tensor * build_inp_embd(ggml_tensor * tok_embd, float tok_scale = 1.0f) const;
     ggml_tensor * build_inp_pos() const;
     ggml_tensor * build_inp_attn_scale() const;
     ggml_tensor * build_inp_out_ids() const;

src/llama-kv-cache-dsv4.cpp

@@ -83,6 +83,7 @@ static llama_ubatch dsv4_build_raw_write_ubatch(const llama_ubatch & ubatch) {
     if (!dsv4_ubatch_has_coupled(ubatch)) {
         return ubatch;
     }
+    ASSERT_EMBD_OR_TOKEN(ubatch);
     if (ubatch.embd) {
         throw std::runtime_error("DSV4 coupled embedding ubatches are not supported");
     }
@@ -164,6 +165,7 @@ static llama_ubatch dsv4_build_raw_write_ubatch(const llama_ubatch & ubatch) {
         /*.seq_id_unq   =*/ data->seq_id_unq.data(),
         /*.seq_idx      =*/ data->seq_idx.data(),
         /*.output       =*/ data->output.data(),
+        /*.type         =*/ nullptr,
         /*.decision_order =*/ nullptr,
         /*.data         =*/ data,
     };

src/llama-kv-cache.cpp

@@ -309,8 +309,8 @@ llama_kv_cache::llama_kv_cache(
         n_embd_head_k_all = other->n_embd_head_k_all;
         n_embd_head_v_all = other->n_embd_head_v_all;
 
-        attn_rot_k = other->attn_rot_k;
-        attn_rot_v = other->attn_rot_v;
+        n_rot_k = other->n_rot_k;
+        n_rot_v = other->n_rot_v;
     } else {
         const char * LLAMA_ATTN_ROT_DISABLE = getenv("LLAMA_ATTN_ROT_DISABLE");
         const bool attn_rot_disable = LLAMA_ATTN_ROT_DISABLE ? atoi(LLAMA_ATTN_ROT_DISABLE) : false;
@@ -318,7 +318,7 @@ llama_kv_cache::llama_kv_cache(
             LLAMA_LOG_WARN("%s: attention rotation force disabled (LLAMA_ATTN_ROT_DISABLE)\n", __func__);
         }
 
-        attn_rot_k =
+        bool attn_rot_k =
             !attn_rot_disable &&
             n_embd_head_k_all > 0 &&
             ggml_is_quantized(type_k) &&
@@ -331,19 +331,37 @@ llama_kv_cache::llama_kv_cache(
             attn_rot_k = true;
         }
 
-        attn_rot_v =
+        const bool attn_rot_v =
             !attn_rot_disable &&
             n_embd_head_v_all > 0 &&
             ggml_is_quantized(type_v) &&
             hparams.n_embd_head_v() % 64 == 0;
+
+        if (attn_rot_k) {
+            n_rot_k = 64;
+
+            // TODO: investigate if using the smallest rotation matrix is beneficial also for K (similar as for V)
+            // ref: https://github.com/ggml-org/llama.cpp/pull/21038#issuecomment-4141323088
+            while (n_embd_head_k_all % (2*n_rot_k) == 0) {
+                n_rot_k *= 2;
+            }
+        }
+        if (attn_rot_v) {
+            n_rot_v = 64;
+            // using smaller rotation matrices for V seems beneficial
+            // ref: https://github.com/ggml-org/llama.cpp/pull/21038#issuecomment-4146397570
+            //while (hparams.n_embd_head_v() % (2*n_rot_v) == 0) {
+            //    n_rot_v *= 2;
+            //}
+        }
     }
 
-    LLAMA_LOG_INFO("%s: attn_rot_k = %d, n_embd_head_k_all = %d\n", __func__, attn_rot_k, n_embd_head_k_all);
-    LLAMA_LOG_INFO("%s: attn_rot_v = %d, n_embd_head_k_all = %d\n", __func__, attn_rot_v, n_embd_head_v_all);
+    LLAMA_LOG_INFO("%s: n_rot_k = %u, n_embd_head_k_all = %d\n", __func__, n_rot_k, n_embd_head_k_all);
+    LLAMA_LOG_INFO("%s: n_rot_v = %u, n_embd_head_k_all = %d\n", __func__, n_rot_v, n_embd_head_v_all);
 
     // pre-compute the haramard matrices and keep them in host memory
     // TODO: in the future, we can make copies in the backend buffers to avoid host -> device transfers
-    if (attn_rot_k || attn_rot_v) {
+    if (n_rot_k || n_rot_v) {
         for (int64_t n = 64; n <= std::max(n_embd_head_k_all, n_embd_head_v_all); n *= 2) {
             attn_rot_hadamard[n] = std::vector<float>(n*n);
 
@@ -1138,7 +1156,9 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
                     ext.y = ubatch.pos[i + ubatch.n_tokens];
                 }
 
-                if (ubatch.token) {
+                const bool is_embd = !ubatch.token || (ubatch.is_mixed() && ubatch.type[i]);
+
+                if (!is_embd) {
                     ext.tok = ubatch.token[i];
                 } else if (hparams.ple_n_heads > 0) {
                     // embd batch (multimodal input) has no token ids, need to pad it with the correct ID for PLE layers
@@ -1447,17 +1467,8 @@ ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama
 ggml_tensor * llama_kv_cache::build_input_k_rot(ggml_context * ctx) const {
     ggml_tensor * res = nullptr;
 
-    if (attn_rot_k) {
-        int nrot = 64;
-
-        // TODO: investigate if using the smallest rotation matrix is beneficial also for K (similar as for V)
-        // ref: https://github.com/ggml-org/llama.cpp/pull/21038#issuecomment-4141323088
-        do {
-            nrot *= 2;
-        } while (n_embd_head_k_all % nrot == 0);
-        nrot /= 2;
-
-        res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nrot, nrot);
+    if (n_rot_k) {
+        res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_rot_k, n_rot_k);
         ggml_set_input(res);
         ggml_set_name(res, "attn_inp_k_rot");
     }
@@ -1468,16 +1479,8 @@ ggml_tensor * llama_kv_cache::build_input_k_rot(ggml_context * ctx) const {
 ggml_tensor * llama_kv_cache::build_input_v_rot(ggml_context * ctx) const {
     ggml_tensor * res = nullptr;
 
-    if (attn_rot_v) {
-        int nrot = 64;
-        // using smaller rotation matrices for V seems beneficial
-        // ref: https://github.com/ggml-org/llama.cpp/pull/21038#issuecomment-4146397570
-        //do {
-        //    nrot *= 2;
-        //} while (hparams.n_embd_head_v() % nrot == 0);
-        //nrot /= 2;
-
-        res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nrot, nrot);
+    if (n_rot_v) {
+        res = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_rot_v, n_rot_v);
         ggml_set_input(res);
         ggml_set_name(res, "attn_inp_v_rot");
     }
@@ -1862,10 +1865,13 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st
 
     // an embd (multimodal) ubatch can repeat one position for a whole image, so positions
     // do not encode the token order; resolve its predecessors by ubatch order instead
+    // same for a mixed ubatch
+    const bool by_order = !ubatch.token || ubatch.is_mixed();
+
     std::vector<uint32_t> ord; // index among the ubatch tokens of the same seq
     std::unordered_map<llama_seq_id, std::vector<uint32_t>> seq_idx;
 
-    if (!ubatch.token) {
+    if (by_order) {
         ord.resize(n_tokens);
         for (uint32_t i = 0; i < n_tokens; ++i) {
             auto & v = seq_idx[ubatch.seq_id[i][0]];
@@ -1883,7 +1889,7 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st
             const llama_pos d = (llama_pos) (n - j);
 
             llama_pos p;
-            if (!ubatch.token) {
+            if (by_order) {
                 const auto & v = seq_idx[seq_id];
                 const int64_t k = (int64_t) ord[i] - d;
                 // k >= 0: an earlier token of this very ubatch; k < 0: before the chunk
@@ -2249,6 +2255,8 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t
 
     io.write(&v_trans, sizeof(v_trans));
     io.write(&n_layer, sizeof(n_layer));
+    io.write(&n_rot_k, sizeof(n_rot_k));
+    io.write(&n_rot_v, sizeof(n_rot_v));
 
     // Iterate and write all the keys first, each row is a cell
     // Get whole range at a time
@@ -2541,9 +2549,13 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
 
     uint32_t v_trans;
     uint32_t n_layer;
+    uint32_t n_rot_k_ref;
+    uint32_t n_rot_v_ref;
 
     io.read(&v_trans, sizeof(v_trans));
     io.read(&n_layer, sizeof(n_layer));
+    io.read(&n_rot_k_ref, sizeof(n_rot_k_ref));
+    io.read(&n_rot_v_ref, sizeof(n_rot_v_ref));
 
     if (n_layer != layers.size()) {
         LLAMA_LOG_ERROR("%s: mismatched layer count (%u instead of %u)\n", __func__, n_layer, (uint32_t) layers.size());
@@ -2560,6 +2572,16 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
         return false;
     }
 
+    if (n_rot_k_ref != n_rot_k) {
+        LLAMA_LOG_ERROR("%s: incompatible key rotation (%u instead of %u)\n", __func__, n_rot_k_ref, n_rot_k);
+        return false;
+    }
+
+    if (n_rot_v_ref != n_rot_v) {
+        LLAMA_LOG_ERROR("%s: incompatible value rotation (%u instead of %u)\n", __func__, n_rot_v_ref, n_rot_v);
+        return false;
+    }
+
     // For each layer, read the keys for each cell, one row is one cell, read as one contiguous block
     for (const auto & layer : layers) {
         const uint32_t il = layer.il;

src/llama-kv-cache.h

@@ -278,8 +278,8 @@ private:
     const uint32_t n_swa = 0;
 
     // env: LLAMA_ATTN_ROT_DISABLE
-    bool attn_rot_k = false;
-    bool attn_rot_v = false;
+    uint32_t n_rot_k = 0;
+    uint32_t n_rot_v = 0;
 
     // if all layers participating in the cache have constant head size, the value is stored here
     // otherwise the value is -1

src/llama-memory-hybrid-idx.cpp

@@ -192,24 +192,21 @@ bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_po
         const llama_pos stale = mem_idx_stale_pos(seq_id, p0);
         mem_idx->seq_rm(seq_id, p0, p1);
         mem_idx_stale_set(seq_id, stale);
-
-        // removing a sequence can free cells another sequence shared, but only this one is marked stale, so the
-        // survivor would keep shared = true and pin cache_safe off forever; stale every sequence to re-derive it
-        if (kpool_layout_shared()) {
-            mem_idx_stale_set(-1, 0);
-        }
     }
 
     return get_mem_attn()->seq_rm(seq_id, p0, p1);
 }
 
 void llama_memory_hybrid_idx::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
+    // only whole sequences are copied: the recurrent state ignores the range, and a shared cell holds a single pool grouping
+    GGML_ASSERT(p0 <= 0 && p1 < 0 && "partial seq_cp is not supported");
+
     llama_memory_hybrid::seq_cp(seq_id_src, seq_id_dst, p0, p1);
 
     if (mem_idx) {
         mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1);
-        // the copy shares cells, which cannot hold two groupings, so both sides drop their cached keys
-        mem_idx_stale_set(seq_id_src, 0);
+        // a whole sequence copy gives the destination the source's pools, rep rows included: the source keeps its
+        // pooled keys, the destination rebuilds its layout and re-pools into the same rows
         mem_idx_stale_set(seq_id_dst, 0);
     }
 }
@@ -219,7 +216,7 @@ void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) {
 
     if (mem_idx) {
         mem_idx->seq_keep(seq_id);
-        // cells shared with the dropped sequences become exclusive again, their keys were never cached
+        // every other sequence loses its cells, so their layouts must rebuild
         mem_idx_stale_set(-1, 0);
     }
 }
@@ -291,10 +288,6 @@ void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_
                 mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn);
                 // the restore rewrites the cells behind the pool layout's back
                 mem_idx_stale_set(seq_id, 0);
-                // it can also change which cells are shared; re-derive sharing for every sequence, as seq_rm does
-                if (kpool_layout_shared()) {
-                    mem_idx_stale_set(-1, 0);
-                }
             }
         }
 
@@ -321,10 +314,6 @@ void llama_memory_hybrid_idx::state_drop(llama_seq_id seq_id) {
     if (mem_idx) {
         mem_idx->state_clear(seq_id);
         mem_idx_stale_set(seq_id, 0);
-        // clearing this sequence can end a sharing the survivor would otherwise keep flagged (see seq_rm)
-        if (kpool_layout_shared()) {
-            mem_idx_stale_set(-1, 0);
-        }
     }
 }
 
@@ -358,26 +347,22 @@ struct llama_memory_hybrid_idx::kpool_layout {
 
         // Where the pool scan stopped, so an append resumes instead of starting over.
         size_t j_next = 0;
-
-        // Whether any cell also carries another sequence, which rules out caching this sequence's pooled keys.
-        bool shared = false;
     };
 
     std::array<seq, LLAMA_MAX_SEQ> seqs;
 
     uint32_t n_pool_real = 0;
-    bool cache_safe      = true;
 };
 
 // Which pools of the layout the current ubatch must re-pool, in the layout's pool order.
 struct llama_memory_hybrid_idx_context::kpool_state {
     std::vector<uint32_t> is_new;
+    std::vector<uint32_t> rep_gen; // per global cell, the generation that last marked a pool with that rep
     uint32_t generation = 0;
 
     uint32_t n_pool_real = 0;
     uint32_t n_new       = 0;
     uint32_t n_new_g     = 1; // graph size of the new pool list, stable across decode steps
-    bool     cache_safe  = true;
 };
 
 namespace {
@@ -404,10 +389,6 @@ const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_lay
     return *kpool_lay;
 }
 
-bool llama_memory_hybrid_idx::kpool_layout_shared() const {
-    return kpool_lay && !kpool_lay->cache_safe;
-}
-
 // Pools are fixed by the positions relative to the sequence's first one, so the layout survives a plain
 // append. A sequence edit can regroup them, and mem_idx_stale tells us it happened.
 const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_layout_update() {
@@ -424,7 +405,6 @@ const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_lay
     const bool     unified     = n_stream_kv == 1;
 
     lay.n_pool_real = 0;
-    lay.cache_safe  = true;
 
     for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
         auto & sq = lay.seqs[s];
@@ -440,10 +420,8 @@ const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_lay
 
         sq.strm = unified ? 0 : mem_idx->get_stream(s);
 
-        size_t n_kept = 0;
         if (mem_idx_stale[s] == POS_CLEAN && !sq.cells.empty() && !sp.empty() &&
                 sq.pos_min == sp.begin()->first) {
-            n_kept = sq.cells.size();
             for (auto it = sp.upper_bound(sq.cells.back()); it != sp.end(); ++it) {
                 sq.cells.push_back(*it);
             }
@@ -455,21 +433,7 @@ const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_lay
             sq.cells.assign(sp.begin(), sp.end());
             sq.pools.clear();
             sq.j_next  = 0;
-            sq.shared  = false;
             sq.pos_min = sp.empty() ? 0 : sp.begin()->first;
-            n_kept     = 0;
-        }
-
-        // sharing starts with a seq_cp; it ends with an edit, or a seq_rm/state_drop/state_read that frees the
-        // shared cells - each stales every sequence so the rebuild above re-derives it, so once set it holds
-        // until then and the rescan can be skipped
-        if (unified && !sq.shared) {
-            for (size_t j = n_kept; j < sq.cells.size(); ++j) {
-                if (cells.seq_count(sq.cells[j].second) > 1) {
-                    sq.shared = true;
-                    break;
-                }
-            }
         }
 
         // Pools start at the first valid token
@@ -504,7 +468,6 @@ const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_lay
         sq.j_next = j;
 
         lay.n_pool_real += (uint32_t) sq.pools.size();
-        lay.cache_safe   = lay.cache_safe && !sq.shared;
     }
 
     return lay;
@@ -664,7 +627,6 @@ llama_memory_hybrid_idx_context::kpool_state llama_memory_hybrid_idx_context::kp
 
     kpool_state st;
     st.n_pool_real = lay.n_pool_real;
-    st.cache_safe  = lay.cache_safe;
 
     return st;
 }
@@ -682,34 +644,40 @@ void llama_memory_hybrid_idx_context::kpool_build_state(const llama_ubatch & uba
     const auto & lay = mem->kpool_layout_get();
     auto & st = *kpool_st;
 
+    const auto *   idx     = mem->get_mem_idx();
+    const uint32_t kv_size = idx->get_size();
+    const uint32_t kpool   = mem->get_kpool();
+
     st.n_pool_real = lay.n_pool_real;
-    st.cache_safe  = lay.cache_safe;
     st.n_new       = 0;
     if (++st.generation == 0) {
-        std::fill(st.is_new.begin(), st.is_new.end(), 0);
+        std::fill(st.is_new.begin(),  st.is_new.end(),  0);
+        std::fill(st.rep_gen.begin(), st.rep_gen.end(), 0);
         st.generation = 1;
     }
     st.is_new.resize(lay.n_pool_real, 0);
+    st.rep_gen.resize((size_t) kv_size*idx->get_n_stream(), 0);
 
-    auto mark = [&](uint32_t ip) {
-        if (st.is_new[ip] != st.generation) {
-            st.is_new[ip] = st.generation;
+    std::array<uint32_t, LLAMA_MAX_SEQ> pool_start;
+
+    // a pool is marked once per rep: sequences sharing cells (a seq_cp, or tokens decoded for several sequences)
+    // share their pools, whose single pooled row they all read through pool_cells, so the scatter rows stay unique
+    auto mark = [&](llama_seq_id s, size_t k) {
+        const auto & sq  = lay.seqs[s];
+        const size_t rep = (size_t) sq.strm*kv_size + sq.cells[sq.pools[k] + kpool - 1].second;
+        if (st.rep_gen[rep] != st.generation) {
+            st.rep_gen[rep] = st.generation;
+            st.is_new[pool_start[s] + k] = st.generation;
             ++st.n_new;
         }
     };
 
-    const uint32_t kpool = mem->get_kpool();
-    std::array<uint32_t, LLAMA_MAX_SEQ> pool_start;
     uint32_t ip = 0;
     for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
         const auto & sq = lay.seqs[s];
         pool_start[s] = ip;
         ip += (uint32_t) sq.pools.size();
 
-        if (!st.cache_safe) {
-            continue;
-        }
-
         // A sequence edit invalidates only pools ending after the edited position.
         const llama_pos stale_from = i_cur == 0 ?
             mem_idx_stale_batch[s] : llama_memory_hybrid_idx::POS_CLEAN;
@@ -720,18 +688,11 @@ void llama_memory_hybrid_idx_context::kpool_build_state(const llama_ubatch & uba
         auto first = std::lower_bound(sq.pools.begin(), sq.pools.end(), stale_from,
                 [&](uint32_t j, llama_pos p) { return sq.cells[j + kpool - 1].first < p; });
         for (auto it = first; it != sq.pools.end(); ++it) {
-            mark(pool_start[s] + (uint32_t) (it - sq.pools.begin()));
+            mark(s, it - sq.pools.begin());
         }
     }
     GGML_ASSERT(ip == st.is_new.size());
 
-    if (!st.cache_safe) {
-        std::fill(st.is_new.begin(), st.is_new.end(), st.generation);
-        st.n_new   = st.n_pool_real;
-        st.n_new_g = std::max(st.n_new, 1u);
-        return;
-    }
-
     // in order mode a token's cell gives its rank, and the rank its pool: positions cannot, as an image shares one
     const bool by_order = mem->get_kpool_by_order();
     const auto *   sinfo = by_order ? &sinfos_kpool[i_cur] : nullptr;
@@ -746,7 +707,7 @@ void llama_memory_hybrid_idx_context::kpool_build_state(const llama_ubatch & uba
                 const int64_t r = kpool_rank(sq.cells, p, sinfo->idxs[i / n_tps][i % n_tps]);
                 GGML_ASSERT(r >= 0);
                 if ((size_t) r / kpool < sq.pools.size()) {
-                    mark(pool_start[s] + (uint32_t) (r / kpool));
+                    mark(s, (size_t) r / kpool);
                 }
                 continue;
             }
@@ -757,15 +718,14 @@ void llama_memory_hybrid_idx_context::kpool_build_state(const llama_ubatch & uba
             }
             --it;
             if (p <= sq.cells[*it + kpool - 1].first) {
-                mark(pool_start[s] + (uint32_t) (it - sq.pools.begin()));
+                mark(s, it - sq.pools.begin());
             }
         }
     }
 
     // a ubatch touches at most t_s/kpool + 1 pools per sequence, pad to that bound so the graph keeps its shape
     // as the count moves; reserve sizes the list for every pool the cache can hold, so never pad past n_pool_max
-    const auto *   idx        = mem->get_mem_idx();
-    const uint32_t n_pool_max = idx->get_size() / kpool * idx->get_n_seq_max();
+    const uint32_t n_pool_max = kv_size / kpool * idx->get_n_seq_max();
     const uint32_t bound = ubatch.n_tokens/kpool + ubatch.n_seqs_unq;
     st.n_new_g = std::max({st.n_new, 1u, std::min({bound, kpool_pad(st.n_pool_real) - 1, n_pool_max})});
 }
@@ -784,10 +744,6 @@ uint32_t llama_memory_hybrid_idx_context::get_n_kpool_new() const {
     return kpool_cur().n_new_g;
 }
 
-bool llama_memory_hybrid_idx_context::get_kpool_cache_safe() const {
-    return kpool_cur().cache_safe;
-}
-
 void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells, ggml_tensor * pool_idxs, ggml_tensor * pool_mask, ggml_tensor * tail_idxs,
         ggml_tensor * gather_mask, bool gather, ggml_tensor * new_pool_idxs, ggml_tensor * new_pool_rep,
         const llama_ubatch * ubatch, ggml_tensor * new_pool_pos) const {
@@ -816,13 +772,11 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells,
     GGML_ASSERT(pool_mask->ne[0] == (int64_t) n_pool && pool_mask->ne[1] == (int64_t) n_tokens);
     GGML_ASSERT(tail_idxs->ne[0] == (int64_t) kpool - 1 && tail_idxs->ne[1] == (int64_t) n_tokens);
     GGML_ASSERT(pool_idxs->ne[0] == (int64_t) kpool && pool_idxs->ne[1] == (int64_t) n_pool);
-    GGML_ASSERT(st.cache_safe == (new_pool_rep != nullptr));
     GGML_ASSERT(ggml_backend_buffer_is_host(new_pool_idxs->buffer));
     GGML_ASSERT(new_pool_idxs->ne[0] == (int64_t) kpool && new_pool_idxs->ne[1] == (int64_t) n_new_g);
-    if (new_pool_rep != nullptr) {
-        GGML_ASSERT(ggml_backend_buffer_is_host(new_pool_rep->buffer));
-        GGML_ASSERT(new_pool_rep->ne[0] == (int64_t) n_new_g);
-    }
+    // the graph always scatters the fresh pooled keys back into the cache, see build_qsa_sel
+    GGML_ASSERT(new_pool_rep != nullptr && ggml_backend_buffer_is_host(new_pool_rep->buffer));
+    GGML_ASSERT(new_pool_rep->ne[0] == (int64_t) n_new_g);
     if (new_pool_pos != nullptr) {
         GGML_ASSERT(ggml_backend_buffer_is_host(new_pool_pos->buffer));
         GGML_ASSERT(new_pool_pos->ne[0] == 4*(int64_t) n_new_g);
@@ -891,7 +845,7 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells,
     int32_t * pcell = (int32_t *) pool_cells->data;
     int32_t * pidx  = (int32_t *) pool_idxs->data;
     int32_t * nidx  = (int32_t *) new_pool_idxs->data;
-    int64_t * nrep  = new_pool_rep != nullptr ? (int64_t *) new_pool_rep->data : nullptr;
+    int64_t * nrep  = (int64_t *) new_pool_rep->data;
     int32_t * npos  = new_pool_pos != nullptr ? (int32_t *) new_pool_pos->data : nullptr;
 
     if (npos != nullptr) {
@@ -924,9 +878,7 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells,
                 for (uint32_t k = 0; k < kpool; ++k) {
                     nidx[(size_t) i_new*kpool + k] = (int32_t) gcell(sq, sq.cells[j + k].second);
                 }
-                if (nrep != nullptr) {
-                    nrep[i_new] = gcell(sq, rep);
-                }
+                nrep[i_new] = gcell(sq, rep);
                 if (npos != nullptr) {
                     // a pooled key is rotated to the M-RoPE position of its first member
                     const uint32_t c = sq.cells[j].second;
@@ -959,9 +911,7 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells,
             for (uint32_t k = 0; k < kpool; ++k) {
                 nidx[(size_t) i*kpool + k] = (int32_t) pad_cell;
             }
-            if (nrep != nullptr) {
-                nrep[i] = pad_cell;
-            }
+            nrep[i] = pad_cell;
         }
     }
 

src/llama-memory-hybrid-idx.h

@@ -121,9 +121,6 @@ private:
     // unique_ptr because kpool_layout is incomplete here
     std::unique_ptr<kpool_layout> kpool_lay;
 
-    // whether the current layout has cells shared between sequences (kpool_layout is incomplete here, so out of line)
-    bool kpool_layout_shared() const;
-
     // seq_id < 0 stales every sequence, p0 < 0 stales the sequence from its first position
     void mem_idx_stale_set(llama_seq_id seq_id, llama_pos p0);
 
@@ -194,7 +191,6 @@ public:
     // glm5-next and qwen4exp, complete pools of kpool cells per sequence, scored as whole pools.
     uint32_t get_n_kpool    () const; // Padded pool count, where the last pool is always unused.
     uint32_t get_n_kpool_new() const; // Pools to re-pool this ubatch, padded to a stable bound, never below 1.
-    bool get_kpool_cache_safe() const;
     kpool_access get_kpool_access(ggml_context * ctx, int32_t il, int64_t n_embd) const;
     ggml_tensor * gather_mla_rows(ggml_context * ctx, ggml_tensor * idxs, int64_t n_rows, int64_t n_embd, int32_t il) const;
     // new_pool_pos (I32 [4*n_new]): M-RoPE position of each new pool's first member, for pooled keys rotated at pooling time

src/models/clef.cpp

@@ -125,7 +125,7 @@ static clef_spans clef_get_spans(const llama_ubatch & ubatch) {
     std::vector<bool> has_option;
 
     // TODO: support multiple sequences
-    bool ok = ubatch.decision_order != nullptr && ubatch.n_seqs_unq == 1;
+    bool ok = ubatch.decision_order != nullptr && ubatch.token != nullptr && ubatch.n_seqs_unq == 1;
 
     const int32_t n_tokens = ubatch.n_tokens;
     for (int32_t i = 0; ok && i < n_tokens;) {
@@ -156,6 +156,11 @@ static clef_spans clef_get_spans(const llama_ubatch & ubatch) {
         i = end;
     }
 
+    // the head reads the token ids of the spans, they cannot be embeddings
+    for (int32_t i = 0; ok && ubatch.is_mixed() && i < n_tokens; i++) {
+        ok = !ubatch.type[i] || ubatch.decision_order[i] == LLAMA_DECISION_ORDER_NONE;
+    }
+
     // each question needs an option
     res.valid = ok && !res.questions.empty() && std::find(has_option.begin(), has_option.end(), false) == has_option.end();
     if (!res.valid) {
@@ -177,8 +182,9 @@ public:
     }
 
     void set_input(const llama_ubatch * ubatch) override {
-        GGML_ASSERT(ubatch->token);
-        ggml_backend_tensor_set(tokens, ubatch->token, 0, n_tokens * sizeof(llama_token));
+        // a batch of embeddings has no token ids, and no usable spans
+        const std::vector<llama_token> no_tokens(ubatch->token ? 0 : n_tokens, 0);
+        ggml_backend_tensor_set(tokens, ubatch->token ? ubatch->token : no_tokens.data(), 0, n_tokens * sizeof(llama_token));
 
         const auto spans = clef_get_spans(*ubatch);
         GGML_ASSERT(spans.questions.size() == n_questions && spans.options.size() == n_options);
@@ -308,12 +314,35 @@ llama_model_clef::graph::graph(const llama_model & model_base, const llm_graph_p
     ggml_build_forward_expand(gf, cur);
 }
 
-// same as build_attn_inp_no_cache(), but the mask is causal even if the batch is processed by the encoder path
-llm_graph_input_attn_no_cache * llama_model_clef::graph::build_attn_inp_causal() {
-    llama_cparams cparams_causal = cparams;
-    cparams_causal.causal_attn = true;
+// causal mask by batch order: the tokens of an image share the same position
+class llm_graph_input_attn_clef : public llm_graph_input_attn_no_cache {
+public:
+    using llm_graph_input_attn_no_cache::llm_graph_input_attn_no_cache;
 
-    auto inp = std::make_unique<llm_graph_input_attn_no_cache>(hparams, cparams_causal);
+    void set_input(const llama_ubatch * ubatch) override {
+        const int64_t n_tokens = ubatch->n_tokens;
+
+        const auto fill_mask = [&](auto * data, auto zero, auto ninf) {
+            for (int64_t i1 = 0; i1 < n_tokens; ++i1) {
+                for (int64_t i0 = 0; i0 < n_tokens; ++i0) {
+                    const bool visible = i0 <= i1 && ubatch->seq_id[i0][0] == ubatch->seq_id[i1][0];
+                    data[i1 * n_tokens + i0] = visible ? zero : ninf;
+                }
+            }
+        };
+
+        GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask->buffer));
+        if (self_kq_mask->type == GGML_TYPE_F16) {
+            fill_mask((ggml_fp16_t *) self_kq_mask->data, ggml_fp32_to_fp16(0.0f), ggml_fp32_to_fp16(-INFINITY));
+        } else {
+            fill_mask((float *) self_kq_mask->data, 0.0f, -INFINITY);
+        }
+    }
+};
+
+// same as build_attn_inp_no_cache(), with a causal mask even if the batch is processed by the encoder path
+llm_graph_input_attn_no_cache * llama_model_clef::graph::build_attn_inp_causal() {
+    auto inp = std::make_unique<llm_graph_input_attn_clef>(hparams, cparams);
 
     const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
 

src/models/cogvlm.cpp

@@ -70,6 +70,7 @@ llama_model_cogvlm::graph::graph(const llama_model & model, const llm_graph_para
     // check ubatch to see if we have input tokens (text)
     // or an input embedding vector (image)
     bool is_text;
+    ASSERT_EMBD_OR_TOKEN(ubatch);
     if (ubatch.token) {
         is_text = true;
     } else {

src/models/cohere2moe.cpp

@@ -317,6 +317,7 @@ llama_model_cohere2moe::graph_mtp::graph_mtp(const llama_model & model, const ll
     // TODO: make static using `ggml_build_forward_select()`
     //       see llm_graph_context::build_inp_embd() for reference
     ggml_tensor * tok_embd;
+    ASSERT_EMBD_OR_TOKEN(ubatch);
     if (ubatch.token) {
         ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
         tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);

src/models/deepseek2.cpp

@@ -221,6 +221,7 @@ llama_model_deepseek2::graph_mtp::graph_mtp(const llama_model & model, const llm
     ggml_set_input(inp->embd);
 
     ggml_tensor * tok_embd;
+    ASSERT_EMBD_OR_TOKEN(ubatch);
     if (ubatch.token) {
         ggml_tensor * tok_embd_w = layer.nextn.embed_tokens
                 ? layer.nextn.embed_tokens

src/models/deepseek32.cpp

@@ -535,6 +535,7 @@ llama_model_deepseek32::graph_mtp::graph_mtp(const llama_model & model, const ll
     ggml_set_input(inp->embd);
 
     ggml_tensor * tok_embd;
+    ASSERT_EMBD_OR_TOKEN(ubatch);
     if (ubatch.token) {
         ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
 

src/models/deepseek4.cpp

@@ -1281,6 +1281,7 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
         ggml_tensor * exp_probs_b = layer.ffn_exp_probs_b;
 
         // may apply exp_probs_b_vl is input is from mtmd
+        ASSERT_EMBD_OR_TOKEN(ubatch);
         const bool is_media = ubatch.embd != nullptr;
         if (is_media) {
             if (layer.ffn_exp_probs_b_vl) {
@@ -1366,6 +1367,7 @@ llama_model_deepseek4::graph_mtp::graph_mtp(const llama_model & model, const llm
     GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
             cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
             "nextn_layer_offset out of range [0, n_layer_nextn)");
+    ASSERT_EMBD_OR_TOKEN(ubatch);
     GGML_ASSERT(ubatch.token && "DEEPSEEK4 MTP requires token input");
 
     const int64_t hc = hparams.dsv4_hc_mult;

src/models/dflash.cpp

@@ -604,6 +604,7 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
     };
 
     // KV cache injection
+    ASSERT_EMBD_OR_TOKEN(ubatch);
     if (ubatch.embd) {
         auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
 
@@ -870,6 +871,7 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
     llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
 
     // KV cache injection: fused target features from the encoder
+    ASSERT_EMBD_OR_TOKEN(ubatch);
     if (ubatch.embd) {
         auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
 

src/models/gemma-embedding.cpp

@@ -77,10 +77,8 @@ llama_model_gemma_embedding::graph::graph(const llama_model & model, const llm_g
     ggml_tensor * cur;
     ggml_tensor * inpL;
 
-    inpL = build_inp_embd(model.tok_embd);
-
     // important: do not normalize weights for raw embeddings input (i.e. encoded image embeddings)
-    inpL = ggml_scale(ctx0, inpL, ubatch.token ? sqrtf(n_embd) : 1.0f);
+    inpL = build_inp_embd(model.tok_embd, sqrtf(n_embd));
     cb(inpL, "inp_scaled", -1);
 
     // inp_pos - contains the positions

src/models/gemma3.cpp

@@ -85,10 +85,8 @@ llama_model_gemma3::graph<iswa>::graph(const llama_model & model, const llm_grap
     ggml_tensor * cur;
     ggml_tensor * inpL;
 
-    inpL = build_inp_embd(model.tok_embd);
-
     // important: do not normalize weights for raw embeddings input (i.e. encoded image embeddings)
-    inpL = ggml_scale(ctx0, inpL, ubatch.token ? sqrtf(n_embd) : 1.0f);
+    inpL = build_inp_embd(model.tok_embd, sqrtf(n_embd));
     cb(inpL, "inp_scaled", -1);
 
     // inp_pos - contains the positions

src/models/gemma3n.cpp

@@ -96,10 +96,8 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par
     ggml_tensor * cur;
     ggml_tensor * inpL;
 
-    inpL = build_inp_embd(model.tok_embd);
-
     // important: do not normalize weights for raw embeddings input (i.e. encoded image embeddings)
-    inpL = ggml_scale(ctx0, inpL, ubatch.token ? sqrtf(n_embd) : 1.0f);
+    inpL = build_inp_embd(model.tok_embd, sqrtf(n_embd));
     cb(inpL, "inp_scaled", -1);
 
     // inp_pos - contains the positions
@@ -325,6 +323,7 @@ ggml_tensor * llama_model_gemma3n::graph::build_inp_per_layer() {
     auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
     ggml_tensor * inp_per_layer;
     float tok_embd_scale = sqrtf((float) n_embd_altup);
+    // mixed ubatch: embd rows have token id 0, same padding row as below
     if (ubatch.token) {
         inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens);
         ggml_set_input(inp->tokens);

src/models/gemma4.cpp

@@ -159,10 +159,8 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
     ggml_tensor * cur;
     ggml_tensor * inpL;
 
-    inpL = build_inp_embd(model.tok_embd);
-
     // important: do not normalize weights for raw embeddings input (i.e. encoded image emdeddings)
-    inpL = ggml_scale(ctx0, inpL, ubatch.token ? sqrtf(n_embd) : 1.0f);
+    inpL = build_inp_embd(model.tok_embd, sqrtf(n_embd));
     cb(inpL, "inp_scaled", -1);
 
     // inp_pos - contains the positions
@@ -473,6 +471,7 @@ ggml_tensor * llama_model_gemma4::graph::build_inp_per_layer() {
 
     ggml_tensor * inp_per_layer;
     float tok_embd_scale = sqrtf((float) n_embd_per_layer);
+    // mixed ubatch: embd rows have token id 0, same padding row as below
     if (ubatch.token) {
         inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens);
         ggml_set_input(inp->tokens);

src/models/glm-dsa.cpp

@@ -579,6 +579,7 @@ llama_model_glm_dsa::graph_mtp::graph_mtp(const llama_model & model, const llm_g
     ggml_set_input(inp->embd);
 
     ggml_tensor * tok_embd;
+    ASSERT_EMBD_OR_TOKEN(ubatch);
     if (ubatch.token) {
         ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
 

src/models/glm4-moe.cpp

@@ -158,6 +158,7 @@ llama_model_glm4_moe::graph_mtp::graph_mtp(const llama_model & model, const llm_
     ggml_set_input(inp->embd);
 
     ggml_tensor * tok_embd;
+    ASSERT_EMBD_OR_TOKEN(ubatch);
     if (ubatch.token) {
         ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
         tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);

src/models/glm5-next.cpp

@@ -263,7 +263,6 @@ public:
         // The scatter mask shape follows n_kv.
         res &= n_kv              == idx->get_n_kv();
         res &= n_new             == std::max(mctx->get_n_kpool_new(), 1u);
-        res &= cache_safe        == mctx->get_kpool_cache_safe();
 
         return res;
     }
@@ -282,7 +281,6 @@ public:
     const uint32_t kpool;
     uint32_t n_new = 0;
     uint32_t n_sel = 0;
-    bool cache_safe = true;
     bool gather = false;
     uint32_t n_kv  = 0;
 };
@@ -296,8 +294,7 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui
     const uint32_t n_kv   = mctx_idx->get_n_kv();
     // a ubatch that completes no pool still builds one dummy entry, so the graph does not
     // change shape every kpool tokens
-    const uint32_t n_new  = std::max(mctx_hyb->get_n_kpool_new(), 1u);
-    const bool cache_safe = mctx_hyb->get_kpool_cache_safe();
+    const uint32_t n_new = std::max(mctx_hyb->get_n_kpool_new(), 1u);
 
     // the fused lightning indexer wants an f16 mask
     const auto type_mask = cparams.fused_lid ? GGML_TYPE_F16 : GGML_TYPE_F32;
@@ -321,14 +318,17 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui
 
     inp->n_kv = n_kv;
 
-    // Gather selected latents for small decode batches when n_kv exceeds n_sel.
+    // Gather selected latents for small batches when the context exceeds the selection width.
     {
         constexpr int64_t max_ub = 16;
 
         const int64_t n_top_pool = std::min<int64_t>(n_pool, hparams.indexer_top_k / kpool);
         const int64_t n_sel      = kpool*n_top_pool + (hparams.indexer_kpool_select_tail ? kpool - 1 : 0);
         inp->n_sel = (uint32_t) n_sel;
-        inp->gather = (int64_t) n_tokens <= max_ub && (int64_t) n_kv > n_sel;
+        // both terms are context constants: n_ubatch bounds every ubatch and top_k + kpool - 1 bounds
+        // n_sel, so the graph shape follows neither n_tokens nor n_kv, which the reserve cannot predict
+        // TODO: remove "gather" logic and everything related. the backends now support sparse attension so this is obsolete
+        inp->gather = (int64_t) cparams.n_ubatch <= max_ub && (int64_t) cparams.n_ctx > hparams.indexer_top_k + kpool - 1;
 
         // Both paths read the slot mask: gather adds it to the scores, scatter maps its dead slots to dump rows.
         inp->gather_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens);
@@ -338,14 +338,12 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui
     }
 
     inp->n_new = n_new;
-    inp->cache_safe = cache_safe;
 
     inp->new_pool_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool, n_new);
     ggml_set_input(inp->new_pool_idxs);
-    if (cache_safe) {
-        inp->new_pool_rep = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_new);
-        ggml_set_input(inp->new_pool_rep);
-    }
+    // one scatter row per new pool, each a distinct rep row (see kpool_build_state)
+    inp->new_pool_rep = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_new);
+    ggml_set_input(inp->new_pool_rep);
 
     return (llm_graph_input_kpool *) res->add_input(std::move(inp));
 }
@@ -814,20 +812,10 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
     pooled_new = ggml_reshape_2d(ctx0, pooled_new, n_embd_indexer, n_new);
     cb(pooled_new, "indexer_pool_k_new", il);
 
-    if (inp_kpool->cache_safe) {
-        // Write before the pool gather.
-        ggml_build_forward_expand(gf, kpool_cache.scatter_pooled(pooled_new, inp_kpool->new_pool_rep));
-    }
-
-    ggml_tensor * pooled = nullptr;
-    if (inp_kpool->cache_safe) {
-        pooled = kpool_cache.gather_pooled(inp_kpool->pool_cells);
-    } else {
-        GGML_ASSERT(n_new <= n_pool);
-        ggml_tensor * pad = ggml_fill(ctx0,
-                ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_indexer, n_pool - n_new), 0.0f);
-        pooled = ggml_concat(ctx0, pooled_new, pad, 1);
-    }
+    // scatter the fresh pooled keys into their rep rows, then gather all n_pool of them by cell:
+    // the older pools come from the rows earlier ubatches wrote
+    ggml_build_forward_expand(gf, kpool_cache.scatter_pooled(pooled_new, inp_kpool->new_pool_rep));
+    ggml_tensor * pooled = kpool_cache.gather_pooled(inp_kpool->pool_cells);
     pooled = ggml_reshape_3d(ctx0, pooled, n_embd_indexer, 1, n_pool);
     cb(pooled, "indexer_pool_k", il);