cuda: add sqrt_softplus in topk-moe for dsv4 (#25896)

With -hfd pointing to a repo shipping speculative sidecars, the draft
resolved to the main model of that repo, since find_best_model()
excludes sidecar files, and the explicit draft plan suppressed the
sidecar discovery on the -hf repo.

The draft plan already discovers its sidecars, they were just never
consumed. Wire them as the draft, following the fallback pattern of
the main plan, so this now works as expected:

llama-server -hf repo -hfd repo --spec-type draft-dflash

.github/workflows/release.yml

@@ -1651,6 +1651,9 @@ jobs:
 
             </details>
 
+            **Website:**
+            - <https://llama.app>
+
             **macOS/iOS:**
             - [macOS Apple Silicon (arm64)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-macos-arm64.tar.gz)
             - macOS Apple Silicon (arm64, KleidiAI enabled) [DISABLED](https://github.com/ggml-org/llama.cpp/pull/23780)

.github/workflows/server-self-hosted.yml

@@ -143,6 +143,24 @@ jobs:
           export LLAMA_ARG_BACKEND_SAMPLING=1
           pytest -v -x -m "not slow"
 
+      - name: Tests (GPUx2)
+        id: server_integration_tests_gpu2
+        if: ${{ !github.event.pull_request }}
+        run: |
+          cd tools/server/tests
+          source venv/bin/activate
+          export GGML_CUDA_DEVICES=2
+          pytest -v -x -m "not slow"
+
+      - name: Tests (GPUx2, backend-sampling)
+        id: server_integration_tests_gpu2_backend_sampling
+        if: ${{ !github.event.pull_request }}
+        run: |
+          cd tools/server/tests
+          source venv/bin/activate
+          export GGML_CUDA_DEVICES=2 LLAMA_ARG_BACKEND_SAMPLING=1
+          pytest -v -x -m "not slow"
+
   server-kleidiai:
     runs-on: ah-ubuntu_22_04-c8g_8x
 

common/arg.cpp

@@ -361,6 +361,14 @@ common_models_handler common_models_handler_init(const common_params & params, l
                                         params.speculative.types.end(),
                                         COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end();
 
+    const bool spec_type_draft_dflash = std::find(params.speculative.types.begin(),
+                                           params.speculative.types.end(),
+                                           COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) != params.speculative.types.end();
+
+    const bool spec_type_draft_eagle3 = std::find(params.speculative.types.begin(),
+                                           params.speculative.types.end(),
+                                           COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3) != params.speculative.types.end();
+
     // only download mmproj if the current example is using it
     bool use_mmproj = false;
     for (const auto & ex : mmproj_examples) {
@@ -373,6 +381,8 @@ common_models_handler common_models_handler_init(const common_params & params, l
     opts.bearer_token    = params.hf_token;
     opts.offline         = params.offline;
     opts.download_mtp    = spec_type_draft_mtp;
+    opts.download_eagle3 = spec_type_draft_eagle3;
+    opts.download_dflash = spec_type_draft_dflash;
     opts.download_mmproj = use_mmproj && !params.no_mmproj
                         && params.mmproj.path.empty() && params.mmproj.url.empty();
 
@@ -517,8 +527,43 @@ void common_models_handler_apply(common_models_handler & handler, common_params
         }
     };
 
+    // when a sidecar type is requested, the draft repo resolves to its sidecar instead of a full model
+    const bool spec_sidecar_found = !plan_spec.mtp.local_path.empty() ||
+                                    !plan_spec.dflash.local_path.empty() ||
+                                    !plan_spec.eagle3.local_path.empty();
+    if (!plan_spec.mtp.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan_spec.mtp, opts, [&]() {
+            // only use the discovered MTP head when no draft path is set yet
+            if (params.speculative.draft.mparams.path.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.mtp);
+            } else {
+                hf_cache::finalize_file(plan_spec.mtp);
+            }
+        });
+    }
+    if (!plan_spec.dflash.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan_spec.dflash, opts, [&]() {
+            // only use the discovered DFlash sidecar when no draft path is set yet
+            if (params.speculative.draft.mparams.path.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.dflash);
+            } else {
+                hf_cache::finalize_file(plan_spec.dflash);
+            }
+        });
+    }
+    if (!plan_spec.eagle3.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan_spec.eagle3, opts, [&]() {
+            // only use the discovered Eagle3 sidecar when no draft path is set yet
+            if (params.speculative.draft.mparams.path.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.eagle3);
+            } else {
+                hf_cache::finalize_file(plan_spec.eagle3);
+            }
+        });
+    }
+
     // handle plan_spec (e.g. --spec-draft-hf)
-    if (!plan_spec.model_files.empty() && !had_spec_url) {
+    if (!plan_spec.model_files.empty() && !had_spec_url && !spec_sidecar_found) {
         add_tasks(plan_spec.model_files, plan_spec.primary, params.speculative.draft.mparams);
         had_spec_url = true;
     }
@@ -546,6 +591,26 @@ void common_models_handler_apply(common_models_handler & handler, common_params
             }
         });
     }
+    if (!plan.dflash.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan.dflash, opts, [&]() {
+            // only fall back to the discovered DFlash sidecar when no draft was explicitly provided
+            if (params.speculative.draft.mparams.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.dflash);
+            } else {
+                hf_cache::finalize_file(plan.dflash);
+            }
+        });
+    }
+    if (!plan.eagle3.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan.eagle3, opts, [&]() {
+            // only fall back to the discovered Eagle3 sidecar when no draft was explicitly provided
+            if (params.speculative.draft.mparams.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.eagle3);
+            } else {
+                hf_cache::finalize_file(plan.eagle3);
+            }
+        });
+    }
     if (!plan.preset.local_path.empty()) {
         tasks.emplace_back(plan.preset, opts, [&]() {
             // if HF repo is a preset repo, we simply run server in router mode with the preset.ini file
@@ -2815,6 +2880,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.speculative.types.push_back(COMMON_SPECULATIVE_TYPE_DRAFT_MTP);
         }
     ).set_examples({LLAMA_EXAMPLE_DOWNLOAD}));
+    add_opt(common_arg(
+        {"--dflash"},
+        "also download the DFlash sidecar, if available (default: unused)",
+        [](common_params & params) {
+            params.speculative.types.push_back(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH);
+        }
+    ).set_examples({LLAMA_EXAMPLE_DOWNLOAD}));
+    add_opt(common_arg(
+        {"--eagle3"},
+        "also download the Eagle3 sidecar, if available (default: unused)",
+        [](common_params & params) {
+            params.speculative.types.push_back(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3);
+        }
+    ).set_examples({LLAMA_EXAMPLE_DOWNLOAD}));
     add_opt(common_arg(
         {"--context-file"}, "FNAME",
         "file to load context from (use comma-separated values to specify multiple files)",

common/download.cpp

@@ -620,6 +620,16 @@ static hf_cache::hf_file find_best_mtp(const hf_cache::hf_files & files,
     return find_best_sibling(files, model, "mtp-");
 }
 
+static hf_cache::hf_file find_best_eagle3(const hf_cache::hf_files & files,
+                                          const std::string        & model) {
+    return find_best_sibling(files, model, "eagle3-");
+}
+
+static hf_cache::hf_file find_best_dflash(const hf_cache::hf_files & files,
+                                          const std::string        & model) {
+    return find_best_sibling(files, model, "dflash-");
+}
+
 static bool gguf_filename_is_model(const std::string & filepath) {
     if (!string_ends_with(filepath, ".gguf")) {
         return false;
@@ -632,7 +642,9 @@ static bool gguf_filename_is_model(const std::string & filepath) {
 
     return filename.find("mmproj")  == std::string::npos &&
            filename.find("imatrix") == std::string::npos &&
-           filename.find("mtp-")    == std::string::npos;
+           filename.find("mtp-")    == std::string::npos &&
+           filename.find("eagle3-") == std::string::npos &&
+           filename.find("dflash-") == std::string::npos;
 }
 
 static hf_cache::hf_file find_best_model(const hf_cache::hf_files & files,
@@ -740,6 +752,12 @@ common_download_hf_plan common_download_get_hf_plan(const common_params_model &
     if (opts.download_mtp) {
         plan.mtp = find_best_mtp(all, primary.path);
     }
+    if (opts.download_dflash) {
+        plan.dflash = find_best_dflash(all, primary.path);
+    }
+    if (opts.download_eagle3) {
+        plan.eagle3 = find_best_eagle3(all, primary.path);
+    }
 
     return plan;
 }
@@ -911,8 +929,10 @@ std::vector<common_cached_model_info> common_list_cached_models() {
     for (const auto & f : files) {
         auto split = get_gguf_split_info(f.path);
         if (split.index != 1 || split.tag.empty() ||
-            split.prefix.find("mmproj") != std::string::npos ||
-            split.prefix.find("mtp-")   != std::string::npos) {
+            split.prefix.find("mmproj")  != std::string::npos ||
+            split.prefix.find("mtp-")    != std::string::npos ||
+            split.prefix.find("eagle3-") != std::string::npos ||
+            split.prefix.find("dflash-") != std::string::npos) {
             continue;
         }
         if (seen.insert(f.repo_id + ":" + split.tag).second) {

common/download.h

@@ -55,8 +55,10 @@ struct common_download_opts {
     std::string bearer_token;
     common_header_list headers;
     bool offline = false;
-    bool download_mmproj = false;
-    bool download_mtp = false;
+    bool download_mmproj  = false;
+    bool download_mtp     = false;
+    bool download_eagle3  = false;
+    bool download_dflash  = false;
     common_download_callback * callback = nullptr;
 };
 
@@ -106,6 +108,8 @@ struct common_download_hf_plan {
     hf_cache::hf_files model_files;
     hf_cache::hf_file mmproj;
     hf_cache::hf_file mtp;
+    hf_cache::hf_file eagle3;
+    hf_cache::hf_file dflash;
     hf_cache::hf_file preset; // if set, only this file is downloaded
 };
 common_download_hf_plan common_download_get_hf_plan(const common_params_model & model, const common_download_opts & opts);

conversion/__init__.py

@@ -31,6 +31,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "BertForSequenceClassification": "bert",
     "BertModel": "bert",
     "BitnetForCausalLM": "bitnet",
+    "BitNetForCausalLM": "bitnet",
     "BloomForCausalLM": "bloom",
     "BloomModel": "bloom",
     "CamembertModel": "bert",

conversion/bitnet.py

@@ -8,7 +8,7 @@ if TYPE_CHECKING:
 from .base import ModelBase, TextModel, gguf
 
 
-@ModelBase.register("BitnetForCausalLM")
+@ModelBase.register("BitnetForCausalLM", "BitNetForCausalLM")
 class BitnetModel(TextModel):
     model_arch = gguf.MODEL_ARCH.BITNET
 

conversion/qwen.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import json
+
 from typing import Any, Callable, Iterable, TYPE_CHECKING
 
 import torch
@@ -641,7 +643,19 @@ class DFlashModel(Qwen3Model):
         logger.info(f"DFlash: Using tokenizer from target model: {self.target_model_dir}")
         original_dir = self.dir_model
         self.dir_model = self.target_model_dir
-        super().set_vocab()
+
+        # Reuse the target model's own vocab handler (e.g. Gemma-4 needs its
+        # own tokenizer logic, not the Qwen default).
+        from . import get_model_class
+        with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f:
+            target_arch = json.load(f)["architectures"][0]
+        target_cls = get_model_class(target_arch)
+
+        if target_cls is not type(self):
+            target_cls.set_vocab(self)  # ty: ignore[unresolved-attribute]
+        else:
+            super().set_vocab()
+
         self.dir_model = original_dir
 
         mask_token_id = self.hparams.get("dflash_config", {}).get("mask_token_id")

docs/backend/OPENCL.md

@@ -47,6 +47,7 @@ The llama.cpp OpenCL backend is designed to enable llama.cpp on **Qualcomm Adren
 | Adreno GPU                            | Status  |
 |:-------------------------------------:|:-------:|
 | Adreno 750 (Snapdragon 8 Gen 3)       | Support |
+| Adreno 810 (Snapdragon 7s Gen 3)      | Support |
 | Adreno 830 (Snapdragon 8 Elite)       | Support |
 | Adreno 840 (Snapdragon 8 Elite Gen 5) | Support |
 | Adreno X1-85 (Snapdragon X Elite)     | Support |

ggml/CMakeLists.txt

@@ -4,7 +4,7 @@ project("ggml" C CXX ASM)
 
 ### GGML Version
 set(GGML_VERSION_MAJOR 0)
-set(GGML_VERSION_MINOR 16)
+set(GGML_VERSION_MINOR 17)
 set(GGML_VERSION_PATCH 0)
 set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}")
 

ggml/include/ggml-cpu.h

@@ -100,6 +100,7 @@ extern "C" {
     GGML_BACKEND_API int ggml_cpu_has_sve        (void);
     GGML_BACKEND_API int ggml_cpu_get_sve_cnt    (void);  // sve vector length in bytes
     GGML_BACKEND_API int ggml_cpu_has_sme        (void);
+    GGML_BACKEND_API int ggml_cpu_has_sme2       (void);
     // other
     GGML_BACKEND_API int ggml_cpu_has_riscv_v    (void);
     GGML_BACKEND_API int ggml_cpu_get_rvv_vlen   (void);  // risc-v vector length in bytes

ggml/include/ggml-rpc.h

@@ -8,10 +8,10 @@ extern "C" {
 
 #define RPC_PROTO_MAJOR_VERSION    4
 #define RPC_PROTO_MINOR_VERSION    0
-#define RPC_PROTO_PATCH_VERSION    2
+#define RPC_PROTO_PATCH_VERSION    3
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

ggml/include/ggml.h

@@ -571,6 +571,9 @@ extern "C" {
         GGML_OP_SOLVE_TRI,
         GGML_OP_GATED_DELTA_NET,
         GGML_OP_LIGHTNING_INDEXER,
+        GGML_OP_DSV4_HC_COMB,
+        GGML_OP_DSV4_HC_PRE,
+        GGML_OP_DSV4_HC_POST,
 
         GGML_OP_UNARY,
 
@@ -2598,6 +2601,45 @@ extern "C" {
         struct ggml_tensor  * weights,
         struct ggml_tensor  * mask);
 
+    // DeepSeek V4 hyper-connections (ref. https://arxiv.org/pdf/2512.24880)
+    // In short these operations are replacements for the original residual connection (x = transformer(x) + x)
+    // using a richer representation through streams.
+    //
+    // hc_comb: mixes [(2 + hc)*hc, n_tokens], scale [3], base [(2 + hc)*hc]
+    //          -> [dst_hc, src_hc, n_tokens]
+    // logits[dst, src, t] = mixes[2*hc + dst + hc*src, t]*scale[2]
+    //                         + base[2*hc + dst + hc*src]
+    // Softmax over dst, add eps, normalize over src, then repeat normalization
+    // over dst followed by src for iterations 1 through n_iter - 1.
+    GGML_API struct ggml_tensor * ggml_dsv4_hc_comb(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * mixes,
+            struct ggml_tensor  * scale,
+            struct ggml_tensor  * base,
+            float                 eps,
+            int32_t               n_iter);
+
+    // hc_pre: x [n_embd, hc, n_tokens], weights [hc, n_tokens] -> [n_embd, n_tokens]
+    //   result[i, t] = sum_h x[i, h, t]*weights[h, t]
+    //
+    GGML_API struct ggml_tensor * ggml_dsv4_hc_pre(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * x,
+            struct ggml_tensor  * weights);
+
+    // hc_post: x [n_embd, n_tokens], residual [n_embd, hc, n_tokens],
+    //          post [hc, n_tokens], comb [dst_hc, src_hc, n_tokens]
+    //          -> [n_embd, hc, n_tokens]
+    //   result[i, dst, t] = x[i, t]*post[dst, t]
+    //                       + sum_src residual[i, src, t]*comb[dst, src, t]
+    //
+    GGML_API struct ggml_tensor * ggml_dsv4_hc_post(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * x,
+            struct ggml_tensor  * residual,
+            struct ggml_tensor  * post,
+            struct ggml_tensor  * comb);
+
     // custom operators
 
     typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata);

ggml/src/ggml-backend-meta.cpp

@@ -984,6 +984,11 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             case GGML_OP_GATED_DELTA_NET: {
                 split_state = handle_gated_delta_net(src_ss);
             } break;
+            case GGML_OP_DSV4_HC_COMB:
+            case GGML_OP_DSV4_HC_PRE:
+            case GGML_OP_DSV4_HC_POST: {
+                split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+            } break;
             case GGML_OP_UNARY: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ false);
             } break;

ggml/src/ggml-blas/ggml-blas.cpp

@@ -1,3 +1,4 @@
+#include "ggml.h"
 #include "ggml-impl.h"
 #include "ggml-blas.h"
 #include "ggml-backend-impl.h"
@@ -415,6 +416,12 @@ static bool ggml_backend_blas_device_supports_op(ggml_backend_dev_t dev, const s
             // TODO: find the optimal value
             const int64_t min_batch = 32;
 
+            // default back to CPU fast path
+            // see: https://github.com/ggml-org/llama.cpp/issues/25565
+            if (ggml_get_op_params_i32(op, 1) == GGML_HINT_SRC0_IS_HADAMARD) {
+                return false;
+            }
+
             return ggml_is_contiguous(src0) &&
                    ggml_is_contiguous(src1) &&
                    src1->type == GGML_TYPE_F32 &&

ggml/src/ggml-cpu/CMakeLists.txt

@@ -678,7 +678,18 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
         endif()
 
         if (NOT SME_ENABLED MATCHES -1)
-            list(APPEND GGML_KLEIDIAI_SOURCES
+            list(APPEND GGML_KLEIDIAI_SME_SOURCES
+                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa.c
+                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa_asm.S
+                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot.c
+                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot_asm.S
+                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.c
+                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa_asm.S)
+            set_source_files_properties(${GGML_KLEIDIAI_SME_SOURCES}
+                PROPERTIES COMPILE_OPTIONS "-fno-tree-vectorize;${ARCH_FLAGS_TEMP}+sve+sve2+sme")
+            list(APPEND GGML_CPU_SOURCES ${GGML_KLEIDIAI_SME_SOURCES})
+
+            list(APPEND GGML_KLEIDIAI_SME2_SOURCES
                 ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot.c
                 ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa.c
                 ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa_asm.S
@@ -698,7 +709,10 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
                 ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.c
                 ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme_asm.S
                 ${KLEIDIAI_SRC}/kai/kai_common_sme_asm.S)
-            set(PRIVATE_ARCH_FLAGS "-fno-tree-vectorize;${PRIVATE_ARCH_FLAGS}+sve+sve2+sme2+fp16")
+            set_source_files_properties(${GGML_KLEIDIAI_SME2_SOURCES}
+                PROPERTIES COMPILE_OPTIONS "-fno-tree-vectorize;${ARCH_FLAGS_TEMP}+sve+sve2+sme2+fp16")
+            list(APPEND GGML_CPU_SOURCES ${GGML_KLEIDIAI_SME2_SOURCES})
+            set(PRIVATE_ARCH_FLAGS "-fno-tree-vectorize;${PRIVATE_ARCH_FLAGS}")
         endif()
 
         if (NOT SVE_ENABLED MATCHES -1)

ggml/src/ggml-cpu/arch/arm/cpu-feats.cpp

@@ -28,6 +28,7 @@ struct aarch64_features {
     bool has_sve2        = false;
     bool has_i8mm        = false;
     bool has_sme         = false;
+    bool has_sme2        = false;
 
     aarch64_features() {
 #if defined(__linux__)
@@ -56,6 +57,10 @@ struct aarch64_features {
             has_sme = static_cast<bool>(oldp);
         }
 
+        if (sysctlbyname("hw.optional.arm.FEAT_SME2", &oldp, &size, NULL, 0) == 0) {
+            has_sme2 = static_cast<bool>(oldp);
+        }
+
         // Apple apparently does not implement SVE yet
 #endif
     }

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

@@ -2064,6 +2064,18 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_lightning_indexer(params, tensor);
             } break;
+        case GGML_OP_DSV4_HC_COMB:
+            {
+                ggml_compute_forward_dsv4_hc_comb(params, tensor);
+            } break;
+        case GGML_OP_DSV4_HC_PRE:
+            {
+                ggml_compute_forward_dsv4_hc_pre(params, tensor);
+            } break;
+        case GGML_OP_DSV4_HC_POST:
+            {
+                ggml_compute_forward_dsv4_hc_post(params, tensor);
+            } break;
         case GGML_OP_MAP_CUSTOM1:
             {
                 ggml_compute_forward_map_custom1(params, tensor);
@@ -2244,6 +2256,9 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_COUNT_EQUAL:
         case GGML_OP_SOLVE_TRI:
         case GGML_OP_GATED_DELTA_NET:
+        case GGML_OP_DSV4_HC_COMB:
+        case GGML_OP_DSV4_HC_PRE:
+        case GGML_OP_DSV4_HC_POST:
             {
                 n_tasks = n_threads;
             } break;
@@ -3792,6 +3807,14 @@ int ggml_cpu_has_sme(void) {
 #endif
 }
 
+int ggml_cpu_has_sme2(void) {
+#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_SME2)
+    return 1;
+#else
+    return 0;
+#endif
+}
+
 void ggml_cpu_init(void) {
     // needed to initialize ggml_time
     {

ggml/src/ggml-cpu/ggml-cpu.cpp

@@ -595,6 +595,9 @@ static ggml_backend_feature * ggml_backend_cpu_get_features(ggml_backend_reg_t r
         if (ggml_cpu_has_sme()) {
             features.push_back({ "SME", "1" });
         }
+        if (ggml_cpu_has_sme2()) {
+            features.push_back({ "SME2", "1" });
+        }
         if (ggml_cpu_has_riscv_v()) {
             features.push_back({ "RISCV_V", "1" });
         }

ggml/src/ggml-cpu/kleidiai/kernels.cpp

@@ -13,6 +13,8 @@
 #include "kai_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa.h"
 #include "kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme2_mopa.h"
 #include "kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme2_dot.h"
+#include "kai_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa.h"
+#include "kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot.h"
 #include "kai_matmul_clamp_f32_qai8dxp1x8_qsi8cxp4x8_1x4_neon_dotprod.h"
 #include "kai_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4x4_1x4_neon_dotprod.h"
 #include "kai_matmul_clamp_f32_qai8dxp4x4_qsi8cxp4x4_16x4_neon_dotprod.h"
@@ -21,6 +23,7 @@
 #include "kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod.h"
 #include "kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.h"
 #include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.h"
+#include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.h"
 
 #include "kai_lhs_pack_bf16p2vlx2_f32_sme.h"
 #include "kai_lhs_pack_f32p2vlx1_f32_sme.h"
@@ -359,7 +362,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
             /* .packed_stride_ex      = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon>,
             /* .pack_func_ex          = */ &rhs_pack_fn12<kai_run_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon>,
         },
-        /* .required_cpu       = */ CPU_FEATURE_SME,
+        /* .required_cpu       = */ CPU_FEATURE_SME2,
         /* .lhs_type           = */ GGML_TYPE_F32,
         /* .rhs_type           = */ GGML_TYPE_Q4_0,
         /* .op_type            = */ GGML_TYPE_F32,
@@ -412,7 +415,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels[] = {
             /* .packed_stride_ex      = */ &rhs_stride_fn1<kai_get_rhs_packed_stride_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme>,
             /* .pack_func_ex          = */ &rhs_pack_fn13<kai_run_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme>,
         },
-        /* .required_cpu       = */ CPU_FEATURE_SME,
+        /* .required_cpu       = */ CPU_FEATURE_SME2,
         /* .lhs_type           = */ GGML_TYPE_F32,
         /* .rhs_type           = */ GGML_TYPE_F16,
         /* .op_type            = */ GGML_TYPE_F32,
@@ -749,6 +752,59 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
             /* .packed_stride_ex      = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
             /* .pack_func_ex          = */ &rhs_pack_scale_fn12<kai_run_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
         },
+        /* .required_cpu       = */ CPU_FEATURE_SME2,
+        /* .lhs_type           = */ GGML_TYPE_F32,
+        /* .rhs_type           = */ GGML_TYPE_Q8_0,
+        /* .op_type            = */ GGML_TYPE_F32,
+    },
+    {
+        /* SME GEMM (pure SME, no SME2 required) */
+        {
+            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+            /* .get_mr                = */ kai_get_mr_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+            /* .get_nr                = */ kai_get_nr_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+            /* .get_kr                = */ kai_get_kr_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+            /* .get_sr                = */ kai_get_sr_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+            /* .get_dst_offset        = */ kai_get_dst_offset_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+            /* .get_dst_size          = */ kai_get_dst_size_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa,
+            /* .get_lhs_offset_ex     = */ &kernel_offs_fn2<kai_get_lhs_packed_offset_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa>,
+            /* .get_rhs_packed_offset_ex = */ &kernel_offs_fn2<kai_get_rhs_packed_offset_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa>,
+            /* .run_kernel_ex         = */ &kernel_run_float_fn10<kai_run_matmul_clamp_f32_qai8dxp1vlx4_qsi8cxp4vlx4_1vlx4vl_sme_mopa>,
+        },
+        /* .gemm_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
+            /* .pack_func_ex          = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
+        },
+        /* SME GEMV (pure SME, no SME2 required) */
+        {
+            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+            /* .get_mr                = */ kai_get_mr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+            /* .get_nr                = */ kai_get_nr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+            /* .get_kr                = */ kai_get_kr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+            /* .get_sr                = */ kai_get_sr_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+            /* .get_dst_offset        = */ kai_get_dst_offset_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+            /* .get_dst_size          = */ kai_get_dst_size_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot,
+            /* .get_lhs_offset_ex     = */ &kernel_offs_fn2<kai_get_lhs_packed_offset_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot>,
+            /* .get_rhs_packed_offset_ex = */ &kernel_offs_fn2<kai_get_rhs_packed_offset_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot>,
+            /* .run_kernel_ex         = */ &kernel_run_float_fn10<kai_run_matmul_clamp_f32_qai8dxp1x4_qsi8cxp4vlx4_1x4vl_sme_dot>,
+        },
+        /* .gemv_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32>,
+            /* .pack_func_ex          = */ &lhs_pack_float_fn9_no_bl<kai_run_lhs_quant_pack_qai8dxp_f32>,
+        },
+        /* .rhs_info = */ {
+            /* .packed_stride         = */ kai_get_rhs_packed_stride_rhs_pack_nxk_qsi8cxp_qsi8cx_neon,
+            /* .to_float              = */ dequantize_row_qsi8cxp,
+            /* .packed_size_ex        = */ &rhs_ps_fn5<kai_get_rhs_packed_size_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+            /* .packed_stride_ex      = */ &rhs_stride_fn4<kai_get_rhs_packed_stride_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+            /* .pack_func_ex          = */ &rhs_pack_scale_fn12<kai_run_rhs_pack_nxk_qsi8cxp_qsi8cx_neon>,
+        },
         /* .required_cpu       = */ CPU_FEATURE_SME,
         /* .lhs_type           = */ GGML_TYPE_F32,
         /* .rhs_type           = */ GGML_TYPE_Q8_0,
@@ -871,7 +927,7 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
 static ggml_kleidiai_kernels ggml_kleidiai_kernels_f32[] = {
 #if defined(__ARM_FEATURE_SME)
     {
-        /* SME GEMM */
+        /* SME2 GEMM */
         {
             /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
             /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
@@ -918,6 +974,59 @@ static ggml_kleidiai_kernels ggml_kleidiai_kernels_f32[] = {
             /* .packed_stride_ex      = */ &rhs_stride_fn1<kai_get_rhs_packed_stride_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
             /* .pack_func_ex          = */ &rhs_pack_fn13<kai_run_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
         },
+        /* .required_cpu       = */ CPU_FEATURE_SME2,
+        /* .lhs_type           = */ GGML_TYPE_F32,
+        /* .rhs_type           = */ GGML_TYPE_F32,
+        /* .op_type            = */ GGML_TYPE_F32,
+    },
+    {
+        /* SME GEMM */
+        {
+            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_mr                = */ kai_get_mr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_nr                = */ kai_get_nr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_kr                = */ kai_get_kr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_sr                = */ kai_get_sr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_dst_offset        = */ kai_get_dst_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_dst_size          = */ kai_get_dst_size_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_lhs_offset_ex     = */ &kernel_offs_fn2<kai_get_lhs_packed_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa>,
+            /* .get_rhs_packed_offset_ex = */ &kernel_offs_fn2<kai_get_rhs_packed_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa>,
+            /* .run_kernel_ex         = */ &kernel_run_fn10<kai_run_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa>,
+        },
+        /* .gemm_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_pack_f32p2vlx1_f32_sme,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_pack_f32p2vlx1_f32_sme>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_pack_f32p2vlx1_f32_sme>,
+            /* .pack_func_ex          = */ &lhs_pack_void_fn9<kai_run_lhs_pack_f32p2vlx1_f32_sme>,
+        },
+        /* SME GEMV */
+        {
+            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_mr                = */ kai_get_mr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_nr                = */ kai_get_nr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_kr                = */ kai_get_kr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_sr                = */ kai_get_sr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_dst_offset        = */ kai_get_dst_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_dst_size          = */ kai_get_dst_size_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa,
+            /* .get_lhs_offset_ex     = */ nullptr,
+            /* .get_rhs_packed_offset_ex = */ nullptr,
+            /* .run_kernel_ex         = */ nullptr,
+        },
+        /* .gemv_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_pack_f32p2vlx1_f32_sme,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_pack_f32p2vlx1_f32_sme>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_pack_f32p2vlx1_f32_sme>,
+            /* .pack_func_ex          = */ &lhs_pack_void_fn9<kai_run_lhs_pack_f32p2vlx1_f32_sme>,
+        },
+        /* .rhs_info = */ {
+            /* .packed_stride         = */ nullptr,
+            /* .to_float              = */ nullptr,
+            /* .packed_size_ex        = */ &rhs_ps_fn2<kai_get_rhs_packed_size_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
+            /* .packed_stride_ex      = */ &rhs_stride_fn1<kai_get_rhs_packed_stride_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
+            /* .pack_func_ex          = */ &rhs_pack_fn13<kai_run_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
+        },
         /* .required_cpu       = */ CPU_FEATURE_SME,
         /* .lhs_type           = */ GGML_TYPE_F32,
         /* .rhs_type           = */ GGML_TYPE_F32,

ggml/src/ggml-cpu/kleidiai/kernels.h

@@ -11,7 +11,8 @@ enum cpu_feature {
     CPU_FEATURE_DOTPROD = 1,
     CPU_FEATURE_I8MM    = 2,
     CPU_FEATURE_SVE     = 4,
-    CPU_FEATURE_SME     = 8
+    CPU_FEATURE_SME     = 8,
+    CPU_FEATURE_SME2    = 16
 };
 
 inline cpu_feature& operator|=(cpu_feature& lhs, cpu_feature rhs) {

ggml/src/ggml-cpu/kleidiai/kleidiai.cpp

@@ -26,6 +26,9 @@
 #include <sys/types.h>
 #include <sys/stat.h>
 #include <unistd.h>
+#ifndef HWCAP2_SME2
+#define HWCAP2_SME2 (1UL << 37)
+#endif
 #elif defined(__APPLE__)
 #include <string_view>
 #include <sys/sysctl.h>
@@ -66,9 +69,15 @@ struct ggml_kleidiai_context {
     int chunk_multiplier;
 } static ctx = { CPU_FEATURE_NONE, nullptr, nullptr, nullptr, 0, -1, 4 };
 
+static inline bool is_sme_family(cpu_feature f) {
+    return (f & (CPU_FEATURE_SME | CPU_FEATURE_SME2)) != CPU_FEATURE_NONE;
+}
+
 static const char* cpu_feature_to_string(cpu_feature f) {
     if (f == CPU_FEATURE_NONE) {
         return "NONE";
+    } else if ((f & CPU_FEATURE_SME2) == CPU_FEATURE_SME2) {
+        return "SME2";
     } else if ((f & CPU_FEATURE_SME) == CPU_FEATURE_SME) {
         return "SME";
     } else if ((f & CPU_FEATURE_SVE) == CPU_FEATURE_SVE) {
@@ -251,6 +260,18 @@ static void init_kleidiai_context(void) {
 
         if (sme_cores > 0) {
             ctx.features |= CPU_FEATURE_SME;
+#if defined(__aarch64__) && defined(__linux__)
+            // ARM guarantees SME2 implies SME, so only check SME2 when SME is enabled.
+            if (getauxval(AT_HWCAP2) & HWCAP2_SME2) {
+                ctx.features |= CPU_FEATURE_SME2;
+            }
+#elif defined(__aarch64__) && defined(__APPLE__)
+            int feat_sme2 = 0;
+            size_t size = sizeof(feat_sme2);
+            if (sysctlbyname("hw.optional.arm.FEAT_SME2", &feat_sme2, &size, NULL, 0) == 0 && feat_sme2) {
+                ctx.features |= CPU_FEATURE_SME2;
+            }
+#endif
         }
 
         // Kernel selection
@@ -279,10 +300,13 @@ static void init_kleidiai_context(void) {
         ctx.sme_thread_cap = (ctx.features & CPU_FEATURE_SME) ? sme_cores : 0;
 
         if (ctx.features & CPU_FEATURE_SME) {
+            const bool has_sme2 = (ctx.features & CPU_FEATURE_SME2) != CPU_FEATURE_NONE;
             if (sme_env_set && sme_env_ok && sme_cores > 0) {
-                GGML_LOG_INFO("kleidiai: SME enabled (GGML_KLEIDIAI_SME=%d override)\n", sme_cores);
+                GGML_LOG_INFO("kleidiai: SME%s enabled (GGML_KLEIDIAI_SME=%d override)\n",
+                    has_sme2 ? "2" : "", sme_cores);
             } else {
-                GGML_LOG_INFO("kleidiai: SME enabled (runtime-detected SME cores=%d)\n", sme_cores);
+                GGML_LOG_INFO("kleidiai: SME%s enabled (runtime-detected SME cores=%d)\n",
+                    has_sme2 ? "2" : "", sme_cores);
             }
         } else {
             GGML_LOG_INFO("kleidiai: SME disabled\n");
@@ -442,8 +466,8 @@ static int kleidiai_collect_kernel_chain_common(
         return count;
     }
 
-    if ((primary->required_cpu & CPU_FEATURE_SME) == CPU_FEATURE_SME) {
-        const cpu_feature fallback_mask = static_cast<cpu_feature>(features & ~CPU_FEATURE_SME);
+    if (is_sme_family(primary->required_cpu)) {
+        const cpu_feature fallback_mask = static_cast<cpu_feature>(features & ~CPU_FEATURE_SME & ~CPU_FEATURE_SME2);
         if (fallback_mask != CPU_FEATURE_NONE) {
             ggml_kleidiai_kernels * fallback = select_fallback(fallback_mask);
             if (fallback && fallback != primary &&
@@ -1054,14 +1078,14 @@ class tensor_traits : public ggml::cpu::tensor_traits {
 
         int sme_slot = -1;
         for (int i = 0; i < runtime_count; ++i) {
-            if ((runtime[i].kernels->required_cpu & CPU_FEATURE_SME) == CPU_FEATURE_SME) {
+            if (is_sme_family(runtime[i].kernels->required_cpu)) {
                 sme_slot = i;
                 break;
             }
         }
         int non_sme_slot = -1;
         for (int i = 0; i < runtime_count; ++i) {
-            if ((runtime[i].kernels->required_cpu & CPU_FEATURE_SME) != CPU_FEATURE_SME) {
+            if (!is_sme_family(runtime[i].kernels->required_cpu)) {
                 non_sme_slot = i;
                 break;
             }
@@ -1099,7 +1123,7 @@ class tensor_traits : public ggml::cpu::tensor_traits {
             // Recompute SME slot based on the collapsed runtime[0]
             sme_slot = -1;
             if (runtime_count > 0 &&
-                (runtime[0].kernels->required_cpu & CPU_FEATURE_SME) == CPU_FEATURE_SME) {
+                is_sme_family(runtime[0].kernels->required_cpu)) {
                 sme_slot = 0;
             }
         }
@@ -1695,6 +1719,7 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
 
             return true;
         }
+
         return false;
     }
 
@@ -1703,6 +1728,20 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
             if (op->src[0]->buffer && op->src[0]->buffer->buft == ggml_backend_cpu_kleidiai_buffer_type()) {
                 return (ggml::cpu::tensor_traits *) op->src[0]->extra;
             } else {
+                // KleidiAI only has kernels for Q4_0 and Q8_0. For a quantized weight of any
+                // other type (K-quants, IQ) it declines the op and returns nullptr below, so
+                // KleidiAI does not accelerate it. Another CPU backend may still take the op,
+                // and this can run during graph planning, so the message says what KleidiAI
+                // did rather than what ends up executing. Warn once per process.
+                if (ggml_is_quantized(op->src[0]->type) &&
+                    op->src[0]->type != GGML_TYPE_Q4_0 && op->src[0]->type != GGML_TYPE_Q8_0) {
+                    static std::atomic<bool> warned(false);
+                    if (!warned.exchange(true)) {
+                        GGML_LOG_WARN("kleidiai: no kernel for tensor type %s, not accelerated by KleidiAI "
+                                      "(kernels available for Q4_0 and Q8_0)\n",
+                                      ggml_type_name(op->src[0]->type));
+                    }
+                }
                 if (op->src[0]->type != GGML_TYPE_F16) {
                     return nullptr;
                 }

ggml/src/ggml-cpu/ops.cpp

@@ -10944,6 +10944,291 @@ void ggml_compute_forward_gated_delta_net(
     }
 }
 
+
+// ggml_compute_forward_dsv4_hc_comb
+
+static void ggml_dsv4_hc_comb_norm_cols(float * comb, float eps) {
+    constexpr int64_t hc = 4;
+
+    for (int64_t idst = 0; idst < hc; ++idst) {
+        float sum = eps;
+        for (int64_t isrc = 0; isrc < hc; ++isrc) {
+            sum += comb[idst + hc*isrc];
+        }
+
+        const float inv_sum = 1.0f / sum;
+        for (int64_t isrc = 0; isrc < hc; ++isrc) {
+            comb[idst + hc*isrc] *= inv_sum;
+        }
+    }
+}
+
+static void ggml_dsv4_hc_comb_norm_rows(float * comb, float eps) {
+    constexpr int64_t hc = 4;
+
+    for (int64_t isrc = 0; isrc < hc; ++isrc) {
+        float sum = eps;
+        for (int64_t idst = 0; idst < hc; ++idst) {
+            sum += comb[idst + hc*isrc];
+        }
+
+        const float inv_sum = 1.0f / sum;
+        for (int64_t idst = 0; idst < hc; ++idst) {
+            comb[idst + hc*isrc] *= inv_sum;
+        }
+    }
+}
+
+static void ggml_compute_forward_dsv4_hc_comb_f32(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * mixes = dst->src[0];
+    const ggml_tensor * scale = dst->src[1];
+    const ggml_tensor * base  = dst->src[2];
+
+    GGML_ASSERT(mixes->type == GGML_TYPE_F32);
+    GGML_ASSERT(scale->type == GGML_TYPE_F32);
+    GGML_ASSERT(base->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+    constexpr int64_t hc = 4;
+    constexpr int64_t comb_offset = 2*hc;
+    constexpr int64_t hc_mix_dim = (2 + hc)*hc;
+
+    const int64_t n_tokens = mixes->ne[1];
+
+    GGML_ASSERT(mixes->ne[0] == hc_mix_dim);
+    GGML_ASSERT(dst->ne[0] == hc);
+    GGML_ASSERT(dst->ne[1] == hc);
+    GGML_ASSERT(dst->ne[2] == n_tokens);
+    GGML_ASSERT(scale->ne[0] >= 3);
+    GGML_ASSERT(base->ne[0] == hc_mix_dim);
+
+    GGML_TENSOR_LOCALS(size_t, nbm, mixes, nb);
+    GGML_TENSOR_LOCALS(size_t, nbs, scale, nb);
+    GGML_TENSOR_LOCALS(size_t, nbb, base,  nb);
+    GGML_TENSOR_LOCALS(size_t, nbd, dst,   nb);
+
+    const float eps = ggml_get_op_params_f32(dst, 0);
+    const int32_t n_iter = ggml_get_op_params_i32(dst, 1);
+    GGML_ASSERT(n_iter > 0);
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    const int64_t dr  = (n_tokens + nth - 1) / nth;
+    const int64_t it0 = dr * ith;
+    const int64_t it1 = MIN(it0 + dr, n_tokens);
+
+    const float scale_comb = *(const float *) ((const char *) scale->data + 2*nbs0);
+
+    for (int64_t it = it0; it < it1; ++it) {
+        float comb[hc*hc];
+
+        for (int64_t isrc = 0; isrc < hc; ++isrc) {
+            float max = -INFINITY;
+            for (int64_t idst = 0; idst < hc; ++idst) {
+                const int64_t idx = idst + hc*isrc;
+                const float xv = *(const float *) ((const char *) mixes->data + (comb_offset + idx)*nbm0 + it*nbm1);
+                const float bv = *(const float *) ((const char *) base->data  + (comb_offset + idx)*nbb0);
+                const float v = xv * scale_comb + bv;
+                comb[idx] = v;
+                max = MAX(max, v);
+            }
+
+            float sum = 0.0f;
+            for (int64_t idst = 0; idst < hc; ++idst) {
+                const int64_t idx = idst + hc*isrc;
+                const float v = expf(comb[idx] - max);
+                comb[idx] = v;
+                sum += v;
+            }
+
+            const float inv_sum = 1.0f / sum;
+            for (int64_t idst = 0; idst < hc; ++idst) {
+                const int64_t idx = idst + hc*isrc;
+                comb[idx] = comb[idx] * inv_sum + eps;
+            }
+        }
+
+        ggml_dsv4_hc_comb_norm_cols(comb, eps);
+        for (int32_t i = 1; i < n_iter; ++i) {
+            ggml_dsv4_hc_comb_norm_rows(comb, eps);
+            ggml_dsv4_hc_comb_norm_cols(comb, eps);
+        }
+
+        for (int64_t isrc = 0; isrc < hc; ++isrc) {
+            for (int64_t idst = 0; idst < hc; ++idst) {
+                const int64_t idx = idst + hc*isrc;
+                *(float *) ((char *) dst->data + idst*nbd0 + isrc*nbd1 + it*nbd2) = comb[idx];
+            }
+        }
+    }
+}
+
+void ggml_compute_forward_dsv4_hc_comb(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+
+    switch (src0->type) {
+        case GGML_TYPE_F32:
+            {
+                ggml_compute_forward_dsv4_hc_comb_f32(params, dst);
+            } break;
+        default:
+            {
+                GGML_ABORT("fatal error");
+            }
+    }
+}
+
+// ggml_compute_forward_dsv4_hc_pre
+
+static void ggml_compute_forward_dsv4_hc_pre_f32(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * x       = dst->src[0];
+    const ggml_tensor * weights = dst->src[1];
+
+    GGML_ASSERT(x->type == GGML_TYPE_F32);
+    GGML_ASSERT(weights->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+    const int64_t n_embd   = x->ne[0];
+    const int64_t hc       = x->ne[1];
+    const int64_t n_tokens = x->ne[2];
+
+    GGML_ASSERT(dst->ne[0] == n_embd);
+    GGML_ASSERT(dst->ne[1] == n_tokens);
+    GGML_ASSERT(weights->ne[0] == hc);
+    GGML_ASSERT(weights->ne[1] == n_tokens);
+
+    GGML_TENSOR_LOCALS(size_t, nbx, x,       nb);
+    GGML_TENSOR_LOCALS(size_t, nbw, weights, nb);
+    GGML_TENSOR_LOCALS(size_t, nbd, dst,     nb);
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    const int64_t nr  = n_embd * n_tokens;
+    const int64_t dr  = (nr + nth - 1) / nth;
+    const int64_t ir0 = dr * ith;
+    const int64_t ir1 = MIN(ir0 + dr, nr);
+
+    for (int64_t ir = ir0; ir < ir1; ++ir) {
+        const int64_t i0 = ir % n_embd;
+        const int64_t it = ir / n_embd;
+
+        float sum = 0.0f;
+        for (int64_t ih = 0; ih < hc; ++ih) {
+            const float xv = *(const float *) ((const char *) x->data       + i0*nbx0 + ih*nbx1 + it*nbx2);
+            const float wv = *(const float *) ((const char *) weights->data + ih*nbw0 + it*nbw1);
+            sum += xv * wv;
+        }
+
+        *(float *) ((char *) dst->data + i0*nbd0 + it*nbd1) = sum;
+    }
+}
+
+void ggml_compute_forward_dsv4_hc_pre(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+
+    switch (src0->type) {
+        case GGML_TYPE_F32:
+            {
+                ggml_compute_forward_dsv4_hc_pre_f32(params, dst);
+            } break;
+        default:
+            {
+                GGML_ABORT("fatal error");
+            }
+    }
+}
+
+// ggml_compute_forward_dsv4_hc_post
+
+static void ggml_compute_forward_dsv4_hc_post_f32(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * x        = dst->src[0];
+    const ggml_tensor * residual = dst->src[1];
+    const ggml_tensor * post     = dst->src[2];
+    const ggml_tensor * comb     = dst->src[3];
+
+    GGML_ASSERT(x->type == GGML_TYPE_F32);
+    GGML_ASSERT(residual->type == GGML_TYPE_F32);
+    GGML_ASSERT(post->type == GGML_TYPE_F32);
+    GGML_ASSERT(comb->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+    const int64_t n_embd   = x->ne[0];
+    const int64_t n_tokens = x->ne[1];
+    const int64_t hc       = residual->ne[1];
+
+    GGML_ASSERT(dst->ne[0] == n_embd);
+    GGML_ASSERT(dst->ne[1] == hc);
+    GGML_ASSERT(dst->ne[2] == n_tokens);
+    GGML_ASSERT(residual->ne[0] == n_embd);
+    GGML_ASSERT(residual->ne[2] == n_tokens);
+    GGML_ASSERT(post->ne[0] == hc);
+    GGML_ASSERT(post->ne[1] == n_tokens);
+    GGML_ASSERT(comb->ne[0] == hc);
+    GGML_ASSERT(comb->ne[1] == hc);
+    GGML_ASSERT(comb->ne[2] == n_tokens);
+
+    GGML_TENSOR_LOCALS(size_t, nbx, x,        nb);
+    GGML_TENSOR_LOCALS(size_t, nbr, residual, nb);
+    GGML_TENSOR_LOCALS(size_t, nbp, post,     nb);
+    GGML_TENSOR_LOCALS(size_t, nbc, comb,     nb);
+    GGML_TENSOR_LOCALS(size_t, nbd, dst,      nb);
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    const int64_t nr  = n_embd * hc * n_tokens;
+    const int64_t dr  = (nr + nth - 1) / nth;
+    const int64_t ir0 = dr * ith;
+    const int64_t ir1 = MIN(ir0 + dr, nr);
+
+    for (int64_t ir = ir0; ir < ir1; ++ir) {
+        const int64_t i0     = ir % n_embd;
+        const int64_t idst   = (ir / n_embd) % hc;
+        const int64_t it     = ir / (n_embd * hc);
+
+        const float xv = *(const float *) ((const char *) x->data    + i0*nbx0 + it*nbx1);
+        const float pv = *(const float *) ((const char *) post->data + idst*nbp0 + it*nbp1);
+
+        float sum = xv * pv;
+        for (int64_t isrc = 0; isrc < hc; ++isrc) {
+            const float rv = *(const float *) ((const char *) residual->data + i0*nbr0 + isrc*nbr1 + it*nbr2);
+            const float cv = *(const float *) ((const char *) comb->data     + idst*nbc0 + isrc*nbc1 + it*nbc2);
+            sum += rv * cv;
+        }
+
+        *(float *) ((char *) dst->data + i0*nbd0 + idst*nbd1 + it*nbd2) = sum;
+    }
+}
+
+void ggml_compute_forward_dsv4_hc_post(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+
+    switch (src0->type) {
+        case GGML_TYPE_F32:
+            {
+                ggml_compute_forward_dsv4_hc_post_f32(params, dst);
+            } break;
+        default:
+            {
+                GGML_ABORT("fatal error");
+            }
+    }
+}
+
 // ggml_compute_forward_rwkv_wkv7
 
 static void ggml_compute_forward_rwkv_wkv7_f32(

ggml/src/ggml-cpu/ops.h

@@ -106,6 +106,9 @@ void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, s
 void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_dsv4_hc_post(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst);

ggml/src/ggml-cuda/common.cuh

@@ -1115,7 +1115,8 @@ struct ggml_cuda_type_traits<GGML_TYPE_IQ3_S> {
 //////////////////////
 
 struct ggml_cuda_device_info {
-    int device_count;
+    int device_count;           // number of (possibly virtual) devices exposed to the rest of ggml
+    int physical_device_count;  // number of physical CUDA devices actually present
 
     struct cuda_device_info {
         int     cc;                             // compute capability
@@ -1128,6 +1129,9 @@ struct ggml_cuda_device_info {
         size_t  total_vram;
         int     warp_size;                      // Number of threads in a dispatch
         bool    supports_cooperative_launch;    // whether cooperative launch is supported
+        int     physical_device;                // backing physical CUDA device for this (virtual) device
+        int     physical_share_count;           // number of (virtual) devices sharing this device's physical GPU
+        int     virtual_index;                  // index of this (virtual) device among those sharing its physical GPU
     };
 
     cuda_device_info devices[GGML_CUDA_MAX_DEVICES] = {};

ggml/src/ggml-cuda/dsv4-hc.cu

@@ -0,0 +1,294 @@
+#include "common.cuh"
+#include "dsv4-hc.cuh"
+
+
+static constexpr int DSV4_HC = 4;
+
+
+static __device__ void dsv4_hc_comb_norm_cols(float * comb, float eps) {
+    for (int idst = 0; idst < DSV4_HC; ++idst) {
+        float sum = eps;
+        for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
+            sum += comb[idst + DSV4_HC*isrc];
+        }
+
+        const float inv_sum = 1.0f / sum;
+        for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
+            comb[idst + DSV4_HC*isrc] *= inv_sum;
+        }
+    }
+}
+
+static __device__ void dsv4_hc_comb_norm_rows(float * comb, float eps) {
+    for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
+        float sum = eps;
+        for (int idst = 0; idst < DSV4_HC; ++idst) {
+            sum += comb[idst + DSV4_HC*isrc];
+        }
+
+        const float inv_sum = 1.0f / sum;
+        for (int idst = 0; idst < DSV4_HC; ++idst) {
+            comb[idst + DSV4_HC*isrc] *= inv_sum;
+        }
+    }
+}
+
+static __global__ void dsv4_hc_comb_f32(
+        const float * mixes,
+        const float * scale,
+        const float * base,
+        float * dst,
+        int64_t n_tokens,
+        int64_t sm0,
+        int64_t sm1,
+        int64_t ss0,
+        int64_t sb0,
+        int64_t sd0,
+        int64_t sd1,
+        int64_t sd2,
+        float eps,
+        int32_t n_iter) {
+    constexpr int comb_offset = 2*DSV4_HC;
+
+    ggml_cuda_pdl_lc();
+    const int64_t it = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (it >= n_tokens) {
+        return;
+    }
+
+    ggml_cuda_pdl_sync();
+
+    const float scale_comb = scale[2*ss0];
+    float comb[DSV4_HC*DSV4_HC];
+
+    for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
+        float max = -INFINITY;
+        for (int idst = 0; idst < DSV4_HC; ++idst) {
+            const int idx = idst + DSV4_HC*isrc;
+            const float v = mixes[(comb_offset + idx)*sm0 + it*sm1] * scale_comb + base[(comb_offset + idx)*sb0];
+            comb[idx] = v;
+            max = fmaxf(max, v);
+        }
+
+        float sum = 0.0f;
+        for (int idst = 0; idst < DSV4_HC; ++idst) {
+            const int idx = idst + DSV4_HC*isrc;
+            const float v = expf(comb[idx] - max);
+            comb[idx] = v;
+            sum += v;
+        }
+
+        const float inv_sum = 1.0f / sum;
+        for (int idst = 0; idst < DSV4_HC; ++idst) {
+            const int idx = idst + DSV4_HC*isrc;
+            comb[idx] = comb[idx] * inv_sum + eps;
+        }
+    }
+
+    dsv4_hc_comb_norm_cols(comb, eps);
+    for (int32_t i = 1; i < n_iter; ++i) {
+        dsv4_hc_comb_norm_rows(comb, eps);
+        dsv4_hc_comb_norm_cols(comb, eps);
+    }
+
+    for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
+        for (int idst = 0; idst < DSV4_HC; ++idst) {
+            const int idx = idst + DSV4_HC*isrc;
+            dst[idst*sd0 + isrc*sd1 + it*sd2] = comb[idx];
+        }
+    }
+}
+
+static __global__ void dsv4_hc_pre_f32(
+        const float * x,
+        const float * weights,
+        float * dst,
+        int64_t n_embd,
+        int64_t hc,
+        int64_t n_tokens,
+        int64_t sx0,
+        int64_t sx1,
+        int64_t sx2,
+        int64_t sw0,
+        int64_t sw1,
+        int64_t sd0,
+        int64_t sd1) {
+    ggml_cuda_pdl_lc();
+    const int64_t ir = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+    const int64_t nr = n_embd * n_tokens;
+
+    if (ir >= nr) {
+        return;
+    }
+
+    ggml_cuda_pdl_sync();
+
+    const int64_t i0 = ir % n_embd;
+    const int64_t it = ir / n_embd;
+
+    float sum = x[i0*sx0 + it*sx2] * weights[it*sw1];
+    for (int64_t ih = 1; ih < hc; ++ih) {
+        const float xv = x[i0*sx0 + ih*sx1 + it*sx2];
+        const float wv = weights[ih*sw0 + it*sw1];
+        sum += xv * wv;
+    }
+
+    dst[i0*sd0 + it*sd1] = sum;
+}
+
+static __global__ void dsv4_hc_post_f32(
+        const float * x,
+        const float * residual,
+        const float * post,
+        const float * comb,
+        float * dst,
+        int64_t n_embd,
+        int64_t hc,
+        int64_t n_tokens,
+        int64_t sx0,
+        int64_t sx1,
+        int64_t sr0,
+        int64_t sr1,
+        int64_t sr2,
+        int64_t sp0,
+        int64_t sp1,
+        int64_t sc0,
+        int64_t sc1,
+        int64_t sc2,
+        int64_t sd0,
+        int64_t sd1,
+        int64_t sd2) {
+    ggml_cuda_pdl_lc();
+    const int64_t ir = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
+    const int64_t nr = n_embd * hc * n_tokens;
+
+    if (ir >= nr) {
+        return;
+    }
+
+    ggml_cuda_pdl_sync();
+
+    const int64_t i0   = ir % n_embd;
+    const int64_t idst = (ir / n_embd) % hc;
+    const int64_t it   = ir / (n_embd * hc);
+
+    float sum = x[i0*sx0 + it*sx1] * post[idst*sp0 + it*sp1];
+    for (int64_t isrc = 0; isrc < hc; ++isrc) {
+        sum += residual[i0*sr0 + isrc*sr1 + it*sr2] * comb[idst*sc0 + isrc*sc1 + it*sc2];
+    }
+
+    dst[i0*sd0 + idst*sd1 + it*sd2] = sum;
+}
+
+void ggml_cuda_op_dsv4_hc_comb(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * mixes = dst->src[0];
+    const ggml_tensor * scale = dst->src[1];
+    const ggml_tensor * base  = dst->src[2];
+
+    GGML_ASSERT(mixes->type == GGML_TYPE_F32);
+    GGML_ASSERT(scale->type == GGML_TYPE_F32);
+    GGML_ASSERT(base->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+    constexpr int64_t hc_mix_dim = (2 + DSV4_HC)*DSV4_HC;
+
+    GGML_ASSERT(mixes->ne[0] == hc_mix_dim);
+    GGML_ASSERT(dst->ne[0] == DSV4_HC);
+    GGML_ASSERT(dst->ne[1] == DSV4_HC);
+    GGML_ASSERT(dst->ne[2] == mixes->ne[1]);
+    GGML_ASSERT(scale->ne[0] >= 3);
+    GGML_ASSERT(base->ne[0] == hc_mix_dim);
+
+    GGML_TENSOR_LOCALS(size_t, nbm, mixes, nb);
+    GGML_TENSOR_LOCALS(size_t, nbs, scale, nb);
+    GGML_TENSOR_LOCALS(size_t, nbb, base,  nb);
+    GGML_TENSOR_LOCALS(size_t, nbd, dst,   nb);
+
+    const int64_t n_tokens = mixes->ne[1];
+    const float eps = ggml_get_op_params_f32(dst, 0);
+    const int32_t n_iter = ggml_get_op_params_i32(dst, 1);
+
+    const int block_size = 256;
+    const dim3 block_dims(block_size, 1, 1);
+    const dim3 grid_dims((n_tokens + block_size - 1) / block_size, 1, 1);
+    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
+
+    ggml_cuda_kernel_launch(dsv4_hc_comb_f32, launch_params,
+            (const float *) mixes->data, (const float *) scale->data, (const float *) base->data, (float *) dst->data,
+            n_tokens,
+            nbm0 / sizeof(float), nbm1 / sizeof(float),
+            nbs0 / sizeof(float),
+            nbb0 / sizeof(float),
+            nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float),
+            eps, n_iter);
+}
+
+void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * x       = dst->src[0];
+    const ggml_tensor * weights = dst->src[1];
+
+    GGML_ASSERT(x->type == GGML_TYPE_F32);
+    GGML_ASSERT(weights->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+    GGML_TENSOR_LOCALS(size_t, nbx, x,       nb);
+    GGML_TENSOR_LOCALS(size_t, nbw, weights, nb);
+    GGML_TENSOR_LOCALS(size_t, nbd, dst,     nb);
+
+    const int64_t n_embd   = x->ne[0];
+    const int64_t hc       = x->ne[1];
+    const int64_t n_tokens = x->ne[2];
+
+    const int block_size = 256;
+    const int64_t nr = n_embd * n_tokens;
+    const dim3 block_dims(block_size, 1, 1);
+    const dim3 grid_dims((nr + block_size - 1) / block_size, 1, 1);
+    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
+
+    ggml_cuda_kernel_launch(dsv4_hc_pre_f32, launch_params,
+            (const float *) x->data, (const float *) weights->data, (float *) dst->data,
+            n_embd, hc, n_tokens,
+            nbx0 / sizeof(float), nbx1 / sizeof(float), nbx2 / sizeof(float),
+            nbw0 / sizeof(float), nbw1 / sizeof(float),
+            nbd0 / sizeof(float), nbd1 / sizeof(float));
+}
+
+void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * x        = dst->src[0];
+    const ggml_tensor * residual = dst->src[1];
+    const ggml_tensor * post     = dst->src[2];
+    const ggml_tensor * comb     = dst->src[3];
+
+    GGML_ASSERT(x->type == GGML_TYPE_F32);
+    GGML_ASSERT(residual->type == GGML_TYPE_F32);
+    GGML_ASSERT(post->type == GGML_TYPE_F32);
+    GGML_ASSERT(comb->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+    GGML_TENSOR_LOCALS(size_t, nbx, x,        nb);
+    GGML_TENSOR_LOCALS(size_t, nbr, residual, nb);
+    GGML_TENSOR_LOCALS(size_t, nbp, post,     nb);
+    GGML_TENSOR_LOCALS(size_t, nbc, comb,     nb);
+    GGML_TENSOR_LOCALS(size_t, nbd, dst,      nb);
+
+    const int64_t n_embd   = x->ne[0];
+    const int64_t n_tokens = x->ne[1];
+    const int64_t hc       = residual->ne[1];
+
+    const int block_size = 256;
+    const int64_t nr = n_embd * hc * n_tokens;
+    const dim3 block_dims(block_size, 1, 1);
+    const dim3 grid_dims((nr + block_size - 1) / block_size, 1, 1);
+    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
+
+    ggml_cuda_kernel_launch(dsv4_hc_post_f32, launch_params,
+            (const float *) x->data, (const float *) residual->data,
+            (const float *) post->data, (const float *) comb->data, (float *) dst->data,
+            n_embd, hc, n_tokens,
+            nbx0 / sizeof(float), nbx1 / sizeof(float),
+            nbr0 / sizeof(float), nbr1 / sizeof(float), nbr2 / sizeof(float),
+            nbp0 / sizeof(float), nbp1 / sizeof(float),
+            nbc0 / sizeof(float), nbc1 / sizeof(float), nbc2 / sizeof(float),
+            nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float));
+}

ggml/src/ggml-cuda/dsv4-hc.cuh

@@ -0,0 +1,6 @@
+#include "common.cuh"
+#include "ggml.h"
+
+void ggml_cuda_op_dsv4_hc_comb(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

ggml/src/ggml-cuda/getrows.cu

@@ -55,27 +55,51 @@ static __global__ void k_get_rows_float(
     dst_t         * GGML_CUDA_RESTRICT dst  = dst_ptr;
     ggml_cuda_pdl_sync();
     for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+        // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
+        const int i10 = blockIdx.x;
+        const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
+        const int i11 = dm.x;
+        const int i12 = dm.y;
+
+        const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
+
+        dst_t * GGML_CUDA_RESTRICT dst_row = dst + i10*s1 + i11*s2 + i12*s3;
+        const src0_t * GGML_CUDA_RESTRICT src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);
+
         for (int64_t i00 = blockIdx.y*blockDim.x + threadIdx.x; i00 < ne00; i00 += gridDim.y*blockDim.x) {
-            // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
-            const int i10 = blockIdx.x;
-            const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
-            const int i11 = dm.x;
-            const int i12 = dm.y;
-
-            if (i00 >= ne00) {
-                return;
-            }
-
-            const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
-
-            dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
-            const src0_t * src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);
-
             dst_row[i00] = ggml_cuda_cast<dst_t>(src0_row[i00]);
         }
     }
 }
 
+template<typename dst_t>
+static __global__ void k_get_rows_float_vec(
+        const dst_t * src0_ptr, const int32_t * src1_ptr, dst_t * dst_ptr,
+        const int64_t ne00v,
+        const int64_t ne11, const uint3 ne12_fdv,
+        const size_t s1, const size_t s2, const size_t s3,
+        const size_t nb01, const size_t nb02, const size_t nb03,
+        const size_t s10, const size_t s11, const size_t s12) {
+
+    ggml_cuda_pdl_lc();
+    ggml_cuda_pdl_sync();
+    for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+        const int i10 = blockIdx.x;
+        const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
+        const int i11 = dm.x;
+        const int i12 = dm.y;
+
+        const int i01 = src1_ptr[i10*s10 + i11*s11 + i12*s12];
+
+        int4       * GGML_CUDA_RESTRICT dst_row  = (int4 *)      (dst_ptr + i10*s1 + i11*s2 + i12*s3);
+        const int4 * GGML_CUDA_RESTRICT src0_row = (const int4 *)((const char *) src0_ptr + i01*nb01 + i11*nb02 + i12*nb03);
+
+        for (int64_t i = blockIdx.y*blockDim.x + threadIdx.x; i < ne00v; i += gridDim.y*blockDim.x) {
+            dst_row[i] = src0_row[i];
+        }
+    }
+}
+
 template<typename grad_t, typename dst_t>
 static __global__ void k_get_rows_back_float(
         const grad_t * __restrict__ grad, const int32_t * __restrict__ rows, dst_t * __restrict__ dst,
@@ -148,8 +172,6 @@ static void get_rows_cuda_float(
         const size_t nb1, const size_t nb2, const size_t nb3,
         cudaStream_t stream) {
     const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
-    const int block_num_y = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
-    const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
 
     // strides in elements
     // const size_t s0 = nb0 / sizeof(dst_t);
@@ -166,6 +188,34 @@ static void get_rows_cuda_float(
     GGML_ASSERT(ne11 <= std::numeric_limits<uint32_t>::max() / ne12);
     const uint3 ne12_fdv = init_fastdiv_values(ne12);
 
+    if constexpr (std::is_same<src0_t, dst_t>::value) {
+        constexpr int VEC = 16 / sizeof(dst_t);
+        const int64_t ne00v = ne00 / VEC;
+        const int64_t vec_block_num_y = (ne00v + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
+        const bool enough_blocks = vec_block_num_y * ne10 * ne11 * ne12 >= 128;
+        const bool can_vec = VEC > 1 && enough_blocks &&
+            (ne00 % VEC == 0) &&
+            (nb01 % 16 == 0) && (nb02 % 16 == 0) && (nb03 % 16 == 0) &&
+            (nb1  % 16 == 0) && (nb2  % 16 == 0) && (nb3  % 16 == 0) &&
+            (((uintptr_t) src0_d) % 16 == 0) && (((uintptr_t) dst_d) % 16 == 0);
+
+        if (can_vec) {
+            const int block_num_y = vec_block_num_y;
+            const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+            const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{block_nums, block_dims, 0, stream};
+            ggml_cuda_kernel_launch(k_get_rows_float_vec<dst_t>, launch_params,
+                (const dst_t *) src0_d, src1_d, dst_d,
+                ne00v, ne11, ne12_fdv,
+                s1, s2, s3,
+                nb01, nb02, nb03,
+                s10, s11, s12);
+            return;
+        }
+    }
+
+    const int block_num_y = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
+    const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+
     const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{block_nums, block_dims, 0, stream};
     ggml_cuda_kernel_launch(k_get_rows_float<src0_t, dst_t>, launch_params,
         src0_d, src1_d, dst_d,

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

@@ -58,6 +58,7 @@
 #include "ggml-cuda/wkv.cuh"
 #include "ggml-cuda/gla.cuh"
 #include "ggml-cuda/gated_delta_net.cuh"
+#include "ggml-cuda/dsv4-hc.cuh"
 #include "ggml-cuda/set.cuh"
 #include "ggml-cuda/set-rows.cuh"
 #include "ggml-cuda/pad_reflect_1d.cuh"
@@ -105,17 +106,27 @@ void ggml_cuda_error(const char * stmt, const char * func, const char * file, in
     GGML_ABORT(GGML_CUDA_NAME " error");
 }
 
+// map a (possibly virtual) device id to the physical CUDA device that backs it
+static int ggml_cuda_get_physical_device(int device) {
+    const ggml_cuda_device_info & info = ggml_cuda_info();
+    GGML_ASSERT(device >= 0 && device < info.device_count);
+    return info.devices[device].physical_device;
+}
+
 // this is faster on Windows
 // probably because the Windows CUDA libraries forget to make this check before invoking the drivers
 void ggml_cuda_set_device(int device) {
+    // translate the (possibly virtual) device id to the physical CUDA device that backs it
+    const int physical_device = ggml_cuda_get_physical_device(device);
+
     int current_device;
     CUDA_CHECK(cudaGetDevice(&current_device));
 
-    if (device == current_device) {
+    if (physical_device == current_device) {
         return;
     }
 
-    CUDA_CHECK(cudaSetDevice(device));
+    CUDA_CHECK(cudaSetDevice(physical_device));
 }
 
 int ggml_cuda_get_device() {
@@ -206,56 +217,102 @@ static int ggml_cuda_parse_id(char devName[]) {
 static ggml_cuda_device_info ggml_cuda_init() {
     ggml_cuda_device_info info = {};
 
-    cudaError_t err = cudaGetDeviceCount(&info.device_count);
+    cudaError_t err = cudaGetDeviceCount(&info.physical_device_count);
     if (err != cudaSuccess) {
         GGML_LOG_ERROR("%s: failed to initialize " GGML_CUDA_NAME ": %s\n", __func__, cudaGetErrorString(err));
         return info;
     }
 
-    GGML_ASSERT(info.device_count <= GGML_CUDA_MAX_DEVICES);
+    GGML_ASSERT(info.physical_device_count <= GGML_CUDA_MAX_DEVICES);
+
+    // by default expose exactly the physical devices; GGML_CUDA_DEVICES can request a different
+    // number of (virtual) devices to emulate multi-GPU systems on a machine with fewer GPUs
+    info.device_count = info.physical_device_count;
+
+    const char * devices_env = getenv("GGML_CUDA_DEVICES");
+    if (devices_env != nullptr && info.physical_device_count > 0) {
+        const int requested = atoi(devices_env);
+        if (requested > 0) {
+            info.device_count = requested;
+        } else {
+            GGML_LOG_WARN("%s: ignoring invalid GGML_CUDA_DEVICES=\"%s\"\n", __func__, devices_env);
+        }
+    }
+
+    if (info.device_count > GGML_CUDA_MAX_DEVICES) {
+        GGML_LOG_WARN("%s: requested %d devices, clamping to GGML_CUDA_MAX_DEVICES=%d\n",
+                      __func__, info.device_count, GGML_CUDA_MAX_DEVICES);
+        info.device_count = GGML_CUDA_MAX_DEVICES;
+    }
+
+    // map each (virtual) device to a backing physical device (round-robin), assign each its index
+    // among the (virtual) devices sharing that physical GPU, and store the per-physical share count
+    int physical_share_count[GGML_CUDA_MAX_DEVICES] = {};
+    GGML_ASSERT(info.device_count == 0 || info.physical_device_count > 0);
+    for (int id = 0; id < info.device_count; ++id) {
+        info.devices[id].physical_device = id % info.physical_device_count;
+        info.devices[id].virtual_index  = physical_share_count[info.devices[id].physical_device]++;
+    }
 
     int64_t total_vram = 0;
-    for (int id = 0; id < info.device_count; ++id) {
+    for (int id = 0; id < info.physical_device_count; ++id) {
         cudaDeviceProp prop;
         CUDA_CHECK(cudaGetDeviceProperties(&prop, id));
         total_vram += prop.totalGlobalMem;
     }
     GGML_LOG_INFO("%s: found %d " GGML_CUDA_NAME " devices (Total VRAM: %zu MiB):\n",
-                  __func__, info.device_count, (size_t)(total_vram / (1024 * 1024)));
+                  __func__, info.physical_device_count, (size_t)(total_vram / (1024 * 1024)));
+    if (info.device_count != info.physical_device_count) {
+        GGML_LOG_INFO("%s: emulating %d virtual device(s) on %d physical device(s) (GGML_CUDA_DEVICES)\n",
+                      __func__, info.device_count, info.physical_device_count);
+    }
     total_vram = 0;
 
     std::vector<std::pair<int, std::string>> turing_devices_without_mma;
     for (int id = 0; id < info.device_count; ++id) {
+        const int physical_id = info.devices[id].physical_device;
+
         int device_vmm = 0;
 
 #if defined(GGML_USE_VMM)
         CUdevice device;
-        CU_CHECK(cuDeviceGet(&device, id));
+        CU_CHECK(cuDeviceGet(&device, physical_id));
         CU_CHECK(cuDeviceGetAttribute(&device_vmm, CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED, device));
 
         if (device_vmm) {
             CUmemAllocationProp alloc_prop = {};
             alloc_prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
             alloc_prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
-            alloc_prop.location.id = id;
+            alloc_prop.location.id = physical_id;
             CU_CHECK(cuMemGetAllocationGranularity(&info.devices[id].vmm_granularity, &alloc_prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED));
         }
 #endif // defined(GGML_USE_VMM)
         info.devices[id].vmm = !!device_vmm;
 
         cudaDeviceProp prop;
-        CUDA_CHECK(cudaGetDeviceProperties(&prop, id));
+        CUDA_CHECK(cudaGetDeviceProperties(&prop, physical_id));
+
+        // a virtual device owns only a share of its physical GPU's memory; report that share so the
+        // logged per-device VRAM sums to the physical total above.
+        GGML_ASSERT(physical_share_count[physical_id] > 0);
+        info.devices[id].physical_share_count = physical_share_count[physical_id];
+        const size_t device_vram = prop.totalGlobalMem / info.devices[id].physical_share_count;
+        const size_t device_vram_mib = device_vram / (1024 * 1024);
 
         info.default_tensor_split[id] = total_vram;
-        total_vram += prop.totalGlobalMem;
+        total_vram += device_vram;
+#if defined(GGML_USE_HIP)
+        info.devices[id].integrated = prop.integrated;
+#else
         info.devices[id].integrated = false; // Temporarily disabled due to issues with corrupted output (e.g. #15034)
+#endif
         info.devices[id].nsm        = prop.multiProcessorCount;
         info.devices[id].smpb       = prop.sharedMemPerBlock;
         info.devices[id].warp_size  = prop.warpSize;
 
 #ifndef GGML_USE_MUSA
         int supports_coop_launch = 0;
-        CUDA_CHECK(cudaDeviceGetAttribute(&supports_coop_launch, cudaDevAttrCooperativeLaunch, id));
+        CUDA_CHECK(cudaDeviceGetAttribute(&supports_coop_launch, cudaDevAttrCooperativeLaunch, physical_id));
         info.devices[id].supports_cooperative_launch = !!supports_coop_launch;
 #else
         info.devices[id].supports_cooperative_launch = false;
@@ -278,7 +335,7 @@ static ggml_cuda_device_info ggml_cuda_init() {
         GGML_LOG_INFO("  Device %d: %s, %s (0x%x), VMM: %s, Wave Size: %d, VRAM: %zu MiB\n",
                       id, prop.name, prop.gcnArchName, info.devices[id].cc & 0xffff,
                       device_vmm ? "yes" : "no", prop.warpSize,
-                      (size_t)(prop.totalGlobalMem / (1024 * 1024)));
+                      device_vram_mib);
 #elif defined(GGML_USE_MUSA)
         // FIXME: Ensure compatibility with varying warp sizes across different MUSA archs.
         info.devices[id].warp_size = 32;
@@ -287,13 +344,13 @@ static ggml_cuda_device_info ggml_cuda_init() {
         info.devices[id].cc += prop.minor * 0x10;
         GGML_LOG_INFO("  Device %d: %s, compute capability %d.%d, VMM: %s, VRAM: %zu MiB\n",
                       id, prop.name, prop.major, prop.minor, device_vmm ? "yes" : "no",
-                      (size_t)(prop.totalGlobalMem / (1024 * 1024)));
+                      device_vram_mib);
 #else
         info.devices[id].smpbo = prop.sharedMemPerBlockOptin;
         info.devices[id].cc = 100*prop.major + 10*prop.minor;
         GGML_LOG_INFO("  Device %d: %s, compute capability %d.%d, VMM: %s, VRAM: %zu MiB\n",
                       id, prop.name, prop.major, prop.minor, device_vmm ? "yes" : "no",
-                      (size_t)(prop.totalGlobalMem / (1024 * 1024)));
+                      device_vram_mib);
         std::string device_name(prop.name);
         if (device_name == "NVIDIA GeForce MX450") {
             turing_devices_without_mma.push_back({ id, device_name });
@@ -308,7 +365,7 @@ static ggml_cuda_device_info ggml_cuda_init() {
         // TODO: Check for future drivers the default scheduling strategy and
         // remove this call again when cudaDeviceScheduleSpin is default.
         if (prop.major == 12 && prop.minor == 1) {
-            CUDA_CHECK(cudaSetDevice(id));
+            CUDA_CHECK(cudaSetDevice(physical_id));
             CUDA_CHECK(cudaSetDeviceFlags(cudaDeviceScheduleSpin));
         }
 
@@ -333,9 +390,9 @@ static ggml_cuda_device_info ggml_cuda_init() {
     // CUBLAS_CHECK(cublasLoggerConfigure(1, 1, 0, nullptr));
 
     if (getenv("GGML_CUDA_P2P") != nullptr) {
-        for (int id = 0; id < info.device_count; ++id) {
-            ggml_cuda_set_device(id);
-            for (int id_other = 0; id_other < info.device_count; ++id_other) {
+        for (int id = 0; id < info.physical_device_count; ++id) {
+            CUDA_CHECK(cudaSetDevice(id));
+            for (int id_other = 0; id_other < info.physical_device_count; ++id_other) {
                 if (id == id_other) {
                     continue;
                 }
@@ -480,6 +537,7 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
     static const size_t CUDA_POOL_VMM_MAX_SIZE = 1ull << 35; // 32 GB
 
     int device;
+    int physical_device;
     CUdeviceptr pool_addr = 0;
     size_t pool_used = 0;
     size_t pool_size = 0;
@@ -490,6 +548,7 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
 
     explicit ggml_cuda_pool_vmm(int device) :
         device(device),
+        physical_device(ggml_cuda_get_physical_device(device)),
         granularity(ggml_cuda_info().devices[device].vmm_granularity) {
     }
 
@@ -525,7 +584,7 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
             CUmemAllocationProp prop = {};
             prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
             prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
-            prop.location.id = device;
+            prop.location.id = physical_device;
             CUmemGenericAllocationHandle handle;
             CU_CHECK(cuMemCreate(&handle, reserve_size, &prop, 0));
 
@@ -554,20 +613,28 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
                 // NCCL implicitly enables peer access (cudaDeviceEnablePeerAccess), and
                 // GGML_CUDA_P2P enables it explicitly. Unlike cudaMalloc buffers, VMM
                 // allocations do not become peer-accessible from that alone, so access
-                // must be granted explicitly here.
+                // must be granted explicitly here. With virtual devices, grant access
+                // on the backing *physical* devices (deduplicated, since several
+                // virtual devices can map to the same physical GPU).
                 std::vector<CUmemAccessDesc> access_descs;
+                bool physical_seen[GGML_CUDA_MAX_DEVICES] = {};
                 const int device_count = ggml_cuda_info().device_count;
                 for (int id = 0; id < device_count; ++id) {
-                    if (id != device) {
+                    const int id_physical = ggml_cuda_get_physical_device(id);
+                    if (id_physical != physical_device) {
                         int can_access_peer = 0;
-                        CUDA_CHECK(cudaDeviceCanAccessPeer(&can_access_peer, id, device));
+                        CUDA_CHECK(cudaDeviceCanAccessPeer(&can_access_peer, id_physical, physical_device));
                         if (!can_access_peer) {
                             continue;
                         }
                     }
+                    if (physical_seen[id_physical]) {
+                        continue;
+                    }
+                    physical_seen[id_physical] = true;
                     CUmemAccessDesc access = {};
                     access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
-                    access.location.id = id;
+                    access.location.id = id_physical;
                     access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
                     access_descs.push_back(access);
                 }
@@ -576,7 +643,7 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
                 // set access for non P2P
                 CUmemAccessDesc access = {};
                 access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
-                access.location.id = device;
+                access.location.id = physical_device;
                 access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
                 CU_CHECK(cuMemSetAccess(start_ptr, reserve_size, &access, 1));
             }
@@ -752,13 +819,17 @@ static bool ggml_backend_cuda_buffer_cpy_tensor(ggml_backend_buffer_t buffer, co
     if (ggml_backend_buffer_is_cuda(src->buffer)) {
         ggml_backend_cuda_buffer_context * src_ctx = (ggml_backend_cuda_buffer_context *)src->buffer->context;
         ggml_backend_cuda_buffer_context * dst_ctx = (ggml_backend_cuda_buffer_context *)dst->buffer->context;
-        if (src_ctx->device == dst_ctx->device) {
+        // compare the backing physical devices: distinct virtual devices may share one physical GPU,
+        // in which case a same-device copy (not a peer copy) is required
+        const int src_physical = ggml_cuda_get_physical_device(src_ctx->device);
+        const int dst_physical = ggml_cuda_get_physical_device(dst_ctx->device);
+        if (src_physical == dst_physical) {
             CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(src), cudaMemcpyDeviceToDevice, cudaStreamPerThread));
         } else {
 #ifdef GGML_CUDA_NO_PEER_COPY
             return false;
 #else
-            CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, dst_ctx->device, src->data, src_ctx->device, ggml_nbytes(src), cudaStreamPerThread));
+            CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, dst_physical, src->data, src_physical, ggml_nbytes(src), cudaStreamPerThread));
 #endif
         }
         CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
@@ -1100,6 +1171,15 @@ static void ggml_backend_cuda_comm_init_internal(ggml_backend_cuda_comm_context
 
 static void ggml_backend_cuda_comm_init_nccl(ggml_backend_cuda_comm_context * ret) {
 #ifdef GGML_USE_NCCL
+    // Disabling NCCL path when CUDA virtual devices are in use since NCCL requires one distinct physical GPU per rank.
+    const ggml_cuda_device_info & info = ggml_cuda_info();
+    if (info.device_count > info.physical_device_count) {
+        GGML_LOG_WARN("NCCL disabled: virtual devices in use; "
+                      "falling back to internal AllReduce\n");
+        ggml_backend_cuda_comm_init_internal(ret);
+        return;
+    }
+
     const size_t n = ret->dev_ids.size();
     ret->comms.resize(n);
     ncclResult_t rc = ncclCommInitAll(ret->comms.data(), (int) n, ret->dev_ids.data());
@@ -2240,6 +2320,15 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_GATED_DELTA_NET:
             ggml_cuda_op_gated_delta_net(ctx, dst);
             break;
+        case GGML_OP_DSV4_HC_COMB:
+            ggml_cuda_op_dsv4_hc_comb(ctx, dst);
+            break;
+        case GGML_OP_DSV4_HC_PRE:
+            ggml_cuda_op_dsv4_hc_pre(ctx, dst);
+            break;
+        case GGML_OP_DSV4_HC_POST:
+            ggml_cuda_op_dsv4_hc_post(ctx, dst);
+            break;
         case GGML_OP_RWKV_WKV7:
             ggml_cuda_op_rwkv_wkv7(ctx, dst);
             break;
@@ -2359,13 +2448,17 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_
 
     if (backend_src != backend_dst) {
         // copy on src stream
-        if (cuda_ctx_src->device == cuda_ctx_dst->device) {
+        // compare the backing physical devices: distinct virtual devices may share one physical GPU,
+        // in which case a same-device copy (not a peer copy) is required
+        const int src_physical = ggml_cuda_get_physical_device(cuda_ctx_src->device);
+        const int dst_physical = ggml_cuda_get_physical_device(cuda_ctx_dst->device);
+        if (src_physical == dst_physical) {
             CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream()));
         } else {
 #ifdef GGML_CUDA_NO_PEER_COPY
             return false;
 #else
-            CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, cuda_ctx_dst->device, src->data, cuda_ctx_src->device, ggml_nbytes(dst), cuda_ctx_src->stream()));
+            CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, dst_physical, src->data, src_physical, ggml_nbytes(dst), cuda_ctx_src->stream()));
 #endif // GGML_CUDA_NO_PEER_COPY
         }
 
@@ -2610,6 +2703,7 @@ static int ggml_cuda_try_gdn_cache_fusion(
 
 static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) {
     args.sigmoid         = false;
+    args.sqrt_softplus   = false;
     args.softmax         = false;
     args.delayed_softmax = false;
     args.prob_bias       = false;
@@ -2623,10 +2717,17 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod
     }
 
     if (nodes[node_idx]->op == GGML_OP_UNARY) {
-        if (ggml_get_unary_op(nodes[node_idx]) != GGML_UNARY_OP_SIGMOID) {
+        const ggml_unary_op unary_op = ggml_get_unary_op(nodes[node_idx]);
+        if (unary_op == GGML_UNARY_OP_SIGMOID) {
+            args.sigmoid = true;
+        } else if (unary_op == GGML_UNARY_OP_SOFTPLUS && node_idx + 1 < n_nodes &&
+                   nodes[node_idx + 1]->op == GGML_OP_SQRT && nodes[node_idx + 1]->src[0] == nodes[node_idx]) {
+            // sqrt(softplus(x)) scoring (DeepSeek-V4)
+            args.sqrt_softplus = true;
+            node_idx++;
+        } else {
             return false;
         }
-        args.sigmoid = true;
     }
 
     if (nodes[node_idx]->op == GGML_OP_ARGSORT) {
@@ -2635,7 +2736,7 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod
 
     node_idx++;
 
-    if (args.sigmoid || args.softmax) {
+    if (args.sigmoid || args.sqrt_softplus || args.softmax) {
         // SOFTMAX -> RESHAPE
         if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_RESHAPE ||
                 nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
@@ -3079,21 +3180,27 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
             const ggml_tensor * scale   = nullptr;
 
             if (!args.delayed_softmax) {
-                ggml_op gating_op = args.sigmoid ? GGML_OP_UNARY : GGML_OP_SOFT_MAX;
-                int     out_nodes[2];  // nodes which can't be elided
+                int out_nodes[2];  // nodes which can't be elided
+
+                if (args.sigmoid) {
+                    ops.insert(ops.end(), { GGML_OP_UNARY });
+                } else if (args.sqrt_softplus) {
+                    ops.insert(ops.end(), { GGML_OP_UNARY, GGML_OP_SQRT });
+                } else {
+                    ops.insert(ops.end(), { GGML_OP_SOFT_MAX });
+                }
+                const int i_probs = i + (int) ops.size() - 1;  // last node of the gating activation
 
                 if (args.prob_bias) {
-                    bias = cgraph->nodes[i + 2]->src[1];
-                    ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW,
+                    bias = cgraph->nodes[i_probs + 2]->src[1];
+                    ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW,
                                             GGML_OP_GET_ROWS });
-                    out_nodes[0] = i + 4;
-                    ids          = cgraph->nodes[i + 4];
+                    out_nodes[0] = i_probs + 4;
                 } else {
-                    ops.insert(ops.end(),
-                               { gating_op, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
-                    out_nodes[0] = i + 3;
-                    ids          = cgraph->nodes[i + 3];
+                    ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
+                    out_nodes[0] = i_probs + 3;
                 }
+                ids = cgraph->nodes[out_nodes[0]];
 
                 if (args.norm) {
                     ops.insert(ops.end(),
@@ -3978,7 +4085,7 @@ static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, co
     ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key);
 
     if (graph->graph == nullptr) {
-        if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_AMPERE) {
+        if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_VOLTA) {
             if (!graph->disable_due_to_gpu_arch) {
                 GGML_LOG_DEBUG("%s: disabling CUDA graphs due to GPU architecture\n", __func__);
             }
@@ -4350,16 +4457,38 @@ int ggml_backend_cuda_get_device_count() {
     return ggml_cuda_info().device_count;
 }
 
-void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size) {
+static std::string ggml_cuda_device_description(int device) {
     cudaDeviceProp prop;
-    CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
-    snprintf(description, description_size, "%s", prop.name);
+    CUDA_CHECK(cudaGetDeviceProperties(&prop, ggml_cuda_get_physical_device(device)));
+
+    const ggml_cuda_device_info & info = ggml_cuda_info();
+    std::string description = prop.name;
+    if (info.device_count > info.physical_device_count) {
+        description += " (physical device " + std::to_string(info.devices[device].physical_device) +
+                       ", virtual device " + std::to_string(info.devices[device].virtual_index) + ")";
+    }
+    return description;
+}
+
+void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size) {
+    snprintf(description, description_size, "%s", ggml_cuda_device_description(device).c_str());
+}
+
+static int ggml_cuda_physical_device_share_count(int device) {
+    const ggml_cuda_device_info & info = ggml_cuda_info();
+    GGML_ASSERT(device >= 0 && device < info.device_count);
+    return info.devices[device].physical_share_count;
 }
 
 void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total) {
     ggml_cuda_set_device(device);
 
     CUDA_CHECK(cudaMemGetInfo(free, total));
+
+    // virtual devices sharing one physical GPU share its memory pool; split it between them
+    const int share_count = ggml_cuda_physical_device_share_count(device);
+    *free  /= share_count;
+    *total /= share_count;
 }
 
 bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size) {
@@ -4510,7 +4639,7 @@ static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t *
 #if defined(__linux__)
     // Check if this is a UMA (Unified Memory Architecture) system
     cudaDeviceProp prop;
-    CUDA_CHECK(cudaGetDeviceProperties(&prop, ctx->device));
+    CUDA_CHECK(cudaGetDeviceProperties(&prop, ggml_cuda_get_physical_device(ctx->device)));
 
     // Check if UMA is explicitly enabled via environment variable
     bool uma_env = getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr;
@@ -4529,13 +4658,17 @@ static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t *
     }
 #endif // defined(__linux__)
 
+    // virtual devices sharing one physical GPU share its memory pool; split it between them
+    const int share_count = ggml_cuda_physical_device_share_count(ctx->device);
+    *free  /= share_count;
+    *total /= share_count;
 }
 
 static enum ggml_backend_dev_type ggml_backend_cuda_device_get_type(ggml_backend_dev_t dev) {
     ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *) dev->context;
 
     cudaDeviceProp prop;
-    CUDA_CHECK(cudaGetDeviceProperties(&prop, ctx->device));
+    CUDA_CHECK(cudaGetDeviceProperties(&prop, ggml_cuda_get_physical_device(ctx->device)));
 
     return prop.integrated
         ? GGML_BACKEND_DEVICE_TYPE_IGPU
@@ -4979,6 +5112,16 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
 #else
             return true;
 #endif // GGML_USE_MUSA
+        case GGML_OP_DSV4_HC_COMB:
+            return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
+                op->src[2]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
+        case GGML_OP_DSV4_HC_PRE:
+            return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
+                op->type == GGML_TYPE_F32;
+        case GGML_OP_DSV4_HC_POST:
+            return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
+                op->src[2]->type == GGML_TYPE_F32 && op->src[3]->type == GGML_TYPE_F32 &&
+                op->type == GGML_TYPE_F32;
         case GGML_OP_FLASH_ATTN_EXT:
             return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op);
         case GGML_OP_CROSS_ENTROPY_LOSS:
@@ -5195,18 +5338,24 @@ ggml_backend_reg_t ggml_backend_cuda_reg() {
             ggml_backend_cuda_reg_context * ctx = new ggml_backend_cuda_reg_context;
             const int min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32;
 
-            for (int i = 0; i < ggml_cuda_info().device_count; i++) {
+            const ggml_cuda_device_info & info = ggml_cuda_info();
+            const bool virtual_devices = info.device_count > info.physical_device_count;
+
+            for (int i = 0; i < info.device_count; i++) {
+                const int physical_id = info.devices[i].physical_device;
+
                 ggml_backend_cuda_device_context * dev_ctx = new ggml_backend_cuda_device_context;
                 dev_ctx->device = i;
                 dev_ctx->name = GGML_CUDA_NAME + std::to_string(i);
-
-                cudaDeviceProp prop;
-                CUDA_CHECK(cudaGetDeviceProperties(&prop, i));
-                dev_ctx->description = prop.name;
+                dev_ctx->description = ggml_cuda_device_description(i);
 
                 char pci_bus_id[32] = {};
-                CUDA_CHECK(cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), i));
+                CUDA_CHECK(cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), physical_id));
                 dev_ctx->pci_bus_id = pci_bus_id;
+                if (virtual_devices) {
+                    // make the pci bus id unique for virtual devices
+                    dev_ctx->pci_bus_id += "-v" + std::to_string(i);
+                }
                 for (char & c : dev_ctx->pci_bus_id) {
                     c = std::tolower(c);
                 }

ggml/src/ggml-cuda/mmf.cu

@@ -85,7 +85,7 @@ void ggml_cuda_mul_mat_f(ggml_backend_cuda_context & ctx, const ggml_tensor * sr
         GGML_ASSERT(sis1 > 0);
 
         ggml_cuda_launch_mm_ids_helper(ids_d, ids_src_compact_dev.get(), ids_dst_compact_dev.get(), expert_bounds_dev.get(),
-            static_cast<int>(n_experts), static_cast<int>(n_tokens), static_cast<int>(n_expert_used), static_cast<int>(ne11), si1, sis1, ctx.stream());
+            static_cast<int>(n_experts), static_cast<int>(n_tokens), static_cast<int>(n_expert_used), static_cast<int>(ne11), si1, sis1, /*write_inverse =*/ false, ctx.stream());
         CUDA_CHECK(cudaGetLastError());
 
         ids_info.ids_src_compact   = ids_src_compact_dev.get();

ggml/src/ggml-cuda/mmid.cu

@@ -27,7 +27,7 @@ template <int n_expert_used_template>
 __launch_bounds__(ggml_cuda_get_physical_warp_size(), 1)
 static __global__ void mm_ids_helper(
         const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
-        const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1) {
+        const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse) {
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
     const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
     const int expert = blockIdx.x;
@@ -98,8 +98,13 @@ static __global__ void mm_ids_helper(
         const mm_ids_helper_store store_it = store[itc];
         const int it       = store_it.it();
         const int iex_used = store_it.iex_used();
-        ids_src1[nex_prev + itc] = it*sis1          + iex_used % nchannels_y;
-        ids_dst [nex_prev + itc] = it*n_expert_used + iex_used;
+        ids_dst[nex_prev + itc] = it*n_expert_used + iex_used;
+        // ids_src1 holds the forward map, or the inverse map (token slot -> compact row) for quant dedup
+        if (write_inverse) {
+            ids_src1[it*n_expert_used + iex_used] = nex_prev + itc;
+        } else {
+            ids_src1[nex_prev + itc] = it*sis1 + iex_used % nchannels_y;
+        }
     }
 
     if (threadIdx.x != 0) {
@@ -118,7 +123,7 @@ static __global__ void mm_ids_helper(
 template <int n_expert_used_template>
 static void launch_mm_ids_helper(
         const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
-        const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, cudaStream_t stream) {
+        const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
     GGML_ASSERT(n_tokens          < (1 << 22) && "too few bits in mm_ids_helper_store");
     GGML_ASSERT(n_expert_used_var < (1 << 10) && "too few bits in mm_ids_helper_store");
 
@@ -132,33 +137,33 @@ static void launch_mm_ids_helper(
     const size_t nbytes_shared = n_tokens*sizeof(mm_ids_helper_store);
     GGML_ASSERT(nbytes_shared <= smpbo);
     mm_ids_helper<n_expert_used_template><<<num_blocks, block_size, nbytes_shared, stream>>>
-        (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1);
+        (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1, write_inverse);
 }
 
 void ggml_cuda_launch_mm_ids_helper(
         const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
-        const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, cudaStream_t stream) {
+        const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
     switch (n_expert_used) {
         case  2:
-            launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
             break;
         case  4:
-            launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
             break;
         case  6:
-            launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
             break;
         case  8:
-            launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
             break;
         case 16:
-            launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
             break;
         case 32:
-            launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
             break;
         default:
-            launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, stream);
+            launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
             break;
     }
 }

ggml/src/ggml-cuda/mmid.cuh

@@ -2,4 +2,4 @@
 
 void ggml_cuda_launch_mm_ids_helper(
         const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds,
-        int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, cudaStream_t stream);
+        int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, bool write_inverse, cudaStream_t stream);

ggml/src/ggml-cuda/mmq.cu

@@ -175,13 +175,17 @@ void ggml_cuda_mul_mat_q(
     ggml_cuda_pool_alloc<int32_t> ids_dst(ctx.pool(), ne_get_rows);
     ggml_cuda_pool_alloc<int32_t> expert_bounds(ctx.pool(), ne02 + 1);
 
+    // gate/up activations are broadcast across experts (ne11 == 1): quantize each token once and
+    // scatter to its slots. ids_src1 then holds the inverse map (token slot -> compact row).
+    const bool dedup_bcast = ne11 == 1 && n_expert_used > 1;
+
     {
         GGML_ASSERT(ids->nb[0] == ggml_element_size(ids));
         const int si1  = ids->nb[1] / ggml_element_size(ids);
         const int sis1 = nb12 / nb11;
 
         ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(),
-            ne02, ne12, n_expert_used, ne11, si1, sis1, stream);
+            ne02, ne12, n_expert_used, ne11, si1, sis1, /*write_inverse =*/ dedup_bcast, stream);
         CUDA_CHECK(cudaGetLastError());
     }
 
@@ -198,7 +202,16 @@ void ggml_cuda_mul_mat_q(
         const int64_t s12 = src1->nb[2] / ts_src1;
         const int64_t s13 = src1->nb[3] / ts_src1;
 
-        if (use_native_fp4) {
+        if (dedup_bcast) {
+            // quantize each token once, scatter its block to all n_expert_used slots
+            if (use_native_fp4) {
+                quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
+                                        /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
+            } else {
+                quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
+                                        /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
+            }
+        } else if (use_native_fp4) {
             quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
                                     ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
         } else {

ggml/src/ggml-cuda/quantize.cu

@@ -75,10 +75,12 @@ __device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) {
 }
 
 
+// scatter: grid over tokens, quantize once, write to all the token's compact rows
+template <bool scatter>
 static __global__ void quantize_mmq_nvfp4(
         const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
         const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
-        const int64_t ne0, const int64_t ne1, const int64_t ne2) {
+        const int64_t ne0, const int64_t ne1, const int64_t ne2, const int n_expert_used) {
 #if defined(BLACKWELL_MMA_AVAILABLE)
 
     const int64_t i0_base = ((int64_t) blockDim.x * blockIdx.y + threadIdx.x) * QK_NVFP4_SUB;
@@ -86,25 +88,25 @@ static __global__ void quantize_mmq_nvfp4(
         return;
     }
 
-    const int64_t i1 = blockIdx.x;
-    const int64_t i2 = blockIdx.z % ne2;
-    const int64_t i3 = blockIdx.z / ne2;
-    const int64_t i01 = ids ? ids[i1] : i1;
     const int64_t k_block = i0_base / QK_FP4_MMQ;
     const int64_t blocks_per_col = (ne0 + QK_FP4_MMQ - 1) / QK_FP4_MMQ;
     if (k_block >= blocks_per_col) {
         return;
     }
-
-    const int64_t ib = blockIdx.z * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x;
-    block_fp4_mmq * y = (block_fp4_mmq *) vy;
-    block_fp4_mmq * yb = y + ib;
-
     const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB;
 
+    int64_t base_idx;
+    if constexpr (scatter) {
+        base_idx = (int64_t) blockIdx.x * s02; // one physical row per token
+    } else {
+        const int64_t i2  = blockIdx.z % ne2;
+        const int64_t i3  = blockIdx.z / ne2;
+        const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
+        base_idx = i3 * s03 + i2 * s02 + i01 * s01;
+    }
+
     float vals_raw[QK_NVFP4_SUB];
     float amax_raw = 0.0f;
-    const int64_t base_idx = i3 * s03 + i2 * s02 + i01 * s01;
 #pragma unroll
     for (int k = 0; k < QK_NVFP4_SUB; k++) {
         const int64_t i00 = i0_base + k;
@@ -160,11 +162,27 @@ static __global__ void quantize_mmq_nvfp4(
         q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 12], inv_scale) << (8 * k + 4);
     }
 
-    uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
-    yqs[2 * sub + 0] = q0;
-    yqs[2 * sub + 1] = q1;
-    reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
+    block_fp4_mmq * y = (block_fp4_mmq *) vy;
+    if constexpr (scatter) {
+#pragma unroll
+        for (int slot = 0; slot < n_expert_used; ++slot) {
+            const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+            block_fp4_mmq * yb = y + (k_block * ne1 + i);
+            uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
+            yqs[2 * sub + 0] = q0;
+            yqs[2 * sub + 1] = q1;
+            reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
+        }
+    } else {
+        block_fp4_mmq * yb = y + (blockIdx.z * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x);
+        uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
+        yqs[2 * sub + 0] = q0;
+        yqs[2 * sub + 1] = q1;
+        reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
+    }
+    GGML_UNUSED(n_expert_used);
 #else
+    GGML_UNUSED(n_expert_used);
     NO_DEVICE_CODE; // This is for Blackwell NVFP4 activations only.
 #endif // defined(BLACKWELL_MMA_AVAILABLE)
 
@@ -172,6 +190,8 @@ static __global__ void quantize_mmq_nvfp4(
 
 // quantize values in the format mxfp4 is stored which is interleaved nibbles
 // i.e. a block a0-a31 is represented as a0a16,a1a17 ...a15a31
+// scatter: grid over tokens, quantize once, write to all the token's compact rows
+template <bool scatter>
 static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
                                           const int32_t * __restrict__ ids,
                                           void * __restrict__ vy,
@@ -181,7 +201,8 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
                                           const int64_t s03,
                                           const int64_t ne0,
                                           const int     ne1,
-                                          const int     ne2) {
+                                          const int     ne2,
+                                          const int     n_expert_used) {
     constexpr int vals_per_scale = 32;
     constexpr int vals_per_warp  = 2 * vals_per_scale;  // Each warp processes 2 blocks of 32 = 64 values
 
@@ -196,30 +217,27 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
         return;
     }
 
-    const int64_t i1 = blockIdx.x;
-    const int64_t i2 = blockIdx.z % ne2;
-    const int64_t i3 = blockIdx.z / ne2;
-
-    ggml_cuda_pdl_sync();
-    const int64_t i01 = ids ? ids[i1] : i1;
-    const int64_t i02 = i2;
-    const int64_t i03 = i3;
-
-    block_fp4_mmq * y = (block_fp4_mmq *) vy;
-
     const int64_t block_fp4_mmq_size = QK_FP4_MMQ;
-    const int64_t ib0                = blockIdx.z * ((int64_t) ne1 * (ne0 / block_fp4_mmq_size));
-    const int64_t ib = ib0 + (warp_start_offset / block_fp4_mmq_size) * ne1 + blockIdx.x;
+    const int64_t k_block            = warp_start_offset / block_fp4_mmq_size;
     const int64_t quad_idx_in_block  = (warp_start_offset % block_fp4_mmq_size) / vals_per_warp;
 
     const int group_id = lane_id_32 / 4;
     const int lane_in_group = lane_id_32 % 4;
     const int base = group_id * 2;
-    char2 * yqs2 = (char2 *) y[ib].qs;
 
-    const int64_t base_pos = i03 * s03 + i02 * s02 + i01 * s01;
+    ggml_cuda_pdl_sync();
+    int64_t base_pos;
+    if constexpr (scatter) {
+        base_pos = (int64_t) blockIdx.x * s02; // one physical row per token
+    } else {
+        const int64_t i2  = blockIdx.z % ne2;
+        const int64_t i3  = blockIdx.z / ne2;
+        const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
+        base_pos = i3 * s03 + i2 * s02 + i01 * s01;
+    }
 
     uint8_t scales[2];
+    char2   packed[2];
 
 #pragma unroll
     for (int b = 0; b < 2; ++b) {
@@ -244,11 +262,8 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
         const float val2 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 1, WARP_SIZE);
         const float val3 = __shfl_sync(0xFFFFFFFF, scaled_val, base + 17, WARP_SIZE);
 
-        if (lane_in_group == 0) {
-            __nv_fp4x4_e2m1 fp4_packed(make_float4(val0, val1, val2, val3));
-
-            yqs2[quad_idx_in_block * 16 + b * 8 + group_id] = *(char2 *) &fp4_packed;
-        }
+        __nv_fp4x4_e2m1 fp4_packed(make_float4(val0, val1, val2, val3));
+        packed[b] = *(char2 *) &fp4_packed;
 #else
         // Fallback: manual FP4 conversion using LUT
         const uint8_t q_val = ggml_cuda_float_to_fp4_e2m1(xi, inv_s);
@@ -258,26 +273,49 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
         const uint8_t q_hi_0 = __shfl_sync(0xFFFFFFFF, q_val, base + 16, WARP_SIZE);
         const uint8_t q_hi_1 = __shfl_sync(0xFFFFFFFF, q_val, base + 17, WARP_SIZE);
 
-        if (lane_in_group == 0) {
-            char2 q;
-            q.x = (q_hi_0 << 4) | q_lo_0;
-            q.y = (q_hi_1 << 4) | q_lo_1;
-            yqs2[quad_idx_in_block * 16 + b * 8 + group_id] = q;
-        }
+        char2 q;
+        q.x = (q_hi_0 << 4) | q_lo_0;
+        q.y = (q_hi_1 << 4) | q_lo_1;
+        packed[b] = q;
 #endif // CUDART_VERSION >= 12080
     }
 
-    if (lane_id_32 == 0) {
-        // Store 2 scales packed into 1 uint32
-        y[ib].d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
+    block_fp4_mmq * y = (block_fp4_mmq *) vy;
+    if constexpr (scatter) {
+#pragma unroll
+        for (int slot = 0; slot < n_expert_used; ++slot) {
+            const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+            block_fp4_mmq * yb = y + (k_block * ne1 + i);
+            char2 * yqs2 = (char2 *) yb->qs;
+            if (lane_in_group == 0) {
+                yqs2[quad_idx_in_block * 16 + 0 * 8 + group_id] = packed[0];
+                yqs2[quad_idx_in_block * 16 + 1 * 8 + group_id] = packed[1];
+            }
+            if (lane_id_32 == 0) {
+                yb->d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
+            }
+        }
+    } else {
+        const int64_t ib0 = blockIdx.z * ((int64_t) ne1 * (ne0 / block_fp4_mmq_size));
+        block_fp4_mmq * yb = y + (ib0 + k_block * ne1 + blockIdx.x);
+        char2 * yqs2 = (char2 *) yb->qs;
+        if (lane_in_group == 0) {
+            yqs2[quad_idx_in_block * 16 + 0 * 8 + group_id] = packed[0];
+            yqs2[quad_idx_in_block * 16 + 1 * 8 + group_id] = packed[1];
+        }
+        if (lane_id_32 == 0) {
+            yb->d4[quad_idx_in_block] = (scales[1] << 8) | scales[0];
+        }
     }
+    GGML_UNUSED(n_expert_used);
 }
 
-template <mmq_q8_1_ds_layout ds_layout>
+// scatter: grid over tokens, quantize once, write to all the token's compact rows
+template <mmq_q8_1_ds_layout ds_layout, bool scatter>
 static __global__ void quantize_mmq_q8_1(
         const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
         const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
-        const int64_t ne0, const int ne1, const int ne2) {
+        const int64_t ne0, const int ne1, const int ne2, const int n_expert_used) {
 
     constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32;
     constexpr int vals_per_sum   = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32;
@@ -288,26 +326,27 @@ static __global__ void quantize_mmq_q8_1(
         return;
     }
 
-    const int64_t i1 = blockIdx.x;
-    const int64_t i2 = blockIdx.z % ne2;
-    const int64_t i3 = blockIdx.z / ne2;
-
     const int64_t i00 = i0;
     ggml_cuda_pdl_sync();
-    const int64_t i01 = ids ? ids[i1] : i1;
-    const int64_t i02 = i2;
-    const int64_t i03 = i3;
+
+    int64_t base_idx;
+    if constexpr (scatter) {
+        base_idx = (int64_t) blockIdx.x * s02; // one physical row per token
+    } else {
+        const int64_t i2  = blockIdx.z % ne2;
+        const int64_t i3  = blockIdx.z / ne2;
+        const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
+        base_idx = i3*s03 + i2*s02 + i01*s01;
+    }
 
     const float4 * x4 = (const float4 *) x;
-
     block_q8_1_mmq * y = (block_q8_1_mmq *) vy;
 
-    const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel
-    const int64_t ib  = ib0 + (i0 / QK8_1_MMQ)*ne1 + blockIdx.x;                    // block index in channel
-    const int64_t iqs = i0 % QK8_1_MMQ;                                             // quant index in block
+    const int64_t k_block = i0 / QK8_1_MMQ; // column block in the channel
+    const int64_t iqs     = i0 % QK8_1_MMQ; // quant index in block
 
     // Load 4 floats per thread and calculate max. abs. value between them:
-    const float4 xi = i0 < ne00 ? x4[(i03*s03 + i02*s02 + i01*s01 + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+    const float4 xi = i0 < ne00 ? x4[(base_idx + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
     float amax = fabsf(xi.x);
     amax = fmaxf(amax, fabsf(xi.y));
     amax = fmaxf(amax, fabsf(xi.z));
@@ -336,40 +375,41 @@ static __global__ void quantize_mmq_q8_1(
     q.y = roundf(xi.y*d_inv);
     q.z = roundf(xi.z*d_inv);
     q.w = roundf(xi.w*d_inv);
-
-    // Write back 4 int8 values as a single 32 bit value for better memory bandwidth:
-    char4 * yqs4 = (char4 *) y[ib].qs;
-    yqs4[iqs/4] = q;
-
-    if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) {
-        if (iqs % 16 != 0 || iqs >= 96) {
-            return;
-        }
-
-        y[ib].d2s6[2 + iqs/16] = sum;
-
-        if (iqs % 64 != 0) {
-            return;
-        }
-
-        const float d = 1.0f / d_inv;
-
-        y[ib].d2s6[iqs/64] = d;
-
-        return;
-    }
-
-    if (iqs % 32 != 0) {
-        return;
-    }
-
     const float d = 1.0f / d_inv;
 
-    if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) {
-        y[ib].ds4[iqs/32] = make_half2(d, sum);
-    } else {
-        y[ib].d4[iqs/32]  = d;
+    // write the block once (normal) or to each of the token's compact rows (scatter)
+    const int nwrite = scatter ? n_expert_used : 1;
+#pragma unroll
+    for (int slot = 0; slot < nwrite; ++slot) {
+        int64_t ib;
+        if constexpr (scatter) {
+            const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+            ib = k_block*ne1 + i;
+        } else {
+            const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel
+            ib = ib0 + k_block*ne1 + blockIdx.x;
+        }
+
+        // Write back 4 int8 values as a single 32 bit value for better memory bandwidth:
+        char4 * yqs4 = (char4 *) y[ib].qs;
+        yqs4[iqs/4] = q;
+
+        if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) {
+            if (iqs % 16 == 0 && iqs < 96) {
+                y[ib].d2s6[2 + iqs/16] = sum;
+                if (iqs % 64 == 0) {
+                    y[ib].d2s6[iqs/64] = d;
+                }
+            }
+        } else if (iqs % 32 == 0) {
+            if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) {
+                y[ib].ds4[iqs/32] = make_half2(d, sum);
+            } else {
+                y[ib].d4[iqs/32]  = d;
+            }
+        }
     }
+    GGML_UNUSED(n_expert_used);
 }
 
 void quantize_row_q8_1_cuda(
@@ -402,16 +442,16 @@ void quantize_mmq_q8_1_cuda(
     const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
     switch (mmq_get_q8_1_ds_layout(type_src0)) {
         case MMQ_Q8_1_DS_LAYOUT_D4:
-            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4>
-                <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, false>
+                <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
             break;
         case MMQ_Q8_1_DS_LAYOUT_DS4:
-            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4>
-                <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, false>
+                <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
             break;
         case MMQ_Q8_1_DS_LAYOUT_D2S6:
-            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6>
-                <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, false>
+                <<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
             break;
         default:
             GGML_ABORT("fatal error");
@@ -419,6 +459,62 @@ void quantize_mmq_q8_1_cuda(
     }
 }
 
+// scatter=true reuses the quant kernel: grid over tokens, ids = inverse map (token slot -> compact row)
+void quantize_scatter_mmq_q8_1_cuda(
+        const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0,
+        const int64_t ne00, const int64_t stride_token, const int64_t ne0,
+        const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) {
+    GGML_ASSERT(ne00 % 4 == 0);
+    GGML_ASSERT(ne0 % QK8_1_MMQ == 0);
+
+    const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
+    const dim3 num_blocks(n_tokens, block_num_y, 1);
+    const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+    switch (mmq_get_q8_1_ds_layout(type_src0)) {
+        case MMQ_Q8_1_DS_LAYOUT_D4:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, true><<<num_blocks, block_size, 0, stream>>>(
+                x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+            break;
+        case MMQ_Q8_1_DS_LAYOUT_DS4:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, true><<<num_blocks, block_size, 0, stream>>>(
+                x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+            break;
+        case MMQ_Q8_1_DS_LAYOUT_D2S6:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, true><<<num_blocks, block_size, 0, stream>>>(
+                x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+            break;
+    }
+}
+
+// scatter=true reuses the quant kernels: grid over tokens, ids = inverse map (token slot -> compact row)
+void quantize_scatter_mmq_fp4_cuda(
+        const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0,
+        const int64_t ne00, const int64_t stride_token, const int64_t ne0,
+        const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) {
+    GGML_ASSERT(ne0 > 0);
+    if (type_src0 == GGML_TYPE_NVFP4) {
+        GGML_ASSERT(ne00 % QK_NVFP4 == 0);
+        constexpr int nvfp4_block_size = 128;
+        const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size);
+        const dim3 block_size(nvfp4_block_size, 1, 1);
+        const dim3 num_blocks(n_tokens, block_num_y, 1);
+        quantize_mmq_nvfp4<true><<<num_blocks, block_size, 0, stream>>>(
+            x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+    } else {
+        GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4);
+        constexpr int nwarps = 8;
+        constexpr int vals_per_block = nwarps * 2 * QK_MXFP4;
+        const int64_t block_num_y = (ne0 + vals_per_block - 1) / vals_per_block;
+        const dim3 block_size(WARP_SIZE, nwarps, 1);
+        const dim3 num_blocks(n_tokens, block_num_y, 1);
+        quantize_mmq_mxfp4<true><<<num_blocks, block_size, 0, stream>>>(
+            x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/(int) nrows_dst, /*ne2=*/1, n_expert_used);
+    }
+}
+
 void quantize_mmq_fp4_cuda(
         const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
         const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
@@ -432,8 +528,8 @@ void quantize_mmq_fp4_cuda(
         const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size);
         const dim3 block_size(nvfp4_block_size, 1, 1);
         const dim3 num_blocks(ne1, block_num_y, ne2 * ne3);
-        quantize_mmq_nvfp4<<<num_blocks, block_size, 0, stream>>>(
-            x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+        quantize_mmq_nvfp4<false><<<num_blocks, block_size, 0, stream>>>(
+            x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
     } else {
         GGML_ASSERT(ne0 % (2 * QK_MXFP4) == 0);
 
@@ -445,6 +541,6 @@ void quantize_mmq_fp4_cuda(
         const dim3    num_blocks(ne1, block_num_y, ne2 * ne3);
         const dim3    block_size(WARP_SIZE, nwarps, 1);
 
-        quantize_mmq_mxfp4<<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2);
+        quantize_mmq_mxfp4<false><<<num_blocks, block_size, 0, stream>>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
     }
 }

ggml/src/ggml-cuda/quantize.cuh

@@ -39,3 +39,28 @@ void quantize_mmq_fp4_cuda(const float *   x,
                              int64_t         ne2,
                              int64_t         ne3,
                              cudaStream_t    stream);
+
+// quantize each token once and scatter the block to its compact rows (via the inverse map)
+void quantize_scatter_mmq_fp4_cuda(const float *   x,
+                                   const int32_t * ids_src1_inv,
+                                   void *          vy,
+                                   ggml_type       type_src0,
+                                   int64_t         ne00,
+                                   int64_t         stride_token,
+                                   int64_t         ne0,
+                                   int64_t         n_tokens,
+                                   int64_t         nrows_dst,
+                                   int             n_expert_used,
+                                   cudaStream_t    stream);
+
+void quantize_scatter_mmq_q8_1_cuda(const float *   x,
+                                    const int32_t * ids_src1_inv,
+                                    void *          vy,
+                                    ggml_type       type_src0,
+                                    int64_t         ne00,
+                                    int64_t         stride_token,
+                                    int64_t         ne0,
+                                    int64_t         n_tokens,
+                                    int64_t         nrows_dst,
+                                    int             n_expert_used,
+                                    cudaStream_t    stream);

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

@@ -8,6 +8,7 @@
 // Kernel config struct - passed by value to CUDA kernel
 struct topk_moe_config {
     bool use_sigmoid;
+    bool use_sqrt_softplus;
     bool with_norm;
     bool delayed_softmax;
 };
@@ -67,6 +68,16 @@ __device__ void sigmoid_warp_inplace(float (&vals)[experts_per_thread], const in
     }
 }
 
+template <int experts_per_thread, bool use_limit>
+__device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], const int limit, const int lane) {
+#pragma unroll
+    for (int i = 0; i < experts_per_thread; i++) {
+        const int  idx    = lane + i * WARP_SIZE;
+        const bool active = !use_limit || (idx < limit);
+        vals[i]           = active ? sqrtf(vals[i] > 20.0f ? vals[i] : logf(1.0f + expf(vals[i]))) : -INFINITY;
+    }
+}
+
 /*
     This kernel does the following:
     1. optionally softmax over the logits per token [n_experts, n_tokens]
@@ -115,6 +126,8 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float *
     if (!config.delayed_softmax) {
         if (config.use_sigmoid) {
            sigmoid_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
+        } else if (config.use_sqrt_softplus) {
+           sqrt_softplus_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
         } else {
            softmax_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
         }
@@ -364,9 +377,10 @@ void ggml_cuda_op_topk_moe(ggml_backend_cuda_context &     ctx,
     }
 
     topk_moe_config config;
-    config.use_sigmoid     = args.sigmoid;
-    config.with_norm       = with_norm;
-    config.delayed_softmax = args.delayed_softmax;
+    config.use_sigmoid       = args.sigmoid;
+    config.use_sqrt_softplus = args.sqrt_softplus;
+    config.with_norm         = with_norm;
+    config.delayed_softmax   = args.delayed_softmax;
 
     if (bias) {
         launch_topk_moe_cuda<true>(ctx, logits_d, weights_d, ids_d, bias_d, n_rows, n_experts, n_expert_used, clamp_val,
@@ -415,7 +429,7 @@ bool ggml_cuda_should_use_topk_moe(const ggml_tensor * gating_op,
     } else if (gating_op->op == GGML_OP_UNARY) {
         ggml_unary_op op = ggml_get_unary_op(gating_op);
 
-        if (op != GGML_UNARY_OP_SIGMOID) {
+        if (op != GGML_UNARY_OP_SIGMOID && op != GGML_UNARY_OP_SOFTPLUS) {
             return false;
         }
     }

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

@@ -5,6 +5,7 @@
 
 struct ggml_cuda_topk_moe_args {
     bool sigmoid{};
+    bool sqrt_softplus{};
     bool softmax{};
     bool delayed_softmax{};
     bool prob_bias{};

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

@@ -21,6 +21,11 @@
 #include <algorithm>
 
 #ifdef _WIN32
+#    define WIN32_LEAN_AND_MEAN
+#    ifndef NOMINMAX
+#       define NOMINMAX
+#    endif
+#    include <windows.h>
 #    include <sal.h>
 #else
 #    include <semaphore.h>
@@ -28,7 +33,9 @@
 #endif
 
 #pragma clang diagnostic ignored "-Wnested-anon-types"
+#pragma clang diagnostic ignored "-Wlanguage-extension-token"
 #pragma clang diagnostic ignored "-Wgnu-anonymous-struct"
+#pragma clang diagnostic ignored "-Wmicrosoft-enum-value"
 
 #include <AEEStdErr.h>
 #include <dspqueue.h>
@@ -134,6 +141,8 @@ static const char * htp_event_name(uint16_t id) {
         case HTP_TRACE_EVT_HVX_FA_K_PREP:  return "HVX_K_PREP";
         case HTP_TRACE_EVT_HVX_FA_V_PREP:  return "HVX_V_PREP";
         case HTP_TRACE_EVT_HMX_COMP:       return "HMX_COMP";
+        case HTP_TRACE_EVT_L2FLUSH:        return "L2FLUSH";
+        case HTP_TRACE_EVT_INIT:           return "INIT";
         default:                           return "UNKNOWN";
     }
 }
@@ -501,6 +510,8 @@ static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t size) {
             }
         }
     }
+
+    GGML_UNUSED(size);
 }
 
 // repack q4_0_tiled tensor into q4_0 data
@@ -554,6 +565,8 @@ static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t size) {
             }
         }
     }
+
+    GGML_UNUSED(size);
 }
 
 // repack q4_1 data into q4_1_tiled tensor
@@ -611,6 +624,8 @@ static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t size) {
             }
         }
     }
+
+    GGML_UNUSED(size);
 }
 
 // repack q4_1_tiled tensor into q4_1 data
@@ -665,6 +680,8 @@ static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t size) {
             }
         }
     }
+
+    GGML_UNUSED(size);
 }
 
 // repack q8_0 data into q8_0_tiled tensor
@@ -711,6 +728,8 @@ static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t size) {
             }
         }
     }
+
+    GGML_UNUSED(size);
 }
 
 // repack q8_0_tiled tensor into q8_0 data
@@ -761,6 +780,8 @@ static void repack_tiled_q8_0(void * data, const ggml_tensor * t, size_t size) {
             }
         }
     }
+
+    GGML_UNUSED(size);
 }
 
 // repack mxfp4 data into mxfp4_tiled tensor
@@ -812,6 +833,8 @@ static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t size)
             }
         }
     }
+
+    GGML_UNUSED(size);
 }
 
 // repack mxfp4_tiled tensor into mxfp4 data
@@ -865,6 +888,8 @@ static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t size)
             }
         }
     }
+
+    GGML_UNUSED(size);
 }
 
 static void ggml_backend_hexagon_buffer_set_tensor(ggml_backend_buffer_t buffer,
@@ -965,11 +990,12 @@ static void ggml_backend_hexagon_buffer_get_tensor(ggml_backend_buffer_t buffer,
 static bool ggml_backend_hexagon_buffer_cpy_tensor(ggml_backend_buffer_t      buffer,
                                                    const struct ggml_tensor * src,
                                                    struct ggml_tensor *       dst) {
+    // we might optimize this later, for now take the slow path (ie get/set_tensor)
+    return false;
+
     GGML_UNUSED(buffer);
     GGML_UNUSED(src);
     GGML_UNUSED(dst);
-    // we might optimize this later, for now take the slow path (ie get/set_tensor)
-    return false;
 }
 
 static void ggml_backend_hexagon_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
@@ -1025,9 +1051,9 @@ static ggml_backend_buffer_t ggml_backend_hexagon_repack_buffer_type_alloc_buffe
     }
 }
 
-static size_t ggml_backend_hexagon_buffer_type_get_alignment(ggml_backend_buffer_type_t buffer_type) {
+static size_t ggml_backend_hexagon_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {
     return 128;  // HVX alignment
-    GGML_UNUSED(buffer_type);
+    GGML_UNUSED(buft);
 }
 
 static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const struct ggml_tensor * t) {
@@ -1039,20 +1065,24 @@ static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffe
         return ggml_row_size(t->type, ne0) * ne1 * ne2 * ne3;
     }
     return ggml_nbytes(t);
+
+    GGML_UNUSED(buft);
 }
 
-static size_t ggml_backend_hexagon_buffer_type_get_max_size(ggml_backend_buffer_type_t buffer_type) {
-    auto * context = static_cast<ggml_backend_hexagon_buffer_type_context *>(buffer_type->context);
+static size_t ggml_backend_hexagon_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) {
+    auto * context = static_cast<ggml_backend_hexagon_buffer_type_context *>(buft->context);
     return context->sess->max_bufsize;
 }
 
 static bool ggml_backend_hexagon_buffer_type_is_host(ggml_backend_buffer_type_t buft) {
     return opt_hostbuf;
+
     GGML_UNUSED(buft);
 }
 
 static bool ggml_backend_hexagon_repack_buffer_type_is_host(ggml_backend_buffer_type_t buft) {
     return false;
+
     GGML_UNUSED(buft);
 }
 
@@ -1098,6 +1128,14 @@ struct ggml_hexagon_opbatch {
     std::unordered_map<const ggml_tensor*, int> t_map; // tensor ptr  to index
     std::unordered_multimap<void*, int>         d_map; // tensor data to index
 
+    struct tensor_range {
+        uint64_t start;
+        uint64_t end;
+        int bi;
+        std::vector<int> tensors;
+    };
+    std::vector<tensor_range> ranges;
+
     unsigned int n_bufs;     // num buffers in the batch
     unsigned int n_tens;     // num tensors ...
     unsigned int n_ops;      // num ops ...
@@ -1117,6 +1155,7 @@ struct ggml_hexagon_opbatch {
         b_map.clear();
         t_map.clear();
         d_map.clear();
+        ranges.clear();
     }
 
     ggml_hexagon_opbatch(ggml_hexagon_session *sess, size_t batch_size, size_t max_vmem) {
@@ -1124,7 +1163,7 @@ struct ggml_hexagon_opbatch {
 
         n_bufs_max = HTP_OP_MAX_BUFS;
         n_ops_max  = batch_size;
-        n_tens_max = n_ops_max + n_ops_max * HTP_OP_MAX_INPUTS;
+        n_tens_max = std::min<size_t>(n_ops_max + n_ops_max * HTP_OP_MAX_INPUTS, HTP_OP_MAX_TENSORS);
 
         b_vmem_max = max_vmem;
 
@@ -1170,6 +1209,71 @@ struct ggml_hexagon_opbatch {
         return bi;
     }
 
+    void add_range(const htp_tensor * h, int ti) {
+        uint64_t t_start = h->data;
+        uint64_t t_end   = t_start + h->size;
+        int      bi      = h->bi;
+
+        int first_match = -1;
+        int unused_idx  = -1;
+        for (size_t i = 0; i < ranges.size(); i++) {
+            if (ranges[i].bi == -1) {
+                unused_idx = i;
+                continue;
+            }
+            if (ranges[i].bi != bi) {
+                continue;
+            }
+            if (ranges[i].start >= t_end || ranges[i].end <= t_start) {
+                continue;
+            }
+
+            if (first_match == -1) {
+                first_match = i;
+                HEX_VERBOSE("ggml-hex: %s range-grow #%d : bi %d [%p, %p) + #%d [%p, %p) -> [%p, %p)\n",
+                    sess->c_name(), (int) i, ranges[i].bi,
+                    (void *) (h_bufs[ranges[i].bi].base + ranges[i].start),
+                    (void *) (h_bufs[ranges[i].bi].base + ranges[i].end),
+                    ti,
+                    (void *) (h_bufs[bi].base + t_start),
+                    (void *) (h_bufs[bi].base + t_end),
+                    (void *) (h_bufs[ranges[i].bi].base + std::min(ranges[i].start, t_start)),
+                    (void *) (h_bufs[ranges[i].bi].base + std::max(ranges[i].end, t_end)));
+
+                ranges[i].start = std::min(ranges[i].start, t_start);
+                ranges[i].end   = std::max(ranges[i].end, t_end);
+                ranges[i].tensors.push_back(ti);
+            } else {
+                HEX_VERBOSE("ggml-hex: %s range-merge #%d [%p, %p) + #%d [%p, %p) -> [%p, %p)\n",
+                    sess->c_name(), first_match,
+                    (void *) (h_bufs[bi].base + ranges[first_match].start),
+                    (void *) (h_bufs[bi].base + ranges[first_match].end),
+                    (int) i,
+                    (void *) (h_bufs[bi].base + ranges[i].start),
+                    (void *) (h_bufs[bi].base + ranges[i].end),
+                    (void *) (h_bufs[bi].base + std::min(ranges[first_match].start, ranges[i].start)),
+                    (void *) (h_bufs[bi].base + std::max(ranges[first_match].end, ranges[i].end)));
+
+                ranges[first_match].start = std::min(ranges[first_match].start, ranges[i].start);
+                ranges[first_match].end   = std::max(ranges[first_match].end, ranges[i].end);
+                ranges[first_match].tensors.insert(
+                    ranges[first_match].tensors.end(),
+                    ranges[i].tensors.begin(),
+                    ranges[i].tensors.end()
+                );
+                ranges[i].bi = -1;
+            }
+        }
+
+        if (first_match == -1) {
+            if (unused_idx != -1) {
+                ranges[unused_idx] = {t_start, t_end, bi, {ti}};
+            } else {
+                ranges.push_back({t_start, t_end, bi, {ti}});
+            }
+        }
+    }
+
     bool same_shape(const htp_tensor * h, const ggml_tensor * t) const {
         int64_t ne0 = t->ne[0];
         int64_t ne1 = t->ne[1];
@@ -1213,6 +1317,7 @@ struct ggml_hexagon_opbatch {
 
         htp_tensor &h = h_tens[ti];
         h.bi    = add_buffer(sbuf);
+        h.ti    = ti;
         h.data  = t_offset;
         h.type  = t->type;
 
@@ -1235,8 +1340,11 @@ struct ggml_hexagon_opbatch {
             h.nb[0] = t->nb[0]; h.nb[1] = t->nb[1]; h.nb[2] = t->nb[2]; h.nb[3] = t->nb[3];
         }
 
+        h.alias = ti;
+        add_range(&h, ti);
+
         h.flags = 0;
-        if (ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+        if (ggml_backend_buffer_get_usage(t->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
             h.flags |= HTP_TENSOR_COMPUTE;
         }
 
@@ -1313,6 +1421,17 @@ struct ggml_hexagon_opbatch {
             o.dst[i] = (i < outputs.size() && outputs[i]) ? add_tensor(outputs[i]) : 0xffff;
         }
     }
+
+    void finalize_ranges() {
+        for (const auto & r : ranges) {
+            if (r.bi == -1) {
+                continue;
+            }
+            for (size_t i = 0; i < r.tensors.size(); i++) {
+                h_tens[r.tensors[i]].alias = r.tensors[(i + 1) % r.tensors.size()];
+            }
+        }
+    }
 };
 
 struct ggml_hexagon_opqueue {
@@ -1571,6 +1690,8 @@ void ggml_hexagon_session::flush_pending(bool all) {
 void ggml_hexagon_session::flush_batch() {
     if (op_batch->empty()) { return; }
 
+    op_batch->finalize_ranges();
+
     htp_opbatch_req req {};
     dspqueue_buffer dbuf{};
 
@@ -1647,7 +1768,7 @@ void ggml_hexagon_session::allocate(int dev_id) noexcept(false) {
 
     GGML_LOG_DEBUG("ggml-hex: %s allocating new session\n", this->name.c_str());
 
-    domain * my_domain = get_domain(this->domain_id);
+    domain * my_domain = htpdrv_get_domain(this->domain_id);
     if (my_domain == NULL) {
         GGML_LOG_ERROR("ggml-hex: unable to get domain struct for CDSP\n");
         throw std::runtime_error("ggml-hex: failed to get CDSP domain (see log for details)");
@@ -1793,16 +1914,6 @@ void ggml_hexagon_session::allocate(int dev_id) noexcept(false) {
         }
     }
 
-    if (opt_profile) {
-        htp_iface_pmu_conf pmu_conf{};
-        std::copy(opt_pmu_evt.begin(), opt_pmu_evt.end(), pmu_conf.events);
-
-        err = htp_iface_profiler(this->handle, opt_profile, &pmu_conf);
-        if (err != 0) {
-            GGML_LOG_ERROR("ggml-hex: failed to enable profiling: 0x%08x\n", (unsigned) err);
-        }
-    }
-
     // Allocate buffers and state for op batching
     this->op_queue = new ggml_hexagon_opqueue(this, opt_opbatch, opt_opqueue);
 
@@ -1821,6 +1932,16 @@ void ggml_hexagon_session::allocate(int dev_id) noexcept(false) {
         throw std::runtime_error("ggml-hex: iface start failed (see log for details)");
     }
     this->valid_iface = true;
+
+    if (opt_profile) {
+        htp_iface_pmu_conf pmu_conf{};
+        std::copy(opt_pmu_evt.begin(), opt_pmu_evt.end(), pmu_conf.events);
+
+        err = htp_iface_profiler(this->handle, opt_profile, &pmu_conf);
+        if (err != 0) {
+            GGML_LOG_ERROR("ggml-hex: failed to enable profiling: 0x%08x\n", (unsigned) err);
+        }
+    }
 }
 
 void ggml_hexagon_session::release() noexcept(true) {
@@ -1929,6 +2050,8 @@ static bool ggml_hexagon_flash_attn_is_hmx_eligible(
     }
 
     return true;
+
+    GGML_UNUSED(sinks);
 }
 
 static bool ggml_hexagon_precompute_flash_attn_params(
@@ -2149,8 +2272,9 @@ static bool ggml_hexagon_supported_gated_delta_net(const struct ggml_hexagon_ses
         return false;
     }
 
-    GGML_UNUSED(sess);
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_matmul_is_hmx_eligible(
@@ -2198,6 +2322,8 @@ static bool ggml_hexagon_matmul_is_hmx_eligible(
     }
 
     return true;
+
+    GGML_UNUSED(dst);
 }
 
 static bool ggml_hexagon_precompute_hmx_mm_params(
@@ -2234,109 +2360,15 @@ static bool ggml_hexagon_precompute_hmx_mm_params(
     if (is_batched_val && wtype == GGML_TYPE_F16 && group_size > 1) {
         // Try grouped path first
         const bool use_dma_activation = (src1->nb[1]/sizeof(float) > (size_t)ne00_padded);
-        size_t best_mblocks = SIZE_MAX;
-        int best_act_threads = 0;
-        size_t best_m_chunk = 0;
-        size_t best_n_chunk = 0;
-        size_t best_vtcm_size = 0;
-
-        int act_threads = n_threads;
-        while (act_threads >= 1) {
-            const size_t f32_scratch_size = use_dma_activation ? hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * ne00_padded * sizeof(float), HTP_MM_HMX_TILE_SIZE) : 0;
-            size_t group_overhead = 256 + f32_scratch_size;
-            size_t group_size_per_n, group_size_per_m, group_size_per_mn;
-            htp_mm_hmx_get_batched_chunk_costs(ne00_padded, group_size, &group_size_per_n, &group_size_per_m, &group_size_per_mn);
-
-            size_t m_chunk_candidate = 0;
-            size_t n_chunk_candidate = 0;
-            size_t vtcm_size_candidate = 0;
-
-            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, ne00_padded, m_chunk_candidate, n_chunk_candidate, group_size, use_dma_activation, pipeline, act_threads);
-                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)) {
-                        best_mblocks = mblocks;
-                        best_act_threads = act_threads;
-                        best_m_chunk = m_chunk_candidate;
-                        best_n_chunk = n_chunk_candidate;
-                        best_vtcm_size = exact_size;
-                    }
-                }
-            }
-            if (act_threads == 1) {
-                act_threads = 0;
-            } else {
-                act_threads /= 2;
-            }
-        }
-
-        if (best_act_threads > 0) {
-            m_chunk = best_m_chunk;
-            n_chunk = best_n_chunk;
-            vtcm_size = best_vtcm_size;
-            act_threads_selected = best_act_threads;
+        if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, ne11, group_size, use_dma_activation, n_threads, pipeline, 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)
-        size_t best_mblocks = SIZE_MAX;
-        int best_act_threads = 0;
-        size_t best_m_chunk = 0;
-        size_t best_n_chunk = 0;
-        size_t best_vtcm_size = 0;
-
-        // For MUL_MAT_ID the kernel runs one 2D matmul per expert, with M equal to the number of rows routed to that expert.
-        // A single expert can receive up to all routed rows (dst->ne[1]*dst->ne[2] = n_expert_used*n_tokens), so size the chunk
-        // search for that upper bound rather than ne12 (token positions only).
-        // We recompute m_chunk per expert against the actual count in the NPU kernel.
-        const int m_id_rows    = (int) ((size_t) dst->ne[1] * dst->ne[2]);
-        const int m_for_chunks = is_matmul_id ? hex_align_up(m_id_rows, 32) : ne11_padded;
-        const int m_for_cost   = is_matmul_id ? m_id_rows : ne11;
-
-        int act_threads = n_threads;
-        while (act_threads >= 1) {
-            const size_t act_f32_size = is_matmul_id ? 0 : hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * ne00_padded * sizeof(float), HTP_MM_HMX_TILE_SIZE);
-            size_t simple_2d_overhead = 256 + act_f32_size;
-            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, ne00_padded, pipeline, aligned_tile_size, &simple_2d_size_per_n, &simple_2d_size_per_m, &simple_2d_size_per_mn);
-
-            size_t m_chunk_candidate = 0;
-            size_t n_chunk_candidate = 0;
-            size_t vtcm_size_candidate = 0;
-
-            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, ne00_padded, m_chunk_candidate, n_chunk_candidate, pipeline, is_matmul_id ? 0 : act_threads, aligned_tile_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)) {
-                        best_mblocks = mblocks;
-                        best_act_threads = act_threads;
-                        best_m_chunk = m_chunk_candidate;
-                        best_n_chunk = n_chunk_candidate;
-                        best_vtcm_size = exact_size;
-                    }
-                }
-            }
-            if (act_threads == 1) {
-                act_threads = 0;
-            } else {
-                act_threads /= 2;
-            }
-        }
-
-        if (best_act_threads > 0) {
-            m_chunk = best_m_chunk;
-            n_chunk = best_n_chunk;
-            vtcm_size = best_vtcm_size;
-            act_threads_selected = best_act_threads;
-        } else {
+        const int m_id_rows = (int) ((size_t) dst->ne[1] * dst->ne[2]);
+        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, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
             return false;
         }
     }
@@ -2352,6 +2384,8 @@ static bool ggml_hexagon_precompute_hmx_mm_params(
     kparams->src1_row_size = (wtype == GGML_TYPE_Q4_1) ? 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_dst_size = 0;
 
@@ -2361,6 +2395,8 @@ static bool ggml_hexagon_precompute_hmx_mm_params(
         kparams->kernel_type = HTP_MM_KERNEL_HMX_2D;
     }
     return true;
+
+    GGML_UNUSED(src0);
 }
 
 static void ggml_hexagon_precompute_hvx_mm_params(
@@ -2955,6 +2991,8 @@ static bool ggml_hexagon_supported_binary(const struct ggml_hexagon_session * se
     }
 
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_add_id(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -2981,6 +3019,8 @@ static bool ggml_hexagon_supported_add_id(const struct ggml_hexagon_session * se
     }
 
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3006,6 +3046,8 @@ static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * ses
     }
 
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_sum_rows(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3025,10 +3067,11 @@ static bool ggml_hexagon_supported_sum_rows(const struct ggml_hexagon_session *
     }
 
     return true;
+
+    GGML_UNUSED(sess);
 }
 
-static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session * sess,
-                                               const struct ggml_tensor *          op) {
+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];
     const struct ggml_tensor * dst  = op;
@@ -3057,6 +3100,8 @@ static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session
     }
 
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_softmax(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3122,6 +3167,8 @@ static bool ggml_hexagon_supported_softmax(const struct ggml_hexagon_session * s
     }
 
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_set_rows(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3142,6 +3189,8 @@ static bool ggml_hexagon_supported_set_rows(const struct ggml_hexagon_session *
     }
 
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3162,6 +3211,8 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session *
     }
 
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_argsort(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3182,6 +3233,8 @@ static bool ggml_hexagon_supported_argsort(const struct ggml_hexagon_session * s
     }
 
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_rope(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3243,6 +3296,8 @@ static bool ggml_hexagon_supported_rope(const struct ggml_hexagon_session * sess
         return false;
     }
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3282,6 +3337,8 @@ static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session *
     }
 
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_pad(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3292,8 +3349,9 @@ static bool ggml_hexagon_supported_pad(const struct ggml_hexagon_session * sess,
         return false;
     }
 
-    GGML_UNUSED(sess);
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_cumsum(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3308,8 +3366,9 @@ static bool ggml_hexagon_supported_cumsum(const struct ggml_hexagon_session * se
         return false;
     }
 
-    GGML_UNUSED(sess);
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_diag(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3331,8 +3390,9 @@ static bool ggml_hexagon_supported_diag(const struct ggml_hexagon_session * sess
         return false;
     }
 
-    GGML_UNUSED(sess);
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_solve_tri(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3364,8 +3424,9 @@ static bool ggml_hexagon_supported_solve_tri(const struct ggml_hexagon_session *
         return false;
     }
 
-    GGML_UNUSED(sess);
     return true;
+
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_tri(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -3415,6 +3476,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
         case GGML_OP_RMS_NORM:        return HTP_OP_RMS_NORM;
         case GGML_OP_CONCAT:          return HTP_OP_CONCAT;
         case GGML_OP_SCALE:           return HTP_OP_SCALE;
+        case GGML_OP_CLAMP:           return HTP_OP_CLAMP;
         case GGML_OP_SQR:             return HTP_OP_SQR;
         case GGML_OP_SQRT:            return HTP_OP_SQRT;
         case GGML_OP_SOFT_MAX:        return HTP_OP_SOFTMAX;
@@ -3812,6 +3874,8 @@ static void ggml_backend_hexagon_graph_optimize(ggml_backend_t backend, ggml_cgr
             }
         }
     }
+
+    GGML_UNUSED(backend);
 }
 
 static struct ggml_backend_i hexagon_backend_i = {
@@ -3930,6 +3994,8 @@ static bool ggml_hexagon_supported_buffers(ggml_hexagon_session *sess, const str
 }
 
 static bool ggml_hexagon_supported_cpy(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
+    GGML_UNUSED(sess);
+
     const struct ggml_tensor * src0 = op->src[0];
     const struct ggml_tensor * dst  = op;
 
@@ -4000,6 +4066,7 @@ static bool ggml_hexagon_supported_concat(const struct ggml_hexagon_session * se
     }
 
     return true;
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_fill(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -4009,8 +4076,8 @@ static bool ggml_hexagon_supported_fill(const struct ggml_hexagon_session * sess
         return false;
     }
 
-    GGML_UNUSED(sess);
     return true;
+    GGML_UNUSED(sess);
 }
 
 static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) {
@@ -4060,6 +4127,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
         case GGML_OP_L2_NORM:
         case GGML_OP_RMS_NORM:
         case GGML_OP_SCALE:
+        case GGML_OP_CLAMP:
             supp = ggml_hexagon_supported_unary(sess, op);
             break;
 
@@ -4293,6 +4361,7 @@ static void * ggml_backend_hexagon_get_proc_address(ggml_backend_reg_t reg, cons
     }
 
     return NULL;
+    GGML_UNUSED(reg);
 }
 
 template<typename T> std::vector<T> str_to_vec(const char* str) {
@@ -4351,10 +4420,18 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
 
     // Init Arch first since it affects other defaults
     if (!str_arch) {
-        int err = get_hex_arch_ver(CDSP_DOMAIN_ID, &opt_arch);
+        int err = htpdrv_get_arch(CDSP_DOMAIN_ID, &opt_arch);
         if (err != 0) {
             GGML_LOG_ERROR("ggml-hex: failed to query HTP version (err %d) defaulting to v73\n", err);
             opt_arch = 73;
+        } else {
+            if (opt_arch < 73) {
+                GGML_LOG_WARN("ggml-hex: Hexagon arch v%d is under supported range, capping at v73\n", opt_arch);
+                opt_arch = 73;
+            } else if (opt_arch > 81) {
+                GGML_LOG_WARN("ggml-hex: Hexagon arch v%d is over supported range, capping at v81\n", opt_arch);
+                opt_arch = 81;
+            }
         }
     } else {
         if (str_arch[0] == 'v' || str_arch[0] == 'V') {

ggml/src/ggml-hexagon/htp-drv.cpp

@@ -1,13 +1,8 @@
-// sample drv interface
-
-#pragma clang diagnostic ignored "-Wgnu-anonymous-struct"
-#pragma clang diagnostic ignored "-Wmissing-prototypes"
-#pragma clang diagnostic ignored "-Wsign-compare"
-
 #include <filesystem>
 #include <set>
 #include <sstream>
 #include <string>
+
 #ifdef _WIN32
 #   define WIN32_LEAN_AND_MEAN
 #   ifndef NOMINMAX
@@ -16,9 +11,17 @@
 #   include <windows.h>
 #   include <winevt.h>
 #else
-#    include <dlfcn.h>
-#    include <unistd.h>
+#   include <dlfcn.h>
+#   include <unistd.h>
 #endif
+
+#pragma clang diagnostic ignored "-Wgnu-anonymous-struct"
+#pragma clang diagnostic ignored "-Wmissing-prototypes"
+#pragma clang diagnostic ignored "-Wsign-compare"
+#pragma clang diagnostic ignored "-Wlanguage-extension-token"
+#pragma clang diagnostic ignored "-Wmicrosoft-enum-value"
+#pragma clang diagnostic ignored "-Wnested-anon-types"
+
 #include "ggml-impl.h"
 #include "htp-drv.h"
 #include "libdl.h"
@@ -359,7 +362,7 @@ int htpdrv_init() {
     return AEE_SUCCESS;
 }
 
-domain * get_domain(int domain_id) {
+domain * htpdrv_get_domain(int domain_id) {
     int i    = 0;
     int size = sizeof(supported_domains) / sizeof(domain);
 
@@ -372,7 +375,7 @@ domain * get_domain(int domain_id) {
     return NULL;
 }
 
-int get_hex_arch_ver(int domain, int * arch) {
+int htpdrv_get_arch(int domain, int * arch) {
     if (!remote_handle_control_pfn) {
         GGML_LOG_ERROR("ggml-hex: remote_handle_control is not supported on this device\n");
         return AEE_EUNSUPPORTEDAPI;
@@ -394,25 +397,7 @@ int get_hex_arch_ver(int domain, int * arch) {
         return err;
     }
 
-    switch (arch_ver.capability & 0xff) {
-        case 0x68:
-            *arch = 68;
-            return 0;
-        case 0x69:
-            *arch = 69;
-            return 0;
-        case 0x73:
-            *arch = 73;
-            return 0;
-        case 0x75:
-            *arch = 75;
-            return 0;
-        case 0x79:
-            *arch = 79;
-            return 0;
-        case 0x81:
-            *arch = 81;
-            return 0;
-    }
-    return -1;
+    uint32_t val = arch_ver.capability & 0xff;
+    *arch = (int) ((val >> 4) * 10 + (val & 0x0f));
+    return 0;
 }

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

@@ -96,17 +96,17 @@ extern "C" {
 HTPDRV_API int htpdrv_init(void);
 
 /**
- * get_domain API: get domain struct from domain value.
+ * htpdrv_get_domain API: get domain struct from domain value.
  *
  * @param[in]  domain value of a domain
  * @return     Returns domain struct of the domain if it is supported or else
  *             returns NULL.
  *
  */
-HTPDRV_API domain * get_domain(int domain_id);
+HTPDRV_API domain * htpdrv_get_domain(int domain_id);
 
 /**
- * get_hex_arch_ver API: query the Hexagon processor architecture version information
+ * htpdrv_get_arch API: query the Hexagon processor architecture version information
  *
  * @param[in]   domain_id value of a domain
  * @param[out]  Arch version (73, 75, ...)
@@ -114,7 +114,7 @@ HTPDRV_API domain * get_domain(int domain_id);
  *              non-zero if error, return value points to the error.
  *
  */
-HTPDRV_API int get_hex_arch_ver(int domain, int * arch);
+HTPDRV_API int htpdrv_get_arch(int domain, int * arch);
 
 #ifdef __cplusplus
 }

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

@@ -17,9 +17,12 @@ set(HTP_LIB ggml-htp-${DSP_VERSION})
 add_library(${HTP_LIB} SHARED
     main.c
     htp_iface_skel.c
-    worker-pool.c
-    hex-dma.c
+    work-queue.c
+    dma-queue.c
     hmx-queue.c
+    htp-tensor.c
+    matmul-ops.c
+    flash-attn-ops.c
     gated-delta-net-ops.c
     binary-ops.c
     unary-ops.c
@@ -38,8 +41,6 @@ add_library(${HTP_LIB} SHARED
     diag-ops.c
     solve-tri-ops.c
     pad-ops.c
-    flash-attn-ops.c
-    matmul-ops.c
     argsort-ops.c
 )
 

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

@@ -16,6 +16,7 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 
 #define htp_act_preamble                                 \
     const struct htp_tensor * src0 = actx->octx->src[0]; \
@@ -766,17 +767,11 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
 }
 
 int op_activations(struct htp_ops_context * octx) {
-    int err = HTP_STATUS_OK;
-
     switch (octx->src[0]->type) {
         case HTP_TYPE_F32:
-            err = execute_op_activations_f32(octx);
-            break;
+            return execute_op_activations_f32(octx);
 
         default:
-            err = HTP_STATUS_NO_SUPPORT;
-            break;
+            return HTP_STATUS_NO_SUPPORT;
     }
-
-    return err;
 }

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

@@ -345,7 +345,7 @@ static void htp_argsort_f32_##ne00##_##order_name(unsigned int n, unsigned int i
     int32_t * indices_buf = (int32_t *) (spad + values_size);                                                  \
     uint32_t nb01 = src0->nb[1];                                                                               \
     uint32_t nb1 = dst->nb[1];                                                                                 \
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[i] : NULL;                                    \
+    struct htp_thread_trace * tr = &octx->ctx->trace[i];                                                       \
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start_row);                                              \
     for (uint32_t r = start_row; r < end_row; r++) {                                                           \
         uint32_t src_offset = r * nb01;                                                                        \
@@ -411,7 +411,7 @@ static void htp_argsort_f32_fallback(unsigned int n, unsigned int i, void * data
     const HVX_Vector ind_init_vec = *(HVX_Vector *)argosrt_ramp_lut;
     const HVX_Vector ind_diff_vec = Q6_V_vsplat_R(32);
 
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[i] : NULL;
+    struct htp_thread_trace * tr = &octx->ctx->trace[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start_row);
 
     for (uint32_t r = start_row; r < end_row; r++) {

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

@@ -16,6 +16,7 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 
 #ifndef MIN
 #define MIN(a, b) ((a) < (b) ? (a) : (b))

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

@@ -9,6 +9,7 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 #include "hvx-types.h"
 #include "hvx-utils.h"
 #include "hex-dma.h"
@@ -255,16 +256,10 @@ int op_cumsum_f32(struct htp_ops_context * octx) {
 int op_cumsum(struct htp_ops_context * octx) {
     const struct htp_tensor * dst = octx->dst;
 
-    int err = HTP_STATUS_OK;
-
     switch (dst->type) {
         case HTP_TYPE_F32:
-            err = op_cumsum_f32(octx);
-            break;
+            return op_cumsum_f32(octx);
         default:
-            err = HTP_STATUS_NO_SUPPORT;
-            break;
+            return HTP_STATUS_NO_SUPPORT;
     }
-
-    return err;
 }

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

@@ -0,0 +1,106 @@
+#include "dma-queue.h"
+
+#include <stdbool.h>
+#include <stdlib.h>
+#include <string.h>
+
+#pragma clang diagnostic ignored "-Wunused-function"
+
+static inline uint32_t pow2_ceil(uint32_t x) {
+    if (x <= 1) {
+        return 1;
+    }
+    int p = 2;
+    x--;
+    while (x >>= 1) {
+        p <<= 1;
+    }
+    return p;
+}
+
+static inline uintptr_t align_up(uintptr_t addr, size_t align) {
+    return (addr + align - 1) & ~(align - 1);
+}
+
+size_t dma_queue_sizeof(size_t capacity) {
+    capacity = pow2_ceil(capacity);
+
+    size_t size_q      = sizeof(dma_queue);
+    size_t offset_r    = align_up(size_q, HEX_L2_LINE_SIZE);
+    size_t size_r      = sizeof(dma_ring);
+    size_t offset_desc = align_up(offset_r + size_r, HEX_L2_LINE_SIZE);
+    size_t size_desc   = capacity * sizeof(dma_descriptor_2d);
+    size_t offset_dptr = align_up(offset_desc + size_desc, HEX_L2_LINE_SIZE);
+    size_t size_dptr   = capacity * sizeof(dma_ptr);
+
+    return offset_dptr + size_dptr;
+}
+
+size_t dma_queue_alignof(void) {
+    return HEX_L2_LINE_SIZE;
+}
+
+dma_queue_t dma_queue_init(void * ptr, size_t capacity, uintptr_t vtcm_base, size_t vtcm_size, struct htp_thread_trace * trace) {
+    capacity = pow2_ceil(capacity);
+
+    size_t size_q      = sizeof(dma_queue);
+    size_t offset_r    = align_up(size_q, HEX_L2_LINE_SIZE);
+    size_t size_r      = sizeof(dma_ring);
+    size_t offset_desc = align_up(offset_r + size_r, HEX_L2_LINE_SIZE);
+    size_t size_desc   = capacity * sizeof(dma_descriptor_2d);
+    size_t offset_dptr = align_up(offset_desc + size_desc, HEX_L2_LINE_SIZE);
+    size_t size_dptr   = capacity * sizeof(dma_ptr);
+
+    size_t total_size = offset_dptr + size_dptr;
+    memset(ptr, 0, total_size);
+
+    dma_queue * q = (dma_queue *) ptr;
+    dma_ring * r  = (dma_ring *) ((uintptr_t) ptr + offset_r);
+
+    q->ring    = r;
+    q->nocache = 0;
+    q->alias   = false;
+
+    r->trace     = trace;
+    r->vtcm_base = vtcm_base;
+    r->vtcm_end  = vtcm_base + vtcm_size;
+    r->capacity  = capacity;
+    r->idx_mask  = capacity - 1;
+    r->push_idx  = 0;
+    r->pop_idx   = 0;
+
+    r->desc = (dma_descriptor_2d *) ((uintptr_t) ptr + offset_desc);
+    r->dptr = (dma_ptr *) ((uintptr_t) ptr + offset_dptr);
+    r->tail = &r->desc[capacity - 1];
+
+    FARF(HIGH, "dma-queue: capacity %u, unified memory size %zu\n", capacity, total_size);
+
+    return q;
+}
+
+void dma_queue_free(dma_queue_t q) {
+    (void) q;
+}
+
+size_t dma_queue_alias_sizeof(void) {
+    return sizeof(dma_queue);
+}
+
+dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q, uint8_t nocache) {
+    dma_queue * q = (dma_queue *) ptr;
+    memset(q, 0, sizeof(dma_queue));
+
+    q->ring    = main_q->ring;
+    q->nocache = nocache;
+    q->alias   = true;
+
+    return q;
+}
+
+void dma_queue_alias_free(dma_queue_t q) {
+    (void) q;
+}
+
+void dma_queue_flush(dma_queue_t q) {
+    while (dma_queue_pop(q).dst != NULL) ;
+}

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

@@ -0,0 +1,396 @@
+#ifndef HTP_DMA_H
+#define HTP_DMA_H
+
+#include <HAP_farf.h>
+#include <hexagon_types.h>
+#include <stdbool.h>
+#include <stdint.h>
+#include "hex-utils.h"
+
+#include "hex-profile.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Define the HW descriptor structs here since the ones in HexSDK are a bit out of date
+typedef struct dma_descriptor_1d_s {
+    void *   next;
+    uint32_t size:24;
+    uint32_t desc_size:2;
+    uint32_t dst_comp:1;
+    uint32_t src_comp:1;
+    uint32_t dst_bypass:1;
+    uint32_t src_bypass:1;
+    uint32_t order:1;
+    uint32_t done:1;
+    void *   src;
+    void *   dst;
+} dma_descriptor_1d;
+
+#if __HVX_ARCH__ < 75
+
+typedef struct dma_descriptor_2d_s {
+    void *   next;
+    uint32_t reserved0:24;
+    uint32_t desc_size:2;
+    uint32_t dst_comp:1;
+    uint32_t src_comp:1;
+    uint32_t dst_bypass:1;
+    uint32_t src_bypass:1;
+    uint32_t order:1;
+    uint32_t done:1;
+    void *   src;
+    void *   dst;
+    uint32_t desc_type:8;
+    uint32_t reserved1:24;
+    uint32_t row_size:16;
+    uint32_t nrows:16;
+    uint32_t src_stride:16;
+    uint32_t dst_stride:16;
+    uint32_t src_offset:16;
+    uint32_t dst_offset:16;
+} dma_descriptor_2d;
+
+#else
+
+typedef struct dma_descriptor_2d_s {
+    void *   next;
+    uint32_t dst_stride:24;
+    uint32_t desc_size:2;
+    uint32_t dst_comp:1;
+    uint32_t src_comp:1;
+    uint32_t dst_bypass:1;
+    uint32_t src_bypass:1;
+    uint32_t order:1;
+    uint32_t done:1;
+    void *   src;
+    void *   dst;
+    uint32_t desc_type:8;
+    uint32_t reserved0:24;
+    uint32_t row_size:24;
+    uint32_t nrows_lo:8;
+    uint32_t nrows_hi:8;
+    uint32_t src_stride:24;
+    uint32_t offset:24;
+    uint32_t reserved1:8;
+} dma_descriptor_2d;
+
+#endif
+
+typedef struct {
+    void       *dst;
+    const void *src;
+} dma_ptr;
+
+typedef struct dma_ring_s dma_ring;
+struct dma_ring_s {
+    dma_descriptor_2d * desc;      // descriptor pointers
+    dma_descriptor_2d * tail;      // tail pointer
+    dma_ptr           * dptr;      // dst/src pointers
+    uint32_t            push_idx;
+    uint32_t            pop_idx;
+    uint32_t            capacity;
+    uint32_t            idx_mask;
+    struct htp_thread_trace * trace;
+    uintptr_t           vtcm_base;
+    uintptr_t           vtcm_end;
+};
+
+typedef struct dma_queue_s dma_queue;
+typedef dma_queue * dma_queue_t;
+
+struct dma_queue_s {
+    dma_ring *          ring;      // Points to the descriptor ring state
+    uint8_t             nocache;   // Queue-specific bypass flag
+    bool                alias;     // When set, dma_queue_delete will not free the ring
+};
+
+void        dma_queue_flush(dma_queue_t q);
+
+size_t      dma_queue_sizeof(size_t capacity);
+size_t      dma_queue_alignof(void);
+dma_queue_t dma_queue_init(void * ptr, size_t capacity, uintptr_t vtcm_base, size_t vtcm_size, struct htp_thread_trace * trace);
+void        dma_queue_free(dma_queue_t q);
+
+size_t      dma_queue_alias_sizeof(void);
+dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q, uint8_t nocache);
+void        dma_queue_alias_free(dma_queue_t q);
+
+// TODO: technically we don't need these and could use Q6_dmstart/wait/etc instead
+// but those do not seem to always compiler properly.
+static inline void dmstart(void * next) {
+    asm volatile(" release(%0):at" : : "r"(next));
+    asm volatile(" dmstart(%0)" : : "r"(next));
+}
+
+static inline void dmlink(void * cur, void * next) {
+    asm volatile(" release(%0):at" : : "r"(next));
+    asm volatile(" dmlink(%0, %1)" : : "r"(cur), "r"(next));
+}
+
+static inline unsigned int dmpoll(void) {
+    unsigned int ret = 0;
+    asm volatile(" %0 = dmpoll" : "=r"(ret) : : "memory");
+    return ret;
+}
+
+static inline unsigned int dmwait(void) {
+    unsigned int ret = 0;
+    asm volatile(" %0 = dmwait" : "=r"(ret) : : "memory");
+    return ret;
+}
+
+static inline dma_ptr dma_make_ptr(void *dst, const void *src)
+{
+    dma_ptr p = { dst, src };
+    return p;
+}
+
+static inline bool dma_is_vtcm(const dma_queue * q, const void * ptr) {
+    return (uintptr_t) ptr >= q->ring->vtcm_base && (uintptr_t) ptr < q->ring->vtcm_end;
+}
+
+static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t size) {
+    dma_ring * r = q->ring;
+    if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
+        FARF(HIGH, "dma-push: queue full\n");
+        return false;
+    }
+
+    dma_descriptor_1d * desc = (dma_descriptor_1d *) &r->desc[r->push_idx];
+    desc->src  = (void *) dptr.src;
+    desc->dst  = (void *) dptr.dst;
+    desc->size = size;
+
+    r->dptr[r->push_idx] = dptr;
+
+    if (size) {
+        desc->next       = NULL;
+        desc->desc_size  = 0; // 1D mode
+        desc->src_bypass = dma_is_vtcm(q, dptr.src) ? 1 : q->nocache;
+        desc->dst_bypass = dma_is_vtcm(q, dptr.dst) ? 1 : q->nocache;
+        desc->order      = 0;
+        desc->done       = 0;
+
+        htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
+        dmlink(r->tail, desc);
+        r->tail = (dma_descriptor_2d *) desc;
+    } else {
+        desc->desc_size = 0;
+        desc->done      = 1;
+    }
+
+    r->push_idx = (r->push_idx + 1) & r->idx_mask;
+    return true;
+}
+
+static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
+    dma_ring * r = q->ring;
+    if (((r->push_idx + 1) & r->idx_mask) == r->pop_idx) {
+        FARF(HIGH, "dma-push: queue full\n");
+        return false;
+    }
+
+    dma_descriptor_2d * desc = &r->desc[r->push_idx];
+
+    desc->next           = NULL;
+    desc->reserved0      = 0;
+    desc->reserved1      = 0;
+    desc->desc_size      = 1; // 2d mode
+    desc->src_bypass     = dma_is_vtcm(q, dptr.src) ? 1 : q->nocache;
+    desc->dst_bypass     = dma_is_vtcm(q, dptr.dst) ? 1 : q->nocache;
+    desc->src_comp       = 0;
+    desc->dst_comp       = 0;
+    desc->order          = 0;
+    desc->done           = 0;
+    desc->src_stride     = src_stride;
+    desc->dst_stride     = dst_stride;
+    desc->src            = (void *) dptr.src;
+    desc->dst            = (void *) dptr.dst;
+    desc->row_size       = row_size;
+
+#if __HVX_ARCH__ < 75
+    desc->desc_type      = 0; // 2d (16-bit) mode
+    desc->nrows          = nrows;
+    desc->src_offset     = 0;
+    desc->dst_offset     = 0;
+#else
+    desc->desc_type      = 9; // 2d (24-bit) mode
+    desc->nrows_lo       = (nrows & 0xff);
+    desc->nrows_hi       = (nrows >> 8);
+    desc->offset         = 0;
+#endif
+
+    r->dptr[r->push_idx] = dptr;
+
+    if (nrows) {
+        htp_trace_event_start(r->trace, HTP_TRACE_EVT_DMA, r->push_idx);
+        dmlink(r->tail, desc);
+        r->tail = desc;
+    } else {
+        desc->done = 1;
+    }
+
+    r->push_idx = (r->push_idx + 1) & r->idx_mask;
+    return true;
+}
+
+static inline dma_ptr dma_queue_pop(dma_queue * q) {
+    dma_ring * r = q->ring;
+    dma_ptr dptr  = { NULL };
+
+    if (r->push_idx == r->pop_idx) {
+        return dptr;
+    }
+
+    dma_descriptor_2d * desc = &r->desc[r->pop_idx];
+
+    // Wait for desc to complete
+    if (!desc->done) {
+        while (!desc->done) {
+            dmpoll();
+        }
+    }
+    htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx);
+
+    dptr = r->dptr[r->pop_idx];
+
+    r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
+    return dptr;
+}
+
+static inline dma_ptr dma_queue_pop_nowait(dma_queue * q) {
+    dma_ring * r = q->ring;
+    dma_ptr dptr  = { NULL };
+
+    if (r->push_idx == r->pop_idx) {
+        return dptr;
+    }
+
+    dptr = r->dptr[r->pop_idx];
+
+    r->pop_idx = (r->pop_idx + 1) & r->idx_mask;
+    return dptr;
+}
+
+static inline bool dma_queue_empty(dma_queue * q) {
+    return q->ring->push_idx == q->ring->pop_idx;
+}
+
+static inline uint32_t dma_queue_depth(dma_queue * q) {
+    return (q->ring->push_idx - q->ring->pop_idx) & q->ring->idx_mask;
+}
+
+static inline uint32_t dma_queue_capacity(dma_queue * q) {
+    return q->ring->capacity;
+}
+
+#if __HVX_ARCH__ < 75
+
+// Overflow-safe DMA push: all 2d descriptor fields (row_size, nrows, src_stride, dst_stride) are 16-bit, max 65535.
+// This version transparently handles values that exceed the 16-bit limit and submits chained DMA transtions.
+
+#define DMA_MAX_FIELD_VAL 65535u
+
+static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, 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(
+            row_size   <= DMA_MAX_FIELD_VAL &&
+            nrows      <= DMA_MAX_FIELD_VAL &&
+            src_stride <= DMA_MAX_FIELD_VAL &&
+            dst_stride <= DMA_MAX_FIELD_VAL, 1)) {
+        return dma_queue_push_single_2d(q, dptr, dst_stride, src_stride, row_size, nrows);
+    }
+
+    // Contiguous block
+    // Use 1d DMA mode which supports sizes up to 24-bits (16MB)
+    if (nrows == 1 || (row_size == src_stride && row_size == dst_stride)) {
+        size_t total = row_size * nrows;
+        return dma_queue_push_single_1d(q, dptr, total);
+    }
+
+    // Stride overflow - fall back to row-by-row.
+    {
+        const uint8_t *src = (const uint8_t *) dptr.src;
+        uint8_t       *dst = (uint8_t *)       dptr.dst;
+        for (size_t r = 0; r < nrows; ++r) {
+            dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
+            if (!dma_queue_push_single_1d(q, p, row_size))
+                return false;
+            if (r + 1 < nrows)
+                dma_queue_pop(q);
+        }
+        return true;
+    }
+}
+
+#else // HVX_ARCH >= 75
+
+static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
+    // On v75 and up we always use 2d 24-bit mode
+    return dma_queue_push_single_2d(q, dptr, dst_stride, src_stride, row_size, nrows);
+}
+
+#endif
+
+static inline bool dma_queue_push_ddr_to_vtcm(dma_queue * q, dma_ptr dptr, size_t dst_row_size, size_t src_row_size, size_t nrows) {
+    return dma_queue_push(q, dptr, dst_row_size, src_row_size, src_row_size, nrows);
+}
+
+static inline bool dma_queue_push_vtcm_to_ddr(dma_queue * q, dma_ptr dptr, size_t dst_row_size, size_t src_row_size, size_t nrows) {
+    return dma_queue_push(q, dptr, dst_row_size, src_row_size, dst_row_size, nrows);
+}
+
+#define DMA_CACHE_MAX_SIZE 256U
+
+typedef struct {
+    uint8_t *base;
+    uint32_t line_size;
+    uint32_t capacity;
+    uint32_t src[DMA_CACHE_MAX_SIZE];
+    uint16_t age[DMA_CACHE_MAX_SIZE];
+} dma_cache;
+
+static inline void dma_cache_init(dma_cache *c, uint8_t *base, uint32_t line_size, uint32_t capacity)
+{
+    c->capacity  = (capacity > DMA_CACHE_MAX_SIZE) ? DMA_CACHE_MAX_SIZE : capacity;
+    c->base      = base;
+    c->line_size = line_size;
+
+    for (unsigned i=0; i < c->capacity; i++) {
+        c->src[i] = 0;
+        c->age[i] = 0;
+    }
+}
+
+static inline bool dma_cache_push(dma_queue *q, dma_cache *c, const uint8_t * src, uint32_t dst_stride, uint32_t src_stride, uint32_t row_size, uint32_t nrows)
+{
+    uint32_t o_idx = 0;
+    uint16_t o_age = 0;
+    uint8_t *  dst = 0;
+
+    for (unsigned i=0; i < c->capacity; i++) {
+        if (c->src[i] == (uint32_t) src) {
+            c->age[i] = 0;
+            dst = c->base + (i * c->line_size); nrows = 0; // dummy dma
+        } else {
+            c->age[i]++;
+            if (c->age[i] > o_age) { o_age = c->age[i]; o_idx = i; }
+        }
+    }
+    if (!dst) {
+        c->age[o_idx] = 0;
+        c->src[o_idx] = (uint32_t) src;
+        dst = c->base + o_idx * c->line_size; // normal nrows dma
+        return dma_queue_push(q, dma_make_ptr(dst, src), dst_stride, src_stride, row_size, nrows);
+    }
+
+    return dma_queue_push_single_1d(q, dma_make_ptr(dst, src), 0);
+}
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif
+
+#endif /* HTP_DMA_H */

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

@@ -24,7 +24,7 @@
 #include "hvx-reduce.h"
 #include "hvx-flash-attn.h"
 #include "htp-vtcm.h"
-#include "worker-pool.h"
+#include "work-queue.h"
 
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
@@ -204,7 +204,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
 
     if (ir0 >= ir1) return;
 
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
     dma_queue * dma = octx->ctx->dma[ith];
 
@@ -486,7 +486,7 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
         return;
     }
 
-    struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
     hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles, (const __fp16 *) args->curr_k, total_rows, factx->DK,
                              args->src_stride, start, end);
@@ -494,7 +494,7 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
 }
 
 static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, void * curr_k, uint32_t kv_start) {
-    worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+    work_queue_t wp = factx->octx->ctx->work_queue;
     uint32_t n = 1;
     if (factx->n_threads > 1 && kv_rows >= factx->n_threads * 2) {
         n = factx->n_threads;
@@ -502,7 +502,7 @@ static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_row
     uint32_t rows_per_t = hex_align_up(hmx_ceil_div(kv_rows, n), 2);
     fa_k_int_args_t args = { factx, kv_rows, src_stride, curr_k, kv_start, rows_per_t };
     if (n > 1) {
-        worker_pool_run_func(wp, fa_k_interleave_thread, &args, n);
+        work_queue_run(wp, fa_k_interleave_thread, &args, n);
     } else {
         fa_k_interleave_thread(1, 0, &args);
     }
@@ -534,7 +534,7 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data)
 
     __fp16 * v_tiles_dst = (__fp16 *) args->v_tiles_dst;
 
-    struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
     hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV,
                              args->src_stride, (uint32_t) args->n_col_tiles, start, end);
@@ -548,7 +548,7 @@ static void fa_phase_v_interleave(struct hmx_fa_context * factx,
                                   void *                  v_tiles_dst,
                                   size_t                  n_col_tiles,
                                   uint32_t                kv_start) {
-    worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+    work_queue_t wp = factx->octx->ctx->work_queue;
     uint32_t n = 1;
     if (factx->n_threads > 1 && kv_rows >= factx->n_threads * 2) {
         n = factx->n_threads;
@@ -556,7 +556,7 @@ static void fa_phase_v_interleave(struct hmx_fa_context * factx,
     uint32_t rows_per_t = hex_align_up(hmx_ceil_div(kv_rows, n), 2);
     fa_v_int_args_t args = { factx, kv_rows, src_stride, v_src, v_tiles_dst, n_col_tiles, kv_start, rows_per_t };
     if (n > 1) {
-        worker_pool_run_func(wp, fa_v_interleave_thread, &args, n);
+        work_queue_run(wp, fa_v_interleave_thread, &args, n);
     } else {
         fa_v_interleave_thread(1, 0, &args);
     }
@@ -589,7 +589,7 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
     const size_t start      = (size_t) i * rows_per_t;
     const size_t end        = hex_smin(start + rows_per_t, factx->g_br);
 
-    struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_Q_PREP, (uint16_t) (args->q_start * G + start));
 
     // Parallel initialization of per-block state
@@ -720,7 +720,7 @@ static void fa_phase_q_load(struct hmx_fa_context *   factx,
                             uint32_t                  kv_head,
                             uint32_t                  ib3,
                             size_t                    n_rows_g) {
-    worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+    work_queue_t wp = factx->octx->ctx->work_queue;
     uint32_t n = 1;
     if (factx->n_threads > 1 && n_rows_g >= (size_t) (factx->n_threads * 2)) {
         n = factx->n_threads;
@@ -739,7 +739,7 @@ static void fa_phase_q_load(struct hmx_fa_context *   factx,
     args.q_transposed = q->nb[1] < q->nb[2];
     atomic_init(&args.barrier, n);
     if (n > 1) {
-        worker_pool_run_func(wp, fa_q_load_thread, &args, n);
+        work_queue_run(wp, fa_q_load_thread, &args, n);
     } else {
         fa_q_load_thread(1, 0, &args);
     }
@@ -772,7 +772,7 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
         return;
     }
 
-    struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
 
     const struct htp_tensor * dst        = args->dst;
@@ -820,7 +820,7 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
         return;
     }
 
-    struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) (args->q_start * G + start));
 
     const struct htp_tensor * dst        = args->dst;
@@ -862,7 +862,7 @@ static void fa_phase_o_store(struct hmx_fa_context *   factx,
                              uint32_t                  kv_head,
                              uint32_t                  ib3,
                              size_t                    n_rows_g) {
-    worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+    work_queue_t wp = factx->octx->ctx->work_queue;
     uint32_t n = 1;
     if (factx->n_threads > 1 && n_rows_g >= (size_t) (factx->n_threads * 2)) {
         n = factx->n_threads;
@@ -871,7 +871,7 @@ static void fa_phase_o_store(struct hmx_fa_context *   factx,
     fa_o_store_args_t args = { factx, dst, o_tile_src, q_start, kv_head, ib3, n_rows_g, rows_per_t };
     worker_callback_t store_fn = factx->is_dst_fp32 ? fa_o_store_thread_f32 : fa_o_store_thread_f16;
     if (n > 1) {
-        worker_pool_run_func(wp, store_fn, &args, n);
+        work_queue_run(wp, store_fn, &args, n);
     } else {
         store_fn(1, 0, &args);
     }
@@ -930,7 +930,7 @@ static inline void fa_softmax_impl(
         return;
     }
 
-    struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
+    struct htp_thread_trace * tr = &factx->octx->ctx->trace[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_SFM, (uint16_t) (args->q_start * G + vec_start * 64));
 
     // Per-thread row scratch: thread i uses bufs at offset i * 2 * stride
@@ -1290,7 +1290,7 @@ static void fa_phase_softmax_and_build_d(struct hmx_fa_context * factx,
                                          fa_softmax_args_t *     sargs,
                                          size_t                  n_row_tiles,
                                          size_t                  n_row_tiles_g_br) {
-    worker_pool_context_t wp = factx->octx->ctx->worker_pool;
+    work_queue_t wp = factx->octx->ctx->work_queue;
     const size_t n_row_vec_cnt = hmx_ceil_div(sargs->n_rows_g, 64);
 
     worker_callback_t softmax_fn = fa_softmax_thread;
@@ -1307,7 +1307,7 @@ static void fa_phase_softmax_and_build_d(struct hmx_fa_context * factx,
     if (factx->n_threads > 1 && n_row_vec_cnt >= 2) {
         uint32_t n_use = (uint32_t) hex_smin((size_t) factx->n_threads, n_row_vec_cnt);
         sargs->thread_div = init_fastdiv_values(n_use);
-        worker_pool_run_func(wp, softmax_fn, sargs, n_use);
+        work_queue_run(wp, softmax_fn, sargs, n_use);
     } else {
         softmax_fn(1, 0, sargs);
     }
@@ -1519,8 +1519,8 @@ static void fa_pop_mask_dma_gqa(dma_queue * dma, uint32_t G) {
 // ============================================================================
 
 int hmx_flash_attn_ext(struct htp_ops_context * octx) {
-    struct htp_thread_trace * tr_hvx = octx->ctx ? &octx->ctx->trace[0] : NULL;
-    struct htp_thread_trace * tr_hmx = octx->ctx ? &octx->ctx->trace[HTP_MAX_NTHREADS] : NULL;
+    struct htp_thread_trace * tr_hvx = &octx->ctx->trace[0];
+    struct htp_thread_trace * tr_hmx = &octx->ctx->trace[HTP_MAX_NTHREADS];
     const struct htp_tensor * q    = octx->src[0];
     const struct htp_tensor * k    = octx->src[1];
     const struct htp_tensor * v    = octx->src[2];
@@ -1735,7 +1735,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                 const size_t k_src_stride = size_k_row_padded / sizeof(__fp16);
                 const size_t v_src_stride = size_v_row_padded / sizeof(__fp16);
 
-                struct hmx_queue * hmx_q = ctx->hmx_queue;
+                hmx_queue_t hmx_q = ctx->hmx_queue;
 
                 if (factx.pipeline) {
                     // Pipeline path
@@ -2084,7 +2084,7 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     }
 
     if (!(octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
-        worker_pool_run_func(octx->ctx->worker_pool, flash_attn_ext_f16_thread, &factx, octx->n_threads);
+        work_queue_run(octx->ctx->work_queue, flash_attn_ext_f16_thread, &factx, octx->n_threads);
     }
 
     return HTP_STATUS_OK;

ggml/src/ggml-hexagon/htp/hex-bitmap.h

@@ -0,0 +1,24 @@
+#ifndef HEX_BITMAP_H
+#define HEX_BITMAP_H
+
+#include <stdint.h>
+#include <stdbool.h>
+#include <string.h>
+
+static inline void bitmap_set(uint32_t * bitmap, uint32_t idx) {
+    bitmap[idx / 32] |= (1U << (idx % 32));
+}
+
+static inline void bitmap_clear(uint32_t * bitmap, uint32_t idx) {
+    bitmap[idx / 32] &= ~(1U << (idx % 32));
+}
+
+static inline bool bitmap_test(const uint32_t * bitmap, uint32_t idx) {
+    return (bitmap[idx / 32] & (1U << (idx % 32))) != 0;
+}
+
+static inline void bitmap_reset(uint32_t * bitmap, size_t size_in_bits) {
+    memset(bitmap, 0, ((size_in_bits + 31) / 32) * sizeof(uint32_t));
+}
+
+#endif // HEX_BITMAP_H

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

@@ -1,63 +0,0 @@
-#include "hex-dma.h"
-
-#include <stdbool.h>
-#include <stdlib.h>
-#include <string.h>
-
-#pragma clang diagnostic ignored "-Wunused-function"
-
-static inline uint32_t pow2_ceil(uint32_t x) {
-    if (x <= 1) {
-        return 1;
-    }
-    int p = 2;
-    x--;
-    while (x >>= 1) {
-        p <<= 1;
-    }
-    return p;
-}
-
-dma_queue * dma_queue_create(size_t capacity) {
-    dma_queue * q = (dma_queue *) memalign(32, sizeof(dma_queue));
-    if (q == NULL) {
-        FARF(ERROR, "%s: failed to allocate DMA queue\n", __FUNCTION__);
-        return NULL;
-    }
-
-    capacity = pow2_ceil(capacity);
-
-    memset(q, 0, sizeof(dma_queue));
-    q->capacity = capacity;
-    q->idx_mask = capacity - 1;
-
-    q->desc = (dma_descriptor_2d *) memalign(64, capacity * sizeof(dma_descriptor_2d));
-    memset(q->desc, 0, capacity * sizeof(dma_descriptor_2d));
-
-    q->dptr = (dma_ptr *) memalign(4, capacity * sizeof(dma_ptr));
-    memset(q->dptr, 0, capacity * sizeof(dma_ptr));
-
-    q->tail = &q->desc[capacity - 1];
-
-    if (!q->desc && !q->dptr) {
-        FARF(ERROR, "%s: failed to allocate DMA queue items\n", __FUNCTION__);
-        return NULL;
-    }
-
-    FARF(HIGH, "dma-queue: capacity %u\n", capacity);
-
-    return q;
-}
-
-void dma_queue_delete(dma_queue * q) {
-    if (!q) {
-        return;
-    }
-    free(q->desc);
-    free(q->dptr);
-    free(q);
-}
-
-void dma_queue_flush(dma_queue * q) {
-    while (dma_queue_pop(q).dst != NULL) ;
-}

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

@@ -1,375 +1,2 @@
-#ifndef HTP_DMA_H
-#define HTP_DMA_H
-
-#include <HAP_farf.h>
-#include <hexagon_types.h>
-#include <stdbool.h>
-#include <stdint.h>
-#include "hex-utils.h"
-
-#include "hex-profile.h"
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-// Define the HW descriptor structs here since the ones in HexSDK are a bit out of date
-typedef struct dma_descriptor_1d_s {
-    void *   next;
-    uint32_t size:24;
-    uint32_t desc_size:2;
-    uint32_t dst_comp:1;
-    uint32_t src_comp:1;
-    uint32_t dst_bypass:1;
-    uint32_t src_bypass:1;
-    uint32_t order:1;
-    uint32_t done:1;
-    void *   src;
-    void *   dst;
-} dma_descriptor_1d;
-
-#if __HVX_ARCH__ < 75
-
-typedef struct dma_descriptor_2d_s {
-    void *   next;
-    uint32_t reserved0:24;
-    uint32_t desc_size:2;
-    uint32_t dst_comp:1;
-    uint32_t src_comp:1;
-    uint32_t dst_bypass:1;
-    uint32_t src_bypass:1;
-    uint32_t order:1;
-    uint32_t done:1;
-    void *   src;
-    void *   dst;
-    uint32_t desc_type:8;
-    uint32_t reserved1:24;
-    uint32_t row_size:16;
-    uint32_t nrows:16;
-    uint32_t src_stride:16;
-    uint32_t dst_stride:16;
-    uint32_t src_offset:16;
-    uint32_t dst_offset:16;
-} dma_descriptor_2d;
-
-#else
-
-typedef struct dma_descriptor_2d_s {
-    void *   next;
-    uint32_t dst_stride:24;
-    uint32_t desc_size:2;
-    uint32_t dst_comp:1;
-    uint32_t src_comp:1;
-    uint32_t dst_bypass:1;
-    uint32_t src_bypass:1;
-    uint32_t order:1;
-    uint32_t done:1;
-    void *   src;
-    void *   dst;
-    uint32_t desc_type:8;
-    uint32_t reserved0:24;
-    uint32_t row_size:24;
-    uint32_t nrows_lo:8;
-    uint32_t nrows_hi:8;
-    uint32_t src_stride:24;
-    uint32_t offset:24;
-    uint32_t reserved1:8;
-} dma_descriptor_2d;
-
-#endif
-
-typedef struct {
-    void       *dst;
-    const void *src;
-} dma_ptr;
-
-typedef struct {
-    dma_descriptor_2d * desc;  // descriptor pointers
-    dma_descriptor_2d * tail;  // tail pointer
-    dma_ptr           * dptr;  // dst/src pointers
-    uint32_t            push_idx;
-    uint32_t            pop_idx;
-    uint32_t            capacity;
-    uint32_t            idx_mask;
-    struct htp_thread_trace * trace;
-} dma_queue;
-
-dma_queue * dma_queue_create(size_t capacity);
-void        dma_queue_delete(dma_queue * q);
-void        dma_queue_flush(dma_queue * q);
-
-// TODO: technically we don't need these and could use Q6_dmstart/wait/etc instead
-// but those do not seem to always compiler properly.
-static inline void dmstart(void * next) {
-    asm volatile(" release(%0):at" : : "r"(next));
-    asm volatile(" dmstart(%0)" : : "r"(next));
-}
-
-static inline void dmlink(void * cur, void * next) {
-    asm volatile(" release(%0):at" : : "r"(next));
-    asm volatile(" dmlink(%0, %1)" : : "r"(cur), "r"(next));
-}
-
-static inline unsigned int dmpoll(void) {
-    unsigned int ret = 0;
-    asm volatile(" %0 = dmpoll" : "=r"(ret) : : "memory");
-    return ret;
-}
-
-static inline unsigned int dmwait(void) {
-    unsigned int ret = 0;
-    asm volatile(" %0 = dmwait" : "=r"(ret) : : "memory");
-    return ret;
-}
-
-static inline dma_ptr dma_make_ptr(void *dst, const void *src)
-{
-    dma_ptr p = { dst, src };
-    return p;
-}
-
-static const uint32_t dma_src_l2_bypass_on = 1;
-static const uint32_t dma_dst_l2_bypass_on = 1;
-
-static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t size) {
-    if (((q->push_idx + 1) & q->idx_mask) == q->pop_idx) {
-        FARF(HIGH, "dma-push: queue full\n");
-        return false;
-    }
-
-    dma_descriptor_1d * desc = (dma_descriptor_1d *) &q->desc[q->push_idx];
-    desc->src  = (void *) dptr.src;
-    desc->dst  = (void *) dptr.dst;
-    desc->size = size;
-
-    q->dptr[q->push_idx] = dptr;
-
-    if (size) {
-        desc->next       = NULL;
-        desc->desc_size  = 0; // 1D mode
-        desc->src_bypass = dma_src_l2_bypass_on;
-        desc->dst_bypass = dma_dst_l2_bypass_on;
-        desc->order      = 0;
-        desc->done       = 0;
-
-        htp_trace_event_start(q->trace, HTP_TRACE_EVT_DMA, q->push_idx);
-        dmlink(q->tail, desc);
-        q->tail = (dma_descriptor_2d *) desc;
-    } else {
-        desc->desc_size = 0;
-        desc->done      = 1;
-    }
-
-    q->push_idx = (q->push_idx + 1) & q->idx_mask;
-    return true;
-}
-
-static inline bool dma_queue_push_single_2d(dma_queue * q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
-    if (((q->push_idx + 1) & q->idx_mask) == q->pop_idx) {
-        FARF(HIGH, "dma-push: queue full\n");
-        return false;
-    }
-
-    dma_descriptor_2d * desc = &q->desc[q->push_idx];
-
-    desc->next           = NULL;
-    desc->reserved0      = 0;
-    desc->reserved1      = 0;
-    desc->desc_size      = 1; // 2d mode
-    desc->src_bypass     = dma_src_l2_bypass_on;
-    desc->dst_bypass     = dma_dst_l2_bypass_on;
-    desc->src_comp       = 0;
-    desc->dst_comp       = 0;
-    desc->order          = 0;
-    desc->done           = 0;
-    desc->src_stride     = src_stride;
-    desc->dst_stride     = dst_stride;
-    desc->src            = (void *) dptr.src;
-    desc->dst            = (void *) dptr.dst;
-    desc->row_size       = row_size;
-
-#if __HVX_ARCH__ < 75
-    desc->desc_type      = 0; // 2d (16-bit) mode
-    desc->nrows          = nrows;
-    desc->src_offset     = 0;
-    desc->dst_offset     = 0;
-#else
-    desc->desc_type      = 9; // 2d (24-bit) mode
-    desc->nrows_lo       = (nrows & 0xff);
-    desc->nrows_hi       = (nrows >> 8);
-    desc->offset         = 0;
-#endif
-
-    q->dptr[q->push_idx] = dptr;
-
-    if (nrows) {
-        htp_trace_event_start(q->trace, HTP_TRACE_EVT_DMA, q->push_idx);
-        dmlink(q->tail, desc);
-        q->tail = desc;
-    } else {
-        desc->done = 1;
-    }
-
-    // FARF(ERROR, "dma-push: i %u row-size %u nrows %d dst %p src %p\n", q->push_idx, row_size, nrows, dptr.dst, dptr.src);
-    q->push_idx = (q->push_idx + 1) & q->idx_mask;
-    return true;
-}
-
-static inline dma_ptr dma_queue_pop(dma_queue * q) {
-    dma_ptr dptr  = { NULL };
-
-    if (q->push_idx == q->pop_idx) {
-        return dptr;
-    }
-
-    dma_descriptor_2d * desc = &q->desc[q->pop_idx];
-
-    // Wait for desc to complete
-    if (!desc->done) {
-        while (!desc->done) {
-            dmpoll();
-        }
-    }
-    htp_trace_event_stop(q->trace, HTP_TRACE_EVT_DMA, q->pop_idx);
-
-    dptr = q->dptr[q->pop_idx];
-
-    // FARF(ERROR, "dma-pop: i %u dst %p src %p\n", q->pop_idx, dptr.dst, dptr.src);
-    q->pop_idx = (q->pop_idx + 1) & q->idx_mask;
-    return dptr;
-}
-
-static inline dma_ptr dma_queue_pop_nowait(dma_queue * q) {
-    dma_ptr dptr  = { NULL };
-
-    if (q->push_idx == q->pop_idx) {
-        return dptr;
-    }
-
-    dptr = q->dptr[q->pop_idx];
-
-    // FARF(ERROR, "dma-pop-nowait: i %u dst %p src %p\n", q->pop_idx, dptr.dst, dptr.src);
-    q->pop_idx = (q->pop_idx + 1) & q->idx_mask;
-    return dptr;
-}
-
-static inline bool dma_queue_empty(dma_queue * q) {
-    return q->push_idx == q->pop_idx;
-}
-
-static inline uint32_t dma_queue_depth(dma_queue * q) {
-    return (q->push_idx - q->pop_idx) & q->idx_mask;
-}
-
-static inline uint32_t dma_queue_capacity(dma_queue * q) {
-    return q->capacity;
-}
-
-#if __HVX_ARCH__ < 75
-
-// Overflow-safe DMA push: all 2d descriptor fields (row_size, nrows, src_stride, dst_stride) are 16-bit, max 65535.
-// This version transparently handles values that exceed the 16-bit limit and submits chained DMA transtions.
-
-#define DMA_MAX_FIELD_VAL 65535u
-
-static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, 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(
-            row_size   <= DMA_MAX_FIELD_VAL &&
-            nrows      <= DMA_MAX_FIELD_VAL &&
-            src_stride <= DMA_MAX_FIELD_VAL &&
-            dst_stride <= DMA_MAX_FIELD_VAL, 1)) {
-        return dma_queue_push_single_2d(q, dptr, dst_stride, src_stride, row_size, nrows);
-    }
-
-    // Contiguous block
-    // Use 1d DMA mode which supports sizes up to 24-bits (16MB)
-    if (nrows == 1 || (row_size == src_stride && row_size == dst_stride)) {
-        size_t total = row_size * nrows;
-        return dma_queue_push_single_1d(q, dptr, total);
-    }
-
-    // Stride overflow — fall back to row-by-row.
-    {
-        const uint8_t *src = (const uint8_t *) dptr.src;
-        uint8_t       *dst = (uint8_t *)       dptr.dst;
-        for (size_t r = 0; r < nrows; ++r) {
-            dma_ptr p = dma_make_ptr(dst + r * dst_stride, src + r * src_stride);
-            if (!dma_queue_push_single_1d(q, p, row_size))
-                return false;
-            if (r + 1 < nrows)
-                dma_queue_pop(q);
-        }
-        return true;
-    }
-}
-
-#else // HVX_ARCH >= 75
-
-static inline bool dma_queue_push(dma_queue *q, dma_ptr dptr, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
-    // On v75 and up we always use 2d 24-bit mode
-    return dma_queue_push_single_2d(q, dptr, dst_stride, src_stride, row_size, nrows);
-}
-
-#endif
-
-static inline bool dma_queue_push_ddr_to_vtcm(dma_queue * q, dma_ptr dptr, size_t dst_row_size, size_t src_row_size, size_t nrows) {
-    return dma_queue_push(q, dptr, dst_row_size, src_row_size, src_row_size, nrows);
-}
-
-static inline bool dma_queue_push_vtcm_to_ddr(dma_queue * q, dma_ptr dptr, size_t dst_row_size, size_t src_row_size, size_t nrows) {
-    return dma_queue_push(q, dptr, dst_row_size, src_row_size, dst_row_size, nrows);
-}
-
-#define DMA_CACHE_MAX_SIZE 256U
-
-typedef struct {
-    uint8_t *base;
-    uint32_t line_size;
-    uint32_t capacity;
-    uint32_t src[DMA_CACHE_MAX_SIZE];
-    uint16_t age[DMA_CACHE_MAX_SIZE];
-} dma_cache;
-
-static inline void dma_cache_init(dma_cache *c, uint8_t *base, uint32_t line_size, uint32_t capacity)
-{
-    c->capacity  = (capacity > DMA_CACHE_MAX_SIZE) ? DMA_CACHE_MAX_SIZE : capacity;
-    c->base      = base;
-    c->line_size = line_size;
-
-    for (unsigned i=0; i < c->capacity; i++) {
-        c->src[i] = 0;
-        c->age[i] = 0;
-    }
-}
-
-static inline bool dma_cache_push(dma_queue *q, dma_cache *c, const uint8_t * src, uint32_t dst_stride, uint32_t src_stride, uint32_t row_size, uint32_t nrows)
-{
-    uint32_t o_idx = 0;
-    uint16_t o_age = 0;
-    uint8_t *  dst = 0;
-
-    for (unsigned i=0; i < c->capacity; i++) {
-        if (c->src[i] == (uint32_t) src) {
-            c->age[i] = 0;
-            dst = c->base + (i * c->line_size); nrows = 0; // dummy dma
-        } else {
-            c->age[i]++;
-            if (c->age[i] > o_age) { o_age = c->age[i]; o_idx = i; }
-        }
-    }
-    if (!dst) {
-        c->age[o_idx] = 0;
-        c->src[o_idx] = (uint32_t) src;
-        dst = c->base + o_idx * c->line_size; // normal nrows dma
-        return dma_queue_push(q, dma_make_ptr(dst, src), dst_stride, src_stride, row_size, nrows);
-    }
-
-    return dma_queue_push_single_1d(q, dma_make_ptr(dst, src), 0);
-}
-
-#ifdef __cplusplus
-}  // extern "C"
-#endif
-
-#endif /* HTP_DMA_H */
+#pragma once
+#include "dma-queue.h"

ggml/src/ggml-hexagon/htp/hex-profile.h

@@ -44,11 +44,11 @@ struct htp_thread_trace {
 };
 
 static inline void htp_trace_event(struct htp_thread_trace * tr, uint16_t id, uint16_t info, uint32_t type) {
-    if (tr && tr->events && tr->count < tr->max_events) {
-        uint32_t idx = tr->count;
-        tr->events[idx].id = id;
-        tr->events[idx].info = info | (type == HTP_TRACE_EVT_STOP ? 0x8000 : 0);
-        tr->events[idx].cycles = (uint32_t) hex_get_cycles();
+    if (tr->count < tr->max_events) {
+        uint32_t i = tr->count;
+        tr->events[i].id     = id;
+        tr->events[i].info   = info | (type == HTP_TRACE_EVT_STOP ? 0x8000 : 0);
+        tr->events[i].cycles = (uint32_t) hex_get_cycles();
         tr->count++;
     }
 }

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

@@ -30,21 +30,26 @@ static inline void hex_l2fetch(const void * p, uint32_t width, uint32_t stride,
     Q6_l2fetch_AP((void *) p, control);
 }
 
-#define HEX_L2_LINE_SIZE  64
-#define HEX_L2_FLUSH_SIZE (128 * 1024)
+static inline void hex_l2fetch_block(const void * addr, size_t size) {
+    if (size == 0) return;
+    const uint32_t width = 16384; // 16KB rows
+    const uint32_t height = (size + width - 1) / width;
+    hex_l2fetch(addr, width, width, height);
+}
+
+#define HEX_L2_LINE_SIZE           128
+#define HEX_L2_BLOCK_SIZE          (HEX_L2_LINE_SIZE * 4) // flush granularity (lines per loop iteration)
+#define HEX_L2_FLUSH_WQ_THRESHOLD  (4 * 1024)
+#define HEX_L2_FLUSH_ALL_THRESHOLD (4 * 1024 * 1024)
 
 static inline void hex_l2flush(void * addr, size_t size) {
-    if (size > HEX_L2_FLUSH_SIZE) {
-        qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
-    } else {
-        const uint32_t s = (uint32_t) addr;
-        const uint32_t e = s + size;
-        for (uint32_t i = s; i < e; i += HEX_L2_LINE_SIZE * 4) {
-            Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 0);
-            Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 1);
-            Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 2);
-            Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 3);
-        }
+    const uint32_t s = ((uint32_t) addr) & ~(HEX_L2_LINE_SIZE - 1);
+    const uint32_t e = (((uint32_t) addr) + size + HEX_L2_LINE_SIZE - 1) & ~(HEX_L2_LINE_SIZE - 1);
+    for (uint32_t i = s; i < e; i += HEX_L2_BLOCK_SIZE) {
+        Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 0);
+        Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 1);
+        Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 2);
+        Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 3);
     }
 }
 

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

@@ -1005,10 +1005,62 @@ static void transfer_activation_row_pair_fp32_to_fp16(
     }
 }
 
+static void transfer_activation_row_pair_fp32_to_fp16_col_chunk(
+        __fp16 *restrict vtcm_dst,
+        const float *restrict row0, // offset by c_first
+        const float *restrict row1, // offset by c_first
+        uint32_t r,
+        uint32_t k_block,
+        uint32_t c_first,
+        uint32_t c_len,
+        uint32_t k_chunk_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
+
+    uint32_t c = 0;
+    for (; c + 32 <= k_chunk_valid; c += 32) {
+        HVX_Vector v0 = Q6_V_vzero();
+        HVX_Vector v1 = Q6_V_vzero();
+        if (row0_valid) v0 = *(const HVX_Vector *)(row0 + c);
+        if (row1_valid) v1 = *(const HVX_Vector *)(row1 + c);
+
+        HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
+
+        uint32_t c0       = (c_first + c) / HTP_MM_HMX_TILE_N_COLS;  // tile column index
+        uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0;
+
+        HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS);
+        tile[r1 / 2]     = v_out;
+    }
+    if (c < c_len) {
+        HVX_Vector v0 = Q6_V_vzero();
+        HVX_Vector v1 = Q6_V_vzero();
+        if (row0_valid) v0 = *(const HVX_Vector *)(row0 + c);
+        if (row1_valid) v1 = *(const HVX_Vector *)(row1 + c);
+
+        uint32_t rem = (k_chunk_valid > c) ? (k_chunk_valid - c) : 0;
+        HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0);
+        v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero());
+        v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero());
+
+        HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
+
+        uint32_t c0       = (c_first + c) / HTP_MM_HMX_TILE_N_COLS;  // tile column index
+        uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0;
+
+        HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS);
+        tile[r1 / 2]     = v_out;
+    }
+}
+
 static void transfer_activation_chunk_fp32_to_fp16_gathered(
             __fp16 *restrict vtcm_dst,
             const float *restrict src,
             uint32_t start_row,
+            uint32_t vtcm_start_row,
             uint32_t n_rows,
             uint32_t k_block,
             const struct mmid_row_mapping *matrix_rows,
@@ -1029,8 +1081,9 @@ static void transfer_activation_chunk_fp32_to_fp16_gathered(
     for (r = 0; r < n_rows_tiled; r += 2) {
         uint32_t r_idx0 = start_row + r + 0;
         uint32_t r_idx1 = start_row + r + 1;
-        uint32_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS;  // tile row index
-        uint32_t r1 = r_idx0 % HTP_MM_HMX_TILE_N_ROWS;  // intra-tile row idx
+        uint32_t lr = vtcm_start_row + r;           // vtcm-local row
+        uint32_t r0 = lr / HTP_MM_HMX_TILE_N_ROWS;  // tile row index
+        uint32_t r1 = lr % HTP_MM_HMX_TILE_N_ROWS;  // intra-tile row idx
 
         struct mmid_row_mapping mapping0 = matrix_rows[cur_a * mapping_stride + r_idx0];
         struct mmid_row_mapping mapping1 = matrix_rows[cur_a * mapping_stride + r_idx1];
@@ -1073,9 +1126,9 @@ static void transfer_activation_chunk_fp32_to_fp16_gathered(
     }
 
     for (; r < n_rows_padded; r += 2) {
-        uint32_t r_idx0 = start_row + r;
-        uint32_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS;  // tile row index
-        uint32_t r1 = r_idx0 % HTP_MM_HMX_TILE_N_ROWS;  // intra-tile row idx
+        uint32_t lr = vtcm_start_row + r;           // vtcm-local row
+        uint32_t r0 = lr / HTP_MM_HMX_TILE_N_ROWS;  // tile row index
+        uint32_t r1 = lr % HTP_MM_HMX_TILE_N_ROWS;  // intra-tile row idx
 
         const bool row0_valid = (start_row + r + 0) < cne1;
         const bool row1_valid = (start_row + r + 1) < cne1;
@@ -1135,6 +1188,7 @@ static void transfer_activation_chunk_fp32_to_fp16_gathered_flat(
             __fp16 *restrict vtcm_dst,
             const float *restrict src,
             uint32_t start_row,
+            uint32_t vtcm_start_row,
             uint32_t n_rows,
             uint32_t k_block,
             const struct mmid_row_mapping *matrix_rows,
@@ -1152,8 +1206,9 @@ static void transfer_activation_chunk_fp32_to_fp16_gathered_flat(
     for (r = 0; r < n_rows_tiled; r += 2) {
         uint32_t r_idx0 = start_row + r + 0;
         uint32_t r_idx1 = start_row + r + 1;
-        uint32_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS;  // tile row index
-        uint32_t r1 = r_idx0 % HTP_MM_HMX_TILE_N_ROWS;  // intra-tile row idx
+        uint32_t lr = vtcm_start_row + r;           // vtcm-local row
+        uint32_t r0 = lr / HTP_MM_HMX_TILE_N_ROWS;  // tile row index
+        uint32_t r1 = lr % HTP_MM_HMX_TILE_N_ROWS;  // intra-tile row idx
 
         struct mmid_row_mapping mapping0 = matrix_rows[cur_a * mapping_stride + r_idx0];
         struct mmid_row_mapping mapping1 = matrix_rows[cur_a * mapping_stride + r_idx1];
@@ -1193,9 +1248,9 @@ static void transfer_activation_chunk_fp32_to_fp16_gathered_flat(
     }
 
     for (; r < n_rows_padded; r += 2) {
-        uint32_t r_idx0 = start_row + r;
-        uint32_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS;  // tile row index
-        uint32_t r1 = r_idx0 % HTP_MM_HMX_TILE_N_ROWS;  // intra-tile row idx
+        uint32_t lr = vtcm_start_row + r;           // vtcm-local row
+        uint32_t r0 = lr / HTP_MM_HMX_TILE_N_ROWS;  // tile row index
+        uint32_t r1 = lr % HTP_MM_HMX_TILE_N_ROWS;  // intra-tile row idx
 
         const bool row0_valid = (start_row + r + 0) < cne1;
         const bool row1_valid = (start_row + r + 1) < cne1;
@@ -1253,6 +1308,7 @@ static void transfer_output_chunk_fp16_to_fp32_scattered(
             float *restrict dst,
             const __fp16 *restrict vtcm_src,
             uint32_t start_row,
+            uint32_t vtcm_start_row,
             uint32_t n_rows,
             uint32_t n_cols,
             const struct mmid_row_mapping *matrix_rows,
@@ -1269,8 +1325,9 @@ static void transfer_output_chunk_fp16_to_fp32_scattered(
     for (size_t r = 0; r < n_rows; r += 2) {
         uint32_t r_idx0 = start_row + r + 0;
         uint32_t r_idx1 = start_row + r + 1;
-        const size_t r0 = r_idx0 / HTP_MM_HMX_TILE_N_ROWS;
-        const size_t r1 = (r_idx0 % HTP_MM_HMX_TILE_N_ROWS) / 2;  // index of the row pair within the tile
+        uint32_t     lr = vtcm_start_row + r;                 // vtcm-local row
+        const size_t r0 = (lr / HTP_MM_HMX_TILE_N_ROWS);
+        const size_t r1 = (lr % HTP_MM_HMX_TILE_N_ROWS) / 2;  // index of the row pair within the tile
         const __fp16 *row_base = vtcm_src + r0 * tile_row_stride;
 
         if (r_idx0 >= cne1) break;

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

@@ -14,7 +14,7 @@
 
 #define QURT_LOWEST_PRIO (254)
 
-static inline void hmx_lock(struct hmx_queue *q)
+static inline void hmx_lock(hmx_queue_t q)
 {
     if (!q->hmx_locked) {
         HAP_compute_res_hmx_lock(q->hap_rctx);
@@ -22,7 +22,7 @@ static inline void hmx_lock(struct hmx_queue *q)
     }
 }
 
-static inline void hmx_unlock(struct hmx_queue *q)
+static inline void hmx_unlock(hmx_queue_t q)
 {
     if (q->hmx_locked) {
         HAP_compute_res_hmx_unlock(q->hap_rctx);
@@ -30,7 +30,7 @@ static inline void hmx_unlock(struct hmx_queue *q)
     }
 }
 
-static inline void hmx_queue_process(struct hmx_queue *q, bool* killed) {
+static inline void hmx_queue_process(hmx_queue_t q, bool* killed) {
     unsigned int ir = atomic_load(&q->idx_read);
 
     while (ir != atomic_load(&q->idx_write)) {
@@ -61,7 +61,7 @@ static inline void hmx_queue_process(struct hmx_queue *q, bool* killed) {
 }
 
 static void hmx_queue_thread(void * arg) {
-    struct hmx_queue * q = (struct hmx_queue *) arg;
+    hmx_queue_t q = (hmx_queue_t) arg;
 
     FARF(HIGH, "hmx-queue-thread: started");
 
@@ -93,34 +93,41 @@ static void hmx_queue_thread(void * arg) {
     FARF(HIGH, "hmx-queue-thread: stopped");
 }
 
-struct hmx_queue * hmx_queue_create(size_t capacity, uint32_t hap_rctx) {
+size_t hmx_queue_sizeof(size_t capacity, uint32_t stack_size) {
     capacity = hex_ceil_pow2(capacity);
+    size_t size_q = hex_align_up(sizeof(struct hmx_queue_s), HEX_L2_LINE_SIZE);
+    size_t size_desc = hex_align_up(capacity * sizeof(struct hmx_queue_desc), HEX_L2_LINE_SIZE);
+    size_t size_stack = stack_size;
+    return size_q + size_desc + size_stack;
+}
+
+size_t hmx_queue_alignof(void) {
+    return HEX_L2_LINE_SIZE;
+}
+
+hmx_queue_t hmx_queue_init(void * ptr, size_t capacity, uint32_t stack_size, uint32_t hap_rctx, struct htp_thread_trace * trace) {
+    capacity = hex_ceil_pow2(capacity);
+    size_t size_q = hex_align_up(sizeof(struct hmx_queue_s), HEX_L2_LINE_SIZE);
+    size_t size_desc = hex_align_up(capacity * sizeof(struct hmx_queue_desc), HEX_L2_LINE_SIZE);
+
+    uint8_t * block = (uint8_t *) ptr;
+
+    hmx_queue_t q = (hmx_queue_t) block; block += size_q;
+    memset(q, 0, sizeof(struct hmx_queue_s));
 
-    struct hmx_queue * q = (struct hmx_queue *) memalign(32, sizeof(struct hmx_queue));
-    if (q == NULL) {
-        FARF(ERROR, "%s: failed to allocate DMA queue\n", __FUNCTION__);
-        return NULL;
-    }
-    memset(q, 0, sizeof(struct hmx_queue));
     q->capacity = capacity;
     q->idx_mask = capacity - 1;
     q->hap_rctx = hap_rctx;
+    q->external_mem = true;
 
-    q->desc = (struct hmx_queue_desc *) memalign(64, capacity * sizeof(struct hmx_queue_desc));
-    if (!q->desc) {
-        FARF(ERROR, "hmx-queue: failed to allocate HMX queue descriptors\n");
-        return NULL;
-    }
+    q->desc = (struct hmx_queue_desc *) block; block += size_desc;
     memset(q->desc, 0, capacity * sizeof(struct hmx_queue_desc));
 
-    const size_t stack_size = HMX_QUEUE_THREAD_STACK_SIZE;
-    q->stack = (unsigned char *) memalign(64, stack_size);
-    if (!q->stack) {
-        FARF(ERROR, "hmx-queue: thread stack allocation failed (%zu bytes)", stack_size);
-        return NULL;
-    }
+    q->stack = block;
     memset(q->stack, 0, stack_size);
 
+    q->trace = trace;
+
     // Match caller thread priority (same pattern as worker-pool.c).
     int prio = qurt_thread_get_priority(qurt_thread_get_id());
     if (prio < 1) {
@@ -148,7 +155,7 @@ struct hmx_queue * hmx_queue_create(size_t capacity, uint32_t hap_rctx) {
     return q;
 }
 
-void hmx_queue_delete(struct hmx_queue * q) {
+void hmx_queue_free(hmx_queue_t q) {
     if (!q) {
         return;
     }
@@ -160,8 +167,4 @@ void hmx_queue_delete(struct hmx_queue * q) {
 
     int status;
     qurt_thread_join(q->thread, &status);
-
-    free(q->desc);
-    free(q->stack);
-    free(q);
 }

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

@@ -17,8 +17,6 @@
 extern "C" {
 #endif
 
-#define HMX_QUEUE_THREAD_STACK_SIZE (16 * 1024)
-
 #if __HVX_ARCH__ > 79
 #define HMX_QUEUE_POLL_COUNT 2000
 #else
@@ -41,7 +39,7 @@ struct hmx_queue_desc {
     atomic_uint      done;
 };
 
-struct hmx_queue {
+struct hmx_queue_s {
     struct hmx_queue_desc * desc;
     atomic_uint      idx_write; // updated by producer (push)
     atomic_uint      idx_read;  // updated by consumer (process)
@@ -55,19 +53,24 @@ struct hmx_queue {
     uint32_t         hap_rctx;
     bool             hmx_locked;
     struct htp_thread_trace * trace;
+    bool             external_mem; // memory owned externally
 };
 
-struct hmx_queue * hmx_queue_create(size_t capacity, uint32_t hap_rctx);
-void hmx_queue_delete(struct hmx_queue * q);
+typedef struct hmx_queue_s * hmx_queue_t;
+
+size_t      hmx_queue_sizeof(size_t capacity, uint32_t stack_size);
+size_t      hmx_queue_alignof(void);
+hmx_queue_t hmx_queue_init(void * ptr, size_t capacity, uint32_t stack_size, uint32_t hap_rctx, struct htp_thread_trace * trace);
+void        hmx_queue_free(hmx_queue_t q);
 
 static inline struct hmx_queue_desc hmx_queue_make_desc(hmx_queue_func func, void * data) {
     struct hmx_queue_desc d = { func, data };
     return d;
 }
 
-static inline bool hmx_queue_push(struct hmx_queue * q, struct hmx_queue_desc d) {
+static inline bool hmx_queue_push(hmx_queue_t q, struct hmx_queue_desc d) {
     unsigned int ir = atomic_load(&q->idx_read);
-    unsigned int iw = q->idx_write;
+    unsigned int iw = atomic_load(&q->idx_write);
 
     if (((iw + 1) & q->idx_mask) == ir) {
         FARF(HIGH, "hmx-queue-push: queue is full\n");
@@ -87,25 +90,25 @@ static inline bool hmx_queue_push(struct hmx_queue * q, struct hmx_queue_desc d)
     return true;
 }
 
-static inline bool hmx_queue_signal(struct hmx_queue *q, enum hmx_queue_signal sig) {
+static inline bool hmx_queue_signal(hmx_queue_t q, enum hmx_queue_signal sig) {
     return hmx_queue_push(q, hmx_queue_make_desc((hmx_queue_func) sig, NULL));
 }
 
-static inline bool hmx_queue_empty(struct hmx_queue * q) {
-    return q->idx_pop == q->idx_write;
+static inline bool hmx_queue_empty(hmx_queue_t q) {
+    return q->idx_pop == atomic_load(&q->idx_write);
 }
 
-static inline uint32_t hmx_queue_depth(struct hmx_queue * q) {
-    return (q->idx_read - q->idx_read) & q->idx_mask;
+static inline uint32_t hmx_queue_depth(hmx_queue_t q) {
+    return (atomic_load(&q->idx_write) - atomic_load(&q->idx_read)) & q->idx_mask;
 }
 
-static inline uint32_t hmx_queue_capacity(struct hmx_queue * q) {
+static inline uint32_t hmx_queue_capacity(hmx_queue_t q) {
     return q->capacity;
 }
 
-static inline struct hmx_queue_desc hmx_queue_pop_one(struct hmx_queue * q) {
+static inline struct hmx_queue_desc hmx_queue_pop_one(hmx_queue_t q) {
     unsigned int ip = q->idx_pop;
-    unsigned int iw = q->idx_write;
+    unsigned int iw = atomic_load(&q->idx_write);
 
     struct hmx_queue_desc rd = { NULL, NULL };
     if (ip == iw) {
@@ -126,7 +129,7 @@ static inline struct hmx_queue_desc hmx_queue_pop_one(struct hmx_queue * q) {
     return rd;
 }
 
-static inline struct hmx_queue_desc hmx_queue_pop(struct hmx_queue * q) {
+static inline struct hmx_queue_desc hmx_queue_pop(hmx_queue_t q) {
     while (1) {
         struct hmx_queue_desc d = hmx_queue_pop_one(q);
 
@@ -138,15 +141,15 @@ static inline struct hmx_queue_desc hmx_queue_pop(struct hmx_queue * q) {
     }
 }
 
-static inline void hmx_queue_flush(struct hmx_queue * q) {
+static inline void hmx_queue_flush(hmx_queue_t q) {
     while (hmx_queue_pop_one(q).func != NULL) ;
 }
 
-static inline void hmx_queue_wakeup(struct hmx_queue * q) {
+static inline void hmx_queue_wakeup(hmx_queue_t q) {
     hmx_queue_signal(q, HMX_QUEUE_WAKEUP);
 }
 
-static inline void hmx_queue_suspend(struct hmx_queue *q) {
+static inline void hmx_queue_suspend(hmx_queue_t q) {
     hmx_queue_signal(q, HMX_QUEUE_SUSPEND);
 }
 

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

@@ -5,7 +5,8 @@
 #include "hmx-queue.h"
 #include "htp-ops.h"
 #include "hex-profile.h"
-#include "worker-pool.h"
+#include "work-queue.h"
+#include "hex-fastdiv.h"
 
 #include <assert.h>
 #include <dspqueue.h>
@@ -52,6 +53,9 @@ struct htp_ops_context {
         const struct htp_tensor * dsts[HTP_OP_MAX_OUTPUTS];
     };
 
+    dma_queue **    src_dma[HTP_OP_MAX_INPUTS];
+    dma_queue **    dst_dma[HTP_OP_MAX_OUTPUTS];
+
     // TODO convert these to an array
     struct htp_spad src0_spad;
     struct htp_spad src1_spad;
@@ -65,11 +69,16 @@ struct htp_ops_context {
 
 // Main context for htp DSP backend
 struct htp_context {
-    dspqueue_t             queue;
-    dma_queue *            dma[HTP_MAX_NTHREADS];
+    dspqueue_t             dsp_queue;
+
     struct htp_mmap        mmap[HTP_MAX_MMAPS];
-    worker_pool_context_t  worker_pool;
+    dma_queue_t            dma[HTP_MAX_NTHREADS];
+    dma_queue_t            dma_cached[HTP_MAX_NTHREADS];
+    work_queue_t           work_queue;
+    hmx_queue_t            hmx_queue;
+
     uint32_t               n_threads;
+    struct fastdiv_values  n_threads_div;
 
     int                    thread_id;
     int                    thread_prio;
@@ -86,6 +95,7 @@ struct htp_context {
     atomic_bool            vtcm_needs_release;
 
     uint64_t               max_vmem;
+    uint32_t               dirty_map[HTP_OP_MAX_TENSORS / 32];
 
     // Persistent DDR scratchpad for MUL_MAT_ID mappings
     void *                 ddr_spad_base;
@@ -93,7 +103,10 @@ struct htp_context {
 
     struct htp_ops_context octx;
 
-    struct hmx_queue *     hmx_queue; // Async HMX queue for pipeline overlap
+    qurt_thread_t          main_thread;
+    void *                 main_stack;
+    atomic_bool            killed;
+    size_t                 footprint;
 };
 
 int op_matmul(struct htp_ops_context * octx);

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

@@ -97,6 +97,7 @@ enum htp_op_code {
     HTP_OP_PAD,
     HTP_OP_NORM,
     HTP_OP_CONCAT,
+    HTP_OP_CLAMP,
 
     HTP_OP_INVALID
 };
@@ -108,8 +109,7 @@ enum htp_op_code {
 #define HTP_OP_MAX_KERN_PARAMS 32
 
 #define HTP_OP_MAX_BUFS    16
-#define HTP_OP_MAX_REQS    256
-#define HTP_OP_MAX_TENSORS (HTP_OP_MAX_REQS * HTP_OP_MAX_INPUTS + HTP_OP_MAX_REQS)
+#define HTP_OP_MAX_TENSORS 8192 // must stay under 64K (uint16)
 
 #define HTP_OP_MAX_VMEM_DEFAULT (3355443200u)
 
@@ -117,16 +117,18 @@ enum htp_op_code {
 
 enum htp_tensor_flags {
     HTP_TENSOR_COMPUTE = (1U << 0), // Tensor buffer temporal compute data (not weights)
-    HTP_TENSOR_FLUSHED = (1U << 1)  // Tensor buffer has been flushed (set by the NPU)
+    HTP_TENSOR_DIRTY   = (1U << 1)  // Tensor buffer is dirty and needs to be flushed
 };
 
 // Tensor descriptor
 struct htp_tensor {
     uint32_t data;                 // Buffer offset in the messages, and data pointer on the NPU
+    uint32_t alias;                // Index of the canonical tensor for this memory buffer
     uint32_t size;                 // Data size in bytes
     uint32_t flags;                // Buffer / tensor flags
-    uint16_t type;                 // Data type
+    uint32_t type;                 // Data type
     uint16_t bi;                   // Buffer index
+    uint16_t ti;                   // Tensor index
     uint32_t ne[HTP_OP_MAX_DIMS];  // Number of elements
     uint32_t nb[HTP_OP_MAX_DIMS];  // Stride in bytes (see ggml.h ggml_tensor)
 };
@@ -169,6 +171,8 @@ enum htp_profiler_mode {
 
 enum htp_trace_event_id {
     HTP_TRACE_EVT_DMA                 = 0,
+    HTP_TRACE_EVT_L2FLUSH             = 1,
+    HTP_TRACE_EVT_INIT                = 2,
 
     HTP_TRACE_EVT_HVX_COMP            = 20,
     HTP_TRACE_EVT_HVX_A_QUANT         = 21,

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

@@ -0,0 +1,204 @@
+#include "htp-tensor.h"
+
+#include <qurt.h>
+#include <qurt_memory.h>
+
+#include "hex-common.h"
+#include "hex-utils.h"
+#include "hex-fastdiv.h"
+#include "hex-profile.h"
+#include "htp-ctx.h"
+#include "work-queue.h"
+
+struct l2flush_task {
+    struct htp_thread_trace * trace;
+    uint32_t start;
+    uint32_t end;
+    uint32_t chunk_size;
+    uint32_t ti;
+};
+
+static void l2flush_thread_worker(unsigned int n, unsigned int i, void * data) {
+    struct l2flush_task * task = (struct l2flush_task *) data;
+    const uint32_t start = task->start;
+    const uint32_t end   = task->end;
+    const uint32_t ti    = task->ti;
+    const uint32_t chunk_size = task->chunk_size;
+
+    const uint32_t thread_s = start + i * chunk_size;
+    if (thread_s >= end) {
+        return;
+    }
+    uint32_t thread_e = thread_s + chunk_size;
+    if (thread_e > end) {
+        thread_e = end;
+    }
+
+    struct htp_thread_trace * tr = &task->trace[i];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, ti);
+    hex_l2flush((void *) (uintptr_t) thread_s, thread_e - thread_s);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, ti);
+}
+
+static void flush_all_dcache(struct htp_context * ctx) {
+    struct htp_thread_trace * tr = &ctx->trace[0];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
+    hex_l2fetch_block(ctx, ctx->footprint);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
+    bitmap_reset(ctx->dirty_map, HTP_OP_MAX_TENSORS);
+}
+
+static void flush_tensor_range(struct htp_context * ctx, const struct htp_tensor * t) {
+    struct htp_thread_trace * tr = &ctx->trace[0];
+
+    if (t->size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
+        struct l2flush_task task;
+        task.start = hex_align_down((size_t) t->data, HEX_L2_LINE_SIZE);
+        task.end   = hex_align_up((size_t) t->data + t->size, HEX_L2_LINE_SIZE);
+        task.ti    = t->ti;
+        task.trace = ctx->trace;
+
+        const uint32_t total_size = task.end - task.start;
+        const uint32_t n_blocks   = (total_size + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
+        const uint32_t blocks_per_thread = fastdiv(n_blocks + ctx->n_threads - 1, &ctx->n_threads_div);
+        task.chunk_size = blocks_per_thread * HEX_L2_BLOCK_SIZE;
+
+        work_queue_run(ctx->work_queue, l2flush_thread_worker, &task, ctx->n_threads);
+    } else {
+        htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
+        hex_l2flush((void *) t->data, t->size);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
+    }
+
+    htp_tensor_make_clean(t, ctx->dirty_map);
+}
+
+void htp_tensor_flush(struct htp_context * ctx, const struct htp_tensor * t) {
+    if (!bitmap_test(ctx->dirty_map, t->ti)) {
+        return;
+    }
+
+    if (t->size > HEX_L2_FLUSH_ALL_THRESHOLD) {
+        flush_all_dcache(ctx);
+        return;
+    }
+
+    flush_tensor_range(ctx, t);
+}
+
+// One dirty tensor's line-aligned range, placed in the flattened global block space.
+struct l2flush_range {
+    uint32_t start;       // line-aligned start address
+    uint32_t end;         // line-aligned end address
+    uint32_t block_first; // global block index of this range's first block
+    uint32_t n_blocks;    // number of HEX_L2_BLOCK_SIZE chunks (last may be partial)
+};
+
+struct l2flush_multi_task {
+    struct htp_thread_trace * trace;
+    struct l2flush_range      ranges[HTP_OP_MAX_INPUTS];
+    uint32_t                  n_ranges;
+    uint32_t                  total_blocks;
+    uint32_t                  blocks_per_thread;
+};
+
+static void l2flush_multi_worker(unsigned int n, unsigned int i, void * data) {
+    (void) n;
+    struct l2flush_multi_task * task = (struct l2flush_multi_task *) data;
+
+    const uint32_t gb_first = i * task->blocks_per_thread;
+    uint32_t       gb_last  = gb_first + task->blocks_per_thread;
+    if (gb_last > task->total_blocks) {
+        gb_last = task->total_blocks;
+    }
+    if (gb_first >= gb_last) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &task->trace[i];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, gb_first);
+
+    for (uint32_t r = 0; r < task->n_ranges; r++) {
+        const struct l2flush_range * rg = &task->ranges[r];
+        const uint32_t rb_first = rg->block_first;
+        const uint32_t rb_last  = rg->block_first + rg->n_blocks;
+
+        const uint32_t lo = gb_first > rb_first ? gb_first : rb_first;
+        const uint32_t hi = gb_last  < rb_last  ? gb_last  : rb_last;
+        if (lo >= hi) {
+            continue;
+        }
+
+        const uint32_t s = rg->start + (lo - rb_first) * HEX_L2_BLOCK_SIZE;
+        uint32_t       e = rg->start + (hi - rb_first) * HEX_L2_BLOCK_SIZE;
+        if (e > rg->end) {
+            e = rg->end;
+        }
+        hex_l2flush((void *) (uintptr_t) s, e - s);
+    }
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, gb_first);
+}
+
+void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
+    uint64_t total_dirty = 0;
+    for (uint32_t i = 0; i < n; i++) {
+        const struct htp_tensor * t = tensors[i];
+        if (t && bitmap_test(ctx->dirty_map, t->ti)) {
+            total_dirty += t->size;
+        }
+    }
+
+    if (total_dirty == 0) {
+        return;
+    }
+
+    if (total_dirty > HEX_L2_FLUSH_ALL_THRESHOLD) {
+        flush_all_dcache(ctx);
+        return;
+    }
+
+    // Aggregate is small enough to walk. Thread it across all dirty ranges at once
+    // when it is worth the dispatch, otherwise flush sequentially.
+    if (total_dirty > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
+        struct l2flush_multi_task task;
+        task.trace    = ctx->trace;
+        task.n_ranges = 0;
+
+        uint32_t block_acc = 0;
+        for (uint32_t i = 0; i < n; i++) {
+            const struct htp_tensor * t = tensors[i];
+            if (!t || !bitmap_test(ctx->dirty_map, t->ti)) {
+                continue;
+            }
+            // Clear as we go: dedups a tensor passed as multiple srcs (e.g. mul(x,x)).
+            htp_tensor_make_clean(t, ctx->dirty_map);
+
+            struct l2flush_range * rg = &task.ranges[task.n_ranges++];
+            rg->start = hex_align_down((size_t) t->data, HEX_L2_LINE_SIZE);
+            rg->end   = hex_align_up((size_t) t->data + t->size, HEX_L2_LINE_SIZE);
+            rg->block_first = block_acc;
+            rg->n_blocks = (rg->end - rg->start + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
+            block_acc += rg->n_blocks;
+        }
+
+        task.total_blocks      = block_acc;
+        task.blocks_per_thread = fastdiv(block_acc + ctx->n_threads - 1, &ctx->n_threads_div);
+
+        work_queue_run(ctx->work_queue, l2flush_multi_worker, &task, ctx->n_threads);
+        return;
+    }
+
+    struct htp_thread_trace * tr = &ctx->trace[0];
+    for (uint32_t i = 0; i < n; i++) {
+        const struct htp_tensor * t = tensors[i];
+        if (!t || !bitmap_test(ctx->dirty_map, t->ti)) {
+            continue;
+        }
+        htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
+        hex_l2flush((void *) t->data, t->size);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, t->ti);
+        htp_tensor_make_clean(t, ctx->dirty_map);
+    }
+}

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

@@ -0,0 +1,36 @@
+#ifndef HTP_TENSOR_H
+#define HTP_TENSOR_H
+
+#include <stdint.h>
+#include "htp-ops.h"
+#include "hex-bitmap.h"
+
+static inline struct htp_tensor * htp_tensor_alias(const struct htp_tensor * t) {
+    return (struct htp_tensor *) (uintptr_t) t->alias;
+}
+
+static inline void * htp_tensor_data(const struct htp_tensor * t) {
+    return (void *) (uintptr_t) t->data;
+}
+
+static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) {
+    return (uint32_t *) &t->flags;
+}
+
+static inline void htp_tensor_make_dirty(const struct htp_tensor * t, uint32_t * dirty_map) {
+    struct htp_tensor * curr = (struct htp_tensor *) t;
+    do {
+        bitmap_set(dirty_map, curr->ti);
+        curr = htp_tensor_alias(curr);
+    } while (curr != t);
+}
+
+static inline void htp_tensor_make_clean(const struct htp_tensor * t, uint32_t * dirty_map) {
+    bitmap_clear(dirty_map, t->ti);
+}
+
+struct htp_context;
+void htp_tensor_flush(struct htp_context * ctx, const struct htp_tensor * t);
+void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
+
+#endif // HTP_TENSOR_H

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

@@ -25,112 +25,44 @@
 
 #define GGML_COMMON_DECL_C
 #include "ggml-common.h"
+#include "hex-bitmap.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
-#include "htp-ops.h"
+#include "htp-tensor.h"
 #include "htp_iface.h"
-#include "worker-pool.h"
+#include "work-queue.h"
+#include "hex-profile.h"
+
+#define HMX_QUEUE_CAPACITY     16
+#define HMX_QUEUE_STACK_SIZE   16384
+#define WORK_QUEUE_CAPACITY    16
+#define WORK_QUEUE_STACK_SIZE  16384
+#define MAIN_THREAD_STACK_SIZE 32768
+
+_Static_assert(WORK_QUEUE_MAX_N_THREADS >= HTP_MAX_NTHREADS,
+               "work-queue thread cap must be >= HTP_MAX_NTHREADS");
+
+struct htp_handle {
+    struct htp_context * ctx;
+};
 
 AEEResult htp_iface_open(const char * uri, remote_handle64 * handle) {
-    struct htp_context * ctx;
-    int                  err = 0;
-
-    ctx = calloc(1, sizeof(*ctx));
-    if (ctx == NULL) {
+    (void) uri;
+    struct htp_handle * h = calloc(1, sizeof(*h));
+    if (h == NULL) {
         return AEE_ENOMEMORY;
     }
 
-    // Use the context structure as the handle
-    *handle = (remote_handle64) ctx;
-
-    // Enable FARF logs
-    HAP_setFARFRuntimeLoggingParams(0xffff, NULL, 0);
-
-    // Set client class
-    {
-        HAP_power_request_t request;
-        memset(&request, 0, sizeof(HAP_power_request_t));
-        request.type    = HAP_power_set_apptype;
-        request.apptype = HAP_POWER_COMPUTE_CLIENT_CLASS;
-
-        if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
-            return err;
-        }
-    }
-
-    {
-        HAP_power_request_t request;
-        memset(&request, 0, sizeof(request));
-
-        request.type                              = HAP_power_set_DCVS_v3;
-        request.dcvs_v3.set_dcvs_enable           = TRUE;
-        request.dcvs_v3.dcvs_enable               = FALSE;
-        request.dcvs_v3.set_bus_params            = TRUE;
-        request.dcvs_v3.bus_params.min_corner     = HAP_DCVS_VCORNER_MAX;
-        request.dcvs_v3.bus_params.max_corner     = HAP_DCVS_VCORNER_MAX;
-        request.dcvs_v3.bus_params.target_corner  = HAP_DCVS_VCORNER_MAX;
-        request.dcvs_v3.set_core_params           = TRUE;
-        request.dcvs_v3.core_params.min_corner    = HAP_DCVS_VCORNER_MAX;
-        request.dcvs_v3.core_params.max_corner    = HAP_DCVS_VCORNER_MAX;
-        request.dcvs_v3.core_params.target_corner = HAP_DCVS_VCORNER_MAX;
-        request.dcvs_v3.set_sleep_disable         = TRUE;
-        request.dcvs_v3.sleep_disable             = TRUE;
-
-#if (__HEXAGON_ARCH__ >= 79)
-        HAP_set_dcvs_v3_protected_bus_corners(&request, 1);
-#endif
-        if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
-            return err;
-        }
-
-        memset(&request, 0, sizeof(request));
-        request.type         = HAP_power_set_HVX;
-        request.hvx.power_up = TRUE;
-        if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
-            return err;
-        }
-    }
-
-#if __HVX_ARCH__ >= 75
-    {
-        // Power on HMX and set HMX clock
-        HAP_power_request_t request;
-        memset(&request, 0, sizeof(HAP_power_request_t));
-        request.type = HAP_power_set_HMX_v2;
-        request.hmx_v2.set_power     = TRUE;
-        request.hmx_v2.power_up      = TRUE;
-        request.hmx_v2.set_clock     = TRUE;
-        request.hmx_v2.target_corner = HAP_DCVS_EXP_VCORNER_MAX;
-        request.hmx_v2.min_corner    = HAP_DCVS_EXP_VCORNER_MAX;
-        request.hmx_v2.max_corner    = HAP_DCVS_EXP_VCORNER_MAX;
-        request.hmx_v2.perf_mode     = HAP_CLK_PERF_HIGH;
-        FARF(ALWAYS, "Setting HMX clock\n");
-        err = HAP_power_set((void *) ctx, &request);
-        if (err != AEE_SUCCESS) {
-            FARF(ERROR, "ggml-hex: error setting HMX clock.");
-            return err;
-        }
-    }
-#else
-    {
-        // Power on HMX
-        HAP_power_request_t request;
-        memset(&request, 0, sizeof(HAP_power_request_t));
-        request.type         = HAP_power_set_HMX;
-        request.hmx.power_up = TRUE;
-        FARF(ALWAYS, "Powering HMX on\n");
-        err = HAP_power_set((void *) ctx, &request);
-        if (err != AEE_SUCCESS) {
-            FARF(ERROR, "ggml-hex: error powering on HMX.");
-            return err;
-        }
-    }
-#endif
-
+    *handle = (remote_handle64) h;
     return AEE_SUCCESS;
 }
 
 AEEResult htp_iface_etm(remote_handle64 handle, uint32_t enable) {
+    struct htp_handle * h = (struct htp_handle *) handle;
+    if (!h) {
+        return AEE_EBADPARM;
+    }
+
     int err = enable ? HAP_user_etm_enable() : HAP_user_etm_disable();
     if (err) {
         if (err == AEE_EVERSIONNOTSUPPORT) {
@@ -143,10 +75,11 @@ AEEResult htp_iface_etm(remote_handle64 handle, uint32_t enable) {
 }
 
 AEEResult htp_iface_profiler(remote_handle64 handle, uint32_t mode, const htp_iface_pmu_conf* pmu_conf) {
-    struct htp_context * ctx = (struct htp_context *) handle;
-    if (!ctx) {
+    struct htp_handle * h = (struct htp_handle *) handle;
+    if (!h || !h->ctx) {
         return AEE_EBADPARM;
     }
+    struct htp_context * ctx = h->ctx;
 
     if (mode == HTP_PROF_PMU) {
         const uint32_t* events = pmu_conf->events;
@@ -179,48 +112,55 @@ AEEResult htp_iface_profiler(remote_handle64 handle, uint32_t mode, const htp_if
 }
 
 AEEResult htp_iface_close(remote_handle64 handle) {
-    struct htp_context * ctx = (struct htp_context *) handle;
-
-    if (!ctx) {
+    struct htp_handle * h = (struct htp_handle *) handle;
+    if (!h) {
         return AEE_EBADPARM;
     }
 
-    if (ctx->queue) {
-        FARF(ERROR, "Closing handle with queue still open");
-        return AEE_EITEMBUSY;
-    }
-
-    // release the mmaps (if any)
-    for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
-        if (ctx->mmap[i].size) {
-#if __HVX_ARCH__ > 73
-            HAP_munmap2((void *) ctx->mmap[i].base, ctx->mmap[i].size);
-#else
-            HAP_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size);
-#endif
-            ctx->mmap[i].size = 0;
-            ctx->mmap[i].base = NULL;
-            ctx->mmap[i].fd   = -1;
+    struct htp_context * ctx = h->ctx;
+    if (ctx) {
+        if (ctx->dsp_queue) {
+            FARF(ERROR, "Closing handle with queue still open");
+            return AEE_EITEMBUSY;
         }
+
+        // release the mmaps (if any)
+        for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
+            if (ctx->mmap[i].size) {
+#if __HVX_ARCH__ > 73
+                HAP_munmap2((void *) ctx->mmap[i].base, ctx->mmap[i].size);
+#else
+                HAP_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size);
+#endif
+                ctx->mmap[i].size = 0;
+                ctx->mmap[i].base = NULL;
+                ctx->mmap[i].fd   = -1;
+            }
+        }
+
+        if (ctx->profiler) {
+            qurt_pmu_enable(1);
+        }
+
+        if (ctx->etm) {
+            HAP_user_etm_disable();
+        }
+
+        // Free the unified block (ctx is the base address of the block)
+        free(ctx);
+        h->ctx = NULL;
     }
 
-    if (ctx->profiler) {
-        qurt_pmu_enable(1);
-    }
-
-    if (ctx->etm) {
-        HAP_user_etm_disable();
-    }
-
-    free(ctx);
+    free(h);
     return AEE_SUCCESS;
 }
 
 AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint32_t size) {
-    struct htp_context * ctx = (struct htp_context *) handle;
-    if (!ctx) {
+    struct htp_handle * h = (struct htp_handle *) handle;
+    if (!h || !h->ctx) {
         return AEE_EBADPARM;
     }
+    struct htp_context * ctx = h->ctx;
 
     // See if we already have this mapping
     for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
@@ -262,10 +202,11 @@ AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint32_t size) {
 }
 
 AEEResult htp_iface_munmap(remote_handle64 handle, uint32 fd) {
-    struct htp_context * ctx = (struct htp_context *) handle;
-    if (!ctx) {
+    struct htp_handle * h = (struct htp_handle *) handle;
+    if (!h || !h->ctx) {
         return AEE_EBADPARM;
     }
+    struct htp_context * ctx = h->ctx;
 
     for (uint32_t i=0; i<HTP_MAX_MMAPS; i++) {
         struct htp_mmap *m = &ctx->mmap[i];
@@ -358,56 +299,44 @@ static void vtcm_free(struct htp_context * ctx) {
     }
 }
 
+static void htp_main_thread(void * context);
 static void htp_packet_callback(dspqueue_t queue, int error, void * context);
 static void htp_error_callback(dspqueue_t queue, int error, void * context);
 
 AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp_queue_id, uint32_t n_hvx, uint32_t n_hmx, uint64_t max_vmem) {
-    struct htp_context * ctx = (struct htp_context *) handle;
-
-    if (!ctx) {
+    struct htp_handle * h = (struct htp_handle *) handle;
+    if (!h) {
         return AEE_EBADPARM;
     }
 
-    if (ctx->queue) {
+    if (h->ctx) {
         FARF(ERROR, "Queue already open");
         return AEE_EITEMBUSY;
     }
 
-    // Import queue created on the CPU
-    int err = dspqueue_import(dsp_queue_id,         // Queue ID from dspqueue_export
-                              htp_packet_callback,  // Packet callback
-                              htp_error_callback,   // Error callback; no errors expected on the DSP
-                              (void *) ctx,         // Callback context
-                              &ctx->queue);
+    // Cache the original FastRPC thread priority, then calculate compute priority
+    int fastrpc_tid  = qurt_thread_get_id();
+    int fastrpc_prio = qurt_thread_get_priority(fastrpc_tid);
+    int main_prio    = fastrpc_prio - 10;
+    if (main_prio < 1) main_prio = 1;
+
+    dspqueue_t dsp_queue = NULL;
+    bool use_callbacks = false;
+
+    // Import queue with NULL callbacks to avoid starting dspueue internal threads
+    int err = dspqueue_import(dsp_queue_id, NULL, NULL, (void *) h, &dsp_queue);
+    if (err == AEE_EBADPARM) {
+        // Fallback for devices that don't support NULL callbacks
+        FARF(HIGH, "dspqueue import with NULL callbacks failed, trying with callbacks");
+        use_callbacks = true;
+        err = dspqueue_import(dsp_queue_id, htp_packet_callback, htp_error_callback, (void *) h, &dsp_queue);
+    }
+
     if (err) {
         FARF(ERROR, "Queue import failed with 0x%08x", (unsigned) err);
         return err;
     }
 
-    ctx->max_vmem    = max_vmem;
-    ctx->thread_id   = qurt_thread_get_id();
-    ctx->thread_prio = qurt_thread_get_priority(ctx->thread_id);
-
-    // allocate VTCM
-    err = vtcm_alloc(ctx);
-    if (err != AEE_SUCCESS) {
-        FARF(ERROR, "Unable to allocate VTCM");
-        return AEE_ENOMEMORY;
-    }
-
-    ctx->hmx_enabled = n_hmx;
-    ctx->hmx_queue   = NULL;
-    if (n_hmx) {
-        ctx->hmx_queue = hmx_queue_create(16, ctx->vtcm_rctx);
-        if (ctx->hmx_queue) {
-            ctx->hmx_queue->trace = &ctx->trace[HTP_MAX_NTHREADS];
-        } else {
-            FARF(ERROR, "hmx-queue-create failed");
-            ctx->hmx_enabled = false;
-        }
-    }
-    FARF(HIGH, "HMX %s (n_hmx=%d)", ctx->hmx_enabled ? "enabled" : "disabled", n_hmx);
-
     qurt_sysenv_max_hthreads_t hw_threads;
     qurt_sysenv_get_max_hw_threads(&hw_threads);
     uint32_t hw_nhvx = (qurt_hvx_get_units() >> 8) & 0xFF;
@@ -422,27 +351,216 @@ AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp
         n_hvx = HTP_MAX_NTHREADS;
     }
 
-    ctx->n_threads = n_hvx;
-    for (int i = 0; i < ctx->n_threads; i++) {
-        ctx->dma[i] = dma_queue_create(256); // queue depth
-        if (ctx->dma[i]) {
-            ctx->dma[i]->trace = &ctx->trace[i];
+    // layout segments of our contiguous block
+
+    // 1. htp_context : sits at the base (block is 4K-aligned via memalign below)
+    size_t offset = sizeof(struct htp_context);
+
+    // 2. main_stack
+    size_t offset_main_stack = 0;
+    size_t size_main_stack   = 0;
+    if (!use_callbacks) {
+        offset_main_stack = hex_align_up(offset, 4096);
+        size_main_stack   = MAIN_THREAD_STACK_SIZE;
+        offset = offset_main_stack + size_main_stack;
+    }
+
+    // 3. work_queue
+    size_t wq_align  = work_queue_alignof();
+    size_t offset_wq = hex_align_up(offset, wq_align);
+    size_t size_wq   = work_queue_sizeof(n_hvx, WORK_QUEUE_CAPACITY, WORK_QUEUE_STACK_SIZE);
+    offset = offset_wq + size_wq;
+
+    // 4. dma_queue
+    size_t dma_align = dma_queue_alignof();
+    size_t offset_dma = hex_align_up(offset, dma_align);
+    size_t size_dma = 0;
+    for (uint32_t i = 0; i < n_hvx; i++) {
+        size_dma  = hex_align_up(size_dma, dma_queue_alignof());
+        size_dma += dma_queue_sizeof(256);
+        size_dma  = hex_align_up(size_dma, dma_queue_alignof());
+        size_dma += dma_queue_alias_sizeof();
+    }
+    offset = offset_dma + size_dma;
+
+    // 5. hmx_queue
+    size_t offset_hmx = 0;
+    size_t size_hmx = 0;
+    if (n_hmx) {
+        size_t hmx_align = hmx_queue_alignof();
+        offset_hmx = hex_align_up(offset, hmx_align);
+        size_hmx   = hmx_queue_sizeof(HMX_QUEUE_CAPACITY, HMX_QUEUE_STACK_SIZE);
+        offset = offset_hmx + size_hmx;
+    }
+
+    size_t footprint = hex_align_up(offset, 128);
+
+    void * block = memalign(4096, footprint);
+    if (!block) {
+        FARF(ERROR, "Unable to allocate unified block of size %zu\n", footprint);
+        dspqueue_close(dsp_queue);
+        return AEE_ENOMEMORY;
+    }
+    memset(block, 0, footprint);
+
+    h->ctx = (struct htp_context *) block;
+    struct htp_context * ctx = h->ctx;
+    ctx->footprint = footprint;
+
+    ctx->thread_id   = fastrpc_tid;
+    ctx->thread_prio = main_prio;
+    ctx->max_vmem    = max_vmem;
+    ctx->dsp_queue   = dsp_queue;
+
+    err = vtcm_alloc(ctx);
+    if (err != AEE_SUCCESS) {
+        FARF(ERROR, "Unable to allocate VTCM");
+        htp_iface_stop(handle);
+        return AEE_ENOMEMORY;
+    }
+
+    HAP_setFARFRuntimeLoggingParams(0xffff, NULL, 0);
+
+    // Set client class
+    {
+        HAP_power_request_t request;
+        memset(&request, 0, sizeof(HAP_power_request_t));
+        request.type    = HAP_power_set_apptype;
+        request.apptype = HAP_POWER_COMPUTE_CLIENT_CLASS;
+
+        if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
+            htp_iface_stop(handle);
+            return err;
         }
     }
 
+    // DCVS setup
+    {
+        HAP_power_request_t request;
+        memset(&request, 0, sizeof(request));
+
+        request.type                              = HAP_power_set_DCVS_v3;
+        request.dcvs_v3.set_dcvs_enable           = TRUE;
+        request.dcvs_v3.dcvs_enable               = FALSE;
+        request.dcvs_v3.set_bus_params            = TRUE;
+        request.dcvs_v3.bus_params.min_corner     = HAP_DCVS_VCORNER_MAX;
+        request.dcvs_v3.bus_params.max_corner     = HAP_DCVS_VCORNER_MAX;
+        request.dcvs_v3.bus_params.target_corner  = HAP_DCVS_VCORNER_MAX;
+        request.dcvs_v3.set_core_params           = TRUE;
+        request.dcvs_v3.core_params.min_corner    = HAP_DCVS_VCORNER_MAX;
+        request.dcvs_v3.core_params.max_corner    = HAP_DCVS_VCORNER_MAX;
+        request.dcvs_v3.core_params.target_corner = HAP_DCVS_VCORNER_MAX;
+        request.dcvs_v3.set_sleep_disable         = TRUE;
+        request.dcvs_v3.sleep_disable             = TRUE;
+
+#if (__HEXAGON_ARCH__ >= 79)
+        HAP_set_dcvs_v3_protected_bus_corners(&request, 1);
+#endif
+        if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
+            htp_iface_stop(handle);
+            return err;
+        }
+
+        memset(&request, 0, sizeof(request));
+        request.type         = HAP_power_set_HVX;
+        request.hvx.power_up = TRUE;
+        if ((err = HAP_power_set((void *) ctx, &request)) != 0) {
+            htp_iface_stop(handle);
+            return err;
+        }
+    }
+
+#if __HVX_ARCH__ >= 75
+    {
+        // Power on HMX and set HMX clock
+        HAP_power_request_t request;
+        memset(&request, 0, sizeof(HAP_power_request_t));
+        request.type = HAP_power_set_HMX_v2;
+        request.hmx_v2.set_power     = TRUE;
+        request.hmx_v2.power_up      = TRUE;
+        request.hmx_v2.set_clock     = TRUE;
+        request.hmx_v2.target_corner = HAP_DCVS_EXP_VCORNER_MAX;
+        request.hmx_v2.min_corner    = HAP_DCVS_EXP_VCORNER_MAX;
+        request.hmx_v2.max_corner    = HAP_DCVS_EXP_VCORNER_MAX;
+        request.hmx_v2.perf_mode     = HAP_CLK_PERF_HIGH;
+        FARF(ALWAYS, "Setting HMX clock\n");
+        err = HAP_power_set((void *) ctx, &request);
+        if (err != AEE_SUCCESS) {
+            FARF(ERROR, "ggml-hex: error setting HMX clock.");
+            htp_iface_stop(handle);
+            return err;
+        }
+    }
+#else
+    {
+        // Power on HMX
+        HAP_power_request_t request;
+        memset(&request, 0, sizeof(HAP_power_request_t));
+        request.type         = HAP_power_set_HMX;
+        request.hmx.power_up = TRUE;
+        FARF(ALWAYS, "Powering HMX on\n");
+        err = HAP_power_set((void *) ctx, &request);
+        if (err != AEE_SUCCESS) {
+            FARF(ERROR, "ggml-hex: error powering on HMX.");
+            htp_iface_stop(handle);
+            return err;
+        }
+    }
+#endif
+
+    ctx->hmx_enabled = n_hmx;
+    ctx->hmx_queue   = NULL;
+    if (n_hmx) {
+        void * hmx_ptr = (void *) ((uintptr_t) block + offset_hmx);
+        ctx->hmx_queue = hmx_queue_init(hmx_ptr, HMX_QUEUE_CAPACITY, HMX_QUEUE_STACK_SIZE, ctx->vtcm_rctx, &ctx->trace[HTP_MAX_NTHREADS]);
+    }
+    FARF(HIGH, "HMX %s (n_hmx=%d)", ctx->hmx_enabled ? "enabled" : "disabled", n_hmx);
+
+    ctx->n_threads = n_hvx;
+    ctx->n_threads_div = init_fastdiv_values(ctx->n_threads);
+
+    // Initialize DMA queues
+    uint8_t * dma_ptr_curr = (uint8_t *) ((uintptr_t) block + offset_dma);
+    size_t size_dma_q = dma_queue_sizeof(256);
+    size_t size_dma_alias = dma_queue_alias_sizeof();
+
+    for (int i = 0; i < ctx->n_threads; i++) {
+        dma_ptr_curr = (uint8_t *) hex_align_up((uintptr_t) dma_ptr_curr, dma_queue_alignof());
+        ctx->dma_cached[i] = dma_queue_init(dma_ptr_curr, 256, (uintptr_t) ctx->vtcm_base, ctx->vtcm_size, &ctx->trace[i]);
+        dma_ptr_curr += size_dma_q;
+
+        dma_ptr_curr = (uint8_t *) hex_align_up((uintptr_t) dma_ptr_curr, dma_queue_alignof());
+        ctx->dma[i] = dma_queue_alias_init(dma_ptr_curr, ctx->dma_cached[i], 1);
+        dma_ptr_curr += size_dma_alias;
+    }
+
     ctx->ddr_spad_size = 512 * 1024; // 512 KB
     ctx->ddr_spad_base = memalign(128, ctx->ddr_spad_size);
 
-    // init worker pool
-    err = worker_pool_init(&ctx->worker_pool, n_hvx);
-    if (err != AEE_SUCCESS) {
-        FARF(ERROR, "Unable to create worker pool");
-        if (ctx->ddr_spad_base) {
-            free(ctx->ddr_spad_base);
-            ctx->ddr_spad_base = NULL;
-            ctx->ddr_spad_size = 0;
+    void * wq_ptr = (void *) ((uintptr_t) block + offset_wq);
+    ctx->work_queue = work_queue_init(wq_ptr, n_hvx, WORK_QUEUE_CAPACITY, WORK_QUEUE_STACK_SIZE);
+
+    ctx->main_stack = NULL;
+    ctx->main_thread = 0;
+    atomic_store(&ctx->killed, false);
+
+    if (!use_callbacks) {
+        // Start main compute thread
+        ctx->main_stack = (void *) ((uintptr_t) block + offset_main_stack);
+
+        qurt_thread_attr_t attr;
+        qurt_thread_attr_init(&attr);
+        qurt_thread_attr_set_stack_addr(&attr, ctx->main_stack);
+        qurt_thread_attr_set_stack_size(&attr, size_main_stack);
+        qurt_thread_attr_set_priority(&attr, main_prio);
+        qurt_thread_attr_set_name(&attr, "htp-main");
+
+        int err_thread = qurt_thread_create(&ctx->main_thread, &attr, htp_main_thread, ctx);
+        if (err_thread) {
+            FARF(ERROR, "Unable to create htp main thread: %d", err_thread);
+            htp_iface_stop(handle);
+            return AEE_ENOMEMORY;
         }
-        return err;
     }
 
     FARF(HIGH, "session %u started: n-hvx %u vtcm-size %zu vtcm-rctx %u n-threads %u thread-id %d thread-prio %d \n",
@@ -452,35 +570,34 @@ AEEResult htp_iface_start(remote_handle64 handle, uint32_t sess_id, uint64_t dsp
 }
 
 AEEResult htp_iface_stop(remote_handle64 handle) {
-    struct htp_context * ctx = (struct htp_context *) handle;
-    if (!ctx) {
+    struct htp_handle * h = (struct htp_handle *) handle;
+    if (!h || !h->ctx) {
         return AEE_EBADPARM;
     }
+    struct htp_context * ctx = h->ctx;
 
-    if (!ctx->queue) {
-        FARF(ERROR, "Queue not open");
-        return AEE_EBADSTATE;
+    if (ctx->main_thread) {
+        atomic_store(&ctx->killed, true);
+        int status;
+        (void) qurt_thread_join(ctx->main_thread, &status);
+        ctx->main_thread = 0;
     }
 
-    // Close queue. dspqueue_close() will also wait for callbacks to finish.
-    int err    = dspqueue_close(ctx->queue);
-    ctx->queue = NULL;
+    int err = dspqueue_close(ctx->dsp_queue); ctx->dsp_queue = NULL;
     if (err != 0) {
         FARF(ERROR, "Queue close failed with 0x%08x", (unsigned) err);
         return err;
     }
 
-    if (ctx->worker_pool) {
-        // Release worker pool
-        worker_pool_release(&ctx->worker_pool);
-    }
+    work_queue_free(ctx->work_queue);
 
     for (int i = 0; i < ctx->n_threads; i++) {
-        dma_queue_delete(ctx->dma[i]);
+        dma_queue_alias_free(ctx->dma[i]);
+        dma_queue_free(ctx->dma_cached[i]);
     }
 
     if (ctx->hmx_queue) {
-        hmx_queue_delete(ctx->hmx_queue);
+        hmx_queue_free(ctx->hmx_queue);
         ctx->hmx_queue = NULL;
     }
     ctx->hmx_enabled = false;
@@ -493,6 +610,9 @@ AEEResult htp_iface_stop(remote_handle64 handle) {
         ctx->ddr_spad_size = 0;
     }
 
+    free(ctx);
+    h->ctx = NULL;
+
     return AEE_SUCCESS;
 }
 
@@ -598,6 +718,7 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_RMS_NORM:
         case HTP_OP_RMS_NORM_MUL:
         case HTP_OP_SCALE:
+        case HTP_OP_CLAMP:
         case HTP_OP_SQR:
         case HTP_OP_SQRT:
         case HTP_OP_UNARY_SOFTPLUS:
@@ -671,8 +792,6 @@ static int execute_op(struct htp_ops_context * octx) {
 
         case HTP_OP_INVALID:
             break;
-
-        // No default to catch missing cases
     }
 
     FARF(ERROR, "Unknown Op %u", octx->op);
@@ -778,12 +897,14 @@ static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uin
     }
 }
 
-static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, uint32_t idx, struct htp_tensor *t) {
+static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t idx, struct htp_tensor *t) {
     uint32_t offset = t->data;
     uint32_t size   = t->size;
     uint32_t bi     = t->bi;
+    uint32_t alias  = t->alias;
 
-    t->data = bufs[bi].base + offset; // update data to the actual pointer
+    t->data  = (uint32_t) (bufs[bi].base + offset);  // update data to the actual pointer
+    t->alias = (uint32_t) (tens + alias);            // update alias to the actual pointer
 
     FARF(HIGH, "prep-tensor #%u: bi %u offset %u size %u data %p : %u:%u:%u:%u", idx, t->bi, offset, t->size, (void*) t->data,
         t->ne[0], t->ne[1], t->ne[3], t->ne[3]);
@@ -791,7 +912,7 @@ static void prep_tensor(struct htp_context *ctx, struct htp_buf_desc *bufs, uint
 
 static void prep_tensors(struct htp_context *ctx, struct htp_buf_desc *bufs, struct htp_tensor *tens, uint32_t n_tens) {
     for (uint32_t i=0; i < n_tens; i++) {
-        prep_tensor(ctx, bufs, i, tens + i);
+        prep_tensor(ctx, bufs, tens, i, tens + i);
     }
 }
 
@@ -805,29 +926,36 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, u
 
     // Prep input tensors
     for (uint32_t i=0; i<HTP_OP_MAX_INPUTS; i++) {
-        struct htp_tensor *src = op->src[i] == 0xffff ? NULL : tens + op->src[i];
-
-        octx->src[i] = src;
-        if (!src) continue;
-
-        if (!(src->flags & HTP_TENSOR_FLUSHED) && (src->flags & HTP_TENSOR_COMPUTE)) {
-            // flush compute buffers on input
-            hex_l2flush((void *) src->data, src->size);
+        uint16_t src_idx = op->src[i];
+        if (src_idx == 0xffff) {
+            octx->src[i]     = NULL;
+            octx->src_dma[i] = NULL;
+            continue;
         }
 
+        struct htp_tensor *src = tens + src_idx;
+        octx->src[i]     = src;
+        octx->src_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma;
+
         FARF(HIGH, "prep-src #%u: data %p size %u : %u:%u:%u:%u", op->src[i], (void*) src->data, src->size,
             src->ne[0], src->ne[1], src->ne[3], src->ne[3]);
     }
 
+    htp_tensor_flush_all(octx->ctx, octx->src, HTP_OP_MAX_INPUTS);
+
     // Prep output tensors
     for (uint32_t i = 0; i < HTP_OP_MAX_OUTPUTS; i++) {
         uint16_t dst_idx = op->dst[i];
         if (dst_idx == 0xffff) {
-            octx->dsts[i] = NULL;
+            octx->dsts[i]    = NULL;
+            octx->dst_dma[i] = NULL;
             continue;
         }
         struct htp_tensor *dst = tens + dst_idx;
-        octx->dsts[i] = dst;
+        octx->dsts[i]    = dst;
+        octx->dst_dma[i] = octx->ctx->dma; // FIXME: ? octx->ctx->dma_cached : octx->ctx->dma;
+
+        htp_tensor_make_dirty(dst, octx->ctx->dirty_map);
 
         FARF(HIGH, "prep-dst[%u] #%u: data %p size %u : %u:%u:%u:%u", i, dst_idx, (void*) dst->data, dst->size,
             dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
@@ -841,34 +969,134 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_tensor *tens, u
     octx->src3_spad.src = NULL;
     octx->dst_spad.src  = NULL;
 
-    // flush buffers on output
-    for (uint32_t i = 0; i < HTP_OP_MAX_OUTPUTS; i++) {
-        if (octx->dsts[i]) {
-            struct htp_tensor *dst = (struct htp_tensor *)octx->dsts[i];
-            hex_l2flush((void *) dst->data, dst->size);
-            dst->flags |= HTP_TENSOR_FLUSHED;
-
-            FARF(HIGH, "post-dst[%u] #%u: data %p size %u : %u:%u:%u:%u", i, op->dst[i], (void*) dst->data, dst->size,
-                dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
-        }
-    }
-
     return status;
 }
 
+static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_req * req, const struct dspqueue_buffer * dbuf) {
+    dspqueue_t queue = ctx->dsp_queue;
+    int err;
+
+    const uint32_t n_bufs = req->n_bufs;
+    const uint32_t n_tens = req->n_tensors;
+    const uint32_t n_ops  = req->n_ops;
+
+    const uint32_t b_size = sizeof(struct htp_buf_desc)  * n_bufs;
+    const uint32_t t_size = sizeof(struct htp_tensor)    * n_tens;
+    const uint32_t o_size = sizeof(struct htp_op_desc)   * n_ops;
+    const uint32_t p_size = sizeof(struct htp_prof_desc) * n_ops;
+    const uint32_t tr_size = (HTP_MAX_NTHREADS + 1) * req->n_traces * sizeof(struct htp_trace_desc);
+
+    if (dbuf->size < b_size + t_size + o_size + p_size + tr_size) {
+        FARF(ERROR, "invalid opbatch memory block size %u (req %u)", dbuf->size, b_size + t_size + o_size + p_size + tr_size);
+        return;
+    }
+
+    FARF(HIGH, "processing opbatch #%u: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", req->id,
+            n_bufs, n_tens, n_ops, req->n_traces, dbuf->size, b_size, t_size, o_size);
+
+    // Clean cache at the start of the batch
+    // We cant trace this part because the trace buffer is setup later
+    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
+    hex_l2fetch_block(ctx, ctx->footprint);
+    bitmap_reset(ctx->dirty_map, HTP_OP_MAX_TENSORS);
+
+    // Setup descriptor pointers
+    uint8_t * m_ptr = dbuf->ptr;
+    struct htp_buf_desc* bufs = (struct htp_buf_desc*)  m_ptr; m_ptr += b_size;
+    struct htp_tensor*   tens = (struct htp_tensor*)    m_ptr; m_ptr += t_size;
+    struct htp_op_desc*   ops = (struct htp_op_desc*)   m_ptr; m_ptr += o_size;
+    struct htp_prof_desc* pds = (struct htp_prof_desc*) m_ptr;
+
+    prep_op_bufs(ctx, bufs, n_bufs);
+    prep_tensors(ctx, bufs, tens, n_tens);
+
+    struct htp_ops_context *octx = &ctx->octx;
+    memset(octx, 0, sizeof(*octx));
+    octx->n_threads = ctx->n_threads;
+    octx->ctx       = ctx;
+
+    memset(ctx->trace, 0, sizeof(ctx->trace));
+    if (ctx->profiler == HTP_PROF_TRACE) {
+        struct htp_trace_desc * trace_events = (struct htp_trace_desc *) (m_ptr + p_size);
+        for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
+            ctx->trace[t].events     = &trace_events[t * req->n_traces];
+            ctx->trace[t].max_events = req->n_traces;
+        }
+    }
+
+    work_queue_wakeup(ctx->work_queue);
+    if (ctx->hmx_queue) {
+        hmx_queue_wakeup(ctx->hmx_queue);
+    }
+
+    int op_status = HTP_STATUS_OK;
+    for (uint32_t i = 0; i < n_ops && op_status == HTP_STATUS_OK; i++) {
+        struct profile_data prof;
+
+        profile_start(ctx->profiler, &prof);
+
+        op_status = proc_op_req(octx, tens, i, &ops[i]);
+
+        profile_stop(ctx->profiler, &prof);
+
+        if (ctx->profiler) {
+            pds[i].opcode = ops[i].opcode;
+            pds[i].usecs  = prof.usecs;
+            pds[i].cycles_start = prof.cycles_start;
+            pds[i].cycles_stop = prof.cycles_stop;
+            for (int j = 0; j < HEX_NUM_PMU_COUNTERS; j++) {
+                pds[i].pmu[j] = prof.pmu_counters[j];
+            }
+        }
+    }
+
+    if (ctx->hmx_queue) {
+        hmx_queue_suspend(ctx->hmx_queue);
+        hmx_queue_flush(ctx->hmx_queue);
+    }
+    work_queue_suspend(ctx->work_queue);
+
+    struct htp_opbatch_rsp rsp;
+    memset(&rsp, 0, sizeof(rsp));
+    rsp.id        = req->id;
+    rsp.status    = op_status;
+    rsp.n_bufs    = n_bufs;
+    rsp.n_tensors = n_tens;
+    rsp.n_ops     = n_ops;
+
+    if (ctx->profiler == HTP_PROF_TRACE) {
+        for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
+            rsp.n_traces[t] = ctx->trace[t].count;
+        }
+    }
+
+    struct dspqueue_buffer write_dbuf = *dbuf;
+    write_dbuf.flags = DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT;
+
+    // Flush remaining dirty tensors at the end of the batch
+    htp_trace_event_start(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
+    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
+    htp_trace_event_stop(&ctx->trace[0], HTP_TRACE_EVT_L2FLUSH, 0);
+
+    err = dspqueue_write(queue, 0, 1, &write_dbuf, sizeof(rsp), (const uint8_t *) &rsp, DSPQUEUE_TIMEOUT_NONE);
+    if (err != 0) {
+        FARF(ERROR, "dspqueue_write failed: 0x%08x", (unsigned) err);
+    }
+}
+
+#define DSPQUEUE_READ_TIMEOUT_USEC 5000
 #define DSPQUEUE_POLL_TIMEOUT_USEC 100
 #define DSPQUEUE_POLL_COUNT        100
 
-static void htp_packet_callback(dspqueue_t queue, int error, void * context) {
-    struct htp_context * ctx = (struct htp_context *) context;
-
+static void process_ops(struct htp_context * ctx) {
+    dspqueue_t queue = ctx->dsp_queue;
     int err;
 
     uint32_t poll_count = DSPQUEUE_POLL_COUNT;
 
     vtcm_acquire(ctx);
 
-    while (!ctx->vtcm_needs_release) {
+    while (!ctx->vtcm_needs_release && !atomic_load(&ctx->killed)) {
         struct htp_opbatch_req req;
         uint32_t r_size = sizeof(req);
 
@@ -898,111 +1126,41 @@ static void htp_packet_callback(dspqueue_t queue, int error, void * context) {
         // Reset poll count for valid requests
         poll_count = DSPQUEUE_POLL_COUNT;
 
-        const uint32_t n_bufs = req.n_bufs;
-        const uint32_t n_tens = req.n_tensors;
-        const uint32_t n_ops  = req.n_ops;
-
-        const uint32_t b_size = sizeof(struct htp_buf_desc)  * n_bufs;
-        const uint32_t t_size = sizeof(struct htp_tensor)    * n_tens;
-        const uint32_t o_size = sizeof(struct htp_op_desc)   * n_ops;
-        const uint32_t p_size = sizeof(struct htp_prof_desc) * n_ops;
-        const uint32_t tr_size = (HTP_MAX_NTHREADS + 1) * req.n_traces * sizeof(struct htp_trace_desc);
-
-        if (dbuf.size < b_size + t_size + o_size + p_size + tr_size) {
-            FARF(ERROR, "invalid opbatch memory block size %u (req %u)", dbuf.size, b_size + t_size + o_size + p_size + tr_size);
-            break;
-        }
-
-        FARF(HIGH, "processing opbatch #%u: n-bufs %u n-tensors %u n-ops %u n-traces %u : m-size %u b-size %u t-size %u o-size %u", req.id,
-                n_bufs, n_tens, n_ops, req.n_traces, dbuf.size, b_size, t_size, o_size);
-
-        // Setup descriptor pointers
-        uint8_t * m_ptr = dbuf.ptr;
-        struct htp_buf_desc* bufs = (struct htp_buf_desc*)  m_ptr; m_ptr += b_size;
-        struct htp_tensor*   tens = (struct htp_tensor*)    m_ptr; m_ptr += t_size;
-        struct htp_op_desc*   ops = (struct htp_op_desc*)   m_ptr; m_ptr += o_size;
-        struct htp_prof_desc* pds = (struct htp_prof_desc*) m_ptr;
-
-        prep_op_bufs(ctx, bufs, n_bufs);
-        prep_tensors(ctx, bufs, tens, n_tens);
-
-        struct htp_ops_context *octx = &ctx->octx;
-        memset(octx, 0, sizeof(*octx));
-        octx->n_threads = ctx->n_threads;
-        octx->ctx       = ctx;
-
-        if (ctx->profiler == HTP_PROF_TRACE) {
-            memset(ctx->trace, 0, sizeof(ctx->trace));
-            struct htp_trace_desc * trace_events = (struct htp_trace_desc *) (m_ptr + p_size);
-            for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
-                ctx->trace[t].events = &trace_events[t * req.n_traces];
-                ctx->trace[t].max_events = req.n_traces;
-            }
-        } else {
-            for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
-                ctx->trace[t].events = NULL;
-                ctx->trace[t].max_events = 0;
-            }
-        }
-
-        int      op_status = HTP_STATUS_OK;
-        uint32_t op_wakeup = n_ops / 2; // half-way throgh the batch
-
-        hmx_queue_wakeup(ctx->hmx_queue);
-
-        for (uint32_t i=0; i < n_ops; i++) {
-            struct profile_data prof;
-
-            if (i == op_wakeup) {
-                dspqueue_write_early_wakeup_noblock(queue, 0, 0);
-            }
-
-            profile_start(ctx->profiler, &prof);
-
-            op_status = proc_op_req(octx, tens, i, &ops[i]);
-
-            profile_stop(ctx->profiler, &prof);
-
-            if (op_status != HTP_STATUS_OK) {
-                break;
-            }
-
-            if (ctx->profiler) {
-                pds[i].opcode = ops[i].opcode;
-                pds[i].usecs  = prof.usecs;
-                pds[i].cycles_start = prof.cycles_start;
-                pds[i].cycles_stop = prof.cycles_stop;
-                for (int j = 0; j < HEX_NUM_PMU_COUNTERS; j++) {
-                    pds[i].pmu[j] = prof.pmu_counters[j];
-                }
-            }
-        }
-
-        hmx_queue_suspend(ctx->hmx_queue);
-
-        struct htp_opbatch_rsp rsp;
-        rsp.id        = req.id;
-        rsp.status    = op_status;
-        rsp.n_bufs    = n_bufs;
-        rsp.n_tensors = n_tens;
-        rsp.n_ops     = n_ops;
-        memset(rsp.pad, 0, sizeof(rsp.pad));
-        if (ctx->profiler == HTP_PROF_TRACE) {
-            for (int t = 0; t <= HTP_MAX_NTHREADS; t++) {
-                rsp.n_traces[t] = ctx->trace[t].count;
-            }
-        } else {
-            memset(rsp.n_traces, 0, sizeof(rsp.n_traces));
-        }
-
-        dbuf.flags = DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT;
-
-        err = dspqueue_write(queue, 0, 1, &dbuf, sizeof(rsp), (const uint8_t *) &rsp, DSPQUEUE_TIMEOUT_NONE);
-        if (err != 0) {
-            FARF(ERROR, "dspqueue_write failed: 0x%08x", (unsigned) err);
-            break;
-        }
+        process_opbatch(ctx, &req, &dbuf);
     }
 
     vtcm_release(ctx);
 }
+
+static void htp_packet_callback(dspqueue_t queue, int error, void * context) {
+    (void) queue;
+    (void) error;
+    struct htp_handle * h = (struct htp_handle *) context;
+    if (h && h->ctx) {
+        process_ops(h->ctx);
+    }
+}
+
+static void htp_main_thread(void * context) {
+    struct htp_context * ctx = (struct htp_context *) context;
+
+    FARF(HIGH, "htp-main-thread: started");
+
+    while (!atomic_load(&ctx->killed)) {
+        uint32_t flags = 0;
+        uint32_t num_buffers = 0;
+        uint32_t message_length = 0;
+
+        int err = dspqueue_peek(ctx->dsp_queue, &flags, &num_buffers, &message_length, 50000);
+        if (err == 0) {
+            process_ops(ctx);
+        } else if (err == AEE_EWOULDBLOCK || err == AEE_EEXPIRED) {
+            continue;
+        } else {
+            FARF(ERROR, "dspqueue_peek failed: 0x%08x", (unsigned) err);
+            break;
+        }
+    }
+
+    FARF(HIGH, "htp-main-thread: stopped");
+}

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

@@ -92,10 +92,10 @@ struct htp_mm_context {
     // Per thread quant tasks
     // Precomputed block-parallel quantization values
     worker_callback_t quant_task_func;
-    uint32_t          quant_ib_first[MAX_NUM_WORKERS];
-    uint32_t          quant_ib_last[MAX_NUM_WORKERS];
-    uint32_t          quant_r[MAX_NUM_WORKERS];
-    uint32_t          quant_c[MAX_NUM_WORKERS];
+    uint32_t          quant_ib_first[WORK_QUEUE_MAX_N_THREADS];
+    uint32_t          quant_ib_last[WORK_QUEUE_MAX_N_THREADS];
+    uint32_t          quant_r[WORK_QUEUE_MAX_N_THREADS];
+    uint32_t          quant_c[WORK_QUEUE_MAX_N_THREADS];
     uint32_t          n_quant_tasks;
     uint32_t          n_quant_rows_per_thread;
     atomic_uint       quant_barrier;
@@ -254,7 +254,7 @@ static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) {
         return;
     }
 
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0_start);
 
     const uint32_t blck_0 = 64;
@@ -309,7 +309,7 @@ static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
     const uint32_t src0_start_row  = src0_nrows_per_thread * ith;                                                                 \
     const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);                                     \
                                                                                                                                   \
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;                                                     \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                        \
                                                                                                                                   \
     const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;                      \
     const uint32_t n_prefetch = kparams->n_prefetch;                                                                              \
@@ -410,7 +410,7 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void
     const uint32_t src0_start_row  = src0_nrows_per_thread * ith;                                                                 \
     const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);                                     \
                                                                                                                                   \
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;                                                     \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                        \
                                                                                                                                   \
     const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;                      \
     const uint32_t n_prefetch = kparams->n_prefetch;                                                                              \
@@ -523,7 +523,7 @@ static void hvx_mm_qkv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith,
     uint8_t * restrict vtcm_src3_ptr = mmctx->vtcm_src3 + mmctx->vtcm_src3_size_per_thread * ith;                                 \
     uint8_t * restrict src1_data = mmctx->vtcm_src1;                                                                              \
                                                                                                                                   \
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;                                                     \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                        \
                                                                                                                                   \
     const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;                      \
     const uint32_t n_prefetch = kparams->n_prefetch;                                                                              \
@@ -699,7 +699,7 @@ static void hvx_mm_ffn_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith,
     uint8_t * restrict vtcm_src2_ptr = mmctx->vtcm_src2 + mmctx->vtcm_src2_size_per_thread * ith;                                 \
     uint8_t * restrict src1_data = mmctx->vtcm_src1;                                                                              \
                                                                                                                                   \
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;                                                     \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                        \
                                                                                                                                   \
     const uint8_t * restrict src0_row = (const uint8_t *) src0->data;                                                             \
     const uint8_t * restrict src2_row = (const uint8_t *) src2->data;                                                             \
@@ -820,7 +820,7 @@ static void name(unsigned int nth, unsigned int ith, void * data) {
         return;                                                                                            \
     }                                                                                                      \
                                                                                                            \
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;                              \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                 \
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first);                                        \
                                                                                                            \
     uint8_t * restrict dst = mmctx->vtcm_src1;                                                             \
@@ -846,7 +846,7 @@ QUANTIZE_IMPL(quantize_f16_f16_flat,   "quantize-f16-f16",        quantize_f16_f
 static void quantize_f32_q8_0_tiled_block(unsigned int nth, unsigned int ith, void * data) {
     struct htp_mm_context * mmctx = data;
     struct htp_ops_context * octx = mmctx->octx;
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]);
 
     const struct htp_tensor * src = octx->src[1];
@@ -870,7 +870,7 @@ static void quantize_f32_q8_0_tiled_block(unsigned int nth, unsigned int ith, vo
 static void quantize_f32_q8_1_tiled_block(unsigned int nth, unsigned int ith, void * data) {
     struct htp_mm_context * mmctx = data;
     struct htp_ops_context * octx = mmctx->octx;
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]);
 
     const struct htp_tensor * src = octx->src[1];
@@ -944,7 +944,7 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) {
     const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
     const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U);
 
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
     const size_t dst_row_size  = nb1;
     const size_t src0_row_size = nb01;
@@ -1040,7 +1040,7 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) {
     const uint32_t src0_start_row  = src0_nrows_per_thread * ith;
     const uint32_t src0_end_row    = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows);
 
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
     const size_t dst_row_size  = nb1;
     const size_t src0_row_size = nb01;
@@ -1155,7 +1155,7 @@ static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) {
         return;
     }
 
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
     const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
     const uint32_t n_prefetch = kparams->n_prefetch;
@@ -1244,7 +1244,7 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) {
         return;
     }
 
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
 
     const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
     const uint32_t n_prefetch = kparams->n_prefetch;
@@ -1338,6 +1338,9 @@ static int hvx_mm_init_vec_dot(struct htp_mm_context * mmctx, enum htp_data_type
 static int hvx_mm_matmul(struct htp_ops_context * octx) {
     htp_matmul_tensors_preamble;
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[0];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
     struct htp_mm_context mmctx_struct = {0};
     struct htp_mm_context * mmctx = &mmctx_struct;
     mmctx->octx = octx;
@@ -1557,9 +1560,6 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
     mmctx->vtcm_src0_stride = src0_row_size_padded;
     mmctx->vtcm_src1_stride = src1_row_size;
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)
-        return HTP_STATUS_OK;
-
     if (need_quant) {
         mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
         mmctx->quant_task_func = quant_task_func;
@@ -1570,8 +1570,9 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
         mmctx->n_quant_tasks = 0;
     }
 
-    const uint32_t n_matmul_jobs = octx->n_threads;
-    worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, n_matmul_jobs);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
+    worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, octx->n_threads);
 
     return HTP_STATUS_OK;
 }
@@ -1874,7 +1875,7 @@ static void hvx_mm_ffn_2d(unsigned int nth, unsigned int ith, void * data) {
 #define DEQUANTIZE_WORKER_LOOP_IMPL(SUFFIX)                                                     \
 static void dequantize_tiled_worker_loop_##SUFFIX(unsigned int n, unsigned int i, void *data) { \
     tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data;                         \
-    struct htp_thread_trace * tr = state->traces ? &state->traces[i] : NULL;                    \
+    struct htp_thread_trace * tr = &state->traces[i];                                           \
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_W_DEQUANT, i);                                  \
     for (unsigned int task_id = i; task_id < (unsigned int)state->n_tasks; task_id += n) {      \
         int start = task_id * state->n_tiles_per_task;                                          \
@@ -1892,7 +1893,7 @@ DEQUANTIZE_WORKER_LOOP_IMPL(q8_0)
 
 static void convert_f16_worker_loop(unsigned int n, unsigned int i, void *data) {
     tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data;
-    struct htp_thread_trace * tr = state->traces ? &state->traces[i] : NULL;
+    struct htp_thread_trace * tr = &state->traces[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_W_DEQUANT, i);
     for (unsigned int task_id = i; task_id < (unsigned int)state->n_tasks; task_id += n) {
         int start = task_id * state->n_tiles_per_task;
@@ -1905,7 +1906,7 @@ static void convert_f16_worker_loop(unsigned int n, unsigned int i, void *data)
 static void quantize_f32_worker_loop(unsigned int n, unsigned int i, void *data) {
     tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data;
 
-    struct htp_thread_trace * tr = state->traces ? &state->traces[i] : NULL;
+    struct htp_thread_trace * tr = &state->traces[i];
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, i);
 
     for (unsigned int task_id = i; task_id < (unsigned int)state->n_tasks; task_id += n) {
@@ -1920,7 +1921,7 @@ static void quantize_f32_worker_loop(unsigned int n, unsigned int i, void *data)
 static void transfer_output_chunk_worker_fn(unsigned int n, unsigned int i, void *data) {
     output_transfer_task_state_t *st = (output_transfer_task_state_t *) data;
 
-    struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL;
+    struct htp_thread_trace * tr = &st->traces[i];
 
     int start_chunk_idx = i * st->n_chunks_per_task;
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, start_chunk_idx);
@@ -1955,6 +1956,170 @@ typedef struct {
     uint32_t             dma_step_rows_shift;
 } activation_transfer_task_state_t;
 
+typedef struct {
+    __fp16                         *dst;
+    const float                    *src;
+    uint32_t                        n_rows;
+    uint32_t                        k_block;
+    uint32_t                        k_stride;
+    uint32_t                        k_valid;
+    uint32_t                        n_col_chunks;
+    struct fastdiv_values           n_threads_div;
+    float                          *vtcm_f32_act;
+    size_t                          vtcm_f32_act_bytes;
+    struct htp_thread_trace        *traces;
+    struct htp_context             *ctx;
+    uint32_t                        dma_step_rows;
+    uint32_t                        dma_step_rows_shift;
+} activation_transfer_col_chunk_state_t;
+
+static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk(
+        dma_queue *dma_q,
+        __fp16 *restrict vtcm_dst,
+        const float *restrict src,
+        uint32_t n_rows,
+        uint32_t k_block,
+        uint32_t k_stride,
+        uint32_t k_chunk_valid,
+        uint32_t c_first,
+        uint32_t c_len,
+        float *thread_f32_act,
+        struct htp_thread_trace *tr,
+        uint32_t dma_step_rows,
+        uint32_t dma_step_rows_shift) {
+
+    const uint32_t R = dma_step_rows;
+    const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS);
+
+    const uint32_t n_steps = n_rows_padded >> dma_step_rows_shift;
+
+    // Push step 0
+    if (n_steps > 0 && n_rows > 0) {
+        uint32_t nrows_to_fetch = hex_smin(n_rows, R);
+        dma_queue_push(dma_q, dma_make_ptr(thread_f32_act, src + c_first),
+                       c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch);
+    }
+    // Push step 1
+    if (n_steps > 1) {
+        uint32_t next_r = R * 1;
+        if (next_r < n_rows) {
+            uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R);
+            const float *next_src = src + next_r * k_stride + c_first;
+            float *next_buf = thread_f32_act + 1 * R * c_len;
+            dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src),
+                           c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch);
+        }
+    }
+    for (uint32_t s = 0; s < n_steps; ++s) {
+        uint32_t r = s << dma_step_rows_shift;
+        float *curr_buf = thread_f32_act;
+
+        if (r < n_rows) {
+            curr_buf = (float *) dma_queue_pop(dma_q).dst;
+        }
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, r);
+        for (uint32_t p = 0; p < (R >> 1); ++p) {
+            uint32_t row_idx = r + (p << 1);
+            float *pair_buf = curr_buf + (p << 1) * c_len;
+            bool r0_valid = ((row_idx + 0) < n_rows);
+            bool r1_valid = ((row_idx + 1) < n_rows);
+
+            transfer_activation_row_pair_fp32_to_fp16_col_chunk(
+                vtcm_dst, pair_buf, pair_buf + c_len, row_idx, k_block, c_first, c_len, k_chunk_valid, r0_valid, r1_valid
+            );
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, r);
+
+        // Push step s + 2
+        uint32_t next_s = s + 2;
+        uint32_t next_r = next_s << dma_step_rows_shift;
+        if (next_r < n_rows) {
+            uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R);
+            const float *next_src = src + next_r * k_stride + c_first;
+            dma_queue_push(dma_q, dma_make_ptr(curr_buf, next_src),
+                           c_len * sizeof(float), k_stride * sizeof(float), k_chunk_valid * sizeof(float), nrows_to_fetch);
+        }
+    }
+}
+
+static void transfer_activation_chunk_fp32_to_fp16_col_chunk(
+        __fp16 *restrict vtcm_dst,
+        const float *restrict src,
+        uint32_t n_rows,
+        uint32_t k_block,
+        uint32_t k_stride,
+        uint32_t c_first,
+        uint32_t c_len,
+        uint32_t k_chunk_valid) {
+    const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS);
+    const uint32_t n_rows_tiled  = (n_rows / HTP_MM_HMX_TILE_N_ROWS) * HTP_MM_HMX_TILE_N_ROWS;
+
+    uint32_t r = 0;
+
+    #pragma unroll(2)
+    for (r = 0; r < n_rows_tiled; r += 2) {
+        const float *ptr_in0 = src + (r + 0) * k_stride + c_first;
+        const float *ptr_in1 = src + (r + 1) * k_stride + c_first;
+
+        transfer_activation_row_pair_fp32_to_fp16_col_chunk(
+            vtcm_dst, ptr_in0, ptr_in1, r, k_block, c_first, c_len, k_chunk_valid, true, true
+        );
+    }
+
+    for (; r < n_rows_padded; r += 2) {
+        const bool row0_valid = r       < n_rows;
+        const bool row1_valid = (r + 1) < n_rows;
+
+        const float *ptr_in0 = row0_valid ? (src + (r + 0) * k_stride + c_first) : NULL;
+        const float *ptr_in1 = row1_valid ? (src + (r + 1) * k_stride + c_first) : NULL;
+
+        transfer_activation_row_pair_fp32_to_fp16_col_chunk(
+            vtcm_dst, ptr_in0, ptr_in1, r, k_block, c_first, c_len, k_chunk_valid, row0_valid, row1_valid
+        );
+    }
+}
+
+static void transfer_activation_chunk_col_chunk_worker_fn(unsigned int n, unsigned int i, void *data) {
+    activation_transfer_col_chunk_state_t *st = (activation_transfer_col_chunk_state_t *) data;
+    struct htp_thread_trace * tr = &st->traces[i];
+
+    uint32_t n_blocks = st->k_block / 32;
+    uint32_t b_first = fastdiv(n_blocks * i, &st->n_threads_div);
+    uint32_t b_last  = fastdiv(n_blocks * (i + 1), &st->n_threads_div);
+    uint32_t c_first = b_first * 32;
+    uint32_t c_last = b_last * 32;
+    uint32_t c_len = c_last - c_first;
+
+    if (c_len == 0) {
+        return;
+    }
+
+    uint32_t k_chunk_valid = 0;
+    if (st->k_valid > c_first) {
+        k_chunk_valid = hex_smin(st->k_valid, c_last) - c_first;
+    }
+
+    __fp16 *dst = st->dst;
+    const float *src = st->src;
+
+    if (st->vtcm_f32_act) {
+        size_t thread_scratch_bytes = hex_align_down(fastdiv(st->vtcm_f32_act_bytes, &st->n_threads_div), 128);
+        float *thread_f32_act = (float *)((char *)st->vtcm_f32_act + i * thread_scratch_bytes);
+
+        transfer_activation_chunk_fp32_to_fp16_dma_pipelined_col_chunk(
+            st->ctx->dma[i], dst, src, st->n_rows, st->k_block, st->k_stride, k_chunk_valid,
+            c_first, c_len, thread_f32_act, tr, st->dma_step_rows, st->dma_step_rows_shift
+        );
+    } else {
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, c_first);
+        transfer_activation_chunk_fp32_to_fp16_col_chunk(
+            dst, src, st->n_rows, st->k_block, st->k_stride, c_first, c_len, k_chunk_valid
+        );
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, c_first);
+    }
+}
+
 static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined(
         dma_queue *dma_q,
         __fp16 *restrict vtcm_dst,
@@ -2024,7 +2189,7 @@ static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined(
 static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i, void *data) {
     activation_transfer_task_state_t *st = (activation_transfer_task_state_t *) data;
 
-    struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL;
+    struct htp_thread_trace * tr = &st->traces[i];
 
     for (unsigned int task_id = i; task_id < (unsigned int)st->n_tasks; task_id += n) {
         int    chunk_idx  = task_id * st->n_chunks_per_task;
@@ -2085,15 +2250,16 @@ typedef struct {
 
 static void transfer_activation_chunk_gathered_worker_fn(unsigned int n, unsigned int i, void *data) {
     activation_transfer_gathered_task_state_t *st = data;
-    struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL;
+    struct htp_thread_trace * tr = &st->traces[i];
     int chunk_idx = i;
     int chunk_size = st->n_chunks_per_task;
-    int start_row = st->start_row + chunk_idx * chunk_size;
+    int vtcm_start_row = chunk_idx * chunk_size;
+    int start_row = st->start_row + vtcm_start_row;
     int n_rows = hex_smin(st->cne1 - start_row, chunk_size);
     if (n_rows > 0) {
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx);
         transfer_activation_chunk_fp32_to_fp16_gathered(
-            st->dst, st->src, start_row, n_rows, st->k_block,
+            st->dst, st->src, start_row, vtcm_start_row, n_rows, st->k_block,
             st->matrix_rows, st->cur_a, st->mapping_stride,
             st->ne11, &st->ne11_div, st->nb11, st->nb12, st->cne1, st->k_valid);
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx);
@@ -2102,15 +2268,16 @@ static void transfer_activation_chunk_gathered_worker_fn(unsigned int n, unsigne
 
 static void transfer_activation_chunk_gathered_worker_flat_fn(unsigned int n, unsigned int i, void *data) {
     activation_transfer_gathered_task_state_t *st = data;
-    struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL;
+    struct htp_thread_trace * tr = &st->traces[i];
     int chunk_idx = i;
     int chunk_size = st->n_chunks_per_task;
-    int start_row = st->start_row + chunk_idx * chunk_size;
+    int vtcm_start_row = chunk_idx * chunk_size;
+    int start_row = st->start_row + vtcm_start_row;
     int n_rows = hex_smin(st->cne1 - start_row, chunk_size);
     if (n_rows > 0) {
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx);
         transfer_activation_chunk_fp32_to_fp16_gathered_flat(
-            st->dst, st->src, start_row, n_rows, st->k_block,
+            st->dst, st->src, start_row, vtcm_start_row, n_rows, st->k_block,
             st->matrix_rows, st->cur_a, st->mapping_stride,
             st->nb12, st->cne1, st->k_valid);
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx);
@@ -2119,15 +2286,16 @@ static void transfer_activation_chunk_gathered_worker_flat_fn(unsigned int n, un
 
 static void transfer_output_chunk_scattered_worker_fn(unsigned int n, unsigned int i, void *data) {
     output_transfer_scattered_task_state_t *st = data;
-    struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL;
+    struct htp_thread_trace * tr = &st->traces[i];
     int chunk_idx = i;
     int chunk_size = st->n_chunks_per_task;
-    int start_row = st->start_row + chunk_idx * chunk_size;
+    int vtcm_start_row = chunk_idx * chunk_size;
+    int start_row = st->start_row + vtcm_start_row;
     int n_rows = hex_smin(st->cne1 - start_row, chunk_size);
     if (n_rows > 0) {
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_O_PROC, chunk_idx);
         transfer_output_chunk_fp16_to_fp32_scattered(
-            st->dst, st->vtcm_src, start_row, n_rows, st->n_cols,
+            st->dst, st->vtcm_src, start_row, vtcm_start_row, n_rows, st->n_cols,
             st->matrix_rows, st->cur_a, st->mapping_stride,
             st->dst_nb1, st->dst_nb2, st->cne1);
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, chunk_idx);
@@ -2210,42 +2378,81 @@ static void transfer_output_chunk_threaded(struct htp_context *ctx, float *dst,
     }
 }
 
-static void transfer_activation_chunk_threaded(
-        struct htp_context *ctx,
-        __fp16 *dst,
-        const float *src,
-        int n_rows,
-        int k_block,
-        int k_stride,
-        int n_threads,
-        int k_valid,
-        float *vtcm_f32_act,
-        size_t vtcm_f32_act_bytes) {
+struct activation_transfer_params {
+    struct htp_context *          ctx;
+    __fp16 *                      dst;
+    const float *                 src;
+    int                           n_rows;
+    int                           k_block;
+    int                           k_stride;
+    int                           n_threads;
+    const struct fastdiv_values * act_threads_div;
+    const struct fastdiv_values * k_div;
+    int                           k_valid;
+    float *                       vtcm_f32_act;
+    size_t                        vtcm_f32_act_bytes;
+};
+
+static void transfer_activation_chunk_threaded(const struct activation_transfer_params * params) {
+    struct htp_context *          ctx                = params->ctx;
+    __fp16 *                      dst                = params->dst;
+    const float *                 src                = params->src;
+    int                           n_rows             = params->n_rows;
+    int                           k_block            = params->k_block;
+    int                           k_stride           = params->k_stride;
+    int                           n_threads          = params->n_threads;
+    const struct fastdiv_values * act_threads_div    = params->act_threads_div;
+    const struct fastdiv_values * k_div              = params->k_div;
+    int                           k_valid            = params->k_valid;
+    float *                       vtcm_f32_act       = params->vtcm_f32_act;
+    size_t                        vtcm_f32_act_bytes = params->vtcm_f32_act_bytes;
+
     if (n_rows <= 0) {
         return;
     }
 
+    const size_t n_tasks = (n_rows + 31) >> 5;
+    if (n_threads > 1 && k_block > 32 && n_tasks < (size_t) n_threads) {
+        // Calculate step rows parameters for column-chunked dma pipelining
+        uint32_t dma_step_rows = 2;
+        uint32_t dma_step_rows_shift = 1;
+        if (vtcm_f32_act && vtcm_f32_act_bytes > 0 && k_block > 0) {
+            size_t thread_scratch_bytes = hex_align_down(fastdiv(vtcm_f32_act_bytes, act_threads_div), 128);
+            size_t thread_scratch_elements = thread_scratch_bytes / sizeof(float);
+            size_t dma_step_rows_max = fastdiv(thread_scratch_elements / 2, k_div);
+            if (dma_step_rows_max >= 4) {
+                dma_step_rows = 4;
+                dma_step_rows_shift = 2;
+            }
+        }
+
+        activation_transfer_col_chunk_state_t col_state;
+        col_state.dst = dst;
+        col_state.src = src;
+        col_state.n_rows = n_rows;
+        col_state.k_block = k_block;
+        col_state.k_stride = k_stride;
+        col_state.k_valid = k_valid;
+        col_state.n_col_chunks = n_threads;
+        col_state.n_threads_div = *act_threads_div;
+        col_state.vtcm_f32_act = vtcm_f32_act;
+        col_state.vtcm_f32_act_bytes = vtcm_f32_act_bytes;
+        col_state.traces = ctx->trace;
+        col_state.ctx = ctx;
+        col_state.dma_step_rows = dma_step_rows;
+        col_state.dma_step_rows_shift = dma_step_rows_shift;
+
+        worker_pool_run_func(ctx->worker_pool, transfer_activation_chunk_col_chunk_worker_fn, &col_state, n_threads);
+        return;
+    }
+
     assert(k_block % HTP_MM_HMX_TILE_N_COLS == 0 && k_stride % HTP_MM_HMX_TILE_N_COLS == 0);
 
     size_t n_tot_chunks      = n_rows;
     size_t n_chunks_per_task = (n_threads == 1) ? n_tot_chunks : 32;  // must be multiple of 32 to ensure correct destination address
 
-    uint32_t dma_step_rows = 2;
-    uint32_t dma_step_rows_shift = 1;
-    if (vtcm_f32_act && vtcm_f32_act_bytes > 0 && k_block > 0) {
-        size_t thread_scratch_elements = vtcm_f32_act_bytes / (n_threads * sizeof(float));
-        size_t dma_step_rows_max = (thread_scratch_elements / 2) / k_block;
-        if (dma_step_rows_max >= 4) {
-            dma_step_rows = 4;
-            dma_step_rows_shift = 2;
-        } else {
-            dma_step_rows = 2;
-            dma_step_rows_shift = 1;
-        }
-    }
-
     activation_transfer_task_state_t state;
-    state.n_tasks            = (n_tot_chunks + n_chunks_per_task - 1) / n_chunks_per_task;
+    state.n_tasks            = (n_threads == 1) ? 1 : hmx_ceil_div(n_tot_chunks, 32);
     state.n_tot_chunks       = n_tot_chunks;
     state.n_chunks_per_task  = n_chunks_per_task;
     state.dst                = dst;
@@ -2258,7 +2465,18 @@ static void transfer_activation_chunk_threaded(
     state.vtcm_f32_act       = vtcm_f32_act;
 
     int active_threads = hex_smin(n_threads, (int)state.n_tasks);
-    state.vtcm_f32_act_bytes_per_thread = (vtcm_f32_act_bytes / active_threads) & ~127u;
+    state.vtcm_f32_act_bytes_per_thread = hex_align_down(vtcm_f32_act_bytes / active_threads, 128);
+
+    uint32_t dma_step_rows = 2;
+    uint32_t dma_step_rows_shift = 1;
+    if (vtcm_f32_act && state.vtcm_f32_act_bytes_per_thread > 0 && k_block > 0) {
+        size_t thread_scratch_elements = state.vtcm_f32_act_bytes_per_thread / sizeof(float);
+        size_t dma_step_rows_max = fastdiv(thread_scratch_elements / 2, k_div);
+        if (dma_step_rows_max >= 4) {
+            dma_step_rows = 4;
+            dma_step_rows_shift = 2;
+        }
+    }
     state.dma_step_rows      = dma_step_rows;
     state.dma_step_rows_shift = dma_step_rows_shift;
 
@@ -2321,9 +2539,14 @@ static int hmx_mm_2d_f32(struct htp_context *ctx,
                                   int pipeline,
                                   int n_threads,
                                   int act_threads,
+                                  const struct fastdiv_values * act_threads_div,
+                                  const struct fastdiv_values * k_div,
                                   int tile_size,
                                   int aligned_tile_size,
                                   int vtcm_size) {
+    struct htp_thread_trace * tr = &ctx->trace[0];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
     if (k % 32 != 0 || n % 32 != 0) { return -1; }
     if (!hex_is_aligned(dst, VLEN) || !hex_is_aligned(activation, VLEN)) { return -1; }
 
@@ -2393,6 +2616,8 @@ static int hmx_mm_2d_f32(struct htp_context *ctx,
 
     int n_chunk_cnt = hmx_ceil_div(n, n_chunk_n_cols);
 
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
     if (pipeline) {
         // --- Asynchronous Pipelined Loop ---
         hmx_matmul_job_t job_slots[2];  // persistent double-buffered job descriptors
@@ -2403,7 +2628,21 @@ static int hmx_mm_2d_f32(struct htp_context *ctx,
             void *vtcm_weight_bufs[2] = { vtcm_scratch0, vtcm_scratch1 };
             void *vtcm_output_bufs[2] = { vtcm_output,   vtcm_scratch2 };
 
-            transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act, L.act_f32_bytes);
+            struct activation_transfer_params act_params = {
+                .ctx = ctx,
+                .dst = vtcm_f16_act,
+                .src = activation + mr * act_stride,
+                .n_rows = (int) n_rows,
+                .k_block = k,
+                .k_stride = act_stride,
+                .n_threads = act_threads,
+                .act_threads_div = act_threads_div,
+                .k_div = k_div,
+                .k_valid = k_valid,
+                .vtcm_f32_act = vtcm_f32_act,
+                .vtcm_f32_act_bytes = L.act_f32_bytes,
+            };
+            transfer_activation_chunk_threaded(&act_params);
 
             // Prologue: push A0 and optionally A1 (if n_chunk_cnt > 1)
             const size_t   n_cols_A0 = hex_smin(n - 0 * n_chunk_n_cols, n_chunk_n_cols);
@@ -2480,7 +2719,21 @@ static int hmx_mm_2d_f32(struct htp_context *ctx,
         for (size_t mr = 0; mr < m; mr += m_chunk_n_rows) {
             const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows);
 
-            transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act, L.act_f32_bytes);
+            struct activation_transfer_params act_params = {
+                .ctx = ctx,
+                .dst = vtcm_f16_act,
+                .src = activation + mr * act_stride,
+                .n_rows = (int) n_rows,
+                .k_block = k,
+                .k_stride = act_stride,
+                .n_threads = act_threads,
+                .act_threads_div = act_threads_div,
+                .k_div = k_div,
+                .k_valid = k_valid,
+                .vtcm_f32_act = vtcm_f32_act,
+                .vtcm_f32_act_bytes = L.act_f32_bytes,
+            };
+            transfer_activation_chunk_threaded(&act_params);
 
             // A0: Pre-fetch the first weight chunk (nc = 0)
             if (n > 0) {
@@ -2570,7 +2823,8 @@ static inline const float *hmx_mm_src2_batch_ptr(const hmx_mm_f16_f32_batched_pa
 
 static int hmx_mm_f16_f32_batched_simple(struct htp_context *ctx,
                                                         const hmx_mm_f16_f32_batched_params_t *params,
-                                                        int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, int vtcm_size) {
+                                                        int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, int vtcm_size,
+                                                        const struct fastdiv_values * act_threads_div, const struct fastdiv_values * k_div) {
     int ret = 0;
     for (int b3 = 0; b3 < params->ne13 && ret == 0; ++b3) {
         for (int b2 = 0; b2 < params->ne12 && ret == 0; ++b2) {
@@ -2582,14 +2836,17 @@ static int hmx_mm_f16_f32_batched_simple(struct htp_context *ctx,
                                            params->act_stride, params->weight_stride * (int)sizeof(__fp16),
                                            HTP_TYPE_F16, params->k, params->dst_stride, params->src2_stride, params->n,
                                            m_chunk, n_chunk, pipeline, n_threads, act_threads,
-                                           0, 0, vtcm_size);
+                                           act_threads_div, k_div, 0, 0, vtcm_size);
         }
     }
     return ret;
 }
 
 static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_batched_params_t *params,
-                               int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, int vtcm_size) {
+                               int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads,
+                               const struct fastdiv_values * act_threads_div,
+                               const struct fastdiv_values * k_div,
+                               int vtcm_size) {
     if (params->act_stride < params->k || params->weight_stride < params->k || params->dst_stride < params->n) { return -1; }
     if (params->ne02 <= 0 || params->ne03 <= 0 || params->ne12 <= 0 || params->ne13 <= 0) { return -1; }
     if (params->ne12 % params->ne02 != 0 || params->ne13 % params->ne03 != 0) { return -1; }
@@ -2604,9 +2861,12 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_
     // Grouped path is only valid if group_size > 1 and it fits within VTCM budget.
     bool run_grouped = (group_size > 1 && (size_t)vtcm_size <= vtcm_budget);
     if (!run_grouped) {
-        return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size);
+        return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size, act_threads_div, k_div);
     }
 
+    struct htp_thread_trace * tr = &ctx->trace[0];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
     const size_t vec_dot_size = params->k * sizeof(__fp16);
 
     const bool use_dma_activation = (params->act_stride > params->k);
@@ -2622,7 +2882,8 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_
 
     if (L.total_bytes > vtcm_budget) {
         FARF(HIGH, "%s: grouped layout overflowed VTCM, falling back to simple batched loop", __func__);
-        return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+        return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size, act_threads_div, k_div);
     }
 
     uint8_t * const base = (uint8_t *) ctx->vtcm_base;
@@ -2644,6 +2905,8 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_
     const size_t fp16_row_bytes   = (size_t) params->k * sizeof(__fp16);
     const size_t weight_row_bytes = (size_t) params->weight_stride * sizeof(__fp16);
 
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
     hmx_matmul_job_t job;
 
     for (int b3 = 0; b3 < params->ne13; ++b3) {
@@ -2662,9 +2925,21 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_
                 for (int g = 0; g < group_size; ++g) {
                     const float *activation_chunk = hmx_mm_activation_batch_ptr(params, b2_base + g, b3) + mr * params->act_stride;
                     __fp16 *vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride;
-                    transfer_activation_chunk_threaded(ctx, vtcm_act_g,
-                                                           activation_chunk, (int) n_rows,
-                                                           params->k, params->act_stride, act_threads, params->k, vtcm_f32_act, L.act_f32_bytes);
+                    struct activation_transfer_params act_params = {
+                        .ctx = ctx,
+                        .dst = vtcm_act_g,
+                        .src = activation_chunk,
+                        .n_rows = (int) n_rows,
+                        .k_block = params->k,
+                        .k_stride = params->act_stride,
+                        .n_threads = act_threads,
+                        .act_threads_div = act_threads_div,
+                        .k_div = k_div,
+                        .k_valid = params->k,
+                        .vtcm_f32_act = vtcm_f32_act,
+                        .vtcm_f32_act_bytes = L.act_f32_bytes,
+                    };
+                    transfer_activation_chunk_threaded(&act_params);
                 }
 
                 // Prologue: Push A0 and A1 (if exists)
@@ -2835,6 +3110,9 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx,
                                          const struct mmid_row_mapping *matrix_rows,
                                          int cur_a,
                                          int mapping_stride) {
+    struct htp_thread_trace * tr = &ctx->trace[0];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
     const int cne1 = m;
     const int m_padded = hex_align_up(m, 32);
 
@@ -2913,6 +3191,8 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx,
 
     hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00));
 
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
     hmx_matmul_job_t job;
 
     for (size_t mr = 0; mr < (size_t) m_padded; mr += m_chunk_n_rows) {
@@ -2980,10 +3260,6 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k
     const int act_stride = (int)(src1->nb[1] / sizeof(float));
     const int wgt_stride = (int)(src0->nb[1] / sizeof(__fp16));
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
-        return HTP_STATUS_OK;
-    }
-
     const float * src2_ptr = NULL;
     uint32_t src2_stride = 0;
     size_t src2_nb2 = 0;
@@ -3027,6 +3303,8 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k
                                      kparams->m_chunk, kparams->n_chunk,
                                      kparams->pipeline, n_threads,
                                      kparams->n_act_threads,
+                                     &kparams->div_n_act_threads,
+                                     &kparams->div_ne00_padded,
                                      kparams->vtcm_size);
     } else {
         ret = hmx_mm_2d_f32(
@@ -3035,6 +3313,8 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k
             (int)(dst->nb[1] / sizeof(float)), src2_stride, (int)dst->ne[0],
             kparams->m_chunk, kparams->n_chunk, kparams->pipeline, n_threads,
             kparams->n_act_threads,
+            &kparams->div_n_act_threads,
+            &kparams->div_ne00_padded,
             kparams->tile_size, kparams->aligned_tile_size, kparams->vtcm_size
         );
     }
@@ -3103,6 +3383,10 @@ static int hvx_mm_matmul_id(
     bool must_free_mapping
 ) {
     htp_matmul_tensors_preamble;
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[0];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
     const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
     const struct htp_tensor * restrict ids = octx->src[2];
     const size_t src0_row_size = nb01;
@@ -3175,8 +3459,9 @@ static int hvx_mm_matmul_id(
     mmctx->n_quant_tasks = n_quant_tasks;
     atomic_init(&mmctx->quant_barrier, n_quant_tasks);
 
-    const uint32_t n_matmul_jobs = octx->n_threads;
-    worker_pool_run_func(octx->ctx->worker_pool, matmul_id_job_func, mmctx, n_matmul_jobs);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
+    worker_pool_run_func(octx->ctx->worker_pool, matmul_id_job_func, mmctx, octx->n_threads);
 
     if (must_free_mapping) free(mapping_buf);
     return HTP_STATUS_OK;
@@ -3185,6 +3470,9 @@ static int hvx_mm_matmul_id(
 int op_matmul_id(struct htp_ops_context * octx) {
     htp_matmul_tensors_preamble;
 
+    struct htp_thread_trace * tr = &octx->ctx->trace[0];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
     struct htp_mm_context mmctx_struct = {0};
     struct htp_mm_context * mmctx = &mmctx_struct;
     mmctx->octx = octx;
@@ -3262,10 +3550,7 @@ int op_matmul_id(struct htp_ops_context * octx) {
         }
     }
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
-        if (must_free_mapping) free(mapping_buf);
-        return HTP_STATUS_OK;
-    }
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
 
     if (kparams->n_hmx) {
         return hmx_mm_op_matmul_id(octx, mmctx, matrix_row_counts, matrix_rows, mapping_buf, must_free_mapping);
@@ -3275,6 +3560,9 @@ int op_matmul_id(struct htp_ops_context * octx) {
 }
 
 int op_matmul_qkv(struct htp_ops_context * octx) {
+    struct htp_thread_trace * tr = &octx->ctx->trace[0];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
     const struct htp_tensor * restrict src0 = octx->src[0]; // Wk
     const struct htp_tensor * restrict src1 = octx->src[1]; // x
     const struct htp_tensor * restrict src2 = octx->src[2]; // Wv
@@ -3379,9 +3667,6 @@ int op_matmul_qkv(struct htp_ops_context * octx) {
     mmctx->vtcm_src3_size_per_thread = L.src3_bytes / octx->n_threads;
     mmctx->vtcm_dst_size_per_thread  = L.dst_bytes / octx->n_threads;
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)
-        return HTP_STATUS_OK;
-
     mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
     mmctx->quant_task_func = quant_task_func;
     mmctx->n_quant_tasks = n_quant_tasks;
@@ -3413,12 +3698,18 @@ int op_matmul_qkv(struct htp_ops_context * octx) {
     } else {
         matmul_job_func = hvx_mm_qkv_2d;
     }
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
     worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, n_matmul_jobs);
 
     return HTP_STATUS_OK;
 }
 
 int op_matmul_ffn(struct htp_ops_context * octx) {
+    struct htp_thread_trace * tr = &octx->ctx->trace[0];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
     const struct htp_tensor * restrict src0 = octx->src[0]; // Wgate
     const struct htp_tensor * restrict src1 = octx->src[1]; // y
     const struct htp_tensor * restrict src2 = octx->src[2]; // Wup
@@ -3516,9 +3807,6 @@ int op_matmul_ffn(struct htp_ops_context * octx) {
     mmctx->vtcm_src2_size_per_thread = L.src2_bytes / octx->n_threads;
     mmctx->vtcm_dst_size_per_thread  = L.dst_bytes / octx->n_threads;
 
-    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)
-        return HTP_STATUS_OK;
-
     mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
     mmctx->quant_task_func = quant_task_func;
     mmctx->n_quant_tasks = n_quant_tasks;
@@ -3550,6 +3838,9 @@ int op_matmul_ffn(struct htp_ops_context * octx) {
     } else {
         matmul_job_func = hvx_mm_ffn_2d;
     }
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0);
+
     worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, n_matmul_jobs);
 
     return HTP_STATUS_OK;

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

@@ -95,6 +95,8 @@ struct htp_mm_kernel_params {
     struct fastdiv_values div_r2;
     struct fastdiv_values div_r3;
     struct fastdiv_values div_ne11;
+    struct fastdiv_values div_n_act_threads;
+    struct fastdiv_values div_ne00_padded;
 };
 
 #if defined(__cplusplus)
@@ -643,6 +645,136 @@ static inline size_t htp_mm_hmx_get_batched_vtcm_size(
     return L.total_bytes;
 }
 
+static inline bool htp_mm_hmx_solve_batched_params(
+    int wtype,
+    uint32_t k,
+    uint32_t ne01_padded,
+    uint32_t ne11,
+    uint32_t group_size,
+    bool use_dma_activation,
+    int n_threads,
+    bool pipeline,
+    size_t vtcm_budget,
+    size_t * m_chunk_out,
+    size_t * n_chunk_out,
+    int * act_threads_out,
+    size_t * vtcm_size_out
+) {
+    size_t best_mblocks = SIZE_MAX;
+    int best_act_threads = 0;
+    size_t best_m_chunk = 0;
+    size_t best_n_chunk = 0;
+    size_t best_vtcm_size = 0;
+
+    int act_threads = n_threads;
+    while (act_threads >= 1) {
+        size_t group_overhead = 256;
+        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);
+
+        size_t m_chunk_candidate = 0;
+        size_t n_chunk_candidate = 0;
+        size_t vtcm_size_candidate = 0;
+
+        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, use_dma_activation, pipeline, act_threads);
+            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)) {
+                    best_mblocks = mblocks;
+                    best_act_threads = act_threads;
+                    best_m_chunk = m_chunk_candidate;
+                    best_n_chunk = n_chunk_candidate;
+                    best_vtcm_size = exact_size;
+                }
+            }
+        }
+        if (act_threads == 1) {
+            act_threads = 0;
+        } else {
+            act_threads /= 2;
+        }
+    }
+
+    if (best_act_threads > 0) {
+        *m_chunk_out = best_m_chunk;
+        *n_chunk_out = best_n_chunk;
+        *vtcm_size_out = best_vtcm_size;
+        *act_threads_out = best_act_threads;
+        return true;
+    }
+    return false;
+}
+
+static inline bool htp_mm_hmx_solve_2d_params(
+    int wtype,
+    uint32_t k,
+    uint32_t m_id_rows,
+    uint32_t ne01_padded,
+    uint32_t ne11_padded,
+    uint32_t m_for_cost,
+    int n_threads,
+    bool pipeline,
+    bool is_matmul_id,
+    uint32_t aligned_tile_size,
+    size_t vtcm_budget,
+    size_t * m_chunk_out,
+    size_t * n_chunk_out,
+    int * act_threads_out,
+    size_t * vtcm_size_out
+) {
+    size_t best_mblocks = SIZE_MAX;
+    int best_act_threads = 0;
+    size_t best_m_chunk = 0;
+    size_t best_n_chunk = 0;
+    size_t best_vtcm_size = 0;
+
+    const int m_for_chunks = is_matmul_id ? hex_align_up(m_id_rows, 32) : ne11_padded;
+
+    int act_threads = n_threads;
+    while (act_threads >= 1) {
+        size_t simple_2d_overhead = 256;
+        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);
+
+        size_t m_chunk_candidate = 0;
+        size_t n_chunk_candidate = 0;
+        size_t vtcm_size_candidate = 0;
+
+        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);
+            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)) {
+                    best_mblocks = mblocks;
+                    best_act_threads = act_threads;
+                    best_m_chunk = m_chunk_candidate;
+                    best_n_chunk = n_chunk_candidate;
+                    best_vtcm_size = exact_size;
+                }
+            }
+        }
+        if (act_threads == 1) {
+            act_threads = 0;
+        } else {
+            act_threads /= 2;
+        }
+    }
+
+    if (best_act_threads > 0) {
+        *m_chunk_out = best_m_chunk;
+        *n_chunk_out = best_n_chunk;
+        *vtcm_size_out = best_vtcm_size;
+        *act_threads_out = best_act_threads;
+        return true;
+    }
+    return false;
+}
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -18,6 +18,7 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 
 // Redefined the rope type constants as we can't include ggml.h
 #define HTP_ROPE_TYPE_NORMAL 0
@@ -712,17 +713,11 @@ static int execute_op_rope_f32(struct htp_ops_context * octx) {
 }
 
 int op_rope(struct htp_ops_context * octx) {
-    int err = HTP_STATUS_OK;
-
     switch (octx->src[0]->type) {
         case HTP_TYPE_F32:
-            err = execute_op_rope_f32(octx);
-            break;
+            return execute_op_rope_f32(octx);
 
         default:
-            err = HTP_STATUS_NO_SUPPORT;
-            break;
+            return HTP_STATUS_NO_SUPPORT;
     }
-
-    return err;
 }

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

@@ -19,6 +19,7 @@
 #include "ggml-common.h"
 #include "htp-ctx.h"
 #include "htp-ops.h"
+#include "htp-tensor.h"
 #include "htp-vtcm.h"
 #include "hex-profile.h"
 
@@ -137,6 +138,24 @@ static void scale_f32(const float * restrict src,
     }
 }
 
+static void clamp_f32(const float * restrict src,
+                      float * restrict dst,
+                      const uint32_t num_rows,
+                      const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+    float min = 0.f;
+    float max = 0.f;
+    memcpy(&min, &op_params[0], sizeof(float));
+    memcpy(&max, &op_params[1], sizeof(float));
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_clamp_scalar_f32(dst_local, src_local, min, max, ne0);
+    }
+}
+
 static void rms_norm_f32(const float * restrict src,
                          float * restrict dst,
                          const uint32_t num_rows,
@@ -397,7 +416,7 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat
     struct htp_ops_context * octx = uctx->octx;                                                                     \
     const struct htp_tensor * src = octx->src[0];                                                                   \
     const struct htp_tensor * dst = octx->dst;                                                                      \
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;                                       \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                          \
                                                                                                                     \
     htp_unary_preamble;                                                                                             \
                                                                                                                     \
@@ -541,6 +560,7 @@ DEFINE_UNARY_TASK(norm,           false, false, norm_f32(src0_vtcm, dst_vtcm, bl
 DEFINE_UNARY_TASK(rms_norm,       false, false, rms_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(rms_norm_mul,   true,  false, rms_norm_mul_f32(src0_vtcm, uctx->broadcast_weight ? (const float *) src1_vtcm_data : src1_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(scale,          false, false, scale_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(clamp,          false, false, clamp_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(sqr,            false, false, sqr_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(sqrt,           false, false, sqrt_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(unary_neg,      false, false, neg_f32(src0_vtcm, dst_vtcm, block_size, uctx))
@@ -558,7 +578,7 @@ static void unary_task_f32_tiled_##NAME(unsigned int nth, unsigned int ith, void
     struct htp_ops_context * octx = uctx->octx;                                                                     \
     const struct htp_tensor * src = octx->src[0];                                                                   \
     const struct htp_tensor * dst = octx->dst;                                                                      \
-    struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[ith] : NULL;                                       \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                          \
                                                                                                                     \
     htp_unary_preamble;                                                                                             \
                                                                                                                     \
@@ -680,6 +700,14 @@ static inline void tile_scale_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm,
     hvx_scale_offset_f32_aa(dst_vtcm, src_vtcm, tw, scale, bias);
 }
 
+static inline void tile_clamp_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw, const int32_t * op_params) {
+    float min = 0.f;
+    float max = 0.f;
+    memcpy(&min, &op_params[0], sizeof(float));
+    memcpy(&max, &op_params[1], sizeof(float));
+    hvx_clamp_scalar_f32(dst_vtcm, src_vtcm, min, max, tw);
+}
+
 static inline void tile_unary_softplus_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw) {
     const float * restrict sf = (const float *) src_vtcm;
     float * restrict df       = (float *) dst_vtcm;
@@ -764,6 +792,7 @@ static inline void tri_apply_tile_f32(const uint8_t * restrict src, uint8_t * re
 }
 
 DEFINE_UNARY_TILED_TASK(scale,          false, tile_scale_f32(dst_vtcm, src_vtcm, tw, op_params))
+DEFINE_UNARY_TILED_TASK(clamp,          false, tile_clamp_f32(dst_vtcm, src_vtcm, tw, op_params))
 DEFINE_UNARY_TILED_TASK(sqr,            false, hvx_sqr_f32_aa(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(sqrt,           false, hvx_sqrt_f32_aa(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(unary_neg,      false, hvx_scale_f32_aa(dst_vtcm, src_vtcm, tw, -1.0f))
@@ -786,6 +815,7 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
         case HTP_OP_RMS_NORM:        op_type = "rmsnorm-f32";      break;
         case HTP_OP_RMS_NORM_MUL:    op_type = "rmsnorm-mul-f32";  break;
         case HTP_OP_SCALE:           op_type = "scale-f32";        break;
+        case HTP_OP_CLAMP:           op_type = "clamp-f32";        break;
         case HTP_OP_SQR:             op_type = "sqr-f32";          break;
         case HTP_OP_SQRT:            op_type = "sqrt-f32";         break;
         case HTP_OP_UNARY_NEG:       op_type = "neg-f32";          break;
@@ -881,6 +911,7 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
         if (col_tile) {
             switch (octx->op) {
                 case HTP_OP_SCALE:           task_func = unary_task_f32_tiled_scale;          break;
+                case HTP_OP_CLAMP:           task_func = unary_task_f32_tiled_clamp;          break;
                 case HTP_OP_SQR:             task_func = unary_task_f32_tiled_sqr;            break;
                 case HTP_OP_SQRT:            task_func = unary_task_f32_tiled_sqrt;           break;
                 case HTP_OP_UNARY_NEG:       task_func = unary_task_f32_tiled_unary_neg;      break;
@@ -897,6 +928,7 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
                 case HTP_OP_RMS_NORM:        task_func = unary_task_f32_rms_norm;             break;
                 case HTP_OP_RMS_NORM_MUL:    task_func = unary_task_f32_rms_norm_mul;         break;
                 case HTP_OP_SCALE:           task_func = unary_task_f32_scale;                break;
+                case HTP_OP_CLAMP:           task_func = unary_task_f32_clamp;                break;
                 case HTP_OP_SQR:             task_func = unary_task_f32_sqr;                  break;
                 case HTP_OP_SQRT:            task_func = unary_task_f32_sqrt;                 break;
                 case HTP_OP_UNARY_NEG:       task_func = unary_task_f32_unary_neg;            break;
@@ -922,17 +954,11 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
 }
 
 int op_unary(struct htp_ops_context * octx) {
-    int err = HTP_STATUS_OK;
-
     switch (octx->src[0]->type) {
         case HTP_TYPE_F32:
-            err = execute_op_unary_f32(octx);
-            break;
+            return execute_op_unary_f32(octx);
 
         default:
-            err = HTP_STATUS_NO_SUPPORT;
-            break;
+            return HTP_STATUS_NO_SUPPORT;
     }
-
-    return err;
 }

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

@@ -41,6 +41,7 @@ _Static_assert(sizeof(struct htp_unary_kernel_params) <= 128, "htp_unary_kernel_
 
 static inline bool htp_op_is_unary(uint32_t opcode) {
     switch (opcode) {
+        case HTP_OP_CLAMP:
         case HTP_OP_NORM:
         case HTP_OP_RMS_NORM:
         case HTP_OP_RMS_NORM_MUL:

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

@@ -0,0 +1,244 @@
+#include "work-queue.h"
+#include "hex-utils.h"
+
+#include <qurt.h>
+#include <qurt_hvx.h>
+
+#include <stdatomic.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include "HAP_farf.h"
+
+#define LOWEST_USABLE_QURT_PRIO (254)
+
+// internal structure kept in thread-local storage per instance of work queue
+typedef struct {
+    work_queue_t  queue;
+    unsigned int  id;
+} worker_context_t;
+
+struct work_queue_task_s {
+    work_queue_func_t func;
+    void *            data;
+    unsigned int      n_threads;
+    atomic_uint       barrier;
+};
+
+// internal structure kept in thread-local storage per instance of work queue
+struct work_queue_s {
+    atomic_uint        seqn;      // seqno used to detect new jobs
+    atomic_uint        idx_read;  // Updated by producer (pop/reclaim)
+    unsigned int       idx_write; // Updated by producer (push)
+    uint32_t           idx_mask;
+    uint32_t           capacity;
+
+    qurt_thread_t      thread[WORK_QUEUE_MAX_N_THREADS];   // thread ID's of the workers
+    worker_context_t   context[WORK_QUEUE_MAX_N_THREADS];  // worker contexts
+    void *             stack[WORK_QUEUE_MAX_N_THREADS];    // thread stack pointers
+    unsigned int       n_threads;                          // total threads (workers + main)
+    unsigned int       n_workers;                          // number of active threads (just workers)
+
+    atomic_bool        active;                             // workers are polling/active
+    atomic_bool        killed;                             // threads need to exit
+    bool               external_mem;                       // memory owned externally
+
+    struct work_queue_task_s queue[] __attribute__((aligned(HEX_L2_LINE_SIZE)));
+};
+
+static void work_queue_thread(void * context) {
+    worker_context_t * me = (worker_context_t *) context;
+    work_queue_t       q  = me->queue;
+
+    FARF(HIGH, "work-queue: thread %u started", me->id);
+
+    unsigned int prev_seqn = 0;
+
+    while (!atomic_load_explicit(&q->killed, memory_order_relaxed)) {
+        unsigned int seqn = atomic_load_explicit(&q->seqn, memory_order_acquire);
+        if (seqn == prev_seqn) {
+            if (atomic_load_explicit(&q->active, memory_order_relaxed)) {
+                hex_pause();
+            } else {
+                qurt_futex_wait(&q->seqn, prev_seqn);
+            }
+            continue;
+        }
+
+        prev_seqn = seqn;
+
+        // Process all active tasks in the queue
+        unsigned int ir = atomic_load_explicit(&q->idx_read, memory_order_relaxed);
+        unsigned int iw = q->idx_write;
+
+        while (ir != iw) {
+            struct work_queue_task_s * task = &q->queue[ir];
+
+            unsigned int n = task->n_threads;
+            unsigned int i = me->id;
+            if (i < n) {
+                task->func(n, i, task->data);
+
+                atomic_fetch_sub_explicit(&task->barrier, 1, memory_order_release);
+            } else {
+                while (atomic_load_explicit(&task->barrier, memory_order_relaxed) > 0) {
+                    hex_pause();
+                }
+            }
+
+            ir = (ir + 1) & q->idx_mask;
+        }
+    }
+
+    FARF(HIGH, "work-queue: thread %u stopped", me->id);
+}
+
+bool work_queue_run_async(work_queue_t q, work_queue_func_t func, void * data, unsigned int n) {
+    if (n > q->n_threads) {
+        FARF(ERROR, "work-queue: invalid number of jobs %u for n-threads %u", n, q->n_threads);
+        return false;
+    }
+
+    unsigned int ir = atomic_load_explicit(&q->idx_read, memory_order_relaxed);
+    unsigned int iw = q->idx_write;
+
+    if (((iw + 1) & q->idx_mask) == ir) {
+        FARF(ERROR, "work-queue-push: queue is full\n");
+        return false;
+    }
+
+    struct work_queue_task_s * task = &q->queue[iw];
+    task->func      = func;
+    task->data      = data;
+    task->n_threads = n;
+    atomic_store_explicit(&task->barrier, n, memory_order_relaxed);
+
+    q->idx_write = (iw + 1) & q->idx_mask;
+
+    // publish job to workers (already awake and polling)
+    atomic_fetch_add_explicit(&q->seqn, 1, memory_order_release);
+
+    // main thread runs job #0
+    func(n, 0, data);
+
+    atomic_fetch_sub_explicit(&task->barrier, 1, memory_order_release);
+
+    while (atomic_load_explicit(&task->barrier, memory_order_relaxed) > 0) {
+        hex_pause();
+    }
+
+    atomic_thread_fence(memory_order_acquire);
+
+    atomic_store_explicit(&q->idx_read, (ir + 1) & q->idx_mask, memory_order_relaxed);
+
+    return true;
+}
+
+size_t work_queue_sizeof(uint32_t n_threads, uint32_t capacity, uint32_t stack_size) {
+    capacity = hex_ceil_pow2(capacity);
+    uint32_t n_workers = n_threads > 1 ? n_threads - 1 : 0;
+    size_t size_stacks = stack_size * n_workers;
+    size_t size_q = hex_align_up(sizeof(struct work_queue_s) + capacity * sizeof(struct work_queue_task_s), HEX_L2_LINE_SIZE);
+    return size_stacks + size_q;
+}
+
+size_t work_queue_alignof(void) {
+    return 4096;
+}
+
+work_queue_t work_queue_init(void * ptr, uint32_t n_threads, uint32_t capacity, uint32_t stack_size) {
+    capacity = hex_ceil_pow2(capacity);
+    uint32_t n_workers = n_threads > 1 ? n_threads - 1 : 0;
+    unsigned char * mem_blob = (unsigned char *) ptr;
+
+    work_queue_t q = (work_queue_t) (mem_blob + stack_size * n_workers);
+    memset(q, 0, sizeof(struct work_queue_s) + capacity * sizeof(struct work_queue_task_s));
+
+    q->n_threads = n_threads;
+    q->n_workers = n_workers;
+    q->external_mem = true;
+    q->capacity = capacity;
+
+    for (unsigned int i = 0; i < n_workers; i++) {
+        q->stack[i]  = mem_blob; mem_blob += stack_size;
+        q->thread[i] = 0;
+        q->context[i].id    = i + 1;
+        q->context[i].queue = q;
+    }
+
+    atomic_init(&q->idx_read, 0);
+    atomic_init(&q->seqn,     0);
+    atomic_init(&q->active,   false);
+    q->idx_write = 0;
+    q->idx_mask  = capacity - 1;
+    q->killed    = 0;
+    for (int i = 0; i < (int) capacity; i++) {
+        atomic_init(&q->queue[i].barrier, 0);
+        q->queue[i].func      = NULL;
+        q->queue[i].data      = NULL;
+        q->queue[i].n_threads = 0;
+    }
+
+    // launch the workers
+    qurt_thread_attr_t attr;
+    qurt_thread_attr_init(&attr);
+
+    for (unsigned int i = 0; i < n_workers; i++) {
+        qurt_thread_attr_set_stack_addr(&attr, q->stack[i]);
+        qurt_thread_attr_set_stack_size(&attr, stack_size);
+
+        char thread_name[32];
+        snprintf(thread_name, sizeof(thread_name), "work-queue:%u", i);
+        qurt_thread_attr_set_name(&attr, thread_name);
+
+        // set up priority - by default, match the creating thread's prio
+        int prio = qurt_thread_get_priority(qurt_thread_get_id());
+        if (prio < 1) {
+            prio = 1;
+        }
+        if (prio > LOWEST_USABLE_QURT_PRIO) {
+            prio = LOWEST_USABLE_QURT_PRIO;
+        }
+
+        qurt_thread_attr_set_priority(&attr, prio);
+
+        int err = qurt_thread_create(&q->thread[i], &attr, work_queue_thread, (void *) &q->context[i]);
+        if (err) {
+            FARF(ERROR, "Could not launch worker threads!");
+            work_queue_free(q);
+            return NULL;
+        }
+    }
+
+    return q;
+}
+
+void work_queue_free(work_queue_t q) {
+    if (!q) { return; }
+
+    atomic_store_explicit(&q->killed,   1, memory_order_relaxed);
+    atomic_fetch_add_explicit(&q->seqn, 1, memory_order_release);
+    qurt_futex_wake(&q->seqn, q->n_workers);
+
+    for (unsigned int i = 0; i < q->n_workers; i++) {
+        if (q->thread[i]) {
+            int status;
+            (void) qurt_thread_join(q->thread[i], &status);
+        }
+    }
+}
+
+void work_queue_wakeup(work_queue_t q) {
+    if (!atomic_load_explicit(&q->active, memory_order_relaxed)) {
+        atomic_store_explicit(&q->active, true, memory_order_release);
+        // Increment seqn and wake workers to transition them out of sleep
+        atomic_fetch_add_explicit(&q->seqn, 1, memory_order_release);
+        qurt_futex_wake(&q->seqn, q->n_workers);
+    }
+}
+
+void work_queue_suspend(work_queue_t q) {
+    atomic_store_explicit(&q->active, false, memory_order_release);
+}

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

@@ -0,0 +1,38 @@
+#ifndef HTP_WORK_QUEUE_H
+#define HTP_WORK_QUEUE_H
+
+#include <stdbool.h>
+#include <stdint.h>
+#include <stddef.h>
+
+typedef void (*work_queue_func_t)(unsigned int n, unsigned int i, void *);
+
+struct work_queue_s;
+typedef struct work_queue_s * work_queue_t;
+
+#define WORK_QUEUE_MAX_N_THREADS      10
+
+size_t       work_queue_sizeof(uint32_t n_threads, uint32_t capacity, uint32_t stack_size);
+size_t       work_queue_alignof(void);
+work_queue_t work_queue_init(void * ptr, uint32_t n_threads, uint32_t capacity, uint32_t stack_size);
+void         work_queue_free(work_queue_t q);
+
+void work_queue_wakeup(work_queue_t q);
+void work_queue_suspend(work_queue_t q);
+
+bool work_queue_run_async(work_queue_t q, work_queue_func_t func, void * data, unsigned int n);
+
+static inline bool work_queue_run(work_queue_t q, work_queue_func_t func, void * data, unsigned int n) {
+    if (n <= 1) {
+        func(n, 0, data);
+        return true;
+    }
+    return work_queue_run_async(q, func, data, n);
+}
+
+// Legacy compatibility
+typedef work_queue_func_t     worker_callback_t;
+#define worker_pool_run_func  work_queue_run
+#define worker_pool           work_queue
+
+#endif  // #ifndef HTP_WORK_QUEUE_H

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

@@ -1,305 +0,0 @@
-#include "worker-pool.h"
-#include "hex-utils.h"
-
-#include <qurt.h>
-#include <qurt_hvx.h>
-
-#include <stdatomic.h>
-#include <stdint.h>
-#include <stdio.h>
-#include <stdlib.h>
-#include <string.h>
-
-#include "HAP_farf.h"
-
-#define LOWEST_USABLE_QURT_PRIO (254)
-
-struct worker_pool_s;
-
-// internal structure kept in thread-local storage per instance of worker pool
-typedef struct {
-    struct worker_pool_s * pool;
-    unsigned int           id;
-} worker_context_t;
-
-// internal structure kept in thread-local storage per instance of worker pool
-typedef struct worker_pool_s {
-    worker_pool_job_t job[MAX_NUM_WORKERS];      // list of job descriptors
-    qurt_thread_t     thread[MAX_NUM_WORKERS];   // thread ID's of the workers
-    worker_context_t  context[MAX_NUM_WORKERS];  // worker contexts
-    void *            stack[MAX_NUM_WORKERS];    // thread stack pointers
-    unsigned int      n_threads;                 // number of workers in this pool
-
-    atomic_uint seqn;                            // seqno used to detect new jobs
-    atomic_uint next_job;                        // next job index
-    atomic_uint n_pending;                       // number of pending jobs
-    atomic_uint n_jobs;                          // number of current jobs
-    atomic_bool killed;                          // threads need to exit
-} worker_pool_t;
-
-static void worker_pool_main(void * context) {
-    worker_context_t * me   = (worker_context_t *) context;
-    worker_pool_t *    pool = me->pool;
-
-    FARF(HIGH, "worker-pool: thread %u started", me->id);
-
-    unsigned int prev_seqn = 0;
-    unsigned int poll_cnt  = WORKER_POOL_POLL_COUNT;
-    while (!atomic_load(&pool->killed)) {
-        unsigned int seqn = atomic_load(&pool->seqn);
-        if (seqn == prev_seqn) {
-            // drop HVX context while spinning
-            if (poll_cnt > 1 && poll_cnt == WORKER_POOL_POLL_COUNT) {
-                qurt_hvx_unlock();
-            }
-            if (--poll_cnt) {
-                hex_pause();
-                continue;
-            }
-            qurt_futex_wait(&pool->seqn, prev_seqn);
-            poll_cnt = WORKER_POOL_POLL_COUNT;
-            continue;
-        }
-
-        prev_seqn = seqn;
-        poll_cnt  = WORKER_POOL_POLL_COUNT;
-
-        // New job
-        unsigned int n = atomic_load(&pool->n_jobs);
-        unsigned int i = atomic_fetch_add(&pool->next_job, 1);
-        if (i >= n) {
-            // Spurious wakeup
-            continue;
-        }
-
-        pool->job[i].func(n, i, pool->job[i].data);
-
-        atomic_fetch_sub(&pool->n_pending, 1);
-    }
-
-    FARF(HIGH, "worker-pool: thread %u stopped", me->id);
-}
-
-AEEResult worker_pool_init_with_stack_size(worker_pool_context_t * context, uint32_t n_threads, uint32_t stack_size) {
-    int err = 0;
-
-    if (NULL == context) {
-        FARF(ERROR, "NULL context passed to worker_pool_init().");
-        return AEE_EBADPARM;
-    }
-
-    // Allocations
-    int size = (stack_size * n_threads) + (sizeof(worker_pool_t));
-
-    unsigned char * mem_blob = (unsigned char *) malloc(size);
-    if (!mem_blob) {
-        FARF(ERROR, "Could not allocate memory for worker pool!!");
-        return AEE_ENOMEMORY;
-    }
-
-    worker_pool_t * me = (worker_pool_t *) (mem_blob + stack_size * n_threads);
-
-    // name for the first worker, useful in debugging threads
-    char name[19];
-    snprintf(name, 12, "0x%8x:", (int) me);
-    strcat(name, "worker0");
-    me->n_threads = n_threads;
-
-    // initializations
-    for (unsigned int i = 0; i < me->n_threads; i++) {
-        me->stack[i]  = NULL;
-        me->thread[i] = 0;
-
-        me->context[i].id   = i;
-        me->context[i].pool = me;
-    }
-
-    // initialize job queue
-    me->n_pending = 0;
-    me->n_jobs    = 0;
-    me->next_job  = 0;
-    me->seqn      = 0;
-    me->killed    = 0;
-
-    // launch the workers
-    qurt_thread_attr_t attr;
-    qurt_thread_attr_init(&attr);
-
-    for (unsigned int i = 0; i < me->n_threads; i++) {
-        // set up stack
-        me->stack[i] = mem_blob;
-        mem_blob += stack_size;
-        qurt_thread_attr_set_stack_addr(&attr, me->stack[i]);
-        qurt_thread_attr_set_stack_size(&attr, stack_size);
-
-        // set up name
-        qurt_thread_attr_set_name(&attr, name);
-        name[17] = (name[17] + 1);
-        // name threads context:worker0, context:worker1, .. (recycle at 9, but num threads should be less than that anyway)
-        if (name[17] > '9') {
-            name[17] = '0';
-        }
-
-        // set up priority - by default, match the creating thread's prio
-        int prio = qurt_thread_get_priority(qurt_thread_get_id());
-
-        if (prio < 1) {
-            prio = 1;
-        }
-        if (prio > LOWEST_USABLE_QURT_PRIO) {
-            prio = LOWEST_USABLE_QURT_PRIO;
-        }
-
-        qurt_thread_attr_set_priority(&attr, prio);
-
-        // launch
-        err = qurt_thread_create(&me->thread[i], &attr, worker_pool_main, (void *) &me->context[i]);
-        if (err) {
-            FARF(ERROR, "Could not launch worker threads!");
-            worker_pool_release((worker_pool_context_t *) &me);
-            return AEE_EQURTTHREADCREATE;
-        }
-    }
-    *context = (worker_pool_context_t *) me;
-    return AEE_SUCCESS;
-}
-
-AEEResult worker_pool_init(worker_pool_context_t * context, uint32_t n_threads) {
-    return worker_pool_init_with_stack_size(context, n_threads, WORKER_THREAD_STACK_SZ);
-}
-
-// clean up worker pool
-void worker_pool_release(worker_pool_context_t * context) {
-    worker_pool_t * me = (worker_pool_t *) *context;
-
-    // if no worker pool exists, return error.
-    if (NULL == me) {
-        return;
-    }
-
-    atomic_store(&me->killed, 1);
-    atomic_fetch_add(&me->seqn, 1);
-    qurt_futex_wake(&me->seqn, me->n_threads);
-
-    // de-initializations
-    for (unsigned int i = 0; i < me->n_threads; i++) {
-        if (me->thread[i]) {
-            int status;
-            (void) qurt_thread_join(me->thread[i], &status);
-        }
-    }
-
-    // free allocated memory (were allocated as a single buffer starting at stack[0])
-    if (me->stack[0]) {
-        free(me->stack[0]);
-    }
-
-    *context = NULL;
-}
-
-// run jobs
-AEEResult worker_pool_run_jobs(worker_pool_context_t context, worker_pool_job_t * job, unsigned int n) {
-    worker_pool_t * me = (worker_pool_t *) context;
-    if (NULL == me) {
-        FARF(ERROR, "worker-pool: invalid context");
-        return AEE_EBADPARM;
-    }
-
-    if (n > me->n_threads) {
-        FARF(ERROR, "worker-pool: invalid number of jobs %u for n-threads %u", n, me->n_threads);
-        return AEE_EBADPARM;
-    }
-
-    memcpy(me->job, job, sizeof(worker_pool_job_t) * n);
-
-    if (n > 1) {
-        atomic_store(&me->next_job, 1);
-        atomic_store(&me->n_jobs, n);
-        atomic_store(&me->n_pending, n - 1);
-
-        // wake up workers
-        atomic_fetch_add(&me->seqn, 1);
-        qurt_futex_wake(&me->seqn, n - 1);
-    }
-
-    // main thread runs job #0
-    me->job[0].func(n, 0, me->job[0].data);
-
-    if (n > 1) {
-        while (atomic_load(&me->n_pending))
-            ;
-    }
-
-    return 0;
-}
-
-// run func
-AEEResult worker_pool_run_func(worker_pool_context_t context, worker_callback_t func, void * data, unsigned int n) {
-    worker_pool_job_t job[n];
-
-    for (unsigned int i = 0; i < n; i++) {
-        job[i].func = func;
-        job[i].data = data;
-    }
-
-    return worker_pool_run_jobs(context, job, n);
-}
-
-AEEResult worker_pool_set_thread_priority(worker_pool_context_t context, unsigned int prio) {
-    worker_pool_t * me = (worker_pool_t *) context;
-
-    // if no worker pool exists, return error.
-    if (!me) {
-        return AEE_ENOMORE;
-    }
-
-    int result = AEE_SUCCESS;
-    if (prio < 1) {
-        prio = 1;
-    }
-    if (prio > LOWEST_USABLE_QURT_PRIO) {
-        prio = LOWEST_USABLE_QURT_PRIO;
-    }
-
-    for (unsigned int i = 0; i < me->n_threads; i++) {
-        int res = qurt_thread_set_priority(me->thread[i], (unsigned short) prio);
-        if (0 != res) {
-            result = AEE_EBADPARM;
-            FARF(ERROR, "QURT failed to set priority of thread %d, ERROR = %d", me->thread[i], res);
-        }
-    }
-
-    return result;
-}
-
-AEEResult worker_pool_retrieve_thread_id(worker_pool_context_t context, unsigned int * tids) {
-    worker_pool_t * me = (worker_pool_t *) context;
-    if (!me) {
-        FARF(ERROR, "worker-pool: invalid context");
-        return AEE_EBADPARM;
-        ;
-    }
-
-    for (int i = 0; i < me->n_threads; i++) {
-        tids[i] = me->thread[i];
-    }
-
-    return AEE_SUCCESS;
-}
-
-AEEResult worker_pool_get_thread_priority(worker_pool_context_t context, unsigned int * prio) {
-    worker_pool_t * me = (worker_pool_t *) context;
-    if (!me) {
-        FARF(ERROR, "worker-pool: invalid context");
-        return AEE_EBADPARM;
-    }
-
-    int priority = qurt_thread_get_priority(me->thread[0]);
-    if (priority > 0) {
-        *prio = priority;
-        return 0;
-    } else {
-        *prio = 0;
-        return AEE_EBADSTATE;
-    }
-}

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

@@ -1,65 +0,0 @@
-#ifndef HTP_WORKER_POOL_H
-#define HTP_WORKER_POOL_H
-
-// MACRO enables function to be visible in shared-library case.
-#define WORKERPOOL_API __attribute__((visibility("default")))
-
-#include <AEEStdDef.h>
-#include <AEEStdErr.h>
-#include <stdint.h>
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-/// signature of callbacks to be invoked by worker threads
-typedef void (*worker_callback_t)(unsigned int n, unsigned int i, void *);
-
-/// Typedef of worker_pool context
-typedef void * worker_pool_context_t;
-
-/// descriptor for requested callback
-typedef struct {
-    worker_callback_t func;
-    void *            data;
-} worker_pool_job_t;
-
-#define WORKER_THREAD_STACK_SZ  (2 * 16384)
-
-/// Maximum supported number of worker threads.
-#define MAX_NUM_WORKERS 10
-
-#if __HVX_ARCH__ > 79
-#define WORKER_POOL_POLL_COUNT  2000
-#else
-#define WORKER_POOL_POLL_COUNT  1
-#endif
-
-// Initialize worker pool.
-WORKERPOOL_API AEEResult worker_pool_init(worker_pool_context_t * context, uint32_t n_threads);
-
-// Initialize worker pool with custom stack size
-WORKERPOOL_API AEEResult worker_pool_init_with_stack_size(worker_pool_context_t * context,
-                                                          uint32_t                n_threads,
-                                                          uint32_t                stack_size);
-
-// Kill worker threads and release worker pool resources
-WORKERPOOL_API void worker_pool_release(worker_pool_context_t * context);
-
-// Run jobs with the worker pool.
-WORKERPOOL_API AEEResult worker_pool_run_jobs(worker_pool_context_t context, worker_pool_job_t * job, unsigned int n);
-
-WORKERPOOL_API AEEResult worker_pool_run_func(worker_pool_context_t context,
-                                              worker_callback_t     func,
-                                              void *                data,
-                                              unsigned int          n);
-
-WORKERPOOL_API AEEResult worker_pool_set_thread_priority(worker_pool_context_t context, unsigned int prio);
-WORKERPOOL_API AEEResult worker_pool_get_thread_priority(worker_pool_context_t context, unsigned int * prio);
-WORKERPOOL_API AEEResult worker_pool_retrieve_thread_id(worker_pool_context_t context, unsigned int * tids);
-
-#ifdef __cplusplus
-}
-#endif
-
-#endif  // #ifndef HTP_WORKER_POOL_H

ggml/src/ggml-opencl/CMakeLists.txt

@@ -211,6 +211,7 @@ set(GGML_OPENCL_KERNELS
     tanh
     exp
     expm1
+    abs
     softplus
     pad
     repeat

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

@@ -264,7 +264,8 @@ static ADRENO_GPU_GEN get_adreno_gpu_gen(const char *device_name) {
         return ADRENO_GPU_GEN::A7X;
     }
 
-    if (strstr(device_name, "830") ||
+    if (strstr(device_name, "810") ||
+        strstr(device_name, "830") ||
         strstr(device_name, "840") ||
         strstr(device_name, "850")) {
         return ADRENO_GPU_GEN::A8X;
@@ -844,6 +845,8 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_exp_f16, kernel_exp_f16_4, kernel_exp_f16_nc;
     cl_kernel kernel_expm1_f32, kernel_expm1_f32_4, kernel_expm1_f32_nc;
     cl_kernel kernel_expm1_f16, kernel_expm1_f16_4, kernel_expm1_f16_nc;
+    cl_kernel kernel_abs_f32, kernel_abs_f32_4, kernel_abs_f32_nc;
+    cl_kernel kernel_abs_f16, kernel_abs_f16_4, kernel_abs_f16_nc;
     cl_kernel kernel_softplus_f32, kernel_softplus_f32_4, kernel_softplus_f32_nc;
     cl_kernel kernel_softplus_f16, kernel_softplus_f16_4, kernel_softplus_f16_nc;
     cl_kernel kernel_upscale;
@@ -873,7 +876,7 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_gemm_moe_q8_1_dp4a_q5k = nullptr;   // generic dp4a MoE GEMM (MOE_QT=5, q5_K), opt-in
     cl_kernel kernel_moe_expand_scale_q5_K = nullptr;    // q5_K 6-bit s[] -> uniform scale[16]/min[8]
     cl_kernel kernel_gemv_moe_q5_k_f32_ns, kernel_gemm_moe_q5_k_f32_ns;
-    cl_kernel kernel_gemv_moe_q6_k_f32_ns, kernel_gemm_moe_q6_k_f32_ns;
+    cl_kernel kernel_gemv_moe_q6_k_f32_ns, kernel_gemm_moe_q6_k_f32_ns, kernel_gemm_moe_q6_k_f32_ns_bin;
     cl_kernel kernel_gemm_moe_q6_k_q8_1_dp4a = nullptr;    // dp4a (int8) q6_K MoE prefill GEMM
     cl_kernel kernel_gemv_moe_mxfp4_f32, kernel_gemm_moe_mxfp4_f32;
     cl_kernel kernel_gemv_moe_mxfp4_f32_ns, kernel_gemm_moe_mxfp4_f32_ns, kernel_gemm_moe_mxfp4_f32_ns_bin;
@@ -2928,6 +2931,27 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
         GGML_LOG_CONT(".");
     }
 
+    // abs
+    {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+        const std::string kernel_src {
+            #include "abs.cl.h"
+        };
+#else
+        const std::string kernel_src = read_file("abs.cl");
+#endif
+        cl_program prog =
+            build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
+        CL_CHECK((backend_ctx->kernel_abs_f32    = clCreateKernel(prog, "kernel_abs_f32", &err), err));
+        CL_CHECK((backend_ctx->kernel_abs_f32_4  = clCreateKernel(prog, "kernel_abs_f32_4", &err), err));
+        CL_CHECK((backend_ctx->kernel_abs_f32_nc = clCreateKernel(prog, "kernel_abs_f32_nc", &err), err));
+        CL_CHECK((backend_ctx->kernel_abs_f16    = clCreateKernel(prog, "kernel_abs_f16", &err), err));
+        CL_CHECK((backend_ctx->kernel_abs_f16_4  = clCreateKernel(prog, "kernel_abs_f16_4", &err), err));
+        CL_CHECK((backend_ctx->kernel_abs_f16_nc = clCreateKernel(prog, "kernel_abs_f16_nc", &err), err));
+        CL_CHECK(clReleaseProgram(prog));
+        GGML_LOG_CONT(".");
+    }
+
     // softplus
     {
 #ifdef GGML_OPENCL_EMBED_KERNELS
@@ -4286,6 +4310,24 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
         GGML_LOG_CONT(".");
     }
 
+    // gemm_moe_q6_k_f32_ns_bin
+    {
+        size_t bin_size = 0;
+        backend_ctx->kernel_gemm_moe_q6_k_f32_ns_bin = nullptr;
+
+        if (use_adreno_bin_kernels(backend_ctx)) {
+            const char * kernel_bin = (const char *)backend_ctx->get_adreno_bin_kernel("gemm_moe_q6_k_f32_ns_ila", &bin_size);
+            if (kernel_bin && bin_size > 0) {
+                cl_program prog =
+                    build_program_from_binary(backend_ctx->context, backend_ctx->device, kernel_bin, CL_moe_compile_opts, bin_size);
+
+                CL_CHECK((backend_ctx->kernel_gemm_moe_q6_k_f32_ns_bin = clCreateKernel(prog, "kernel_gemm_moe_q6_k_f32_ns_ila", &err), err));
+                CL_CHECK(clReleaseProgram(prog));
+                GGML_LOG_CONT(".");
+            }
+        }
+    }
+
     // gemm_moe_q6_k_q8_1_dp4a (dp4a q6_K MoE prefill GEMM)
     if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
@@ -5988,7 +6030,8 @@ static void transpose_2d(
     cl_kernel kernel,
     cl_mem src, cl_mem dst, size_t size,
     cl_int stride, cl_int rows,
-    bool blocking = true
+    bool blocking = true,
+    bool auto_local = false // let driver pick local size for non-uniform workgroups
 ) {
     static ggml_cl_buffer buf;
 
@@ -6014,7 +6057,7 @@ static void transpose_2d(
     size_t local_size[3] = {64, 1, 1};
     size_t global_size[3] = {(size_t)stride, (size_t)rows, 1};;
     CL_CHECK(clEnqueueNDRangeKernel(backend_ctx->queue, kernel, 3, NULL,
-        global_size, local_size, 0, NULL, NULL));
+        global_size, auto_local ? NULL : local_size, 0, NULL, NULL));
 
     if (blocking) {
         CL_CHECK(clEnqueueCopyBuffer(backend_ctx->queue, trans, dst, 0, 0, size, 0, NULL, &evt));
@@ -6031,10 +6074,11 @@ static void transpose_2d_as_8b(
     ggml_backend_opencl_context * backend_ctx,
     cl_mem src, cl_mem dst, size_t size,
     cl_int stride, cl_int rows,
-    bool blocking = true
+    bool blocking = true,
+    bool auto_local = false
 ) {
     transpose_2d(backend_ctx, backend_ctx->kernel_transpose_8_buf,
-        src, dst, size, stride, rows, blocking);
+        src, dst, size, stride, rows, blocking, auto_local);
 }
 
 static void transpose_2d_as_16b(
@@ -7152,6 +7196,8 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
                     return op->src[0]->type == GGML_TYPE_F32;
                 case GGML_UNARY_OP_EXPM1:
                     return op->src[0]->type == GGML_TYPE_F32;
+                case GGML_UNARY_OP_ABS:
+                    return op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16;
                 case GGML_UNARY_OP_SOFTPLUS:
                     return op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16;
                 default:
@@ -9028,6 +9074,9 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer,
             transpose_2d_as_16b(backend_ctx, extra->q, extra->q, size_q, K/4, M);
             transpose_2d_as_16b(backend_ctx, extra->d, extra->d, size_d, K/256, M);
             transpose_2d_as_16b(backend_ctx, extra->dm, extra->dm, size_dm, K/256, M);
+
+            // Transpose s as uchar
+            transpose_2d_as_8b(backend_ctx, extra->s, extra->s, size_s, K/256*12, M, true, true);
         }
 #endif // GGML_OPENCL_USE_ADRENO_KERNELS
         return;
@@ -10196,23 +10245,27 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
             size_t size_q  = ggml_nelements(tensor)/ggml_blck_size(tensor->type)*ggml_blck_size(tensor->type)/2;
             size_t size_d  = ggml_nelements(tensor)/ggml_blck_size(tensor->type)*sizeof(ggml_fp16_t);
             size_t size_dm = ggml_nelements(tensor)/ggml_blck_size(tensor->type)*sizeof(ggml_fp16_t);
+            size_t size_s  = ggml_nelements(tensor)/ggml_blck_size(tensor->type)*12;
 
             static ggml_cl_buffer buf_trans_q;
             static ggml_cl_buffer buf_trans_d;
             static ggml_cl_buffer buf_trans_dm;
+            static ggml_cl_buffer buf_trans_s;
 
             buf_trans_q.allocate(backend_ctx->context, size_q);
             buf_trans_d.allocate(backend_ctx->context, size_d);
             buf_trans_dm.allocate(backend_ctx->context, size_dm);
+            buf_trans_s.allocate(backend_ctx->context, size_s);
 
-            // Transpose q, d, dm back
+            // Transpose q, d, dm, s back
             transpose_2d_as_16b(backend_ctx, extra->q,  buf_trans_q.buffer,  size_q,  M, K/4);
             transpose_2d_as_16b(backend_ctx, extra->d,  buf_trans_d.buffer,  size_d,  M, K/256);
             transpose_2d_as_16b(backend_ctx, extra->dm, buf_trans_dm.buffer, size_dm, M, K/256);
+            transpose_2d_as_8b (backend_ctx, extra->s,  buf_trans_s.buffer,  size_s,  M, K/256*12, true, true);
 
             cl_kernel kernel = backend_ctx->kernel_restore_block_q4_K_noshuffle;
             CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &buf_trans_q.buffer));
-            CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &extra->s));
+            CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &buf_trans_s.buffer));
             CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &buf_trans_d.buffer));
             CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_mem), &buf_trans_dm.buffer));
             CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &data_device));
@@ -13386,6 +13439,102 @@ static void ggml_cl_expm1(ggml_backend_t backend, const ggml_tensor * src0, cons
     }
 }
 
+static void ggml_cl_abs(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    GGML_ASSERT(src0);
+    GGML_ASSERT(src0->extra);
+    GGML_ASSERT(dst);
+    GGML_ASSERT(dst->extra);
+
+    UNUSED(src1);
+
+    ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context;
+
+    ggml_tensor_extra_cl * extra0 = (ggml_tensor_extra_cl *)src0->extra;
+    ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra;
+
+    cl_ulong offset0 = extra0->offset + src0->view_offs;
+    cl_ulong offsetd = extrad->offset + dst->view_offs;
+
+    const int ne00 = src0->ne[0];
+    const int ne01 = src0->ne[1];
+    const int ne02 = src0->ne[2];
+    const int ne03 = src0->ne[3];
+
+    const cl_ulong nb00 = src0->nb[0];
+    const cl_ulong nb01 = src0->nb[1];
+    const cl_ulong nb02 = src0->nb[2];
+    const cl_ulong nb03 = src0->nb[3];
+
+    const cl_ulong nb0 = dst->nb[0];
+    const cl_ulong nb1 = dst->nb[1];
+    const cl_ulong nb2 = dst->nb[2];
+    const cl_ulong nb3 = dst->nb[3];
+
+    cl_kernel kernel;
+
+    if (ggml_is_contiguous(src0)) {
+        // Handle contiguous input
+        int n = ggml_nelements(dst);
+        if (n % 4 == 0) {
+            if (src0->type == GGML_TYPE_F32) {
+                kernel = backend_ctx->kernel_abs_f32_4;
+            } else {
+                kernel = backend_ctx->kernel_abs_f16_4;
+            }
+            n /= 4;
+        } else {
+            if (src0->type == GGML_TYPE_F32) {
+                kernel = backend_ctx->kernel_abs_f32;
+            } else {
+                kernel = backend_ctx->kernel_abs_f16;
+            }
+        }
+
+        CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem),   &extra0->data_device));
+        CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset0));
+        CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem),   &extrad->data_device));
+        CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offsetd));
+
+        size_t global_work_size[] = {(size_t)n, 1, 1};
+        size_t local_work_size[] = {64, 1, 1};
+
+        size_t * local_work_size_ptr = local_work_size;
+        if (n % 64 != 0 && !backend_ctx->non_uniform_workgroups) {
+            local_work_size_ptr = nullptr;
+        }
+
+        backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size_ptr, dst);
+    } else {
+        // Handle non-contiguous input
+        if (src0->type == GGML_TYPE_F32) {
+            kernel = backend_ctx->kernel_abs_f32_nc;
+        } else {
+            kernel = backend_ctx->kernel_abs_f16_nc;
+        }
+
+        CL_CHECK(clSetKernelArg(kernel,  0, sizeof(cl_mem),   &extra0->data_device));
+        CL_CHECK(clSetKernelArg(kernel,  1, sizeof(cl_ulong), &offset0));
+        CL_CHECK(clSetKernelArg(kernel,  2, sizeof(cl_mem),   &extrad->data_device));
+        CL_CHECK(clSetKernelArg(kernel,  3, sizeof(cl_ulong), &offsetd));
+        CL_CHECK(clSetKernelArg(kernel,  4, sizeof(int),      &ne00));
+        CL_CHECK(clSetKernelArg(kernel,  5, sizeof(cl_ulong), &nb00));
+        CL_CHECK(clSetKernelArg(kernel,  6, sizeof(cl_ulong), &nb01));
+        CL_CHECK(clSetKernelArg(kernel,  7, sizeof(cl_ulong), &nb02));
+        CL_CHECK(clSetKernelArg(kernel,  8, sizeof(cl_ulong), &nb03));
+        CL_CHECK(clSetKernelArg(kernel,  9, sizeof(cl_ulong), &nb0));
+        CL_CHECK(clSetKernelArg(kernel, 10, sizeof(cl_ulong), &nb1));
+        CL_CHECK(clSetKernelArg(kernel, 11, sizeof(cl_ulong), &nb2));
+        CL_CHECK(clSetKernelArg(kernel, 12, sizeof(cl_ulong), &nb3));
+
+        int nth = 64;
+
+        size_t global_work_size[] = {(size_t)ne01*nth, (size_t)ne02, (size_t)ne03};
+        size_t local_work_size[] = {(size_t)nth, 1, 1};
+
+        backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+    }
+}
+
 static void ggml_cl_softplus(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     GGML_ASSERT(src0);
     GGML_ASSERT(src0->extra);
@@ -16936,9 +17085,6 @@ static void ggml_cl_mul_mat_q8_0_f32_adreno(ggml_backend_t backend, const ggml_t
     cl_ulong offset1 = extra1->offset + src1->view_offs;
     cl_ulong offsetd = extrad->offset + dst->view_offs;
 
-    GGML_ASSERT(src1->view_offs == 0);
-    GGML_ASSERT(dst->view_offs == 0);
-
     const int  ne00 = src0->ne[0];
     const int  ne01 = src0->ne[1];
     const int  ne02 = src0->ne[2];
@@ -16999,9 +17145,9 @@ static void ggml_cl_mul_mat_q8_0_f32_adreno(ggml_backend_t backend, const ggml_t
         CL_CHECK(clSetKernelArg(kernel,  0, sizeof(cl_mem),   &q_img));
         CL_CHECK(clSetKernelArg(kernel,  1, sizeof(cl_mem),   &extra0_q8_0->d));
         CL_CHECK(clSetKernelArg(kernel,  2, sizeof(cl_mem),   &b_img));
-        CL_CHECK(clSetKernelArg(kernel,  3, sizeof(cl_ulong), &extra1->offset));
+        CL_CHECK(clSetKernelArg(kernel,  3, sizeof(cl_ulong), &offset1));
         CL_CHECK(clSetKernelArg(kernel,  4, sizeof(cl_mem),   &extrad->data_device));
-        CL_CHECK(clSetKernelArg(kernel,  5, sizeof(cl_ulong), &extrad->offset));
+        CL_CHECK(clSetKernelArg(kernel,  5, sizeof(cl_ulong), &offsetd));
         CL_CHECK(clSetKernelArg(kernel,  6, sizeof(int),      &ne00));
         CL_CHECK(clSetKernelArg(kernel,  7, sizeof(int),      &ne01));
         CL_CHECK(clSetKernelArg(kernel,  8, sizeof(int),      &ne02));
@@ -18396,6 +18542,26 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
 
     GGML_ASSERT(ne00 == ne10);
 
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+    // adreno GEMM/GEMV kernels do not support broadcast, assuming ne2 and ne3 are 1 for src1
+    // so we handle broadcast here
+    if ((ne12 > 1 || ne13 > 1) && ne02 == 1 && ne03 == 1 &&
+        src0t != GGML_TYPE_F16 && src0t != GGML_TYPE_F32) {
+        for (int i13 = 0; i13 < ne13; ++i13) {
+            for (int i12 = 0; i12 < ne12; ++i12) {
+                ggml_tensor s1 = *src1;
+                s1.ne[2] = 1; s1.ne[3] = 1;
+                s1.view_offs = src1->view_offs + (size_t)i12*nb12 + (size_t)i13*nb13;
+                ggml_tensor d = *dst;
+                d.ne[2] = 1; d.ne[3] = 1;
+                d.view_offs = dst->view_offs + (size_t)i12*nb2 + (size_t)i13*nb3;
+                ggml_cl_mul_mat(backend, src0, &s1, &d);
+            }
+        }
+        return;
+    }
+#endif
+
     int nth0 = 32;
     int nth1 = 1;
     int nrows = 1;
@@ -22160,8 +22326,8 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
                         bool use_moe_dp4a = q5kmdp4a_on
                             && backend_ctx->kernel_gemm_moe_q8_1_dp4a_q5k != nullptr
                             && extra0_q5_K->scale != nullptr;
-                        // bin kernel takes precedence
-                        use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q4_k_f32_ns_bin == nullptr;
+                        // dot prod has to be available
+                        use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
 
                         if (use_moe_dp4a) {
                             const size_t tok_slots = (size_t)max_post_router_tile * n_tile_size;
@@ -22379,6 +22545,9 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
 
                 } else { // for gemm
                     kernel = backend_ctx->kernel_gemm_moe_q6_k_f32_ns;
+                    if (backend_ctx->kernel_gemm_moe_q6_k_f32_ns_bin) {
+                        kernel = backend_ctx->kernel_gemm_moe_q6_k_f32_ns_bin;
+                    }
 
                     // Reorder router if called from test-backend-ops or when new router is generated.
                     // Otherwise reuse the reordered result from previous mul_mat_id call.
@@ -22399,6 +22568,8 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
                                                             || backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E);
                     // dot prod has to be available
                     use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
+                    // bin kernel takes precedence
+                    use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q6_k_f32_ns_bin == nullptr;
 
                     cl_buffer_region region;
                     region.origin = 0;
@@ -22435,6 +22606,11 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
                         CL_CHECK(status);
                         cl_image_format image_format_buf_src1 = {CL_RGBA, CL_FLOAT};
                         cl_image_desc image_desc_buf_src1 = {CL_MEM_OBJECT_IMAGE1D_BUFFER, static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size / 4), 0,0,0,0,0,0,0, {buf_src1_reordered}};
+                        if (backend_ctx->kernel_gemm_moe_q6_k_f32_ns_bin) {
+                            // bin kernel uses slightly different image format
+                            image_format_buf_src1 = {CL_R, CL_FLOAT};
+                            image_desc_buf_src1.image_width = static_cast<size_t>(ne00 * max_post_router_tile * n_tile_size);
+                        }
                         image_src1_reordered = clCreateImage(backend_ctx->context, CL_MEM_READ_ONLY, &image_format_buf_src1, &image_desc_buf_src1, NULL, &status);
                         CL_CHECK(status);
 
@@ -24407,6 +24583,12 @@ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor
                     }
                     func = ggml_cl_expm1;
                     break;
+                case GGML_UNARY_OP_ABS:
+                    if (!any_on_device) {
+                        return false;
+                    }
+                    func = ggml_cl_abs;
+                    break;
                 case GGML_UNARY_OP_SOFTPLUS:
                     if (!any_on_device) {
                         return false;

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

@@ -0,0 +1,113 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
+//------------------------------------------------------------------------------
+// abs
+//------------------------------------------------------------------------------
+
+kernel void kernel_abs_f32(
+        global const float * src0,
+        ulong                offset0,
+        global       float * dst,
+        ulong                offsetd
+) {
+    src0 = (global float*)((global char*)src0 + offset0);
+    dst  = (global float*)((global char*)dst + offsetd);
+
+    dst[get_global_id(0)] = fabs(src0[get_global_id(0)]);
+}
+
+kernel void kernel_abs_f32_4(
+        global const float4 * src0,
+        ulong                 offset0,
+        global       float4 * dst,
+        ulong                 offsetd
+) {
+    src0 = (global float4*)((global char*)src0 + offset0);
+    dst  = (global float4*)((global char*)dst + offsetd);
+
+    dst[get_global_id(0)] = fabs(src0[get_global_id(0)]);
+}
+
+kernel void kernel_abs_f16(
+        global const half * src0,
+        ulong               offset0,
+        global       half * dst,
+        ulong               offsetd
+) {
+    src0 = (global half*)((global char*)src0 + offset0);
+    dst  = (global half*)((global char*)dst + offsetd);
+
+    dst[get_global_id(0)] = fabs(src0[get_global_id(0)]);
+}
+
+kernel void kernel_abs_f16_4(
+        global const half4 * src0,
+        ulong                offset0,
+        global       half4 * dst,
+        ulong                offsetd
+) {
+    src0 = (global half4*)((global char*)src0 + offset0);
+    dst  = (global half4*)((global char*)dst + offsetd);
+
+    dst[get_global_id(0)] = fabs(src0[get_global_id(0)]);
+}
+
+kernel void kernel_abs_f32_nc(
+        global const char * src0,
+        ulong               offset0,
+        global       char * dst,
+        ulong               offsetd,
+        int   ne00,
+        ulong nb00,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        ulong nb0,
+        ulong nb1,
+        ulong nb2,
+        ulong nb3
+) {
+    src0 = src0 + offset0;
+    dst  = dst + offsetd;
+
+    const int i3 = get_group_id(2);
+    const int i2 = get_group_id(1);
+    const int i1 = get_group_id(0);
+
+    for (int i0 = get_local_id(0); i0 < ne00; i0 += get_local_size(0)) {
+        global const float * x = (global const float *)(src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
+        global       float * y = (global       float *)(dst  + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
+
+        *y = fabs(*x);
+    }
+}
+
+kernel void kernel_abs_f16_nc(
+        global const char * src0,
+        ulong               offset0,
+        global       char * dst,
+        ulong               offsetd,
+        int   ne00,
+        ulong nb00,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        ulong nb0,
+        ulong nb1,
+        ulong nb2,
+        ulong nb3
+) {
+    src0 = src0 + offset0;
+    dst  = dst + offsetd;
+
+    const int i3 = get_group_id(2);
+    const int i2 = get_group_id(1);
+    const int i1 = get_group_id(0);
+
+    for (int i0 = get_local_id(0); i0 < ne00; i0 += get_local_size(0)) {
+        global const half * x = (global const half *)(src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
+        global       half * y = (global       half *)(dst  + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
+
+        *y = fabs(*x);
+    }
+}

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

@@ -37,15 +37,17 @@ static inline float e8m0_to_fp32(uchar x) {
 // One token's dp4a dot (8 uints = 32 K elems) + mxfp4 block-scale epilogue.
 // blk_scale already carries the 0.5 factor (== 0.5 * 2^e).
 #define MOE_MXFP4_DP4A_T(t) do {                                     \
+        uint4 a0 = vload4(0, &sh_qa[t][0]);                          \
+        uint4 a1 = vload4(0, &sh_qa[t][4]);                          \
         int raw = 0;                                                 \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[0], sh_qa[t][0], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[1], sh_qa[t][1], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[2], sh_qa[t][2], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[3], sh_qa[t][3], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[4], sh_qa[t][4], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[5], sh_qa[t][5], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[6], sh_qa[t][6], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[7], sh_qa[t][7], raw); \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[0], a0.s0, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[1], a0.s1, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[2], a0.s2, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[3], a0.s3, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[4], a1.s0, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[5], a1.s1, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[6], a1.s2, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[7], a1.s3, raw);       \
         acc[t] += blk_scale * (float)sh_d[t] * (float)raw;           \
     } while (0)
 
@@ -133,11 +135,13 @@ kernel void kernel_gemm_moe_mxfp4_q8_1_dp4a(
         qw[6] = mxfp4_pack((ushort)(r3));        qw[7] = mxfp4_pack((ushort)(r3 >> 16));
 
         // cooperatively stage the n_real-token x 32-K int8 activations
-        const uint stage_lim = (uint)n_real * 8;
-        for (uint idx = lid; idx < stage_lim; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
+        // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores.
+        const uint vlim = (uint)n_real * 2;
+        for (uint idx = lid; idx < vlim; idx += 64) {
+            const uint t = idx >> 1;
+            const uint h = (idx & 1) << 2;   // 0 or 4
+            uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]);
+            vstore4(v, 0, &sh_qa[t][h]);
         }
         if (lid < (uint)n_real) {
             sh_d[lid] = src1_da[(col + lid) * num_blocks + sub];

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

@@ -17,15 +17,17 @@
 
 // One token's dp4a dot (8 uints = 32 K elems) + q4_0 scale/zero-point epilogue.
 #define MOE_Q40_DP4A_T(t) do {                                       \
+        uint4 a0 = vload4(0, &sh_qa[t][0]);                          \
+        uint4 a1 = vload4(0, &sh_qa[t][4]);                          \
         int raw = 0;                                                 \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[0], sh_qa[t][0], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[1], sh_qa[t][1], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[2], sh_qa[t][2], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[3], sh_qa[t][3], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[4], sh_qa[t][4], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[5], sh_qa[t][5], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[6], sh_qa[t][6], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[7], sh_qa[t][7], raw); \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[0], a0.s0, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[1], a0.s1, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[2], a0.s2, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[3], a0.s3, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[4], a1.s0, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[5], a1.s1, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[6], a1.s2, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[7], a1.s3, raw);       \
         acc[t] += d_val * ((float)sh_d[t] * (float)raw - 8.0f * (float)sh_s[t]); \
     } while (0)
 
@@ -112,11 +114,13 @@ kernel void kernel_gemm_moe_q4_0_q8_1_dp4a(
         qw[6] = EXP4(r3);        qw[7] = EXP4(r3 >> 16);
 
         // cooperatively stage the n_real-token x 32-K int8 activations
-        const uint stage_lim = (uint)n_real * 8;
-        for (uint idx = lid; idx < stage_lim; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
+        // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores.
+        const uint vlim = (uint)n_real * 2;
+        for (uint idx = lid; idx < vlim; idx += 64) {
+            const uint t = idx >> 1;
+            const uint h = (idx & 1) << 2;   // 0 or 4
+            uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]);
+            vstore4(v, 0, &sh_qa[t][h]);
         }
         if (lid < (uint)n_real) {
             sh_d[lid] = src1_da[(col + lid) * num_blocks + sub];

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

@@ -37,16 +37,21 @@ inline void get_scale_min_k4(
                   (((uint)((u) & 0xF000u)) << 12) )
 
 // One token's dp4a dot (8 uints = 32 K elems) + q4_K scale/min epilogue into acc[t].
+// The 8 activation uints are read as two 128-bit uint4 loads staged to private (Adreno
+// wants 128-bit local reads, and a __local operand fed straight to the dp4a builtin is
+// slower and can miscompile).
 #define MOE_Q4K_DP4A_T(t) do {                                       \
+        uint4 a0 = vload4(0, &sh_qa[t][0]);                          \
+        uint4 a1 = vload4(0, &sh_qa[t][4]);                          \
         int raw = 0;                                                 \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[0], sh_qa[t][0], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[1], sh_qa[t][1], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[2], sh_qa[t][2], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[3], sh_qa[t][3], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[4], sh_qa[t][4], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[5], sh_qa[t][5], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[6], sh_qa[t][6], raw); \
-        raw = dot_acc_sat_4x8packed_ss_int(qw[7], sh_qa[t][7], raw); \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[0], a0.s0, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[1], a0.s1, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[2], a0.s2, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[3], a0.s3, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[4], a1.s0, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[5], a1.s1, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[6], a1.s2, raw);       \
+        raw = dot_acc_sat_4x8packed_ss_int(qw[7], a1.s3, raw);       \
         acc[t] += scale * (float)sh_d[t] * (float)raw - minv * (float)sh_s[t]; \
     } while (0)
 
@@ -145,12 +150,14 @@ kernel void kernel_gemm_moe_q4_k_q8_1_dp4a(
         qw[4] = EXP4(r2);        qw[5] = EXP4(r2 >> 16);
         qw[6] = EXP4(r3);        qw[7] = EXP4(r3 >> 16);
 
-        // --- cooperatively stage the n_real-token x 32-K int8 activations to LDS ---
-        const uint stage_lim = (uint)n_real * 8;
-        for (uint idx = lid; idx < stage_lim; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
+        // cooperatively stage the n_real-token x 32-K int8 activations to lm
+        // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores.
+        const uint vlim = (uint)n_real * 2;
+        for (uint idx = lid; idx < vlim; idx += 64) {
+            const uint t = idx >> 1;
+            const uint h = (idx & 1) << 2;   // 0 or 4
+            uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]);
+            vstore4(v, 0, &sh_qa[t][h]);
         }
         if (lid < (uint)n_real) {
             sh_d[lid] = src1_da[(col + lid) * ne00_b + sub];

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

@@ -41,8 +41,10 @@ inline int dp4a_q6(uint qw0, uint qw1, uint qw2, uint qw3,
 
 // One token's q6_K dp4a dot (two halves, per-16 scales) + epilogue into acc[t].
 #define MOE_Q6K_DP4A_T(t) do {                                                                            \
-        const int raw1 = dp4a_q6(qw[0], qw[1], qw[2], qw[3], sh_qa[t][0], sh_qa[t][1], sh_qa[t][2], sh_qa[t][3]); \
-        const int raw2 = dp4a_q6(qw[4], qw[5], qw[6], qw[7], sh_qa[t][4], sh_qa[t][5], sh_qa[t][6], sh_qa[t][7]); \
+        uint4 a0 = vload4(0, &sh_qa[t][0]);                                                               \
+        uint4 a1 = vload4(0, &sh_qa[t][4]);                                                               \
+        const int raw1 = dp4a_q6(qw[0], qw[1], qw[2], qw[3], a0.s0, a0.s1, a0.s2, a0.s3);                 \
+        const int raw2 = dp4a_q6(qw[4], qw[5], qw[6], qw[7], a1.s0, a1.s1, a1.s2, a1.s3);                 \
         const float a_d = (float)sh_d[t];                                                                 \
         acc[t] += scale0 * a_d * (float)raw1 + scale1 * a_d * (float)raw2;                                \
     } while (0)
@@ -144,11 +146,13 @@ kernel void kernel_gemm_moe_q6_k_q8_1_dp4a(
         qw[6] = SIGN6(EXP4(r3)       | EXP2((qh2 >> 16) & 0xFFu));
         qw[7] = SIGN6(EXP4(r3 >> 16) | EXP2((qh2 >> 24) & 0xFFu));
 
-        const uint stage_lim = (uint)n_real * 8;
-        for (uint idx = lid; idx < stage_lim; idx += 64) {
-            const uint t = idx >> 3;
-            const uint u = idx & 7;
-            sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
+        // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores.
+        const uint vlim = (uint)n_real * 2;
+        for (uint idx = lid; idx < vlim; idx += 64) {
+            const uint t = idx >> 1;
+            const uint h = (idx & 1) << 2;   // 0 or 4
+            uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]);
+            vstore4(v, 0, &sh_qa[t][h]);
         }
         if (lid < (uint)n_real) {
             sh_d[lid] = src1_da[(col + lid) * ne00_b + sub];

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

@@ -102,8 +102,10 @@ inline int dp4a4(uint w0,uint w1,uint w2,uint w3,uint a0,uint a1,uint a2,uint a3
 
 // One token's two-half dp4a + uniform scale/min epilogue into acc[t].
 #define MOE_DP4A_T(t) do {                                                                  \
-        const int raw1 = dp4a4(qw[0],qw[1],qw[2],qw[3], sh_qa[t][0],sh_qa[t][1],sh_qa[t][2],sh_qa[t][3]); \
-        const int raw2 = dp4a4(qw[4],qw[5],qw[6],qw[7], sh_qa[t][4],sh_qa[t][5],sh_qa[t][6],sh_qa[t][7]); \
+        uint4 a0 = vload4(0, &sh_qa[t][0]);                                                 \
+        uint4 a1 = vload4(0, &sh_qa[t][4]);                                                 \
+        const int raw1 = dp4a4(qw[0],qw[1],qw[2],qw[3], a0.s0,a0.s1,a0.s2,a0.s3);           \
+        const int raw2 = dp4a4(qw[4],qw[5],qw[6],qw[7], a1.s0,a1.s1,a1.s2,a1.s3);           \
         const float a_d = (float)sh_d[t];                                                   \
         acc[t] += sc0*a_d*(float)raw1 + sc1*a_d*(float)raw2 - mn*(float)sh_s[t];             \
     } while (0)
@@ -178,10 +180,13 @@ kernel void kernel_gemm_moe_q8_1_dp4a(
 
         LOAD_QW(step, sub)
 
-        const uint stage_lim = (uint)n_real * 8;
-        for (uint idx = lid; idx < stage_lim; idx += 64) {
-            const uint t = idx >> 3, u = idx & 7;
-            sh_qa[t][u] = src1_qa[(col + t) * ne00_u + (step >> 2) + u];
+        // Stage each token's 8 activation uints as two 128-bit uint4 loads/stores.
+        const uint vlim = (uint)n_real * 2;
+        for (uint idx = lid; idx < vlim; idx += 64) {
+            const uint t = idx >> 1;
+            const uint h = (idx & 1) << 2;   // 0 or 4
+            uint4 v = vload4(0, &src1_qa[(col + t) * ne00_u + (step >> 2) + h]);
+            vstore4(v, 0, &sh_qa[t][h]);
         }
         if (lid < (uint)n_real) {
             sh_d[lid] = src1_da[(col + lid) * ne00_b + sub];

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

@@ -8,9 +8,11 @@
 #define QK_K         256
 #define K_SCALE_SIZE 12
 
+// scales are transposed: consecutive codes of a row are `stride` apart
 inline void get_scale_min_k4(
     int j,
     global const uchar * q,
+    int stride,
     uchar * d,
     uchar * m,
     uchar mask_d6,
@@ -18,11 +20,11 @@ inline void get_scale_min_k4(
     uchar mask_hi2
 ) {
     if (j < 4) {
-        *d = q[j]   & mask_d6;
-        *m = q[j+4] & mask_d6;
+        *d = q[j*stride]     & mask_d6;
+        *m = q[(j+4)*stride] & mask_d6;
     } else {
-        *d = (q[j+4] & mask_d4) | ((q[j-4] & mask_hi2) >> 2);
-        *m = ((q[j+4] >> 4) & mask_d4) | ((q[j]   & mask_hi2) >> 2);
+        *d = (q[(j+4)*stride] & mask_d4) | ((q[(j-4)*stride] & mask_hi2) >> 2);
+        *m = ((q[(j+4)*stride] >> 4) & mask_d4) | ((q[j*stride] & mask_hi2) >> 2);
     }
 }
 
@@ -55,7 +57,6 @@ kernel void kernel_gemm_noshuffle_q4_k_f32(
     half8 B;
     half4 dequantized_weights;
 
-    int num_blocks_K = k / QK_K;
 
     global const ushort * weight_ptr = src0_q + gx_2;
     global const half   * d_ptr      = src0_d  + gx_2;
@@ -68,16 +69,16 @@ kernel void kernel_gemm_noshuffle_q4_k_f32(
         half4 d  = vload4(0, d_ptr  + sb_idx * m);
         half4 dm = vload4(0, dm_ptr + sb_idx * m);
 
-        global const uchar * sc0 = src0_s + (gx_2+0) * num_blocks_K * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
-        global const uchar * sc1 = src0_s + (gx_2+1) * num_blocks_K * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
-        global const uchar * sc2 = src0_s + (gx_2+2) * num_blocks_K * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
-        global const uchar * sc3 = src0_s + (gx_2+3) * num_blocks_K * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
+        global const uchar * sc0 = src0_s + sb_idx * K_SCALE_SIZE * m + (gx_2+0);
+        global const uchar * sc1 = sc0 + 1;
+        global const uchar * sc2 = sc0 + 2;
+        global const uchar * sc3 = sc0 + 3;
 
         uchar sv0, mn0, sv1, mn1, sv2, mn2, sv3, mn3;
-        get_scale_min_k4(sub_idx, sc0, &sv0, &mn0, mask_d6, mask_d4, mask_hi2);
-        get_scale_min_k4(sub_idx, sc1, &sv1, &mn1, mask_d6, mask_d4, mask_hi2);
-        get_scale_min_k4(sub_idx, sc2, &sv2, &mn2, mask_d6, mask_d4, mask_hi2);
-        get_scale_min_k4(sub_idx, sc3, &sv3, &mn3, mask_d6, mask_d4, mask_hi2);
+        get_scale_min_k4(sub_idx, sc0, m, &sv0, &mn0, mask_d6, mask_d4, mask_hi2);
+        get_scale_min_k4(sub_idx, sc1, m, &sv1, &mn1, mask_d6, mask_d4, mask_hi2);
+        get_scale_min_k4(sub_idx, sc2, m, &sv2, &mn2, mask_d6, mask_d4, mask_hi2);
+        get_scale_min_k4(sub_idx, sc3, m, &sv3, &mn3, mask_d6, mask_d4, mask_hi2);
 
         half4 scale = convert_half4(convert_float4(d)  * convert_float4((uchar4)(sv0, sv1, sv2, sv3)));
         half4 mval  = convert_half4(convert_float4(dm) * convert_float4((uchar4)(mn0, mn1, mn2, mn3)));

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

@@ -10,9 +10,11 @@
 #define QK_K 256
 #define K_SCALE_SIZE 12
 
+// scales are transposed: consecutive codes of a row are `stride` apart
 inline void get_scale_min_k4(
     int j,
     global const uchar * q,
+    uint stride,
     uchar * d,
     uchar * m,
     uchar mask_d6,
@@ -20,11 +22,11 @@ inline void get_scale_min_k4(
     uchar mask_hi2
 ) {
     if (j < 4) {
-        *d = q[j]   & mask_d6;
-        *m = q[j+4] & mask_d6;
+        *d = q[j*stride]     & mask_d6;
+        *m = q[(j+4)*stride] & mask_d6;
     } else {
-        *d = (q[j+4] & mask_d4) | ((q[j-4] & mask_hi2) >> 2);
-        *m = ((q[j+4] >> 4) & mask_d4) | ((q[j]   & mask_hi2) >> 2);
+        *d = (q[(j+4)*stride] & mask_d4) | ((q[(j-4)*stride] & mask_hi2) >> 2);
+        *m = ((q[(j+4)*stride] >> 4) & mask_d4) | ((q[j*stride] & mask_hi2) >> 2);
     }
 }
 
@@ -79,7 +81,6 @@ kernel void kernel_gemm_noshuffle_q4_k_q8_1_dp4a(
     const bool row_valid = row < (uint)m;
     const uint rrow     = row_valid ? row : 0;  // clamp OOB rows; their writes are masked
 
-    const uint num_superblocks = (uint)k / QK_K;
     const uint k_u = (uint)k >> 2;   // K in uint (int8x4) units
     const uint k_b = (uint)k >> 5;   // blocks-of-32 along K
 
@@ -101,9 +102,9 @@ kernel void kernel_gemm_noshuffle_q4_k_q8_1_dp4a(
         // weight scale/min for this WI's row, this subblock
         const float dd  = (float)src0_d [rrow + sb_idx * m];
         const float dmm = (float)src0_dm[rrow + sb_idx * m];
-        global const uchar * sc = src0_s + rrow * num_superblocks * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
+        global const uchar * sc = src0_s + sb_idx * K_SCALE_SIZE * (uint)m + rrow;
         uchar sv, mn;
-        get_scale_min_k4(sub_idx, sc, &sv, &mn, mask_d6, mask_d4, mask_hi2);
+        get_scale_min_k4(sub_idx, sc, (uint)m, &sv, &mn, mask_d6, mask_d4, mask_hi2);
         const float scale = dd  * (float)sv;
         const float minv  = dmm * (float)mn;
 
@@ -202,7 +203,6 @@ kernel void kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg(
 
     const uint k_u = (uint)k >> 2;   // K in uint (int8x4) units
     const uint k_b = (uint)k >> 5;   // blocks-of-32 along K
-    const uint num_superblocks = (uint)k / QK_K;
 
     __local uint sh_qa[TILESIZE_N][8];
     __local half sh_d[TILESIZE_N];
@@ -220,9 +220,9 @@ kernel void kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg(
 
         const float dd  = (float)src0_d [rrow + sb_idx * m];
         const float dmm = (float)src0_dm[rrow + sb_idx * m];
-        global const uchar * sc = src0_s + rrow * num_superblocks * K_SCALE_SIZE + sb_idx * K_SCALE_SIZE;
+        global const uchar * sc = src0_s + sb_idx * K_SCALE_SIZE * (uint)m + rrow;
         uchar sv, mn;
-        get_scale_min_k4(sub_idx, sc, &sv, &mn, mask_d6, mask_d4, mask_hi2);
+        get_scale_min_k4(sub_idx, sc, (uint)m, &sv, &mn, mask_d6, mask_d4, mask_hi2);
         const float scale = dd  * (float)sv;
         const float minv  = dmm * (float)mn;
 

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

@@ -11,9 +11,11 @@
 #define NSUBGROUPS 4
 #define SUBGROUP_SIZE 64
 
+// scales are transposed: consecutive codes of a row are `stride` apart
 inline void get_scale_min_k4(
     int j,
     global const uchar * q,
+    uint stride,
     uchar * d,
     uchar * m,
     uchar mask_d6,
@@ -21,11 +23,11 @@ inline void get_scale_min_k4(
     uchar mask_hi2
 ) {
     if (j < 4) {
-        *d = q[j]   & mask_d6;
-        *m = q[j+4] & mask_d6;
+        *d = q[j*stride]     & mask_d6;
+        *m = q[(j+4)*stride] & mask_d6;
     } else {
-        *d = (q[j+4] & mask_d4) | ((q[j-4] & mask_hi2) >> 2);
-        *m = ((q[j+4] >> 4) & mask_d4) | ((q[j]   & mask_hi2) >> 2);
+        *d = (q[(j+4)*stride] & mask_d4) | ((q[(j-4)*stride] & mask_hi2) >> 2);
+        *m = ((q[(j+4)*stride] >> 4) & mask_d4) | ((q[j*stride] & mask_hi2) >> 2);
     }
 }
 
@@ -232,7 +234,6 @@ kernel void kernel_gemv_noshuffle_q4_k_f32(
 
     uint LINE_STRIDE_A  = M / 2;
     uint BLOCK_STRIDE_A = NSUBGROUPS * M;
-    uint scales_per_row = (K / QK_K) * 12;
 
     private uint4     regA;
     private half2     regS;
@@ -248,12 +249,12 @@ kernel void kernel_gemv_noshuffle_q4_k_f32(
         half2 d   = src0_d[gid + sb * LINE_STRIDE_A];
         half2 dm  = src0_m[gid + sb * LINE_STRIDE_A];
 
-        global const uchar * sc0 = src0_s + 2 * gid * scales_per_row + sb * 12;
-        global const uchar * sc1 = src0_s + (2 * gid + 1) * scales_per_row + sb * 12;
+        global const uchar * sc0 = src0_s + sb * 12 * M + 2 * gid;
+        global const uchar * sc1 = sc0 + 1;
 
         uchar sv0, mn0, sv1, mn1;
-        get_scale_min_k4(j, sc0, &sv0, &mn0, mask_d6, mask_d4, mask_hi2);
-        get_scale_min_k4(j, sc1, &sv1, &mn1, mask_d6, mask_d4, mask_hi2);
+        get_scale_min_k4(j, sc0, M, &sv0, &mn0, mask_d6, mask_d4, mask_hi2);
+        get_scale_min_k4(j, sc1, M, &sv1, &mn1, mask_d6, mask_d4, mask_hi2);
 
         regS = convert_half2(convert_float2(d)  * convert_float2((uchar2)(sv0, sv1)));
         regM = convert_half2(convert_float2(dm) * convert_float2((uchar2)(mn0, mn1)));

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

@@ -153,18 +153,27 @@ kernel void kernel_mul_mv_q4_K_f32_flat(
 
             global ushort * q2 = q1 + 32;
 
-            float4 acc1 = {0.f, 0.f, 0.f, 0.f};
-            float4 acc2 = {0.f, 0.f, 0.f, 0.f};
-            for (int i = 0; i < 8; i += 2) {
-                acc1.s0 += yl[i+0] * (q1[i/2] & 0x000F);
-                acc1.s1 += yl[i+1] * (q1[i/2] & 0x0F00);
-                acc1.s2 += yl[i+8] * (q1[i/2] & 0x00F0);
-                acc1.s3 += yl[i+9] * (q1[i/2] & 0xF000);
-                acc2.s0 += yh[i+0] * (q2[i/2] & 0x000F);
-                acc2.s1 += yh[i+1] * (q2[i/2] & 0x0F00);
-                acc2.s2 += yh[i+8] * (q2[i/2] & 0x00F0);
-                acc2.s3 += yh[i+9] * (q2[i/2] & 0xF000);
-            }
+            // Load the 4 q1 / 4 q2 quant ushorts as 2 uints each. 16-bit integer ops are
+            // disproportionately slow on the A7X (E031.41) compiler; keeping the dequant
+            // operands in 32-bit registers avoids the ushort path. q1/q2 are 4-byte aligned
+            // (ib*128 + (32*iq+8*ir) bytes; q1 += blk*128 bytes/row). Math is unchanged:
+            // w & 0x0F00 on the low/high halves equals the original ushort mask value.
+            global uint * q1u = (global uint *)q1;
+            global uint * q2u = (global uint *)q2;
+            uint a0 = q1u[0], a1 = q1u[1];
+            uint b0 = q2u[0], b1 = q2u[1];
+            uint w0 = a0 & 0xFFFF, w1 = a0 >> 16, w2 = a1 & 0xFFFF, w3 = a1 >> 16;
+            uint v0 = b0 & 0xFFFF, v1 = b0 >> 16, v2 = b1 & 0xFFFF, v3 = b1 >> 16;
+
+            float4 acc1, acc2;
+            acc1.s0 = yl[0]*(w0&0x000F) + yl[ 2]*(w1&0x000F) + yl[ 4]*(w2&0x000F) + yl[ 6]*(w3&0x000F);
+            acc1.s1 = yl[1]*(w0&0x0F00) + yl[ 3]*(w1&0x0F00) + yl[ 5]*(w2&0x0F00) + yl[ 7]*(w3&0x0F00);
+            acc1.s2 = yl[8]*(w0&0x00F0) + yl[10]*(w1&0x00F0) + yl[12]*(w2&0x00F0) + yl[14]*(w3&0x00F0);
+            acc1.s3 = yl[9]*(w0&0xF000) + yl[11]*(w1&0xF000) + yl[13]*(w2&0xF000) + yl[15]*(w3&0xF000);
+            acc2.s0 = yh[0]*(v0&0x000F) + yh[ 2]*(v1&0x000F) + yh[ 4]*(v2&0x000F) + yh[ 6]*(v3&0x000F);
+            acc2.s1 = yh[1]*(v0&0x0F00) + yh[ 3]*(v1&0x0F00) + yh[ 5]*(v2&0x0F00) + yh[ 7]*(v3&0x0F00);
+            acc2.s2 = yh[8]*(v0&0x00F0) + yh[10]*(v1&0x00F0) + yh[12]*(v2&0x00F0) + yh[14]*(v3&0x00F0);
+            acc2.s3 = yh[9]*(v0&0xF000) + yh[11]*(v1&0xF000) + yh[13]*(v2&0xF000) + yh[15]*(v3&0xF000);
 
             float dall = *d;
             float dmin = *dm;

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

@@ -159,18 +159,59 @@ kernel void kernel_mul_mv_q5_K_f32_flat(
 
             global ushort * q2 = q1 + 32;
 
-            float4 acc1 = {0.f, 0.f, 0.f, 0.f};
-            float4 acc2 = {0.f, 0.f, 0.f, 0.f};
-            for (int i = 0; i < 8; i += 2) {
-                acc1.s0 += yl[i+0] * ((q1[i/2] & 0x000F) + (qh[i+0] & u1_lo ? 16.f       : 0.f));
-                acc1.s1 += yl[i+1] * ((q1[i/2] & 0x0F00) + (qh[i+1] & u1_lo ? 16.f*256.f : 0.f));
-                acc1.s2 += yl[i+8] * ((q1[i/2] & 0x00F0) + (qh[i+0] & u2_lo ? 16.f*16.f  : 0.f));
-                acc1.s3 += yl[i+9] * ((q1[i/2] & 0xF000) + (qh[i+1] & u2_lo ? 16.f*4096.f: 0.f));
-                acc2.s0 += yh[i+0] * ((q2[i/2] & 0x000F) + (qh[i+0] & u1_hi ? 16.f       : 0.f));
-                acc2.s1 += yh[i+1] * ((q2[i/2] & 0x0F00) + (qh[i+1] & u1_hi ? 16.f*256.f : 0.f));
-                acc2.s2 += yh[i+8] * ((q2[i/2] & 0x00F0) + (qh[i+0] & u2_hi ? 16.f*16.f  : 0.f));
-                acc2.s3 += yh[i+9] * ((q2[i/2] & 0xF000) + (qh[i+1] & u2_hi ? 16.f*4096.f: 0.f));
-            }
+            // Load the 4 q1 / 4 q2 quant ushorts as 2 uints each. 16-bit integer ops are
+            // disproportionately slow on the A7X (E031.41) compiler; keeping the dequant
+            // operands in 32-bit registers avoids the ushort path (same fix as q4_K flat).
+            // q1/q2 are 4-byte aligned; w & 0x0F00 on the low/high half of a uint equals the
+            // original ushort mask value, so the math is unchanged. The qh high-bit term is
+            // byte-indexed (qh[0..7]) and left as-is.
+            global uint * q1u = (global uint *)q1;
+            global uint * q2u = (global uint *)q2;
+            uint a0 = q1u[0], a1 = q1u[1], b0 = q2u[0], b1 = q2u[1];
+            uint w0 = a0 & 0xFFFF, w1 = a0 >> 16, w2 = a1 & 0xFFFF, w3 = a1 >> 16;
+            uint v0 = b0 & 0xFFFF, v1 = b0 >> 16, v2 = b1 & 0xFFFF, v3 = b1 >> 16;
+
+            float4 acc1, acc2;
+            acc1.s0 =
+                yl[0]*((w0&0x000F)+(qh[0]&u1_lo?16.f:0.f)) +
+                yl[2]*((w1&0x000F)+(qh[2]&u1_lo?16.f:0.f)) +
+                yl[4]*((w2&0x000F)+(qh[4]&u1_lo?16.f:0.f)) +
+                yl[6]*((w3&0x000F)+(qh[6]&u1_lo?16.f:0.f));
+            acc1.s1 =
+                yl[1]*((w0&0x0F00)+(qh[1]&u1_lo?16.f*256.f:0.f)) +
+                yl[3]*((w1&0x0F00)+(qh[3]&u1_lo?16.f*256.f:0.f)) +
+                yl[5]*((w2&0x0F00)+(qh[5]&u1_lo?16.f*256.f:0.f)) +
+                yl[7]*((w3&0x0F00)+(qh[7]&u1_lo?16.f*256.f:0.f));
+            acc1.s2 =
+                yl[ 8]*((w0&0x00F0)+(qh[0]&u2_lo?16.f*16.f:0.f)) +
+                yl[10]*((w1&0x00F0)+(qh[2]&u2_lo?16.f*16.f:0.f)) +
+                yl[12]*((w2&0x00F0)+(qh[4]&u2_lo?16.f*16.f:0.f)) +
+                yl[14]*((w3&0x00F0)+(qh[6]&u2_lo?16.f*16.f:0.f));
+            acc1.s3 =
+                yl[ 9]*((w0&0xF000)+(qh[1]&u2_lo?16.f*4096.f:0.f)) +
+                yl[11]*((w1&0xF000)+(qh[3]&u2_lo?16.f*4096.f:0.f)) +
+                yl[13]*((w2&0xF000)+(qh[5]&u2_lo?16.f*4096.f:0.f)) +
+                yl[15]*((w3&0xF000)+(qh[7]&u2_lo?16.f*4096.f:0.f));
+            acc2.s0 =
+                yh[0]*((v0&0x000F)+(qh[0]&u1_hi?16.f:0.f)) +
+                yh[2]*((v1&0x000F)+(qh[2]&u1_hi?16.f:0.f)) +
+                yh[4]*((v2&0x000F)+(qh[4]&u1_hi?16.f:0.f)) +
+                yh[6]*((v3&0x000F)+(qh[6]&u1_hi?16.f:0.f));
+            acc2.s1 =
+                yh[1]*((v0&0x0F00)+(qh[1]&u1_hi?16.f*256.f:0.f)) +
+                yh[3]*((v1&0x0F00)+(qh[3]&u1_hi?16.f*256.f:0.f)) +
+                yh[5]*((v2&0x0F00)+(qh[5]&u1_hi?16.f*256.f:0.f)) +
+                yh[7]*((v3&0x0F00)+(qh[7]&u1_hi?16.f*256.f:0.f));
+            acc2.s2 =
+                yh[ 8]*((v0&0x00F0)+(qh[0]&u2_hi?16.f*16.f:0.f)) +
+                yh[10]*((v1&0x00F0)+(qh[2]&u2_hi?16.f*16.f:0.f)) +
+                yh[12]*((v2&0x00F0)+(qh[4]&u2_hi?16.f*16.f:0.f)) +
+                yh[14]*((v3&0x00F0)+(qh[6]&u2_hi?16.f*16.f:0.f));
+            acc2.s3 =
+                yh[ 9]*((v0&0xF000)+(qh[1]&u2_hi?16.f*4096.f:0.f)) +
+                yh[11]*((v1&0xF000)+(qh[3]&u2_hi?16.f*4096.f:0.f)) +
+                yh[13]*((v2&0xF000)+(qh[5]&u2_hi?16.f*4096.f:0.f)) +
+                yh[15]*((v3&0xF000)+(qh[7]&u2_hi?16.f*4096.f:0.f));
 
             float dall = *d;
             float dmin = *dm;

ggml/src/ggml-openvino/ggml-openvino.cpp

@@ -1378,3 +1378,5 @@ GGML_BACKEND_API ggml_backend_reg_t ggml_backend_openvino_reg(void) {
 
     return ®
 }
+
+GGML_BACKEND_DL_IMPL(ggml_backend_openvino_reg)

ggml/src/ggml-sycl/dmmv.cpp

@@ -266,7 +266,7 @@ static void dequantize_mul_mat_vec_q2_k(const void *__restrict__ vx,
 
     const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
                     item_ct1.get_local_id(1);
-    if (row > nrows) return;
+    if (row >= nrows) return;
 
     const int num_blocks_per_row = ncols / QK_K;
     const int ib0 = row*num_blocks_per_row;
@@ -387,7 +387,7 @@ static void dequantize_mul_mat_vec_q2_k_reorder(const void *__restrict__ vx,
 
     const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
                     item_ct1.get_local_id(1);
-    if (row > nrows) return;
+    if (row >= nrows) return;
 
     const int num_blocks_per_row = ncols / QK_K;
     const int ib0 = row*num_blocks_per_row;
@@ -483,7 +483,7 @@ static void dequantize_mul_mat_vec_q3_k(const void *__restrict__ vx,
 
     const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
                     item_ct1.get_local_id(1);
-    if (row > nrows) return;
+    if (row >= nrows) return;
 
     const int num_blocks_per_row = ncols / QK_K;
     const int ib0 = row*num_blocks_per_row;
@@ -595,7 +595,7 @@ static void dequantize_mul_mat_vec_q3_k_reorder(const void *__restrict__ vx,
 
     const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
                     item_ct1.get_local_id(1);
-    if (row > nrows) return;
+    if (row >= nrows) return;
 
     const int num_blocks_per_row = ncols / QK_K;
     const int ib0 = row*num_blocks_per_row;
@@ -693,7 +693,7 @@ static void dequantize_mul_mat_vec_q4_k(const void *__restrict__ vx,
 
     const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
                     item_ct1.get_local_id(1);
-    if (row > nrows) return;
+    if (row >= nrows) return;
     const int num_blocks_per_row = ncols / QK_K;
     const int ib0 = row*num_blocks_per_row;
 
@@ -841,7 +841,7 @@ static void dequantize_mul_mat_vec_q4_k_reorder(const void *__restrict__ vx,
 
     const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
                     item_ct1.get_local_id(1);
-    if (row > nrows) return;
+    if (row >= nrows) return;
     const int num_blocks_per_row = ncols / QK_K;
     const int ib0 = row*num_blocks_per_row;
 
@@ -994,10 +994,12 @@ static void dequantize_mul_mat_vec_q4_k_reorder(const void *__restrict__ vx,
 static void dequantize_mul_mat_vec_q5_k(const void *__restrict__ vx,
                                         const float *__restrict__ yy,
                                         float *__restrict__ dst,
-                                        const int ncols,
+                                        const int ncols, int nrows,
                                         const sycl::nd_item<3> &item_ct1) {
 
-    const int row = item_ct1.get_group(2);
+    const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
+                    item_ct1.get_local_id(1);
+    if (row >= nrows) return;
     const int num_blocks_per_row = ncols / QK_K;
     const int ib0 = row*num_blocks_per_row;
 
@@ -1126,7 +1128,9 @@ static void dequantize_mul_mat_vec_q5_k_reorder(const void *__restrict__ vx,
                                                 const int ncols, int nrows,
                                                 const sycl::nd_item<3> &item_ct1) {
 
-    const int row = item_ct1.get_group(2);
+    const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
+                    item_ct1.get_local_id(1);
+    if (row >= nrows) return;
     const int num_blocks_per_row = ncols / QK_K;
     const int ib0 = row*num_blocks_per_row;
 
@@ -1148,19 +1152,13 @@ static void dequantize_mul_mat_vec_q5_k_reorder(const void *__restrict__ vx,
     const int tid = item_ct1.get_local_id(2) / 2; // 0...15
     const int ix = item_ct1.get_local_id(2) % 2;
 
-    const int il  = tid/4;     // 0...3
-    const int ir  = tid - 4*il;// 0...3
+    const int il_base = tid/4;      // 0...3
+    const int ir  = tid - 4*il_base;// 0...3
     const int n   = 2;
 
-    const int im = il/2;  // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224
-    const int in = il%2;
+    const int in = il_base%2;
 
     const int l0 = n*(2*ir + in);
-    const int q_offset = 32*im + l0;
-    const int y_offset = 64*im + l0;
-
-    const uint8_t hm1  = 1 << (2*im);
-    const uint8_t hm2  = hm1 << 4;
 
     uint16_t aux[4];
     const uint8_t * sc = (const uint8_t *)aux;
@@ -1171,52 +1169,60 @@ static void dequantize_mul_mat_vec_q5_k_reorder(const void *__restrict__ vx,
     for (int i = ix; i < num_blocks_per_row; i += 2) {
         const int bi = ib0 + i;
 
-        const uint8_t * ql1 = qs_base + bi * (QK_K / 2) + q_offset;
         const uint8_t * qh  = qh_base + bi * (QK_K / 8) + l0;
-        const float   * y1  = yy + i*QK_K + y_offset;
-        const float   * y2  = y1 + 128;
-
         const sycl::half2 dm_val = dm_base[bi];
         const float dall = dm_val[0];
         const float dmin = dm_val[1];
 
-        const uint16_t * a = (const uint16_t *)(scales_base + bi * K_SCALE_SIZE);
-        aux[0] = a[im+0] & kmask1;
-        aux[1] = a[im+2] & kmask1;
-        aux[2] = ((a[im+4] >> 0) & kmask2) | ((a[im+0] & kmask3) >> 2);
-        aux[3] = ((a[im+4] >> 4) & kmask2) | ((a[im+2] & kmask3) >> 2);
+        for (int im = 0; im < 2; ++im) {
+            const int q_offset = 32*im + l0;
+            const int y_offset = 64*im + l0;
 
-        sycl::float4 sum = {0.f, 0.f, 0.f, 0.f};
-        float smin = 0;
-        const uint16_t * q1 = (const uint16_t *)ql1;
-        const uint16_t * q2 = q1 + 32;
-        q16[0] = q1[0] & 0x0f0f;
-        q16[1] = q1[8] & 0x0f0f;
-        q16[2] = (q1[0] >> 4) & 0x0f0f;
-        q16[3] = (q1[8] >> 4) & 0x0f0f;
-        q16[4] = q2[0] & 0x0f0f;
-        q16[5] = q2[8] & 0x0f0f;
-        q16[6] = (q2[0] >> 4) & 0x0f0f;
-        q16[7] = (q2[8] >> 4) & 0x0f0f;
-        for (int l = 0; l < n; ++l) {
-            sum.x() +=
-                y1[l + 0] * (q4[l + 0] + (qh[l + 0] & (hm1 << 0) ? 16 : 0)) +
-                y1[l + 16] * (q4[l + 2] + (qh[l + 16] & (hm1 << 0) ? 16 : 0));
-            sum.y() +=
-                y1[l + 32] * (q4[l + 4] + (qh[l + 0] & (hm1 << 1) ? 16 : 0)) +
-                y1[l + 48] * (q4[l + 6] + (qh[l + 16] & (hm1 << 1) ? 16 : 0));
-            sum.z() +=
-                y2[l + 0] * (q4[l + 8] + (qh[l + 0] & (hm2 << 0) ? 16 : 0)) +
-                y2[l + 16] * (q4[l + 10] + (qh[l + 16] & (hm2 << 0) ? 16 : 0));
-            sum.w() +=
-                y2[l + 32] * (q4[l + 12] + (qh[l + 0] & (hm2 << 1) ? 16 : 0)) +
-                y2[l + 48] * (q4[l + 14] + (qh[l + 16] & (hm2 << 1) ? 16 : 0));
-            smin += (y1[l] + y1[l+16]) * sc[2] + (y1[l+32] + y1[l+48]) * sc[3]
-                  + (y2[l] + y2[l+16]) * sc[6] + (y2[l+32] + y2[l+48]) * sc[7];
+            const uint8_t hm1  = 1 << (2*im);
+            const uint8_t hm2  = hm1 << 4;
+
+            const uint8_t * ql1 = qs_base + bi * (QK_K / 2) + q_offset;
+            const float   * y1  = yy + i*QK_K + y_offset;
+            const float   * y2  = y1 + 128;
+
+            const uint16_t * a = (const uint16_t *)(scales_base + bi * K_SCALE_SIZE);
+            aux[0] = a[im+0] & kmask1;
+            aux[1] = a[im+2] & kmask1;
+            aux[2] = ((a[im+4] >> 0) & kmask2) | ((a[im+0] & kmask3) >> 2);
+            aux[3] = ((a[im+4] >> 4) & kmask2) | ((a[im+2] & kmask3) >> 2);
+
+            sycl::float4 sum = {0.f, 0.f, 0.f, 0.f};
+            float smin = 0;
+            const uint16_t * q1 = (const uint16_t *)ql1;
+            const uint16_t * q2 = q1 + 32;
+            q16[0] = q1[0] & 0x0f0f;
+            q16[1] = q1[8] & 0x0f0f;
+            q16[2] = (q1[0] >> 4) & 0x0f0f;
+            q16[3] = (q1[8] >> 4) & 0x0f0f;
+            q16[4] = q2[0] & 0x0f0f;
+            q16[5] = q2[8] & 0x0f0f;
+            q16[6] = (q2[0] >> 4) & 0x0f0f;
+            q16[7] = (q2[8] >> 4) & 0x0f0f;
+            for (int l = 0; l < n; ++l) {
+                sum.x() +=
+                    y1[l + 0] * (q4[l + 0] + (qh[l + 0] & (hm1 << 0) ? 16 : 0)) +
+                    y1[l + 16] * (q4[l + 2] + (qh[l + 16] & (hm1 << 0) ? 16 : 0));
+                sum.y() +=
+                    y1[l + 32] * (q4[l + 4] + (qh[l + 0] & (hm1 << 1) ? 16 : 0)) +
+                    y1[l + 48] * (q4[l + 6] + (qh[l + 16] & (hm1 << 1) ? 16 : 0));
+                sum.z() +=
+                    y2[l + 0] * (q4[l + 8] + (qh[l + 0] & (hm2 << 0) ? 16 : 0)) +
+                    y2[l + 16] * (q4[l + 10] + (qh[l + 16] & (hm2 << 0) ? 16 : 0));
+                sum.w() +=
+                    y2[l + 32] * (q4[l + 12] + (qh[l + 0] & (hm2 << 1) ? 16 : 0)) +
+                    y2[l + 48] * (q4[l + 14] + (qh[l + 16] & (hm2 << 1) ? 16 : 0));
+                smin += (y1[l] + y1[l+16]) * sc[2] + (y1[l+32] + y1[l+48]) * sc[3]
+                      + (y2[l] + y2[l+16]) * sc[6] + (y2[l+32] + y2[l+48]) * sc[7];
+            }
+            tmp += dall * (sum.x() * sc[0] + sum.y() * sc[1] + sum.z() * sc[4] +
+                           sum.w() * sc[5]) -
+                   dmin * smin;
         }
-        tmp += dall * (sum.x() * sc[0] + sum.y() * sc[1] + sum.z() * sc[4] +
-                       sum.w() * sc[5]) -
-               dmin * smin;
     }
 #else
     // The reordered Q5_K layout is only produced for QK_K == 256.
@@ -1241,7 +1247,7 @@ static void dequantize_mul_mat_vec_q6_k(const void * __restrict__ vx, const floa
 
     const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
                     item_ct1.get_local_id(1);
-    if (row > nrows) return;
+    if (row >= nrows) return;
 
     const int num_blocks_per_row = ncols / QK_K;
     const int ib0 = row*num_blocks_per_row;
@@ -1358,7 +1364,7 @@ static void dequantize_mul_mat_vec_q6_k_reorder(const void * __restrict__ vx, co
 
     const int row = item_ct1.get_group(2) * item_ct1.get_local_range(1) +
                     item_ct1.get_local_id(1);
-    if (row > nrows) return;
+    if (row >= nrows) return;
 
     const int num_blocks_per_row = ncols / QK_K;
     const int ib0 = row*num_blocks_per_row;
@@ -1831,11 +1837,14 @@ static void dequantize_mul_mat_vec_q5_K_sycl(const void *vx, const float *y,
                                              const int nrows,
                                              dpct::queue_ptr stream) {
     GGML_ASSERT(ncols % QK_K == 0);
-    const sycl::range<3> block_dims(1, 1, WARP_SIZE);
+    const int ny = 2 / K_QUANTS_PER_ITERATION;
+    const int block_num_y = (nrows + ny - 1) / ny;
+    const sycl::range<3> block_nums(1, 1, block_num_y);
+    const sycl::range<3> block_dims(1, ny, WARP_SIZE);
     stream->parallel_for(
-        sycl::nd_range<3>(sycl::range<3>(1, 1, nrows) * block_dims, block_dims),
+        sycl::nd_range<3>(block_nums * block_dims, block_dims),
         [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-            dequantize_mul_mat_vec_q5_k(vx, y, dst, ncols, item_ct1);
+            dequantize_mul_mat_vec_q5_k(vx, y, dst, ncols, nrows, item_ct1);
         });
 }
 

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

@@ -156,6 +156,24 @@ typedef struct VkPhysicalDeviceShaderFloat8FeaturesEXT {
 } VkPhysicalDeviceShaderFloat8FeaturesEXT;
 #endif
 
+#ifndef VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME
+#define VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME "VK_KHR_internally_synchronized_queues"
+#define VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_INTERNALLY_SYNCHRONIZED_QUEUES_FEATURES_KHR ((VkStructureType)1000504000)
+#define VK_DEVICE_QUEUE_CREATE_INTERNALLY_SYNCHRONIZED_BIT_KHR ((VkDeviceQueueCreateFlagBits)0x00000004)
+
+// Compile-time constant guaranteed; no runtime initialization overhead
+static constexpr vk::DeviceQueueCreateFlagBits eInternallySynchronizedKHR =
+    static_cast<vk::DeviceQueueCreateFlagBits>(0x00000004);
+
+typedef struct VkPhysicalDeviceInternallySynchronizedQueuesFeaturesKHR {
+    VkStructureType    sType;
+    void* pNext;
+    VkBool32           internallySynchronizedQueues;
+} VkPhysicalDeviceInternallySynchronizedQueuesFeaturesKHR;
+#else
+static constexpr vk::DeviceQueueCreateFlagBits eInternallySynchronizedKHR = vk::DeviceQueueCreateFlagBits::eInternallySynchronizedKHR;
+#endif
+
 #define ROUNDUP_POW2(M, N) (((M) + (N) - 1) & ~((N) - 1))
 #define CEIL_DIV(M, N) (((M) / (N)) + (((M) % (N)) != 0))
 static bool is_pow2(uint32_t x) { return x > 1 && (x & (x-1)) == 0; }
@@ -285,27 +303,41 @@ struct vk_command_pool {
 };
 
 // Prevent simultaneous submissions to the same queue.
-// This could be per vk_queue if we stopped having two vk_queue structures
-// sharing the same vk::Queue.
-static std::mutex queue_mutex;
+struct vk_queue_handle {
+    vk::Queue queue;
+    virtual void submit(vk::ArrayProxy<const vk::SubmitInfo> submits, vk::Fence fence) = 0;
+    virtual void lock()   {}   // no-op by default (internally synchronized case)
+    virtual void unlock() {}
+    virtual ~vk_queue_handle() = default;
+};
+
+struct vk_queue_handle_synchronized : vk_queue_handle {
+    std::mutex mutex;
+    void submit(vk::ArrayProxy<const vk::SubmitInfo> submits, vk::Fence fence) override {
+        std::lock_guard<std::mutex> guard(mutex);
+        queue.submit(submits, fence);
+    }
+    void lock()   override { mutex.lock(); }
+    void unlock() override { mutex.unlock(); }
+};
+
+struct vk_queue_handle_unsynchronized : vk_queue_handle {
+    void submit(vk::ArrayProxy<const vk::SubmitInfo> submits, vk::Fence fence) override {
+        // Driver guarantees internal synchronization via VK_KHR_internally_synchronized_queues
+        queue.submit(submits, fence);
+    }
+    // lock()/unlock() inherited no-ops
+};
 
 struct vk_queue {
     uint32_t queue_family_index;
-    vk::Queue queue;
+    std::shared_ptr<vk_queue_handle> handle;
 
     vk_command_pool cmd_pool;
 
     vk::PipelineStageFlags stage_flags;
 
     bool transfer_only;
-
-    // copy everything except the cmd_pool
-    void copyFrom(vk_queue &other) {
-        queue_family_index = other.queue_family_index;
-        queue = other.queue;
-        stage_flags = other.stage_flags;
-        transfer_only = other.transfer_only;
-    }
 };
 
 static const char * ggml_backend_vk_buffer_type_name(ggml_backend_buffer_type_t buft);
@@ -712,11 +744,12 @@ struct vk_device_struct {
     uint32_t vendor_id;
     vk::DriverId driver_id;
     vk_device_architecture architecture;
-    vk_queue compute_queue;
-    vk_queue transfer_queue;
+    std::unique_ptr<vk_queue> compute_queue;
+    std::unique_ptr<vk_queue> transfer_queue;
     bool single_queue;
     bool support_async;
     bool async_use_transfer_queue;
+    bool has_internally_synchronized_queues = false;
     uint32_t subgroup_size;
     uint32_t subgroup_size_log2;
     uint32_t shader_core_count;
@@ -1019,8 +1052,13 @@ struct vk_device_struct {
 
         ggml_vk_destroy_buffer(sync_staging);
 
-        compute_queue.cmd_pool.destroy(device);
-        transfer_queue.cmd_pool.destroy(device);
+        if (compute_queue) compute_queue->cmd_pool.destroy(device);
+        if (transfer_queue) transfer_queue->cmd_pool.destroy(device);
+
+        // Explicitly clear to ensure queues drop their shared_ptrs to handles
+        // before the Vulkan logical device instance is destroyed
+        compute_queue.reset();
+        transfer_queue.reset();
 
         for (auto& pipeline : all_pipelines) {
             if (pipeline.expired()) {
@@ -2909,8 +2947,7 @@ static vk_command_buffer* ggml_vk_create_cmd_buffer(vk_device& device, vk_comman
 static void ggml_vk_submit(vk_context& ctx, vk::Fence fence) {
     if (ctx->seqs.empty()) {
         if (fence) {
-            std::lock_guard<std::mutex> guard(queue_mutex);
-            ctx->p->q->queue.submit({}, fence);
+            ctx->p->q->handle->submit({}, fence);
         }
         return;
     }
@@ -2979,8 +3016,7 @@ static void ggml_vk_submit(vk_context& ctx, vk::Fence fence) {
         }
     }
 
-    std::lock_guard<std::mutex> guard(queue_mutex);
-    ctx->p->q->queue.submit(submit_infos, fence);
+    ctx->p->q->handle->submit(submit_infos, fence);
 
     ctx->seqs.clear();
 }
@@ -3031,18 +3067,44 @@ static uint32_t ggml_vk_find_queue_family_index(std::vector<vk::QueueFamilyPrope
     abort();
 }
 
-static void ggml_vk_create_queue(vk_device& device, vk_queue& q, uint32_t queue_family_index, uint32_t queue_index, vk::PipelineStageFlags&& stage_flags, bool transfer_only) {
+static std::unique_ptr<vk_queue> ggml_vk_create_queue(vk_device& device, uint32_t queue_family_index, uint32_t queue_index, vk::PipelineStageFlags&& stage_flags, bool transfer_only) {
     VK_LOG_DEBUG("ggml_vk_create_queue()");
     std::lock_guard<std::recursive_mutex> guard(device->mutex);
 
-    q.queue_family_index = queue_family_index;
-    q.transfer_only = transfer_only;
+    auto q = std::make_unique<vk_queue>();
+    q->queue_family_index = queue_family_index;
+    q->transfer_only = transfer_only;
 
-    q.cmd_pool.init(device, &q);
+    std::shared_ptr<vk_queue_handle> h;
+    vk::DeviceQueueInfo2 queue_info2{};
+    queue_info2.queueFamilyIndex = queue_family_index;
+    queue_info2.queueIndex = queue_index;
 
-    q.queue = device->device.getQueue(queue_family_index, queue_index);
+    if (device->has_internally_synchronized_queues) {
+        h = std::make_shared<vk_queue_handle_unsynchronized>();
+        queue_info2.flags = eInternallySynchronizedKHR;
+    } else {
+        h = std::make_shared<vk_queue_handle_synchronized>();
+    }
 
-    q.stage_flags = stage_flags;
+    h->queue = device->device.getQueue2(queue_info2);
+    q->handle = h;
+
+    q->cmd_pool.init(device, q.get());
+
+    q->stage_flags = stage_flags;
+    return q;
+}
+
+static std::unique_ptr<vk_queue> ggml_vk_create_aliased_queue(vk_device& device, const std::unique_ptr<vk_queue>& source) {
+    std::lock_guard<std::recursive_mutex> guard(device->mutex);
+    auto q = std::make_unique<vk_queue>();
+    q->handle = source->handle;
+    q->queue_family_index = source->queue_family_index;
+    q->stage_flags = source->stage_flags;
+    q->transfer_only = source->transfer_only;
+    q->cmd_pool.init(device, q.get());
+    return q;
 }
 
 static vk_context ggml_vk_create_context(ggml_backend_vk_context * ctx, vk_command_pool& p) {
@@ -3107,11 +3169,11 @@ static void ggml_vk_queue_command_pools_cleanup(vk_device& device) {
     // Arbitrary frequency to cleanup/reuse command buffers
     static constexpr uint32_t cleanup_frequency = 10;
 
-    if (device->compute_queue.cmd_pool.buffers_in_use() >= cleanup_frequency) {
-        ggml_vk_command_pool_cleanup(device, device->compute_queue.cmd_pool);
+    if (device->compute_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) {
+        ggml_vk_command_pool_cleanup(device, device->compute_queue->cmd_pool);
     }
-    if (device->transfer_queue.cmd_pool.buffers_in_use() >= cleanup_frequency) {
-        ggml_vk_command_pool_cleanup(device, device->transfer_queue.cmd_pool);
+    if (device->transfer_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) {
+        ggml_vk_command_pool_cleanup(device, device->transfer_queue->cmd_pool);
     }
 }
 
@@ -3735,6 +3797,7 @@ static bool ggml_vk_matmul_int_shmem_support(const vk_device& device, const std:
 
     uint32_t block_a_size = 0;
     switch (src0_type) {
+        case GGML_TYPE_Q2_0:    block_a_size = std430_size({{32, 4}, {fp_size,  fp_align}});                  break; // qs[8] + dm
         case GGML_TYPE_Q4_0:    block_a_size = std430_size({{16, 4}, {fp_size,  fp_align}});                  break; // qs[16/4] + dm
         case GGML_TYPE_Q4_1:    block_a_size = std430_size({{16, 4}, {fp2_size, fp2_align}});                 break; // qs[16/4] + dm(vec2)
         case GGML_TYPE_Q5_0:    block_a_size = std430_size({{16, 4}, {4, 4}, {fp_size,  fp_align}});          break; // qs[16/4] + qh + dm
@@ -4339,6 +4402,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         }
 #endif
         CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q1_0], matmul_q1_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
+        CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_0], matmul_q2_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
         CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0], matmul_q4_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
         CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_1], matmul_q4_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
         CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_0], matmul_q5_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3)
@@ -4378,6 +4442,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         }
 #endif
         CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
+        CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
         CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
         CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
         CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5)
@@ -4448,6 +4513,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 #endif
 
         CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q1_0], matmul_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
+        CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_0], matmul_q2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
         CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0], matmul_q4_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
         CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1], matmul_q4_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
         CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0], matmul_q5_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, );
@@ -4492,6 +4558,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 #endif
 
         CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
+        CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
         CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
         CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
         CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id);
@@ -4581,6 +4648,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         CREATE_MM_NODOT2(GGML_TYPE_BF16, pipeline_matmul_bf16, matmul_bf16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0);
 
         CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q1_0], matmul_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
+        CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_0], matmul_q2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
         CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0], matmul_q4_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
         CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1], matmul_q4_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
         CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0], matmul_q5_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
@@ -4605,6 +4673,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 
 #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
         if (device->integer_dot_product) {
+            CREATE_MMQ(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q2_0], matmul_q2_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0);
             CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_0], matmul_q4_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0);
             CREATE_MMQ(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_1], matmul_q4_1_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0);
             CREATE_MMQ(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q5_0], matmul_q5_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0);
@@ -4627,6 +4696,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
             CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16_f32, matmul_id_subgroup_f16_f32, wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16);
             CREATE_MM_NODOT2(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_subgroup_bf16, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16);
             CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
+            CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
             CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
             CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
             CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
@@ -4651,6 +4721,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 
 #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
             if (device->integer_dot_product) {
+                CREATE_MMQ(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_q8_1, mmq_wg_denoms, warptile_mmqid_int,   vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
                 CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_q8_1, mmq_wg_denoms, warptile_mmqid_int,   vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
                 CREATE_MMQ(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_q8_1, mmq_wg_denoms, warptile_mmqid_int,   vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
                 CREATE_MMQ(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_q8_1, mmq_wg_denoms, warptile_mmqid_int,   vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
@@ -4672,6 +4743,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
             CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
             CREATE_MM_NODOT2(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
             CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_q1_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
+            CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_q2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
             CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_q4_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
             CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_q4_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
             CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_q5_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
@@ -4696,6 +4768,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 
 #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
             if (device->integer_dot_product) {
+                CREATE_MMQ(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q2_0], matmul_id_q2_0_q8_1, mmq_wg_denoms, warptile_mmqid_int,   vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
                 CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_0], matmul_id_q4_0_q8_1, mmq_wg_denoms, warptile_mmqid_int,   vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
                 CREATE_MMQ(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_1], matmul_id_q4_1_q8_1, mmq_wg_denoms, warptile_mmqid_int,   vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
                 CREATE_MMQ(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q5_0], matmul_id_q5_0_q8_1, mmq_wg_denoms, warptile_mmqid_int,   vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
@@ -4748,6 +4821,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_bf16, matmul_bf16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0);
 
         CREATE_MM(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q1_0].f32acc, matmul_q1_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
+        CREATE_MM(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_0].f32acc, matmul_q2_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
         CREATE_MM(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f32acc, matmul_q4_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
         CREATE_MM(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f32acc, matmul_q4_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
         CREATE_MM(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f32acc, matmul_q5_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0);
@@ -4773,6 +4847,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 
 #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
         if (device->integer_dot_product) {
+            CREATE_MMQ(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q2_0].f32acc, matmul_q2_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, );
             CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_0].f32acc, matmul_q4_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, );
             CREATE_MMQ(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_1].f32acc, matmul_q4_1_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, );
             CREATE_MMQ(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q5_0].f32acc, matmul_q5_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, );
@@ -4794,6 +4869,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
             CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_subgroup_bf16, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16);
 
             CREATE_MM(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0].f32acc, matmul_id_subgroup_q1_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
+            CREATE_MM(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0].f32acc, matmul_id_subgroup_q2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
             CREATE_MM(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f32acc, matmul_id_subgroup_q4_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
             CREATE_MM(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f32acc, matmul_id_subgroup_q4_1_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
             CREATE_MM(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f32acc, matmul_id_subgroup_q5_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size);
@@ -4822,6 +4898,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
             CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
 
             CREATE_MM(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0].f32acc, matmul_id_q1_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
+            CREATE_MM(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0].f32acc, matmul_id_q2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
             CREATE_MM(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f32acc, matmul_id_q4_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
             CREATE_MM(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f32acc, matmul_id_q4_1_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
             CREATE_MM(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f32acc, matmul_id_q5_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0);
@@ -4924,6 +5001,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_F16 ][i], "mul_mat_vec_f16_f32_f32",  arr_dmmv_f16_f32_f32_len[reduc],  arr_dmmv_f16_f32_f32_data[reduc],  "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, i+1}, 1, false, use_subgroups, force_subgroup_size);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_BF16][i], "mul_mat_vec_bf16_f32_f32", arr_dmmv_bf16_f32_f32_len[reduc], arr_dmmv_bf16_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, i+1}, 1, false, use_subgroups, force_subgroup_size);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q1_0][i], "mul_mat_vec_q1_0_f32_f32", arr_dmmv_q1_0_f32_f32_len[reduc], arr_dmmv_q1_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q2_0][i], "mul_mat_vec_q2_0_f32_f32", arr_dmmv_q2_0_f32_f32_len[reduc], arr_dmmv_q2_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f32_f32", arr_dmmv_q4_0_f32_f32_len[reduc], arr_dmmv_q4_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f32_f32", arr_dmmv_q4_1_f32_f32_len[reduc], arr_dmmv_q4_1_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f32_f32", arr_dmmv_q5_0_f32_f32_len[reduc], arr_dmmv_q5_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
@@ -4950,6 +5028,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_F16 ][i], "mul_mat_vec_f16_f16_f32",  arr_dmmv_f16_f16_f32_len[reduc],  arr_dmmv_f16_f16_f32_data[reduc],  "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, i+1}, 1, false, use_subgroups, force_subgroup_size);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_BF16][i], "mul_mat_vec_bf16_f16_f32", arr_dmmv_bf16_f16_f32_len[reduc], arr_dmmv_bf16_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, i+1}, 1, false, use_subgroups, force_subgroup_size);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q1_0][i], "mul_mat_vec_q1_0_f16_f32", arr_dmmv_q1_0_f16_f32_len[reduc], arr_dmmv_q1_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q2_0][i], "mul_mat_vec_q2_0_f16_f32", arr_dmmv_q2_0_f16_f32_len[reduc], arr_dmmv_q2_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f16_f32", arr_dmmv_q4_0_f16_f32_len[reduc], arr_dmmv_q4_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f16_f32", arr_dmmv_q4_1_f16_f32_len[reduc], arr_dmmv_q4_1_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_f16_f32", arr_dmmv_q5_0_f16_f32_len[reduc], arr_dmmv_q5_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size);
@@ -4977,6 +5056,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
                 const uint32_t subgroup_size_int = (device->vendor_id == VK_VENDOR_ID_INTEL && device->subgroup_size_control) ? device->subgroup_min_size : device->subgroup_size;
                 const uint32_t wg_size_subgroup_int = (w == DMMV_WG_SIZE_SUBGROUP) ? subgroup_size_int : (subgroup_size_int * 4);
 
+                ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q2_0][i], "mul_mat_vec_q2_0_q8_1_f32", arr_dmmv_q2_0_q8_1_f32_len[reduc], arr_dmmv_q2_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 2*rm_kq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
                 ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_q8_1_f32", arr_dmmv_q4_0_q8_1_f32_len[reduc], arr_dmmv_q4_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
                 ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_q8_1_f32", arr_dmmv_q4_1_q8_1_f32_len[reduc], arr_dmmv_q4_1_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
                 ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q5_0][i], "mul_mat_vec_q5_0_q8_1_f32", arr_dmmv_q5_0_q8_1_f32_len[reduc], arr_dmmv_q5_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int, i+1}, 1, true, use_subgroups, subgroup_size_int);
@@ -5002,6 +5082,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_F16 ], "mul_mat_vec_id_f16_f32",        arr_dmmv_id_f16_f32_f32_len[reduc],     arr_dmmv_id_f16_f32_f32_data[reduc],     "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {wg_size_subgroup, 2}, 1, false, use_subgroups, force_subgroup_size);
         ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_BF16], "mul_mat_vec_id_bf16_f32",       arr_dmmv_id_bf16_f32_f32_len[reduc],    arr_dmmv_id_bf16_f32_f32_data[reduc],    "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {wg_size_subgroup, 2}, 1, false, use_subgroups, force_subgroup_size);
         ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q1_0], "mul_mat_vec_id_q1_0_f32",       arr_dmmv_id_q1_0_f32_f32_len[reduc],    arr_dmmv_id_q1_0_f32_f32_data[reduc],    "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size);
+        ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q2_0], "mul_mat_vec_id_q2_0_f32",       arr_dmmv_id_q2_0_f32_f32_len[reduc],    arr_dmmv_id_q2_0_f32_f32_data[reduc],    "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size);
         ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q4_0], "mul_mat_vec_id_q4_0_f32",       arr_dmmv_id_q4_0_f32_f32_len[reduc],    arr_dmmv_id_q4_0_f32_f32_data[reduc],    "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size);
         ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q4_1], "mul_mat_vec_id_q4_1_f32",       arr_dmmv_id_q4_1_f32_f32_len[reduc],    arr_dmmv_id_q4_1_f32_f32_data[reduc],    "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size);
         ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q5_0], "mul_mat_vec_id_q5_0_f32",       arr_dmmv_id_q5_0_f32_f32_len[reduc],    arr_dmmv_id_q5_0_f32_f32_data[reduc],    "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size);
@@ -5029,6 +5110,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
             const uint32_t subgroup_size_int = (device->vendor_id == VK_VENDOR_ID_INTEL && device->subgroup_size_control) ? device->subgroup_min_size : device->subgroup_size;
             const uint32_t wg_size_subgroup_int = (w == DMMV_WG_SIZE_SUBGROUP) ? subgroup_size_int : (subgroup_size_int * 4);
 
+            ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q2_0], "mul_mat_vec_id_q2_0_q8_1_f32", arr_dmmv_id_q2_0_q8_1_f32_len[reduc], arr_dmmv_id_q2_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_kq_int, 1, 1}, {wg_size_subgroup_int, 2*rm_kq_int}, 1, true, use_subgroups, subgroup_size_int);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q4_0], "mul_mat_vec_id_q4_0_q8_1_f32", arr_dmmv_id_q4_0_q8_1_f32_len[reduc], arr_dmmv_id_q4_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int}, 1, true, use_subgroups, subgroup_size_int);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q4_1], "mul_mat_vec_id_q4_1_q8_1_f32", arr_dmmv_id_q4_1_q8_1_f32_len[reduc], arr_dmmv_id_q4_1_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int}, 1, true, use_subgroups, subgroup_size_int);
             ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q5_0], "mul_mat_vec_id_q5_0_q8_1_f32", arr_dmmv_id_q5_0_q8_1_f32_len[reduc], arr_dmmv_id_q5_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq_int, 1, 1}, {wg_size_subgroup_int, 1*rm_stdq_int}, 1, true, use_subgroups, subgroup_size_int);
@@ -5061,6 +5143,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     // dequant shaders
     ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_F32 ], "f32_to_f16",   dequant_f32_len,  dequant_f32_data,  "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q1_0], "dequant_q1_0", dequant_q1_0_len, dequant_q1_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 8, 1, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_0], "dequant_q2_0", dequant_q2_0_len, dequant_q2_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_0], "dequant_q4_0", dequant_q4_0_len, dequant_q4_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_1], "dequant_q4_1", dequant_q4_1_len, dequant_q4_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_0], "dequant_q5_0", dequant_q5_0_len, dequant_q5_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1);
@@ -5088,6 +5171,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_F16 ], "get_rows_f16",  get_rows_f16_len,  get_rows_f16_data,  "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_BF16], "get_rows_bf16", get_rows_bf16_len, get_rows_bf16_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q1_0], "get_rows_q1_0", get_rows_q1_0_len, get_rows_q1_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q2_0], "get_rows_q2_0", get_rows_q2_0_len, get_rows_q2_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_0], "get_rows_q4_0", get_rows_q4_0_len, get_rows_q4_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_1], "get_rows_q4_1", get_rows_q4_1_len, get_rows_q4_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_0], "get_rows_q5_0", get_rows_q5_0_len, get_rows_q5_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
@@ -5115,6 +5199,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_F16 ], "get_rows_f16_f32",  get_rows_f16_f32_len,  get_rows_f16_f32_data,  "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_BF16], "get_rows_bf16_f32", get_rows_bf16_f32_len, get_rows_bf16_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), { 512, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q1_0], "get_rows_q1_0_f32", get_rows_q1_0_f32_len, get_rows_q1_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q2_0], "get_rows_q2_0_f32", get_rows_q2_0_f32_len, get_rows_q2_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_0], "get_rows_q4_0_f32", get_rows_q4_0_f32_len, get_rows_q4_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_1], "get_rows_q4_1_f32", get_rows_q4_1_f32_len, get_rows_q4_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_0], "get_rows_q5_0_f32", get_rows_q5_0_f32_len, get_rows_q5_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1);
@@ -5199,6 +5284,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_16, "cpy_transpose_16", cpy_transpose_16_len, cpy_transpose_16_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1);
 
     ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q1_0], "cpy_f32_q1_0", cpy_f32_q1_0_len, cpy_f32_q1_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q2_0], "cpy_f32_q2_0", cpy_f32_q2_0_len, cpy_f32_q2_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q4_0], "cpy_f32_q4_0", cpy_f32_q4_0_len, cpy_f32_q4_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q4_1], "cpy_f32_q4_1", cpy_f32_q4_1_len, cpy_f32_q4_1_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q5_0], "cpy_f32_q5_0", cpy_f32_q5_0_len, cpy_f32_q5_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
@@ -5211,6 +5297,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_F16],  "set_rows_" #src "_f16" #itype,  set_rows_ ## src ## _f16 ## itype ## _len,  set_rows_ ## src ## _f16 ## itype ## _data,  "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
         ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_BF16], "set_rows_" #src "_bf16" #itype, set_rows_ ## src ## _bf16 ## itype ## _len, set_rows_ ## src ## _bf16 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
         ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q1_0], "set_rows_" #src "_q1_0" #itype, set_rows_ ## src ## _q1_0 ## itype ## _len, set_rows_ ## src ## _q1_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q2_0], "set_rows_" #src "_q2_0" #itype, set_rows_ ## src ## _q2_0 ## itype ## _len, set_rows_ ## src ## _q2_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
         ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q4_0], "set_rows_" #src "_q4_0" #itype, set_rows_ ## src ## _q4_0 ## itype ## _len, set_rows_ ## src ## _q4_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
         ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q4_1], "set_rows_" #src "_q4_1" #itype, set_rows_ ## src ## _q4_1 ## itype ## _len, set_rows_ ## src ## _q4_1 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
         ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q5_0], "set_rows_" #src "_q5_0" #itype, set_rows_ ## src ## _q5_0 ## itype ## _len, set_rows_ ## src ## _q5_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
@@ -5226,6 +5313,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 
 
     ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q1_0], "cpy_q1_0_f32", cpy_q1_0_f32_len, cpy_q1_0_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q1_0), 1, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q2_0], "cpy_q2_0_f32", cpy_q2_0_f32_len, cpy_q2_0_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q2_0), 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q4_0], "cpy_q4_0_f32", cpy_q4_0_f32_len, cpy_q4_0_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q4_0), 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q4_1], "cpy_q4_1_f32", cpy_q4_1_f32_len, cpy_q4_1_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q4_1), 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q5_0], "cpy_q5_0_f32", cpy_q5_0_f32_len, cpy_q5_0_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q5_0), 1, 1}, {}, 1);
@@ -5860,6 +5948,7 @@ static vk_device ggml_vk_get_device(size_t idx) {
         bool coopmat2_support = false;
         bool coopmat2_decode_vector_support = false;
         bool pipeline_executable_properties_support = false;
+        bool internally_sync_support = false;
         device->coopmat_support = false;
         device->integer_dot_product = false;
         device->shader_64b_indexing = false;
@@ -5931,6 +6020,8 @@ static vk_device ggml_vk_get_device(size_t idx) {
             } else if (strcmp("VK_EXT_shader_64bit_indexing", properties.extensionName) == 0) {
                 device->shader_64b_indexing = true;
 #endif
+            } else if (strcmp(VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME, properties.extensionName) == 0) {
+                internally_sync_support = true;
             }
         }
 
@@ -6117,14 +6208,6 @@ static vk_device ggml_vk_get_device(size_t idx) {
         device->single_queue = compute_queue_family_index == transfer_queue_family_index && queue_family_props[compute_queue_family_index].queueCount == 1;
 
         std::vector<vk::DeviceQueueCreateInfo> device_queue_create_infos;
-        if (compute_queue_family_index != transfer_queue_family_index) {
-            device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities});
-            device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), transfer_queue_family_index, 1, priorities + 1});
-        } else if(!device->single_queue) {
-            device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 2, priorities});
-        } else {
-            device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities});
-        }
         vk::DeviceCreateInfo device_create_info{};
         std::vector<const char *> device_extensions;
         vk::PhysicalDeviceFeatures device_features = device->physical_device.getFeatures();
@@ -6146,6 +6229,17 @@ static vk_device ggml_vk_get_device(size_t idx) {
 
         last_struct = (VkBaseOutStructure *)&vk12_features;
 
+        VkPhysicalDeviceInternallySynchronizedQueuesFeaturesKHR internally_synchronized_queues_features{};
+        internally_synchronized_queues_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_INTERNALLY_SYNCHRONIZED_QUEUES_FEATURES_KHR;
+        internally_synchronized_queues_features.pNext = nullptr;
+        internally_synchronized_queues_features.internallySynchronizedQueues = VK_FALSE;
+
+        if (internally_sync_support) {
+            last_struct->pNext = (VkBaseOutStructure *)&internally_synchronized_queues_features;
+            last_struct = (VkBaseOutStructure *)&internally_synchronized_queues_features;
+            device_extensions.push_back(VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME);
+        }
+
         VkPhysicalDevicePipelineRobustnessFeaturesEXT pl_robustness_features;
         pl_robustness_features.pNext = nullptr;
         pl_robustness_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PIPELINE_ROBUSTNESS_FEATURES_EXT;
@@ -6284,6 +6378,23 @@ static vk_device ggml_vk_get_device(size_t idx) {
 
         vkGetPhysicalDeviceFeatures2(device->physical_device, &device_features2);
 
+        device->has_internally_synchronized_queues = internally_synchronized_queues_features.internallySynchronizedQueues;
+
+        // Build queue create infos only after querying whether internally synchronized queues are enabled.
+        // getQueue2() later uses the same flag, so creation/retrieval must stay consistent.
+        vk::DeviceQueueCreateFlags queue_flags = device->has_internally_synchronized_queues ?
+                                                eInternallySynchronizedKHR :
+                                                vk::DeviceQueueCreateFlags();
+
+        if (compute_queue_family_index != transfer_queue_family_index) {
+            device_queue_create_infos.push_back({queue_flags, compute_queue_family_index, 1, priorities});
+            device_queue_create_infos.push_back({queue_flags, transfer_queue_family_index, 1, priorities + 1});
+        } else if(!device->single_queue) {
+            device_queue_create_infos.push_back({queue_flags, compute_queue_family_index, 2, priorities});
+        } else {
+            device_queue_create_infos.push_back({queue_flags, compute_queue_family_index, 1, priorities});
+        }
+
         device->pipeline_executable_properties_support = pipeline_executable_properties_support;
 
         device->fp16 = device->fp16 && vk12_features.shaderFloat16;
@@ -6566,7 +6677,7 @@ static vk_device ggml_vk_get_device(size_t idx) {
         device->device = device->physical_device.createDevice(device_create_info);
 
         // Queues
-        ggml_vk_create_queue(device, device->compute_queue, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false);
+        device->compute_queue = ggml_vk_create_queue(device, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false);
 
         // Shaders
         // Disable matmul tile sizes early if performance low or not supported
@@ -6668,13 +6779,11 @@ static vk_device ggml_vk_get_device(size_t idx) {
 
         if (!device->single_queue) {
             const uint32_t transfer_queue_index = compute_queue_family_index == transfer_queue_family_index ? 1 : 0;
-            ggml_vk_create_queue(device, device->transfer_queue, transfer_queue_family_index, transfer_queue_index, { vk::PipelineStageFlagBits::eTransfer }, true);
+            device->transfer_queue = ggml_vk_create_queue(device, transfer_queue_family_index, transfer_queue_index, { vk::PipelineStageFlagBits::eTransfer }, true);
 
             device->async_use_transfer_queue = prefers_transfer_queue || (getenv("GGML_VK_ASYNC_USE_TRANSFER_QUEUE") != nullptr);
         } else {
-            // TODO: Use pointer or reference to avoid copy
-            device->transfer_queue.copyFrom(device->compute_queue);
-            device->transfer_queue.cmd_pool.init(device, &device->transfer_queue);
+            device->transfer_queue = ggml_vk_create_aliased_queue(device, device->compute_queue);
 
             device->async_use_transfer_queue = false;
         }
@@ -7237,7 +7346,7 @@ static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) {
     ctx->fence = ctx->device->device.createFence({});
     ctx->almost_ready_fence = ctx->device->device.createFence({});
 
-    ctx->compute_cmd_pool.init(ctx->device, &ctx->device->compute_queue);
+    ctx->compute_cmd_pool.init(ctx->device, ctx->device->compute_queue.get());
     if (ctx->device->async_use_transfer_queue) {
         vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eTimeline, 0 };
         vk::SemaphoreCreateInfo ci{};
@@ -7245,7 +7354,7 @@ static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) {
         ctx->transfer_semaphore.s = ctx->device->device.createSemaphore(ci);
         ctx->transfer_semaphore.value = 0;
 
-        ctx->transfer_cmd_pool.init(ctx->device, &ctx->device->transfer_queue);
+        ctx->transfer_cmd_pool.init(ctx->device, ctx->device->transfer_queue.get());
     }
 
     if (vk_perf_logger_enabled) {
@@ -7265,6 +7374,7 @@ static vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type
     switch (type) {
         case GGML_TYPE_F32:
         case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
         case GGML_TYPE_Q5_0:
@@ -7338,6 +7448,7 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_conte
 
     switch (src0_type) {
         case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
         case GGML_TYPE_Q5_0:
@@ -7381,6 +7492,7 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context *
 
     if (b_type == GGML_TYPE_Q8_1) {
         switch (a_type) {
+            case GGML_TYPE_Q2_0:
             case GGML_TYPE_Q4_0:
             case GGML_TYPE_Q4_1:
             case GGML_TYPE_Q5_0:
@@ -7405,6 +7517,7 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context *
         case GGML_TYPE_F16:
         case GGML_TYPE_BF16:
         case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
         case GGML_TYPE_Q5_0:
@@ -7497,6 +7610,7 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_co
 
     switch (src0_type) {
         case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
         case GGML_TYPE_Q5_0:
@@ -7543,6 +7657,7 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec_id(ggml_backend_vk_context
 
     if (b_type == GGML_TYPE_Q8_1) {
         switch (a_type) {
+            case GGML_TYPE_Q2_0:
             case GGML_TYPE_Q4_0:
             case GGML_TYPE_Q4_1:
             case GGML_TYPE_Q5_0:
@@ -7567,6 +7682,7 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec_id(ggml_backend_vk_context
         case GGML_TYPE_F16:
         case GGML_TYPE_BF16:
         case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
         case GGML_TYPE_Q5_0:
@@ -7829,6 +7945,20 @@ static vk_context ggml_vk_get_compute_ctx(ggml_backend_vk_context * ctx) {
     return result;
 }
 
+static vk_context ggml_vk_get_transfer_ctx(ggml_backend_vk_context * ctx) {
+    vk_context result;
+    if (!ctx->transfer_ctx.expired()) {
+        result = ctx->transfer_ctx.lock();
+    } else {
+        result = ggml_vk_create_context(ctx, ctx->transfer_cmd_pool);
+
+        ctx->transfer_ctx = result;
+        ggml_vk_ctx_begin(ctx->device, result);
+    }
+
+    return result;
+}
+
 // Submit any pending transfer queue work and signal the transfer semaphore.
 // The next compute context created via ggml_vk_get_compute_ctx will wait on this semaphore.
 // Returns true if work was submitted.
@@ -8079,7 +8209,7 @@ static void ggml_vk_buffer_write_2d(vk_buffer& dst, size_t offset, const void *
     } else {
         std::lock_guard<std::recursive_mutex> guard(dst->device->mutex);
 
-        vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue.cmd_pool);
+        vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue->cmd_pool);
         ggml_vk_ctx_begin(dst->device, subctx);
         bool ret = ggml_vk_buffer_write_2d_async(subctx, dst, offset, src, spitch, dpitch, width, height, true);
         GGML_ASSERT(ret);
@@ -8194,7 +8324,7 @@ static void ggml_vk_buffer_read_2d(vk_buffer& src, size_t offset, void * dst, si
         GGML_ASSERT(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent);
 
         std::lock_guard<std::recursive_mutex> guard(src->device->mutex);
-        vk_context subctx = ggml_vk_create_temporary_context(src->device->compute_queue.cmd_pool);
+        vk_context subctx = ggml_vk_create_temporary_context(src->device->compute_queue->cmd_pool);
         ggml_vk_ctx_begin(src->device, subctx);
         subctx->s->buffer->buf.pipelineBarrier(
             vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer,
@@ -8220,7 +8350,7 @@ static void ggml_vk_buffer_read_2d(vk_buffer& src, size_t offset, void * dst, si
     } else {
         std::lock_guard<std::recursive_mutex> guard(src->device->mutex);
 
-        vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue.cmd_pool);
+        vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue->cmd_pool);
         ggml_vk_ctx_begin(src->device, subctx);
         bool ret = ggml_vk_buffer_read_2d_async(subctx, src, offset, dst, spitch, dpitch, width, height, true);
         GGML_ASSERT(ret);
@@ -8257,7 +8387,7 @@ static void ggml_vk_buffer_copy(vk_buffer& dst, size_t dst_offset, vk_buffer& sr
         std::lock_guard<std::recursive_mutex> guard(src->device->mutex);
         VK_LOG_DEBUG("ggml_vk_buffer_copy(SINGLE_DEVICE, " << size << ")");
         // Copy within the device
-        vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue.cmd_pool);
+        vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue->cmd_pool);
         ggml_vk_ctx_begin(src->device, subctx);
         ggml_vk_buffer_copy_async(subctx, dst, dst_offset, src, src_offset, size);
         ggml_vk_ctx_end(subctx);
@@ -8300,7 +8430,7 @@ static void ggml_vk_buffer_memset(vk_buffer& dst, size_t offset, uint32_t c, siz
     }
 
     std::lock_guard<std::recursive_mutex> guard(dst->device->mutex);
-    vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue.cmd_pool);
+    vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue->cmd_pool);
     ggml_vk_ctx_begin(dst->device, subctx);
     subctx->s->buffer->buf.fillBuffer(dst->buffer, offset, size, c);
     ggml_vk_ctx_end(subctx);
@@ -8589,6 +8719,7 @@ static vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const
     if (src->type == GGML_TYPE_F32) {
         switch (to) {
         case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
         case GGML_TYPE_Q5_0:
@@ -8604,6 +8735,7 @@ static vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const
     if (to == GGML_TYPE_F32) {
         switch (src->type) {
         case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q2_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
         case GGML_TYPE_Q5_0:
@@ -9030,7 +9162,7 @@ static bool ggml_vk_should_use_mmvq(const vk_device& device, uint32_t m, uint32_
     // Quantization overhead is not worth it for small k
     switch (device->vendor_id) {
     case VK_VENDOR_ID_NVIDIA:
-        if (src0_type == GGML_TYPE_Q2_K || src0_type == GGML_TYPE_Q3_K || src0_type == GGML_TYPE_IQ1_S || src0_type == GGML_TYPE_IQ1_M) {
+        if (src0_type == GGML_TYPE_Q2_0 || src0_type == GGML_TYPE_Q2_K || src0_type == GGML_TYPE_Q3_K || src0_type == GGML_TYPE_IQ1_S || src0_type == GGML_TYPE_IQ1_M) {
             return true;
         }
 
@@ -9058,7 +9190,7 @@ static bool ggml_vk_should_use_mmvq(const vk_device& device, uint32_t m, uint32_
         }
     case VK_VENDOR_ID_INTEL:
         if (device->architecture == vk_device_architecture::INTEL_XE2) {
-            if (src0_type == GGML_TYPE_Q2_K || src0_type == GGML_TYPE_Q3_K || src0_type == GGML_TYPE_Q6_K) {
+            if (src0_type == GGML_TYPE_Q2_0 || src0_type == GGML_TYPE_Q2_K || src0_type == GGML_TYPE_Q3_K || src0_type == GGML_TYPE_Q6_K) {
                 return true;
             }
         }
@@ -15633,13 +15765,7 @@ static void ggml_backend_vk_set_tensor_2d_async(ggml_backend_t backend, ggml_ten
     vk_context cpy_ctx;
 
     if (ctx->device->async_use_transfer_queue) {
-        if (ctx->transfer_ctx.expired()) {
-            cpy_ctx = ggml_vk_create_context(ctx, ctx->transfer_cmd_pool);
-            ctx->transfer_ctx = cpy_ctx;
-            ggml_vk_ctx_begin(ctx->device, cpy_ctx);
-        } else {
-            cpy_ctx = ctx->transfer_ctx.lock();
-        }
+        cpy_ctx = ggml_vk_get_transfer_ctx(ctx);
     } else {
         cpy_ctx = ggml_vk_get_compute_ctx(ctx);
     }
@@ -15785,13 +15911,7 @@ static bool ggml_backend_vk_cpy_tensor_async(ggml_backend_t backend_src, ggml_ba
 
         vk_context cpy_ctx;
         if (ctx->device->async_use_transfer_queue) {
-            if (ctx->transfer_ctx.expired()) {
-                cpy_ctx = ggml_vk_create_context(ctx, ctx->transfer_cmd_pool);
-                ctx->transfer_ctx = cpy_ctx;
-                ggml_vk_ctx_begin(ctx->device, cpy_ctx);
-            } else {
-                cpy_ctx = ctx->transfer_ctx.lock();
-            }
+            cpy_ctx = ggml_vk_get_transfer_ctx(ctx);
         } else {
             cpy_ctx = ggml_vk_get_compute_ctx(ctx);
         }
@@ -15838,19 +15958,17 @@ static void ggml_vk_synchronize(ggml_backend_vk_context * ctx) {
                 1, &ctx->transfer_semaphore.value,
                 0, nullptr,
             };
-            vk::PipelineStageFlags stage = ctx->device->transfer_queue.stage_flags;
+            vk::PipelineStageFlags stage = ctx->device->transfer_queue->stage_flags;
             vk::SubmitInfo si{
                 1, &ctx->transfer_semaphore.s, &stage,
                 0, nullptr,
                 0, nullptr,
             };
             si.setPNext(&tl_info);
-            std::lock_guard<std::mutex> guard(queue_mutex);
-            ctx->device->compute_queue.queue.submit({ si }, ctx->fence);
+            ctx->device->compute_queue->handle->submit({ si }, ctx->fence);
             ctx->transfer_semaphore_last_submitted = ctx->transfer_semaphore.value;
         } else {
-            std::lock_guard<std::mutex> guard(queue_mutex);
-            ctx->device->compute_queue.queue.submit({}, ctx->fence);
+            ctx->device->compute_queue->handle->submit({}, ctx->fence);
         }
         ggml_vk_wait_for_fence(ctx);
         ctx->submit_pending = false;
@@ -16412,7 +16530,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
         vk::DebugUtilsLabelEXT dul = {};
         dul.pLabelName = "ggml_backend_vk_graph_compute";
         dul.color = std::array<float,4>{1.0f, 1.0f, 1.0f, 1.0f};
-        vk_instance.pfn_vkQueueBeginDebugUtilsLabelEXT(ctx->device->compute_queue.queue, reinterpret_cast<VkDebugUtilsLabelEXT*>(&dul));
+
+        std::lock_guard<vk_queue_handle> guard(*ctx->device->compute_queue->handle);
+        vk_instance.pfn_vkQueueBeginDebugUtilsLabelEXT(ctx->device->compute_queue->handle->queue, reinterpret_cast<VkDebugUtilsLabelEXT*>(&dul));
     }
 
     ctx->prealloc_size_add_rms_partials_offset = 0;
@@ -17123,6 +17243,11 @@ static void ggml_backend_vk_event_wait(ggml_backend_t backend, ggml_backend_even
     if (vkev->has_event) {
         // Wait for latest event
         ggml_vk_wait_events(compute_ctx, { vkev->event });
+
+        if (ctx->device->async_use_transfer_queue) {
+            vk_context transfer_ctx = ggml_vk_get_transfer_ctx(ctx);
+            transfer_ctx->s->wait_semaphores.push_back(vkev->tl_semaphore);
+        }
     }
 }
 
@@ -17422,6 +17547,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
                     case GGML_TYPE_F16:
                     case GGML_TYPE_BF16:
                     case GGML_TYPE_Q1_0:
+                    case GGML_TYPE_Q2_0:
                     case GGML_TYPE_Q4_0:
                     case GGML_TYPE_Q4_1:
                     case GGML_TYPE_Q5_0:
@@ -17527,6 +17653,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
                     case GGML_TYPE_F16:
                     case GGML_TYPE_BF16:
                     case GGML_TYPE_Q1_0:
+                    case GGML_TYPE_Q2_0:
                     case GGML_TYPE_Q4_0:
                     case GGML_TYPE_Q4_1:
                     case GGML_TYPE_Q5_0:
@@ -17567,6 +17694,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
                     case GGML_TYPE_F16:
                     case GGML_TYPE_BF16:
                     case GGML_TYPE_Q1_0:
+                    case GGML_TYPE_Q2_0:
                     case GGML_TYPE_Q4_0:
                     case GGML_TYPE_Q4_1:
                     case GGML_TYPE_Q5_0:
@@ -17591,6 +17719,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
                     case GGML_TYPE_F16:
                     case GGML_TYPE_BF16:
                     case GGML_TYPE_Q1_0:
+                    case GGML_TYPE_Q2_0:
                     case GGML_TYPE_Q4_0:
                     case GGML_TYPE_Q4_1:
                     case GGML_TYPE_Q5_0:
@@ -17607,6 +17736,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
                     case GGML_TYPE_F16:
                     case GGML_TYPE_BF16:
                     case GGML_TYPE_Q1_0:
+                    case GGML_TYPE_Q2_0:
                     case GGML_TYPE_Q4_0:
                     case GGML_TYPE_Q4_1:
                     case GGML_TYPE_Q5_0:

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

@@ -208,6 +208,36 @@ void quantize(uint dst_idx, uint src_idx)
 }
 #endif
 
+#if defined(DATA_A_Q2_0)
+uint quantize_q2_0(float x)
+{
+    const int q = int(x >= 0.0f ? floor(x + 0.5f) : ceil(x - 0.5f)) + 1;
+    return uint(clamp(q, 0, 3));
+}
+
+void quantize(uint dst_idx, uint src_idx)
+{
+    float amax = 0.0f;
+
+    [[unroll]] for (int j = 0; j < QUANT_K_Q2_0; ++j) {
+        amax = max(amax, abs(float(data_s[src_idx + j])));
+    }
+
+    const float d = amax;
+    const float id = d != 0.0f ? 1.0f / d : 0.0f;
+
+    data_q[dst_idx].d = float16_t(d);
+
+    [[unroll]] for (int j = 0; j < QUANT_K_Q2_0 / 4; ++j) {
+        const uint q0 = quantize_q2_0(float(data_s[src_idx + 4*j    ]) * id);
+        const uint q1 = quantize_q2_0(float(data_s[src_idx + 4*j + 1]) * id);
+        const uint q2 = quantize_q2_0(float(data_s[src_idx + 4*j + 2]) * id);
+        const uint q3 = quantize_q2_0(float(data_s[src_idx + 4*j + 3]) * id);
+        data_q[dst_idx].qs[j] = uint8_t(q0 | (q1 << 2u) | (q2 << 4u) | (q3 << 6u));
+    }
+}
+#endif
+
 #if defined(DATA_A_IQ4_NL)
 uint best_index(float x) {
     if (x <= kvalues_iq4nl[0]) return 0;

ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl

@@ -143,6 +143,17 @@ vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
 }
 #endif
 
+#if defined(DATA_A_Q2_0)
+vec2 dequantize(uint ib, uint iqs, uint a_offset) {
+    const uint bits = uint(data_a[a_offset + ib].qs[iqs / 4u]) >> (2u * (iqs % 4u));
+    return vec2(bits & 3u, (bits >> 2u) & 3u) - 1.0f;
+}
+vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
+    const uint bits = uint(data_a[a_offset + ib].qs[iqs / 4u]);
+    return vec4(bits & 3u, (bits >> 2u) & 3u, (bits >> 4u) & 3u, bits >> 6u) - 1.0f;
+}
+#endif
+
 #if defined(DATA_A_IQ1_S)
 vec2 dequantize(uint ib, uint iqs, uint a_offset) {
     const uint ib32 = iqs / 32;
@@ -547,7 +558,7 @@ vec2 get_dm(uint ib, uint a_offset) {
 }
 #endif
 
-#if defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ1_S) || defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S) || defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S) || defined(DATA_A_IQ4_XS) || defined(DATA_A_IQ4_NL)
+#if defined(DATA_A_Q2_0) || defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ1_S) || defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S) || defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S) || defined(DATA_A_IQ4_XS) || defined(DATA_A_IQ4_NL)
 vec2 get_dm(uint ib, uint a_offset) {
     return vec2(float(data_a[a_offset + ib].d), 0);
 }

ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl

@@ -46,6 +46,26 @@ f16vec4 dequantFuncQ1_0_v(const in decodeBufQ1_0 bl, const in uint blockCoords[2
         (qs_nib & 8u) != 0u ? d : md);
 }
 
+layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ2_0 {
+   block_q2_0 block;
+};
+
+float16_t dequantFuncQ2_0(const in decodeBufQ2_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
+{
+    const float16_t d = bl.block.d;
+    const uint idx = coordInBlock[1];
+    const uint bits = uint(bl.block.qs[idx >> 2]) >> (2u * (idx & 3u));
+    return (float16_t(bits & 3u) - float16_t(1.0)) * d;
+}
+
+f16vec4 dequantFuncQ2_0_v(const in decodeBufQ2_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2])
+{
+    const float16_t d = bl.block.d;
+    const uint idx = coordInBlock[1];
+    const uint bits = uint(bl.block.qs[idx >> 2]);
+    return f16vec4((vec4(bits & 3u, (bits >> 2u) & 3u, (bits >> 4u) & 3u, bits >> 6u) - 1.0f) * float(d));
+}
+
 layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufQ4_0 {
    block_q4_0_packed16 block;
 };
@@ -1330,6 +1350,9 @@ f16vec4 dequantFuncNVFP4_v(const in decodeBufNVFP4 bl, const in uint blockCoords
 #if defined(DATA_A_Q1_0)
 #define dequantFuncA dequantFuncQ1_0
 #define dequantFuncA_v dequantFuncQ1_0_v
+#elif defined(DATA_A_Q2_0)
+#define dequantFuncA dequantFuncQ2_0
+#define dequantFuncA_v dequantFuncQ2_0_v
 #elif defined(DATA_A_Q4_0)
 #define dequantFuncA dequantFuncQ4_0
 #define dequantFuncA_v dequantFuncQ4_0_v

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

@@ -0,0 +1,29 @@
+#version 450
+
+#include "dequant_head.glsl"
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {block_q2_0 data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+    const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64;
+
+    const uint tid = gl_LocalInvocationID.x % 64;
+    const uint il  = tid / 4;
+    const uint ir  = tid % 4;
+    const uint ib = 4*i + ir;
+    if (ib >= p.nel / QUANT_K_Q2_0) {
+        return;
+    }
+
+    const uint b_idx = 256*i + QUANT_K_Q2_0*ir + 4*il;
+    const uint bits = uint(data_a[ib].qs[il]);
+    const float d = float(data_a[ib].d);
+
+    data_b[b_idx    ] = D_TYPE(d * (float(bits & 3u) - 1.0f));
+    data_b[b_idx + 1] = D_TYPE(d * (float((bits >> 2u) & 3u) - 1.0f));
+    data_b[b_idx + 2] = D_TYPE(d * (float((bits >> 4u) & 3u) - 1.0f));
+    data_b[b_idx + 3] = D_TYPE(d * (float(bits >> 6u) - 1.0f));
+}

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

@@ -11,10 +11,10 @@
 
 layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
 
-#if defined(DATA_A_QUANT_LEGACY) || defined(DATA_A_MXFP4)
-#define K_PER_ITER 8
-#elif defined(DATA_A_QUANT_K)
+#if defined(DATA_A_Q2_0) || defined(DATA_A_QUANT_K)
 #define K_PER_ITER 16
+#elif defined(DATA_A_QUANT_LEGACY) || defined(DATA_A_MXFP4)
+#define K_PER_ITER 8
 #elif defined(DATA_A_IQ1_S) || defined(DATA_A_IQ1_M)
 #define K_PER_ITER 32
 #else

ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq_funcs.glsl

@@ -4,7 +4,11 @@
 
 #include "types.glsl"
 
-#if defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ1_S) || defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S) || defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S) || defined(DATA_A_IQ4_XS) || defined(DATA_A_IQ4_NL)
+#if defined(DATA_A_Q2_0)
+FLOAT_TYPE get_dm(uint ib) {
+    return FLOAT_TYPE(data_a[ib / 2].d);
+}
+#elif defined(DATA_A_Q4_0) || defined(DATA_A_Q5_0) || defined(DATA_A_Q8_0) || defined(DATA_A_IQ1_S) || defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S) || defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S) || defined(DATA_A_IQ4_XS) || defined(DATA_A_IQ4_NL)
 FLOAT_TYPE get_dm(uint ib) {
     return FLOAT_TYPE(data_a[ib].d);
 }
@@ -30,6 +34,27 @@ FLOAT_TYPEV2 get_dm(uint ib) {
 #endif
 
 // Each iqs value maps to a 32-bit integer
+#if defined(DATA_A_Q2_0)
+uint unpack_q2_0(uint bits) {
+    // Move bit pairs [1:0], [3:2], [5:4], [7:6] to [1:0], [9:8], [17:16], [25:24].
+    bits &= 0xffu;
+    bits = (bits | (bits << 12u)) & 0x000f000fu;
+    return (bits | (bits << 6u)) & 0x03030303u;
+}
+
+i32vec4 repack4(uint ib, uint iqs) {
+    const uint qs_idx = (ib & 1u) * 4u + iqs * 2u;
+    const uint bits = pack32(u16vec2(data_a_packed16[ib / 2].qs[qs_idx],
+                                     data_a_packed16[ib / 2].qs[qs_idx + 1]));
+    return i32vec4(unpack_q2_0(bits), unpack_q2_0(bits >> 8u),
+                   unpack_q2_0(bits >> 16u), unpack_q2_0(bits >> 24u));
+}
+
+FLOAT_TYPE mul_q8_1(const int32_t q_sum, const float da, const vec2 dsb, const int32_t sum_divisor) {
+    return FLOAT_TYPE(da * (float(q_sum) * dsb.x - dsb.y / float(sum_divisor)));
+}
+#endif
+
 #if defined(DATA_A_Q4_0)
 // 2-byte loads for Q4_0 blocks (18 bytes)
 i32vec2 repack(uint ib, uint iqs) {
@@ -132,7 +157,19 @@ FLOAT_TYPE mul_q8_1(const int32_t q_sum, const float da, const vec2 dsb, const i
 }
 #endif
 
-#if defined(DATA_A_QUANT_LEGACY) || defined(DATA_A_MXFP4)
+#if defined(DATA_A_Q2_0)
+FLOAT_TYPE mmvq_dot_product(const uint ib_a, const uint iqs) {
+    int32_t q_sum = 0;
+    const i32vec4 qs_a = repack4(ib_a, iqs);
+    q_sum += dotPacked4x8EXT(qs_a.x, cache_b_qs[0]);
+    q_sum += dotPacked4x8EXT(qs_a.y, cache_b_qs[1]);
+    q_sum += dotPacked4x8EXT(qs_a.z, cache_b_qs[2]);
+    q_sum += dotPacked4x8EXT(qs_a.w, cache_b_qs[3]);
+
+    // 16 quants per call => divide sums by 32/16 = 2
+    return mul_q8_1(q_sum, get_dm(ib_a), cache_b_ds, 2);
+}
+#elif defined(DATA_A_QUANT_LEGACY) || defined(DATA_A_MXFP4)
 FLOAT_TYPE mmvq_dot_product(const uint ib_a, const uint iqs) {
     int32_t q_sum = 0;
 #if QUANT_R == 2

ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl

@@ -151,6 +151,18 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
             buf_a[buf_idx + 1] = FLOAT_TYPEV2((bits & 0x04u) != 0u ? d : -d, (bits & 0x08u) != 0u ? d : -d);
             buf_a[buf_idx + 2] = FLOAT_TYPEV2((bits & 0x10u) != 0u ? d : -d, (bits & 0x20u) != 0u ? d : -d);
             buf_a[buf_idx + 3] = FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d);
+#elif defined(DATA_A_Q2_0)
+            const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
+            const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
+
+            const uint ib = idx / 16;
+            const uint iqs = idx & 0xfu;
+
+            const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib].d);
+            const uint bits = uint(data_a[ib].qs[iqs]);
+
+            buf_a[buf_idx    ] = d * (FLOAT_TYPEV2(bits & 3u, (bits >> 2u) & 3u) - FLOAT_TYPEV2(1.0f));
+            buf_a[buf_idx + 1] = d * (FLOAT_TYPEV2((bits >> 4u) & 3u, bits >> 6u) - FLOAT_TYPEV2(1.0f));
 #elif defined(DATA_A_Q2_K)
             const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
             const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;

ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl

@@ -6,6 +6,40 @@
 
 // Each iqs value maps to a 32-bit integer
 
+#if defined(DATA_A_Q2_0)
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
+    const uint block_idx = ib / 2;
+    const uint byte_idx = (ib & 1u) * 8u + iqs;
+    const uint bits = uint(data_a[block_idx].qs[byte_idx]);
+    buf_a[buf_ib].qs[iqs] = pack32(i8vec4(
+        int8_t(bits & 3u),
+        int8_t((bits >> 2u) & 3u),
+        int8_t((bits >> 4u) & 3u),
+        int8_t(bits >> 6u)));
+
+    if (iqs == 0) {
+        buf_a[buf_ib].dm = FLOAT_TYPE(data_a[block_idx].d);
+    }
+}
+
+void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
+    cache_a[reg_ib].dm = buf_a[buf_ib].dm;
+
+    [[unroll]] for (uint iqs = 0; iqs < 8; ++iqs) {
+        cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs];
+    }
+}
+
+ACC_TYPE mmq_dot_product(const uint ib_a) {
+    int32_t q_sum = 0;
+    [[unroll]] for (uint iqs = 0; iqs < 8; ++iqs) {
+        q_sum += dotPacked4x8EXT(cache_a[ib_a].qs[iqs], cache_b.qs[iqs]);
+    }
+
+    return ACC_TYPE(float(cache_a[ib_a].dm) * (float(q_sum) * float(cache_b.ds.x) - float(cache_b.ds.y)));
+}
+#endif
+
 #if defined(DATA_A_Q4_0) || defined(DATA_A_Q4_1)
 // 2-byte loads for Q4_0 blocks (18 bytes)
 // 4-byte loads for Q4_1 blocks (20 bytes)

ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_shmem_types.glsl

@@ -13,6 +13,12 @@ struct block_a_cache {
     uint32_t qs[16/4];
     FLOAT_TYPE dm;
 };
+#elif defined(DATA_A_Q2_0)
+#define QUANT_R_MMQ 1
+struct block_a_cache {
+    int32_t qs[8];
+    FLOAT_TYPE dm;
+};
 #elif defined(DATA_A_Q4_1)
 #define QUANT_R_MMQ 2
 struct block_a_cache {

ggml/src/ggml-vulkan/vulkan-shaders/types.glsl

@@ -211,6 +211,30 @@ struct block_q1_0
 #define A_TYPE block_q1_0
 #endif
 
+#define QUANT_K_Q2_0 64
+#define QUANT_R_Q2_0 1
+
+struct block_q2_0
+{
+    float16_t d;
+    uint8_t qs[QUANT_K_Q2_0 / 4];
+};
+
+struct block_q2_0_packed16
+{
+    float16_t d;
+    uint16_t qs[QUANT_K_Q2_0 / 8];
+};
+
+#if defined(DATA_A_Q2_0)
+#define QUANT_K QUANT_K_Q2_0
+#define QUANT_R QUANT_R_Q2_0
+#define QUANT_AUXF 1
+#define A_TYPE block_q2_0
+#define A_TYPE_PACKED16 block_q2_0_packed16
+#define DATA_A_QUANT_LEGACY
+#endif
+
 #define QUANT_K_Q8_1 32
 #define QUANT_R_Q8_1 1
 

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

@@ -50,6 +50,7 @@ const std::vector<std::string> type_names = {
     "f32",
     "f16",
     "q1_0",
+    "q2_0",
     "q4_0",
     "q4_1",
     "q5_0",
@@ -232,7 +233,7 @@ bool is_quantized_type(const std::string& type_name) {
 }
 
 bool is_legacy_quant(const std::string& type_name) {
-    return type_name == "q4_0" || type_name == "q4_1" || type_name == "q5_0" || type_name == "q5_1" || type_name == "q8_0";
+    return type_name == "q2_0" || type_name == "q4_0" || type_name == "q4_1" || type_name == "q5_0" || type_name == "q5_1" || type_name == "q8_0";
 }
 
 bool is_k_quant(const std::string& type_name) {
@@ -583,7 +584,7 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c
         std::string load_vec_quant = "2";
         if ((tname == "q1_0") || (tname == "q4_0") || (tname == "q4_1") || (tname == "q5_1") || (tname == "iq1_s") || (tname == "iq1_m") || (tname == "iq2_xxs") || (tname == "iq2_xs") || (tname == "iq2_s"))
             load_vec_quant = "8";
-        else if ((tname == "q5_0") || (tname == "q8_0") || (tname == "q2_k") || (tname == "q4_k") || (tname == "q5_k") || (tname == "iq3_xxs") || (tname == "iq3_s") || (tname == "iq4_xs") || (tname == "iq4_nl") || (tname == "mxfp4") || (tname == "nvfp4"))
+        else if ((tname == "q2_0") || (tname == "q5_0") || (tname == "q8_0") || (tname == "q2_k") || (tname == "q4_k") || (tname == "q5_k") || (tname == "iq3_xxs") || (tname == "iq3_s") || (tname == "iq4_xs") || (tname == "iq4_nl") || (tname == "mxfp4") || (tname == "nvfp4"))
             load_vec_quant = "4";
 
         if (tname == "bf16") {
@@ -823,13 +824,13 @@ void process_shaders() {
     string_to_spv("cpy_transpose_16", "copy_transpose.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}});
     string_to_spv("cpy_transpose_32", "copy_transpose.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}});
 
-    for (std::string t : {"q1_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
+    for (std::string t : {"q1_0", "q2_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
         string_to_spv("cpy_f32_" + t, "copy_to_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"S_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
         string_to_spv("cpy_" + t + "_f32", "copy_from_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
     }
 
     for (auto src : {std::pair{"f32", "float"}, std::pair{"f16", "float16_t"}}) {
-        for (std::string dst : {"f32", "f16", "bf16", "q1_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
+        for (std::string dst : {"f32", "f16", "bf16", "q1_0", "q2_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
             string_to_spv("set_rows_" + std::string(src.first) + "_" + dst + "_i32", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(dst), "1"}, {"B_TYPE", "uint"}, {"B_SIZE", "32"}, {"S_TYPE", src.second}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
             string_to_spv("set_rows_" + std::string(src.first) + "_" + dst + "_i64", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(dst), "1"}, {"B_TYPE", "uvec2"}, {"B_SIZE", "64"}, {"S_TYPE", src.second}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
         }

ggml/src/ggml.c

@@ -1080,6 +1080,9 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "SOLVE_TRI",
     "GATED_DELTA_NET",
     "LIGHTNING_INDEXER",
+    "DSV4_HC_COMB",
+    "DSV4_HC_PRE",
+    "DSV4_HC_POST",
 
     "UNARY",
 
@@ -1097,7 +1100,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "GLU",
 };
 
-static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98");
+static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101");
 
 static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "none",
@@ -1192,6 +1195,9 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "A X = B, A triangular, solve X",
     "gated_delta_net(q, k, v, g, beta, s)",
     "lightning_indexer(q, k, weights, mask)",
+    "dsv4_hc_comb(mixes, scale, base)",
+    "dsv4_hc_pre(x, weights)",
+    "dsv4_hc_post(x, residual, post, comb)",
 
     "unary(x)",
 
@@ -1209,7 +1215,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "glu(x)",
 };
 
-static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98");
+static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101");
 
 static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
 
@@ -5437,6 +5443,7 @@ struct ggml_tensor * ggml_flash_attn_ext(
     return result;
 }
 
+
 void ggml_flash_attn_ext_set_prec(
         struct ggml_tensor * a,
         enum ggml_prec       prec) {
@@ -6337,6 +6344,132 @@ struct ggml_tensor * ggml_lightning_indexer(
     return result;
 }
 
+// ggml_dsv4_hc_comb
+
+struct ggml_tensor * ggml_dsv4_hc_comb(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * mixes,
+        struct ggml_tensor  * scale,
+        struct ggml_tensor  * base,
+        float                 eps,
+        int32_t               n_iter) {
+    GGML_ASSERT(mixes->type == GGML_TYPE_F32);
+    GGML_ASSERT(scale->type == GGML_TYPE_F32);
+    GGML_ASSERT(base->type == GGML_TYPE_F32);
+    GGML_ASSERT(n_iter > 0);
+
+    const int64_t hc_mix_dim = mixes->ne[0];
+    const int64_t n_tokens   = mixes->ne[1];
+
+    int64_t hc = 0;
+    for (int64_t i = 1; i*i + 2*i <= hc_mix_dim; ++i) {
+        if ((2 + i)*i == hc_mix_dim) {
+            hc = i;
+            break;
+        }
+    }
+
+    GGML_ASSERT(hc > 0);
+    GGML_ASSERT(hc == 4);
+    GGML_ASSERT(mixes->ne[2] == 1);
+    GGML_ASSERT(mixes->ne[3] == 1);
+    GGML_ASSERT(scale->ne[0] >= 3);
+    GGML_ASSERT(scale->ne[1] == 1);
+    GGML_ASSERT(scale->ne[2] == 1);
+    GGML_ASSERT(scale->ne[3] == 1);
+    GGML_ASSERT(base->ne[0] == hc_mix_dim);
+    GGML_ASSERT(base->ne[1] == 1);
+    GGML_ASSERT(base->ne[2] == 1);
+    GGML_ASSERT(base->ne[3] == 1);
+
+    struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hc, hc, n_tokens);
+
+    ggml_set_op_params_f32(result, 0, eps);
+    ggml_set_op_params_i32(result, 1, n_iter);
+
+    result->op     = GGML_OP_DSV4_HC_COMB;
+    result->src[0] = mixes;
+    result->src[1] = scale;
+    result->src[2] = base;
+
+    return result;
+}
+
+// ggml_dsv4_hc_pre
+
+struct ggml_tensor * ggml_dsv4_hc_pre(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * x,
+        struct ggml_tensor  * weights) {
+    GGML_ASSERT(x->type == GGML_TYPE_F32);
+    GGML_ASSERT(weights->type == GGML_TYPE_F32);
+
+    const int64_t n_embd   = x->ne[0];
+    const int64_t hc       = x->ne[1];
+    const int64_t n_tokens = x->ne[2];
+
+    GGML_ASSERT(hc > 0);
+    GGML_ASSERT(x->ne[3] == 1);
+    GGML_ASSERT(weights->ne[0] == hc);
+    GGML_ASSERT(weights->ne[1] == n_tokens);
+    GGML_ASSERT(weights->ne[2] == 1);
+    GGML_ASSERT(weights->ne[3] == 1);
+
+    struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens);
+
+    result->op     = GGML_OP_DSV4_HC_PRE;
+    result->src[0] = x;
+    result->src[1] = weights;
+
+    return result;
+}
+
+// ggml_dsv4_hc_post
+
+struct ggml_tensor * ggml_dsv4_hc_post(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * x,
+        struct ggml_tensor  * residual,
+        struct ggml_tensor  * post,
+        struct ggml_tensor  * comb) {
+    GGML_ASSERT(x->type == GGML_TYPE_F32);
+    GGML_ASSERT(residual->type == GGML_TYPE_F32);
+    GGML_ASSERT(post->type == GGML_TYPE_F32);
+    GGML_ASSERT(comb->type == GGML_TYPE_F32);
+
+    const int64_t n_embd   = x->ne[0];
+    const int64_t n_tokens = x->ne[1];
+    const int64_t hc       = residual->ne[1];
+
+    GGML_ASSERT(hc > 0);
+    GGML_ASSERT(x->ne[2] == 1);
+    GGML_ASSERT(x->ne[3] == 1);
+
+    GGML_ASSERT(residual->ne[0] == n_embd);
+    GGML_ASSERT(residual->ne[2] == n_tokens);
+    GGML_ASSERT(residual->ne[3] == 1);
+
+    GGML_ASSERT(post->ne[0] == hc);
+    GGML_ASSERT(post->ne[1] == n_tokens);
+    GGML_ASSERT(post->ne[2] == 1);
+    GGML_ASSERT(post->ne[3] == 1);
+
+    GGML_ASSERT(comb->ne[0] == hc);
+    GGML_ASSERT(comb->ne[1] == hc);
+    GGML_ASSERT(comb->ne[2] == n_tokens);
+    GGML_ASSERT(comb->ne[3] == 1);
+
+    struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, hc, n_tokens);
+
+    result->op     = GGML_OP_DSV4_HC_POST;
+    result->src[0] = x;
+    result->src[1] = residual;
+    result->src[2] = post;
+    result->src[3] = comb;
+
+    return result;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 struct ggml_hash_set ggml_hash_set_new(size_t size) {

scripts/snapdragon/ggml-hexagon-profile.py

@@ -458,15 +458,11 @@ def main():
 
     if args.timeline:
         logger.info(f"\n# ASCII Timing {args.timeline.capitalize()}\n")
-        printed_cnt = 0
         for op in ops:
             if args.timeline == "summary":
                 print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'], op.get('evt_val'))
             elif args.timeline == "diagram":
                 print_ascii_timeline(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events'], op.get('evt_val'))
-            printed_cnt += 1
-            if printed_cnt >= args.top:
-                break
     else:
         generate_report(ops, args.top, overrides, args.sort, pmu_name=final_pmu_name)
 

scripts/sync-ggml.last

@@ -1 +1 @@
-eaa0a74fa768bb72da623a61d9da3d436053ea91
+9be313313c8ecb9488911bd64550190e3ed80f38