llama-batch: fix allowed decreasing pos in a seq (#25449)

* hex-fa: refactor kernel param compute to use common layout builder

* hmx: add explicit compiler barriers to make hmx funcs more robust

* hex-vtcm: more generic vtcm layout builder for mm and flash-attn kernels

* hex-hmx: unroll inner kernels

* hex-hmx: use inline asm instead of intrinsics to avoid compiler issues

* hex-hmx: define inline asm macros and simplify code

* hex-hmx: replace leftover intrinsics

* hmx-fa: minor cleanup for hmx asm

* hmx-mm: move per-task stucts out of the kernels header

* hmx-mm: simplify core_dot_chunk

* hmx-mm: simplify inner loops that call hmx instructions

* hmx-mm: proper instrumentation for activation prep work for dma pipelined version

* hmx-mm: update a-prep loop for better prefetch

* hex-vtcm: improved vtcm layout alloc for mm to support overlapping areas

* hmx-mm: reduce the number of act fetch tows to 4 for now, going larger doesnt help here

* hex-hmx: always use hmx-queue in all modes

* hmx-mm: update comments and minor formatting

* hmx-mm: further improve synchro fallback path to prefetch the weights earlier

* hex-fa: further pipeline improvements (earlier prefetch)

* hmx-mm: cleanup dma pipelines to use dst cached in the queue

* hmx-fa: minor cleanup and opts for fa dma pipelines

* hmx-fa: optimize q-prep stage with dma and unrolling

* hmx-fa: use o_tile size from layout instead of computing it

* hmx-mm: cleanup types and size handling

* hmx-mm: replace divs with fastdiv in qprep loops

* hmx-fa: minor update/formatting to q_tile handling

* hmx-fa: cleanup the layout to avoid overpadding

* hmx-fa: simplified and improved cost mode for hmx fa solver that uses vtcm layout funcs

* hmx-queue: add support queue wakeup and make suspend async to avoid hmx-lock latency

* hex-hmx: move queue wakeup / suspend to the op-batch level

* hex-threads: add hybrid polling to workpool

* hex-mm: fix trailing spaces

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

@@ -9,6 +9,8 @@ on:
       '.github/workflows/hip-quality-check.yml',
       '**/*.cu',
       '**/*.cuh',
+      'ggml/src/ggml-hip/CMakeLists.txt',
+      'ggml/src/ggml-cuda/vendors/hip.h',
       'scripts/hip/gcn-cdna-vgpr-check.py'
     ]
 
@@ -18,6 +20,8 @@ on:
       '.github/workflows/hip-quality-check.yml',
       '**/*.cu',
       '**/*.cuh',
+      'ggml/src/ggml-hip/CMakeLists.txt',
+      'ggml/src/ggml-cuda/vendors/hip.h',
       'scripts/hip/gcn-cdna-vgpr-check.py'
     ]
 

common/arg.cpp

@@ -27,6 +27,7 @@
 #include <cinttypes>
 #include <climits>
 #include <cstdarg>
+#include <filesystem>
 #include <fstream>
 #include <list>
 #include <regex>
@@ -718,9 +719,8 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
 
         // model is required (except for server)
         // TODO @ngxson : maybe show a list of available models in CLI in this case
-        if (params.model.path.empty()
-                && !params.usage
-                && !params.completion) {
+        bool can_skip_model = params.usage || params.completion || !params.server_base.empty();
+        if (!can_skip_model && params.model.path.empty()) {
             throw std::invalid_argument("error: --model is required\n");
         }
     }
@@ -1240,6 +1240,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.completion = true;
         }
     ));
+    add_opt(common_arg(
+        {"--server-base"}, "URL",
+        string_format("connect to this server instead of starting a new one, example: 'http://localhost:8080' (default: none)"),
+        [](common_params & params, const std::string & value) {
+            params.server_base = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_CLI}));
     add_opt(common_arg(
         {"--verbose-prompt"},
         string_format("print a verbose prompt before generation (default: %s)", params.verbose_prompt ? "true" : "false"),
@@ -3451,9 +3458,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     ).set_env("LLAMA_ARG_LOG_FILE"));
     add_opt(common_arg(
         {"--log-prompts-dir"}, "PATH",
-        "Log prompts to directory (only used for debugging, default: disabled)",
+        "Log prompts to directory (auto-created if not present; only used for debugging, default: disabled)",
         [](common_params & params, const std::string & value) {
             params.path_prompts_log_dir = value;
+            std::error_code ec;
+            std::filesystem::create_directories(value, ec);
+            if (ec) {
+                fprintf(stderr, "warning: failed to create prompts-log-dir '%s': %s\n", value.c_str(), ec.message().c_str());
+            }
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
     add_opt(common_arg(

common/common.cpp

@@ -55,6 +55,10 @@
 #include <pwd.h>
 #endif
 
+#if defined(_AIX)
+#include <sys/systemcfg.h>
+#endif
+
 #if defined(_MSC_VER)
 #pragma warning(disable: 4244 4267) // possible loss of data
 #endif
@@ -72,7 +76,16 @@ common_time_meas::~common_time_meas() {
 //
 
 int32_t common_cpu_get_num_physical_cores() {
-#ifdef __linux__
+#if defined(_AIX)
+    int32_t logical_cpus = _system_configuration.ncpus;
+    int32_t smt_threads = _system_configuration.smt_threads;
+    if (smt_threads > 0) {
+        return static_cast<int32_t>(logical_cpus / smt_threads);
+    }
+    if (logical_cpus > 0) {
+        return static_cast<int32_t>(logical_cpus);
+    }
+#elif defined(__linux__)
     // enumerate the set of thread siblings, num entries is num cores
     std::unordered_set<std::string> siblings;
     for (uint32_t cpu=0; cpu < UINT32_MAX; ++cpu) {
@@ -202,6 +215,14 @@ int32_t common_cpu_get_num_math() {
             }
         }
     }
+#elif defined(__powerpc64__) || defined(__powerpc__)
+    int32_t smt_factor = 1;
+    int phy_cpus = common_cpu_get_num_physical_cores();
+    int logical_cpus = sysconf(_SC_NPROCESSORS_ONLN);
+    if (phy_cpus > 0 && logical_cpus > phy_cpus) {
+        smt_factor = logical_cpus / phy_cpus;
+    }
+    return phy_cpus * std::min(smt_factor, 2);
 #endif
     return common_cpu_get_num_physical_cores();
 }

common/common.h

@@ -14,6 +14,7 @@
 #include <vector>
 #include <map>
 #include <algorithm>
+#include <fstream>
 
 #if defined(_WIN32) && !defined(_WIN32_WINNT)
 #define _WIN32_WINNT 0x0A00
@@ -643,6 +644,9 @@ struct common_params {
 
     std::map<std::string, std::string> default_template_kwargs;
 
+    // CLI params
+    std::string server_base; // if set, connect to this server instead of starting a new one
+
     // UI configs
     bool ui = true;
     bool ui_mcp_proxy = false;

common/http.h

@@ -2,6 +2,16 @@
 
 #include <cpp-httplib/httplib.h>
 
+#ifdef _WIN32
+#include <winsock2.h>
+#include <windows.h>
+#else
+#include <sys/socket.h>
+#include <netinet/in.h>
+#include <arpa/inet.h>
+#include <unistd.h>
+#endif
+
 struct common_http_url {
     std::string scheme;
     std::string user;
@@ -119,3 +129,63 @@ static std::pair<httplib::Client, common_http_url> common_http_client(const std:
 static std::string common_http_show_masked_url(const common_http_url & parts) {
     return parts.scheme + "://" + (parts.user.empty() ? "" : "****:****@") + common_http_format_host(parts.host) + parts.path;
 }
+
+static int common_http_get_free_port() {
+#ifdef _WIN32
+    WSADATA wsaData;
+    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
+        return -1;
+    }
+    typedef SOCKET native_socket_t;
+#define INVALID_SOCKET_VAL INVALID_SOCKET
+#define CLOSE_SOCKET(s) closesocket(s)
+#else
+    typedef int native_socket_t;
+#define INVALID_SOCKET_VAL -1
+#define CLOSE_SOCKET(s) close(s)
+#endif
+
+    native_socket_t sock = socket(AF_INET, SOCK_STREAM, 0);
+    if (sock == INVALID_SOCKET_VAL) {
+#ifdef _WIN32
+        WSACleanup();
+#endif
+        return -1;
+    }
+
+    struct sockaddr_in serv_addr;
+    std::memset(&serv_addr, 0, sizeof(serv_addr));
+    serv_addr.sin_family = AF_INET;
+    serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
+    serv_addr.sin_port = htons(0);
+
+    if (bind(sock, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) != 0) {
+        CLOSE_SOCKET(sock);
+#ifdef _WIN32
+        WSACleanup();
+#endif
+        return -1;
+    }
+
+#ifdef _WIN32
+    int namelen = sizeof(serv_addr);
+#else
+    socklen_t namelen = sizeof(serv_addr);
+#endif
+    if (getsockname(sock, (struct sockaddr*)&serv_addr, &namelen) != 0) {
+        CLOSE_SOCKET(sock);
+#ifdef _WIN32
+        WSACleanup();
+#endif
+        return -1;
+    }
+
+    int port = ntohs(serv_addr.sin_port);
+
+    CLOSE_SOCKET(sock);
+#ifdef _WIN32
+    WSACleanup();
+#endif
+
+    return port;
+}

common/ngram-map.cpp

@@ -125,6 +125,16 @@ void common_ngram_map_begin(
     LOG_DBG("%s: begin, idx_last_draft=%zu, new begin=%zu, #keys=%zu\n", __func__,
             map.idx_last_check, size_begin, map.keys.size());
 
+    size_t idx_begin_cleanup = map.size_last_begin;
+    if (idx_begin_cleanup > size_begin) {
+        if (size_begin > (size_t) map.size_key + map.size_value) {
+            idx_begin_cleanup = size_begin - map.size_key - map.size_value;
+        } else {
+            idx_begin_cleanup = 0;
+        }
+        LOG_INF("%s: shrink cleanup begin: %zu -> %zu\n", __func__, map.size_last_begin, idx_begin_cleanup);
+    }
+
     size_t count_map_entries_upd = 0;
     if (!map.key_map.empty() && size_begin < map.idx_last_check) {
         if (map.show_key_map_stats) {
@@ -150,27 +160,23 @@ void common_ngram_map_begin(
         // Update the map from hash to key index (clear outdated entries).
         for (size_t i = 0; i < map.key_map.size(); ++i) {
             uint32_t key_idx = map.key_map[i];
-            if (key_idx >= map.size_last_begin) {
+            if (key_idx != 0 && key_idx >= idx_begin_cleanup) {
                 map.key_map[i] = 0;
                 count_map_entries_upd++;
             }
         }
-        map.key_map_last_idx = (map.size_last_begin > 0) ? map.size_last_begin - 1 : 0;
+        map.key_map_last_idx = (idx_begin_cleanup > 0) ? (uint32_t) (idx_begin_cleanup - 1) : 0;
     }
 
     if (size_begin < map.idx_last_check && !map.keys.empty()) {
-        // The next token generation will start at index size_begin.
-        // The tokens between map.size_last_begin and size_begin are no longer valid.
-        //
-        // Refresh map: Remove all entries with index >= map.size_last_begin.
         size_t count_keys = map.keys.size();
         size_t count_keys_del = 0;
         size_t count_values_del = 0;
         for (int32_t i = map.keys.size() - 1; i >= 0; --i) {
             common_ngram_map_key & key = map.keys[i];
-            if (key.key_idx >= map.size_last_begin) {
+            if (key.key_idx >= idx_begin_cleanup) {
                 // Delete the key.
-                LOG_DBG("%s: delete key %d at index %zu (>= size_last_begin=%zu)\n", __func__, i, key.key_idx, map.size_last_begin);
+                LOG_DBG("%s: delete key %d at index %zu (>= idx_begin_cleanup=%zu)\n", __func__, i, key.key_idx, idx_begin_cleanup);
                 map.keys.erase(map.keys.begin() + i);
                 count_keys_del++;
                 continue;
@@ -182,7 +188,7 @@ void common_ngram_map_begin(
             // Check the indices of the values.
             for (int16_t j = COMMON_NGRAM_MAX_VALUES - 1; j >= 0; --j) {
                 common_ngram_map_value & value = key.values[j];
-                if (value.value_idx >= map.size_last_begin) {
+                if (value.value_idx != 0 && value.value_idx >= idx_begin_cleanup) {
                     // Delete the value.
                     count_values_del++;
 

common/speculative.cpp

@@ -2221,6 +2221,112 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) {
     return n_max;
 }
 
+common_params common_base_params_to_speculative(const common_params & params) {
+    const bool has_draft = params.speculative.has_dft();
+
+    const auto & params_spec = params.speculative.draft;
+    common_params result = params;
+
+    if (has_draft) {
+        result.devices               = params_spec.devices;
+        result.model                 = params_spec.mparams;
+        result.n_gpu_layers          = params_spec.n_gpu_layers;
+        result.tensor_buft_overrides = params_spec.tensor_buft_overrides;
+
+        if (params_spec.cpuparams.n_threads > 0) {
+            result.cpuparams.n_threads       = params_spec.cpuparams.n_threads;
+            result.cpuparams_batch.n_threads = params_spec.cpuparams_batch.n_threads;
+        }
+    }
+
+    result.cache_type_k  = params_spec.cache_type_k;
+    result.cache_type_v  = params_spec.cache_type_v;
+    result.n_outputs_max = params.n_parallel;
+
+    return result;
+}
+
+struct common_speculative_init_result::impl {
+    impl() = default;
+    ~impl() = default;
+
+    // note: the order in which model, context, etc. are declared matters because their destructors will be called bottom-to-top
+    llama_model_ptr   model;
+    llama_context_ptr context;
+};
+
+common_speculative_init_result::common_speculative_init_result(
+    common_params & params,
+      llama_model * model_tgt,
+    llama_context * ctx_tgt) :
+    pimpl(new impl{}) {
+    const bool has_draft = params.speculative.has_dft();
+    const bool spec_mtp = std::find(params.speculative.types.begin(),
+                                    params.speculative.types.end(),
+                                    COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end();
+    GGML_ASSERT(has_draft || spec_mtp);
+
+    auto mparams = common_model_params_to_llama(params);
+    auto cparams = common_context_params_to_llama(params);
+
+    if (spec_mtp) {
+        cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
+    }
+
+    // note: for small models maybe we can set this to the maximum possible draft from all speculative types
+    //       the extra memory for small models is likely negligible?
+    cparams.n_rs_seq  = 0;
+    cparams.ctx_other = ctx_tgt;
+
+    std::string model_path;
+    if (has_draft) {
+        model_path = params.speculative.draft.mparams.path;
+        LOG_TRC("%s: loading draft model '%s'\n", __func__, model_path.c_str());
+
+        llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams);
+        if (model_dft == NULL) {
+            LOG_ERR("%s: failed to load draft model, '%s'\n", __func__, model_path.c_str());
+            return;
+        }
+
+        pimpl->model.reset(model_dft);
+
+        llama_context * ctx_dft = llama_init_from_model(model_dft, cparams);
+        if (ctx_dft == nullptr) {
+            LOG_ERR("%s: failed to create MTP context\n", __func__);
+            return;
+        }
+
+        pimpl->context.reset(ctx_dft);
+    } else if (spec_mtp) {
+        model_path = params.model.path;
+
+        LOG_TRC("%s: creating MTP draft context against the target model '%s'\n", __func__, model_path.c_str());
+
+        llama_context * ctx_dft = llama_init_from_model(model_tgt, cparams);
+        if (ctx_dft == nullptr) {
+            LOG_ERR("%s: failed to create MTP context\n", __func__);
+            return;
+        }
+
+        pimpl->context.reset(ctx_dft);
+    }
+}
+
+common_speculative_init_result::~common_speculative_init_result() = default;
+
+llama_model * common_speculative_init_result::model() {
+    return pimpl->model.get();
+}
+
+llama_context * common_speculative_init_result::context() {
+    return pimpl->context.get();
+}
+
+common_speculative_init_result_ptr common_speculative_init_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt) {
+    return std::make_unique<common_speculative_init_result>(params, model_tgt, ctx_tgt);
+}
+
 // initialization of the speculative decoding system
 //
 common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq) {

common/speculative.h

@@ -23,6 +23,8 @@ std::string common_speculative_type_to_str(enum common_speculative_type type);
 // return the max number of draft tokens based on the speculative parameters
 int32_t common_speculative_n_max(const common_params_speculative * spec);
 
+common_params common_base_params_to_speculative(const common_params & params);
+
 common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq);
 
 void common_speculative_free(common_speculative * spec);
@@ -80,3 +82,19 @@ struct common_speculative_deleter {
 };
 
 typedef std::unique_ptr<common_speculative, common_speculative_deleter> common_speculative_ptr;
+
+struct common_speculative_init_result {
+    common_speculative_init_result(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt);
+    ~common_speculative_init_result();
+
+    llama_model   * model();
+    llama_context * context();
+
+private:
+    struct impl;
+    std::unique_ptr<impl> pimpl;
+};
+
+using common_speculative_init_result_ptr = std::unique_ptr<common_speculative_init_result>;
+
+common_speculative_init_result_ptr common_speculative_init_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt);

docs/backend/SYCL.md

@@ -790,10 +790,10 @@ use 1 SYCL GPUs: [0] with Max compute units:512
 | GGML_SYCL_DEBUG   | 0 (default) or 1 | Enable log function by macro: GGML_SYCL_DEBUG                                                                             |
 | GGML_SYCL_DEV2DEV_MEMCPY | 0 (default) or 1 | Choose the SYCL or L0 API in dev2dev memory copy.<br>Value: <br>*  0: SYCL API (default)<br>* 1: L0 API -- L0 API is found to lead to abnormal crash in some case. This debug flag is used to check the issue.|
 | GGML_SYCL_ENABLE_FLASH_ATTN | 1 (default) or 0| Enable Flash-Attention. It can reduce memory usage. The performance impact depends on the LLM.|
-| GGML_SYCL_DISABLE_OPT | 0 (default) or 1 | Disable optimize features for Intel GPUs. (Recommended to 1 for Intel devices older than Gen 10) |
-| GGML_SYCL_DISABLE_GRAPH | 0 or 1 (default) | Disable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. |
+| GGML_SYCL_ENABLE_OPT | 0 or 1 (default)| Enable optimize features for Intel GPUs. (Recommended to 0 for Intel devices older than Gen 10) |
+| GGML_SYCL_ENABLE_GRAPH | 0 (default) or 1 | Enable running computations through SYCL Graphs feature. Disabled by default because SYCL Graph is still on development, no better performance. |
 | GGML_SYCL_USE_LEVEL_ZERO_API | 1 (default) or 0 | Use Level Zero API for device memory allocation instead of SYCL. Reduces system RAM usage on Intel dGPUs by avoiding DMA-buf/TTM host memory staging. Requires GGML_SYCL_SUPPORT_LEVEL_ZERO_API=ON at build time. SYCL backend always runs on Level Zero running time even if it's set as OFF (The SYCL api will be usage for memory allocation).|
-| GGML_SYCL_DISABLE_DNN | 0 (default) or 1 | Disable running computations through oneDNN and always use oneMKL. |
+| GGML_SYCL_ENABLE_DNN | 0 or 1 (default)| Enable running computations through oneDNN and always use oneMKL. |
 | GGML_SYCL_ENABLE_VMM | 0 or 1 (default) | Enable the virtual-memory device pool. |
 | ZES_ENABLE_SYSMAN | 0 (default) or 1 | Support to get free memory of GPU by sycl::aspect::ext_intel_free_memory.<br>Recommended to use when --split-mode = layer |
 | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. |
@@ -807,7 +807,7 @@ Pass these via `CXXFLAGS` or add a one-off `#define` to enable a flag on the spo
 |-----------------|----------------------------------------------------------------------------------|
 | DEBUG_SYCL_POOL | Enable device memory pool logging on teardown. Useful for profiling allocations. |
 | DEBUG_SYCL_MALLOC | Enable verbose per-call logging of device pool alloc/free operations. |
-
+| GGML_SYCL_SUPPORT_VMM | Support to building with VMM code. Default is Yes. |
 
 ## Design Rule
 

docs/build.md

@@ -270,13 +270,10 @@ The environment variable [`CUDA_SCALE_LAUNCH_QUEUES`](https://docs.nvidia.com/cu
 
 Consider setting `CUDA_SCALE_LAUNCH_QUEUES=4x`, which increases the CUDA command buffer to 4 times its default size. This optimization is particularly beneficial for **Multi-GPU setups with pipeline parallelism**, where it significantly improves prompt processing throughput by allowing more operations to be enqueued across GPUs.
 
-#### GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F
+#### GGML_CUDA_CUBLAS_COMPUTE_TYPE
 
-Use `GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F` environment variable to use FP32 compute type on all GPUs in FP16 cuBLAS for preventing possible numerical overflows in exchange for slower prompt processing (small impact on RTX PRO/Datacenter products and significant on GeForce products).
-
-#### GGML_CUDA_FORCE_CUBLAS_COMPUTE_16F
-
-Use `GGML_CUDA_FORCE_CUBLAS_COMPUTE_16F` environment variable to force use FP16 compute type (instead of default FP32) in FP16 cuBLAS for V100, CDNA and RDNA4.
+Override default, speed-optimized compute types for cuBLAS matrix multiplications.
+Legal values: `auto`, `f16`, `fp16`, `bf16`, `f32`, `fp32`.
 
 ### Unified Memory
 

docs/ops.md

@@ -21,12 +21,12 @@ Legend:
 |                           ADD_ID | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           ARANGE | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ |
+|                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             CEIL | ❌ | ❌ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            CLAMP | ❌ | ✅ | ✅ | ✅ | ✅ | 🟡 | ✅ | 🟡 | ✅ | ❌ | ❌ |
-|                        COL2IM_1D | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                        COL2IM_1D | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                           CONCAT | ❌ | ✅ | ✅ | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             CONT | ❌ | 🟡 | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ❌ | ❌ |
+|                             CONT | ❌ | 🟡 | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                          CONV_2D | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                          CONV_3D | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
@@ -35,8 +35,8 @@ Legend:
 |                              COS | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | 🟡 | ✅ | ❌ | ❌ |
 |                      COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              CPY | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
-|               CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-|          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|               CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                           CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             DIAG | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -70,7 +70,7 @@ Legend:
 |                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
 |                 MUL_MAT_HADAMARD | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | 🟡 | ❌ |
+|                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ❌ |
 |                              NEG | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |

examples/llama-eval/llama-eval.py

@@ -362,7 +362,7 @@ class EvalState:
             case = cases.get(task_id, {})
             status = case.get("status", "pending")
             expected = case.get("expected", "")
-            answer = case.get("answer", "") if status == "ok" else ""
+            answer = case.get("answer") or "" if status == "ok" else ""
             is_correct = case.get("correct", False) if status == "ok" else False
             response = case.get("response", "") or ""
             prompt = case.get("prompt", "") or ""
@@ -647,7 +647,7 @@ class EvalState:
             question, prompt, expected = self.get_case(i)
             case = cases.get(task_id, {})
             status = case.get("status", "pending")
-            answer = case.get("answer", "N/A") if status == "ok" else "N/A"
+            answer = case.get("answer") or "N/A" if status == "ok" else "N/A"
             tokens = case.get("tokens")
             tokens_str = str(tokens) if tokens is not None else "N/A"
             tps_gen = case.get("tps_gen")

ggml/include/ggml-cuda.h

@@ -30,9 +30,6 @@ GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int de
 // conduct allreduce operation between devices
 GGML_BACKEND_API bool ggml_backend_cuda_allreduce_tensor(ggml_backend_t * backends, struct ggml_tensor ** tensors, size_t n_backends);
 
-// split tensor buffer that splits matrices by rows across multiple devices
-GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_split_buffer_type(int main_device, const float * tensor_split);
-
 // pinned host buffer for use with the CPU backend for faster copies between CPU and GPU
 GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type(void);
 

ggml/include/ggml.h

@@ -429,7 +429,8 @@ extern "C" {
         GGML_TYPE_MXFP4   = 39, // MXFP4 (1 block)
         GGML_TYPE_NVFP4   = 40, // NVFP4 (4 blocks, E4M3 scale)
         GGML_TYPE_Q1_0    = 41,
-        GGML_TYPE_COUNT   = 42,
+        GGML_TYPE_Q2_0    = 42,
+        GGML_TYPE_COUNT   = 43,
     };
 
     // precision
@@ -473,6 +474,7 @@ extern "C" {
         GGML_FTYPE_MOSTLY_MXFP4   = 25, // except 1d tensors
         GGML_FTYPE_MOSTLY_NVFP4   = 26, // except 1d tensors
         GGML_FTYPE_MOSTLY_Q1_0    = 27, // except 1d tensors
+        GGML_FTYPE_MOSTLY_Q2_0    = 28, // except 1d tensors
     };
 
     // available tensor operations:

ggml/src/ggml-common.h

@@ -96,6 +96,9 @@ typedef sycl::half2 ggml_half2;
 #define QI1_0 (QK1_0 / 32)
 #define QR1_0 1
 
+#define QI2_0 (QK2_0 / 32)
+#define QR2_0 1
+
 
 #define QI4_0 (QK4_0 / (4 * QR4_0))
 #define QR4_0 2
@@ -181,6 +184,13 @@ typedef struct {
 } block_q1_0;
 static_assert(sizeof(block_q1_0) == sizeof(ggml_half) + QK1_0 / 8, "wrong q1_0 block size/padding");
 
+#define QK2_0 64
+typedef struct {
+    ggml_half d;              // delta (scale)
+    uint8_t qs[QK2_0 / 4];   // 2 bits per element
+} block_q2_0;
+static_assert(sizeof(block_q2_0) == sizeof(ggml_half) + QK2_0 / 4, "wrong q2_0 block size/padding");
+
 #define QK4_0 32
 typedef struct {
     ggml_half d;           // delta

ggml/src/ggml-cpu/arch-fallback.h

@@ -17,6 +17,7 @@
 #define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0
 #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
 #define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
 #define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
 #define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
 #define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K
@@ -82,6 +83,7 @@
 #define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
 #elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64)
 // quants.c
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
 // repack.cpp
 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
@@ -113,6 +115,7 @@
 #define quantize_row_q8_K_generic quantize_row_q8_K
 #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
 #define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
 #define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
 #define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
 #define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K
@@ -162,6 +165,7 @@
 #define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0
 #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
 #define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
 // repack.cpp
 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8
@@ -202,6 +206,7 @@
 #elif defined(__riscv)
 // quants.c
 #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
 // repack.cpp
 #define ggml_quantize_mat_q8_0_4x1_generic ggml_quantize_mat_q8_0_4x1
 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
@@ -243,6 +248,7 @@
 #define quantize_row_q8_K_generic quantize_row_q8_K
 #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
 #define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
 #define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K
 #define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K
 #define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K
@@ -306,6 +312,7 @@
 #define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0
 #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0
 #define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0
+#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
 // repack.cpp
 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8

ggml/src/ggml-cpu/arch/arm/quants.c

@@ -219,6 +219,80 @@ void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi
 #endif
 }
 
+void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK2_0;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q2_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    float sumf = 0.0f;
+
+#if defined(__ARM_NEON)
+    // Replicate pattern: each byte repeated 4 times
+    static const uint8_t tbl_idx_lo[16] = {0,0,0,0, 1,1,1,1, 2,2,2,2, 3,3,3,3};
+    static const uint8_t tbl_idx_hi[16] = {4,4,4,4, 5,5,5,5, 6,6,6,6, 7,7,7,7};
+    // Right-shift amounts: 0,2,4,6 repeated for each group of 4
+    static const int8_t shift_vals[16] = {0,-2,-4,-6, 0,-2,-4,-6, 0,-2,-4,-6, 0,-2,-4,-6};
+
+    const uint8x16_t idx_lo  = vld1q_u8(tbl_idx_lo);
+    const uint8x16_t idx_hi  = vld1q_u8(tbl_idx_hi);
+    const int8x16_t  shifts  = vld1q_s8(shift_vals);
+    const uint8x16_t mask2   = vdupq_n_u8(0x03);
+    const int8x16_t  one     = vdupq_n_s8(1);
+
+    float32x4_t sumv = vdupq_n_f32(0.0f);
+
+    for (int i = 0; i < nb; i++) {
+        const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d);
+
+        // group 64: one Q2_0 block (64 weights) maps to two Q8_0 blocks (2 * 32 = 64)
+        for (int k = 0; k < 2; k++) {
+            const block_q8_0 * GGML_RESTRICT yb = &y[i * 2 + k];
+            const float d1 = GGML_CPU_FP16_TO_FP32(yb->d);
+
+            // Load 8 bytes of packed 2-bit values
+            const uint8x8_t raw = vld1_u8(&x[i].qs[k * 8]);
+            const uint8x16_t raw16 = vcombine_u8(raw, raw);
+
+            // First 16 elements: replicate bytes 0-3, shift, mask, subtract 1
+            uint8x16_t bytes0 = vqtbl1q_u8(raw16, idx_lo);
+            int8x16_t qv0 = vsubq_s8(
+                vreinterpretq_s8_u8(vandq_u8(vshlq_u8(bytes0, shifts), mask2)),
+                one);
+
+            // Second 16 elements: replicate bytes 4-7, shift, mask, subtract 1
+            uint8x16_t bytes1 = vqtbl1q_u8(raw16, idx_hi);
+            int8x16_t qv1 = vsubq_s8(
+                vreinterpretq_s8_u8(vandq_u8(vshlq_u8(bytes1, shifts), mask2)),
+                one);
+
+            // Load Q8_0 values and dot product
+            const int8x16_t y0 = vld1q_s8(yb->qs);
+            const int8x16_t y1 = vld1q_s8(yb->qs + 16);
+
+            int32x4_t p0 = ggml_vdotq_s32(vdupq_n_s32(0), qv0, y0);
+            int32x4_t p1 = ggml_vdotq_s32(p0, qv1, y1);
+
+            sumv = vmlaq_n_f32(sumv, vcvtq_f32_s32(p1), d0 * d1);
+        }
+    }
+
+    sumf = vaddvq_f32(sumv);
+#else
+    ggml_vec_dot_q2_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
+    return;
+#endif
+
+    *s = sumf;
+}
 
 void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
     const int qk = QK8_0;
@@ -812,10 +886,10 @@ void ggml_vec_dot_nvfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
         const float dy0 = GGML_CPU_FP16_TO_FP32(y[2*ib].d);
         const float dy1 = GGML_CPU_FP16_TO_FP32(y[2*ib+1].d);
         const float32x4_t nvsc = {
-            ggml_ue4m3_to_fp32(x[ib].d[0]),
-            ggml_ue4m3_to_fp32(x[ib].d[1]),
-            ggml_ue4m3_to_fp32(x[ib].d[2]),
-            ggml_ue4m3_to_fp32(x[ib].d[3])
+            GGML_CPU_UE4M3_TO_FP32(x[ib].d[0]),
+            GGML_CPU_UE4M3_TO_FP32(x[ib].d[1]),
+            GGML_CPU_UE4M3_TO_FP32(x[ib].d[2]),
+            GGML_CPU_UE4M3_TO_FP32(x[ib].d[3])
         };
         const float32x4_t scales = vmulq_f32(nvsc, (float32x4_t){dy0, dy0, dy1, dy1});
 

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

@@ -230,6 +230,12 @@ static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = {
         .vec_dot_type             = GGML_TYPE_Q8_0,
         .nrows                    = 1,
     },
+    [GGML_TYPE_Q2_0] = {
+        .from_float               = quantize_row_q2_0,
+        .vec_dot                  = ggml_vec_dot_q2_0_q8_0,
+        .vec_dot_type             = GGML_TYPE_Q8_0,
+        .nrows                    = 1,
+    },
     [GGML_TYPE_Q4_0] = {
         .from_float               = quantize_row_q4_0,
         .vec_dot                  = ggml_vec_dot_q4_0_q8_0,

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

@@ -2321,24 +2321,28 @@ class tinyBLAS_Q0_PPC {
     }
 
     void matmul(int64_t m, int64_t n) {
-    #if defined(_AIX) || defined(__BIG_ENDIAN__)
-        mnpack(0, m, 0, n);
-    #else
-        const int64_t mc = 64;
-        const int64_t kc = 64;
+        int64_t mc = 64;
         int64_t nc = 64;
+        int64_t kc = 64;
+        int64_t n_chunk = 64;
+    #if defined(_AIX) || defined(__BIG_ENDIAN__)
+        mc = 32;
+        nc = 32;
+        kc = 32;
+        n_chunk = 32
+    #endif
         int64_t n_aligned = 0;
-        if (n % 64 == 0) {
+        if (n % n_chunk == 0) {
             n_aligned = n;
         } else if (n == 4) {
             n_aligned = 4;
-        } else if (n < 64) {
+        } else if (n < n_chunk) {
             n_aligned = (n / 8) * 8;
         } else {
-            n_aligned = (n / 64) * 64;
+            n_aligned = (n / n_chunk) * n_chunk;
         }
         if (n_aligned > 0) {
-            if (n_aligned % 64 == 0)      nc = 64;
+            if (n_aligned % n_chunk == 0) nc = n_chunk;
             else if (n_aligned == n)      nc = n;
             else if (n_aligned % 32 == 0) nc = 32;
             else if (n_aligned % 24 == 0) nc = 24;
@@ -2354,7 +2358,6 @@ class tinyBLAS_Q0_PPC {
         } else {
             mnpack(0, m, 0, n);
         }
-    #endif
     }
 
   private:
@@ -3195,16 +3198,19 @@ class tinyBLAS_PPC {
     }
 
     void matmul(int64_t m, int64_t n) {
+        int64_t mc = 256;
+        int64_t nc = 256;
+        int64_t kc = 256;
     #if defined(_AIX) || defined(__BIG_ENDIAN__)
-        mnpack(0, m, 0, n);
-    #else
-        int64_t mc = 256; int64_t nc = 256; int64_t kc = 256;
+        mc = 128;
+        nc = 128;
+        kc = 128;
+    #endif
         if (m % mc == 0 && n % nc == 0 && k % kc == 0) {
             matmul_tiled(m, n, mc, nc, kc);
         } else {
             mnpack(0, m, 0, n);
         }
-    #endif
     }
 
   private:

ggml/src/ggml-cpu/ops.cpp

@@ -665,6 +665,7 @@ void ggml_compute_forward_add(
                 ggml_compute_forward_add_non_quantized(params, dst);
             } break;
         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:
@@ -1115,6 +1116,7 @@ void ggml_compute_forward_add1(
                 }
             } break;
         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:
@@ -1245,6 +1247,7 @@ void ggml_compute_forward_acc(
         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:
@@ -4454,6 +4457,7 @@ void ggml_compute_forward_out_prod(
 
     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:
@@ -4730,6 +4734,7 @@ void ggml_compute_forward_set(
         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:
@@ -4954,6 +4959,7 @@ void ggml_compute_forward_get_rows(
 
     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:
@@ -5019,8 +5025,8 @@ void ggml_compute_forward_get_rows(
     //}
 }
 
-template<typename idx_t>
-static void ggml_compute_forward_set_rows_f32(
+template<typename src_t, typename idx_t>
+static void ggml_compute_forward_set_rows_impl(
         const ggml_compute_params * params,
               ggml_tensor * dst) {
 
@@ -5035,7 +5041,7 @@ static void ggml_compute_forward_set_rows_f32(
     assert(ne0  == nc);
     assert(ne2  == ne02);
     assert(ne3  == ne03);
-    assert(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src0->type == GGML_TYPE_F32 || (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16));
     assert(ne02 % ne11 == 0);
     assert(ne03 % ne12 == 0);
 
@@ -5049,6 +5055,8 @@ static void ggml_compute_forward_set_rows_f32(
     const int64_t ir0 = dr*ith;
     const int64_t ir1 = std::min(ir0 + dr, nr);
 
+    const size_t rs = ggml_row_size(src0->type, nc);
+
     ggml_from_float_t const from_float = ggml_get_type_traits_cpu(dst->type)->from_float;
 
     for (int64_t i03 = 0; i03 < ne03; ++i03) {
@@ -5062,9 +5070,18 @@ static void ggml_compute_forward_set_rows_f32(
 
                 GGML_ASSERT(i1 >= 0 && i1 < ne1);
 
-                from_float(
-                        (const float *) ((char *) src0->data +  i*nb01 + i02*nb02 + i03*nb03),
-                                        ((char *)  dst->data + i1*nb1  + i02*nb2  + i03*nb3), nc);
+                if constexpr (std::is_same_v<src_t, float>) {
+                    from_float(
+                            (const float *) ((char *) src0->data +  i*nb01 + i02*nb02 + i03*nb03),
+                                            ((char *)  dst->data + i1*nb1  + i02*nb2  + i03*nb3), nc);
+                } else if constexpr (std::is_same_v<src_t, ggml_fp16_t>) {
+                    memcpy(
+                                            ((char *)  dst->data + i1*nb1  + i02*nb2  + i03*nb3),
+                                            ((char *) src0->data +  i*nb01 + i02*nb02 + i03*nb03),
+                                            rs);
+                } else {
+                    GGML_ABORT("src0->type = %d (%s) not supported", src0->type, ggml_type_name(src0->type));
+                }
             }
         }
     }
@@ -5081,13 +5098,27 @@ void ggml_compute_forward_set_rows(
         case GGML_TYPE_F32:
             {
                 if (src1->type == GGML_TYPE_I64) {
-                    ggml_compute_forward_set_rows_f32<int64_t>(params, dst);
+                    ggml_compute_forward_set_rows_impl<float, int64_t>(params, dst);
                 } else if (src1->type == GGML_TYPE_I32) {
-                    ggml_compute_forward_set_rows_f32<int32_t>(params, dst);
+                    ggml_compute_forward_set_rows_impl<float, int32_t>(params, dst);
                 } else {
                     GGML_ABORT("src1->type = %d (%s) not supported", src1->type, ggml_type_name(src1->type));
                 }
             } break;
+        case GGML_TYPE_F16:
+            {
+                if (dst->type == GGML_TYPE_F16) {
+                    if (src1->type == GGML_TYPE_I64) {
+                        ggml_compute_forward_set_rows_impl<ggml_fp16_t, int64_t>(params, dst);
+                    } else if (src1->type == GGML_TYPE_I32) {
+                        ggml_compute_forward_set_rows_impl<ggml_fp16_t, int32_t>(params, dst);
+                    } else {
+                        GGML_ABORT("src1->type = %d (%s) not supported", src1->type, ggml_type_name(src1->type));
+                    }
+                } else {
+                    GGML_ABORT("dst->type = %d (%s) not supported with src0->type = %d (%s)", dst->type, ggml_type_name(dst->type), src0->type, ggml_type_name(src0->type));
+                }
+            } break;
         default:
             {
                 GGML_ABORT("src0->type = %d (%s) not supported", src0->type, ggml_type_name(src0->type));
@@ -5680,6 +5711,7 @@ void ggml_compute_forward_clamp(
             } break;
         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-cpu/quants.c

@@ -26,6 +26,10 @@ void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, in
     quantize_row_q1_0_ref(x, y, k);
 }
 
+void quantize_row_q2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_q2_0_ref(x, y, k);
+}
+
 void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
     quantize_row_q4_0_ref(x, y, k);
 }
@@ -170,6 +174,53 @@ void ggml_vec_dot_q1_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, c
     *s = sumf;
 }
 
+void ggml_vec_dot_q2_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK2_0;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q2_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+    float sumf = 0.0f;
+
+    for (int i = 0; i < nb; i++) {
+        const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d);
+
+        float sumi = 0.0f;
+
+        // group 64: one Q2_0 block (64 weights) maps to two Q8_0 blocks (2 * 32 = 64)
+        for (int k = 0; k < 2; k++) {
+            const block_q8_0 * GGML_RESTRICT yb = &y[i * 2 + k];
+            const float d1 = GGML_CPU_FP16_TO_FP32(yb->d);
+            int sumi_block = 0;
+
+            const uint8_t * GGML_RESTRICT qs = &x[i].qs[k * 8];
+            const int8_t  * GGML_RESTRICT qy = yb->qs;
+
+            for (int b = 0; b < 8; ++b) {
+                const uint8_t byte = qs[b];
+                // Extract 4 two-bit values, map {0,1,2,3} -> {-1,0,1,2}
+                sumi_block += ((int)((byte >> 0) & 3) - 1) * qy[b*4 + 0];
+                sumi_block += ((int)((byte >> 2) & 3) - 1) * qy[b*4 + 1];
+                sumi_block += ((int)((byte >> 4) & 3) - 1) * qy[b*4 + 2];
+                sumi_block += ((int)((byte >> 6) & 3) - 1) * qy[b*4 + 3];
+            }
+
+            sumi += d1 * sumi_block;
+        }
+
+        sumf += d0 * sumi;
+    }
+
+    *s = sumf;
+}
 
 void ggml_vec_dot_q4_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
     const int qk = QK8_0;

ggml/src/ggml-cpu/quants.h

@@ -13,6 +13,7 @@ extern "C" {
 
 // Quantization
 void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_q2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
 void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
 void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
 void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
@@ -38,6 +39,7 @@ void quantize_row_iq4_xs (const float * GGML_RESTRICT x, void * GGML_RESTRICT y,
 
 // Dot product
 void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
 void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
 void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
 void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
@@ -71,6 +73,7 @@ void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRI
 void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k);
 void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
 void ggml_vec_dot_q1_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_q2_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
 void ggml_vec_dot_q4_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
 void ggml_vec_dot_q4_1_q8_1_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
 void ggml_vec_dot_q5_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);

ggml/src/ggml-cpu/simd-gemm.h

@@ -78,7 +78,7 @@ static void simd_gemm(
             for (int64_t i = 0; i < GEMM_RM; i++) {
                 float a = C[i * N + jj];
                 for (int64_t kk = 0; kk < K; kk++) {
-                    a += A[i + kk] * B[kk * N + jj];
+                    a += A[i * K + kk] * B[kk * N + jj];
                 }
                 C[i * N + jj] = a;
             }

ggml/src/ggml-cpu/simd-mappings.h

@@ -131,8 +131,8 @@ extern float ggml_table_f32_ue4m3[1 << 8];
 #define GGML_CPU_E8M0_TO_FP32_HALF(x) GGML_E8M0_TO_FP32_HALF(x)
 #endif
 
-// Use lookup table for UE4M3 on x86 (faster than bit manipulation)
-#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
+// Use lookup table for UE4M3 on x86 and ARM (faster than bit manipulation)
+#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__ARM_NEON)
 #define GGML_CPU_UE4M3_TO_FP32(x) ggml_table_f32_ue4m3[(uint8_t)(x)]
 #else
 #define GGML_CPU_UE4M3_TO_FP32(x) ggml_ue4m3_to_fp32(x)

ggml/src/ggml-cuda/common.cuh

@@ -1505,12 +1505,16 @@ struct ggml_cuda_mm_fusion_args_host {
     const ggml_tensor * x_bias = nullptr;
     const ggml_tensor * gate = nullptr;
     const ggml_tensor * gate_bias = nullptr;
+    const ggml_tensor * x_scale = nullptr;
+    const ggml_tensor * gate_scale = nullptr;
     ggml_glu_op glu_op;
 };
 struct ggml_cuda_mm_fusion_args_device {
     const void * x_bias = nullptr;
     const void * gate = nullptr;
     const void * gate_bias = nullptr;
+    const void * x_scale = nullptr;
+    const void * gate_scale = nullptr;
     ggml_glu_op glu_op;
 };
 

ggml/src/ggml-cuda/convert.cu

@@ -104,8 +104,8 @@ static __global__ void dequantize_block_q4_0(const void * __restrict__ vx, dst_t
     const uint8_t * q = x->qs + 4*il;
 
     for (int l = 0; l < 4; ++l) {
-        y[l+ 0] = d * (q[l] & 0xF) + dm;
-        y[l+16] = d * (q[l] >>  4) + dm;
+        y[l+ 0] = ggml_cuda_cast<dst_t>(d * (q[l] & 0xF) + dm);
+        y[l+16] = ggml_cuda_cast<dst_t>(d * (q[l] >>  4) + dm);
     }
 }
 
@@ -131,8 +131,8 @@ static __global__ void dequantize_block_q4_1(const void * __restrict__ vx, dst_t
     const uint8_t * q = x->qs + 4*il;
 
     for (int l = 0; l < 4; ++l) {
-        y[l+ 0] = d.x * (q[l] & 0xF) + d.y;
-        y[l+16] = d.x * (q[l] >>  4) + d.y;
+        y[l+ 0] = ggml_cuda_cast<dst_t>(d.x * (q[l] & 0xF) + d.y);
+        y[l+16] = ggml_cuda_cast<dst_t>(d.x * (q[l] >>  4) + d.y);
     }
 }
 
@@ -154,10 +154,10 @@ static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t
 
     float dall = __low2half(x[i].dm);
     float dmin = __high2half(x[i].dm);
-    y[l+ 0] = dall * (x[i].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[i].scales[is+0] >> 4);
-    y[l+32] = dall * (x[i].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[i].scales[is+2] >> 4);
-    y[l+64] = dall * (x[i].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[i].scales[is+4] >> 4);
-    y[l+96] = dall * (x[i].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[i].scales[is+6] >> 4);
+    y[l+ 0] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[i].scales[is+0] >> 4));
+    y[l+32] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[i].scales[is+2] >> 4));
+    y[l+64] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[i].scales[is+4] >> 4));
+    y[l+96] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[i].scales[is+6] >> 4));
 }
 
 template<typename dst_t>
@@ -188,7 +188,9 @@ static __global__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t
     const uint8_t * q = x[i].qs + 32*n;
     const uint8_t * hm = x[i].hmask;
 
-    for (int l = l0; l < l0+4; ++l) y[l] = dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4));
+    for (int l = l0; l < l0+4; ++l) {
+        y[l] = ggml_cuda_cast<dst_t>(dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)));
+    }
 }
 
 static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) {
@@ -226,8 +228,8 @@ static __global__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t
     get_scale_min_k4(is + 1, x[i].scales, sc, m);
     const float d2 = dall * sc; const float m2 = dmin * m;
     for (int l = 0; l < n; ++l) {
-        y[l + 0] = d1 * (q[l] & 0xF) - m1;
-        y[l +32] = d2 * (q[l] >>  4) - m2;
+        y[l + 0] = ggml_cuda_cast<dst_t>(d1 * (q[l] & 0xF) - m1);
+        y[l +32] = ggml_cuda_cast<dst_t>(d2 * (q[l] >>  4) - m2);
     }
 }
 
@@ -258,11 +260,11 @@ static __global__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t
     const float d2 = dall * sc; const float m2 = dmin * m;
 
     uint8_t   hm  = 1 << (2*il);
-    y[ 0] = d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1;
-    y[ 1] = d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1;
+    y[ 0] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1);
+    y[ 1] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1);
     hm <<= 1;
-    y[32] = d2 * ((ql[ 0] >>  4) + (qh[ 0] & hm ? 16 : 0)) - m2;
-    y[33] = d2 * ((ql[ 1] >>  4) + (qh[ 1] & hm ? 16 : 0)) - m2;
+    y[32] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 0] >>  4) + (qh[ 0] & hm ? 16 : 0)) - m2);
+    y[33] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 1] >>  4) + (qh[ 1] & hm ? 16 : 0)) - m2);
 }
 
 template<typename dst_t>
@@ -285,10 +287,10 @@ static __global__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t
     const uint8_t   qh = x[i].qh[32*ip + il];
     const int8_t  * sc = x[i].scales + is;
 
-    y[ 0] = d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32);
-    y[32] = d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32);
-    y[64] = d * sc[4] * ((int8_t)((ql[ 0]  >> 4) | (((qh >> 4) & 3) << 4)) - 32);
-    y[96] = d * sc[6] * ((int8_t)((ql[32]  >> 4) | (((qh >> 6) & 3) << 4)) - 32);
+    y[ 0] = ggml_cuda_cast<dst_t>(d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32));
+    y[32] = ggml_cuda_cast<dst_t>(d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32));
+    y[64] = ggml_cuda_cast<dst_t>(d * sc[4] * ((int8_t)((ql[ 0]  >> 4) | (((qh >> 4) & 3) << 4)) - 32));
+    y[96] = ggml_cuda_cast<dst_t>(d * sc[6] * ((int8_t)((ql[32]  >> 4) | (((qh >> 6) & 3) << 4)) - 32));
 }
 
 template<typename dst_t>
@@ -307,7 +309,9 @@ static __global__ void dequantize_block_iq2_xxs(const void * __restrict__ vx, ds
     const uint32_t aux32 = q2[2] | (q2[3] << 16);
     const float d = (float)x[i].d * (0.5f + (aux32 >> 28)) * 0.25f;
     const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
-    for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+    }
 }
 
 template<typename dst_t>
@@ -324,7 +328,9 @@ static __global__ void dequantize_block_iq2_xs(const void * __restrict__ vx, dst
     const uint8_t  * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511));
     const float d = (float)x[i].d * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
     const uint8_t signs = ksigns_iq2xs[q2[il] >> 9];
-    for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+    }
 }
 
 template<typename dst_t>
@@ -340,7 +346,9 @@ static __global__ void dequantize_block_iq2_s(const void * __restrict__ vx, dst_
     const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[i].qs[4*ib+il] | ((x[i].qh[ib] << (8-2*il)) & 0x300)));
     const float d = (float)x[i].d * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
     const uint8_t signs = x[i].qs[QK_K/8+4*ib+il];
-    for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+    }
 }
 
 template<typename dst_t>
@@ -361,8 +369,8 @@ static __global__ void dequantize_block_iq3_xxs(const void * __restrict__ vx, ds
     const float d = (float)x[i].d * (0.5f + (aux32 >> 28)) * 0.5f;
     const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
     for (int j = 0; j < 4; ++j) {
-        y[j+0] = d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
-        y[j+4] = d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f);
+        y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
+        y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
     }
 }
 
@@ -382,8 +390,8 @@ static __global__ void dequantize_block_iq3_s(const void * __restrict__ vx, dst_
     const float d = (float)x[i].d * (1 + 2*((x[i].scales[ib/2] >> 4*(ib%2)) & 0xf));
     const uint8_t signs = x[i].signs[4*ib + il];
     for (int j = 0; j < 4; ++j) {
-        y[j+0] = d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
-        y[j+4] = d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f);
+        y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
+        y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
     }
 }
 
@@ -404,7 +412,7 @@ static __global__ void dequantize_block_iq1_s(const void * __restrict__ vx, dst_
     grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
     grid32[0] &= 0x0f0f0f0f;
     for (int j = 0; j < 8; ++j) {
-        y[j] = d * (q[j] + delta);
+        y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
     }
 }
 
@@ -429,7 +437,7 @@ static __global__ void dequantize_block_iq1_m(const void * __restrict__ vx, dst_
     grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
     grid32[0] &= 0x0f0f0f0f;
     for (int j = 0; j < 8; ++j) {
-        y[j] = d * (q[j] + delta);
+        y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
     }
 }
 
@@ -446,8 +454,8 @@ static __global__ void dequantize_block_iq4_nl(const void * __restrict__ vx, dst
     const uint8_t  * q4 = x[ib].qs + 4*il;
     const float d = (float)x[ib].d;
     for (int j = 0; j < 4; ++j) {
-        y[j+ 0] = d * kvalues_iq4nl[q4[j] & 0xf];
-        y[j+16] = d * kvalues_iq4nl[q4[j] >>  4];
+        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
+        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >>  4]);
     }
 }
 
@@ -463,8 +471,8 @@ static __global__ void dequantize_block_iq4_xs(const void * __restrict__ vx, dst
     const uint8_t  * q4 = x[i].qs + 16*ib + 4*il;
     const float d = (float)x[i].d * ((((x[i].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[i].scales_h >> 2*ib) & 3) << 4)) - 32);
     for (int j = 0; j < 4; ++j) {
-        y[j+ 0] = d * kvalues_iq4nl[q4[j] & 0xf];
-        y[j+16] = d * kvalues_iq4nl[q4[j] >>  4];
+        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
+        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >>  4]);
     }
 }
 
@@ -481,8 +489,8 @@ static __global__ void dequantize_block_mxfp4(const void * __restrict__ vx, dst_
     const uint8_t  * q4 = x[ib].qs + 4*il;
     const float d = ggml_cuda_e8m0_to_fp32(x[ib].e);
     for (int j = 0; j < 4; ++j) {
-        y[j+ 0] = d * kvalues_mxfp4[q4[j] & 0xf]*0.5f;
-        y[j+16] = d * kvalues_mxfp4[q4[j] >>  4]*0.5f;
+        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] & 0xf]*0.5f);
+        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] >>  4]*0.5f);
     }
 }
 
@@ -700,6 +708,50 @@ static void convert_unary_cont_cuda(const void * vx, dst_t * y, const int64_t k,
 
 to_bf16_cuda_t ggml_get_to_bf16_cuda(ggml_type type) {
     switch (type) {
+        case GGML_TYPE_Q1_0:
+            return dequantize_block_cont_cuda<QK1_0, QR1_0, dequantize_q1_0>;
+        case GGML_TYPE_Q4_0:
+            return dequantize_row_q4_0_cuda;
+        case GGML_TYPE_Q4_1:
+            return dequantize_row_q4_1_cuda;
+        case GGML_TYPE_Q5_0:
+            return dequantize_block_cont_cuda<QK5_0, QR5_0, dequantize_q5_0>;
+        case GGML_TYPE_Q5_1:
+            return dequantize_block_cont_cuda<QK5_1, QR5_1, dequantize_q5_1>;
+        case GGML_TYPE_Q8_0:
+            return dequantize_block_cont_cuda<QK8_0, QR8_0, dequantize_q8_0>;
+        case GGML_TYPE_Q2_K:
+            return dequantize_row_q2_K_cuda;
+        case GGML_TYPE_Q3_K:
+            return dequantize_row_q3_K_cuda;
+        case GGML_TYPE_Q4_K:
+            return dequantize_row_q4_K_cuda;
+        case GGML_TYPE_Q5_K:
+            return dequantize_row_q5_K_cuda;
+        case GGML_TYPE_Q6_K:
+            return dequantize_row_q6_K_cuda;
+        case GGML_TYPE_IQ2_XXS:
+            return dequantize_row_iq2_xxs_cuda;
+        case GGML_TYPE_IQ2_XS:
+            return dequantize_row_iq2_xs_cuda;
+        case GGML_TYPE_IQ2_S:
+            return dequantize_row_iq2_s_cuda;
+        case GGML_TYPE_IQ3_XXS:
+            return dequantize_row_iq3_xxs_cuda;
+        case GGML_TYPE_IQ1_S:
+            return dequantize_row_iq1_s_cuda;
+        case GGML_TYPE_IQ1_M:
+            return dequantize_row_iq1_m_cuda;
+        case GGML_TYPE_IQ4_NL:
+            return dequantize_row_iq4_nl_cuda;
+        case GGML_TYPE_IQ4_XS:
+            return dequantize_row_iq4_xs_cuda;
+        case GGML_TYPE_IQ3_S:
+            return dequantize_row_iq3_s_cuda;
+        case GGML_TYPE_MXFP4:
+            return dequantize_row_mxfp4_cuda;
+        case GGML_TYPE_NVFP4:
+            return dequantize_row_nvfp4_cuda;
         case GGML_TYPE_F32:
             return convert_unary_cont_cuda<float>;
         case GGML_TYPE_F16:

ggml/src/ggml-cuda/fattn.cu

@@ -337,6 +337,26 @@ enum best_fattn_kernel {
     BEST_FATTN_KERNEL_MMA_F16  = 400,
 };
 
+static bool ggml_cuda_fattn_kv_type_supported(ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_F32:
+        case GGML_TYPE_F16:
+            return true;
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q5_1:
+#ifndef GGML_CUDA_FA_ALL_QUANTS
+            return false;
+#endif // GGML_CUDA_FA_ALL_QUANTS
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q8_0:
+        case GGML_TYPE_BF16:
+            return true;
+        default:
+            return false;
+    }
+}
+
 static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const ggml_tensor * dst) {
 #ifndef FLASH_ATTN_AVAILABLE
     GGML_UNUSED(device); GGML_UNUSED(dst);
@@ -427,22 +447,8 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
     }
 #endif // GGML_CUDA_FA_ALL_QUANTS
 
-    switch (K->type) {
-        case GGML_TYPE_F32:
-        case GGML_TYPE_F16:
-            break;
-        case GGML_TYPE_Q4_1:
-        case GGML_TYPE_Q5_0:
-        case GGML_TYPE_Q5_1:
-#ifndef GGML_CUDA_FA_ALL_QUANTS
-            return BEST_FATTN_KERNEL_NONE;
-#endif // GGML_CUDA_FA_ALL_QUANTS
-        case GGML_TYPE_Q4_0:
-        case GGML_TYPE_Q8_0:
-        case GGML_TYPE_BF16:
-            break;
-        default:
-            return BEST_FATTN_KERNEL_NONE;
+    if (!ggml_cuda_fattn_kv_type_supported(K->type) || !ggml_cuda_fattn_kv_type_supported(V->type)) {
+        return BEST_FATTN_KERNEL_NONE;
     }
 
     if (mask && mask->ne[2] != 1) {

ggml/src/ggml-cuda/mmvq.cu

@@ -278,6 +278,9 @@ int get_mmvq_mmid_max_batch(ggml_type type, int cc) {
 }
 
 bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) {
+    if (!ggml_is_quantized(type)) {
+        return false;
+    }
     if (GGML_CUDA_CC_IS_CDNA(cc)) {
         if (GGML_CUDA_CC_IS_CDNA1(cc)) {
             switch (type) {
@@ -518,9 +521,13 @@ static __global__ void mul_mat_vec_q(
     bool use_gate = false;
     bool use_bias = false;
     bool use_gate_bias = false;
+    bool use_scale = false;
+    bool use_gate_scale = false;
     [[maybe_unused]] const void * vgate = nullptr;
     const float * x_bias = nullptr;
     const float * gate_bias = nullptr;
+    const float * x_scale = nullptr;
+    const float * gate_scale = nullptr;
     ggml_glu_op active_glu;
 
     if constexpr (has_fusion) {
@@ -531,34 +538,47 @@ static __global__ void mul_mat_vec_q(
         x_bias        = (const float *) fusion.x_bias;
         gate_bias     = (const float *) fusion.gate_bias;
         active_glu    = fusion.glu_op;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            use_scale      = fusion.x_scale    != nullptr;
+            use_gate_scale = fusion.gate_scale != nullptr && use_gate;
+            x_scale        = (const float *) fusion.x_scale;
+            gate_scale     = (const float *) fusion.gate_scale;
+        }
     }
 
 
     [[maybe_unused]] float x_biases[ncols_dst]    = { 0.0f };
     [[maybe_unused]] float gate_biases[ncols_dst] = { 0.0f };
+    [[maybe_unused]] float x_scales;
+    [[maybe_unused]] float gate_scales;
     if constexpr (has_fusion) {
+        // 1. Hide latency by prefetching bias, gates and scales here
+        // 2. load only on threads that won't die after partial sum calculation
         const uint32_t channel_bias = ids ? channel_x : channel_dst;
-        if (use_bias) {
-            x_bias = x_bias + sample_dst*stride_sample_dst + channel_bias*stride_channel_dst + row0;
-            // 1. Hide latency by prefetching bias and gate here
-            // 2. load only on threads that won't die after partial sum calculation
-            if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 &&
-                (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
+        if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 &&
+            (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
+            if (use_bias) {
+                x_bias = x_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0;
 #pragma unroll
                 for (int j = 0; j < ncols_dst; ++j) {
                     x_biases[j] = x_bias[j * stride_col_dst + threadIdx.x];
                 }
             }
-        }
-        if (use_gate_bias) {
-            gate_bias = gate_bias + sample_dst*stride_sample_dst + channel_bias*stride_channel_dst + row0;
-            if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 &&
-                (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
+            if (use_gate_bias) {
+                gate_bias = gate_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0;
 #pragma unroll
                 for (int j = 0; j < ncols_dst; ++j) {
                     gate_biases[j] = gate_bias[j * stride_col_dst + threadIdx.x];
                 }
             }
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                if (use_scale) {
+                    x_scales = x_scale[ids ? channel_x : 0];
+                }
+                if (use_gate_scale) {
+                    gate_scales = gate_scale[ids ? channel_x : 0];
+                }
+            }
         }
     }
 
@@ -640,11 +660,21 @@ static __global__ void mul_mat_vec_q(
         if (threadIdx.x < rows_per_cuda_block && (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
             float result = tmp[j][threadIdx.x];
             if constexpr (has_fusion) {
+                if constexpr (type == GGML_TYPE_NVFP4) {
+                    if (use_scale) {
+                        result *= x_scales;
+                    }
+                }
                 if (use_bias) {
                     result += x_biases[j];
                 }
                 if (use_gate) {
                     float gate_value = tmp_gate[j][threadIdx.x];
+                    if constexpr (type == GGML_TYPE_NVFP4) {
+                        if (use_gate_scale) {
+                            gate_value *= gate_scales;
+                        }
+                    }
                     if (use_gate_bias) {
                         gate_value += gate_biases[j];
                     }
@@ -670,7 +700,10 @@ static __global__ void mul_mat_vec_q(
     }
 
     if constexpr (!has_fusion) {
-        GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, active_glu, gate_bias, x_bias, tmp_gate);
+        GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, use_scale, use_gate_scale, active_glu, gate_bias, x_bias, x_scale, gate_scale, tmp_gate);
+    }
+    if constexpr (type != GGML_TYPE_NVFP4) {
+        GGML_UNUSED_VARS(use_scale, use_gate_scale, x_scale, gate_scale, x_scales, gate_scales);
     }
 }
 
@@ -766,7 +799,8 @@ static void mul_mat_vec_q_switch_fusion(
         const dim3 & block_nums, const dim3 & block_dims, const int nbytes_shared,
         const uint32_t ids_stride, cudaStream_t stream) {
 
-    const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr;
+    const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr ||
+                            fusion.x_scale != nullptr || fusion.gate_scale != nullptr;
     if constexpr (c_ncols_dst == 1) {
         if (has_fusion) {
             const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream);
@@ -831,7 +865,6 @@ static void mul_mat_vec_q_switch_ncols_dst(
     const int warp_size = ggml_cuda_info().devices[device].warp_size;
     const mmvq_parameter_table_id table_id  = get_device_table_id(cc);
 
-    const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr;
     const bool has_ids = ids != nullptr;
 
     const auto should_use_small_k = [&](int c_ncols_dst) {
@@ -970,8 +1003,6 @@ static void mul_mat_vec_q_switch_ncols_dst(
             GGML_ABORT("fatal error");
             break;
     }
-
-    GGML_UNUSED(has_fusion);
 }
 static void mul_mat_vec_q_switch_type(
         const void * vx, const ggml_type type_x, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
@@ -1151,6 +1182,9 @@ void ggml_cuda_mul_mat_vec_q(
     if (fusion) {
         GGML_ASSERT( !ids || dst->ne[2] == 1);
         GGML_ASSERT(  ids || dst->ne[1] == 1);
+        // Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is
+        // non-negligible for some models such as gpt-oss-20b
+        GGML_ASSERT((fusion->x_scale == nullptr && fusion->gate_scale == nullptr) || src0->type == GGML_TYPE_NVFP4);
 
         if (fusion->x_bias) {
             GGML_ASSERT(fusion->x_bias->type == GGML_TYPE_F32);
@@ -1168,6 +1202,18 @@ void ggml_cuda_mul_mat_vec_q(
             GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]);
             fusion_local.gate_bias = fusion->gate_bias->data;
         }
+        if (fusion->x_scale) {
+            GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32);
+            GGML_ASSERT(ggml_is_contiguous(fusion->x_scale));
+            GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1));
+            fusion_local.x_scale = fusion->x_scale->data;
+        }
+        if (fusion->gate_scale) {
+            GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32);
+            GGML_ASSERT(ggml_is_contiguous(fusion->gate_scale));
+            GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1));
+            fusion_local.gate_scale = fusion->gate_scale->data;
+        }
         fusion_local.glu_op = fusion->glu_op;
     }
 

ggml/src/ggml-cuda/set-rows.cu

@@ -322,17 +322,77 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s
     }
 }
 
+template<>
+void set_rows_cuda<half, int32_t>(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    const half    * src0_d = (const half *)src0->data;
+    const int32_t * src1_d = (const int32_t *)src1->data;
+
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    cudaStream_t stream = ctx.stream();
+
+
+    if (dst->type == GGML_TYPE_F16) {
+        set_rows_cuda(
+            src0_d, src1_d, (half*)dst->data,
+            ne00, ne01, ne02, ne03,
+            ne10, ne11, ne12, ne13,
+            nb01, nb02, nb03,
+            nb10, nb11, nb12,
+            nb1, nb2, nb3,
+            stream
+        );
+    } else {
+        GGML_ABORT("unsupported type %s", ggml_type_name(dst->type));
+    }
+}
+
+template<>
+void set_rows_cuda<half, int64_t>(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    const half    * src0_d = (const half *)src0->data;
+    const int64_t * src1_d = (const int64_t *)src1->data;
+
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    cudaStream_t stream = ctx.stream();
+
+
+    if (dst->type == GGML_TYPE_F16) {
+        set_rows_cuda(
+            src0_d, src1_d, (half*)dst->data,
+            ne00, ne01, ne02, ne03,
+            ne10, ne11, ne12, ne13,
+            nb01, nb02, nb03,
+            nb10, nb11, nb12,
+            nb1, nb2, nb3,
+            stream
+        );
+    } else {
+        GGML_ABORT("unsupported type %s", ggml_type_name(dst->type));
+    }
+}
+
 
 void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const ggml_tensor * src0 = dst->src[0];
     const ggml_tensor * src1 = dst->src[1];
 
-    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src0->type == GGML_TYPE_F32 || (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16));
     GGML_ASSERT(src1->type == GGML_TYPE_I64 || src1->type == GGML_TYPE_I32);
 
-    if (src1->type == GGML_TYPE_I64) {
-        set_rows_cuda<float, int64_t>(ctx, src0, src1, dst);
+    if (src0->type == GGML_TYPE_F32) {
+        if (src1->type == GGML_TYPE_I64) {
+            set_rows_cuda<float, int64_t>(ctx, src0, src1, dst);
+        } else {
+            set_rows_cuda<float, int32_t>(ctx, src0, src1, dst);
+        }
+    } else if (src0->type == GGML_TYPE_F16) {
+        if (src1->type == GGML_TYPE_I64) {
+            set_rows_cuda<half, int64_t>(ctx, src0, src1, dst);
+        } else {
+            set_rows_cuda<half, int32_t>(ctx, src0, src1, dst);
+        }
     } else {
-        set_rows_cuda<float, int32_t>(ctx, src0, src1, dst);
+        GGML_ABORT("unsupported type %s", ggml_type_name(src0->type));
     }
 }

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

@@ -2028,10 +2028,10 @@ static bool ggml_hexagon_precompute_flash_attn_params(
     kparams->u.hvx.size_v_row_padded = size_v_row_padded;
     kparams->u.hvx.src0_div21 = init_fastdiv_values(q->ne[2] * q->ne[1]);
     kparams->u.hvx.src0_div1 = init_fastdiv_values(q->ne[1]);
-    kparams->u.hvx.broadcast_rk2 = init_fastdiv_values(q->ne[2]/k->ne[2]);
-    kparams->u.hvx.broadcast_rk3 = init_fastdiv_values(q->ne[3]/k->ne[3]);
-    kparams->u.hvx.broadcast_rv2 = init_fastdiv_values(q->ne[2]/v->ne[2]);
-    kparams->u.hvx.broadcast_rv3 = init_fastdiv_values(q->ne[3]/v->ne[3]);
+    kparams->broadcast_rk2 = init_fastdiv_values(q->ne[2]/k->ne[2]);
+    kparams->broadcast_rk3 = init_fastdiv_values(q->ne[3]/k->ne[3]);
+    kparams->broadcast_rv2 = init_fastdiv_values(q->ne[2]/v->ne[2]);
+    kparams->broadcast_rv3 = init_fastdiv_values(q->ne[3]/v->ne[3]);
     if (mask) {
         kparams->src3_div2 = init_fastdiv_values(mask->ne[2]);
         kparams->src3_div3 = init_fastdiv_values(mask->ne[3]);
@@ -2385,31 +2385,30 @@ static void ggml_hexagon_precompute_hvx_mm_params(
             kparams->kernel_type   = (src1_nrows < (int) sess->n_threads) ? HTP_MM_KERNEL_HVX_QUANT_BLOCK : HTP_MM_KERNEL_HVX_QUANT_ROW;
             kparams->src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
 
-            size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0;
+            struct htp_mm_hvx_vtcm_layout L;
             uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
             uint32_t best_n_prefetch = 2;
-            size_t total_size = 0;
             for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
-                total_size = htp_mm_hvx_id_get_vtcm_sizes(
-                    wtype, ne10, src1_nrows, sess->n_threads, src0->nb[1], d,
-                    &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+                htp_mm_hvx_vtcm_layout_build(
+                    &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+                    0, src0->nb[1], 0, d, true, false, false
                 );
-                if (total_size <= vtcm_budget) {
+                if (L.total_bytes <= vtcm_budget) {
                     best_n_prefetch = d;
                     break;
                 }
             }
-            if (best_n_prefetch == 2 && total_size > vtcm_budget) {
-                total_size = htp_mm_hvx_id_get_vtcm_sizes(
-                    wtype, ne10, src1_nrows, sess->n_threads, src0->nb[1], 2,
-                    &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+            if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) {
+                htp_mm_hvx_vtcm_layout_build(
+                    &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+                    0, src0->nb[1], 0, 2, true, false, false
                 );
             }
             kparams->n_prefetch = best_n_prefetch;
-            kparams->vtcm_size      = total_size;
-            kparams->vtcm_src0_size = vtcm_src0_size;
-            kparams->vtcm_src1_size = vtcm_src1_size;
-            kparams->vtcm_dst_size  = vtcm_dst_size;
+            kparams->vtcm_size      = L.total_bytes;
+            kparams->vtcm_src0_size = L.src0_bytes;
+            kparams->vtcm_src1_size = L.src1_bytes;
+            kparams->vtcm_dst_size  = L.dst_bytes;
         } else {
             bool try_tiled = (k_align && opt_mm_select >= 2);
             if (try_tiled) {
@@ -2420,37 +2419,36 @@ static void ggml_hexagon_precompute_hvx_mm_params(
                     kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
                 }
 
+                struct htp_mm_hvx_vtcm_layout L;
                 uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
                 uint32_t best_n_prefetch = 2;
-                size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0;
-                size_t total_size = 0;
                 for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
-                    total_size = htp_mm_hvx_get_vtcm_sizes(
-                        kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                        dst->nb[1], src0->nb[1], src1->nb[1], d, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+                    htp_mm_hvx_vtcm_layout_build(
+                        &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+                        dst->nb[1], src0->nb[1], src1->nb[1], d, false, false, false
                     );
-                    if (total_size <= vtcm_budget) {
+                    if (L.total_bytes <= vtcm_budget) {
                         best_n_prefetch = d;
                         break;
                     }
                 }
-                if (best_n_prefetch == 2 && total_size > vtcm_budget) {
-                    total_size = htp_mm_hvx_get_vtcm_sizes(
-                        kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                        dst->nb[1], src0->nb[1], src1->nb[1], 2, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+                if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) {
+                    htp_mm_hvx_vtcm_layout_build(
+                        &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+                        dst->nb[1], src0->nb[1], src1->nb[1], 2, false, false, false
                     );
                 }
 
                 kparams->n_prefetch = best_n_prefetch;
 
-                if (total_size <= vtcm_budget) {
-                    kparams->vtcm_size = total_size;
-                    kparams->vtcm_src0_size = vtcm_src0_size;
-                    kparams->vtcm_src1_size = vtcm_src1_size;
-                    kparams->vtcm_dst_size = vtcm_dst_size;
+                if (L.total_bytes <= vtcm_budget) {
+                    kparams->vtcm_size = L.total_bytes;
+                    kparams->vtcm_src0_size = L.src0_bytes;
+                    kparams->vtcm_src1_size = L.src1_bytes;
+                    kparams->vtcm_dst_size = L.dst_bytes;
                     goto done_quant;
                 }
-                HEX_VERBOSE("ggml-hex: %s HVX tiled path VTCM size needed (%zu) > budget (%zu), falling back to HVX flat\n", sess->name.c_str(), total_size, vtcm_budget);
+                HEX_VERBOSE("ggml-hex: %s HVX tiled path VTCM size needed (%zu) > budget (%zu), falling back to HVX flat\n", sess->name.c_str(), L.total_bytes, vtcm_budget);
             }
 
             // Flat HVX fallback
@@ -2458,17 +2456,17 @@ static void ggml_hexagon_precompute_hvx_mm_params(
                 kparams->src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
                 kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT;
 
-                size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0;
-                size_t total_size = htp_mm_hvx_get_vtcm_sizes(
-                    kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                    dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+                struct htp_mm_hvx_vtcm_layout L;
+                htp_mm_hvx_vtcm_layout_build(
+                    &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+                    dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
                 );
 
                 kparams->n_prefetch = 16;
-                kparams->vtcm_size = total_size;
-                kparams->vtcm_src0_size = vtcm_src0_size;
-                kparams->vtcm_src1_size = vtcm_src1_size;
-                kparams->vtcm_dst_size = vtcm_dst_size;
+                kparams->vtcm_size = L.total_bytes;
+                kparams->vtcm_src0_size = L.src0_bytes;
+                kparams->vtcm_src1_size = L.src1_bytes;
+                kparams->vtcm_dst_size = L.dst_bytes;
             }
         }
 
@@ -2478,19 +2476,19 @@ static void ggml_hexagon_precompute_hvx_mm_params(
         const bool is_batched  = (ne02 > 1) || (ne03 > 1);
         const bool is_permuted = ggml_is_permuted(src0) || ggml_is_permuted(src1);
 
-        size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0;
-        size_t vtcm_size = htp_mm_hvx_get_vtcm_sizes(
-            HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
-            dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+        struct htp_mm_hvx_vtcm_layout L;
+        htp_mm_hvx_vtcm_layout_build(
+            &L, HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
+            dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
         );
 
-        if (!is_batched && !is_permuted && vtcm_size <= vtcm_budget) {
+        if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) {
             kparams->kernel_type = HTP_MM_KERNEL_HVX_F16_F16_VTCM;
             kparams->src1_row_size = hex_round_up(ne10 * 2, 128);
-            kparams->vtcm_size = vtcm_size;
-            kparams->vtcm_src0_size = vtcm_src0_size;
-            kparams->vtcm_src1_size = vtcm_src1_size;
-            kparams->vtcm_dst_size = vtcm_dst_size;
+            kparams->vtcm_size = L.total_bytes;
+            kparams->vtcm_src0_size = L.src0_bytes;
+            kparams->vtcm_src1_size = L.src1_bytes;
+            kparams->vtcm_dst_size = L.dst_bytes;
             kparams->n_prefetch = 16;
         } else {
             if (src1->type == GGML_TYPE_F32) {
@@ -2499,14 +2497,14 @@ static void ggml_hexagon_precompute_hvx_mm_params(
                 kparams->kernel_type = HTP_MM_KERNEL_HVX_F16_F16_DDR;
             }
             kparams->src1_row_size = src1->nb[1];
-            size_t ddr_size = htp_mm_hvx_get_vtcm_sizes(
-                kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+            htp_mm_hvx_vtcm_layout_build(
+                &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+                dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
             );
-            kparams->vtcm_size = ddr_size;
-            kparams->vtcm_src0_size = vtcm_src0_size;
-            kparams->vtcm_src1_size = vtcm_src1_size;
-            kparams->vtcm_dst_size = vtcm_dst_size;
+            kparams->vtcm_size = L.total_bytes;
+            kparams->vtcm_src0_size = L.src0_bytes;
+            kparams->vtcm_src1_size = L.src1_bytes;
+            kparams->vtcm_dst_size = L.dst_bytes;
             kparams->n_prefetch = 16;
         }
     } else {
@@ -2514,31 +2512,31 @@ static void ggml_hexagon_precompute_hvx_mm_params(
         const bool is_batched  = (ne02 > 1) || (ne03 > 1);
         const bool is_permuted = ggml_is_permuted(src0) || ggml_is_permuted(src1);
 
-        size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0;
-        size_t vtcm_size = htp_mm_hvx_get_vtcm_sizes(
-            HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
-            dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+        struct htp_mm_hvx_vtcm_layout L;
+        htp_mm_hvx_vtcm_layout_build(
+            &L, HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
+            dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
         );
 
-        if (!is_batched && !is_permuted && vtcm_size <= vtcm_budget) {
+        if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) {
             kparams->kernel_type = HTP_MM_KERNEL_HVX_F32_F32_VTCM;
             kparams->src1_row_size = hex_round_up(ne10 * 4, 128);
-            kparams->vtcm_size = vtcm_size;
-            kparams->vtcm_src0_size = vtcm_src0_size;
-            kparams->vtcm_src1_size = vtcm_src1_size;
-            kparams->vtcm_dst_size = vtcm_dst_size;
+            kparams->vtcm_size = L.total_bytes;
+            kparams->vtcm_src0_size = L.src0_bytes;
+            kparams->vtcm_src1_size = L.src1_bytes;
+            kparams->vtcm_dst_size = L.dst_bytes;
             kparams->n_prefetch = 16;
         } else {
             kparams->kernel_type = HTP_MM_KERNEL_HVX_F32_F32_DDR;
             kparams->src1_row_size = src1->nb[1];
-            size_t ddr_size = htp_mm_hvx_get_vtcm_sizes(
-                kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
-                dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+            htp_mm_hvx_vtcm_layout_build(
+                &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+                dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
             );
-            kparams->vtcm_size = ddr_size;
-            kparams->vtcm_src0_size = vtcm_src0_size;
-            kparams->vtcm_src1_size = vtcm_src1_size;
-            kparams->vtcm_dst_size = vtcm_dst_size;
+            kparams->vtcm_size = L.total_bytes;
+            kparams->vtcm_src0_size = L.src0_bytes;
+            kparams->vtcm_src1_size = L.src1_bytes;
+            kparams->vtcm_dst_size = L.dst_bytes;
             kparams->n_prefetch = 16;
         }
     }
@@ -2608,80 +2606,57 @@ static void ggml_hexagon_precompute_fused_qkv_params(
     const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3];
     const size_t src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
     const size_t src0_row_size = src0->nb[1];
-    const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);
 
-    size_t src0_sz_per_thread = 0;
-    size_t src2_sz_per_thread = 0;
-    size_t src3_sz_per_thread = 0;
     uint32_t best_n_prefetch = 16;
 
-    size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * sess->n_threads;
-
     if (is_repack) {
-        uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
-        uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32;
-        uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
-        size_t src1_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128);
-        size_t src1_sz = src1_sz_per_thread;
-
         const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
         best_n_prefetch = 2;
         for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
-            size_t repacked_vtcm_size = hex_round_up(d * tile_row_size, 128);
-            size_t src0_sz = repacked_vtcm_size * sess->n_threads;
-            size_t src2_sz = hex_round_up(d * tile_row_size, 128) * sess->n_threads;
-            size_t src3_sz = hex_round_up(d * tile_row_size, 128) * sess->n_threads;
-            size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + src3_sz + quant_scratch_size;
-
-            if (tiled_vtcm_size <= sess->vtcm_size) {
+            struct htp_mm_hvx_vtcm_layout L;
+            htp_mm_hvx_vtcm_layout_build(
+                &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
+                0, src0_row_size, src1_row_size, d, false, true, false
+            );
+            if (L.total_bytes <= sess->vtcm_size) {
                 best_n_prefetch = d;
-                src0_sz_per_thread = repacked_vtcm_size;
-                src2_sz_per_thread = hex_round_up(d * tile_row_size, 128);
-                src3_sz_per_thread = hex_round_up(d * tile_row_size, 128);
                 break;
             }
         }
-        if (best_n_prefetch == 2 && src0_sz_per_thread == 0) {
-            size_t repacked_vtcm_size = hex_round_up(2 * tile_row_size, 128);
-            src0_sz_per_thread = repacked_vtcm_size;
-            src2_sz_per_thread = hex_round_up(2 * tile_row_size, 128);
-            src3_sz_per_thread = hex_round_up(2 * tile_row_size, 128);
-        }
-    } else {
-        best_n_prefetch = 16;
-        src0_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128);
-        src2_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128);
-        src3_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128);
     }
 
-    size_t src1_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128);
-
-    size_t src0_sz = src0_sz_per_thread * sess->n_threads;
-    size_t src1_sz = src1_sz_per_thread;
-    size_t src2_sz = src2_sz_per_thread * sess->n_threads;
-    size_t src3_sz = src3_sz_per_thread * sess->n_threads;
-
-    size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + src3_sz + quant_scratch_size;
+    struct htp_mm_hvx_vtcm_layout L;
     bool try_tiled = (opt_mm_select >= 2);
-    if (try_tiled && tiled_vtcm_size <= sess->vtcm_size) {
+
+    // Test tiled first
+    htp_mm_hvx_vtcm_layout_build(
+        &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
+        0, src0_row_size, src1_row_size, best_n_prefetch, false, true, false
+    );
+
+    if (try_tiled && L.total_bytes <= sess->vtcm_size) {
         kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
-        kparams->vtcm_src0_size = src0_sz;
-        kparams->vtcm_src1_size = src1_sz;
-        kparams->vtcm_src2_size = src2_sz;
-        kparams->vtcm_src3_size = src3_sz;
-        kparams->vtcm_dst_size  = quant_scratch_size;
-        kparams->vtcm_size      = tiled_vtcm_size;
+        kparams->vtcm_src0_size = L.src0_bytes;
+        kparams->vtcm_src1_size = L.src1_bytes;
+        kparams->vtcm_src2_size = L.src2_bytes;
+        kparams->vtcm_src3_size = L.src3_bytes;
+        kparams->vtcm_dst_size  = L.dst_bytes;
+        kparams->vtcm_size      = L.total_bytes;
         kparams->n_prefetch     = best_n_prefetch;
     } else {
         kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT;
         size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
-        size_t flat_src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128);
-        kparams->vtcm_src0_size = src0_sz;
-        kparams->vtcm_src1_size = flat_src1_sz;
-        kparams->vtcm_src2_size = src2_sz;
-        kparams->vtcm_src3_size = src3_sz;
-        kparams->vtcm_dst_size  = quant_scratch_size;
-        kparams->vtcm_size      = src0_sz + flat_src1_sz + src2_sz + src3_sz + quant_scratch_size;
+
+        htp_mm_hvx_vtcm_layout_build(
+            &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
+            0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, true, false
+        );
+        kparams->vtcm_src0_size = L.src0_bytes;
+        kparams->vtcm_src1_size = L.src1_bytes;
+        kparams->vtcm_src2_size = L.src2_bytes;
+        kparams->vtcm_src3_size = L.src3_bytes;
+        kparams->vtcm_dst_size  = L.dst_bytes;
+        kparams->vtcm_size      = L.total_bytes;
         kparams->n_prefetch     = best_n_prefetch;
     }
 }
@@ -2701,72 +2676,55 @@ static void ggml_hexagon_precompute_fused_ffn_params(
     const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3];
     const size_t src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
     const size_t src0_row_size = src0->nb[1];
-    const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);
 
-    size_t src0_sz_per_thread = 0;
-    size_t src2_sz_per_thread = 0;
     uint32_t best_n_prefetch = 16;
 
-    size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * sess->n_threads;
-
     if (is_repack) {
-        uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
-        uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32;
-        uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
-        size_t src1_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128);
-        size_t src1_sz = src1_sz_per_thread;
-
         const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
         best_n_prefetch = 2;
         for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
-            size_t repacked_vtcm_size = hex_round_up(d * tile_row_size, 128);
-            size_t src0_sz = repacked_vtcm_size * sess->n_threads;
-            size_t src2_sz = hex_round_up(d * tile_row_size, 128) * sess->n_threads;
-            size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + quant_scratch_size;
-
-            if (tiled_vtcm_size <= sess->vtcm_size) {
+            struct htp_mm_hvx_vtcm_layout L;
+            htp_mm_hvx_vtcm_layout_build(
+                &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
+                0, src0_row_size, src1_row_size, d, false, false, true
+            );
+            if (L.total_bytes <= sess->vtcm_size) {
                 best_n_prefetch = d;
-                src0_sz_per_thread = repacked_vtcm_size;
-                src2_sz_per_thread = hex_round_up(d * tile_row_size, 128);
                 break;
             }
         }
-        if (best_n_prefetch == 2 && src0_sz_per_thread == 0) {
-            size_t repacked_vtcm_size = hex_round_up(2 * tile_row_size, 128);
-            src0_sz_per_thread = repacked_vtcm_size;
-            src2_sz_per_thread = hex_round_up(2 * tile_row_size, 128);
-        }
-    } else {
-        best_n_prefetch = 16;
-        src0_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128);
-        src2_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128);
     }
 
-    size_t src1_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128);
-
-    size_t src0_sz = src0_sz_per_thread * sess->n_threads;
-    size_t src1_sz = src1_sz_per_thread;
-    size_t src2_sz = src2_sz_per_thread * sess->n_threads;
-
-    size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + quant_scratch_size;
+    struct htp_mm_hvx_vtcm_layout L;
     bool try_tiled = (opt_mm_select >= 2);
-    if (try_tiled && tiled_vtcm_size <= sess->vtcm_size) {
+
+    // Test tiled first
+    htp_mm_hvx_vtcm_layout_build(
+        &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
+        0, src0_row_size, src1_row_size, best_n_prefetch, false, false, true
+    );
+
+    if (try_tiled && L.total_bytes <= sess->vtcm_size) {
         kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
-        kparams->vtcm_src0_size = src0_sz;
-        kparams->vtcm_src1_size = src1_sz;
-        kparams->vtcm_src2_size = src2_sz;
-        kparams->vtcm_dst_size  = quant_scratch_size;
-        kparams->vtcm_size      = tiled_vtcm_size;
+        kparams->vtcm_src0_size = L.src0_bytes;
+        kparams->vtcm_src1_size = L.src1_bytes;
+        kparams->vtcm_src2_size = L.src2_bytes;
+        kparams->vtcm_dst_size  = L.dst_bytes;
+        kparams->vtcm_size      = L.total_bytes;
         kparams->n_prefetch     = best_n_prefetch;
     } else {
         kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT;
         size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
-        size_t flat_src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128);
-        kparams->vtcm_src0_size = src0_sz;
-        kparams->vtcm_src1_size = flat_src1_sz;
-        kparams->vtcm_src2_size = src2_sz;
-        kparams->vtcm_dst_size  = quant_scratch_size;
-        kparams->vtcm_size      = src0_sz + flat_src1_sz + src2_sz + quant_scratch_size;
+
+        htp_mm_hvx_vtcm_layout_build(
+            &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
+            0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, false, true
+        );
+        kparams->vtcm_src0_size = L.src0_bytes;
+        kparams->vtcm_src1_size = L.src1_bytes;
+        kparams->vtcm_src2_size = L.src2_bytes;
+        kparams->vtcm_dst_size  = L.dst_bytes;
+        kparams->vtcm_size      = L.total_bytes;
         kparams->n_prefetch     = best_n_prefetch;
     }
 }

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

@@ -20,6 +20,7 @@ add_library(${HTP_LIB} SHARED
     worker-pool.c
     hex-dma.c
     hmx-queue.c
+    gated-delta-net-ops.c
     binary-ops.c
     unary-ops.c
     sum-rows-ops.c
@@ -37,10 +38,9 @@ add_library(${HTP_LIB} SHARED
     concat-ops.c
     diag-ops.c
     solve-tri-ops.c
-    gated-delta-net-ops.c
     pad-ops.c
-    matmul-ops.c
     flash-attn-ops.c
+    matmul-ops.c
 )
 
 target_compile_definitions(${HTP_LIB} PRIVATE

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

@@ -4,7 +4,7 @@
 #include "hexagon_protos.h"
 #include "hvx_hexagon_protos.h"
 #include "hex-dma.h"
-#include "vtcm-utils.h"
+#include "htp-vtcm.h"
 #include "hvx-utils.h"
 #include "hex-fastdiv.h"
 #include <string.h>

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

@@ -8,6 +8,7 @@
 #include <HAP_perf.h>
 #include <math.h>
 #include <stdbool.h>
+#include <stdatomic.h>
 #include <stddef.h>
 #include <stdint.h>
 #include <string.h>
@@ -22,7 +23,7 @@
 #include "hvx-copy.h"
 #include "hvx-reduce.h"
 #include "hvx-flash-attn.h"
-#include "vtcm-utils.h"
+#include "htp-vtcm.h"
 #include "worker-pool.h"
 
 #define GGML_COMMON_DECL_C
@@ -142,6 +143,10 @@ struct hmx_fa_context {
     __fp16 *     vtcm_slopes;          // ALiBi slopes [g_br]
     size_t       row_buf_stride;       // HVX vectors per row buffer (Bc/64)
     size_t       mask_buf_row_stride;  // elements (__fp16) per row in mask buffer
+    size_t       q_tile_bytes;
+    size_t       o_tile_bytes;
+    size_t       col_vec_bytes;
+    size_t       d_tile_bytes;
     bool         mask_broadcast;       // true when mask->ne[2] == 1 (head-independent, single 2D DMA)
     dma_cache    m_cache;
 };
@@ -463,7 +468,7 @@ typedef struct {
     struct hmx_fa_context * factx;
     uint32_t                kv_rows;
     size_t                  src_stride;
-    size_t                  buf_idx;
+    void *                  curr_k;
     uint32_t                kv_start;
     uint32_t                rows_per_t;
 } fa_k_int_args_t;
@@ -483,19 +488,19 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
 
     struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
     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, factx->vtcm_k_fp16[args->buf_idx], total_rows, factx->DK,
+    hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles, (const __fp16 *) args->curr_k, total_rows, factx->DK,
                              args->src_stride, start, end);
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
 }
 
-static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, size_t buf_idx, uint32_t kv_start) {
+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;
     uint32_t n = 1;
     if (factx->n_threads > 1 && kv_rows >= factx->n_threads * 2) {
         n = factx->n_threads;
     }
     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, buf_idx, kv_start, rows_per_t };
+    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);
     } else {
@@ -507,7 +512,8 @@ typedef struct {
     struct hmx_fa_context * factx;
     uint32_t                kv_rows;
     size_t                  src_stride;
-    size_t                  buf_idx;
+    void *                  v_src;
+    void *                  v_tiles_dst;
     size_t                  n_col_tiles;
     uint32_t                kv_start;
     uint32_t                rows_per_t;
@@ -526,11 +532,11 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data)
         return;
     }
 
-    __fp16 * v_tiles_dest = factx->pipeline ? factx->vtcm_v_tiles[args->buf_idx] : factx->vtcm_v_tiles[0];
+    __fp16 * v_tiles_dst = (__fp16 *) args->v_tiles_dst;
 
     struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
     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_dest, factx->vtcm_v_fp16[args->buf_idx], total_rows, factx->DV,
+    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);
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
 }
@@ -538,7 +544,8 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data)
 static void fa_phase_v_interleave(struct hmx_fa_context * factx,
                                   uint32_t                kv_rows,
                                   size_t                  src_stride,
-                                  size_t                  buf_idx,
+                                  void *                  v_src,
+                                  void *                  v_tiles_dst,
                                   size_t                  n_col_tiles,
                                   uint32_t                kv_start) {
     worker_pool_context_t wp = factx->octx->ctx->worker_pool;
@@ -547,7 +554,7 @@ static void fa_phase_v_interleave(struct hmx_fa_context * factx,
         n = factx->n_threads;
     }
     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, buf_idx, n_col_tiles, kv_start, rows_per_t };
+    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);
     } else {
@@ -563,6 +570,9 @@ typedef struct {
     uint32_t                  ib3;
     size_t                    n_rows_g;
     size_t                    rows_per_t;
+    size_t                    n_rows_q;
+    bool                      q_transposed;
+    atomic_uint               barrier;
 } fa_q_load_args_t;
 
 static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
@@ -587,9 +597,8 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
         const uint32_t g_br = factx->g_br;
         const uint32_t DV   = factx->DV;
 
-        const size_t col_vec_bytes = hex_align_up(g_br * sizeof(float), 256);
-        const size_t d_tile_bytes  = hex_align_up(g_br * g_br * sizeof(__fp16), 4096);
-        const size_t o_tile_bytes  = hex_align_up(g_br * DV * sizeof(__fp16), 4096);
+        const size_t col_vec_bytes = factx->col_vec_bytes;
+        const size_t d_tile_bytes  = factx->d_tile_bytes;
 
         // Initialize vtcm_l_vec & vtcm_m_vec
         const size_t l_bytes_per_t = hex_align_up(col_vec_bytes / n, 128);
@@ -643,72 +652,63 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
         if (d_start < d_tile_bytes) {
             hvx_splat_u8_a((char *) factx->vtcm_d_tiles + d_start, 0, d_end - d_start);
         }
+    }
 
-        // Initialize vtcm_o_tiles[0] to 0
-        __fp16 * o_tile_prev       = factx->vtcm_o_tiles[0];
+    if (start < factx->g_br) {
+        const struct htp_tensor * q       = args->q;
+        const uint32_t            q_start = args->q_start;
+        const uint32_t            kv_head = args->kv_head;
+        const uint32_t            ib3     = args->ib3;
+
+        assert(factx->DK == factx->DV);
+
+        const size_t o_tile_bytes = factx->o_tile_bytes;
+        const bool use_q_dma = (2 * o_tile_bytes >= factx->g_br * DK * (factx->is_q_fp32 ? 4 : 2));
+
+        __fp16 * q_tiles = factx->vtcm_q_tiles;
+        if (use_q_dma) {
+            const size_t g_rows_end = hex_smin(end, n_rows_g);
+            const uint32_t d_limit = factx->is_q_fp32 ? DK / 32 : DK / 64;
+
+            uint8_t * q_flat  = (uint8_t *) factx->vtcm_o_tiles[0];
+            if (factx->is_q_fp32) {
+                switch (d_limit) {
+                case 2:  hmx_fa_q_prep_fp32_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
+                case 4:  hmx_fa_q_prep_fp32_d4(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
+                default: hmx_fa_q_prep_fp32(   q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); break;
+                }
+            } else {
+                switch (d_limit) {
+                case 1:  hmx_fa_q_prep_fp16_d1(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
+                case 2:  hmx_fa_q_prep_fp16_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
+                default: hmx_fa_q_prep_fp16(   q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); break;
+                }
+            }
+        } else {
+            // Fallback: direct-from-DDR/L2 path
+            hmx_fa_q_prep_fallback(q_tiles, q->data, q->nb[1], q->nb[2], q->nb[3],
+                                   q_start, kv_head, ib3, start, end, n_rows_g, G, DK, factx->is_q_fp32, &factx->div_G);
+        }
+    }
+
+    // Synchronize threads before zeroing out vtcm_o_tiles[0] to prevent race condition
+    if (n > 1) {
+        atomic_fetch_sub(&args->barrier, 1);
+        while (atomic_load(&args->barrier) > 0) {
+            // spin wait
+        }
+    }
+
+    // Zero out vtcm_o_tiles[0] as it was used as temp_q_vtcm
+    {
+        const uint32_t g_br = factx->g_br;
+        const uint32_t DV   = factx->DV;
+        const size_t o_tile_bytes  = factx->o_tile_bytes;
         const size_t o_bytes_per_t = hex_align_up(o_tile_bytes / n, 128);
         const size_t o_start       = i * o_bytes_per_t;
         const size_t o_end         = hex_smin(o_start + o_bytes_per_t, o_tile_bytes);
         if (o_start < o_tile_bytes) {
-            hvx_splat_u8_a((char *) o_tile_prev + o_start, 0, o_end - o_start);
-        }
-    }
-
-    if (start >= factx->g_br) {
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_Q_PREP, (uint16_t) (args->q_start * G + start));
-        return;
-    }
-
-    const struct htp_tensor * q       = args->q;
-    const uint32_t            q_start = args->q_start;
-    const uint32_t            kv_head = args->kv_head;
-    const uint32_t            ib3     = args->ib3;
-
-    for (size_t r = start; r < end; r += 2) {
-        const size_t q_idx0 = fastdiv(r + 0, &factx->div_G);
-        const size_t h_idx0 = fastmodulo(r + 0, G, &factx->div_G);
-        const size_t q_idx1 = fastdiv(r + 1, &factx->div_G);
-        const size_t h_idx1 = fastmodulo(r + 1, G, &factx->div_G);
-
-        const uint8_t * q_ptr0 = (r + 0 < n_rows_g) ? ((const uint8_t *) q->data + (q_start + q_idx0) * q->nb[1] +
-                                                      (kv_head * G + h_idx0) * q->nb[2] + ib3 * q->nb[3]) :
-                                                      NULL;
-        const uint8_t * q_ptr1 = (r + 1 < n_rows_g) ? ((const uint8_t *) q->data + (q_start + q_idx1) * q->nb[1] +
-                                                      (kv_head * G + h_idx1) * q->nb[2] + ib3 * q->nb[3]) :
-                                                      NULL;
-
-        size_t   r0       = r / HMX_FP16_TILE_N_ROWS;
-        size_t   r1       = r % HMX_FP16_TILE_N_ROWS;
-        __fp16 * out_base = factx->vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
-
-        if (factx->is_q_fp32) {
-            const HVX_Vector * pv_in0 = q_ptr0 ? (const HVX_Vector *) q_ptr0 : NULL;
-            const HVX_Vector * pv_in1 = q_ptr1 ? (const HVX_Vector *) q_ptr1 : NULL;
-
-            for (uint32_t d = 0; d < DK / 32; ++d) {
-                HVX_Vector v0   = pv_in0 ? pv_in0[d] : Q6_V_vzero();
-                HVX_Vector v1   = pv_in1 ? pv_in1[d] : Q6_V_vzero();
-                HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
-
-                HVX_Vector * out_tile = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
-                out_tile[r1 / 2]      = v_hf;
-            }
-        } else {
-            const HVX_Vector * pv_in0 = q_ptr0 ? (const HVX_Vector *) q_ptr0 : NULL;
-            const HVX_Vector * pv_in1 = q_ptr1 ? (const HVX_Vector *) q_ptr1 : NULL;
-
-            for (uint32_t d = 0; d < DK / 64; ++d) {
-                HVX_Vector     v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
-                HVX_Vector     v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
-                HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
-
-                __fp16 *     out_dual_tile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
-                HVX_Vector * pv_out0       = ((HVX_Vector *) out_dual_tile) + r1 / 2;
-                HVX_Vector * pv_out1       = pv_out0 + 16;
-
-                *pv_out0 = Q6_V_lo_W(vp);
-                *pv_out1 = Q6_V_hi_W(vp);
-            }
+            hvx_splat_u8_a((char *) factx->vtcm_o_tiles[0] + o_start, 0, o_end - o_start);
         }
     }
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_Q_PREP, (uint16_t) (args->q_start * G + start));
@@ -726,7 +726,18 @@ static void fa_phase_q_load(struct hmx_fa_context *   factx,
         n = factx->n_threads;
     }
     size_t rows_per_t = hex_align_up(hmx_ceil_div(factx->g_br, n), 2);
-    fa_q_load_args_t args = { factx, q, q_start, kv_head, ib3, n_rows_g, rows_per_t };
+    const uint32_t n_rows_q = hex_smin(factx->Br, factx->neq1 - q_start);
+    fa_q_load_args_t args;
+    args.factx = factx;
+    args.q = q;
+    args.q_start = q_start;
+    args.kv_head = kv_head;
+    args.ib3 = ib3;
+    args.n_rows_g = n_rows_g;
+    args.rows_per_t = rows_per_t;
+    args.n_rows_q = n_rows_q;
+    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);
     } else {
@@ -798,11 +809,10 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
     fa_o_store_args_t *     args  = (fa_o_store_args_t *) data;
     struct hmx_fa_context * factx = args->factx;
 
-    const size_t n_rows_g = args->n_rows_g;
-    const size_t G        = factx->G;
-    const size_t DV       = factx->DV;
-
+    const size_t n_rows_g   = args->n_rows_g;
     const size_t rows_per_t = args->rows_per_t;
+    const size_t G          = factx->G;
+    const size_t DV         = factx->DV;
     const size_t start      = (size_t) i * rows_per_t;
     const size_t end        = hex_smin(start + rows_per_t, n_rows_g);
 
@@ -831,10 +841,10 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
         const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV;
 
         for (uint32_t d = 0; d < DV / 64; ++d) {
-            const __fp16 *     in_dual_tile = tile_row_base + d * HMX_FP16_TILE_N_ELMS * 2;
-            const HVX_Vector * pv_in0       = ((const HVX_Vector *) in_dual_tile) + r1 / 2;
-            const HVX_Vector * pv_in1       = pv_in0 + 16;
-            HVX_VectorPair     vp           = Q6_W_vdeal_VVR(*pv_in1, *pv_in0, -2);
+            const __fp16 *     in_dtile = tile_row_base + d * HMX_FP16_TILE_N_ELMS * 2;
+            const HVX_Vector * pv_in0   = ((const HVX_Vector *) in_dtile) + r1 / 2;
+            const HVX_Vector * pv_in1   = pv_in0 + 16;
+            HVX_VectorPair     vp       = Q6_W_vdeal_VVR(*pv_in1, *pv_in0, -2);
             if (r1 % 2 == 0) {
                 *(HVX_UVector *) (out + d * 64) = Q6_V_lo_W(vp);
             } else {
@@ -957,14 +967,14 @@ static inline void fa_softmax_impl(
             if (has_softcap) {
                 const HVX_Vector v_cap = hvx_vec_splat_f16(factx->logit_softcap);
                 for (size_t c = 0; c < kv_rows; c += 64) {
-                    size_t             ci           = c / 64;
-                    const __fp16 *     in_dual_tile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
-                    const HVX_Vector * pv_s_in0     = ((const HVX_Vector *) in_dual_tile) + r1 / 2;
-                    const HVX_Vector * pv_s_in1     = pv_s_in0 + 16;
+                    size_t             ci       = c / 64;
+                    const __fp16 *     in_dtile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
+                    const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dtile) + r1 / 2;
+                    const HVX_Vector * pv_s_in1 = pv_s_in0 + 16;
 
-                    HVX_VectorPair vp_s_dual_row = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2);
-                    HVX_Vector     v_s_row0      = Q6_V_lo_W(vp_s_dual_row);
-                    HVX_Vector     v_s_row1      = Q6_V_hi_W(vp_s_dual_row);
+                    HVX_VectorPair vp_s_drow = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2);
+                    HVX_Vector     v_s_row0  = Q6_V_lo_W(vp_s_drow);
+                    HVX_Vector     v_s_row1  = Q6_V_hi_W(vp_s_drow);
 
                     HVX_Vector t0   = hvx_vec_tanh_f16(v_s_row0);
                     my_row_buf0[ci] = hvx_vec_mul_f16_f16(t0, v_cap);
@@ -974,14 +984,14 @@ static inline void fa_softmax_impl(
                 }
             } else {
                 for (size_t c = 0; c < kv_rows; c += 64) {
-                    size_t             ci           = c / 64;
-                    const __fp16 *     in_dual_tile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
-                    const HVX_Vector * pv_s_in0     = ((const HVX_Vector *) in_dual_tile) + r1 / 2;
-                    const HVX_Vector * pv_s_in1     = pv_s_in0 + 16;
+                    size_t             ci       = c / 64;
+                    const __fp16 *     in_dtile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
+                    const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dtile) + r1 / 2;
+                    const HVX_Vector * pv_s_in1 = pv_s_in0 + 16;
 
-                    HVX_VectorPair vp_s_dual_row = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2);
-                    my_row_buf0[ci]              = Q6_V_lo_W(vp_s_dual_row);
-                    my_row_buf1[ci]              = Q6_V_hi_W(vp_s_dual_row);
+                    HVX_VectorPair vp_s_drow = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2);
+                    my_row_buf0[ci]          = Q6_V_lo_W(vp_s_drow);
+                    my_row_buf1[ci]          = Q6_V_hi_W(vp_s_drow);
                 }
             }
 
@@ -1118,9 +1128,9 @@ static inline void fa_softmax_impl(
 
                 HVX_Vector v_p_row0_hf  = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m0));
                 HVX_Vector v_p_row1_hf  = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m1));
-                __fp16 *     out_dual_tile = p_st_base + (c / 64) * HMX_FP16_TILE_N_ELMS * 2;
-                HVX_Vector * pv_p_out0     = ((HVX_Vector *) out_dual_tile) + r1 / 2;
-                HVX_Vector * pv_p_out1     = pv_p_out0 + 16;
+                __fp16 *     out_dtile  = p_st_base + ci * HMX_FP16_TILE_N_ELMS * 2;
+                HVX_Vector * pv_p_out0  = ((HVX_Vector *) out_dtile) + r1 / 2;
+                HVX_Vector * pv_p_out1  = pv_p_out0 + 16;
 
                 HVX_VectorPair vp_p_dual = Q6_W_vshuff_VVR(v_p_row1_hf, v_p_row0_hf, -2);
                 *pv_p_out0               = Q6_V_lo_W(vp_p_dual);
@@ -1150,7 +1160,7 @@ static inline void fa_softmax_impl(
         }
 
         // Inline fa_ml_update_and_build_d for this vector (lock-free and in parallel)
-        HVX_VectorPair rowmax_acc_pair = hvx_vec_f16_to_f32(rowmax_acc_v);
+        HVX_VectorPair rowmax_acc_pair    = hvx_vec_f16_to_f32(rowmax_acc_v);
         HVX_Vector     v_rowmax_acc_f32_0 = Q6_V_lo_W(rowmax_acc_pair);
         HVX_Vector     v_rowmax_acc_f32_1 = Q6_V_hi_W(rowmax_acc_pair);
 
@@ -1160,7 +1170,7 @@ static inline void fa_softmax_impl(
         HVX_Vector v_m_diff0 = HVX_OP_SUB_F32(m_prev_v0, v_m_curr0);
         HVX_Vector v_m_diff1 = HVX_OP_SUB_F32(m_prev_v1, v_m_curr1);
 
-        HVX_Vector v_m_diff_f16 = hvx_vec_f32_to_f16(v_m_diff0, v_m_diff1);
+        HVX_Vector v_m_diff_f16   = hvx_vec_f32_to_f16(v_m_diff0, v_m_diff1);
         HVX_Vector exp_m_diff_f16 = hvx_vec_exp2_f16(v_m_diff_f16);
 
         HVX_VectorPair exp_m_diff_pair = hvx_vec_f16_to_f32(exp_m_diff_f16);
@@ -1331,14 +1341,17 @@ static void hmx_fa_qk_dot_worker(void * data) {
     __builtin_assume(n_col_tiles > 0);
     __builtin_assume(n_dot_tiles > 0);
 
-    Q6_bias_mxmem2_A((void *) job->hmx_scales);
+    asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)job->hmx_scales));
+    const size_t dot_stride = n_dot_tiles * HMX_FP16_TILE_N_ELMS;
     for (size_t r = 0; r < n_row_tiles; ++r) {
-        for (size_t c = 0; c < n_col_tiles; ++c) {
-            const __fp16 * row_tiles = q_tiles + r * HMX_FP16_TILE_N_ROWS * n_dot_tiles * HMX_FP16_TILE_N_COLS;
-            const __fp16 * col_tiles = k_tiles + c * HMX_FP16_TILE_N_COLS * n_dot_tiles * HMX_FP16_TILE_N_COLS;
-            __fp16 *       out_tile  = s_tiles + (r * n_tiles_per_bc + c) * HMX_FP16_TILE_N_ELMS;
+        const __fp16 * row_tiles = q_tiles + r * dot_stride;
+        const __fp16 * col_tiles = k_tiles;
+        __fp16 *       out_tile  = s_tiles + r * n_tiles_per_bc * HMX_FP16_TILE_N_ELMS;
 
+        for (size_t c = 0; c < n_col_tiles; ++c) {
             hmx_fa_qk_dot_tile(row_tiles, col_tiles, out_tile, n_dot_tiles);
+            col_tiles += dot_stride;
+            out_tile  += HMX_FP16_TILE_N_ELMS;
         }
     }
 }
@@ -1373,17 +1386,21 @@ static void hmx_fa_o_update_worker(void * data) {
     __builtin_assume(n_col_tiles > 0);
     __builtin_assume(DV_tiles > 0);
 
-    Q6_bias_mxmem2_A((void *) job->hmx_scales);
+    asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)job->hmx_scales));
+    const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS;
+    const size_t v_stride = n_tiles_per_bc * HMX_FP16_TILE_N_ELMS;
     for (size_t r = 0; r < n_row_tiles; ++r) {
-        for (size_t c = 0; c < DV_tiles; ++c) {
-            // D[r,r] @ O_prev[r,c] — only the diagonal tile
-            const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
-            const __fp16 * o_rc   = o_prev + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
-            const __fp16 * p_tile_in = p_tiles + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
-            const __fp16 * v_tile_in = v_tiles + (c * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
-            __fp16 * o_tile_out = o_curr + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
+        const __fp16 * d_diag     = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+        const __fp16 * p_tile_in  = p_tiles + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
+        const __fp16 * o_rc       = o_prev + r * HMX_FP16_TILE_N_ELMS;
+        const __fp16 * v_tile_in  = v_tiles;
+        __fp16       * o_tile_out = o_curr + r * HMX_FP16_TILE_N_ELMS;
 
+        for (size_t c = 0; c < DV_tiles; ++c) {
             hmx_fa_o_update_tile(d_diag, o_rc, p_tile_in, v_tile_in, o_tile_out, n_col_tiles);
+            o_rc       += o_stride;
+            v_tile_in  += v_stride;
+            o_tile_out += o_stride;
         }
     }
 }
@@ -1409,14 +1426,17 @@ static void hmx_fa_o_norm_worker(void * data) {
     __builtin_assume(n_row_tiles > 0);
     __builtin_assume(DV_tiles > 0);
 
-    Q6_bias_mxmem2_A((void *) job->hmx_scales);
+    asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)job->hmx_scales));
+    const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS;
     for (size_t r = 0; r < n_row_tiles; ++r) {
-        for (size_t c = 0; c < DV_tiles; ++c) {
-            const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
-            const __fp16 * o_rc   = o_prev + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
-            __fp16 *       o_out  = o_curr + (r * DV_tiles + c) * HMX_FP16_TILE_N_ELMS;
+        const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+        const __fp16 * o_rc = o_prev + r * HMX_FP16_TILE_N_ELMS;
+        __fp16 *       o_out = o_curr + r * DV_tiles * HMX_FP16_TILE_N_ELMS;
 
+        for (size_t c = 0; c < DV_tiles; ++c) {
             hmx_fa_o_norm_tile(d_diag, o_rc, o_out);
+            o_rc  += o_stride;
+            o_out += HMX_FP16_TILE_N_ELMS;
         }
     }
 }
@@ -1475,7 +1495,7 @@ static void fa_push_mask_dma_gqa(
     uint32_t                  G,
     uint32_t                  m_line_bytes,
     uint32_t                  kv_rows,
-    uint32_t                  n_q_rows,
+    uint32_t                  n_rows_q,
     struct hmx_fa_context *   factx
 ) {
     for (uint32_t g = 0; g < G; ++g) {
@@ -1484,7 +1504,7 @@ static void fa_push_mask_dma_gqa(
         const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] +
                                  im2 * mask->nb[2] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
         uint8_t * ms_dst = (uint8_t *) factx->vtcm_mask_buf + g * m_line_bytes;
-        dma_queue_push(dma, dma_make_ptr(ms_dst, ms_src), G * m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_q_rows);
+        dma_queue_push(dma, dma_make_ptr(ms_dst, ms_src), G * m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
     }
 }
 
@@ -1582,62 +1602,57 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     const uint32_t G = factx.G;
 
     // ======== VTCM allocation (GQA-aware) ========
+    // K/V row sizes drive the DMA descriptors (not the VTCM layout) and are used
+    // throughout the KV loop below.
     const size_t size_k_row        = DK * sizeof(__fp16);
     const size_t size_v_row        = DV * sizeof(__fp16);
     const size_t size_k_row_padded = hex_round_up(size_k_row, 128);
     const size_t size_v_row_padded = hex_round_up(size_v_row, 128);
 
-    const size_t q_tile_bytes  = hex_align_up(g_br * DK * sizeof(__fp16), 4096);
-    const size_t o_tile_bytes  = hex_align_up(g_br * DV * sizeof(__fp16), 4096);
-    const size_t k_dma_bytes   = hex_align_up(Bc * size_k_row_padded, 4096);
-    const size_t v_dma_bytes   = hex_align_up(Bc * size_v_row_padded, 4096);
-    const size_t k_tile_bytes  = hex_align_up(Bc * DK * sizeof(__fp16), 4096);
-    const size_t v_tile_bytes  = hex_align_up(Bc * DV * sizeof(__fp16), 4096);
-    const size_t s_tile_bytes  = hex_align_up(g_br * Bc * sizeof(__fp16), 4096);
-    const size_t d_tile_bytes  = hex_align_up(g_br * g_br * sizeof(__fp16), 4096);
-    const size_t col_vec_bytes = hex_align_up(g_br * sizeof(float), 256);
-    const size_t row_vec_bytes = hex_align_up(Bc * sizeof(__fp16), 256);
-    const size_t m_line_bytes  = hex_align_up(Bc * sizeof(__fp16), 128);
-    const size_t m_buf_bytes   = hex_align_up(Br * m_line_bytes, 4096) * HMX_FA_DMA_CACHE_SIZE;
-    const size_t slopes_bytes  = hex_align_up(g_br * sizeof(__fp16), 128);
+    // Build the VTCM layout once (shared with the host estimator) and place every
+    // scratch buffer at its computed offset.
+    struct hmx_fa_vtcm_layout L;
+    hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline);
 
-    uint8_t * vtcm_cur = ctx->vtcm_base;
-
-    factx.vtcm_q_tiles        = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, q_tile_bytes);
-    factx.vtcm_o_tiles[0]     = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, o_tile_bytes);
-    factx.vtcm_o_tiles[1]     = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, o_tile_bytes);
-    factx.vtcm_k_fp16[0]      = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, k_dma_bytes);
-    factx.vtcm_k_fp16[1]      = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, k_dma_bytes);
-    factx.vtcm_v_fp16[0]      = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_dma_bytes);
-    factx.vtcm_v_fp16[1]      = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_dma_bytes);
-    factx.vtcm_k_tiles        = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, k_tile_bytes);
-    factx.vtcm_v_tiles[0]     = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_tile_bytes);
-    if (pipeline) {
-        factx.vtcm_v_tiles[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_tile_bytes);
-    } else {
-        factx.vtcm_v_tiles[1] = NULL;
-    }
-    factx.vtcm_s_tiles        = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, s_tile_bytes);
-    factx.vtcm_p_tiles        = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, s_tile_bytes);
-    factx.vtcm_d_tiles        = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, d_tile_bytes);
-    factx.vtcm_m_vec          = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
-    factx.vtcm_l_vec          = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
-    factx.vtcm_s_rowmax       = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
-    factx.vtcm_p_rowsum       = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
-    factx.vtcm_row_bufs       = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, row_vec_bytes * 2 * n_threads);
-    factx.row_buf_stride      = row_vec_bytes / sizeof(HVX_Vector);
-    factx.vtcm_hmx_scales_id  = vtcm_seq_alloc(&vtcm_cur, 256);
-    factx.vtcm_hmx_scales_qk  = vtcm_seq_alloc(&vtcm_cur, 256);
-    factx.vtcm_mask_buf       = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, m_buf_bytes);
-    factx.mask_buf_row_stride = m_line_bytes / sizeof(__fp16);
-    factx.vtcm_slopes         = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, slopes_bytes);
-
-    dma_cache_init(&factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, hex_align_up(Br * m_line_bytes, 4096), HMX_FA_DMA_CACHE_SIZE);
-
-    if ((size_t) (vtcm_cur - ctx->vtcm_base) > ctx->vtcm_size) {
+    if (L.total_bytes > ctx->vtcm_size) {
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
+    uint8_t * const base = ctx->vtcm_base;
+
+    factx.vtcm_q_tiles        = VTCM_LAYOUT_PTR(__fp16, base, L.off_q_tiles);
+    factx.vtcm_o_tiles[0]     = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[0]);
+    factx.vtcm_o_tiles[1]     = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[1]);
+    factx.vtcm_k_fp16[0]      = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_fp16[0]);
+    factx.vtcm_k_fp16[1]      = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_fp16[1]);
+    factx.vtcm_v_fp16[0]      = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[0]);
+    factx.vtcm_v_fp16[1]      = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[1]);
+    factx.vtcm_k_tiles        = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_tiles);
+    factx.vtcm_v_tiles[0]     = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_tiles[0]);
+    factx.vtcm_v_tiles[1]     = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_v_tiles[1], pipeline);
+    factx.vtcm_s_tiles        = VTCM_LAYOUT_PTR(__fp16, base, L.off_s_tiles);
+    factx.vtcm_p_tiles        = VTCM_LAYOUT_PTR(__fp16, base, L.off_p_tiles);
+    factx.vtcm_d_tiles        = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_tiles);
+    factx.vtcm_m_vec          = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_m_vec);
+    factx.vtcm_l_vec          = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_l_vec);
+    factx.vtcm_s_rowmax       = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_s_rowmax);
+    factx.vtcm_p_rowsum       = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_p_rowsum);
+    factx.vtcm_row_bufs       = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_row_bufs);
+    factx.row_buf_stride      = L.row_buf_stride;
+    factx.vtcm_hmx_scales_id  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_hmx_scales_id);
+    factx.vtcm_hmx_scales_qk  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_hmx_scales_qk);
+    factx.vtcm_mask_buf       = VTCM_LAYOUT_PTR(__fp16, base, L.off_mask_buf);
+    factx.mask_buf_row_stride = L.mask_buf_row_stride;
+    factx.q_tile_bytes        = L.q_tile_bytes;
+    factx.o_tile_bytes        = L.o_tile_bytes;
+    factx.col_vec_bytes       = L.col_vec_bytes;
+    factx.d_tile_bytes        = L.d_tile_bytes;
+    factx.vtcm_slopes         = VTCM_LAYOUT_PTR(__fp16, base, L.off_slopes);
+
+    const size_t m_line_bytes = L.m_line_bytes;  // used by the mask DMAs in the KV loop
+
+    dma_cache_init(&factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, L.m_buf_slot_bytes, HMX_FA_DMA_CACHE_SIZE);
+
     // ======== Initialize HMX output scales ========
     hmx_init_column_scales(factx.vtcm_hmx_scales_id, Q6_V_vsplat_R(0x3c00)); // 1.0
     hmx_init_column_scales(factx.vtcm_hmx_scales_qk, hvx_vec_splat_f16(factx.scale));
@@ -1655,11 +1670,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
 
     const size_t qo_element_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
 
-    // ======== HMX lock strategy ========
-    if (!factx.pipeline) {
-        HAP_compute_res_hmx_lock(ctx->vtcm_rctx);
-    }
-
     // ======== Reusable job descriptors for pipeline ========
     hmx_fa_qk_job_t       qk_job;
     hmx_fa_o_update_job_t ou_job;
@@ -1669,28 +1679,44 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     for (uint32_t ib3 = 0; ib3 < neq3; ++ib3) {
         const uint32_t im3 = mask ? fastmodulo(ib3, mask->ne[3], &factx.src3_div3) : 0;
         for (uint32_t q_start = 0; q_start < neq1; q_start += Br) {
-            const uint32_t n_q_rows    = hex_smin(Br, neq1 - q_start);
-            const size_t   n_rows_g    = n_q_rows * G;
+            const uint32_t n_rows_q    = hex_smin(Br, neq1 - q_start);
+            const size_t   n_rows_g    = n_rows_q * G;
             const size_t   g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
             const size_t   n_row_tiles = g_br_actual / HMX_FP16_TILE_N_ROWS;
 
             for (uint32_t kv_head = 0; kv_head < n_kv_heads; ++kv_head) {
                 const uint32_t ik2 = kv_head;
-                const uint32_t ik3 = ib3 / (neq3 / k->ne[3]);
+                const uint32_t ik3 = fastdiv(ib3, &kparams->broadcast_rk3);
                 const uint32_t iv2 = kv_head;
-                const uint32_t iv3 = ib3 / (neq3 / v->ne[3]);
+                const uint32_t iv3 = fastdiv(ib3, &kparams->broadcast_rv3);
 
-                // Prefetch first KV block
+                // 1. Push Q DMA (if Q DMA is used)
+                const size_t o_tile_bytes = factx.o_tile_bytes;
+                const bool use_q_dma = (2 * o_tile_bytes >= factx.g_br * factx.DK * (factx.is_q_fp32 ? 4 : 2));
+                if (use_q_dma) {
+                    const bool q_transposed = q->nb[1] < q->nb[2];
+                    const uint8_t * q_ptr = (const uint8_t *) q->data + q_start * q->nb[1] + (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3];
+                    const size_t el_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
+                    const size_t q_row_bytes = q_transposed ? n_rows_q * factx.DK * el_size : factx.G * factx.DK * el_size;
+                    const size_t src_stride  = q_transposed ? q->nb[2] : q->nb[1];
+                    const size_t n_rows      = q_transposed ? factx.G : n_rows_q;
+                    dma_queue_push(dma, dma_make_ptr(factx.vtcm_o_tiles[0], q_ptr), q_row_bytes, hex_smax(src_stride, q_row_bytes), q_row_bytes, n_rows);
+                }
+
+                // 2. Prefetch first KV block
                 if (factx.n_kv_blocks > 0) {
                     const uint32_t kv_rows0 = hex_smin(Bc, nek1);
 
                     const uint8_t * k_src = (const uint8_t *) k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
-                    dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1],
-                                   size_k_row, kv_rows0);
+                    dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
 
                     const uint8_t * v_src = (const uint8_t *) v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
-                    dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1],
-                                   size_v_row, kv_rows0);
+                    dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
+                }
+
+                // 3. Pop Q DMA (blocks until Q is loaded)
+                if (use_q_dma) {
+                    dma_queue_pop(dma);
                 }
 
                 // ---- Load Q block & Initialize per-block state ----
@@ -1709,12 +1735,10 @@ 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);
 
-                if (factx.pipeline) {
-                    // ==================================================================
-                    // Pipeline path
-                    // ==================================================================
-                    struct hmx_queue * hmx_q = ctx->hmx_queue;
+                struct hmx_queue * hmx_q = ctx->hmx_queue;
 
+                if (factx.pipeline) {
+                    // Pipeline path
                     for (uint32_t kv_blk = 0; kv_blk < factx.n_kv_blocks; ++kv_blk) {
                         const uint32_t kv_start    = kv_blk * Bc;
                         const uint32_t kv_rows     = hex_smin(Bc, nek1 - kv_start);
@@ -1724,15 +1748,22 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         if (mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
                                 const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
-                                dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_q_rows);
+                                dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
                             } else {
-                                fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_q_rows, &factx);
+                                fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
                             }
                         }
 
-                        // Wait for current KV DMA
-                        dma_queue_pop(dma);  // K
-                        dma_queue_pop(dma);  // V
+                        // Prefetch next KV block early
+                        if (kv_blk + 1 < factx.n_kv_blocks) {
+                            const uint32_t  prefetch_start = (kv_blk + 1) * Bc;
+                            const uint32_t  prefetch_rows  = hex_smin(Bc, nek1 - prefetch_start);
+                            const size_t    prefetch_buf   = 1 - buf_idx;
+                            const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
+                            dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
+                            const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
+                            dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
+                        }
 
                         // ---- Phase 1: K_int ----
                         if (kv_blk > 0) {
@@ -1749,7 +1780,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             ou_job.DV               = DV;
                             hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
                         }
-                        fa_phase_k_interleave(&factx, kv_rows, k_src_stride, buf_idx, kv_start);
+
+                        // Wait for current K DMA and interleave
+                        void * curr_k = dma_queue_pop(dma).dst;
+                        fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start);
 
                         // ---- Phase 2: qk_dot ----
                         qk_job.q_tiles        = factx.vtcm_q_tiles;
@@ -1762,16 +1796,9 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         qk_job.hmx_scales     = factx.vtcm_hmx_scales_qk;
                         hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job));
 
-                        if (kv_blk + 1 < factx.n_kv_blocks) {
-                            const uint32_t  prefetch_start = (kv_blk + 1) * Bc;
-                            const uint32_t  prefetch_rows  = hex_smin(Bc, nek1 - prefetch_start);
-                            const size_t    prefetch_buf   = 1 - buf_idx;
-                            const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
-                            dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
-                            const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
-                            dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
-                        }
-                        fa_phase_v_interleave(&factx, kv_rows, v_src_stride, buf_idx, n_tiles_per_bc, kv_start);
+                        // Wait for current V DMA and interleave
+                        void * curr_v = dma_queue_pop(dma).dst;
+                        fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[buf_idx], n_tiles_per_bc, kv_start);
 
                         if (kv_blk > 0) {
                             hmx_queue_pop(hmx_q);
@@ -1838,24 +1865,21 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                     }
 
                 } else {
-                    // ==================================================================
                     // Fallback path
-                    // ==================================================================
                     for (uint32_t kv_blk = 0; kv_blk < factx.n_kv_blocks; ++kv_blk) {
                         const uint32_t kv_start    = kv_blk * Bc;
                         const uint32_t kv_rows     = hex_smin(Bc, nek1 - kv_start);
                         const size_t   n_col_tiles = hmx_ceil_div(kv_rows, HMX_FP16_TILE_N_COLS);
-                        dma_queue_pop(dma);  // K
-                        dma_queue_pop(dma);  // V
 
                         if (mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
                                 const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
-                                dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_q_rows);
+                                dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
                             } else {
-                                fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_q_rows, &factx);
+                                fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
                             }
                         }
+
                         if (kv_blk + 1 < factx.n_kv_blocks) {
                             const uint32_t  prefetch_start = (kv_blk + 1) * Bc;
                             const uint32_t  prefetch_rows  = hex_smin(Bc, nek1 - prefetch_start);
@@ -1865,31 +1889,29 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
                             dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
                         }
-                        fa_phase_k_interleave(&factx, kv_rows, k_src_stride, buf_idx, kv_start);
+
+                        // Wait for current K DMA and interleave
+                        void * curr_k = dma_queue_pop(dma).dst;
+                        fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start);
 
                         {
-                            const size_t n_dot_tiles       = (size_t) (DK / 32);
-                            const __fp16 * restrict q_base = factx.vtcm_q_tiles;
-                            const __fp16 * restrict k_base = factx.vtcm_k_tiles;
-                            __fp16 * restrict s_base       = factx.vtcm_s_tiles;
-                            __builtin_assume(n_row_tiles > 0);
-                            __builtin_assume(n_col_tiles > 0);
-                            __builtin_assume(n_dot_tiles > 0);
+                            qk_job.q_tiles        = factx.vtcm_q_tiles;
+                            qk_job.k_tiles        = factx.vtcm_k_tiles;
+                            qk_job.s_tiles        = factx.vtcm_s_tiles;
+                            qk_job.n_row_tiles    = n_row_tiles;
+                            qk_job.n_col_tiles    = n_col_tiles;
+                            qk_job.n_dot_tiles    = (size_t) (DK / 32);
+                            qk_job.n_tiles_per_bc = n_tiles_per_bc;
+                            qk_job.hmx_scales     = factx.vtcm_hmx_scales_qk;
 
-                            htp_trace_event_start(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
-                            Q6_bias_mxmem2_A((void *) factx.vtcm_hmx_scales_qk);
-                            for (size_t r = 0; r < n_row_tiles; ++r) {
-                                for (size_t c = 0; c < n_col_tiles; ++c) {
-                                    const __fp16 * row_tiles = q_base + r * HMX_FP16_TILE_N_ROWS * DK;
-                                    const __fp16 * col_tiles = k_base + c * HMX_FP16_TILE_N_COLS * DK;
-                                    __fp16 *       out_tile  = s_base + (r * n_tiles_per_bc + c) * HMX_FP16_TILE_N_ELMS;
-
-                                    hmx_fa_qk_dot_tile(row_tiles, col_tiles, out_tile, n_dot_tiles);
-                                }
-                            }
-                            htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
+                            hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job));
+                            hmx_queue_pop(ctx->hmx_queue);
                         }
 
+                        // Wait for current V DMA and interleave
+                        void * curr_v = dma_queue_pop(dma).dst;
+                        fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[0], n_tiles_per_bc, kv_start);
+
                         // ---- Phase 3: softmax + build_D ----
                         __fp16 * current_mask_vtcm = NULL;
                         if (mask) {
@@ -1922,33 +1944,23 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         sargs.mask_vtcm_row_stride = factx.mask_buf_row_stride;
                         sargs.slopes               = factx.vtcm_slopes;
                         fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br);
-                        fa_phase_v_interleave(&factx, kv_rows, v_src_stride, buf_idx, n_tiles_per_bc, kv_start);
 
                         {
-                            const size_t DV_tiles           = (size_t) (DV / 32);
-                            const __fp16 * restrict d_base  = factx.vtcm_d_tiles;
-                            const __fp16 * restrict p_base  = factx.vtcm_p_tiles;
-                            const __fp16 * restrict v_base  = factx.vtcm_v_tiles[0];
-                            const __fp16 * restrict op_base = o_tile_prev;
-                            __fp16 * restrict oc_base       = o_tile_curr;
-                            __builtin_assume(n_row_tiles > 0);
-                            __builtin_assume(n_col_tiles > 0);
-                            __builtin_assume(DV_tiles > 0);
+                            ou_job.o_curr           = o_tile_curr;
+                            ou_job.o_prev           = o_tile_prev;
+                            ou_job.p_tiles          = factx.vtcm_p_tiles;
+                            ou_job.v_tiles          = factx.vtcm_v_tiles[0];
+                            ou_job.d_tiles          = factx.vtcm_d_tiles;
+                            ou_job.hmx_scales       = factx.vtcm_hmx_scales_id;
+                            ou_job.n_row_tiles      = n_row_tiles;
+                            ou_job.n_col_tiles      = n_col_tiles;
+                            ou_job.n_row_tiles_g_br = n_row_tiles_g_br;
+                            ou_job.n_tiles_per_bc   = n_tiles_per_bc;
+                            ou_job.DV               = DV;
 
-                            htp_trace_event_start(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
-                            Q6_bias_mxmem2_A((void *) factx.vtcm_hmx_scales_id);
-                            for (size_t r = 0; r < n_row_tiles; ++r) {
-                                for (size_t c = 0; c < DV_tiles; ++c) {
-                                    const __fp16 * d_diag = d_base  + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
-                                    const __fp16 * o_rc   = op_base + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
-                                    const __fp16 * p_tile_in = p_base + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
-                                    const __fp16 * v_tile_in = v_base + (c * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
-                                    __fp16 * o_tile_out = oc_base + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
+                            hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
+                            hmx_queue_pop(ctx->hmx_queue);
 
-                                    hmx_fa_o_update_tile(d_diag, o_rc, p_tile_in, v_tile_in, o_tile_out, n_col_tiles);
-                                }
-                            }
-                            htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
                             hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
                         }
 
@@ -1962,37 +1974,15 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                     fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
                     htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
 
-                    if (factx.pipeline) {
-                        on_job.o_curr           = o_tile_curr;
-                        on_job.o_prev           = o_tile_prev;
-                        on_job.d_tiles          = factx.vtcm_d_tiles;
-                        on_job.hmx_scales       = factx.vtcm_hmx_scales_id;
-                        on_job.n_row_tiles      = n_row_tiles;
-                        on_job.n_row_tiles_g_br = n_row_tiles_g_br;
-                        on_job.DV               = DV;
-                        hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_norm_worker, &on_job));
-                        hmx_queue_pop(ctx->hmx_queue);
-                    } else {
-                        const size_t DV_tiles           = (size_t) (DV / 32);
-                        const __fp16 * restrict d_base  = factx.vtcm_d_tiles;
-                        const __fp16 * restrict op_base = o_tile_prev;
-                        __fp16 * restrict oc_base       = o_tile_curr;
-                        __builtin_assume(n_row_tiles > 0);
-                        __builtin_assume(DV_tiles > 0);
-
-                        htp_trace_event_start(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
-                        Q6_bias_mxmem2_A((void *) factx.vtcm_hmx_scales_id);
-                        for (size_t r = 0; r < n_row_tiles; ++r) {
-                            for (size_t c = 0; c < DV_tiles; ++c) {
-                                const __fp16 * d_diag = d_base  + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
-                                const __fp16 * o_rc   = op_base + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
-                                __fp16 *       o_out  = oc_base + (r * DV_tiles + c) * HMX_FP16_TILE_N_ELMS;
-
-                                hmx_fa_o_norm_tile(d_diag, o_rc, o_out);
-                            }
-                        }
-                        htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
-                    }
+                    on_job.o_curr           = o_tile_curr;
+                    on_job.o_prev           = o_tile_prev;
+                    on_job.d_tiles          = factx.vtcm_d_tiles;
+                    on_job.hmx_scales       = factx.vtcm_hmx_scales_id;
+                    on_job.n_row_tiles      = n_row_tiles;
+                    on_job.n_row_tiles_g_br = n_row_tiles_g_br;
+                    on_job.DV               = DV;
+                    hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_norm_worker, &on_job));
+                    hmx_queue_pop(ctx->hmx_queue);
                 }
 
                 // ---- Store O block ----
@@ -2001,12 +1991,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
         }
     }
 
-    if (factx.pipeline) {
-        hmx_queue_suspend(ctx->hmx_queue);
-    } else {
-        HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
-    }
-
     return HTP_STATUS_OK;
 }
 
@@ -2040,10 +2024,10 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
     factx.src0_div21 = kparams->u.hvx.src0_div21;
     factx.src0_div1  = kparams->u.hvx.src0_div1;
 
-    factx.broadcast_rk2 = kparams->u.hvx.broadcast_rk2;
-    factx.broadcast_rk3 = kparams->u.hvx.broadcast_rk3;
-    factx.broadcast_rv2 = kparams->u.hvx.broadcast_rv2;
-    factx.broadcast_rv3 = kparams->u.hvx.broadcast_rv3;
+    factx.broadcast_rk2 = kparams->broadcast_rk2;
+    factx.broadcast_rk3 = kparams->broadcast_rk3;
+    factx.broadcast_rv2 = kparams->broadcast_rv2;
+    factx.broadcast_rv3 = kparams->broadcast_rv3;
 
     if (mask) {
         factx.src3_div2 = kparams->src3_div2;

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

@@ -7,19 +7,23 @@
 
 #include "hex-fastdiv.h"
 #include "hex-common.h"
+#include "htp-vtcm.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 // Tile constants (mirrored from hmx-utils.h for use on host side if needed)
+#define HTP_FA_HMX_TILE_SIZE   2048
+#define HMX_FP16_TILE_SIZE     2048
 #define HMX_FP16_TILE_N_ROWS   32
 #define HMX_FP16_TILE_N_COLS   32
 #define HMX_FP16_TILE_N_ELMS   1024
-#define HMX_FP16_TILE_SIZE     2048
+
 #define HVX_FA_DMA_CACHE_SIZE  128
 #define HMX_FA_DMA_CACHE_SIZE  4
 
+
 #define HTP_FA_M_INITIAL_VAL  -10000.0f
 
 enum htp_fa_kernel_type {
@@ -54,6 +58,11 @@ struct htp_fa_kernel_params {
     struct fastdiv_values src3_div2;
     struct fastdiv_values src3_div3;
 
+    struct fastdiv_values broadcast_rk2;
+    struct fastdiv_values broadcast_rk3;
+    struct fastdiv_values broadcast_rv2;
+    struct fastdiv_values broadcast_rv3;
+
     union {
         struct {
             uint32_t g_br;
@@ -69,10 +78,6 @@ struct htp_fa_kernel_params {
             uint32_t size_v_row_padded;
             struct fastdiv_values src0_div21;
             struct fastdiv_values src0_div1;
-            struct fastdiv_values broadcast_rk2;
-            struct fastdiv_values broadcast_rk3;
-            struct fastdiv_values broadcast_rv2;
-            struct fastdiv_values broadcast_rv3;
         } hvx;
     } u;
 };
@@ -81,39 +86,124 @@ struct htp_fa_kernel_params {
 static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params is too large for kernel_params blob");
 #endif
 
-// Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration.
-// g_br = hex_align_up(gqa_factor * Br, 32) replaces Br for all Q/O/S/P/D dimensions.
-// Layout: Q + O_ping + O_pong + K_dma*2 + V_dma*2 + K_tile + V_tile + S + P + D + vectors + scales
-// Mask is DMA'd into a VTCM buffer (Br rows per KV block) to avoid DDR reads in softmax.
-static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline) {
+// VTCM region layout for the HMX flash-attention kernel.
+//
+// Single source of truth for both the host (which needs the total size to pick a
+// (Br, Bc) tiling that fits the VTCM budget) and the device (which needs the actual
+// byte offsets to place each scratch buffer). Building the layout once and reading
+// offsets/total from it makes host estimate and device allocation impossible to
+// desync -- previously they were duplicated formulas in two files and drifted.
+//
+// All fields are byte offsets / byte sizes -- no HVX_Vector type is named here so the
+// header stays host-includable. The device casts (base + off_*) to the proper type.
+// An offset of 0 marks a region that is not allocated for this configuration (only
+// off_v_tiles[1], which exists only when pipelining); the device sets such pointers NULL.
+struct hmx_fa_vtcm_layout {
+    // Byte offsets from vtcm_base for each region.
+    size_t off_q_tiles;
+    size_t off_o_tiles[2];
+    size_t off_k_fp16[2];
+    size_t off_v_fp16[2];
+    size_t off_k_tiles;
+    size_t off_v_tiles[2];     // [1] allocated only when pipeline, else 0
+    size_t off_s_tiles;
+    size_t off_p_tiles;
+    size_t off_d_tiles;
+    size_t off_m_vec;
+    size_t off_l_vec;
+    size_t off_s_rowmax;
+    size_t off_p_rowsum;
+    size_t off_row_bufs;
+    size_t off_hmx_scales_id;
+    size_t off_hmx_scales_qk;
+    size_t off_mask_buf;
+    size_t off_slopes;
+
+    // Region byte sizes reused by the device at runtime (not just for allocation).
+    size_t q_tile_bytes;
+    size_t o_tile_bytes;
+    size_t s_tile_bytes;       // S and P tiles (same size)
+    size_t d_tile_bytes;
+    size_t m_line_bytes;       // one mask row
+    size_t m_buf_slot_bytes;   // one dma_cache slot = align_up(Br * m_line_bytes, 4096)
+    size_t col_vec_bytes;
+
+    // Derived strides.
+    size_t row_buf_stride;       // HVX vectors (128B) per row buffer
+    size_t mask_buf_row_stride;  // __fp16 elements per row in the mask buffer
+
+    bool   pipeline;
+    size_t total_bytes;
+};
+
+// Build the VTCM layout.
+
+static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
+                                       size_t gqa_factor, size_t DK, size_t DV,
+                                       size_t Br, size_t Bc, size_t n_threads, bool pipeline) {
     const size_t g_br         = hex_align_up(gqa_factor * Br, HMX_FP16_TILE_N_ROWS);
-    const size_t q_tile_size  = hex_align_up(g_br * DK * sizeof(__fp16), 4096);    // Q:  [g_br, DK]
-    const size_t o_tile_size  = hex_align_up(g_br * DV * sizeof(__fp16), 4096);    // O:  [g_br, DV] x2 ping-pong
-    const size_t k_dma_size   = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 4096);      // K DMA: [Bc, DK] x2 double-buf
-    const size_t v_dma_size   = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 4096);      // V DMA: [Bc, DV] x2 double-buf
-    const size_t k_tile_size  = hex_align_up(Bc * DK * sizeof(__fp16), 4096);      // K tiles: [Bc, DK] interleaved
-    const size_t v_tile_size  = hex_align_up(Bc * DV * sizeof(__fp16), 4096);      // V tiles: [Bc, DV] interleaved
-    const size_t s_tile_size  = hex_align_up(g_br * Bc * sizeof(__fp16), 4096);    // S/P:[g_br, Bc]
-    const size_t d_tile_size  = hex_align_up(g_br * g_br * sizeof(__fp16), 4096);  // D:  [g_br, g_br]
-    const size_t col_vec_size = hex_align_up(g_br * sizeof(float), 256);          // m, l, etc.
-    const size_t row_vec_size = hex_align_up(Bc * sizeof(__fp16), 256);
-    const size_t m_line_size  = hex_align_up(Bc * sizeof(__fp16), 128);
-    const size_t m_buf_size   = hex_align_up(Br * m_line_size, 4096) * HMX_FA_DMA_CACHE_SIZE;
+    const size_t q_tile_size  = hex_align_up(g_br * DK   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+    const size_t o_tile_size  = hex_align_up(g_br * DV   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+    const size_t k_tile_size  = hex_align_up(Bc   * DK   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+    const size_t v_tile_size  = hex_align_up(Bc   * DV   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+    const size_t s_tile_size  = hex_align_up(g_br * Bc   * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+    const size_t d_tile_size  = hex_align_up(g_br * g_br * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+
+    const size_t k_dma_size   = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 128);
+    const size_t v_dma_size   = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 128);
+    const size_t col_vec_size = hex_align_up(g_br * sizeof(float),  256);
+    const size_t row_vec_size = hex_align_up(Bc   * sizeof(__fp16), 256);
+    const size_t m_line_size  = hex_align_up(Bc   * sizeof(__fp16), 128);
+    const size_t m_buf_slot   = hex_align_up(Br * m_line_size, 256);
+    const size_t m_buf_size   = m_buf_slot * HMX_FA_DMA_CACHE_SIZE;
     const size_t slopes_size  = hex_align_up(g_br * sizeof(__fp16), 128);
 
-    return   q_tile_size * 1               // Q tiles
-           + o_tile_size * 2               // O ping-pong
-           + k_dma_size  * 2               // K DMA x2
-           + v_dma_size  * 2               // V DMA x2
-           + k_tile_size * 1               // K tiles
-           + v_tile_size * (pipeline ? 2 : 1) // V tiles (double-buffered if pipelining)
-           + s_tile_size * 2               // S + P
-           + d_tile_size * 1               // D (diagonal matrix)
-           + col_vec_size * 4              // m_vec, l_vec, s_rowmax, p_rowsum
-           + row_vec_size * 2 * n_threads  // per-thread softmax row scratch
-           + m_buf_size * 1                // mask VTCM buffer [Br rows]
-           + slopes_size                   // Slopes
-           + 256 * 2;                      // HMX scales (id + qk)
+    size_t off = 0;
+
+    // Section 1: HMX Tiled Buffers (FA_HMX_TILE_SIZE = 2KB Aligned)
+    VTCM_LAYOUT_ALLOC(off, off_q_tiles,       q_tile_size);
+    VTCM_LAYOUT_ALLOC(off, off_o_tiles[0],    o_tile_size);
+    VTCM_LAYOUT_ALLOC(off, off_o_tiles[1],    o_tile_size);
+    VTCM_LAYOUT_ALLOC(off, off_k_tiles,       k_tile_size);
+    VTCM_LAYOUT_ALLOC(off, off_v_tiles[0],    v_tile_size);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_v_tiles[1], v_tile_size, pipeline);
+    VTCM_LAYOUT_ALLOC(off, off_s_tiles,       s_tile_size);
+    VTCM_LAYOUT_ALLOC(off, off_p_tiles,       s_tile_size);
+    VTCM_LAYOUT_ALLOC(off, off_d_tiles,       d_tile_size);
+
+    // Section 2: HVX/DMA flat and vector buffers (128B / 256B Aligned)
+    VTCM_LAYOUT_ALLOC(off, off_k_fp16[0],     k_dma_size);
+    VTCM_LAYOUT_ALLOC(off, off_k_fp16[1],     k_dma_size);
+    VTCM_LAYOUT_ALLOC(off, off_v_fp16[0],     v_dma_size);
+    VTCM_LAYOUT_ALLOC(off, off_v_fp16[1],     v_dma_size);
+    VTCM_LAYOUT_ALLOC(off, off_m_vec,         col_vec_size);
+    VTCM_LAYOUT_ALLOC(off, off_l_vec,         col_vec_size);
+    VTCM_LAYOUT_ALLOC(off, off_s_rowmax,      col_vec_size);
+    VTCM_LAYOUT_ALLOC(off, off_p_rowsum,      col_vec_size);
+    VTCM_LAYOUT_ALLOC(off, off_row_bufs,      row_vec_size * 2 * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_hmx_scales_id, 256);
+    VTCM_LAYOUT_ALLOC(off, off_hmx_scales_qk, 256);
+    VTCM_LAYOUT_ALLOC(off, off_mask_buf,      m_buf_size);
+    VTCM_LAYOUT_ALLOC(off, off_slopes,        slopes_size);
+
+    L->q_tile_bytes        = q_tile_size;
+    L->o_tile_bytes        = o_tile_size;
+    L->col_vec_bytes       = col_vec_size;
+    L->s_tile_bytes        = s_tile_size;
+    L->d_tile_bytes        = d_tile_size;
+    L->m_line_bytes        = m_line_size;
+    L->m_buf_slot_bytes    = m_buf_slot;
+    L->row_buf_stride      = row_vec_size / 128;
+    L->mask_buf_row_stride = m_line_size / sizeof(__fp16);
+    L->pipeline            = pipeline;
+    L->total_bytes         = off;
+}
+
+// Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration.
+static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline) {
+    struct hmx_fa_vtcm_layout L;
+    hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline);
+    return L.total_bytes;
 }
 
 #define FA_HVX_BLOCK_SIZE 64
@@ -153,23 +243,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
     const size_t T       = HMX_FP16_TILE_N_ROWS;  // 32
     const size_t br_unit = hmx_ceil_div(T, gqa_factor);
     const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2;  // 64
-    const size_t fp16    = sizeof(__fp16);
     const bool   can_pipeline = (kv_len >= FA_MIN_KV_BLOCKS * bc_unit && n_threads >= 2);
 
-    // Approximate per-unit VTCM costs (without per-buffer alignment padding).
-    const size_t per_gbr  = (DK + 2 * DV) * fp16 + 4 * sizeof(float);  // Q + O*2 + 4 col vectors
-    const size_t per_gbr2 = fp16;                             // D diagonal matrix
-    const size_t per_bc =
-        3 * DK * fp16 + (can_pipeline ? 4 : 3) * DV * fp16 + 2 * n_threads * fp16;          // K/V DMA x2 + tiles + row bufs
-    const size_t per_gbr_bc = 2 * fp16;                       // S + P
-
-    const size_t overhead = 256 * 2 + 13 * 4096;
-
-    if (vtcm_budget <= overhead) {
-        return -1;
-    }
-    const size_t usable = vtcm_budget - overhead;
-
     // Br_max: largest Br aligned to br_unit that does not exceed qo_len.
     const size_t Br_max = qo_len >= br_unit ? hex_align_down(qo_len, br_unit) : br_unit;
 
@@ -185,51 +260,26 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
     size_t best_Br = 0, best_Bc = 0;
 
     for (size_t Br = Br_max; Br >= br_unit; Br -= br_unit) {
-        const size_t g_br = hex_align_up(gqa_factor * Br, T);
+        // Try all Bc candidates from Bc_limit down to bc_unit
+        for (size_t Bc = Bc_limit; Bc >= bc_unit; Bc -= bc_unit) {
+            size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline);
+            if (vtcm_needed <= vtcm_budget) {
+                // This Bc fits for this Br!
+                const size_t q_blocks  = (qo_len + Br - 1) / Br;
+                const size_t kv_blocks = (kv_len + Bc - 1) / Bc;
+                const size_t cost      = q_blocks * (c_q_fixed + kv_blocks * c_iter_fixed);
+                const size_t mn        = Br * Bc;
 
-        // g_br-dependent VTCM cost: g_br * per_gbr + g_br*g_br * per_gbr2
-        const size_t gbr_cost = g_br * per_gbr + g_br * g_br * per_gbr2;
-        if (gbr_cost >= usable) {
-            if (Br == br_unit) {
+                if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
+                    best_cost = cost;
+                    best_mn   = mn;
+                    best_Br   = Br;
+                    best_Bc   = Bc;
+                }
+                // Since we iterate Bc from largest to smallest, this is the largest Bc that fits
+                // for this Br. We can break to the next Br.
                 break;
             }
-            continue;
-        }
-
-        // Analytically solve for max Bc:
-        //   remain >= Bc * (per_bc + g_br * per_gbr_bc + Br * fp16 * HMX_FA_DMA_CACHE_SIZE)
-        // The Br * fp16 term accounts for the VTCM mask buffer [Br * Bc].
-        const size_t remain   = usable - gbr_cost;
-        const size_t bc_denom = per_bc + g_br * per_gbr_bc + Br * fp16 * HMX_FA_DMA_CACHE_SIZE;
-        size_t       Bc       = hex_smin(hex_align_down(remain / bc_denom, bc_unit), Bc_limit);
-        if (Bc < bc_unit) {
-            if (Br == br_unit) {
-                break;
-            }
-            continue;
-        }
-
-        // Exact VTCM verification (alignment padding may push over budget)
-        while (Bc >= bc_unit && hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline) > vtcm_budget) {
-            Bc -= bc_unit;
-        }
-        if (Bc < bc_unit) {
-            if (Br == br_unit) {
-                break;
-            }
-            continue;
-        }
-
-        const size_t q_blocks  = (qo_len + Br - 1) / Br;
-        const size_t kv_blocks = (kv_len + Bc - 1) / Bc;
-        const size_t cost      = q_blocks * (c_q_fixed + kv_blocks * c_iter_fixed);
-        const size_t mn        = Br * Bc;
-
-        if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
-            best_cost = cost;
-            best_mn   = mn;
-            best_Br   = Br;
-            best_Bc   = Bc;
         }
 
         if (Br == br_unit) {
@@ -237,7 +287,7 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
         }
     }
 
-    if (best_Br == 0) {
+    if (best_Br == 0 || best_Bc == 0) {
         return -1;
     }
 

ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h

@@ -6,6 +6,7 @@
 #include <stdbool.h>
 #include "hvx-utils.h"
 #include "hmx-utils.h"
+#include "hex-fastdiv.h"
 
 // HMX-specific parameters, offsets and inner kernels for Flash Attention
 
@@ -47,22 +48,75 @@ static const int16_t d_tile_scatter_offsets[64] __attribute__((aligned(128))) =
 };
 // Inner HMX tile computation kernels
 
-static inline void hmx_fa_qk_dot_tile(
+static void hmx_fa_qk_dot_tile(
     const __fp16 * row_tiles,
     const __fp16 * col_tiles,
     __fp16 *       out_tile,
     size_t         n_dot_tiles
 ) {
-    for (size_t k = 0; k < n_dot_tiles; ++k) {
-        Q6_activation_hf_mxmem_RR((unsigned int) row_tiles, 2047);
-        Q6_weight_hf_mxmem_RR((unsigned int) col_tiles, 2047);
-        row_tiles += HMX_FP16_TILE_N_ELMS;
-        col_tiles += HMX_FP16_TILE_N_ELMS;
+    if (n_dot_tiles == 2) {
+        asm volatile(
+            HMX_LOAD_MPY_F16("%1", "%2", "%0")
+            HMX_LOAD_MPY_F16("%3", "%4", "%0")
+            :
+            : "r"(2047),
+              "r"(row_tiles + 0 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 0 * HMX_FP16_TILE_N_ELMS),
+              "r"(row_tiles + 1 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 1 * HMX_FP16_TILE_N_ELMS)
+        );
+    } else if (n_dot_tiles == 4) {
+        asm volatile(
+            HMX_LOAD_MPY_F16("%1", "%2", "%0")
+            HMX_LOAD_MPY_F16("%3", "%4", "%0")
+            HMX_LOAD_MPY_F16("%5", "%6", "%0")
+            HMX_LOAD_MPY_F16("%7", "%8", "%0")
+            :
+            : "r"(2047),
+              "r"(row_tiles + 0 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 0 * HMX_FP16_TILE_N_ELMS),
+              "r"(row_tiles + 1 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 1 * HMX_FP16_TILE_N_ELMS),
+              "r"(row_tiles + 2 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 2 * HMX_FP16_TILE_N_ELMS),
+              "r"(row_tiles + 3 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 3 * HMX_FP16_TILE_N_ELMS)
+        );
+    } else if (n_dot_tiles == 8) {
+        asm volatile(
+            HMX_LOAD_MPY_F16("%1", "%2", "%0")
+            HMX_LOAD_MPY_F16("%3", "%4", "%0")
+            HMX_LOAD_MPY_F16("%5", "%6", "%0")
+            HMX_LOAD_MPY_F16("%7", "%8", "%0")
+            HMX_LOAD_MPY_F16("%9", "%10", "%0")
+            HMX_LOAD_MPY_F16("%11", "%12", "%0")
+            HMX_LOAD_MPY_F16("%13", "%14", "%0")
+            HMX_LOAD_MPY_F16("%15", "%16", "%0")
+            :
+            : "r"(2047),
+              "r"(row_tiles + 0 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 0 * HMX_FP16_TILE_N_ELMS),
+              "r"(row_tiles + 1 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 1 * HMX_FP16_TILE_N_ELMS),
+              "r"(row_tiles + 2 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 2 * HMX_FP16_TILE_N_ELMS),
+              "r"(row_tiles + 3 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 3 * HMX_FP16_TILE_N_ELMS),
+              "r"(row_tiles + 4 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 4 * HMX_FP16_TILE_N_ELMS),
+              "r"(row_tiles + 5 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 5 * HMX_FP16_TILE_N_ELMS),
+              "r"(row_tiles + 6 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 6 * HMX_FP16_TILE_N_ELMS),
+              "r"(row_tiles + 7 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 7 * HMX_FP16_TILE_N_ELMS)
+        );
+    } else {
+        for (size_t k = 0; k < n_dot_tiles; ++k) {
+            asm volatile(
+                HMX_LOAD_MPY_F16("%1", "%2", "%0")
+                :
+                : "r"(2047), "r"(row_tiles), "r"(col_tiles)
+            );
+            row_tiles += HMX_FP16_TILE_N_ELMS;
+            col_tiles += HMX_FP16_TILE_N_ELMS;
+        }
     }
-    Q6_mxmem_AR_after_hf(out_tile, 0);
+    asm volatile(
+        HMX_STORE_AFTER_F16("%0", "%1")
+        :
+        : "r"(out_tile), "r"(0)
+        : "memory"
+    );
 }
 
-static inline void hmx_fa_o_update_tile(
+static void hmx_fa_o_update_tile(
     const __fp16 * d_diag,
     const __fp16 * o_rc,
     const __fp16 * p_tile_in,
@@ -70,17 +124,71 @@ static inline void hmx_fa_o_update_tile(
     __fp16 *       o_tile_out,
     size_t         n_col_tiles
 ) {
-    Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047);
-    Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047);
-
-    for (size_t k = 0; k < n_col_tiles; ++k) {
-        Q6_activation_hf_mxmem_RR((unsigned int) p_tile_in, 2047);
-        Q6_weight_hf_mxmem_RR((unsigned int) v_tile_in, 2047);
-        p_tile_in += HMX_FP16_TILE_N_ELMS;
-        v_tile_in += HMX_FP16_TILE_N_ELMS;
+    asm volatile(
+        HMX_LOAD_MPY_F16("%1", "%2", "%0")
+        :
+        : "r"(2047), "r"(d_diag), "r"(o_rc)
+    );
+    if (n_col_tiles == 2) {
+        asm volatile(
+            HMX_LOAD_MPY_F16("%1", "%2", "%0")
+            HMX_LOAD_MPY_F16("%3", "%4", "%0")
+            :
+            : "r"(2047),
+              "r"(p_tile_in + 0 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 0 * HMX_FP16_TILE_N_ELMS),
+              "r"(p_tile_in + 1 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 1 * HMX_FP16_TILE_N_ELMS)
+        );
+    } else if (n_col_tiles == 4) {
+        asm volatile(
+            HMX_LOAD_MPY_F16("%1", "%2", "%0")
+            HMX_LOAD_MPY_F16("%3", "%4", "%0")
+            HMX_LOAD_MPY_F16("%5", "%6", "%0")
+            HMX_LOAD_MPY_F16("%7", "%8", "%0")
+            :
+            : "r"(2047),
+              "r"(p_tile_in + 0 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 0 * HMX_FP16_TILE_N_ELMS),
+              "r"(p_tile_in + 1 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 1 * HMX_FP16_TILE_N_ELMS),
+              "r"(p_tile_in + 2 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 2 * HMX_FP16_TILE_N_ELMS),
+              "r"(p_tile_in + 3 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 3 * HMX_FP16_TILE_N_ELMS)
+        );
+    } else if (n_col_tiles == 8) {
+        asm volatile(
+            HMX_LOAD_MPY_F16("%1", "%2", "%0")
+            HMX_LOAD_MPY_F16("%3", "%4", "%0")
+            HMX_LOAD_MPY_F16("%5", "%6", "%0")
+            HMX_LOAD_MPY_F16("%7", "%8", "%0")
+            HMX_LOAD_MPY_F16("%9", "%10", "%0")
+            HMX_LOAD_MPY_F16("%11", "%12", "%0")
+            HMX_LOAD_MPY_F16("%13", "%14", "%0")
+            HMX_LOAD_MPY_F16("%15", "%16", "%0")
+            :
+            : "r"(2047),
+              "r"(p_tile_in + 0 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 0 * HMX_FP16_TILE_N_ELMS),
+              "r"(p_tile_in + 1 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 1 * HMX_FP16_TILE_N_ELMS),
+              "r"(p_tile_in + 2 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 2 * HMX_FP16_TILE_N_ELMS),
+              "r"(p_tile_in + 3 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 3 * HMX_FP16_TILE_N_ELMS),
+              "r"(p_tile_in + 4 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 4 * HMX_FP16_TILE_N_ELMS),
+              "r"(p_tile_in + 5 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 5 * HMX_FP16_TILE_N_ELMS),
+              "r"(p_tile_in + 6 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 6 * HMX_FP16_TILE_N_ELMS),
+              "r"(p_tile_in + 7 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 7 * HMX_FP16_TILE_N_ELMS)
+        );
+    } else {
+        for (size_t k = 0; k < n_col_tiles; ++k) {
+            asm volatile(
+                HMX_LOAD_MPY_F16("%1", "%2", "%0")
+                :
+                : "r"(2047), "r"(p_tile_in), "r"(v_tile_in)
+            );
+            p_tile_in += HMX_FP16_TILE_N_ELMS;
+            v_tile_in += HMX_FP16_TILE_N_ELMS;
+        }
     }
-
-    Q6_mxmem_AR_after_hf(o_tile_out, 0);
+    asm volatile(
+        HMX_STORE_AFTER_F16("%0", "%1")
+        :
+        : "r"(o_tile_out), "r"(0)
+        : "memory"
+    );
 }
 
 static inline void hmx_fa_o_norm_tile(
@@ -88,9 +196,360 @@ static inline void hmx_fa_o_norm_tile(
     const __fp16 * o_rc,
     __fp16 *       o_out
 ) {
-    Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047);
-    Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047);
-    Q6_mxmem_AR_after_hf(o_out, 0);
+    asm volatile(
+        HMX_LOAD_MPY_F16("%1", "%2", "%0")
+        :
+        : "r"(2047), "r"(d_diag), "r"(o_rc)
+    );
+    asm volatile(
+        HMX_STORE_AFTER_F16("%0", "%1")
+        :
+        : "r"(o_out), "r"(0)
+        : "memory"
+    );
+}
+
+static inline void hmx_fa_q_prep_fp32_d2(
+    __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm,
+    size_t start, size_t end, size_t g_rows_end,
+    size_t DK, size_t G, size_t n_rows_q,
+    const struct fastdiv_values * div_G, bool q_transposed
+) {
+    for (size_t r = start; r < end; r += 2) {
+        size_t   r0       = r / HMX_FP16_TILE_N_ROWS;
+        size_t   r1       = r % HMX_FP16_TILE_N_ROWS;
+        __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+        if (r >= g_rows_end) {
+            ((HVX_Vector *) (out_base + 0 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero();
+            ((HVX_Vector *) (out_base + 1 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero();
+            continue;
+        }
+
+        const size_t q_idx0 = fastdiv(r + 0, div_G);
+        const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+        const size_t q_idx1 = fastdiv(r + 1, div_G);
+        const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+        const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+        const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+        const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(float));
+        const HVX_Vector * pv_in1 = (r + 1 < g_rows_end)
+            ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(float))
+            : NULL;
+
+        {
+            HVX_Vector v0 = pv_in0[0];
+            HVX_Vector v1 = pv_in1 ? pv_in1[0] : Q6_V_vzero();
+            HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+            ((HVX_Vector *) (out_base + 0 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = v_hf;
+        }
+        {
+            HVX_Vector v0 = pv_in0[1];
+            HVX_Vector v1 = pv_in1 ? pv_in1[1] : Q6_V_vzero();
+            HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+            ((HVX_Vector *) (out_base + 1 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = v_hf;
+        }
+    }
+}
+
+static inline void hmx_fa_q_prep_fp32_d4(
+    __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm,
+    size_t start, size_t end, size_t g_rows_end,
+    size_t DK, size_t G, size_t n_rows_q,
+    const struct fastdiv_values * div_G, bool q_transposed
+) {
+    for (size_t r = start; r < end; r += 2) {
+        size_t   r0       = r / HMX_FP16_TILE_N_ROWS;
+        size_t   r1       = r % HMX_FP16_TILE_N_ROWS;
+        __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+        if (r >= g_rows_end) {
+            for (uint32_t d = 0; d < 4; ++d) {
+                ((HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero();
+            }
+            continue;
+        }
+
+        const size_t q_idx0 = fastdiv(r + 0, div_G);
+        const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+        const size_t q_idx1 = fastdiv(r + 1, div_G);
+        const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+        const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+        const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+        const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(float));
+        const HVX_Vector * pv_in1 = (r + 1 < g_rows_end)
+            ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(float))
+            : NULL;
+
+        for (uint32_t d = 0; d < 4; ++d) {
+            HVX_Vector v0 = pv_in0[d];
+            HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+            HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+            ((HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS))[r1 / 2] = v_hf;
+        }
+    }
+}
+
+static inline void hmx_fa_q_prep_fp32(
+    __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm,
+    size_t start, size_t end, size_t g_rows_end,
+    size_t DK, size_t G, size_t n_rows_q,
+    const struct fastdiv_values * div_G, uint32_t d_limit, bool q_transposed
+) {
+    for (size_t r = start; r < end; r += 2) {
+        size_t   r0       = r / HMX_FP16_TILE_N_ROWS;
+        size_t   r1       = r % HMX_FP16_TILE_N_ROWS;
+        __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+        if (r >= g_rows_end) {
+            for (uint32_t d = 0; d < d_limit; ++d) {
+                ((HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero();
+            }
+            continue;
+        }
+
+        const size_t q_idx0 = fastdiv(r + 0, div_G);
+        const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+        const size_t q_idx1 = fastdiv(r + 1, div_G);
+        const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+        const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+        const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+        const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(float));
+        const HVX_Vector * pv_in1 = (r + 1 < g_rows_end)
+            ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(float))
+            : NULL;
+
+        for (uint32_t d = 0; d < d_limit; ++d) {
+            HVX_Vector v0   = pv_in0[d];
+            HVX_Vector v1   = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+            HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+
+            HVX_Vector * out_tile = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
+            out_tile[r1 / 2]      = v_hf;
+        }
+    }
+}
+
+static inline void hmx_fa_q_prep_fp16_d1(
+    __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm,
+    size_t start, size_t end, size_t g_rows_end,
+    size_t DK, size_t G, size_t n_rows_q,
+    const struct fastdiv_values * div_G, bool q_transposed
+) {
+    for (size_t r = start; r < end; r += 2) {
+        size_t   r0       = r / HMX_FP16_TILE_N_ROWS;
+        size_t   r1       = r % HMX_FP16_TILE_N_ROWS;
+        __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+        if (r >= g_rows_end) {
+            __fp16 *     out_dtile = out_base + 0 * HMX_FP16_TILE_N_ELMS * 2;
+            HVX_Vector * pv_out0   = ((HVX_Vector *) out_dtile) + r1 / 2;
+            HVX_Vector * pv_out1   = pv_out0 + 16;
+            *pv_out0 = Q6_V_vzero();
+            *pv_out1 = Q6_V_vzero();
+            continue;
+        }
+
+        const size_t q_idx0 = fastdiv(r + 0, div_G);
+        const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+        const size_t q_idx1 = fastdiv(r + 1, div_G);
+        const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+        const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+        const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+        const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(__fp16));
+        const HVX_Vector * pv_in1 = (r + 1 < g_rows_end)
+            ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(__fp16))
+            : NULL;
+
+        HVX_Vector     v0 = pv_in0[0];
+        HVX_Vector     v1 = pv_in1 ? pv_in1[0] : Q6_V_vzero();
+        HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+        __fp16 *     out_dtile = out_base + 0 * HMX_FP16_TILE_N_ELMS * 2;
+        HVX_Vector * pv_out0   = ((HVX_Vector *) out_dtile) + r1 / 2;
+        HVX_Vector * pv_out1   = pv_out0 + 16;
+
+        *pv_out0 = Q6_V_lo_W(vp);
+        *pv_out1 = Q6_V_hi_W(vp);
+    }
+}
+
+static inline void hmx_fa_q_prep_fp16_d2(
+    __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm,
+    size_t start, size_t end, size_t g_rows_end,
+    size_t DK, size_t G, size_t n_rows_q,
+    const struct fastdiv_values * div_G, bool q_transposed
+) {
+    for (size_t r = start; r < end; r += 2) {
+        size_t   r0       = r / HMX_FP16_TILE_N_ROWS;
+        size_t   r1       = r % HMX_FP16_TILE_N_ROWS;
+        __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+        if (r >= g_rows_end) {
+            for (uint32_t d = 0; d < 2; ++d) {
+                __fp16 *     out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
+                HVX_Vector * pv_out0   = ((HVX_Vector *) out_dtile) + r1 / 2;
+                HVX_Vector * pv_out1   = pv_out0 + 16;
+                *pv_out0 = Q6_V_vzero();
+                *pv_out1 = Q6_V_vzero();
+            }
+            continue;
+        }
+
+        const size_t q_idx0 = fastdiv(r + 0, div_G);
+        const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+        const size_t q_idx1 = fastdiv(r + 1, div_G);
+        const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+        const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+        const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+        const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(__fp16));
+        const HVX_Vector * pv_in1 = (r + 1 < g_rows_end)
+            ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(__fp16))
+            : NULL;
+
+        {
+            HVX_Vector     v0 = pv_in0[0];
+            HVX_Vector     v1 = pv_in1 ? pv_in1[0] : Q6_V_vzero();
+            HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+            __fp16 *     out_dtile = out_base + 0 * HMX_FP16_TILE_N_ELMS * 2;
+            HVX_Vector * pv_out0   = ((HVX_Vector *) out_dtile) + r1 / 2;
+            HVX_Vector * pv_out1   = pv_out0 + 16;
+
+            *pv_out0 = Q6_V_lo_W(vp);
+            *pv_out1 = Q6_V_hi_W(vp);
+        }
+        {
+            HVX_Vector     v0 = pv_in0[1];
+            HVX_Vector     v1 = pv_in1 ? pv_in1[1] : Q6_V_vzero();
+            HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+            __fp16 *     out_dtile = out_base + 1 * HMX_FP16_TILE_N_ELMS * 2;
+            HVX_Vector * pv_out0   = ((HVX_Vector *) out_dtile) + r1 / 2;
+            HVX_Vector * pv_out1   = pv_out0 + 16;
+
+            *pv_out0 = Q6_V_lo_W(vp);
+            *pv_out1 = Q6_V_hi_W(vp);
+        }
+    }
+}
+
+static inline void hmx_fa_q_prep_fp16(
+    __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm,
+    size_t start, size_t end, size_t g_rows_end,
+    size_t DK, size_t G, size_t n_rows_q,
+    const struct fastdiv_values * div_G, uint32_t d_limit, bool q_transposed
+) {
+    for (size_t r = start; r < end; r += 2) {
+        size_t   r0       = r / HMX_FP16_TILE_N_ROWS;
+        size_t   r1       = r % HMX_FP16_TILE_N_ROWS;
+        __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+        if (r >= g_rows_end) {
+            for (uint32_t d = 0; d < d_limit; ++d) {
+                __fp16 *     out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
+                HVX_Vector * pv_out0   = ((HVX_Vector *) out_dtile) + r1 / 2;
+                HVX_Vector * pv_out1   = pv_out0 + 16;
+                *pv_out0 = Q6_V_vzero();
+                *pv_out1 = Q6_V_vzero();
+            }
+            continue;
+        }
+
+        const size_t q_idx0 = fastdiv(r + 0, div_G);
+        const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+        const size_t q_idx1 = fastdiv(r + 1, div_G);
+        const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+        const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+        const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+        const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(__fp16));
+        const HVX_Vector * pv_in1 = (r + 1 < g_rows_end)
+            ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(__fp16))
+            : NULL;
+
+        for (uint32_t d = 0; d < d_limit; ++d) {
+            HVX_Vector     v0 = pv_in0[d];
+            HVX_Vector     v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+            HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+            __fp16 *     out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
+            HVX_Vector * pv_out0   = ((HVX_Vector *) out_dtile) + r1 / 2;
+            HVX_Vector * pv_out1   = pv_out0 + 16;
+
+            *pv_out0 = Q6_V_lo_W(vp);
+            *pv_out1 = Q6_V_hi_W(vp);
+        }
+    }
+}
+
+
+static inline void hmx_fa_q_prep_fallback(
+    __fp16 * vtcm_q_tiles, uintptr_t q_data,
+    size_t q_nb1, size_t q_nb2, size_t q_nb3,
+    uint32_t q_start, uint32_t kv_head, uint32_t ib3,
+    size_t start, size_t end, size_t n_rows_g,
+    size_t G, size_t DK, bool is_q_fp32,
+    const struct fastdiv_values * div_G
+) {
+    for (size_t r = start; r < end; r += 2) {
+        const size_t q_idx0 = fastdiv(r + 0, div_G);
+        const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+        const size_t q_idx1 = fastdiv(r + 1, div_G);
+        const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+        const uint8_t * q_ptr0 = (r + 0 < n_rows_g) ? ((const uint8_t *) q_data + (q_start + q_idx0) * q_nb1 +
+                                                      (kv_head * G + h_idx0) * q_nb2 + ib3 * q_nb3) :
+                                                      NULL;
+        const uint8_t * q_ptr1 = (r + 1 < n_rows_g) ? ((const uint8_t *) q_data + (q_start + q_idx1) * q_nb1 +
+                                                      (kv_head * G + h_idx1) * q_nb2 + ib3 * q_nb3) :
+                                                      NULL;
+
+        size_t   r0       = r / HMX_FP16_TILE_N_ROWS;
+        size_t   r1       = r % HMX_FP16_TILE_N_ROWS;
+        __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+        if (is_q_fp32) {
+            const HVX_UVector * pv_in0 = q_ptr0 ? (const HVX_UVector *) q_ptr0 : NULL;
+            const HVX_UVector * pv_in1 = q_ptr1 ? (const HVX_UVector *) q_ptr1 : NULL;
+
+            for (uint32_t d = 0; d < DK / 32; ++d) {
+                HVX_Vector v0   = pv_in0 ? pv_in0[d] : Q6_V_vzero();
+                HVX_Vector v1   = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+                HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+
+                HVX_Vector * out_tile = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
+                out_tile[r1 / 2]      = v_hf;
+            }
+        } else {
+            const HVX_UVector * pv_in0 = q_ptr0 ? (const HVX_UVector *) q_ptr0 : NULL;
+            const HVX_UVector * pv_in1 = q_ptr1 ? (const HVX_UVector *) q_ptr1 : NULL;
+
+            for (uint32_t d = 0; d < DK / 64; ++d) {
+                HVX_Vector     v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
+                HVX_Vector     v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+                HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+                __fp16 *     out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
+                HVX_Vector * pv_out0   = ((HVX_Vector *) out_dtile) + r1 / 2;
+                HVX_Vector * pv_out1   = pv_out0 + 16;
+
+                *pv_out0 = Q6_V_lo_W(vp);
+                *pv_out1 = Q6_V_hi_W(vp);
+            }
+        }
+    }
 }
 
 #endif /* HMX_FA_KERNELS_H */

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

@@ -506,7 +506,8 @@ static void dequantize_tiled_weight_to_fp16_task_q8_0(
     }
 }
 
-static void convert_f16_weight_to_fp16_tiles_task(
+static __attribute__((noinline))
+void convert_f16_weight_to_fp16_tiles_task(
         const tiled_dequantize_state_t *state,
         uint32_t start_tile, uint32_t end_tile) {
 
@@ -543,17 +544,13 @@ static void convert_f16_weight_to_fp16_tiles_task(
                 Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v1);
                 v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step);
             }
-            (void) *(volatile HVX_Vector *)(tile_base);
         }
         ++t; ++kt;
     }
-
-    if (start_tile < end_tile) {
-        (void) *(volatile HVX_Vector *)(state->dst + (end_tile - 1) * HTP_MM_HMX_TILE_N_ELMS);
-    }
 }
 
-static void quantize_f32_weight_to_fp16_tiles_task(
+static __attribute__((noinline))
+void quantize_f32_weight_to_fp16_tiles_task(
         const tiled_dequantize_state_t *state,
         uint32_t start_tile, uint32_t end_tile) {
 
@@ -594,120 +591,178 @@ static void quantize_f32_weight_to_fp16_tiles_task(
                 Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v_out_hi);
                 v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step);
             }
-            (void) *(volatile HVX_Vector *)(tile_base);
         }
         ++t; ++kt;
     }
-
-    if (start_tile < end_tile) {
-        (void) *(volatile HVX_Vector *)(state->dst + (end_tile - 1) * HTP_MM_HMX_TILE_N_ELMS);
-    }
 }
 
 // --- End tiled dequantizers ---
 
-// requires external HMX lock
-static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales,
+// dot-chunk functions require external HMX lock
+
+static void core_dot_chunk_fp16_short(__fp16 *restrict output, const __fp16 *restrict activation,
+                                const __fp16 *restrict weight, const __fp16 *restrict scales,
                                 uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) {
     __builtin_assume(n_row_tiles > 0);
     __builtin_assume(n_col_tiles > 0);
     __builtin_assume(n_dot_tiles > 0);
+    __builtin_assume(n_dot_tiles <= 32);
+
+    asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)scales));
+
+    const size_t dot_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
+    const uint32_t range = 2048u * n_dot_tiles - 1;
 
-    Q6_bias_mxmem2_A((void *)scales);
     for (uint32_t r = 0; r < n_row_tiles; ++r) {
+        const __fp16 *row_base = activation + r * dot_stride;
+        const __fp16 *col_base = weight;
+        __fp16 *out_tile = output + r * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS;
+
         for (size_t c = 0; c < n_col_tiles; ++c) {
-            Q6_mxclracc_hf();
-
-            const __fp16 *row_tiles = activation + r * n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
-            const __fp16 *col_tiles = weight + c * n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
-
-            for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) {
-                k_block = hex_smin(n_dot_tiles - k, 32);
-                const uint32_t range = 2048u * (uint32_t)k_block - 1;
-                Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range);
-                Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range);
-                row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS;
-                col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS;
-            }
-
-            __fp16 *out_tile = output + (r * n_col_tiles + c) * HTP_MM_HMX_TILE_N_ELMS;
-            Q6_mxmem_AR_after_hf(out_tile, 0);
+            asm volatile(HMX_CLRACC_F16());
+            asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_base), "r"(col_base));
+            asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(out_tile), "r"(0) : "memory");
+            col_base += dot_stride;
+            out_tile += HTP_MM_HMX_TILE_N_ELMS;
         }
     }
 }
 
-// C += AB
-static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b,
+static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation,
+                          const __fp16 *restrict weight, const __fp16 *restrict scales,
+                          uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) {
+    if (n_dot_tiles <= 32) {
+        core_dot_chunk_fp16_short(output, activation, weight, scales, n_row_tiles, n_col_tiles, n_dot_tiles);
+        return;
+    }
+    __builtin_assume(n_row_tiles > 0);
+    __builtin_assume(n_col_tiles > 0);
+    __builtin_assume(n_dot_tiles > 32);
+
+    asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)scales));
+
+    const size_t dot_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
+
+    for (uint32_t r = 0; r < n_row_tiles; ++r) {
+        const __fp16 *row_base = activation + r * dot_stride;
+        const __fp16 *col_base = weight;
+        __fp16 *out_tile = output + r * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS;
+
+        for (size_t c = 0; c < n_col_tiles; ++c) {
+            const __fp16 *row_tiles = row_base;
+            const __fp16 *col_tiles = col_base;
+
+            asm volatile(HMX_CLRACC_F16());
+
+            const uint32_t n_loops = n_dot_tiles / 32;
+            const uint32_t rem = n_dot_tiles % 32;
+
+            for (uint32_t l = 0; l < n_loops; ++l) {
+                asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(65535), "r"(row_tiles), "r"(col_tiles));
+                row_tiles += 32 * HTP_MM_HMX_TILE_N_ELMS;
+                col_tiles += 32 * HTP_MM_HMX_TILE_N_ELMS;
+            }
+
+            if (rem > 0) {
+                const uint32_t range = 2048u * rem - 1;
+                asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_tiles), "r"(col_tiles));
+            }
+
+            asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(out_tile), "r"(0) : "memory");
+
+            col_base += dot_stride;
+            out_tile += HTP_MM_HMX_TILE_N_ELMS;
+        }
+    }
+}
+
+static void core_mma_chunk_fp16_short(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b,
                                 const __fp16 *restrict col_scales, const __fp16 *restrict eye_tile,
                                 uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles, bool zero_init) {
     __builtin_assume(n_row_tiles > 0);
     __builtin_assume(n_col_tiles > 0);
     __builtin_assume(n_dot_tiles > 0);
+    __builtin_assume(n_dot_tiles <= 32);
 
-    Q6_bias_mxmem2_A((void *)col_scales);
+    asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)col_scales));
 
     const size_t dot_tile_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
+    const uint32_t range = 2048u * n_dot_tiles - 1;
+
     for (size_t i = 0; i < n_row_tiles; ++i) {
         const __fp16 *row_base = a + i * dot_tile_stride;
         __fp16 *res_base = c + i * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS;
+        const __fp16 *col_base = b;
+        __fp16 *accum_tile = res_base;
+
         for (size_t j = 0; j < n_col_tiles; ++j) {
-            Q6_mxclracc_hf();
+            asm volatile(HMX_CLRACC_F16());
 
-            const __fp16 *col_tiles = b + j * dot_tile_stride;
-            const __fp16 *row_tiles = row_base;
-            __fp16 *accum_tile = res_base + j * HTP_MM_HMX_TILE_N_ELMS;
             if (!zero_init) {
-                Q6_activation_hf_mxmem_RR((unsigned int)accum_tile, 2047);
-                Q6_weight_hf_mxmem_RR((unsigned int)eye_tile, 2047);
+                asm volatile(HMX_LOAD_MPY_F16("%1", "%2", "%0") : : "r"(2047), "r"(accum_tile), "r"(eye_tile));
             }
 
-            for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) {
-                k_block = hex_smin(n_dot_tiles - k, 32);
-                const uint32_t range = 2048u * k_block - 1;
-                Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range);
-                Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range);
-                row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS;
-                col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS;
-            }
+            asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_base), "r"(col_base));
 
-            Q6_mxmem_AR_after_hf(accum_tile, 0);
+            asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(accum_tile), "r"(0) : "memory");
+
+            col_base   += dot_tile_stride;
+            accum_tile += HTP_MM_HMX_TILE_N_ELMS;
         }
     }
 }
 
-// --- Async HMX matmul job (for pipeline overlap) ---
+static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b,
+                          const __fp16 *restrict col_scales, const __fp16 *restrict eye_tile,
+                          uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles, bool zero_init) {
+    if (n_dot_tiles <= 32) {
+        core_mma_chunk_fp16_short(c, a, b, col_scales, eye_tile, n_row_tiles, n_col_tiles, n_dot_tiles, zero_init);
+        return;
+    }
+    __builtin_assume(n_row_tiles > 0);
+    __builtin_assume(n_col_tiles > 0);
+    __builtin_assume(n_dot_tiles > 32);
 
-typedef struct {
-    __fp16 *       output;
-    const __fp16 * activation;
-    const __fp16 * weight;
-    const __fp16 * scales;
-    uint32_t       n_row_tiles;
-    uint32_t       n_col_tiles;
-    uint32_t       n_dot_tiles;
-} hmx_matmul_job_t;
+    asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)col_scales));
 
-static void hmx_matmul_worker_fn(void * data) {
-    hmx_matmul_job_t * job = (hmx_matmul_job_t *) data;
-    FARF(HIGH, "hmx-mm-job: n_row_tiles %u n_col_tiles %u n_dot_tiles %u", job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles);
-    core_dot_chunk_fp16(job->output, job->activation, job->weight, job->scales, job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles);
-}
+    const size_t dot_tile_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
 
-static inline void hmx_matmul_job_init(hmx_matmul_job_t * job,
-                                       __fp16 *           output,
-                                       const __fp16 *     activation,
-                                       const __fp16 *     weight,
-                                       const __fp16 *     scales,
-                                       uint32_t           n_row_tiles,
-                                       uint32_t           n_col_tiles,
-                                       uint32_t           n_dot_tiles) {
-    job->output      = output;
-    job->activation  = activation;
-    job->weight      = weight;
-    job->scales      = scales;
-    job->n_row_tiles = n_row_tiles;
-    job->n_col_tiles = n_col_tiles;
-    job->n_dot_tiles = n_dot_tiles;
+    for (size_t i = 0; i < n_row_tiles; ++i) {
+        const __fp16 *row_base = a + i * dot_tile_stride;
+        __fp16 *res_base = c + i * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS;
+        const __fp16 *col_base = b;
+        __fp16 *accum_tile = res_base;
+
+        for (size_t j = 0; j < n_col_tiles; ++j) {
+            const __fp16 *col_tiles = col_base;
+            const __fp16 *row_tiles = row_base;
+
+            asm volatile(HMX_CLRACC_F16());
+
+            if (!zero_init) {
+                asm volatile(HMX_LOAD_MPY_F16("%1", "%2", "%0") : : "r"(2047), "r"(accum_tile), "r"(eye_tile));
+            }
+
+            const uint32_t n_loops = n_dot_tiles / 32;
+            const uint32_t rem = n_dot_tiles % 32;
+
+            for (uint32_t l = 0; l < n_loops; ++l) {
+                asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(65535), "r"(row_tiles), "r"(col_tiles));
+                row_tiles += 32 * HTP_MM_HMX_TILE_N_ELMS;
+                col_tiles += 32 * HTP_MM_HMX_TILE_N_ELMS;
+            }
+
+            if (rem > 0) {
+                const uint32_t range = 2048u * rem - 1;
+                asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_tiles), "r"(col_tiles));
+            }
+
+            asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(accum_tile), "r"(0) : "memory");
+
+            col_base += dot_tile_stride;
+            accum_tile += HTP_MM_HMX_TILE_N_ELMS;
+        }
+    }
 }
 
 // output : fp16 -> f32p
@@ -901,148 +956,55 @@ static void transfer_activation_chunk_fp32_to_fp16(__fp16 *restrict vtcm_dst, co
     }
 }
 
-typedef struct {
-    __fp16      *dst;
-    const float *src;
-    uint32_t     n_tasks;
-    uint32_t     n_tot_chunks;
-    uint32_t     n_chunks_per_task;
-    uint32_t     k_block;
-    uint32_t     k_stride;
-    uint32_t     k_valid;
-    struct htp_thread_trace * traces;
-    struct htp_context * ctx;
-    float              * vtcm_f32_act;
-} activation_transfer_task_state_t;
-
-static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined(
-        dma_queue *dma_q,
+static void transfer_activation_row_pair_fp32_to_fp16(
         __fp16 *restrict vtcm_dst,
-        const float *restrict src,
-        uint32_t n_rows,
+        const float *restrict row0,
+        const float *restrict row1,
+        uint32_t r,
         uint32_t k_block,
-        uint32_t k_stride,
         uint32_t k_valid,
-        float *thread_f32_act) {
+        bool row0_valid,
+        bool row1_valid) {
 
-    const uint32_t R = HTP_MM_DMA_ACT_ROWS_PER_STEP;
-    const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS);
+    uint32_t r0 = r / HTP_MM_HMX_TILE_N_ROWS;  // tile row index
+    uint32_t r1 = r % HTP_MM_HMX_TILE_N_ROWS;  // intra-tile row idx
 
-    const uint32_t n_steps = n_rows_padded / R;
+    uint32_t c = 0;
+    for (; c + 32 <= k_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);
 
-    // pre-fetch 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),
-                       k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch);
+        HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
+
+        uint32_t c0       = 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 < k_block) {
+        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);
 
-    for (uint32_t s = 0; s < n_steps; ++s) {
-        uint32_t r = R * s;
-        float *curr_buf = thread_f32_act + (s % 2) * R * k_block;
+        uint32_t rem = k_valid - c;
+        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());
 
-        if (r < n_rows) {
-            dma_queue_pop(dma_q);
-        }
+        HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
 
-        uint32_t next_s = s + 1;
-        uint32_t next_r = R * next_s;
-        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;
-            float *next_buf = thread_f32_act + (next_s % 2) * R * k_block;
-            dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src),
-                           k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch);
-        }
+        uint32_t c0       = c / HTP_MM_HMX_TILE_N_COLS;  // tile column index
+        uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0;
 
-        #pragma unroll
-        for (uint32_t i = 0; i < HTP_MM_DMA_ACT_ROWS_PER_STEP; i += 2) {
-            uint32_t curr_r = r + i;
-            const bool row0_valid = (curr_r < n_rows);
-            const bool row1_valid = (curr_r + 1) < n_rows;
-
-            const float *ptr_in0 = curr_buf + i * k_block;
-            const float *ptr_in1 = curr_buf + (i + 1) * k_block;
-
-            uint32_t c = 0;
-            for (; c + 32 <= k_valid; c += 32) {
-                HVX_Vector v0 = Q6_V_vzero();
-                HVX_Vector v1 = Q6_V_vzero();
-                if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c);
-                if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c);
-
-                HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
-
-                uint32_t r0       = curr_r / HTP_MM_HMX_TILE_N_ROWS;  // tile row index
-                uint32_t r1       = curr_r % HTP_MM_HMX_TILE_N_ROWS;  // intra-tile row idx
-                uint32_t c0       = 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 < k_block) {
-                HVX_Vector v0 = Q6_V_vzero();
-                HVX_Vector v1 = Q6_V_vzero();
-                if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c);
-                if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c);
-
-                uint32_t rem = k_valid - c;
-                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 r0       = curr_r / HTP_MM_HMX_TILE_N_ROWS;  // tile row index
-                uint32_t r1       = curr_r % HTP_MM_HMX_TILE_N_ROWS;  // intra-tile row idx
-                uint32_t c0       = 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;
-            }
-        }
+        HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS);
+        tile[r1 / 2]     = v_out;
     }
 }
 
-typedef struct {
-    const struct mmid_row_mapping  *matrix_rows;
-    __fp16                         *dst;
-    const float                    *src;
-    uint32_t                        n_tasks;
-    uint32_t                        n_tot_chunks;
-    uint32_t                        n_chunks_per_task;
-    uint32_t                        k_block;
-    uint32_t                        cur_a;
-    uint32_t                        mapping_stride;
-    uint32_t                        ne11;
-    struct fastdiv_values           ne11_div;
-    size_t                          nb11;
-    size_t                          nb12;
-    uint32_t                        start_row;
-    uint32_t                        cne1;
-    uint32_t                        k_valid;
-    struct htp_thread_trace        *traces;
-} activation_transfer_gathered_task_state_t;
-
-typedef struct {
-    const struct mmid_row_mapping  *matrix_rows;
-    const __fp16                   *vtcm_src;
-    float                          *dst;
-    uint32_t                        n_tasks;
-    uint32_t                        n_tot_chunks;
-    uint32_t                        n_chunks_per_task;
-    uint32_t                        n_cols;
-    uint32_t                        cur_a;
-    uint32_t                        mapping_stride;
-    size_t                          dst_nb1;
-    size_t                          dst_nb2;
-    uint32_t                        start_row;
-    uint32_t                        cne1;
-    struct htp_thread_trace        *traces;
-} output_transfer_scattered_task_state_t;
-
 static void transfer_activation_chunk_fp32_to_fp16_gathered(
             __fp16 *restrict vtcm_dst,
             const float *restrict src,

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

@@ -6,6 +6,7 @@
 
 #include <qurt_thread.h>
 #include <qurt_futex.h>
+#include <qurt_hvx.h>
 
 #include <HAP_compute_res.h>
 
@@ -42,6 +43,7 @@ static inline void hmx_queue_process(struct hmx_queue *q, bool* killed) {
                 case HMX_QUEUE_NOOP:    /* noop */;     break;
                 case HMX_QUEUE_KILL:    *killed = true; break;
                 case HMX_QUEUE_SUSPEND: hmx_unlock(q);  break;
+                case HMX_QUEUE_WAKEUP:  hmx_lock(q);    break;
                 default:
                     hmx_lock(q);
                     htp_trace_event_start(q->trace, HTP_TRACE_EVT_HMX_COMP, ir);
@@ -70,9 +72,14 @@ static void hmx_queue_thread(void * arg) {
     while (!killed) {
         unsigned int seqn = atomic_load(&q->seqn);
         if (seqn == prev_seqn) {
+            // drop HVX context while spinning
+            if (poll_cnt > 1 && poll_cnt == HMX_QUEUE_POLL_COUNT) {
+                qurt_hvx_unlock();
+            }
             if (--poll_cnt) { hex_pause(); continue; }
             FARF(HIGH, "hmx-queue-thread: sleeping");
             qurt_futex_wait(&q->seqn, prev_seqn);
+            poll_cnt = HMX_QUEUE_POLL_COUNT;
             continue;
         }
         prev_seqn = seqn;

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

@@ -18,13 +18,19 @@ extern "C" {
 #endif
 
 #define HMX_QUEUE_THREAD_STACK_SIZE (16 * 1024)
-#define HMX_QUEUE_POLL_COUNT        2000
+
+#if __HVX_ARCH__ > 79
+#define HMX_QUEUE_POLL_COUNT 2000
+#else
+#define HMX_QUEUE_POLL_COUNT 1
+#endif
 
 typedef void (*hmx_queue_func)(void *);
 
 // Dummy funcs used as signals
 enum hmx_queue_signal {
     HMX_QUEUE_NOOP = 0, // aka NULL
+    HMX_QUEUE_WAKEUP,
     HMX_QUEUE_SUSPEND,
     HMX_QUEUE_KILL
 };
@@ -97,7 +103,7 @@ static inline uint32_t hmx_queue_capacity(struct hmx_queue * q) {
     return q->capacity;
 }
 
-static inline struct hmx_queue_desc hmx_queue_pop(struct hmx_queue * q) {
+static inline struct hmx_queue_desc hmx_queue_pop_one(struct hmx_queue * q) {
     unsigned int ip = q->idx_pop;
     unsigned int iw = q->idx_write;
 
@@ -120,13 +126,28 @@ static inline struct hmx_queue_desc hmx_queue_pop(struct hmx_queue * q) {
     return rd;
 }
 
+static inline struct hmx_queue_desc hmx_queue_pop(struct hmx_queue * q) {
+    while (1) {
+        struct hmx_queue_desc d = hmx_queue_pop_one(q);
+
+        uint32_t sig = (uint32_t) d.func;
+        if (sig && sig <= HMX_QUEUE_KILL)
+            continue;
+
+        return d;
+    }
+}
+
 static inline void hmx_queue_flush(struct hmx_queue * q) {
-    while (hmx_queue_pop(q).func != NULL) ;
+    while (hmx_queue_pop_one(q).func != NULL) ;
+}
+
+static inline void hmx_queue_wakeup(struct hmx_queue * q) {
+    hmx_queue_signal(q, HMX_QUEUE_WAKEUP);
 }
 
 static inline void hmx_queue_suspend(struct hmx_queue *q) {
     hmx_queue_signal(q, HMX_QUEUE_SUSPEND);
-    hmx_queue_flush(q);
 }
 
 #ifdef __cplusplus

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

@@ -197,4 +197,26 @@ static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out,
     }
 }
 
+// --- HMX inline asm macros for load-store packetization ---
+#define HMX_LOAD_MPY_F16(act, wt, range) \
+    "{\n" \
+    "    activation.hf = mxmem(" act ", " range ")\n" \
+    "    weight.hf = mxmem(" wt ", " range ")\n" \
+    "}\n"
+
+#define HMX_LOAD_MPY_DEEP_F16(act, wt, range) \
+    "{\n" \
+    "    activation.hf = mxmem(" act ", " range "):deep\n" \
+    "    weight.hf = mxmem(" wt ", " range ")\n" \
+    "}\n"
+
+#define HMX_STORE_AFTER_F16(out, scale_reg) \
+    "mxmem(" out ", " scale_reg "):after.hf = acc\n"
+
+#define HMX_SET_BIAS(scales) \
+    "bias = mxmem2(" scales ")\n"
+
+#define HMX_CLRACC_F16() \
+    "mxclracc.hf\n"
+
 #endif // HMX_UTILS_H

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

@@ -0,0 +1,19 @@
+#ifndef HTP_VTCM_H
+#define HTP_VTCM_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+static inline uint8_t *vtcm_seq_alloc(uint8_t **vtcm_ptr, size_t size) {
+    uint8_t *p = *vtcm_ptr;
+    *vtcm_ptr += size;
+    return p;
+}
+
+#define VTCM_LAYOUT_ALLOC(off, field, sz) do { (L)->field = (off); (off) += (sz); } while (0)
+#define VTCM_LAYOUT_ALLOC_OPTIONAL(off, field, sz, cond) do { if (cond) { VTCM_LAYOUT_ALLOC(off, field, sz); } else { (L)->field = 0; } } while (0)
+
+#define VTCM_LAYOUT_PTR(type, base, offset) ((type *)((uint8_t *)(base) + (offset)))
+#define VTCM_LAYOUT_PTR_OPTIONAL(type, base, offset, cond) ((cond) ? VTCM_LAYOUT_PTR(type, base, offset) : NULL)
+
+#endif // HTP_VTCM_H

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

@@ -948,6 +948,8 @@ static void htp_packet_callback(dspqueue_t queue, int error, void * context) {
         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;
 
@@ -976,6 +978,8 @@ static void htp_packet_callback(dspqueue_t queue, int error, void * context) {
             }
         }
 
+        hmx_queue_suspend(ctx->hmx_queue);
+
         struct htp_opbatch_rsp rsp;
         rsp.id        = req.id;
         rsp.status    = op_status;

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

@@ -20,7 +20,7 @@
 #include "htp-ctx.h"
 #include "htp-ops.h"
 #include "matmul-ops.h"
-#include "vtcm-utils.h"
+#include "htp-vtcm.h"
 
 static void hvx_tensor_add_f32_grid(
     const struct htp_tensor * restrict dst,
@@ -1514,37 +1514,26 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
             break;
     }
 
-    size_t src0_sz = 0, src1_sz = 0, dst_sz = 0;
-    if (kparams->vtcm_src0_size > 0 || kparams->vtcm_src1_size > 0 || kparams->vtcm_dst_size > 0) {
-        src0_sz = kparams->vtcm_src0_size;
-        src1_sz = kparams->vtcm_src1_size;
-        dst_sz  = kparams->vtcm_dst_size;
-    } else {
-        const uint32_t n_prefetch = kparams->n_prefetch;
-        assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0);
-        htp_mm_hvx_get_vtcm_sizes(
-            kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
-            dst_row_size, src0_row_size, src1_row_size, n_prefetch,
-            &src0_sz, &src1_sz, &dst_sz
-        );
-    }
+    struct htp_mm_hvx_vtcm_layout L;
+    htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
+                                 dst_row_size, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, false);
 
     if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM ||
         kparams->kernel_type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
         kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW ||
         kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
-        mmctx->vtcm_src1_size_per_thread = src1_sz;
+        mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
     } else {
-        mmctx->vtcm_src1_size_per_thread = src1_sz / octx->n_threads;
+        mmctx->vtcm_src1_size_per_thread = L.src1_bytes / octx->n_threads;
     }
 
-    mmctx->vtcm_src0_size_per_thread = src0_sz / octx->n_threads;
-    mmctx->vtcm_dst_size_per_thread  = dst_sz  / octx->n_threads;
+    mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads;
+    mmctx->vtcm_dst_size_per_thread  = L.dst_bytes / octx->n_threads;
 
-    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : (src1_sz + src0_sz + dst_sz);
+    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
 
     FARF(HIGH, "matmul-%s : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu (%zu)\n", mmctx->type,
-         src0_sz, src1_sz, dst_sz, vtcm_size);
+         L.src0_bytes, L.src1_bytes, L.dst_bytes, vtcm_size);
 
     FARF(HIGH, "matmul-%s : %ux%ux%ux%u * %ux%ux%ux%u-> %ux%ux%ux%u (0x%p, 0x%p, 0x%p)\n", mmctx->type, src0->ne[0],
          src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0],
@@ -1556,10 +1545,10 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
-    mmctx->vtcm_src1 = vtcm_seq_alloc(&vtcm_ptr, src1_sz);
-    mmctx->vtcm_src0 = vtcm_seq_alloc(&vtcm_ptr, src0_sz);
-    mmctx->vtcm_dst  = vtcm_seq_alloc(&vtcm_ptr, dst_sz);
+    uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
+    mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+    mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    mmctx->vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
 
     octx->src1_spad.src  = NULL;
     octx->src0_spad.src  = NULL;
@@ -1948,14 +1937,95 @@ static void transfer_output_chunk_worker_fn(unsigned int n, unsigned int i, void
     htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, start_chunk_idx);
 }
 
+typedef struct {
+    const struct mmid_row_mapping  *matrix_rows;
+    __fp16      *dst;
+    const float *src;
+    uint32_t     n_tasks;
+    uint32_t     n_tot_chunks;
+    uint32_t     n_chunks_per_task;
+    uint32_t     k_block;
+    uint32_t     k_stride;
+    uint32_t     k_valid;
+    struct htp_thread_trace * traces;
+    struct htp_context * ctx;
+    float              * vtcm_f32_act;
+    size_t               vtcm_f32_act_bytes_per_thread;
+    uint32_t             dma_step_rows;
+    uint32_t             dma_step_rows_shift;
+} activation_transfer_task_state_t;
+
+static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined(
+        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_valid,
+        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),
+                       k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch);
+    }
+    // Push step 1 (if valid)
+    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;
+            float *next_buf = thread_f32_act + 1 * R * k_block;
+            dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src),
+                           k_block * sizeof(float), k_stride * sizeof(float), k_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) * k_block;
+            bool r0_valid = ((row_idx + 0) < n_rows);
+            bool r1_valid = ((row_idx + 1) < n_rows);
+
+            transfer_activation_row_pair_fp32_to_fp16(vtcm_dst, pair_buf, pair_buf + k_block, row_idx, k_block, k_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;
+            dma_queue_push(dma_q, dma_make_ptr(curr_buf, next_src),
+                           k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch);
+        }
+    }
+}
+
 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;
 
-    int start_chunk_idx = i * st->n_chunks_per_task;
-    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, start_chunk_idx);
-
     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;
         size_t chunk_size = hex_smin(st->n_tot_chunks - chunk_idx, st->n_chunks_per_task);
@@ -1964,18 +2034,55 @@ static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i,
         const float *src = st->src + chunk_idx * st->k_stride;
 
         if (st->vtcm_f32_act) {
-            float *thread_f32_act = st->vtcm_f32_act + i * HTP_MM_DMA_ACT_MULTIPLIER * st->k_block;
+            float *thread_f32_act = (float *)((char *)st->vtcm_f32_act + i * st->vtcm_f32_act_bytes_per_thread);
             transfer_activation_chunk_fp32_to_fp16_dma_pipelined(
-                st->ctx->dma[i], dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid, thread_f32_act
+                st->ctx->dma[i], dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid, 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, chunk_idx);
             transfer_activation_chunk_fp32_to_fp16(dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid);
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx);
         }
     }
-
-    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, start_chunk_idx);
 }
 
+typedef struct {
+    const struct mmid_row_mapping  *matrix_rows;
+    __fp16                         *dst;
+    const float                    *src;
+    uint32_t                        n_tasks;
+    uint32_t                        n_tot_chunks;
+    uint32_t                        n_chunks_per_task;
+    uint32_t                        k_block;
+    uint32_t                        cur_a;
+    uint32_t                        mapping_stride;
+    uint32_t                        ne11;
+    struct fastdiv_values           ne11_div;
+    size_t                          nb11;
+    size_t                          nb12;
+    uint32_t                        start_row;
+    uint32_t                        cne1;
+    uint32_t                        k_valid;
+    struct htp_thread_trace        *traces;
+} activation_transfer_gathered_task_state_t;
+
+typedef struct {
+    const struct mmid_row_mapping  *matrix_rows;
+    const __fp16                   *vtcm_src;
+    float                          *dst;
+    uint32_t                        n_tasks;
+    uint32_t                        n_tot_chunks;
+    uint32_t                        n_chunks_per_task;
+    uint32_t                        n_cols;
+    uint32_t                        cur_a;
+    uint32_t                        mapping_stride;
+    size_t                          dst_nb1;
+    size_t                          dst_nb2;
+    uint32_t                        start_row;
+    uint32_t                        cne1;
+    struct htp_thread_trace        *traces;
+} output_transfer_scattered_task_state_t;
+
 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;
@@ -2112,32 +2219,89 @@ static void transfer_activation_chunk_threaded(
         int k_stride,
         int n_threads,
         int k_valid,
-        float *vtcm_f32_act) {
+        float *vtcm_f32_act,
+        size_t vtcm_f32_act_bytes) {
+    if (n_rows <= 0) {
+        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_tot_chunks      = n_tot_chunks;
-    state.n_chunks_per_task = n_chunks_per_task;
-    state.dst               = dst;
-    state.src               = src;
-    state.k_block           = k_block;
-    state.k_stride          = k_stride;
-    state.k_valid           = k_valid;
-    state.traces            = ctx->trace;
-    state.ctx               = ctx;
-    state.vtcm_f32_act      = vtcm_f32_act;
+    state.n_tasks            = (n_tot_chunks + n_chunks_per_task - 1) / n_chunks_per_task;
+    state.n_tot_chunks       = n_tot_chunks;
+    state.n_chunks_per_task  = n_chunks_per_task;
+    state.dst                = dst;
+    state.src                = src;
+    state.k_block            = k_block;
+    state.k_stride           = k_stride;
+    state.k_valid            = k_valid;
+    state.traces             = ctx->trace;
+    state.ctx                = ctx;
+    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.dma_step_rows      = dma_step_rows;
+    state.dma_step_rows_shift = dma_step_rows_shift;
 
     if (state.n_tasks == 1 || n_threads == 1) {
         transfer_activation_chunk_worker_fn(1, 0, &state);
     } else {
-        int n_tasks = hex_smin((int) state.n_tasks, n_threads);
-        worker_pool_run_func(ctx->worker_pool, transfer_activation_chunk_worker_fn, &state, n_tasks);
+        worker_pool_run_func(ctx->worker_pool, transfer_activation_chunk_worker_fn, &state, active_threads);
     }
 }
+// --- Async HMX matmul job (for pipeline overlap) ---
+
+typedef struct {
+    __fp16 *       output;
+    const __fp16 * activation;
+    const __fp16 * weight;
+    const __fp16 * scales;
+    uint32_t       n_row_tiles;
+    uint32_t       n_col_tiles;
+    uint32_t       n_dot_tiles;
+} hmx_matmul_job_t;
+
+static void hmx_matmul_worker_fn(void * data) {
+    hmx_matmul_job_t * job = (hmx_matmul_job_t *) data;
+    FARF(HIGH, "hmx-mm-job: n_row_tiles %u n_col_tiles %u n_dot_tiles %u", job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles);
+    core_dot_chunk_fp16(job->output, job->activation, job->weight, job->scales, job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles);
+}
+
+static inline void hmx_matmul_job_init(hmx_matmul_job_t * job,
+                                       __fp16 *           output,
+                                       const __fp16 *     activation,
+                                       const __fp16 *     weight,
+                                       const __fp16 *     scales,
+                                       uint32_t           n_row_tiles,
+                                       uint32_t           n_col_tiles,
+                                       uint32_t           n_dot_tiles) {
+    job->output      = output;
+    job->activation  = activation;
+    job->weight      = weight;
+    job->scales      = scales;
+    job->n_row_tiles = n_row_tiles;
+    job->n_col_tiles = n_col_tiles;
+    job->n_dot_tiles = n_dot_tiles;
+}
 
 static int hmx_mm_2d_f32(struct htp_context *ctx,
                                   float *restrict dst,
@@ -2198,48 +2362,33 @@ static int hmx_mm_2d_f32(struct htp_context *ctx,
 
     const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0;
 
-    const size_t act_f32_size     = hex_align_up((size_t)act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), HTP_MM_HMX_TILE_SIZE);
+    struct htp_mm_hmx_vtcm_layout L;
+    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, false, pipeline, act_threads, aligned_tile_size);
 
-    const size_t weight_area_size = is_quant
-        ? hex_align_up((n_chunk_n_cols / 32) * n_k_tiles * aligned_tile_size, HTP_MM_HMX_TILE_SIZE)
-        : hex_align_up(n_chunk_n_cols * row_stride, HTP_MM_HMX_TILE_SIZE);
-    const size_t act_area_size    = hex_align_up(m_chunk_n_rows * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
-    const size_t output_area_size = hex_align_up(m_chunk_n_rows * n_chunk_n_cols * sizeof(__fp16), HTP_MM_HMX_TILE_SIZE);
-
-    size_t scratch0_size, scratch1_size, scratch2_size;
-    scratch0_size = hex_align_up(n_chunk_n_cols * vec_dot_size, HTP_MM_HMX_TILE_SIZE);  // dequant buf 0
-    scratch1_size = pipeline ? scratch0_size : 0;                                 // dequant buf 1
-    scratch2_size = pipeline ? output_area_size : 0;                              // output  buf 1
-
-    uint8_t *vtcm_ptr        = (uint8_t *) ctx->vtcm_base;
-    __fp16  *vtcm_weight_raw[2] = { NULL, NULL };
-    if (weight_area_size) {
-        if (pipeline) {
-            vtcm_weight_raw[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size);
-            vtcm_weight_raw[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size);
-        } else {
-            vtcm_weight_raw[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size);
-        }
-    }
-
-    __fp16  *vtcm_f16_act    = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, act_area_size);
-    float   *vtcm_f32_act    = (float *)  vtcm_seq_alloc(&vtcm_ptr, act_f32_size);
-    __fp16  *vtcm_output     = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, output_area_size);
-    void    *vtcm_scratch0   = vtcm_seq_alloc(&vtcm_ptr, scratch0_size);
-    void    *vtcm_scratch1   = scratch1_size ? vtcm_seq_alloc(&vtcm_ptr, scratch1_size) : NULL;
-    void    *vtcm_scratch2   = scratch2_size ? vtcm_seq_alloc(&vtcm_ptr, scratch2_size) : NULL;
-    __fp16  *vtcm_scales     = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, 256);
-
-    vtcm_used = vtcm_ptr - (uint8_t *) ctx->vtcm_base;
+    vtcm_used = L.total_bytes;
     if (vtcm_used > vtcm_budget) {
         FARF(ERROR, "hmx-mm-2d-precomputed: VTCM overflow: used %zu budget %zu, m %d k %d n %d mc %zu nc %zu",
              vtcm_used, vtcm_budget, m, k, n, m_chunk_n_rows, n_chunk_n_cols);
         return -1;
     }
 
+    uint8_t * const base = (uint8_t *) ctx->vtcm_base;
+    __fp16  *vtcm_weight_raw[2] = {
+        VTCM_LAYOUT_PTR(__fp16, base, L.off_weight[0]),
+        VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_weight[1], pipeline)
+    };
+
+    __fp16  *vtcm_f16_act    = VTCM_LAYOUT_PTR(__fp16, base, L.off_act);
+    float   *vtcm_f32_act    = VTCM_LAYOUT_PTR(float, base, L.off_act_f32);
+    __fp16  *vtcm_output     = VTCM_LAYOUT_PTR(__fp16, base, L.off_dst[0]);
+    void    *vtcm_scratch0   = VTCM_LAYOUT_PTR(void, base, L.off_scratch[0]);
+    void    *vtcm_scratch1   = VTCM_LAYOUT_PTR_OPTIONAL(void, base, L.off_scratch[1], pipeline);
+    void    *vtcm_scratch2   = VTCM_LAYOUT_PTR_OPTIONAL(void, base, L.off_dst[1], pipeline);
+    __fp16  *vtcm_scales     = VTCM_LAYOUT_PTR(__fp16, base, L.off_scales);
+
     hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00));  // scale: 1.0, bias: 0.0 in FP16
 
-    FARF(HIGH, "hmx-mm-2d-precomputed: standard : m %d k %d n %d wtype %d mc %zu nc %zu vtcm %zu/%zu",
+    FARF(HIGH, "hmx-mm-2d: m %d k %d n %d wtype %d mc %zu nc %zu vtcm %zu/%zu",
          m, k, n, weight_type, m_chunk_n_rows, n_chunk_n_cols, vtcm_used, vtcm_budget);
 
     int n_chunk_cnt = hmx_ceil_div(n, n_chunk_n_cols);
@@ -2254,107 +2403,118 @@ 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);
+            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);
 
             // 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);
+            const size_t   n_cols_A0 = hex_smin(n - 0 * n_chunk_n_cols, n_chunk_n_cols);
             const uint32_t height_A0 = is_quant ? (n_cols_A0 / 32) * n_k_tiles : n_cols_A0;
             dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight),
                            dma_dst_stride, dma_src_stride, dma_width_bytes, height_A0);
 
             if (1 < n_chunk_cnt) {
-                const size_t n_cols_A1 = hex_smin(n - 1 * n_chunk_n_cols, n_chunk_n_cols);
+                const size_t   n_cols_A1 = hex_smin(n - 1 * n_chunk_n_cols, n_chunk_n_cols);
                 const uint32_t height_A1 = is_quant ? (n_cols_A1 / 32) * n_k_tiles : n_cols_A1;
                 dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[1], weight + n_chunk_n_cols * weight_stride),
                                dma_dst_stride, dma_src_stride, dma_width_bytes, height_A1);
             }
 
-            // pop A0 -> dequantize A0 -> submit C0
-            dma_queue_pop(ctx->dma[0]);
-            dequantize_tiled_weight_chunk_to_fp16_tiles(
-                ctx, vtcm_weight_bufs[0], vtcm_weight_raw[0],
-                n_cols_A0, k, row_stride, weight_type,
-                n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads);
-
-            hmx_matmul_job_init(&job_slots[0], (__fp16 *) vtcm_output_bufs[0], (__fp16 *) vtcm_f16_act,
-                                (__fp16 *) vtcm_weight_bufs[0], vtcm_scales,
-                                hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS),
-                                hmx_ceil_div(n_cols_A0, HTP_MM_HMX_TILE_N_COLS), k / HTP_MM_HMX_TILE_N_ROWS);
-            hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job_slots[0]));
-
-            // Main loop: pop/dequantize A_{i+1} -> push A_{i+2} -> submit C_{i+1} -> wait C_i and store D_i
+            // Main loop: pop A_i -> dequantize A_i -> push A_{i+2} -> submit C_i -> wait C_{i-1} and store D_{i-1}
             for (int i = 0; i < n_chunk_cnt; ++i) {
                 const size_t nc    = i * n_chunk_n_cols;
-                const size_t nc_p1 = nc + 1 * n_chunk_n_cols;
                 const size_t nc_p2 = nc + 2 * n_chunk_n_cols;
 
                 const size_t n_cols    = hex_smin(n - nc, n_chunk_n_cols);
-                const size_t n_cols_p1 = hex_smin(n - nc_p1, n_chunk_n_cols);
                 const size_t n_cols_p2 = hex_smin(n - nc_p2, n_chunk_n_cols);
 
-                // 1. pop A_{i+1} and dequantize it (if i+1 < n_chunk_cnt)
-                if (i + 1 < n_chunk_cnt) {
-                    dma_queue_pop(ctx->dma[0]);
-                    dequantize_tiled_weight_chunk_to_fp16_tiles(
-                        ctx, vtcm_weight_bufs[(i + 1) % 2], vtcm_weight_raw[(i + 1) % 2],
-                        n_cols_p1, k, row_stride, weight_type,
-                        n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads);
-                }
+                // 1. pop A_i
+                void * curr_raw = dma_queue_pop(ctx->dma[0]).dst;
 
-                // 2. push A_{i+2} (if i+2 < n_chunk_cnt)
+                // 2. dequantize A_i
+                dequantize_tiled_weight_chunk_to_fp16_tiles(
+                    ctx, vtcm_weight_bufs[i % 2], curr_raw,
+                    n_cols, k, row_stride, weight_type,
+                    n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads);
+
+                // 3. push A_{i+2} (if i+2 < n_chunk_cnt)
                 if (i + 2 < n_chunk_cnt) {
                     const uint32_t height_p2 = is_quant ? (n_cols_p2 / 32) * n_k_tiles : n_cols_p2;
-                    dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[(i + 2) % 2], weight + nc_p2 * weight_stride),
+                    dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_p2 * weight_stride),
                                    dma_dst_stride, dma_src_stride, dma_width_bytes, height_p2);
                 }
 
-                // 3. submit C_{i+1} (if i+1 < n_chunk_cnt)
-                if (i + 1 < n_chunk_cnt) {
-                    hmx_matmul_job_init(&job_slots[(i + 1) % 2], (__fp16 *) vtcm_output_bufs[(i + 1) % 2],
-                                        (__fp16 *) vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[(i + 1) % 2],
-                                        vtcm_scales, hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS),
-                                        hmx_ceil_div(n_cols_p1, HTP_MM_HMX_TILE_N_COLS), k / HTP_MM_HMX_TILE_N_ROWS);
-                    hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job_slots[(i + 1) % 2]));
-                }
+                // 4. submit C_i
+                hmx_matmul_job_init(&job_slots[i % 2], (__fp16 *) vtcm_output_bufs[i % 2],
+                                    (__fp16 *) vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[i % 2],
+                                    vtcm_scales, hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS),
+                                    hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS), k / HTP_MM_HMX_TILE_N_ROWS);
+                hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job_slots[i % 2]));
 
-                // 4. wait C_i and store D_i (multi-thread HVX, parallel with C_{i+1})
-                hmx_queue_pop(ctx->hmx_queue);
-                float *output_chunk = dst + (mr * dst_stride + nc);
-                const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc) : NULL;
-                int chunk_dst_cols = dst_cols - (int)nc;
-                if (chunk_dst_cols > 0) {
-                    transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[i % 2], n_rows, n_cols, dst_stride, src2_stride, chunk_dst_cols, n_threads);
+                // 5. wait C_{i-1} and store D_{i-1} (multi-thread HVX, parallel with C_i)
+                if (i > 0) {
+                    hmx_queue_pop(ctx->hmx_queue);
+                    const size_t nc_prev = (i - 1) * n_chunk_n_cols;
+                    const size_t n_cols_prev = hex_smin(n - nc_prev, n_chunk_n_cols);
+                    float *output_chunk = dst + (mr * dst_stride + nc_prev);
+                    const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc_prev) : NULL;
+                    int chunk_dst_cols = dst_cols - (int)nc_prev;
+                    if (chunk_dst_cols > 0) {
+                        transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[(i - 1) % 2], n_rows, n_cols_prev, dst_stride, src2_stride, chunk_dst_cols, n_threads);
+                    }
                 }
             }
+
+            // Epilogue: wait C_{last} and store D_{last}
+            hmx_queue_pop(ctx->hmx_queue);
+            const size_t nc_last = (n_chunk_cnt - 1) * n_chunk_n_cols;
+            const size_t n_cols_last = hex_smin(n - nc_last, n_chunk_n_cols);
+            float *output_chunk = dst + (mr * dst_stride + nc_last);
+            const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc_last) : NULL;
+            int chunk_dst_cols = dst_cols - (int)nc_last;
+            if (chunk_dst_cols > 0) {
+                transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[(n_chunk_cnt - 1) % 2], n_rows, n_cols_last, dst_stride, src2_stride, chunk_dst_cols, n_threads);
+            }
         }
-        hmx_queue_suspend(ctx->hmx_queue);
     } else {
-        // --- Synchronous Un-pipelined loop (m <= 32 or fallback) ---
-        HAP_compute_res_hmx_lock(ctx->vtcm_rctx);
+        // --- Synchronous loop (m <= 32 or fallback) ---
+        hmx_matmul_job_t job;
         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);
+            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);
+
+            // A0: Pre-fetch the first weight chunk (nc = 0)
+            if (n > 0) {
+                const size_t n_cols = hex_smin(n, n_chunk_n_cols);
+                const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols;
+                dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height);
+            }
 
             for (size_t nc = 0; nc < n; nc += n_chunk_n_cols) {
                 const size_t n_cols = hex_smin(n - nc, n_chunk_n_cols);
                 const size_t n_row_tiles = hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS);
                 const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS);
 
-                // A: Weight DMA (Synchronous)
-                const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols;
-                dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight + nc * weight_stride),
-                               dma_dst_stride, dma_src_stride, dma_width_bytes, height);
-                dma_queue_pop(ctx->dma[0]);
+                // A: Wait for weight DMA
+                void * curr_raw = dma_queue_pop(ctx->dma[0]).dst;
 
                 // B: Weight Dequantize (Threaded)
                 dequantize_tiled_weight_chunk_to_fp16_tiles(
-                    ctx, vtcm_scratch0, vtcm_weight_raw[0],
+                    ctx, vtcm_scratch0, curr_raw,
                     n_cols, k, row_stride, weight_type,
                     n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads);
 
-                // C: HMX Compute (Synchronous)
-                core_dot_chunk_fp16(vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS);
+                // Start weight DMA for the next chunk early
+                const size_t nc_next = nc + n_chunk_n_cols;
+                if (nc_next < n) {
+                    const size_t n_cols_next = hex_smin(n - nc_next, n_chunk_n_cols);
+                    const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next;
+                    dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next);
+                }
+
+                // C: HMX Compute (Queue-based)
+                hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS);
+                hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job));
+                hmx_queue_pop(ctx->hmx_queue);
 
                 // D: Output Store
                 float *output_chunk = dst + (mr * dst_stride + nc);
@@ -2365,7 +2525,6 @@ static int hmx_mm_2d_f32(struct htp_context *ctx,
                 }
             }
         }
-        HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
     }
 
     return 0;
@@ -2458,37 +2617,34 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_
     size_t n_chunk_n_cols = n_chunk;
     size_t vtcm_used = vtcm_size;
 
-    const size_t act_head_stride      = m_chunk_n_rows * (size_t) params->k;  // fp16 elements between heads
-    const size_t weight_area_size     = hex_align_up(n_chunk_n_cols * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
-    const size_t activation_area_size = hex_align_up(group_size * m_chunk_n_rows * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
-    const size_t output_area_size     = hex_align_up(m_chunk_n_rows * n_chunk_n_cols * sizeof(__fp16), HTP_MM_HMX_TILE_SIZE);
-    const size_t scratch_area_size    = hex_align_up(n_chunk_n_cols * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
+    struct htp_mm_hmx_vtcm_layout L;
+    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, HTP_TYPE_F16, params->k, m_chunk_n_rows, n_chunk_n_cols, group_size, use_dma_activation, false, act_threads, 0);
 
-    uint8_t *vtcm_ptr        = (uint8_t *) ctx->vtcm_base;
-    __fp16  *vtcm_weight     = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size);
-    __fp16  *vtcm_f16_act    = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, activation_area_size);
-    __fp16  *vtcm_output     = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, output_area_size);
-    void    *vtcm_scratch0   = vtcm_seq_alloc(&vtcm_ptr, scratch_area_size);
-    void    *vtcm_scratch1   = vtcm_seq_alloc(&vtcm_ptr, scratch_area_size);
-    __fp16  *vtcm_scales     = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, 256);
-    float   *vtcm_f32_act    = use_dma_activation ? (float *) vtcm_seq_alloc(&vtcm_ptr, f32_scratch_size) : NULL;
-
-    if ((size_t) (vtcm_ptr - (uint8_t *) ctx->vtcm_base) > vtcm_budget) {
+    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);
     }
 
+    uint8_t * const base = (uint8_t *) ctx->vtcm_base;
+    __fp16  *vtcm_weight     = VTCM_LAYOUT_PTR(__fp16, base, L.off_weight[0]);
+    __fp16  *vtcm_f16_act    = VTCM_LAYOUT_PTR(__fp16, base, L.off_act);
+    __fp16  *vtcm_output     = VTCM_LAYOUT_PTR(__fp16, base, L.off_dst[0]);
+    void    *vtcm_scratch0   = VTCM_LAYOUT_PTR(void, base, L.off_scratch[0]);
+    void    *vtcm_scratch1   = VTCM_LAYOUT_PTR(void, base, L.off_scratch[1]);
+    __fp16  *vtcm_scales     = VTCM_LAYOUT_PTR(__fp16, base, L.off_scales);
+    float   *vtcm_f32_act    = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_act_f32, use_dma_activation);
+
     hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00));  // scale: 1.0, bias: 0.0 in FP16
 
     FARF(HIGH, "%s: grouped path m=%d k=%d n=%d group=%d streams=%d mc=%zu nc=%zu vtcm=%zu/%zu",
             __func__, params->m, params->k, params->n, group_size, params->ne13,
             m_chunk_n_rows, n_chunk_n_cols,
-            (size_t) (vtcm_ptr - (uint8_t *) ctx->vtcm_base), vtcm_budget);
+            L.total_bytes, vtcm_budget);
 
     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);
 
-    HAP_compute_res_hmx_lock(ctx->vtcm_rctx);
+    hmx_matmul_job_t job;
 
     for (int b3 = 0; b3 < params->ne13; ++b3) {
         for (int b2_base = 0; b2_base < params->ne12; b2_base += group_size) {
@@ -2505,58 +2661,59 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_
                 // thrashing from HVX loads at large strides.
                 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 * act_head_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);
+                                                           params->k, params->act_stride, act_threads, params->k, vtcm_f32_act, L.act_f32_bytes);
                 }
 
-                void *buf_curr = vtcm_scratch0;
-                void *buf_next = vtcm_scratch1;
-
+                // Prologue: Push A0 and A1 (if exists)
                 {
                     const size_t n_cols_first = hex_smin((size_t) params->n, n_chunk_n_cols);
-                    dma_queue_push(ctx->dma[0], dma_make_ptr(buf_curr, weight_group),
+                    dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_scratch0, weight_group),
                                       fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_first);
                 }
+                if (n_chunk_n_cols < (size_t) params->n) {
+                    const size_t n_cols_second = hex_smin((size_t) params->n - n_chunk_n_cols, n_chunk_n_cols);
+                    dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_scratch1, weight_group + params->weight_stride),
+                                      fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_second);
+                }
 
                 for (size_t nc = 0; nc < (size_t) params->n; nc += n_chunk_n_cols) {
-                    const size_t n_cols = hex_smin((size_t) params->n - nc, n_chunk_n_cols);
+                    const size_t n_cols      = hex_smin((size_t) params->n - nc, n_chunk_n_cols);
                     const size_t n_col_tiles = hmx_ceil_div((int) n_cols, HTP_MM_HMX_TILE_N_COLS);
 
                     {
-                        dma_queue_pop(ctx->dma[0]);
+                        void * curr_raw = dma_queue_pop(ctx->dma[0]).dst;
 
-                        const size_t nc_next = nc + n_chunk_n_cols;
+                        hmx_interleave_rows_to_tiles(vtcm_weight, (const __fp16 *) curr_raw, n_cols, params->k, params->k, 0, n_cols);
+
+                        const size_t nc_next = nc + n_chunk_n_cols * 2;
                         if (nc_next < (size_t) params->n) {
                             const size_t n_cols_next = hex_smin((size_t) params->n - nc_next, n_chunk_n_cols);
                             const __fp16 *next_weight_chunk = weight_group + nc_next * params->weight_stride;
 
-                            dma_queue_push(ctx->dma[0], dma_make_ptr(buf_next, next_weight_chunk),
+                            dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, next_weight_chunk),
                                               fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_next);
                         }
-
-                        hmx_interleave_rows_to_tiles(vtcm_weight, (const __fp16 *) buf_curr, n_cols, params->k, params->k, 0, n_cols);
-                        hex_swap_ptr(&buf_curr, &buf_next);
                     }
 
                     // Reuse the interleaved weight for every q_head in this GQA group
                     for (int g = 0; g < group_size; ++g) {
-                        struct htp_thread_trace * tr = &ctx->trace[HTP_MAX_NTHREADS];
-                        htp_trace_event_start(tr, HTP_TRACE_EVT_HMX_COMP, g);
                         {
-                            const __fp16 * vtcm_act_g = vtcm_f16_act + (size_t) g * act_head_stride;
-                            core_dot_chunk_fp16(vtcm_output, vtcm_act_g, vtcm_weight, vtcm_scales, n_row_tiles, n_col_tiles,
-                                                params->k / 32);
+                            const __fp16 * vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride;
+                            hmx_matmul_job_init(&job, vtcm_output, vtcm_act_g, vtcm_weight, vtcm_scales, n_row_tiles, n_col_tiles, params->k / 32);
+                            hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job));
+                            hmx_queue_pop(ctx->hmx_queue);
                         }
-                        htp_trace_event_stop(tr, HTP_TRACE_EVT_HMX_COMP, g);
 
                         {
                             float *output = hmx_mm_dst_batch_ptr(params, b2_base + g, b3) + mr * params->dst_stride + nc;
                             const float *src2_chunk = params->src2 ? (hmx_mm_src2_batch_ptr(params, b2_base + g, b3) + mr * params->src2_stride + nc) : NULL;
                             int chunk_dst_cols = params->n - (int)nc;
                             if (chunk_dst_cols > 0) {
-                                transfer_output_chunk_threaded(ctx, output, src2_chunk, vtcm_output, (int) n_rows, (int) n_cols, params->dst_stride, params->src2_stride, chunk_dst_cols, ctx->n_threads);
+                                transfer_output_chunk_threaded(ctx, output, src2_chunk, vtcm_output, (int) n_rows, (int) n_cols,
+                                                               params->dst_stride, params->src2_stride, chunk_dst_cols, ctx->n_threads);
                             }
                         }
                     }
@@ -2565,8 +2722,6 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_
         }
     }
 
-    HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
-
     return 0;
 }
 
@@ -2758,7 +2913,7 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx,
 
     hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00));
 
-    HAP_compute_res_hmx_lock(ctx->vtcm_rctx);
+    hmx_matmul_job_t job;
 
     for (size_t mr = 0; mr < (size_t) m_padded; mr += m_chunk_n_rows) {
         const size_t n_rows = hex_smin(m_padded - mr, m_chunk_n_rows);
@@ -2768,37 +2923,52 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx,
             ctx, vtcm_f16_act, activation, (int) mr, (int) n_rows, k,
             matrix_rows, cur_a, mapping_stride, ne11, act_nb1, act_nb2, cne1, n_threads, k_valid);
 
+        // A0: Pre-fetch the first weight chunk (nc = 0)
+        if (n > 0) {
+            const size_t n_cols = hex_smin((size_t) n, n_chunk_n_cols);
+            const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols;
+            dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight, weight),
+                           dma_dst_stride, dma_src_stride, dma_width_bytes, height);
+        }
+
         for (size_t nc = 0; nc < (size_t) n; nc += n_chunk_n_cols) {
             const size_t n_cols = hex_smin((size_t) n - nc, n_chunk_n_cols);
             const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS);
 
-            const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols;
-            dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight, weight + nc * weight_stride),
-                           dma_dst_stride, dma_src_stride, dma_width_bytes, height);
-            dma_queue_pop(ctx->dma[0]);
+            // A: Wait for weight DMA
+            void * curr_raw = dma_queue_pop(ctx->dma[0]).dst;
 
+            // B: Weight Dequantize (Threaded)
             dequantize_tiled_weight_chunk_to_fp16_tiles(
-                ctx, vtcm_scratch0, vtcm_weight,
+                ctx, vtcm_scratch0, curr_raw,
                 n_cols, k, row_stride, weight_type,
                 n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads
             );
 
-            struct htp_thread_trace * tr = &ctx->trace[HTP_MAX_NTHREADS];
-            htp_trace_event_start(tr, HTP_TRACE_EVT_HMX_COMP, nc);
-            core_dot_chunk_fp16(vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS);
-            htp_trace_event_stop(tr, HTP_TRACE_EVT_HMX_COMP, nc);
+            // Start weight DMA for the next chunk early
+            const size_t nc_next = nc + n_chunk_n_cols;
+            if (nc_next < (size_t) n) {
+                const size_t n_cols_next = hex_smin((size_t) n - nc_next, n_chunk_n_cols);
+                const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next;
+                dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_next * weight_stride),
+                               dma_dst_stride, dma_src_stride, dma_width_bytes, height_next);
+            }
 
+            // C: HMX Compute (Queue-based)
+            hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS);
+            hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job));
+            hmx_queue_pop(ctx->hmx_queue);
+
+            // D: Output Store
             transfer_output_chunk_scattered_threaded(
                 ctx, dst + nc, vtcm_output, (int) mr, (int) n_rows, (int) n_cols,
                 matrix_rows, cur_a, mapping_stride, dst_nb1, dst_nb2, cne1, n_threads);
         }
     }
 
-    HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
     return 0;
 }
 
-
 // --- Dispatchers and Public Entry Points ---
 
 static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_kernel_params * kparams) {
@@ -2960,22 +3130,14 @@ static int hvx_mm_matmul_id(
     }
     size_t src1_row_size  = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
 
-    // Scratchpad sizes are computed on the host (htp_mm_hvx_id_get_vtcm_sizes) and passed in.
-    // The ID layout is routing-independent, so the host has exact visibility -- consume it here
-    // rather than recomputing, to keep host budgeting and device allocation in lockstep.
-    size_t src0_sz = kparams->vtcm_src0_size;
-    size_t src1_sz = kparams->vtcm_src1_size;
-    size_t src2_sz = 0; // mapping lives in DDR
-    size_t dst_sz  = kparams->vtcm_dst_size;
-    size_t vtcm_size = kparams->vtcm_size;
+    struct htp_mm_hvx_vtcm_layout L;
+    htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
+                                 0, src0_row_size, src1_row_size, kparams->n_prefetch, true, false, false);
 
-    size_t src0_sz_per_thread = src0_sz / octx->n_threads;
-    size_t src1_sz_per_thread = src1_sz;
-    size_t src2_sz_per_thread = 0;
-    size_t dst_sz_per_thread  = dst_sz / octx->n_threads;
+    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
 
-    FARF(HIGH, "matmul-id-%s : src0-spad-size %zu src1-spad-size %zu src2-spad-size %zu dst-spad-size %zu (%zu)\n", mmctx->type,
-         src0_sz, src1_sz, src2_sz, dst_sz, vtcm_size);
+    FARF(HIGH, "matmul-id-%s : src0-spad-size %zu src1-spad-size %zu src2-spad-size 0 dst-spad-size %zu (%zu)\n", mmctx->type,
+         L.src0_bytes, L.src1_bytes, L.dst_bytes, vtcm_size);
 
     FARF(HIGH, "matmul-id-%s : %ux%ux%ux%u * %ux%ux%ux%u (%ux%ux%ux%u) -> %ux%ux%ux%u (0x%p, 0x%p, 0x%p)\n", mmctx->type,
          src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3],
@@ -2989,11 +3151,11 @@ static int hvx_mm_matmul_id(
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
-    mmctx->vtcm_src1 = vtcm_seq_alloc(&vtcm_ptr, src1_sz);
-    mmctx->vtcm_src0 = vtcm_seq_alloc(&vtcm_ptr, src0_sz);
-    mmctx->vtcm_src2 = vtcm_seq_alloc(&vtcm_ptr, src2_sz);
-    mmctx->vtcm_dst  = vtcm_seq_alloc(&vtcm_ptr, dst_sz);
+    uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
+    mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+    mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    mmctx->vtcm_src2 = NULL;
+    mmctx->vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
 
     octx->src1_spad.src  = NULL;
     octx->src0_spad.src  = NULL;
@@ -3003,10 +3165,10 @@ static int hvx_mm_matmul_id(
     mmctx->vtcm_src0_stride = src0_row_size_padded;
     mmctx->vtcm_src1_stride = src1_row_size;
 
-    mmctx->vtcm_src0_size_per_thread = src0_sz_per_thread;
-    mmctx->vtcm_src1_size_per_thread = src1_sz_per_thread;
-    mmctx->vtcm_src2_size_per_thread = src2_sz_per_thread;
-    mmctx->vtcm_dst_size_per_thread  = dst_sz_per_thread;
+    mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads;
+    mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
+    mmctx->vtcm_src2_size_per_thread = 0;
+    mmctx->vtcm_dst_size_per_thread  = L.dst_bytes / octx->n_threads;
 
     mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
     mmctx->quant_task_func = quant_task_func;
@@ -3181,19 +3343,11 @@ int op_matmul_qkv(struct htp_ops_context * octx) {
         src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(src1->ne[0]) : htp_mm_q8_0_tiled_row_size(src1->ne[0]);
     }
 
-    // Set up scratchpads using precomputed sizes from the host
-    size_t src0_sz = kparams->vtcm_src0_size;
-    size_t src1_sz = kparams->vtcm_src1_size;
-    size_t src2_sz = kparams->vtcm_src2_size;
-    size_t src3_sz = kparams->vtcm_src3_size;
-    size_t dst_sz  = kparams->vtcm_dst_size;
-    size_t vtcm_size = kparams->vtcm_size;
+    struct htp_mm_hvx_vtcm_layout L;
+    htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads,
+                                 0, src0_row_size, src1_row_size, kparams->n_prefetch, false, true, false);
 
-    size_t src0_sz_per_thread = src0_sz / octx->n_threads;
-    size_t src1_sz_per_thread = src1_sz;
-    size_t src2_sz_per_thread = src2_sz / octx->n_threads;
-    size_t src3_sz_per_thread = src3_sz / octx->n_threads;
-    size_t dst_sz_per_thread  = dst_sz / octx->n_threads;
+    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
 
     if (octx->ctx->vtcm_size < vtcm_size) {
         FARF(ERROR, "matmul-qkv: current VTCM reservation %zu is too small, needed %zu\n",
@@ -3201,12 +3355,12 @@ int op_matmul_qkv(struct htp_ops_context * octx) {
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
-    mmctx->vtcm_src1 = vtcm_seq_alloc(&vtcm_ptr, src1_sz);
-    mmctx->vtcm_src0 = vtcm_seq_alloc(&vtcm_ptr, src0_sz);
-    mmctx->vtcm_src2 = vtcm_seq_alloc(&vtcm_ptr, src2_sz);
-    mmctx->vtcm_src3 = vtcm_seq_alloc(&vtcm_ptr, src3_sz);
-    mmctx->vtcm_dst  = vtcm_seq_alloc(&vtcm_ptr, dst_sz);
+    uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
+    mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+    mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2);
+    mmctx->vtcm_src3 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src3);
+    mmctx->vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
 
     octx->src1_spad.src  = NULL;
     octx->src0_spad.src  = NULL;
@@ -3219,11 +3373,11 @@ int op_matmul_qkv(struct htp_ops_context * octx) {
     mmctx->vtcm_src3_stride = is_repacked ? 0 : src0_row_size_padded;
     mmctx->vtcm_src1_stride = src1_row_size;
 
-    mmctx->vtcm_src0_size_per_thread = src0_sz_per_thread;
-    mmctx->vtcm_src1_size_per_thread = src1_sz_per_thread;
-    mmctx->vtcm_src2_size_per_thread = src2_sz_per_thread;
-    mmctx->vtcm_src3_size_per_thread = src3_sz_per_thread;
-    mmctx->vtcm_dst_size_per_thread  = dst_sz_per_thread;
+    mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads;
+    mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
+    mmctx->vtcm_src2_size_per_thread = L.src2_bytes / octx->n_threads;
+    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;
@@ -3331,28 +3485,22 @@ int op_matmul_ffn(struct htp_ops_context * octx) {
         src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(src1->ne[0]) : htp_mm_q8_0_tiled_row_size(src1->ne[0]);
     }
 
-    // Set up scratchpads using precomputed sizes from the host
-    size_t src0_sz = kparams->vtcm_src0_size;
-    size_t src1_sz = kparams->vtcm_src1_size;
-    size_t src2_sz = kparams->vtcm_src2_size;
-    size_t dst_sz  = kparams->vtcm_dst_size;
-    size_t vtcm_size = kparams->vtcm_size;
+    struct htp_mm_hvx_vtcm_layout L;
+    htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads,
+                                 0, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, true);
 
-    size_t src0_sz_per_thread = src0_sz / octx->n_threads;
-    size_t src1_sz_per_thread = src1_sz;
-    size_t src2_sz_per_thread = src2_sz / octx->n_threads;
-    size_t dst_sz_per_thread  = dst_sz / octx->n_threads;
+    size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
 
     if (octx->ctx->vtcm_size < vtcm_size) {
         FARF(ERROR, "matmul-ffn: current VTCM reservation %zu is too small, needed %zu\n", octx->ctx->vtcm_size, vtcm_size);
         return HTP_STATUS_VTCM_TOO_SMALL;
     }
 
-    uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
-    mmctx->vtcm_src1 = vtcm_seq_alloc(&vtcm_ptr, src1_sz);
-    mmctx->vtcm_src0 = vtcm_seq_alloc(&vtcm_ptr, src0_sz);
-    mmctx->vtcm_src2 = vtcm_seq_alloc(&vtcm_ptr, src2_sz);
-    mmctx->vtcm_dst  = vtcm_seq_alloc(&vtcm_ptr, dst_sz);
+    uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
+    mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+    mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+    mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2);
+    mmctx->vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
 
     octx->src1_spad.src  = NULL;
     octx->src0_spad.src  = NULL;
@@ -3363,10 +3511,10 @@ int op_matmul_ffn(struct htp_ops_context * octx) {
     mmctx->vtcm_src2_stride = is_repacked ? 0 : src0_row_size_padded;
     mmctx->vtcm_src1_stride = src1_row_size;
 
-    mmctx->vtcm_src0_size_per_thread = src0_sz_per_thread;
-    mmctx->vtcm_src1_size_per_thread = src1_sz_per_thread;
-    mmctx->vtcm_src2_size_per_thread = src2_sz_per_thread;
-    mmctx->vtcm_dst_size_per_thread  = dst_sz_per_thread;
+    mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads;
+    mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
+    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;

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

@@ -6,6 +6,7 @@
 #include "htp-ops.h"
 #include "hex-fastdiv.h"
 #include "hex-common.h"
+#include "htp-vtcm.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -44,7 +45,7 @@ extern "C" {
 
 // --- DMA Activation Transfer Configuration ---
 #define HTP_MM_DMA_ACT_ROWS_PER_STEP 2
-#define HTP_MM_DMA_ACT_MULTIPLIER    4
+#define HTP_MM_DMA_ACT_MULTIPLIER    (2 * HTP_MM_DMA_ACT_ROWS_PER_STEP)
 
 enum htp_mm_kernel_type {
     HTP_MM_KERNEL_UNSUPPORTED = 0,
@@ -295,197 +296,351 @@ static inline void htp_mm_hmx_get_batched_chunk_costs(
     *size_per_mn_out = sizeof(uint16_t);
 }
 
-static inline size_t htp_mm_hmx_get_2d_vtcm_size(
-    int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size
+struct htp_mm_hmx_vtcm_layout {
+    // Byte offsets from vtcm_base for each region
+    size_t off_weight[2];     // [1] is only used when pipelined
+    size_t off_act;
+    size_t off_act_f32;       // fp32 activation conversion scratch
+    size_t off_dst[2];        // [1] is only used when pipelined
+    size_t off_scratch[2];    // dequantization scratch pads
+    size_t off_scales;        // HMX scales (256 bytes)
+
+    // Cached sizes of regions for HMX kernel use
+    size_t weight_area_bytes;
+    size_t act_area_bytes;
+    size_t act_f32_bytes;
+    size_t output_area_bytes;
+    size_t scratch_bytes[2];
+    size_t act_head_stride;
+
+    size_t total_bytes;
+};
+
+struct htp_mm_hvx_vtcm_layout {
+    // Byte offsets from vtcm_base for each region
+    size_t off_src1;          // vtcm_src1 (activation)
+    size_t off_src0;          // vtcm_src0 (weight/Wk)
+    size_t off_src2;          // vtcm_src2 (Wq / fused only)
+    size_t off_src3;          // vtcm_src3 (Wv / fused only)
+    size_t off_dst;           // vtcm_dst (output scratch)
+
+    // Cached sizes
+    size_t src0_bytes;
+    size_t src1_bytes;
+    size_t src2_bytes;
+    size_t src3_bytes;
+    size_t dst_bytes;
+
+    size_t total_bytes;
+};
+
+static inline void htp_mm_hmx_vtcm_layout_build(
+    struct htp_mm_hmx_vtcm_layout * L,
+    int kernel_type,
+    int wtype,
+    uint32_t k,
+    size_t mc,
+    size_t nc,
+    uint32_t group_size,
+    bool use_dma_activation,
+    bool pipeline,
+    uint32_t act_threads,
+    uint32_t aligned_tile_size
 ) {
-    const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS;
-    const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32);
-    const size_t row_stride = htp_mm_get_tiled_row_stride(wtype, k);
-    const size_t vec_dot_size = k * sizeof(uint16_t);
+    size_t off = 0;
 
-    const size_t act_f32_size = htp_mm_round_up(act_threads * 4 * k * sizeof(float), HTP_MM_HMX_TILE_SIZE);
-    size_t weight_area_size = is_quant
-        ? htp_mm_round_up((nc / 32) * n_k_tiles * aligned_tile_size, HTP_MM_HMX_TILE_SIZE)
-        : htp_mm_round_up(nc * row_stride, HTP_MM_HMX_TILE_SIZE);
-    if (pipeline) {
-        weight_area_size *= 2;
+    if (kernel_type == HTP_MM_KERNEL_HMX_F16_BATCHED) {
+        const size_t vec_dot_size     = k * sizeof(uint16_t);
+        const size_t act_head_stride   = mc * k;
+        const size_t weight_area_size  = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
+        const size_t activation_area_size = hex_align_up(group_size * act_head_stride * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE);
+        const size_t output_area_size  = hex_align_up(group_size * mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE);
+        const size_t scratch_area_size = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
+        const size_t min_f32_size = use_dma_activation
+            ? hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), 128) : 0;
+
+        // Group A: Permanent activation tiles and scales
+        size_t off_group_a = 0;
+        VTCM_LAYOUT_ALLOC(off_group_a, off_act, activation_area_size);
+        VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data
+
+        // Group B: Compute-only buffers (starts at off_group_a)
+        size_t off_group_b = off_group_a;
+        VTCM_LAYOUT_ALLOC(off_group_b, off_weight[0], weight_area_size);
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_weight[1], weight_area_size, false);
+        VTCM_LAYOUT_ALLOC(off_group_b, off_dst[0], output_area_size);
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_dst[1], output_area_size, false);
+        VTCM_LAYOUT_ALLOC(off_group_b, off_scratch[0], scratch_area_size);
+        VTCM_LAYOUT_ALLOC(off_group_b, off_scratch[1], scratch_area_size);
+
+        const size_t group_b_size = off_group_b - off_group_a;
+
+        // Group C: Activation prep temporary buffer (overlaps Group B, starting at off_group_a)
+        const size_t max_f32_size = act_threads * 64 * k * sizeof(float);
+        const size_t act_f32_size = use_dma_activation
+            ? hex_align_up(hex_smin(max_f32_size, hex_smax(min_f32_size, group_b_size)), 128) : 0;
+        size_t off_group_c = off_group_a;
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_c, off_act_f32, act_f32_size, use_dma_activation);
+
+        const size_t group_c_size = off_group_c - off_group_a;
+
+        L->weight_area_bytes = weight_area_size;
+        L->act_area_bytes    = activation_area_size;
+        L->act_f32_bytes     = act_f32_size;
+        L->output_area_bytes = output_area_size;
+        L->scratch_bytes[0]  = scratch_area_size;
+        L->scratch_bytes[1]  = scratch_area_size;
+        L->act_head_stride   = act_head_stride;
+
+        off = off_group_a + hex_smax(group_b_size, group_c_size);
+    } else {
+        // HTP_MM_KERNEL_HMX_2D
+        const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32);
+        const size_t row_stride = htp_mm_get_tiled_row_stride(wtype, k);
+        const size_t vec_dot_size = k * sizeof(uint16_t);
+        const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS;
+
+        const size_t min_f32_size = hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), 128);
+        const size_t weight_area_size = is_quant
+            ? hex_align_up((nc / 32) * n_k_tiles * aligned_tile_size, HTP_MM_HMX_TILE_SIZE)
+            : hex_align_up(nc * row_stride, HTP_MM_HMX_TILE_SIZE);
+        const size_t act_area_size    = hex_align_up(mc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
+        const size_t output_area_size = hex_align_up(mc * nc * sizeof(__fp16), HTP_MM_HMX_TILE_SIZE);
+
+        const size_t scratch0_size = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
+        const size_t scratch1_size = pipeline ? scratch0_size : 0;
+
+        // Group A:  Scales and activation tiles (must not overlap with Group B or C)
+        size_t off_group_a = 0;
+        VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data
+        VTCM_LAYOUT_ALLOC(off_group_a, off_act, act_area_size);
+
+        // Group B: Compute-only buffers (starts at off_group_a)
+        size_t off_group_b = off_group_a;
+        VTCM_LAYOUT_ALLOC(off_group_b, off_weight[0], weight_area_size);
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_weight[1], weight_area_size, pipeline);
+        VTCM_LAYOUT_ALLOC(off_group_b, off_dst[0], output_area_size);
+        VTCM_LAYOUT_ALLOC(off_group_b, off_scratch[0], scratch0_size);
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_scratch[1], scratch0_size, pipeline);
+        VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_dst[1], output_area_size, pipeline);
+
+        const size_t group_b_size = off_group_b - off_group_a;
+
+        // Group C: Activation prep temporary buffer (overlaps Group B, starting at off_group_a)
+        const size_t max_f32_size = act_threads * 64 * k * sizeof(float);
+        const size_t act_f32_size = hex_align_up(hex_smin(max_f32_size, hex_smax(min_f32_size, group_b_size)), 128);
+        size_t off_group_c = off_group_a;
+        VTCM_LAYOUT_ALLOC(off_group_c, off_act_f32, act_f32_size);
+
+        const size_t group_c_size = off_group_c - off_group_a;
+
+        L->weight_area_bytes = weight_area_size;
+        L->act_area_bytes    = act_area_size;
+        L->act_f32_bytes     = act_f32_size;
+        L->output_area_bytes = output_area_size;
+        L->scratch_bytes[0]  = scratch0_size;
+        L->scratch_bytes[1]  = scratch1_size;
+        L->act_head_stride   = 0;
+
+        off = off_group_a + hex_smax(group_b_size, group_c_size);
     }
-    const size_t act_area_size    = htp_mm_round_up(mc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
-    const size_t output_area_size = htp_mm_round_up(mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE);
 
-    size_t scratch0_size = htp_mm_round_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
-    size_t scratch1_size = pipeline ? scratch0_size : 0;
-    size_t scratch2_size = pipeline ? output_area_size : 0;
-
-    return weight_area_size + act_area_size + act_f32_size + output_area_size +
-           scratch0_size + scratch1_size + scratch2_size + 256;
+    L->total_bytes = off;
 }
 
-static inline size_t htp_mm_hmx_get_batched_vtcm_size(
-    int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool use_dma_activation, bool pipeline, uint32_t act_threads) {
-    (void)wtype;
-    (void)pipeline;
-    const size_t vec_dot_size     = k * sizeof(uint16_t);
-    const size_t f32_scratch_size = use_dma_activation
-        ? htp_mm_round_up(act_threads * 4 * k * sizeof(float), HTP_MM_HMX_TILE_SIZE) : 0;
-
-    const size_t act_head_stride   = mc * k;
-    const size_t weight_area_size  = htp_mm_round_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
-    const size_t act_area_size     = htp_mm_round_up(group_size * act_head_stride * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE);
-    const size_t output_area_size  = htp_mm_round_up(group_size * mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE);
-    const size_t scratch_area_size = htp_mm_round_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
-
-    return weight_area_size + act_area_size + output_area_size +
-           2 * scratch_area_size + 256 + f32_scratch_size;
-}
-
-static inline size_t htp_mm_hvx_get_vtcm_sizes(
+static inline void htp_mm_hvx_vtcm_layout_build(
+    struct htp_mm_hvx_vtcm_layout * L,
     int kernel_type,
     int wtype,
     uint32_t ne10,       // k
-    uint32_t src1_nrows, // m_total (or act_nrows)
+    uint32_t src1_nrows, // m_total
     uint32_t n_threads,
     size_t dst_row_size,
     size_t src0_row_size,
     size_t src1_row_size,
     uint32_t n_prefetch,
-    size_t * vtcm_src0_size_out,
-    size_t * vtcm_src1_size_out,
-    size_t * vtcm_dst_size_out
+    bool is_matmul_id,
+    bool is_fused_qkv,
+    bool is_fused_ffn
 ) {
-    size_t vtcm_src0_size = 0;
-    size_t vtcm_src1_size = 0;
-    size_t vtcm_dst_size  = 0;
+    size_t src0_sz = 0;
+    size_t src1_sz = 0;
+    size_t src2_sz = 0;
+    size_t src3_sz = 0;
+    size_t dst_sz  = 0;
 
     const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 ||
                             wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL ||
                             wtype == HTP_TYPE_MXFP4);
 
-    const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
-    const size_t dst_nrows = (src1_nrows > 1) ? 0 : 1;
+    if (is_fused_qkv || is_fused_ffn) {
+        const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);
+        const size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads;
 
-    switch (kernel_type) {
-        case HTP_MM_KERNEL_HVX_F16_F16_VTCM: {
-            size_t f16_src1_row_size = htp_mm_round_up(ne10 * 2, 128);
-            vtcm_src1_size = htp_mm_round_up(f16_src1_row_size * src1_nrows, 256);
-            vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
-            vtcm_dst_size  = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
-            break;
-        }
-        case HTP_MM_KERNEL_HVX_F16_F32_DDR:
-        case HTP_MM_KERNEL_HVX_F16_F16_DDR:
-        case HTP_MM_KERNEL_HVX_F32_F32_DDR:
-        case HTP_MM_KERNEL_HVX_F32_F16_DDR: {
-            vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size, 256) * n_threads;
-            vtcm_src1_size = htp_mm_round_up(n_prefetch * src1_row_size, 256) * n_threads;
-            vtcm_dst_size  = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
-            break;
-        }
-        case HTP_MM_KERNEL_HVX_F32_F32_VTCM: {
-            size_t f32_src1_row_size = htp_mm_round_up(ne10 * 4, 128);
-            vtcm_src1_size = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256);
-            vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
-            vtcm_dst_size  = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
-            break;
-        }
-        case HTP_MM_KERNEL_HVX_QUANT_BLOCK:
-        case HTP_MM_KERNEL_HVX_QUANT_ROW: {
-            size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
+        size_t src0_sz_per_thread = 0;
+        size_t src2_sz_per_thread = 0;
+        size_t src3_sz_per_thread = 0;
 
-            vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
-            vtcm_src1_size = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
+        if (is_repack) {
+            uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
+            uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32;
+            uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
 
-            vtcm_src0_size = vtcm_src0_size * n_threads;
-
-            if (is_repack) {
-                uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
-                uint32_t n_k_tiles = ne10 / 32;
-                uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
-                size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
-                vtcm_src0_size = repacked_vtcm_size * n_threads;
+            src0_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128);
+            src2_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128);
+            if (is_fused_qkv) {
+                src3_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128);
             }
-
-            size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
-            size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
-            if (dst_size_per_thread < quant_scratch_size_per_thread) {
-                dst_size_per_thread = quant_scratch_size_per_thread;
+        } else {
+            src0_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128);
+            src2_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128);
+            if (is_fused_qkv) {
+                src3_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128);
             }
-            vtcm_dst_size = dst_size_per_thread * n_threads;
-            break;
         }
-        case HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT: {
-            size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
 
-            vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
-            vtcm_src1_size = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
+        size_t flat_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
+        size_t tiled_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
 
-            vtcm_src0_size = vtcm_src0_size * n_threads;
-
-            if (is_repack) {
-                uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
-                uint32_t n_k_tiles = ne10 / 32;
-                uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
-                size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
-                vtcm_src0_size = repacked_vtcm_size * n_threads;
-            }
-
-            size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
-            size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
-            if (dst_size_per_thread < quant_scratch_size_per_thread) {
-                dst_size_per_thread = quant_scratch_size_per_thread;
-            }
-            vtcm_dst_size = dst_size_per_thread * n_threads;
-            break;
+        if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) {
+            src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128);
+        } else {
+            src1_sz = hex_round_up(tiled_src1_row_size * src1_nrows, 128);
+        }
+
+        src0_sz = src0_sz_per_thread * n_threads;
+        src2_sz = src2_sz_per_thread * n_threads;
+        src3_sz = src3_sz_per_thread * n_threads;
+        dst_sz  = quant_scratch_size;
+    } else if (is_matmul_id) {
+        const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
+        const size_t src1_row_size_tiled = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10)
+                                                                    : htp_mm_q8_0_tiled_row_size(ne10);
+
+        size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
+        src1_sz                   = htp_mm_round_up(src1_row_size_tiled * src1_nrows, 256);
+
+        if (is_repack) {
+            const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
+            const uint32_t n_k_tiles         = ne10 / 32;
+            const uint32_t tile_row_size     = n_k_tiles * aligned_tile_size;
+            size_t repacked_vtcm_size        = htp_mm_round_up(n_prefetch * tile_row_size, 256);
+            src0_sz_per_thread               = repacked_vtcm_size;
+        }
+
+        src0_sz = src0_sz_per_thread * n_threads;
+        dst_sz  = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads;
+    } else {
+        const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
+        const size_t dst_nrows = (src1_nrows > 1) ? 0 : 1;
+
+        switch (kernel_type) {
+            case HTP_MM_KERNEL_HVX_F16_F16_VTCM: {
+                size_t f16_src1_row_size = htp_mm_round_up(ne10 * 2, 128);
+                src1_sz = htp_mm_round_up(f16_src1_row_size * src1_nrows, 256);
+                src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
+                dst_sz  = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
+                break;
+            }
+            case HTP_MM_KERNEL_HVX_F16_F32_DDR:
+            case HTP_MM_KERNEL_HVX_F16_F16_DDR:
+            case HTP_MM_KERNEL_HVX_F32_F32_DDR:
+            case HTP_MM_KERNEL_HVX_F32_F16_DDR: {
+                src0_sz = htp_mm_round_up(n_prefetch * src0_row_size, 256) * n_threads;
+                src1_sz = htp_mm_round_up(n_prefetch * src1_row_size, 256) * n_threads;
+                dst_sz  = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
+                break;
+            }
+            case HTP_MM_KERNEL_HVX_F32_F32_VTCM: {
+                size_t f32_src1_row_size = htp_mm_round_up(ne10 * 4, 128);
+                src1_sz = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256);
+                src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
+                dst_sz  = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
+                break;
+            }
+            case HTP_MM_KERNEL_HVX_QUANT_BLOCK:
+            case HTP_MM_KERNEL_HVX_QUANT_ROW: {
+                size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
+
+                src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
+                src1_sz = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
+
+                src0_sz = src0_sz * n_threads;
+
+                if (is_repack) {
+                    uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
+                    uint32_t n_k_tiles = ne10 / 32;
+                    uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
+                    size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
+                    src0_sz = repacked_vtcm_size * n_threads;
+                }
+
+                size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+                size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
+                if (dst_size_per_thread < quant_scratch_size_per_thread) {
+                    dst_size_per_thread = quant_scratch_size_per_thread;
+                }
+                dst_sz = dst_size_per_thread * n_threads;
+                break;
+            }
+            case HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT: {
+                size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
+
+                src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
+                src1_sz = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
+
+                src0_sz = src0_sz * n_threads;
+
+                if (is_repack) {
+                    uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
+                    uint32_t n_k_tiles = ne10 / 32;
+                    uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
+                    size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
+                    src0_sz = repacked_vtcm_size * n_threads;
+                }
+
+                size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
+                size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
+                if (dst_size_per_thread < quant_scratch_size_per_thread) {
+                    dst_size_per_thread = quant_scratch_size_per_thread;
+                }
+                dst_sz = dst_size_per_thread * n_threads;
+                break;
+            }
+            default:
+                break;
         }
-        default:
-            break;
     }
 
-    *vtcm_src0_size_out = vtcm_src0_size;
-    *vtcm_src1_size_out = vtcm_src1_size;
-    *vtcm_dst_size_out  = vtcm_dst_size;
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src1, src1_sz);
+    VTCM_LAYOUT_ALLOC(off, off_src0, src0_sz);
+    VTCM_LAYOUT_ALLOC(off, off_src2, src2_sz);
+    VTCM_LAYOUT_ALLOC(off, off_src3, src3_sz);
+    VTCM_LAYOUT_ALLOC(off, off_dst,  dst_sz);
 
-    return vtcm_src0_size + vtcm_src1_size + vtcm_dst_size;
+    L->src0_bytes = src0_sz;
+    L->src1_bytes = src1_sz;
+    L->src2_bytes = src2_sz;
+    L->src3_bytes = src3_sz;
+    L->dst_bytes  = dst_sz;
+    L->total_bytes = off;
 }
 
-static inline size_t htp_mm_hvx_id_get_vtcm_sizes(
-    int wtype,
-    uint32_t ne10,                // k
-    uint32_t src1_nrows,
-    uint32_t n_threads,
-    size_t src0_row_size,    // nb01
-    uint32_t n_prefetch,
-    size_t * vtcm_src0_size_out,
-    size_t * vtcm_src1_size_out,
-    size_t * vtcm_dst_size_out
+static inline size_t htp_mm_hmx_get_2d_vtcm_size(
+    int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size
 ) {
-    const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 ||
-                            wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL ||
-                            wtype == HTP_TYPE_MXFP4);
+    struct htp_mm_hmx_vtcm_layout L;
+    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, false, pipeline, act_threads, aligned_tile_size);
+    return L.total_bytes;
+}
 
-    const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
-    const size_t src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10)
-                                                          : htp_mm_q8_0_tiled_row_size(ne10);
-
-    size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
-    size_t src1_sz            = htp_mm_round_up(src1_row_size * src1_nrows, 256);
-
-    if (is_repack) {
-        const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
-        const uint32_t n_k_tiles    = ne10 / 32;
-        const uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
-        size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
-        src0_sz_per_thread = repacked_vtcm_size;
-    }
-
-    const size_t vtcm_src0_size = src0_sz_per_thread * n_threads;
-    const size_t vtcm_dst_size  = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads;
-
-    *vtcm_src0_size_out = vtcm_src0_size;
-    *vtcm_src1_size_out = src1_sz;
-    *vtcm_dst_size_out  = vtcm_dst_size;
-
-    return vtcm_src0_size + src1_sz + vtcm_dst_size;
+static inline size_t htp_mm_hmx_get_batched_vtcm_size(
+    int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool use_dma_activation, bool pipeline, uint32_t act_threads) {
+    (void)pipeline;
+    struct htp_mm_hmx_vtcm_layout L;
+    htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, use_dma_activation, false, act_threads, 0);
+    return L.total_bytes;
 }
 
 #ifdef __cplusplus

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

@@ -1,16 +0,0 @@
-#ifndef VTCM_UTILS_H
-#define VTCM_UTILS_H
-
-#include "hex-utils.h"
-
-#include <assert.h>
-#include <stdint.h>
-#include <hexagon_types.h>
-
-static inline uint8_t *vtcm_seq_alloc(uint8_t **vtcm_ptr, size_t size) {
-    uint8_t *p = *vtcm_ptr;
-    *vtcm_ptr += size;
-    return p;
-}
-
-#endif // VTCM_UTILS_H

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

@@ -1,6 +1,9 @@
 #include "worker-pool.h"
+#include "hex-utils.h"
 
 #include <qurt.h>
+#include <qurt_hvx.h>
+
 #include <stdatomic.h>
 #include <stdint.h>
 #include <stdio.h>
@@ -9,7 +12,6 @@
 
 #include "HAP_farf.h"
 
-#define WORKER_THREAD_STACK_SZ  (2 * 16384)
 #define LOWEST_USABLE_QURT_PRIO (254)
 
 struct worker_pool_s;
@@ -42,17 +44,27 @@ static void worker_pool_main(void * context) {
     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) {
-            // Nothing to do
+            // 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;
         }
 
-        // New job
         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) {

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

@@ -24,9 +24,17 @@ typedef struct {
     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);
 

ggml/src/ggml-hip/CMakeLists.txt

@@ -156,4 +156,4 @@ endif()
 
 target_link_libraries(ggml-hip PRIVATE ggml-base hip::host roc::rocblas roc::hipblas)
 
-target_compile_options(ggml-hip PRIVATE "$<$<COMPILE_LANGUAGE:HIP>:-ffast-math>")
+target_compile_options(ggml-hip PRIVATE "$<$<COMPILE_LANGUAGE:HIP>:-ffast-math;-fno-finite-math-only>")

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

@@ -160,11 +160,15 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows(ggml_me
     return res;
 }
 
-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows(ggml_metal_library_t lib, ggml_type tidx, ggml_type tdst) {
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows(ggml_metal_library_t lib, const ggml_tensor * op) {
     char base[256];
     char name[256];
 
-    snprintf(base, 256, "kernel_set_rows_%s_%s", ggml_type_name(tdst), ggml_type_name(tidx));
+    const auto tsrc = op->src[0]->type;
+    const auto tidx = op->src[1]->type;
+    const auto tdst = op->type;
+
+    snprintf(base, 256, "kernel_set_rows_%s_%s_%s", ggml_type_name(tsrc), ggml_type_name(tidx), ggml_type_name(tdst));
     snprintf(name, 256, "%s", base);
 
     ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
@@ -1800,6 +1804,26 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_1
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_col2im_1d(ggml_metal_library_t lib, const ggml_tensor * op) {
+    assert(op->op == GGML_OP_COL2IM_1D);
+
+    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
+    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_BF16);
+
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_col2im_1d_%s", ggml_type_name(op->src[0]->type));
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+    }
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d(ggml_metal_library_t lib, const ggml_tensor * op) {
     assert(op->op == GGML_OP_CONV_TRANSPOSE_2D);
 

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

@@ -112,7 +112,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cpy
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_1d           (ggml_metal_library_t lib, const struct ggml_tensor * op, enum ggml_op_pool op_pool);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_2d           (ggml_metal_library_t lib, const struct ggml_tensor * op, enum ggml_op_pool op_pool);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows          (ggml_metal_library_t lib, enum ggml_type tsrc);
-struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows          (ggml_metal_library_t lib, enum ggml_type tidx, enum ggml_type tdst);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows          (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_diag              (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_repeat            (ggml_metal_library_t lib, enum ggml_type tsrc);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_concat            (ggml_metal_library_t lib, enum ggml_type tsrc);
@@ -150,6 +150,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rope
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_im2col            (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_1d (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_col2im_1d         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d           (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_3d           (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_upscale           (ggml_metal_library_t lib, const struct ggml_tensor * op);

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

@@ -1157,6 +1157,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                 (op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32) &&
                 op->src[1]->type == GGML_TYPE_F32 &&
                 op->type == GGML_TYPE_F32;
+        case GGML_OP_COL2IM_1D:
+            return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_BF16) &&
+                op->type == op->src[0]->type &&
+                ggml_is_contiguous(op->src[0]) &&
+                ggml_is_contiguous(op);
         case GGML_OP_CONV_3D:
             return ggml_is_contiguous(op->src[0]) &&
                    ggml_is_contiguous(op->src[1]) &&
@@ -1329,7 +1334,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
             return op->src[0]->type != GGML_TYPE_NVFP4;
         case GGML_OP_SET_ROWS:
             {
-                if (op->src[0]->type != GGML_TYPE_F32) {
+                if (op->src[0]->type != GGML_TYPE_F32 && op->src[0]->type != GGML_TYPE_F16) {
                     return false;
                 }
 

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

@@ -603,6 +603,16 @@ typedef struct {
     uint64_t nb1;
 } ggml_metal_kargs_conv_transpose_1d;
 
+typedef struct {
+    int32_t  T_in;
+    int32_t  T_out;
+    int32_t  OC;
+    int32_t  K;
+    int32_t  K_OC;
+    int32_t  s0;
+    int32_t  p0;
+} ggml_metal_kargs_col2im_1d;
+
 typedef struct {
     int32_t  IC;
     int32_t  IH;

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

@@ -395,6 +395,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
             {
                 n_fuse = ggml_metal_op_conv_transpose_2d(ctx, idx);
             } break;
+        case GGML_OP_COL2IM_1D:
+            {
+                n_fuse = ggml_metal_op_col2im_1d(ctx, idx);
+            } break;
         case GGML_OP_CONV_3D:
             {
                 n_fuse = ggml_metal_op_conv_3d(ctx, idx);
@@ -1198,7 +1202,7 @@ int ggml_metal_op_set_rows(ggml_metal_op_t ctx, int idx) {
     GGML_TENSOR_LOCALS( int32_t, ne,  op,         ne);
     GGML_TENSOR_LOCALS(uint64_t, nb,  op,         nb);
 
-    auto pipeline = ggml_metal_library_get_pipeline_set_rows(lib, op->src[1]->type, op->type);
+    auto pipeline = ggml_metal_library_get_pipeline_set_rows(lib, op);
 
     const int32_t nk0 = ne0/ggml_blck_size(op->type);
 
@@ -3854,6 +3858,47 @@ int ggml_metal_op_conv_transpose_1d(ggml_metal_op_t ctx, int idx) {
     return 1;
 }
 
+int ggml_metal_op_col2im_1d(ggml_metal_op_t ctx, int idx) {
+    ggml_tensor * op = ctx->node(idx);
+
+    ggml_metal_library_t lib = ctx->lib;
+    ggml_metal_encoder_t enc = ctx->enc;
+
+    const int32_t s0 = ((const int32_t *)(op->op_params))[0];
+    const int32_t OC = ((const int32_t *)(op->op_params))[1];
+    const int32_t p0 = ((const int32_t *)(op->op_params))[2];
+
+    const int32_t K_OC  = (int32_t) op->src[0]->ne[0];
+    const int32_t T_in  = (int32_t) op->src[0]->ne[1];
+    const int32_t K     = K_OC / OC;
+    const int32_t T_out = (int32_t) op->ne[0];
+
+    ggml_metal_kargs_col2im_1d args = {
+        /*.T_in  =*/ T_in,
+        /*.T_out =*/ T_out,
+        /*.OC    =*/ OC,
+        /*.K     =*/ K,
+        /*.K_OC  =*/ K_OC,
+        /*.s0    =*/ s0,
+        /*.p0    =*/ p0,
+    };
+
+    auto pipeline = ggml_metal_library_get_pipeline_col2im_1d(lib, op);
+
+    const int total = T_out * OC;
+    const int nth   = 256;
+    const int ntg   = (total + nth - 1) / nth;
+
+    ggml_metal_encoder_set_pipeline(enc, pipeline);
+    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         2);
+
+    ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1);
+
+    return 1;
+}
+
 int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 

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

@@ -78,6 +78,7 @@ int ggml_metal_op_conv_2d           (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_conv_3d           (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_conv_transpose_1d (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_conv_transpose_2d (ggml_metal_op_t ctx, int idx);
+int ggml_metal_op_col2im_1d         (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_upscale           (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_pad               (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_pad_reflect_1d    (ggml_metal_op_t ctx, int idx);

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

@@ -42,6 +42,8 @@ typedef matrix<bfloat, 4, 4> bfloat4x4;
 typedef matrix<bfloat, 2, 4> bfloat2x4;
 #endif
 
+#define QK_NL 16
+
 constexpr constant static float kvalues_iq4nl_f[16] = {
     -127.f, -104.f, -83.f, -65.f, -49.f, -35.f, -22.f, -10.f, 1.f, 13.f, 25.f, 38.f, 53.f, 69.f, 89.f, 113.f
 };
@@ -4977,6 +4979,49 @@ kernel void kernel_conv_transpose_1d<half>(
     uint3    tgpg[[threadgroups_per_grid]]);
 
 
+template <typename T>
+kernel void kernel_col2im_1d(
+        constant ggml_metal_kargs_col2im_1d & args,
+        device const T * col,
+        device       T * dst,
+        uint         tgpig [[threadgroup_position_in_grid]],
+        uint         tpitg [[thread_position_in_threadgroup]],
+        uint         ntg   [[threads_per_threadgroup]]) {
+
+    const int idx = tgpig * ntg + tpitg;
+    if (idx >= args.T_out * args.OC) {
+        return;
+    }
+
+    const int t_out = idx % args.T_out;
+    const int oc    = idx / args.T_out;
+    const int t_abs = t_out + args.p0;  // absolute position in uncropped signal
+
+    int t_in_min = (t_abs - args.K + args.s0) / args.s0;  // ceil((t_abs - K + 1) / s0)
+    if (t_in_min < 0) {
+        t_in_min = 0;
+    }
+    int t_in_max = t_abs / args.s0;
+    if (t_in_max >= args.T_in) {
+        t_in_max = args.T_in - 1;
+    }
+
+    float sum = 0.0f;
+    for (int t_in = t_in_min; t_in <= t_in_max; t_in++) {
+        const int k = t_abs - t_in * args.s0;
+        sum += float(col[(oc * args.K + k) + t_in * args.K_OC]);
+    }
+
+    dst[t_out + oc * args.T_out] = T(sum);
+}
+
+template [[host_name("kernel_col2im_1d_f32")]]  kernel void kernel_col2im_1d<float>(constant ggml_metal_kargs_col2im_1d &, device const float *, device float *, uint, uint, uint);
+template [[host_name("kernel_col2im_1d_f16")]]  kernel void kernel_col2im_1d<half>(constant ggml_metal_kargs_col2im_1d &, device const half *, device half *, uint, uint, uint);
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_col2im_1d_bf16")]] kernel void kernel_col2im_1d<bfloat>(constant ggml_metal_kargs_col2im_1d &, device const bfloat *, device bfloat *, uint, uint, uint);
+#endif
+
+
 typedef void (conv_transpose_2d_t)(
         constant ggml_metal_kargs_conv_transpose_2d & args,
         device const float * src0,
@@ -9343,7 +9388,40 @@ kernel void kernel_get_rows_f(
     }
 }
 
-template<typename TI, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
+typedef decltype(kernel_get_rows_f<float, float>) get_rows_f_t;
+
+template [[host_name("kernel_get_rows_f32")]]  kernel get_rows_f_t kernel_get_rows_f<float, float>;
+template [[host_name("kernel_get_rows_f16")]]  kernel get_rows_f_t kernel_get_rows_f<half,  float>;
+template [[host_name("kernel_get_rows_i32")]]  kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f<bfloat, float>;
+#endif
+
+typedef decltype(kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>) get_rows_q_t;
+
+template [[host_name("kernel_get_rows_q1_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q1_0,    8, dequantize_q1_0>;
+template [[host_name("kernel_get_rows_q4_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_0,    2, dequantize_q4_0>;
+template [[host_name("kernel_get_rows_q4_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_1,    2, dequantize_q4_1>;
+template [[host_name("kernel_get_rows_q5_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_0,    2, dequantize_q5_0>;
+template [[host_name("kernel_get_rows_q5_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_1,    2, dequantize_q5_1>;
+template [[host_name("kernel_get_rows_q8_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q8_0,    2, dequantize_q8_0>;
+template [[host_name("kernel_get_rows_mxfp4")]]   kernel get_rows_q_t kernel_get_rows_q<block_mxfp4,   2, dequantize_mxfp4>;
+template [[host_name("kernel_get_rows_q2_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q2_K,    QK_NL, dequantize_q2_K>;
+template [[host_name("kernel_get_rows_q3_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q3_K,    QK_NL, dequantize_q3_K>;
+template [[host_name("kernel_get_rows_q4_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_K,    QK_NL, dequantize_q4_K>;
+template [[host_name("kernel_get_rows_q5_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_K,    QK_NL, dequantize_q5_K>;
+template [[host_name("kernel_get_rows_q6_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q6_K,    QK_NL, dequantize_q6_K>;
+template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
+template [[host_name("kernel_get_rows_iq2_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq2_xs,  QK_NL, dequantize_iq2_xs>;
+template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_xxs, QK_NL, dequantize_iq3_xxs>;
+template [[host_name("kernel_get_rows_iq3_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq3_s,   QK_NL, dequantize_iq3_s>;
+template [[host_name("kernel_get_rows_iq2_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq2_s,   QK_NL, dequantize_iq2_s>;
+template [[host_name("kernel_get_rows_iq1_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_s,   QK_NL, dequantize_iq1_s>;
+template [[host_name("kernel_get_rows_iq1_m")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_m,   QK_NL, dequantize_iq1_m>;
+template [[host_name("kernel_get_rows_iq4_nl")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_nl,  2,     dequantize_iq4_nl>;
+template [[host_name("kernel_get_rows_iq4_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_xs,  QK_NL, dequantize_iq4_xs>;
+
+template<typename TS, typename TI, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
 kernel void kernel_set_rows_q32(
         constant ggml_metal_kargs_set_rows & args,
         device const  void * src0,
@@ -9367,14 +9445,14 @@ kernel void kernel_set_rows_q32(
     const TI      i1  = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
 
           device block_q * dst_row = (      device block_q *) ((      device char *) dst  +  i1*args.nb1  + i02*args.nb2  + i03*args.nb3);
-    const device float   * src_row = (const device float   *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
+    const device TS      * src_row = (const device TS      *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
 
     for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
         quantize_func(src_row + 32*ind, dst_row[ind]);
     }
 }
 
-template<typename T, typename TI>
+template<typename TS, typename TI, typename TD>
 kernel void kernel_set_rows_f(
         constant ggml_metal_kargs_set_rows & args,
         device const  void * src0,
@@ -9397,14 +9475,47 @@ kernel void kernel_set_rows_f(
     const int32_t i10 = i01;
     const TI      i1  = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0];
 
-          device T     * dst_row = (      device T     *) ((      device char *) dst  +  i1*args.nb1  + i02*args.nb2  + i03*args.nb3);
-    const device float * src_row = (const device float *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
+          device TD * dst_row = (      device TD *) ((      device char *) dst  +  i1*args.nb1  + i02*args.nb2  + i03*args.nb3);
+    const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
 
     for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) {
-        dst_row[ind] = (T) src_row[ind];
+        dst_row[ind] = (TD) src_row[ind];
     }
 }
 
+typedef decltype(kernel_set_rows_f<float, int64_t, float>) set_rows_f_t;
+
+template [[host_name("kernel_set_rows_f32_i64_f32")]]   kernel set_rows_f_t kernel_set_rows_f<float, int64_t, float>;
+template [[host_name("kernel_set_rows_f32_i32_f32")]]   kernel set_rows_f_t kernel_set_rows_f<float, int32_t, float>;
+template [[host_name("kernel_set_rows_f32_i64_f16")]]   kernel set_rows_f_t kernel_set_rows_f<float, int64_t, half>;
+template [[host_name("kernel_set_rows_f32_i32_f16")]]   kernel set_rows_f_t kernel_set_rows_f<float, int32_t, half>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_set_rows_f32_i64_bf16")]]  kernel set_rows_f_t kernel_set_rows_f<float, int64_t, bfloat>;
+template [[host_name("kernel_set_rows_f32_i32_bf16")]]  kernel set_rows_f_t kernel_set_rows_f<float, int32_t, bfloat>;
+#endif
+
+template [[host_name("kernel_set_rows_f16_i64_f16")]]   kernel set_rows_f_t kernel_set_rows_f<half, int64_t, half>;
+template [[host_name("kernel_set_rows_f16_i32_f16")]]   kernel set_rows_f_t kernel_set_rows_f<half, int32_t, half>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_set_rows_bf16_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int64_t, bfloat>;
+template [[host_name("kernel_set_rows_bf16_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int32_t, bfloat>;
+#endif
+
+typedef decltype(kernel_set_rows_q32<float, int64_t, block_q8_0, quantize_q8_0>) set_rows_q32_t;
+
+template [[host_name("kernel_set_rows_f32_i64_q8_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q8_0,   quantize_q8_0>;
+template [[host_name("kernel_set_rows_f32_i32_q8_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q8_0,   quantize_q8_0>;
+template [[host_name("kernel_set_rows_f32_i64_q4_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_0,   quantize_q4_0>;
+template [[host_name("kernel_set_rows_f32_i32_q4_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_0,   quantize_q4_0>;
+template [[host_name("kernel_set_rows_f32_i64_q4_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_1,   quantize_q4_1>;
+template [[host_name("kernel_set_rows_f32_i32_q4_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_1,   quantize_q4_1>;
+template [[host_name("kernel_set_rows_f32_i64_q5_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_0,   quantize_q5_0>;
+template [[host_name("kernel_set_rows_f32_i32_q5_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_0,   quantize_q5_0>;
+template [[host_name("kernel_set_rows_f32_i64_q5_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_1,   quantize_q5_1>;
+template [[host_name("kernel_set_rows_f32_i32_q5_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_1,   quantize_q5_1>;
+template [[host_name("kernel_set_rows_f32_i64_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_iq4_nl, quantize_iq4_nl>;
+template [[host_name("kernel_set_rows_f32_i32_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_iq4_nl, quantize_iq4_nl>;
+
 kernel void kernel_diag_f32(
         constant ggml_metal_kargs_diag & args,
         device   const char * src0,
@@ -10147,75 +10258,6 @@ kernel void kernel_mul_mm_id(
     }
 }
 
-#define QK_NL 16
-
-//
-// get rows
-//
-
-typedef decltype(kernel_get_rows_f<float, float>) get_rows_f_t;
-
-template [[host_name("kernel_get_rows_f32")]]  kernel get_rows_f_t kernel_get_rows_f<float, float>;
-template [[host_name("kernel_get_rows_f16")]]  kernel get_rows_f_t kernel_get_rows_f<half,  float>;
-template [[host_name("kernel_get_rows_i32")]]  kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f<bfloat, float>;
-#endif
-
-typedef decltype(kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>) get_rows_q_t;
-
-template [[host_name("kernel_get_rows_q1_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q1_0,    8, dequantize_q1_0>;
-template [[host_name("kernel_get_rows_q4_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_0,    2, dequantize_q4_0>;
-template [[host_name("kernel_get_rows_q4_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_1,    2, dequantize_q4_1>;
-template [[host_name("kernel_get_rows_q5_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_0,    2, dequantize_q5_0>;
-template [[host_name("kernel_get_rows_q5_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_1,    2, dequantize_q5_1>;
-template [[host_name("kernel_get_rows_q8_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q8_0,    2, dequantize_q8_0>;
-template [[host_name("kernel_get_rows_mxfp4")]]   kernel get_rows_q_t kernel_get_rows_q<block_mxfp4,   2, dequantize_mxfp4>;
-template [[host_name("kernel_get_rows_q2_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q2_K,    QK_NL, dequantize_q2_K>;
-template [[host_name("kernel_get_rows_q3_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q3_K,    QK_NL, dequantize_q3_K>;
-template [[host_name("kernel_get_rows_q4_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_K,    QK_NL, dequantize_q4_K>;
-template [[host_name("kernel_get_rows_q5_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_K,    QK_NL, dequantize_q5_K>;
-template [[host_name("kernel_get_rows_q6_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q6_K,    QK_NL, dequantize_q6_K>;
-template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
-template [[host_name("kernel_get_rows_iq2_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq2_xs,  QK_NL, dequantize_iq2_xs>;
-template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_xxs, QK_NL, dequantize_iq3_xxs>;
-template [[host_name("kernel_get_rows_iq3_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq3_s,   QK_NL, dequantize_iq3_s>;
-template [[host_name("kernel_get_rows_iq2_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq2_s,   QK_NL, dequantize_iq2_s>;
-template [[host_name("kernel_get_rows_iq1_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_s,   QK_NL, dequantize_iq1_s>;
-template [[host_name("kernel_get_rows_iq1_m")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_m,   QK_NL, dequantize_iq1_m>;
-template [[host_name("kernel_get_rows_iq4_nl")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_nl,  2,     dequantize_iq4_nl>;
-template [[host_name("kernel_get_rows_iq4_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_xs,  QK_NL, dequantize_iq4_xs>;
-
-//
-// set rows
-//
-
-typedef decltype(kernel_set_rows_f<float, int64_t>) set_rows_f_t;
-
-template [[host_name("kernel_set_rows_f32_i64")]]  kernel set_rows_f_t kernel_set_rows_f<float, int64_t>;
-template [[host_name("kernel_set_rows_f32_i32")]]  kernel set_rows_f_t kernel_set_rows_f<float, int32_t>;
-template [[host_name("kernel_set_rows_f16_i64")]]  kernel set_rows_f_t kernel_set_rows_f<half, int64_t>;
-template [[host_name("kernel_set_rows_f16_i32")]]  kernel set_rows_f_t kernel_set_rows_f<half, int32_t>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_set_rows_bf16_i64")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int64_t>;
-template [[host_name("kernel_set_rows_bf16_i32")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int32_t>;
-#endif
-
-typedef decltype(kernel_set_rows_q32<int64_t, block_q8_0, quantize_q8_0>) set_rows_q32_t;
-
-template [[host_name("kernel_set_rows_q8_0_i64")]]   kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q8_0,   quantize_q8_0>;
-template [[host_name("kernel_set_rows_q8_0_i32")]]   kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q8_0,   quantize_q8_0>;
-template [[host_name("kernel_set_rows_q4_0_i64")]]   kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q4_0,   quantize_q4_0>;
-template [[host_name("kernel_set_rows_q4_0_i32")]]   kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q4_0,   quantize_q4_0>;
-template [[host_name("kernel_set_rows_q4_1_i64")]]   kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q4_1,   quantize_q4_1>;
-template [[host_name("kernel_set_rows_q4_1_i32")]]   kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q4_1,   quantize_q4_1>;
-template [[host_name("kernel_set_rows_q5_0_i64")]]   kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q5_0,   quantize_q5_0>;
-template [[host_name("kernel_set_rows_q5_0_i32")]]   kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q5_0,   quantize_q5_0>;
-template [[host_name("kernel_set_rows_q5_1_i64")]]   kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q5_1,   quantize_q5_1>;
-template [[host_name("kernel_set_rows_q5_1_i32")]]   kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q5_1,   quantize_q5_1>;
-template [[host_name("kernel_set_rows_iq4_nl_i64")]] kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_iq4_nl, quantize_iq4_nl>;
-template [[host_name("kernel_set_rows_iq4_nl_i32")]] kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_iq4_nl, quantize_iq4_nl>;
-
 //
 // matrix-matrix multiplication
 //

ggml/src/ggml-opencl/fa_tune.h

@@ -20,6 +20,7 @@ static const ggml_opencl_fa_dim g_fa_dims_adreno_default[] = {
     {192, 128, 16, 16, 1, 0},
     {192, 192, 16, 16, 1, 0},
     {256, 256, 16, 16, 16, 0},
+    {512, 512,  8, 16, 64, 0},
 };
 
 struct ggml_opencl_fa_dim_table {

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

@@ -10,7 +10,12 @@
 #define DK_VEC (DK/4)
 #define DV_VEC (DV/4)
 #define WG_SIZE (BLOCK_M)
-#define Q1_WG_SIZE 64
+// q1 reduces over a Q1_WG_SIZE-wide WG via work-group barriers; the launch WG
+// must match. Defaults to the Adreno sg (64); host passes -D FA_SG=32 on Intel.
+#ifndef FA_SG
+#define FA_SG 64
+#endif
+#define Q1_WG_SIZE FA_SG
 
 // The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
 // infinite operand can cause undefined behavior and miscompilation for exp.

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

@@ -11,7 +11,12 @@
 #define DK_VEC (DK/4)
 #define DV_VEC (DV/4)
 #define WG_SIZE (BLOCK_M)
-#define Q1_WG_SIZE 64
+// q1 reduces over a Q1_WG_SIZE-wide WG via work-group barriers; the launch WG
+// must match. Defaults to the Adreno sg (64); host passes -D FA_SG=32 on Intel.
+#ifndef FA_SG
+#define FA_SG 64
+#endif
+#define Q1_WG_SIZE FA_SG
 
 // The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
 // infinite operand can cause undefined behavior and miscompilation for exp.
@@ -114,6 +119,15 @@ __kernel void flash_attn_f32(
     __local DATA_TYPE4 l_v[BLOCK_N][DV_VEC];
 
     for (int k_start = 0; k_start < n_kv; k_start += BLOCK_N) {
+#if FA_SG < 64
+        // WAR on l_k/l_v: threads with my_query_row >= n_q skip the compute below
+        // (continue) and would race ahead to reload the tiles while active threads
+        // still read them. A single 64-wide Adreno subgroup (WG == sg) runs lockstep
+        // and hides this; a WG that spans multiple narrower subgroups (Intel sg=32)
+        // corrupts the result. All threads reach this each iteration (no-op on the
+        // first), so it does not diverge with the continue. Compiled out at sg=64.
+        barrier(CLK_LOCAL_MEM_FENCE);
+#endif
         for (int i = tid; i < BLOCK_N * DK_VEC; i += WG_SIZE) {
             const int row = i / DK_VEC;
             const int col = i % DK_VEC;

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

@@ -27,7 +27,11 @@
 
 #define DK_VEC (DK/4)
 #define DV_VEC (DV/4)
-#define Q1_WG_SIZE 64
+
+#ifndef FA_SG
+#define FA_SG 64
+#endif
+#define Q1_WG_SIZE FA_SG
 
 // The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
 // infinite operand can cause undefined behavior and miscompilation for exp.
@@ -365,6 +369,263 @@ __kernel void flash_attn_f32_q4_0_q1(
     }
 }
 
+#ifdef cl_intel_subgroups
+#pragma OPENCL EXTENSION cl_intel_subgroups : enable
+#else
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#endif
+
+#ifdef cl_qcom_reqd_sub_group_size
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define REQD_SUBGROUP_SIZE_64 __attribute__((qcom_reqd_sub_group_size("half")))
+#else
+#define REQD_SUBGROUP_SIZE_64
+#endif
+
+#define VEC_NSG          4
+#define VEC_WG_SIZE      (Q1_WG_SIZE * VEC_NSG)
+#define Q1V_DV_PER_THREAD ((DV_VEC + Q1_WG_SIZE - 1) / Q1_WG_SIZE)
+
+// Dequant one float4 lane (0..7) from a q4_0 block.
+// Lanes 0..3 → low nibbles of qs[0..15], lanes 4..7 → high nibbles.
+inline float4 dequant_q4_0_lane(const global char * block_ptr, int lane) {
+    const float d = vload_half(0, (const global half *)block_ptr);
+    const global uchar * qs = (const global uchar *)(block_ptr + 2);
+    const int g     = lane & 3;
+    const int shift = (lane < 4) ? 0 : 4;
+    return d * (float4)((float)((qs[g*4+0] >> shift) & 0x0F) - 8.0f,
+                        (float)((qs[g*4+1] >> shift) & 0x0F) - 8.0f,
+                        (float)((qs[g*4+2] >> shift) & 0x0F) - 8.0f,
+                        (float)((qs[g*4+3] >> shift) & 0x0F) - 8.0f);
+}
+
+REQD_SUBGROUP_SIZE_64
+__kernel void flash_attn_f32_q4_0_q1_vec(
+    const global void * q_void, ulong q_offset,
+    const global void * k_void, ulong k_offset,
+    const global void * v_void, ulong v_offset,
+    global void * o_void, ulong o_offset,
+    const float scale,
+    const int n_q,
+    const int n_kv,
+    const int is_causal,
+    const int n_head,
+    const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
+    const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
+    const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
+    const ulong o_nb1, const ulong o_nb2, const ulong o_nb3,
+    const float max_bias,
+    const float m0,
+    const float m1,
+    const int n_head_log2,
+    const float logit_softcap,
+    const int n_head_kv,
+    const global void* mask_void,
+    const ulong mask_offset,
+    const ulong mask_nb1,
+    const ulong mask_nb2,
+    const ulong mask_nb3,
+    const int mask_ne2,
+    const int mask_ne3,
+    const global void* sinks_void,
+    const ulong sinks_offset
+) {
+    const int tid             = get_local_id(0);
+    const int sgid            = tid / Q1_WG_SIZE;
+    const int tid_sg          = tid % Q1_WG_SIZE;
+    const int head_batch_idx  = get_global_id(1);
+
+    const int batch_idx = head_batch_idx / n_head;
+    const int head_idx  = head_batch_idx % n_head;
+
+    const int gqa_ratio   = n_head / n_head_kv;
+    const int head_kv_idx = head_idx / gqa_ratio;
+
+    const global char * q_base = (const global char *) q_void + q_offset;
+    const global char * k_base = (const global char *) k_void + k_offset;
+    const global char * v_base = (const global char *) v_void + v_offset;
+    global       char * o_base = (global       char *) o_void + o_offset;
+
+    const global char * mask_base = NULL;
+    if (mask_void != NULL) {
+        const int mask_head_idx  = head_idx  % mask_ne2;
+        const int mask_batch_idx = batch_idx % mask_ne3;
+        mask_base = (const global char *) mask_void + mask_offset +
+                    mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2;
+    }
+
+    __local ACC_TYPE4 q_shared[DK_VEC];
+    {
+        const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2;
+        const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
+        for (int i = tid; i < DK_VEC; i += VEC_WG_SIZE) {
+            q_shared[i] = CONVERT_Q_ACC4(q_ptr[i]);
+        }
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+#ifdef FA_HAVE_INT_DOT
+    // quantize Q to int8-packed uints + per-block (qd, q_sum) once per WG for dp4a
+    // one thread per Q block, remaining threads idle this step
+    __local uint  q_packed_shared[DK_Q4_BLOCKS * 8];
+    __local float q_d_shared[DK_Q4_BLOCKS];
+    __local int   q_sum_shared[DK_Q4_BLOCKS];
+    if (tid < DK_Q4_BLOCKS) {
+        ACC_TYPE4 q_block[8];
+        #pragma unroll
+        for (int i = 0; i < 8; ++i) q_block[i] = q_shared[tid * 8 + i];
+        uint packed[8];
+        q4_q_block_info info = quant_q_block_int8_packed_q4(q_block, packed);
+        #pragma unroll
+        for (int i = 0; i < 8; ++i) q_packed_shared[tid * 8 + i] = packed[i];
+        q_d_shared[tid]   = info.qd;
+        q_sum_shared[tid] = info.q_sum;
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+#endif
+
+    const float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
+
+    const global ACC_TYPE * sinks_ptr = NULL;
+    if (sinks_void != NULL) {
+        sinks_ptr = (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset);
+    }
+
+    ACC_TYPE4 o_acc[Q1V_DV_PER_THREAD];
+    #pragma unroll
+    for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[i] = (ACC_TYPE4)(0.0f);
+
+    ACC_TYPE m_i = FA_M_INIT;
+    ACC_TYPE l_i = 0.0f;
+
+    const int kv_per_sg = (n_kv + VEC_NSG - 1) / VEC_NSG;
+    const int kv_start  = sgid * kv_per_sg;
+    const int kv_end    = min(n_kv, kv_start + kv_per_sg);
+
+    for (int k_idx = kv_start; k_idx < kv_end; ++k_idx) {
+        const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
+        const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
+
+#ifdef FA_HAVE_INT_DOT
+        // per-lane dp4a: each lane packs 4 raw q4_0 nibbles into a uint,
+        // then dot_acc_sat_4x8packed_ss_int against the matching uint.
+        ACC_TYPE lane_contrib = 0.0f;
+        for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
+            const int block_idx     = qk / 8;
+            const int lane_in_block = qk % 8;
+            const int g             = lane_in_block & 3;
+            const int shift         = (lane_in_block < 4) ? 0 : 4;
+            const global char *  k_block = k_row + block_idx * Q4_0_BLOCK_SIZE;
+            const float kd = vload_half(0, (const global half *)k_block);
+            const global uchar * k_qs = (const global uchar *)(k_block + 2);
+            const uchar b0 = k_qs[g*4 + 0];
+            const uchar b1 = k_qs[g*4 + 1];
+            const uchar b2 = k_qs[g*4 + 2];
+            const uchar b3 = k_qs[g*4 + 3];
+            const uint k_packed = ((uint)((b0 >> shift) & 0x0F))       |
+                                  ((uint)((b1 >> shift) & 0x0F)) <<  8 |
+                                  ((uint)((b2 >> shift) & 0x0F)) << 16 |
+                                  ((uint)((b3 >> shift) & 0x0F)) << 24;
+            const uint q_packed_lane = q_packed_shared[block_idx * 8 + lane_in_block];
+            const int  raw_dot = dot_acc_sat_4x8packed_ss_int(q_packed_lane, k_packed, 0);
+            const float qd          = q_d_shared[block_idx];
+            const float block_scale = qd * kd;
+            float contrib = (float)raw_dot * block_scale;
+            if (lane_in_block == 0) {
+                // block bias correction is per-block
+                const int q_sum_b = q_sum_shared[block_idx];
+                contrib -= 8.0f * block_scale * (float)q_sum_b;
+            }
+            lane_contrib += contrib;
+        }
+        ACC_TYPE score = sub_group_reduce_add(lane_contrib) * scale;
+#else
+        ACC_TYPE4 dot4 = (ACC_TYPE4)(0.0f);
+        for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
+            const int block_idx = qk / 8;
+            const int lane      = qk % 8;
+            const float4 k_v = dequant_q4_0_lane(k_row + block_idx * Q4_0_BLOCK_SIZE, lane);
+            dot4 = mad(q_shared[qk], k_v, dot4);
+        }
+        ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3;
+        ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale;
+#endif
+
+        if (mask_base != NULL) {
+            const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
+            score += slope * (ACC_TYPE) mask_ptr[k_idx];
+        }
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }
+
+        const ACC_TYPE m_new      = max(m_i, score);
+        const ACC_TYPE scale_prev = native_exp(m_i - m_new);
+        const ACC_TYPE p          = native_exp(score - m_new);
+
+        int idx = 0;
+        for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
+            const int block_idx = dv / 8;
+            const int lane      = dv % 8;
+            const float4 v_v = dequant_q4_0_lane(v_row + block_idx * Q4_0_BLOCK_SIZE, lane);
+            o_acc[idx] = mad(p, v_v, o_acc[idx] * scale_prev);
+        }
+        l_i = l_i * scale_prev + p;
+        m_i = m_new;
+    }
+
+    __local ACC_TYPE  sg_m[VEC_NSG];
+    __local ACC_TYPE  sg_l[VEC_NSG];
+    __local ACC_TYPE4 sg_o[VEC_NSG][DV_VEC];
+
+    if (tid_sg == 0) {
+        sg_m[sgid] = m_i;
+        sg_l[sgid] = l_i;
+    }
+    {
+        int idx = 0;
+        for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
+            sg_o[sgid][dv] = o_acc[idx];
+        }
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    if (sgid == 0) {
+        ACC_TYPE m_final = sg_m[0];
+        #pragma unroll
+        for (int s = 1; s < VEC_NSG; ++s) {
+            m_final = max(m_final, sg_m[s]);
+        }
+        if (sinks_ptr != NULL) {
+            m_final = max(m_final, sinks_ptr[head_idx]);
+        }
+
+        ACC_TYPE l_final = 0.0f;
+        #pragma unroll
+        for (int s = 0; s < VEC_NSG; ++s) {
+            l_final += sg_l[s] * native_exp(sg_m[s] - m_final);
+        }
+        if (sinks_ptr != NULL) {
+            l_final += native_exp(sinks_ptr[head_idx] - m_final);
+        }
+        const ACC_TYPE l_inv = (l_final > 0.0f) ? (1.0f / l_final) : 0.0f;
+
+        const ulong o_row_offset = batch_idx * o_nb3 + head_idx * o_nb1;
+        global O_DATA_TYPE4 * o_row = (global O_DATA_TYPE4 *) (o_base + o_row_offset);
+
+        int idx = 0;
+        for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
+            ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
+            #pragma unroll
+            for (int s = 0; s < VEC_NSG; ++s) {
+                const ACC_TYPE alpha = native_exp(sg_m[s] - m_final);
+                o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv], o_merged);
+            }
+            o_row[dv] = CONVERT_O_DATA4(o_merged * l_inv);
+        }
+    }
+}
+
 // Flash-decoding split pass for q4_0 KV. Merge kernel is type-agnostic and
 // shared with the f16/q8_0 FA kernels.
 #define FA_PARTIAL_FLOATS (2 + DV)
@@ -583,6 +844,319 @@ __kernel void flash_attn_f32_q4_0_q1_split(
 #define WG_SIZE             BLOCK_M
 #endif
 
+#ifndef MQ_GQA
+#define MQ_GQA 4
+#endif
+#ifndef MQ_NSG_SPLIT
+#define MQ_NSG_SPLIT 4
+#endif
+#define MQ_SPLIT_WG_SIZE_Q4 (Q1_WG_SIZE * MQ_NSG_SPLIT)
+
+REQD_SUBGROUP_SIZE_64
+__kernel void flash_attn_f32_q4_0_q1_vec_mq_split(
+    const global void * q_void, ulong q_offset,
+    const global void * k_void, ulong k_offset,
+    const global void * v_void, ulong v_offset,
+    const float scale,
+    const int n_q,
+    const int n_kv,
+    const int n_head,
+    const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
+    const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
+    const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
+    const float max_bias,
+    const float m0,
+    const float m1,
+    const int n_head_log2,
+    const float logit_softcap,
+    const int n_head_kv,
+    const global void * mask_void,
+    const ulong mask_offset,
+    const ulong mask_nb1,
+    const ulong mask_nb2,
+    const ulong mask_nb3,
+    const int mask_ne2,
+    const int mask_ne3,
+    global float * partial_void,
+    const int n_splits,
+    const int kv_per_split
+) {
+    const int tid              = get_local_id(0);
+    const int sgid             = tid / Q1_WG_SIZE;
+    const int tid_sg           = tid % Q1_WG_SIZE;
+    const int kvhead_batch_idx = get_global_id(1);
+    const int split_q_idx      = get_global_id(2);
+    const int split_idx        = split_q_idx % n_splits;
+    const int q_idx            = split_q_idx / n_splits;
+
+    const int batch_idx   = kvhead_batch_idx / n_head_kv;
+    const int head_kv_idx = kvhead_batch_idx % n_head_kv;
+
+    const int kv_start = split_idx * kv_per_split;
+    const int kv_end   = min(kv_start + kv_per_split, n_kv);
+
+    const ulong record_stride = (ulong) FA_PARTIAL_FLOATS;
+
+    if (kv_start >= kv_end) {
+        if (tid == 0) {
+            #pragma unroll
+            for (int h = 0; h < MQ_GQA; ++h) {
+                const int head_idx = head_kv_idx * MQ_GQA + h;
+                const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
+                                       * n_splits + split_idx);
+                global float * rec = partial_void + rec_idx * record_stride;
+                rec[0] = FA_M_INIT;
+                rec[1] = 0.0f;
+            }
+        }
+        return;
+    }
+
+    const global char * q_base = (const global char *) q_void + q_offset;
+    const global char * k_base = (const global char *) k_void + k_offset;
+    const global char * v_base = (const global char *) v_void + v_offset;
+
+    __local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
+    for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_SPLIT_WG_SIZE_Q4) {
+        const int h        = i / DK_VEC;
+        const int k        = i % DK_VEC;
+        const int head_idx = head_kv_idx * MQ_GQA + h;
+        const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1;
+        const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
+        q_shared[h * DK_VEC + k] = CONVERT_Q_ACC4(q_ptr[k]);
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+#ifdef FA_HAVE_INT_DOT
+    __local uint  q_packed_shared[MQ_GQA * DK_Q4_BLOCKS * 8];
+    __local float q_d_shared[MQ_GQA * DK_Q4_BLOCKS];
+    __local int   q_sum_shared[MQ_GQA * DK_Q4_BLOCKS];
+    {
+        const int active = MQ_GQA * DK_Q4_BLOCKS;
+        if (tid < active) {
+            const int h        = tid / DK_Q4_BLOCKS;
+            const int block_id = tid % DK_Q4_BLOCKS;
+            ACC_TYPE4 q_block[8];
+            #pragma unroll
+            for (int i = 0; i < 8; ++i) q_block[i] = q_shared[h * DK_VEC + block_id * 8 + i];
+            uint packed[8];
+            q4_q_block_info info = quant_q_block_int8_packed_q4(q_block, packed);
+            #pragma unroll
+            for (int i = 0; i < 8; ++i) q_packed_shared[(h * DK_Q4_BLOCKS + block_id) * 8 + i] = packed[i];
+            q_d_shared[h * DK_Q4_BLOCKS + block_id]   = info.qd;
+            q_sum_shared[h * DK_Q4_BLOCKS + block_id] = info.q_sum;
+        }
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+#endif
+
+    float slope[MQ_GQA];
+    #pragma unroll
+    for (int h = 0; h < MQ_GQA; ++h) {
+        slope[h] = get_alibi_slope(max_bias, head_kv_idx * MQ_GQA + h, n_head_log2, m0, m1);
+    }
+
+    const global char * mask_base[MQ_GQA];
+    if (mask_void != NULL) {
+        const int mask_batch_idx = batch_idx % mask_ne3;
+        const global char * mask_base_b = (const global char *) mask_void + mask_offset +
+                                          mask_batch_idx * mask_nb3 +
+                                          (ulong) q_idx * mask_nb1;
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) {
+            const int head_idx      = head_kv_idx * MQ_GQA + h;
+            const int mask_head_idx = head_idx % mask_ne2;
+            mask_base[h] = mask_base_b + mask_head_idx * mask_nb2;
+        }
+    } else {
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) mask_base[h] = NULL;
+    }
+
+    ACC_TYPE4 o_acc[MQ_GQA][Q1V_DV_PER_THREAD];
+    ACC_TYPE  m_i[MQ_GQA];
+    ACC_TYPE  l_i[MQ_GQA];
+    #pragma unroll
+    for (int h = 0; h < MQ_GQA; ++h) {
+        m_i[h] = FA_M_INIT;
+        l_i[h] = 0.0f;
+        #pragma unroll
+        for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[h][i] = (ACC_TYPE4)(0.0f);
+    }
+
+    const int kv_len    = kv_end - kv_start;
+    const int kv_per_sg = (kv_len + MQ_NSG_SPLIT - 1) / MQ_NSG_SPLIT;
+    const int kv_lo     = kv_start + sgid * kv_per_sg;
+    const int kv_hi     = min(kv_end, kv_lo + kv_per_sg);
+
+    for (int k_idx = kv_lo; k_idx < kv_hi; ++k_idx) {
+        const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
+        const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
+
+#ifdef FA_HAVE_INT_DOT
+        ACC_TYPE lane_contrib[MQ_GQA];
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) lane_contrib[h] = 0.0f;
+
+        for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
+            const int block_idx     = qk / 8;
+            const int lane_in_block = qk % 8;
+            const int g             = lane_in_block & 3;
+            const int shift         = (lane_in_block < 4) ? 0 : 4;
+            const global char *  k_block = k_row + block_idx * Q4_0_BLOCK_SIZE;
+            const float kd = vload_half(0, (const global half *)k_block);
+            const global uchar * k_qs = (const global uchar *)(k_block + 2);
+            const uchar b0 = k_qs[g*4 + 0];
+            const uchar b1 = k_qs[g*4 + 1];
+            const uchar b2 = k_qs[g*4 + 2];
+            const uchar b3 = k_qs[g*4 + 3];
+            const uint k_packed = ((uint)((b0 >> shift) & 0x0F))       |
+                                  ((uint)((b1 >> shift) & 0x0F)) <<  8 |
+                                  ((uint)((b2 >> shift) & 0x0F)) << 16 |
+                                  ((uint)((b3 >> shift) & 0x0F)) << 24;
+            #pragma unroll
+            for (int h = 0; h < MQ_GQA; ++h) {
+                const uint  q_packed_lane = q_packed_shared[(h * DK_Q4_BLOCKS + block_idx) * 8 + lane_in_block];
+                const int   raw_dot       = dot_acc_sat_4x8packed_ss_int(q_packed_lane, k_packed, 0);
+                const float qd            = q_d_shared[h * DK_Q4_BLOCKS + block_idx];
+                const float block_scale   = qd * kd;
+                float contrib = (float) raw_dot * block_scale;
+                if (lane_in_block == 0) {
+                    const int q_sum_b = q_sum_shared[h * DK_Q4_BLOCKS + block_idx];
+                    contrib -= 8.0f * block_scale * (float) q_sum_b;
+                }
+                lane_contrib[h] += contrib;
+            }
+        }
+
+        ACC_TYPE score[MQ_GQA];
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) {
+            ACC_TYPE s = sub_group_reduce_add(lane_contrib[h]) * scale;
+            if (mask_base[h] != NULL) {
+                const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
+                s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
+            }
+            if (logit_softcap > 0.0f) {
+                s = logit_softcap * tanh(s / logit_softcap);
+            }
+            score[h] = s;
+        }
+#else
+        // fallback float-dequant K dot
+        ACC_TYPE4 dot4[MQ_GQA];
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
+
+        for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
+            const int block_idx = qk / 8;
+            const int lane      = qk % 8;
+            const float4 k_v = dequant_q4_0_lane(k_row + block_idx * Q4_0_BLOCK_SIZE, lane);
+            #pragma unroll
+            for (int h = 0; h < MQ_GQA; ++h) {
+                dot4[h] = mad(q_shared[h * DK_VEC + qk], k_v, dot4[h]);
+            }
+        }
+
+        ACC_TYPE score[MQ_GQA];
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) {
+            const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
+            ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
+            if (mask_base[h] != NULL) {
+                const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
+                s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
+            }
+            if (logit_softcap > 0.0f) {
+                s = logit_softcap * tanh(s / logit_softcap);
+            }
+            score[h] = s;
+        }
+#endif
+
+        ACC_TYPE p_h[MQ_GQA];
+        ACC_TYPE sp_h[MQ_GQA];
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) {
+            const ACC_TYPE m_new = max(m_i[h], score[h]);
+            sp_h[h] = native_exp(m_i[h] - m_new);
+            p_h[h]  = native_exp(score[h] - m_new);
+            l_i[h]  = l_i[h] * sp_h[h] + p_h[h];
+            m_i[h]  = m_new;
+        }
+
+        int idx = 0;
+        for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
+            const int block_idx = dv / 8;
+            const int lane      = dv % 8;
+            const float4 v_v = dequant_q4_0_lane(v_row + block_idx * Q4_0_BLOCK_SIZE, lane);
+            #pragma unroll
+            for (int h = 0; h < MQ_GQA; ++h) {
+                o_acc[h][idx] = mad(p_h[h], v_v, o_acc[h][idx] * sp_h[h]);
+            }
+        }
+    }
+
+    // per-h cross-subgroup merge
+    __local ACC_TYPE  sg_m[MQ_GQA][MQ_NSG_SPLIT];
+    __local ACC_TYPE  sg_l[MQ_GQA][MQ_NSG_SPLIT];
+    __local ACC_TYPE4 sg_o[MQ_NSG_SPLIT][DV_VEC];
+
+    if (tid_sg == 0) {
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) {
+            sg_m[h][sgid] = m_i[h];
+            sg_l[h][sgid] = l_i[h];
+        }
+    }
+
+    #pragma unroll
+    for (int h = 0; h < MQ_GQA; ++h) {
+        {
+            int idx = 0;
+            for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE, ++idx) {
+                sg_o[sgid][dv_idx] = o_acc[h][idx];
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+        if (sgid == 0) {
+            const int head_idx = head_kv_idx * MQ_GQA + h;
+
+            ACC_TYPE m_c = sg_m[h][0];
+            #pragma unroll
+            for (int s = 1; s < MQ_NSG_SPLIT; ++s) {
+                m_c = max(m_c, sg_m[h][s]);
+            }
+            ACC_TYPE l_c = 0.0f;
+            #pragma unroll
+            for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
+                l_c += sg_l[h][s] * native_exp(sg_m[h][s] - m_c);
+            }
+
+            const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
+                                   * n_splits + split_idx);
+            global float  * rec   = partial_void + rec_idx * record_stride;
+            global float4 * rec_o = (global float4 *) (rec + 2);
+
+            if (tid_sg == 0) {
+                rec[0] = (float) m_c;
+                rec[1] = (float) l_c;
+            }
+            for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE) {
+                ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
+                #pragma unroll
+                for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
+                    const ACC_TYPE alpha = native_exp(sg_m[h][s] - m_c);
+                    o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv_idx], o_merged);
+                }
+                rec_o[dv_idx] = o_merged;
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+}
+
 __kernel void flash_attn_f32_q4_0(
     const global void * q_void, ulong q_offset,
     const global void * k_void, ulong k_offset,

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

@@ -24,7 +24,11 @@
 
 #define DK_VEC (DK/4)
 #define DV_VEC (DV/4)
-#define Q1_WG_SIZE 64
+
+#ifndef FA_SG
+#define FA_SG 64
+#endif
+#define Q1_WG_SIZE FA_SG
 
 // The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
 // infinite operand can cause undefined behavior and miscompilation for exp.
@@ -310,6 +314,201 @@ __kernel void flash_attn_f32_q8_0_q1(
     }
 }
 
+#ifdef cl_intel_subgroups
+#pragma OPENCL EXTENSION cl_intel_subgroups : enable
+#else
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#endif
+
+#ifdef cl_qcom_reqd_sub_group_size
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+#define REQD_SUBGROUP_SIZE_64 __attribute__((qcom_reqd_sub_group_size("half")))
+#else
+#define REQD_SUBGROUP_SIZE_64
+#endif
+
+#define VEC_NSG          4
+#define VEC_WG_SIZE      (Q1_WG_SIZE * VEC_NSG)
+#define Q1V_DV_PER_THREAD ((DV_VEC + Q1_WG_SIZE - 1) / Q1_WG_SIZE)
+
+inline float4 dequant_q8_0_lane(const global char * block_ptr, int lane) {
+    const float d = vload_half(0, (const global half *)block_ptr);
+    const global char * qs = block_ptr + 2 + lane * 4;
+    return d * (float4)((float)qs[0], (float)qs[1], (float)qs[2], (float)qs[3]);
+}
+
+REQD_SUBGROUP_SIZE_64
+__kernel void flash_attn_f32_q8_0_q1_vec(
+    const global void * q_void, ulong q_offset,
+    const global void * k_void, ulong k_offset,
+    const global void * v_void, ulong v_offset,
+    global void * o_void, ulong o_offset,
+    const float scale,
+    const int n_q,
+    const int n_kv,
+    const int is_causal,
+    const int n_head,
+    const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
+    const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
+    const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
+    const ulong o_nb1, const ulong o_nb2, const ulong o_nb3,
+    const float max_bias,
+    const float m0,
+    const float m1,
+    const int n_head_log2,
+    const float logit_softcap,
+    const int n_head_kv,
+    const global void* mask_void,
+    const ulong mask_offset,
+    const ulong mask_nb1,
+    const ulong mask_nb2,
+    const ulong mask_nb3,
+    const int mask_ne2,
+    const int mask_ne3,
+    const global void* sinks_void,
+    const ulong sinks_offset
+) {
+    const int tid             = get_local_id(0);
+    const int sgid            = tid / Q1_WG_SIZE;
+    const int tid_sg          = tid % Q1_WG_SIZE;
+    const int head_batch_idx  = get_global_id(1);
+
+    const int batch_idx = head_batch_idx / n_head;
+    const int head_idx  = head_batch_idx % n_head;
+
+    const int gqa_ratio   = n_head / n_head_kv;
+    const int head_kv_idx = head_idx / gqa_ratio;
+
+    const global char * q_base = (const global char *) q_void + q_offset;
+    const global char * k_base = (const global char *) k_void + k_offset;
+    const global char * v_base = (const global char *) v_void + v_offset;
+    global       char * o_base = (global       char *) o_void + o_offset;
+
+    const global char * mask_base = NULL;
+    if (mask_void != NULL) {
+        const int mask_head_idx  = head_idx  % mask_ne2;
+        const int mask_batch_idx = batch_idx % mask_ne3;
+        mask_base = (const global char *) mask_void + mask_offset +
+                    mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2;
+    }
+
+    __local ACC_TYPE4 q_shared[DK_VEC];
+    {
+        const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2;
+        const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
+        for (int i = tid; i < DK_VEC; i += VEC_WG_SIZE) {
+            q_shared[i] = CONVERT_Q_ACC4(q_ptr[i]);
+        }
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    const float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
+
+    const global ACC_TYPE * sinks_ptr = NULL;
+    if (sinks_void != NULL) {
+        sinks_ptr = (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset);
+    }
+
+    ACC_TYPE4 o_acc[Q1V_DV_PER_THREAD];
+    #pragma unroll
+    for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[i] = (ACC_TYPE4)(0.0f);
+
+    ACC_TYPE m_i = FA_M_INIT;
+    ACC_TYPE l_i = 0.0f;
+
+    const int kv_per_sg = (n_kv + VEC_NSG - 1) / VEC_NSG;
+    const int kv_start  = sgid * kv_per_sg;
+    const int kv_end    = min(n_kv, kv_start + kv_per_sg);
+
+    for (int k_idx = kv_start; k_idx < kv_end; ++k_idx) {
+        const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
+        const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
+
+        ACC_TYPE4 dot4 = (ACC_TYPE4)(0.0f);
+        for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
+            const int block_idx = qk / 8;
+            const int lane      = qk % 8;
+            const float4 k_v = dequant_q8_0_lane(k_row + block_idx * Q8_0_BLOCK_SIZE, lane);
+            dot4 = mad(q_shared[qk], k_v, dot4);
+        }
+        ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3;
+        ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale;
+
+        if (mask_base != NULL) {
+            const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
+            score += slope * (ACC_TYPE) mask_ptr[k_idx];
+        }
+        if (logit_softcap > 0.0f) {
+            score = logit_softcap * tanh(score / logit_softcap);
+        }
+
+        const ACC_TYPE m_new      = max(m_i, score);
+        const ACC_TYPE scale_prev = native_exp(m_i - m_new);
+        const ACC_TYPE p          = native_exp(score - m_new);
+
+        int idx = 0;
+        for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
+            const int block_idx = dv / 8;
+            const int lane      = dv % 8;
+            const float4 v_v = dequant_q8_0_lane(v_row + block_idx * Q8_0_BLOCK_SIZE, lane);
+            o_acc[idx] = mad(p, v_v, o_acc[idx] * scale_prev);
+        }
+        l_i = l_i * scale_prev + p;
+        m_i = m_new;
+    }
+
+    __local ACC_TYPE  sg_m[VEC_NSG];
+    __local ACC_TYPE  sg_l[VEC_NSG];
+    __local ACC_TYPE4 sg_o[VEC_NSG][DV_VEC];
+
+    if (tid_sg == 0) {
+        sg_m[sgid] = m_i;
+        sg_l[sgid] = l_i;
+    }
+    {
+        int idx = 0;
+        for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
+            sg_o[sgid][dv] = o_acc[idx];
+        }
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    if (sgid == 0) {
+        ACC_TYPE m_final = sg_m[0];
+        #pragma unroll
+        for (int s = 1; s < VEC_NSG; ++s) {
+            m_final = max(m_final, sg_m[s]);
+        }
+        if (sinks_ptr != NULL) {
+            m_final = max(m_final, sinks_ptr[head_idx]);
+        }
+
+        ACC_TYPE l_final = 0.0f;
+        #pragma unroll
+        for (int s = 0; s < VEC_NSG; ++s) {
+            l_final += sg_l[s] * native_exp(sg_m[s] - m_final);
+        }
+        if (sinks_ptr != NULL) {
+            l_final += native_exp(sinks_ptr[head_idx] - m_final);
+        }
+        const ACC_TYPE l_inv = (l_final > 0.0f) ? (1.0f / l_final) : 0.0f;
+
+        const ulong o_row_offset = batch_idx * o_nb3 + head_idx * o_nb1;
+        global O_DATA_TYPE4 * o_row = (global O_DATA_TYPE4 *) (o_base + o_row_offset);
+
+        int idx = 0;
+        for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
+            ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
+            #pragma unroll
+            for (int s = 0; s < VEC_NSG; ++s) {
+                const ACC_TYPE alpha = native_exp(sg_m[s] - m_final);
+                o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv], o_merged);
+            }
+            o_row[dv] = CONVERT_O_DATA4(o_merged * l_inv);
+        }
+    }
+}
+
 // Flash-decoding split pass for q8_0 KV. Partial record: [m, l, O[DV]].
 // Merge kernel from flash_attn_f32_f16.cl is type-agnostic and reused.
 #define FA_PARTIAL_FLOATS (2 + DV)
@@ -533,6 +732,244 @@ __kernel void flash_attn_f32_q8_0_q1_split(
 #define FA_V_STRATEGY 0
 #endif
 
+#ifndef MQ_GQA
+#define MQ_GQA 4
+#endif
+#ifndef MQ_NSG_SPLIT
+#define MQ_NSG_SPLIT 4
+#endif
+#define MQ_SPLIT_WG_SIZE_Q8 (Q1_WG_SIZE * MQ_NSG_SPLIT)
+
+REQD_SUBGROUP_SIZE_64
+__kernel void flash_attn_f32_q8_0_q1_vec_mq_split(
+    const global void * q_void, ulong q_offset,
+    const global void * k_void, ulong k_offset,
+    const global void * v_void, ulong v_offset,
+    const float scale,
+    const int n_q,
+    const int n_kv,
+    const int n_head,
+    const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
+    const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
+    const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
+    const float max_bias,
+    const float m0,
+    const float m1,
+    const int n_head_log2,
+    const float logit_softcap,
+    const int n_head_kv,
+    const global void * mask_void,
+    const ulong mask_offset,
+    const ulong mask_nb1,
+    const ulong mask_nb2,
+    const ulong mask_nb3,
+    const int mask_ne2,
+    const int mask_ne3,
+    global float * partial_void,
+    const int n_splits,
+    const int kv_per_split
+) {
+    const int tid              = get_local_id(0);
+    const int sgid             = tid / Q1_WG_SIZE;
+    const int tid_sg           = tid % Q1_WG_SIZE;
+    const int kvhead_batch_idx = get_global_id(1);
+    const int split_q_idx      = get_global_id(2);
+    const int split_idx        = split_q_idx % n_splits;
+    const int q_idx            = split_q_idx / n_splits;
+
+    const int batch_idx   = kvhead_batch_idx / n_head_kv;
+    const int head_kv_idx = kvhead_batch_idx % n_head_kv;
+
+    const int kv_start = split_idx * kv_per_split;
+    const int kv_end   = min(kv_start + kv_per_split, n_kv);
+
+    const ulong record_stride = (ulong) FA_PARTIAL_FLOATS;
+
+    if (kv_start >= kv_end) {
+        // Empty split — write sentinel for each of the MQ_GQA Q-heads.
+        if (tid == 0) {
+            #pragma unroll
+            for (int h = 0; h < MQ_GQA; ++h) {
+                const int head_idx = head_kv_idx * MQ_GQA + h;
+                const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
+                                       * n_splits + split_idx);
+                global float * rec = partial_void + rec_idx * record_stride;
+                rec[0] = FA_M_INIT;
+                rec[1] = 0.0f;
+            }
+        }
+        return;
+    }
+
+    const global char * q_base = (const global char *) q_void + q_offset;
+    const global char * k_base = (const global char *) k_void + k_offset;
+    const global char * v_base = (const global char *) v_void + v_offset;
+
+    __local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
+    for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_SPLIT_WG_SIZE_Q8) {
+        const int h        = i / DK_VEC;
+        const int k        = i % DK_VEC;
+        const int head_idx = head_kv_idx * MQ_GQA + h;
+        const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1;
+        const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
+        q_shared[h * DK_VEC + k] = CONVERT_Q_ACC4(q_ptr[k]);
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    float slope[MQ_GQA];
+    #pragma unroll
+    for (int h = 0; h < MQ_GQA; ++h) {
+        slope[h] = get_alibi_slope(max_bias, head_kv_idx * MQ_GQA + h, n_head_log2, m0, m1);
+    }
+
+    const global char * mask_base[MQ_GQA];
+    if (mask_void != NULL) {
+        const int mask_batch_idx = batch_idx % mask_ne3;
+        const global char * mask_base_b = (const global char *) mask_void + mask_offset +
+                                          mask_batch_idx * mask_nb3 +
+                                          (ulong) q_idx * mask_nb1;
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) {
+            const int head_idx      = head_kv_idx * MQ_GQA + h;
+            const int mask_head_idx = head_idx % mask_ne2;
+            mask_base[h] = mask_base_b + mask_head_idx * mask_nb2;
+        }
+    } else {
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) mask_base[h] = NULL;
+    }
+
+    ACC_TYPE4 o_acc[MQ_GQA][Q1V_DV_PER_THREAD];
+    ACC_TYPE  m_i[MQ_GQA];
+    ACC_TYPE  l_i[MQ_GQA];
+    #pragma unroll
+    for (int h = 0; h < MQ_GQA; ++h) {
+        m_i[h] = FA_M_INIT;
+        l_i[h] = 0.0f;
+        #pragma unroll
+        for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[h][i] = (ACC_TYPE4)(0.0f);
+    }
+
+    const int kv_len    = kv_end - kv_start;
+    const int kv_per_sg = (kv_len + MQ_NSG_SPLIT - 1) / MQ_NSG_SPLIT;
+    const int kv_lo     = kv_start + sgid * kv_per_sg;
+    const int kv_hi     = min(kv_end, kv_lo + kv_per_sg);
+
+    for (int k_idx = kv_lo; k_idx < kv_hi; ++k_idx) {
+        const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
+        const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
+
+        ACC_TYPE4 dot4[MQ_GQA];
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
+
+        for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
+            const int block_idx = qk / 8;
+            const int lane      = qk % 8;
+            const float4 k_v = dequant_q8_0_lane(k_row + block_idx * Q8_0_BLOCK_SIZE, lane);
+            #pragma unroll
+            for (int h = 0; h < MQ_GQA; ++h) {
+                dot4[h] = mad(q_shared[h * DK_VEC + qk], k_v, dot4[h]);
+            }
+        }
+
+        ACC_TYPE score[MQ_GQA];
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) {
+            const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
+            ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
+            if (mask_base[h] != NULL) {
+                const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
+                s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
+            }
+            if (logit_softcap > 0.0f) {
+                s = logit_softcap * tanh(s / logit_softcap);
+            }
+            score[h] = s;
+        }
+
+        ACC_TYPE p_h[MQ_GQA];
+        ACC_TYPE sp_h[MQ_GQA];
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) {
+            const ACC_TYPE m_new = max(m_i[h], score[h]);
+            sp_h[h] = native_exp(m_i[h] - m_new);
+            p_h[h]  = native_exp(score[h] - m_new);
+            l_i[h]  = l_i[h] * sp_h[h] + p_h[h];
+            m_i[h]  = m_new;
+        }
+
+        int idx = 0;
+        for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
+            const int block_idx = dv / 8;
+            const int lane      = dv % 8;
+            const float4 v_v = dequant_q8_0_lane(v_row + block_idx * Q8_0_BLOCK_SIZE, lane);
+            #pragma unroll
+            for (int h = 0; h < MQ_GQA; ++h) {
+                o_acc[h][idx] = mad(p_h[h], v_v, o_acc[h][idx] * sp_h[h]);
+            }
+        }
+    }
+
+    __local ACC_TYPE  sg_m[MQ_GQA][MQ_NSG_SPLIT];
+    __local ACC_TYPE  sg_l[MQ_GQA][MQ_NSG_SPLIT];
+    __local ACC_TYPE4 sg_o[MQ_NSG_SPLIT][DV_VEC];
+
+    if (tid_sg == 0) {
+        #pragma unroll
+        for (int h = 0; h < MQ_GQA; ++h) {
+            sg_m[h][sgid] = m_i[h];
+            sg_l[h][sgid] = l_i[h];
+        }
+    }
+
+    #pragma unroll
+    for (int h = 0; h < MQ_GQA; ++h) {
+        {
+            int idx = 0;
+            for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE, ++idx) {
+                sg_o[sgid][dv_idx] = o_acc[h][idx];
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+        if (sgid == 0) {
+            const int head_idx = head_kv_idx * MQ_GQA + h;
+
+            ACC_TYPE m_c = sg_m[h][0];
+            #pragma unroll
+            for (int s = 1; s < MQ_NSG_SPLIT; ++s) {
+                m_c = max(m_c, sg_m[h][s]);
+            }
+            ACC_TYPE l_c = 0.0f;
+            #pragma unroll
+            for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
+                l_c += sg_l[h][s] * native_exp(sg_m[h][s] - m_c);
+            }
+
+            const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
+                                   * n_splits + split_idx);
+            global float  * rec   = partial_void + rec_idx * record_stride;
+            global float4 * rec_o = (global float4 *) (rec + 2);
+
+            if (tid_sg == 0) {
+                rec[0] = (float) m_c;
+                rec[1] = (float) l_c;
+            }
+            for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE) {
+                ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
+                #pragma unroll
+                for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
+                    const ACC_TYPE alpha = native_exp(sg_m[h][s] - m_c);
+                    o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv_idx], o_merged);
+                }
+                rec_o[dv_idx] = o_merged;
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+}
+
 __kernel void flash_attn_f32_q8_0(
     const global void * q_void, ulong q_offset,
     const global void * k_void, ulong k_offset,

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

@@ -18,6 +18,14 @@
 #define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
 #endif
 
+#ifdef cl_khr_subgroup_shuffle
+#pragma OPENCL EXTENSION cl_khr_subgroup_shuffle : enable
+#define HAS_SUBGROUP_SHUFFLE 1
+#elif defined(cl_qcom_subgroup_shuffle)
+#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
+#define HAS_SUBGROUP_SHUFFLE 1
+#endif
+
 // Assumes row size (ne00) is a multiple of 4
 #ifdef ADRENO_GPU
 REQD_SUBGROUP_SIZE_64
@@ -378,3 +386,848 @@ kernel void kernel_mul_mat_f16_f32_l4_dr_lq(
     }
 }
 #endif // ADRENO_GPU
+
+#define N_ROWS_PER_WG 8
+#define N_OUTS_PER_WG 8
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_x8(
+        global char * src0,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne02,
+        ulong nb00,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne10,
+        int ne11,
+        int ne12,
+        ulong nb10,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src0 = (global char *)((global char *)src0 + offset0);
+    src1 = (global char *)((global char *)src1 + offset1);
+    dst  = (global float*)((global char *)dst  + offsetd);
+
+    const int sgs_lid = get_sub_group_local_id();
+    const int sgs_sz  = get_max_sub_group_size();
+
+    const int r0_base = get_group_id(0) * N_ROWS_PER_WG;
+    const int im      = get_group_id(2);
+
+    const int i12 = im % ne12;
+    const int i13 = im / ne12;
+
+    const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
+    global float4 * y4 = (global float4 *)(src1 + offset_src1);
+
+    __local float4 q_loc[64];   // ne00/4 max for sub_group_size 64
+    if (sgs_lid < ne00 / 4) {
+        q_loc[sgs_lid] = y4[sgs_lid];
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    #pragma unroll
+    for (int dr = 0; dr < N_ROWS_PER_WG; ++dr) {
+        const int r0 = r0_base + dr;
+        if (r0 >= ne01) return;
+
+        const ulong offset_src0 = r0 * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
+        global half4 * x4 = (global half4 *)(src0 + offset_src0);
+
+        float sumf = 0.0f;
+        for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
+            const half4   k4 = x4[i];
+            const float4  q  = q_loc[i];
+            sumf += convert_float(k4.s0) * q.s0
+                  + convert_float(k4.s1) * q.s1
+                  + convert_float(k4.s2) * q.s2
+                  + convert_float(k4.s3) * q.s3;
+        }
+
+        const float all_sum = sub_group_reduce_add(sumf);
+        if (sgs_lid == 0) {
+            dst[im * ne1 * ne0 + r0] = all_sum;  // ne11 == 1, so r1==0
+        }
+    }
+}
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_y8(
+        global char * src0,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne02,
+        ulong nb00,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne10,
+        int ne11,
+        int ne12,
+        ulong nb10,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src0 = (global char *)((global char *)src0 + offset0);
+    src1 = (global char *)((global char *)src1 + offset1);
+    dst  = (global float*)((global char *)dst  + offsetd);
+
+    const int sgs_lid = get_sub_group_local_id();
+    const int sgs_sz  = get_max_sub_group_size();
+
+    const int r0_base = get_group_id(0) * N_OUTS_PER_WG;
+    const int im      = get_group_id(2);
+
+    const int i12 = im % ne12;
+    const int i13 = im / ne12;
+
+    const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
+    global float4 * y4 = (global float4 *)(src1 + offset_src1);
+
+    global half4 * x4_o[N_OUTS_PER_WG];
+    #pragma unroll
+    for (int o = 0; o < N_OUTS_PER_WG; ++o) {
+        const int r0 = r0_base + o;
+        const int r0c = (r0 < ne01) ? r0 : 0;
+        const ulong off = r0c * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
+        x4_o[o] = (global half4 *)(src0 + off);
+    }
+
+    float sum[N_OUTS_PER_WG] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f };
+
+    for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
+        const float4 q4 = y4[i];
+        #pragma unroll
+        for (int o = 0; o < N_OUTS_PER_WG; ++o) {
+            const half4 v4 = x4_o[o][i];
+            sum[o] += convert_float(v4.s0) * q4.s0
+                    + convert_float(v4.s1) * q4.s1
+                    + convert_float(v4.s2) * q4.s2
+                    + convert_float(v4.s3) * q4.s3;
+        }
+    }
+
+    #pragma unroll
+    for (int o = 0; o < N_OUTS_PER_WG; ++o) {
+        const int r0 = r0_base + o;
+        const float s = sub_group_reduce_add(sum[o]);
+        if (sgs_lid == 0 && r0 < ne01) {
+            dst[im * ne1 * ne0 + r0] = s;
+        }
+    }
+}
+
+#define N_OUTS_PAIR  8
+#define N_PAIRS_PAIR (N_OUTS_PAIR / 2)
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_x8_pair(
+        global char * src0,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne02,
+        ulong nb00,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne10,
+        int ne11,
+        int ne12,
+        ulong nb10,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src0 = (global char *)((global char *)src0 + offset0);
+    src1 = (global char *)((global char *)src1 + offset1);
+    dst  = (global float*)((global char *)dst  + offsetd);
+
+    const int sgs_lid = get_sub_group_local_id();
+    const int half_id = sgs_lid >> 5;     // 0 = lower half, 1 = upper half
+    const int lane_h  = sgs_lid & 31;     // lane 0..31 within half
+
+    const int r0_base = get_group_id(0) * N_OUTS_PAIR;
+    const int im      = get_group_id(2);
+
+    const int i12 = im % ne12;
+    const int i13 = im / ne12;
+
+    const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
+    global float4 * y4 = (global float4 *)(src1 + offset_src1);
+
+    __local float4 q_loc[64];   // ne00/4 max for sub_group_size 64
+    if (sgs_lid < ne00 / 4) {
+        q_loc[sgs_lid] = y4[sgs_lid];
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    const int dk_vec = ne00 / 4;
+
+    #pragma unroll
+    for (int p = 0; p < N_PAIRS_PAIR; ++p) {
+        const int r0 = r0_base + 2 * p + half_id;
+
+        const ulong offset_src0 = r0 * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
+        global half4 * x4 = (global half4 *)(src0 + offset_src0);
+
+        float sumf = 0.0f;
+        for (int i = lane_h; i < dk_vec; i += 32) {
+            const half4  k4 = x4[i];
+            const float4 q  = q_loc[i];
+            sumf += convert_float(k4.s0) * q.s0
+                  + convert_float(k4.s1) * q.s1
+                  + convert_float(k4.s2) * q.s2
+                  + convert_float(k4.s3) * q.s3;
+        }
+
+        sumf += sub_group_shuffle_xor(sumf, 16);
+        sumf += sub_group_shuffle_xor(sumf, 8);
+        sumf += sub_group_shuffle_xor(sumf, 4);
+        sumf += sub_group_shuffle_xor(sumf, 2);
+        sumf += sub_group_shuffle_xor(sumf, 1);
+
+        if (lane_h == 0) {
+            dst[im * ne1 * ne0 + r0] = sumf;
+        }
+    }
+}
+
+#define N_K_ROWS_GQA   16
+#define GQA_RATIO_GQA  8
+#define LANES_PER_QH   8    // 64 / GQA_RATIO_GQA
+#define DK_VEC_GQA     32   // DK / 4 for DK=128
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_x8_gqa4(
+        global char * src0,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne02,
+        ulong nb00,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne10,
+        int ne11,
+        int ne12,
+        ulong nb10,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src0 = (global char *)((global char *)src0 + offset0);
+    src1 = (global char *)((global char *)src1 + offset1);
+    dst  = (global float*)((global char *)dst  + offsetd);
+
+    const int sgs_lid = get_sub_group_local_id();
+    const int q_id    = sgs_lid >> 3;       // 0..7: which Q-head (8 per WG)
+    const int lane_q  = sgs_lid & 7;        // 0..7: lane within Q-head partition
+
+    const int r0_base = get_group_id(0) * N_K_ROWS_GQA;
+    const int im_kv   = get_group_id(2);
+
+    const int i02 = im_kv % ne02;           // K-head index (also K2 batch)
+    const int i03 = im_kv / ne02;           // n13 batch index
+
+    const int q_head_lo = i02 * GQA_RATIO_GQA;
+
+    __local float4 q_loc[GQA_RATIO_GQA * DK_VEC_GQA];   // 4 × 32 = 128 float4
+    #pragma unroll
+    for (int qh = 0; qh < GQA_RATIO_GQA; ++qh) {
+        const int qh_idx = q_head_lo + qh;
+        global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
+
+        if (sgs_lid < DK_VEC_GQA) {
+            q_loc[qh * DK_VEC_GQA + sgs_lid] = y4[sgs_lid];
+        }
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    // K base offset for this WG. All 8 K-rows × 4 Q-heads share this K-head.
+    const ulong offset_src0_base = (i02) * nb02 + (i03 / r3) * nb03;
+
+    #pragma unroll
+    for (int dr = 0; dr < N_K_ROWS_GQA; ++dr) {
+        const int r0 = r0_base + dr;
+
+        const ulong offset_src0 = r0 * nb01 + offset_src0_base;
+        global half4 * x4 = (global half4 *)(src0 + offset_src0);
+
+        float sumf = 0.0f;
+        #pragma unroll
+        for (int t = 0; t < 4; ++t) {
+            const int i = lane_q + t * LANES_PER_QH;   // 8, 16, 24-step
+            const half4  k4 = x4[i];
+            const float4 q  = q_loc[q_id * DK_VEC_GQA + i];
+            sumf += convert_float(k4.s0) * q.s0
+                  + convert_float(k4.s1) * q.s1
+                  + convert_float(k4.s2) * q.s2
+                  + convert_float(k4.s3) * q.s3;
+        }
+
+        sumf += sub_group_shuffle_xor(sumf, 4);
+        sumf += sub_group_shuffle_xor(sumf, 2);
+        sumf += sub_group_shuffle_xor(sumf, 1);
+
+        if (lane_q == 0) {
+            const int im_out = i03 * ne12 + (q_head_lo + q_id);
+            dst[im_out * ne1 * ne0 + r0] = sumf;
+        }
+    }
+}
+
+#define N_DV_ROWS_Y8GQA  8
+#define GQA_RATIO_Y8GQA  8
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_y8_gqa(
+        global char * src0,
+        ulong offset0,
+        global char * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne02,
+        ulong nb00,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne10,
+        int ne11,
+        int ne12,
+        ulong nb10,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src0 = (global char *)((global char *)src0 + offset0);
+    src1 = (global char *)((global char *)src1 + offset1);
+    dst  = (global float*)((global char *)dst  + offsetd);
+
+    const int sgs_lid = get_sub_group_local_id();
+    const int sgs_sz  = get_max_sub_group_size();
+
+    const int r0_base = get_group_id(0) * N_DV_ROWS_Y8GQA;
+    const int im_kv   = get_group_id(2);
+
+    const int i02 = im_kv % ne02;           // K-head index
+    const int i03 = im_kv / ne02;           // n13 batch index
+
+    // GQA Q-heads sharing this K-head.
+    const int q_head_lo = i02 * GQA_RATIO_Y8GQA;
+
+    global float4 * y4_q[GQA_RATIO_Y8GQA];
+    #pragma unroll
+    for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+        const int qh_idx = q_head_lo + qh;
+        y4_q[qh] = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
+    }
+
+    global half4 * x4_o[N_DV_ROWS_Y8GQA];
+    #pragma unroll
+    for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+        const int r0 = r0_base + o;
+        const int r0c = (r0 < ne01) ? r0 : 0;
+        const ulong off = r0c * nb01 + (i02) * nb02 + (i03 / r3) * nb03;
+        x4_o[o] = (global half4 *)(src0 + off);
+    }
+
+    float sum[N_DV_ROWS_Y8GQA][GQA_RATIO_Y8GQA] = { {0.0f} };
+
+    for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
+        // load 8 V values (one per DV row), same K-head, K-pos = i.
+        half4 v[N_DV_ROWS_Y8GQA];
+        #pragma unroll
+        for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+            v[o] = x4_o[o][i];
+        }
+
+        // load 8 softmax values (one per Q-head).
+        float4 q[GQA_RATIO_Y8GQA];
+        #pragma unroll
+        for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+            q[qh] = y4_q[qh][i];
+        }
+
+        #pragma unroll
+        for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+            const float4 vf = (float4)(convert_float(v[o].s0),
+                                       convert_float(v[o].s1),
+                                       convert_float(v[o].s2),
+                                       convert_float(v[o].s3));
+            #pragma unroll
+            for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+                sum[o][qh] += vf.s0 * q[qh].s0
+                            + vf.s1 * q[qh].s1
+                            + vf.s2 * q[qh].s2
+                            + vf.s3 * q[qh].s3;
+            }
+        }
+    }
+
+    #pragma unroll
+    for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+        const int r0 = r0_base + o;
+        #pragma unroll
+        for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+            const float s = sub_group_reduce_add(sum[o][qh]);
+            if (sgs_lid == 0 && r0 < ne01) {
+                const int im_out = i03 * ne12 + (q_head_lo + qh);
+                dst[im_out * ne1 * ne0 + r0] = s;
+            }
+        }
+    }
+}
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_x8_gqa4_img(
+        __read_only image1d_buffer_t src0_img,
+        global char * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne02,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne10,
+        int ne11,
+        int ne12,
+        ulong nb10,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src1 = (global char *)((global char *)src1 + offset1);
+    dst  = (global float*)((global char *)dst  + offsetd);
+
+    const int sgs_lid = get_sub_group_local_id();
+    const int q_id    = sgs_lid >> 3;       // 0..7: which Q-head (8 per WG)
+    const int lane_q  = sgs_lid & 7;        // 0..7: lane within Q-head partition
+
+    const int r0_base = get_group_id(0) * N_K_ROWS_GQA;
+    const int im_kv   = get_group_id(2);
+
+    const int i02 = im_kv % ne02;
+    const int i03 = im_kv / ne02;
+
+    const int q_head_lo = i02 * GQA_RATIO_GQA;
+
+    __local float4 q_loc[GQA_RATIO_GQA * DK_VEC_GQA];
+    #pragma unroll
+    for (int qh = 0; qh < GQA_RATIO_GQA; ++qh) {
+        const int qh_idx = q_head_lo + qh;
+        global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
+        if (sgs_lid < DK_VEC_GQA) {
+            q_loc[qh * DK_VEC_GQA + sgs_lid] = y4[sgs_lid];
+        }
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    const int pitch_px_row  = (int)(nb01 >> 4);
+    const int pitch_px_head = (int)(nb02 >> 4);
+    const int pitch_px_n13  = (int)(nb03 >> 4);
+
+    const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
+
+    #pragma unroll
+    for (int dr = 0; dr < N_K_ROWS_GQA; ++dr) {
+        const int r0 = r0_base + dr;
+        const int row_px_base = r0 * pitch_px_row + head_px_base;
+
+        float sumf = 0.0f;
+        #pragma unroll
+        for (int t = 0; t < 2; ++t) {
+            const int p = lane_q + t * LANES_PER_QH;          // pixel idx in row, 0..15
+            const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
+            const int   i0 = 2 * p;                            // first half4 idx
+            const float4 qa = q_loc[q_id * DK_VEC_GQA + i0    ];
+            const float4 qb = q_loc[q_id * DK_VEC_GQA + i0 + 1];
+            sumf += convert_float(k8.s0) * qa.s0
+                  + convert_float(k8.s1) * qa.s1
+                  + convert_float(k8.s2) * qa.s2
+                  + convert_float(k8.s3) * qa.s3
+                  + convert_float(k8.s4) * qb.s0
+                  + convert_float(k8.s5) * qb.s1
+                  + convert_float(k8.s6) * qb.s2
+                  + convert_float(k8.s7) * qb.s3;
+        }
+
+        sumf += sub_group_shuffle_xor(sumf, 4);
+        sumf += sub_group_shuffle_xor(sumf, 2);
+        sumf += sub_group_shuffle_xor(sumf, 1);
+
+        if (lane_q == 0) {
+            const int im_out = i03 * ne12 + (q_head_lo + q_id);
+            dst[im_out * ne1 * ne0 + r0] = sumf;
+        }
+    }
+}
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_y8_gqa_img(
+        __read_only image1d_buffer_t src0_img,
+        global char * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne02,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne10,
+        int ne11,
+        int ne12,
+        ulong nb10,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src1 = (global char *)((global char *)src1 + offset1);
+    dst  = (global float*)((global char *)dst  + offsetd);
+
+    const int sgs_lid = get_sub_group_local_id();
+    const int sgs_sz  = get_max_sub_group_size();
+
+    const int r0_base = get_group_id(0) * N_DV_ROWS_Y8GQA;
+    const int im_kv   = get_group_id(2);
+
+    const int i02 = im_kv % ne02;
+    const int i03 = im_kv / ne02;
+
+    const int q_head_lo = i02 * GQA_RATIO_Y8GQA;
+
+    // Q (= softmax(KQ)) base pointers per Q-head
+    global float4 * y4_q[GQA_RATIO_Y8GQA];
+    #pragma unroll
+    for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+        const int qh_idx = q_head_lo + qh;
+        y4_q[qh] = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
+    }
+
+    const int pitch_px_row  = (int)(nb01 >> 3);
+    const int pitch_px_head = (int)(nb02 >> 3);
+    const int pitch_px_n13  = (int)(nb03 >> 3);
+
+    const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
+
+    // per-DV-row pixel base
+    int row_px_base[N_DV_ROWS_Y8GQA];
+    #pragma unroll
+    for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+        const int r0  = r0_base + o;
+        const int r0c = (r0 < ne01) ? r0 : 0;
+        row_px_base[o] = r0c * pitch_px_row + head_px_base;
+    }
+
+    float sum[N_DV_ROWS_Y8GQA][GQA_RATIO_Y8GQA] = { {0.0f} };
+
+    for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
+        half4 v[N_DV_ROWS_Y8GQA];
+
+        #pragma unroll
+        for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+            v[o] = read_imageh(src0_img, row_px_base[o] + i);
+        }
+
+        float4 q[GQA_RATIO_Y8GQA];
+        #pragma unroll
+        for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+            q[qh] = y4_q[qh][i];
+        }
+        // 64 mads.
+        #pragma unroll
+        for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+            const float4 vf = (float4)(convert_float(v[o].s0),
+                                       convert_float(v[o].s1),
+                                       convert_float(v[o].s2),
+                                       convert_float(v[o].s3));
+            #pragma unroll
+            for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+                sum[o][qh] += vf.s0 * q[qh].s0
+                            + vf.s1 * q[qh].s1
+                            + vf.s2 * q[qh].s2
+                            + vf.s3 * q[qh].s3;
+            }
+        }
+    }
+
+    #pragma unroll
+    for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+        const int r0 = r0_base + o;
+        #pragma unroll
+        for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+            const float s = sub_group_reduce_add(sum[o][qh]);
+            if (sgs_lid == 0 && r0 < ne01) {
+                const int im_out = i03 * ne12 + (q_head_lo + qh);
+                dst[im_out * ne1 * ne0 + r0] = s;
+            }
+        }
+    }
+}
+
+#define N_K_ROWS_GQA_R4   16
+#define GQA_RATIO_R4      4
+#define LANES_PER_QH_R4   16    // = 64 / GQA_RATIO_R4
+#define DK_VEC_R4         32    // DK / 4 for DK=128
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_x8_gqa_r4_img(
+        __read_only image1d_buffer_t src0_img,
+        global char * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne02,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne10,
+        int ne11,
+        int ne12,
+        ulong nb10,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src1 = (global char *)((global char *)src1 + offset1);
+    dst  = (global float*)((global char *)dst  + offsetd);
+
+    const int sgs_lid = get_sub_group_local_id();
+    const int q_id    = sgs_lid >> 4;       // 0..3
+    const int lane_q  = sgs_lid & 15;       // 0..15
+
+    const int r0_base = get_group_id(0) * N_K_ROWS_GQA_R4;
+    const int im_kv   = get_group_id(2);
+
+    const int i02 = im_kv % ne02;
+    const int i03 = im_kv / ne02;
+
+    const int q_head_lo = i02 * GQA_RATIO_R4;
+
+    __local float4 q_loc[GQA_RATIO_R4 * DK_VEC_R4];
+    #pragma unroll
+    for (int qh = 0; qh < GQA_RATIO_R4; ++qh) {
+        const int qh_idx = q_head_lo + qh;
+        global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
+        if (sgs_lid < DK_VEC_R4) {
+            q_loc[qh * DK_VEC_R4 + sgs_lid] = y4[sgs_lid];
+        }
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    const int pitch_px_row  = (int)(nb01 >> 4);
+    const int pitch_px_head = (int)(nb02 >> 4);
+    const int pitch_px_n13  = (int)(nb03 >> 4);
+
+    const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
+
+    #pragma unroll
+    for (int dr = 0; dr < N_K_ROWS_GQA_R4; ++dr) {
+        const int r0 = r0_base + dr;
+        const int row_px_base = r0 * pitch_px_row + head_px_base;
+
+        const int p = lane_q;
+        const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
+        const int   i0 = 2 * p;
+        const float4 qa = q_loc[q_id * DK_VEC_R4 + i0    ];
+        const float4 qb = q_loc[q_id * DK_VEC_R4 + i0 + 1];
+
+        float sumf =
+              convert_float(k8.s0) * qa.s0
+            + convert_float(k8.s1) * qa.s1
+            + convert_float(k8.s2) * qa.s2
+            + convert_float(k8.s3) * qa.s3
+            + convert_float(k8.s4) * qb.s0
+            + convert_float(k8.s5) * qb.s1
+            + convert_float(k8.s6) * qb.s2
+            + convert_float(k8.s7) * qb.s3;
+
+        sumf += sub_group_shuffle_xor(sumf, 8);
+        sumf += sub_group_shuffle_xor(sumf, 4);
+        sumf += sub_group_shuffle_xor(sumf, 2);
+        sumf += sub_group_shuffle_xor(sumf, 1);
+
+        if (lane_q == 0) {
+            const int im_out = i03 * ne12 + (q_head_lo + q_id);
+            dst[im_out * ne1 * ne0 + r0] = sumf;
+        }
+    }
+}
+
+#define N_K_ROWS_GQA_R2_DK256   16
+#define GQA_RATIO_R2            2
+#define LANES_PER_QH_R2         32    // = 64 / GQA_RATIO_R2
+#define DK_VEC_DK256            64    // DK / 4 for DK=256
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_x8_gqa_r2_dk256_img(
+        __read_only image1d_buffer_t src0_img,
+        global char * src1,
+        ulong offset1,
+        global float * dst,
+        ulong offsetd,
+        int ne00,
+        int ne01,
+        int ne02,
+        ulong nb01,
+        ulong nb02,
+        ulong nb03,
+        int ne10,
+        int ne11,
+        int ne12,
+        ulong nb10,
+        ulong nb11,
+        ulong nb12,
+        ulong nb13,
+        int ne0,
+        int ne1,
+        int r2,
+        int r3
+) {
+    src1 = (global char *)((global char *)src1 + offset1);
+    dst  = (global float*)((global char *)dst  + offsetd);
+
+    const int sgs_lid = get_sub_group_local_id();
+    const int q_id    = sgs_lid >> 5;       // 0..1
+    const int lane_q  = sgs_lid & 31;       // 0..31
+
+    const int r0_base = get_group_id(0) * N_K_ROWS_GQA_R2_DK256;
+    const int im_kv   = get_group_id(2);
+
+    const int i02 = im_kv % ne02;
+    const int i03 = im_kv / ne02;
+
+    const int q_head_lo = i02 * GQA_RATIO_R2;
+
+    __local float4 q_loc[GQA_RATIO_R2 * DK_VEC_DK256];
+    #pragma unroll
+    for (int qh = 0; qh < GQA_RATIO_R2; ++qh) {
+        const int qh_idx = q_head_lo + qh;
+        global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
+        q_loc[qh * DK_VEC_DK256 + sgs_lid] = y4[sgs_lid];
+    }
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    const int pitch_px_row  = (int)(nb01 >> 4);
+    const int pitch_px_head = (int)(nb02 >> 4);
+    const int pitch_px_n13  = (int)(nb03 >> 4);
+
+    const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
+
+    #pragma unroll
+    for (int dr = 0; dr < N_K_ROWS_GQA_R2_DK256; ++dr) {
+        const int r0 = r0_base + dr;
+        const int row_px_base = r0 * pitch_px_row + head_px_base;
+
+        const int p = lane_q;
+        const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
+        const int   i0 = 2 * p;
+        const float4 qa = q_loc[q_id * DK_VEC_DK256 + i0    ];
+        const float4 qb = q_loc[q_id * DK_VEC_DK256 + i0 + 1];
+
+        float sumf =
+              convert_float(k8.s0) * qa.s0
+            + convert_float(k8.s1) * qa.s1
+            + convert_float(k8.s2) * qa.s2
+            + convert_float(k8.s3) * qa.s3
+            + convert_float(k8.s4) * qb.s0
+            + convert_float(k8.s5) * qb.s1
+            + convert_float(k8.s6) * qb.s2
+            + convert_float(k8.s7) * qb.s3;
+
+        sumf += sub_group_shuffle_xor(sumf, 16);
+        sumf += sub_group_shuffle_xor(sumf, 8);
+        sumf += sub_group_shuffle_xor(sumf, 4);
+        sumf += sub_group_shuffle_xor(sumf, 2);
+        sumf += sub_group_shuffle_xor(sumf, 1);
+
+        if (lane_q == 0) {
+            const int im_out = i03 * ne12 + (q_head_lo + q_id);
+            dst[im_out * ne1 * ne0 + r0] = sumf;
+        }
+    }
+}

ggml/src/ggml-quants.c

@@ -71,6 +71,44 @@ void quantize_row_q1_0_ref(const float * GGML_RESTRICT x, block_q1_0 * GGML_REST
     }
 }
 
+void quantize_row_q2_0_ref(const float * GGML_RESTRICT x, block_q2_0 * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK2_0;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        // Compute scale as max absolute value in the block
+        float amax = 0.0f;
+        for (int j = 0; j < qk; j++) {
+            const float a = fabsf(x[i*qk + j]);
+            if (a > amax) amax = a;
+        }
+        const float d = amax;
+        const float id = d > 0.0f ? 1.0f / d : 0.0f;
+
+        y[i].d = GGML_FP32_TO_FP16(d);
+
+        // Clear quant bytes
+        for (int j = 0; j < qk / 4; ++j) {
+            y[i].qs[j] = 0;
+        }
+
+        // Encode 2-bit values: round(w/d) clamped to [-1, 2], then add 1
+        // 00 (-1) = -scale, 01 (0) = 0, 10 (+1) = +scale, 11 (+2) = 2*scale
+        for (int j = 0; j < qk; ++j) {
+            const float w = x[i*qk + j];
+            int q = (int)roundf(w * id) + 1;
+            if (q < 0) q = 0;
+            if (q > 3) q = 3;
+            const int byte_index = j / 4;
+            const int bit_offset = (j % 4) * 2;
+            y[i].qs[byte_index] |= ((uint8_t)q << bit_offset);
+        }
+    }
+}
+
 // reference implementation for deterministic creation of model files
 void quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t k) {
     static const int qk = QK4_0;
@@ -398,6 +436,26 @@ void dequantize_row_q1_0(const block_q1_0 * GGML_RESTRICT x, float * GGML_RESTRI
     }
 }
 
+void dequantize_row_q2_0(const block_q2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
+    static const int qk = QK2_0;
+
+    assert(k % qk == 0);
+
+    const int nb = k / qk;
+
+    for (int i = 0; i < nb; i++) {
+        const float d = GGML_FP16_TO_FP32(x[i].d);
+
+        for (int j = 0; j < qk; ++j) {
+            const int byte_index = j / 4;
+            const int bit_offset = (j % 4) * 2;
+            const uint8_t q = (x[i].qs[byte_index] >> bit_offset) & 0x03;
+            // 00=-1, 01=0, 10=+1, 11=+2
+            y[i*qk + j] = ((int)q - 1) * d;
+        }
+    }
+}
+
 void dequantize_row_q4_0(const block_q4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) {
     static const int qk = QK4_0;
 
@@ -2052,6 +2110,20 @@ size_t quantize_q1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst,
     return nrow * row_size;
 }
 
+size_t quantize_q2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
+    if (!quant_weights) {
+        quantize_row_q2_0_ref(src, dst, (int64_t)nrow*n_per_row);
+        return nrow * ggml_row_size(GGML_TYPE_Q2_0, n_per_row);
+    }
+    size_t row_size = ggml_row_size(GGML_TYPE_Q2_0, n_per_row);
+    char * qrow = (char *)dst;
+    for (int64_t row = 0; row < nrow; ++row) {
+        quantize_row_q2_0_ref(src, (block_q2_0*)qrow, n_per_row);
+        src += n_per_row;
+        qrow += row_size;
+    }
+    return nrow * row_size;
+}
 
 size_t quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
     if (!quant_weights) {
@@ -5461,6 +5533,10 @@ bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbyte
             {
                 VALIDATE_ROW_DATA_D_F16_IMPL(block_q1_0, data, nb);
             } break;
+        case GGML_TYPE_Q2_0:
+            {
+                VALIDATE_ROW_DATA_D_F16_IMPL(block_q2_0, data, nb);
+            } break;
         case GGML_TYPE_Q4_0:
             {
                 VALIDATE_ROW_DATA_D_F16_IMPL(block_q4_0, data, nb);

ggml/src/ggml-quants.h

@@ -15,6 +15,7 @@ extern "C" {
 
 // Quantization
 GGML_API void quantize_row_q1_0_ref(const float * GGML_RESTRICT x, block_q1_0 * GGML_RESTRICT y, int64_t k);
+GGML_API void quantize_row_q2_0_ref(const float * GGML_RESTRICT x, block_q2_0 * GGML_RESTRICT y, int64_t k);
 GGML_API void quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t k);
 GGML_API void quantize_row_q4_1_ref(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t k);
 GGML_API void quantize_row_q5_0_ref(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t k);
@@ -43,6 +44,7 @@ GGML_API void quantize_row_iq2_s_ref  (const float * GGML_RESTRICT x, block_iq2_
 
 // Dequantization
 GGML_API void dequantize_row_q1_0(const block_q1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
+GGML_API void dequantize_row_q2_0(const block_q2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
 GGML_API void dequantize_row_q4_0(const block_q4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
 GGML_API void dequantize_row_q4_1(const block_q4_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
 GGML_API void dequantize_row_q5_0(const block_q5_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k);
@@ -93,6 +95,7 @@ GGML_API size_t quantize_q4_K(const float * GGML_RESTRICT src, void * GGML_RESTR
 GGML_API size_t quantize_q5_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
 GGML_API size_t quantize_q6_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
 GGML_API size_t quantize_q1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
+GGML_API size_t quantize_q2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
 GGML_API size_t quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
 GGML_API size_t quantize_q4_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);
 GGML_API size_t quantize_q5_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix);

ggml/src/ggml-sycl/backend.hpp

@@ -14,6 +14,7 @@
 #define GGML_SYCL_BACKEND_HPP
 
 #include "binbcast.hpp"
+#include "col2im-1d.hpp"
 #include "common.hpp"
 #include "concat.hpp"
 #include "conv.hpp"

ggml/src/ggml-sycl/col2im-1d.cpp

@@ -0,0 +1,102 @@
+#include "col2im-1d.hpp"
+
+template <typename T>
+static void col2im_1d_sycl(
+        const T * col,
+        T * dst,
+        const int T_in,
+        const sycl::uint3 T_out_fd,
+        const int K,
+        const int K_OC,
+        const int32_t s0,
+        const int32_t p0,
+        const int total,
+        dpct::queue_ptr stream) {
+
+    const uint32_t block_size = SYCL_COL2IM_1D_BLOCK_SIZE;
+    const uint32_t num_blocks = (uint32_t) ((total + block_size - 1) / block_size);
+
+    stream->parallel_for(
+        sycl::nd_range<3>(
+            sycl::range<3>(1, 1, num_blocks * block_size),
+            sycl::range<3>(1, 1, block_size)),
+        [=](sycl::nd_item<3> item_ct1) {
+            const int idx = (int) item_ct1.get_global_id(2);
+            if (idx >= total) {
+                return;
+            }
+
+            const sycl::uint2 qr = fast_div_modulo((uint32_t) idx, T_out_fd);
+            const int oc    = (int) qr.x();
+            const int t_out = (int) qr.y();
+            const int t_abs = t_out + p0;
+
+            int t_in_min = (t_abs - K + s0) / s0;
+            if (t_in_min < 0) {
+                t_in_min = 0;
+            }
+            int t_in_max = t_abs / s0;
+            if (t_in_max >= T_in) {
+                t_in_max = T_in - 1;
+            }
+
+            float sum = 0.0f;
+            for (int t_in = t_in_min; t_in <= t_in_max; ++t_in) {
+                const int k = t_abs - t_in * s0;
+                sum += static_cast<float>(col[(oc * K + k) + t_in * K_OC]);
+            }
+
+            dst[idx] = static_cast<T>(sum);
+        });
+}
+
+void ggml_sycl_op_col2im_1d(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+
+    GGML_ASSERT(src0 != nullptr);
+    GGML_ASSERT(ggml_is_contiguous(src0));
+    GGML_ASSERT(src0->type == dst->type);
+
+    const int32_t s0 = ((const int32_t *) dst->op_params)[0];
+    const int32_t OC = ((const int32_t *) dst->op_params)[1];
+    const int32_t p0 = ((const int32_t *) dst->op_params)[2];
+
+    const int K_OC  = (int) src0->ne[0];
+    const int T_in  = (int) src0->ne[1];
+    const int K     = K_OC / OC;
+    const int T_out = (int) dst->ne[0];
+
+    GGML_ASSERT(OC > 0);
+    GGML_ASSERT(K_OC % OC == 0);
+
+    const sycl::uint3 T_out_fd = init_fastdiv_values((uint32_t) T_out);
+
+    const int total = T_out * OC;
+
+    dpct::queue_ptr stream = ctx.stream();
+
+    switch (src0->type) {
+        case GGML_TYPE_F32:
+            col2im_1d_sycl<float>(
+                (const float *) src0->data,
+                (float *) dst->data,
+                T_in, T_out_fd, K, K_OC, s0, p0, total, stream);
+            break;
+        case GGML_TYPE_F16:
+            col2im_1d_sycl<sycl::half>(
+                (const sycl::half *) src0->data,
+                (sycl::half *) dst->data,
+                T_in, T_out_fd, K, K_OC, s0, p0, total, stream);
+            break;
+#ifdef GGML_SYCL_HAS_BF16
+        case GGML_TYPE_BF16:
+            col2im_1d_sycl<sycl::ext::oneapi::bfloat16>(
+                (const sycl::ext::oneapi::bfloat16 *) src0->data,
+                (sycl::ext::oneapi::bfloat16 *) dst->data,
+                T_in, T_out_fd, K, K_OC, s0, p0, total, stream);
+            break;
+#endif
+        default:
+            GGML_ABORT("col2im_1d: unsupported type %d", src0->type);
+    }
+}

ggml/src/ggml-sycl/col2im-1d.hpp

@@ -0,0 +1,8 @@
+#ifndef GGML_SYCL_COL2IM_1D_HPP
+#define GGML_SYCL_COL2IM_1D_HPP
+
+#include "common.hpp"
+
+void ggml_sycl_op_col2im_1d(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
+
+#endif // GGML_SYCL_COL2IM_1D_HPP

ggml/src/ggml-sycl/common.hpp

@@ -59,7 +59,7 @@ void ggml_sycl_host_free(void* ptr);
 
 
 extern int g_ggml_sycl_debug;
-extern int g_ggml_sycl_disable_optimize;
+extern int g_ggml_sycl_enable_optimize;
 extern int g_ggml_sycl_prioritize_dmmv;
 extern int g_ggml_sycl_enable_flash_attention;
 extern int g_ggml_sycl_dev2dev_memcpy;

ggml/src/ggml-sycl/cpy.cpp

@@ -1,6 +1,7 @@
 #include "cpy.hpp"
 
 #include <float.h>
+#include <vector>
 
 #include "dequantize.hpp"
 #include "ggml-sycl/common.hpp"
@@ -50,6 +51,57 @@ static void cpy_1_i32_i32(const char * cxi, char * cdsti) {
     *dsti = *xi;
 }
 
+static void cpy_1_f32_i32(const char * cxi, char * cdsti) {
+    const float * xi     = (const float *) cxi;
+    int32_t *     dsti   = (int32_t *) cdsti;
+
+    *dsti = (int32_t) *xi;
+}
+
+static void cpy_1_i32_f32(const char * cxi, char * cdsti) {
+    const int32_t * xi   = (const int32_t *) cxi;
+    float *         dsti = (float *) cdsti;
+
+    *dsti = (float) *xi;
+}
+
+#ifdef GGML_SYCL_HAS_BF16
+static void cpy_1_f32_bf16(const char * cxi, char * cdsti) {
+    const float * xi   = (const float *) cxi;
+    sycl::ext::oneapi::bfloat16 *  dsti = (sycl::ext::oneapi::bfloat16 *) cdsti;
+
+    *dsti = sycl::ext::oneapi::bfloat16(*xi);
+}
+
+static void cpy_1_bf16_f32(const char * cxi, char * cdsti) {
+    const sycl::ext::oneapi::bfloat16 * xi   = (const sycl::ext::oneapi::bfloat16 *) cxi;
+    float *            dsti = (float *) cdsti;
+
+    *dsti = static_cast<float>(*xi);
+}
+
+static void cpy_1_bf16_bf16(const char * cxi, char * cdsti) {
+    const sycl::ext::oneapi::bfloat16 * xi   = (const sycl::ext::oneapi::bfloat16 *) cxi;
+    sycl::ext::oneapi::bfloat16 *       dsti = (sycl::ext::oneapi::bfloat16 *) cdsti;
+
+    *dsti = *xi;
+}
+
+static void cpy_1_f16_bf16(const char * cxi, char * cdsti) {
+    const sycl::half * xi   = (const sycl::half *) cxi;
+    sycl::ext::oneapi::bfloat16 *  dsti = (sycl::ext::oneapi::bfloat16 *) cdsti;
+
+    *dsti = sycl::ext::oneapi::bfloat16(static_cast<float>(*xi));
+}
+
+static void cpy_1_bf16_f16(const char * cxi, char * cdsti) {
+    const sycl::ext::oneapi::bfloat16 * xi   = (const sycl::ext::oneapi::bfloat16 *) cxi;
+    sycl::half *            dsti = (sycl::half *) cdsti;
+
+    *dsti = sycl::half(static_cast<float>(*xi));
+}
+#endif
+
 template <cpy_kernel_t cpy_1>
 static void cpy_f32_f16(const char * cx, char * cdst, const int ne, const int ne00, const int ne01, const int ne02,
                         const int nb00, const int nb01, const int nb02, const int nb03, const int ne10, const int ne11,
@@ -247,6 +299,38 @@ static void ggml_cpy_f32_f16_sycl(const char * cx, char * cdst, const int ne, co
     }
 }
 
+static void ggml_cpy_f32_i32_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                  const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                  const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                  const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = (ne + SYCL_CPY_BLOCK_SIZE - 1) / SYCL_CPY_BLOCK_SIZE;
+    {
+        stream->parallel_for(
+            sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
+                              sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+            [=](sycl::nd_item<3> item_ct1) {
+                cpy_f32_f16<cpy_1_f32_i32>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
+                                           nb10, nb11, nb12, nb13, item_ct1);
+            });
+    }
+}
+
+static void ggml_cpy_i32_f32_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                  const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                  const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                  const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = (ne + SYCL_CPY_BLOCK_SIZE - 1) / SYCL_CPY_BLOCK_SIZE;
+    {
+        stream->parallel_for(
+            sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
+                              sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+            [=](sycl::nd_item<3> item_ct1) {
+                cpy_f32_f16<cpy_1_i32_f32>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
+                                           nb10, nb11, nb12, nb13, item_ct1);
+            });
+    }
+}
+
 static void ggml_cpy_f32_q8_0_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
                                    const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
                                    const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
@@ -376,6 +460,19 @@ static void ggml_cpy_q5_1_f32_sycl(const char * cx, char * cdst, const int ne, c
         });
 }
 
+static void ggml_cpy_mxfp4_f32_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                    const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                    const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                    const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ne;
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_f32<cpy_blck_q_f32<dequantize_mxfp4, QK_MXFP4>, QK_MXFP4>(cx, cdst, ne, ne00, ne01, ne02, nb00,
+                                                                             nb01, nb02, nb03, ne10, ne11, ne12,
+                                                                             nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
 static void ggml_cpy_f32_iq4_nl_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
                                      const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
                                      const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
@@ -389,6 +486,269 @@ static void ggml_cpy_f32_iq4_nl_sycl(const char * cx, char * cdst, const int ne,
         });
 }
 
+static void cpy_blck_f16_q4_0(const char * cxi, char * cdsti) {
+    const sycl::half * xi = (const sycl::half *) cxi;
+    float xf[QK4_0];
+
+    for (int j = 0; j < QK4_0; ++j) {
+        xf[j] = (float) xi[j];
+    }
+
+    cpy_blck_f32_q4_0((const char *) xf, cdsti);
+}
+
+static void cpy_blck_f16_q4_1(const char * cxi, char * cdsti) {
+    const sycl::half * xi = (const sycl::half *) cxi;
+    float xf[QK4_1];
+
+    for (int j = 0; j < QK4_1; ++j) {
+        xf[j] = (float) xi[j];
+    }
+
+    cpy_blck_f32_q4_1((const char *) xf, cdsti);
+}
+
+static void cpy_blck_f16_q5_0(const char * cxi, char * cdsti) {
+    const sycl::half * xi = (const sycl::half *) cxi;
+    float xf[QK5_0];
+
+    for (int j = 0; j < QK5_0; ++j) {
+        xf[j] = (float) xi[j];
+    }
+
+    cpy_blck_f32_q5_0((const char *) xf, cdsti);
+}
+
+static void ggml_cpy_f16_q4_0_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    GGML_ASSERT(ne % QK4_0 == 0);
+    const int num_blocks = ne / QK4_0;
+    stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
+                         [=](sycl::nd_item<3> item_ct1) {
+                             cpy_f32_q<cpy_blck_f16_q4_0, QK4_0>(cx, cdst, ne, ne00, ne01, ne02,
+                                                                 nb00, nb01, nb02, nb03,
+                                                                 ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+                         });
+}
+
+static void ggml_cpy_f16_q4_1_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    GGML_ASSERT(ne % QK4_1 == 0);
+    const int num_blocks = ne / QK4_1;
+    stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
+                         [=](sycl::nd_item<3> item_ct1) {
+                             cpy_f32_q<cpy_blck_f16_q4_1, QK4_1>(cx, cdst, ne, ne00, ne01, ne02,
+                                                                 nb00, nb01, nb02, nb03,
+                                                                 ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+                         });
+}
+
+static void ggml_cpy_f16_q5_0_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    GGML_ASSERT(ne % QK5_0 == 0);
+    const int num_blocks = ne / QK5_0;
+    stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks), sycl::range<3>(1, 1, 1)),
+                         [=](sycl::nd_item<3> item_ct1) {
+                             cpy_f32_q<cpy_blck_f16_q5_0, QK5_0>(cx, cdst, ne, ne00, ne01, ne02,
+                                                                 nb00, nb01, nb02, nb03,
+                                                                 ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+                         });
+}
+
+static bool ggml_sycl_is_quantized_type(enum ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q5_1:
+        case GGML_TYPE_Q8_0:
+        case GGML_TYPE_MXFP4:
+        case GGML_TYPE_NVFP4:
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
+        case GGML_TYPE_IQ2_XXS:
+        case GGML_TYPE_IQ2_XS:
+        case GGML_TYPE_IQ2_S:
+        case GGML_TYPE_IQ3_XXS:
+        case GGML_TYPE_IQ3_S:
+        case GGML_TYPE_IQ1_S:
+        case GGML_TYPE_IQ1_M:
+        case GGML_TYPE_IQ4_NL:
+        case GGML_TYPE_IQ4_XS:
+            return true;
+        default:
+            return false;
+    }
+}
+
+static bool ggml_sycl_can_quantize_rows_sycl(enum ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_Q1_0:
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q5_1:
+        case GGML_TYPE_Q8_0:
+        case GGML_TYPE_MXFP4:
+        case GGML_TYPE_NVFP4:
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
+        case GGML_TYPE_IQ4_NL:
+        case GGML_TYPE_IQ4_XS:
+            return true;
+        default:
+            return false;
+    }
+}
+
+template <typename SrcScalar>
+static inline float ggml_sycl_src_to_f32(const SrcScalar & x) {
+    return (float) x;
+}
+
+#ifdef GGML_SYCL_HAS_BF16
+template <>
+inline float ggml_sycl_src_to_f32<sycl::ext::oneapi::bfloat16>(const sycl::ext::oneapi::bfloat16 & x) {
+    return static_cast<float>(x);
+}
+
+template <>
+inline float ggml_sycl_src_to_f32<ggml_bf16_t>(const ggml_bf16_t & x) {
+    union {
+        uint32_t u32;
+        float f32;
+    } value;
+
+    value.u32 = (uint32_t) x.bits << 16;
+    return value.f32;
+}
+#endif
+
+template <typename SrcScalar, cpy_kernel_t quantize_block, int qk>
+static void ggml_sycl_quantize_rows_q(const char * cx, char * cdst, const int64_t ne,
+                                      const int64_t ne00, const int64_t ne01, const int64_t ne02,
+                                      const size_t nb00, const size_t nb01, const size_t nb02, const size_t nb03,
+                                      const int64_t ne10, const int64_t ne11, const int64_t ne12,
+                                      const size_t nb10, const size_t nb11, const size_t nb12, const size_t nb13,
+                                      queue_ptr stream) {
+    GGML_ASSERT(ne % qk == 0);
+    GGML_ASSERT(ne00 % qk == 0);
+
+    const int64_t total_blocks = ne / qk;
+    constexpr int block_size = 256;
+    const int64_t grid_size = ceil_div(total_blocks, (int64_t) block_size);
+
+    stream->parallel_for(sycl::nd_range<1>(grid_size * block_size, block_size), [=](sycl::nd_item<1> item_ct1) {
+        const int64_t block_idx = item_ct1.get_global_linear_id();
+        if (block_idx >= total_blocks) {
+            return;
+        }
+
+        const int64_t i = block_idx * qk;
+
+        const int64_t i03 = i / (ne00 * ne01 * ne02);
+        const int64_t i02 = (i - i03 * ne00 * ne01 * ne02) / (ne00 * ne01);
+        const int64_t i01 = (i - i03 * ne00 * ne01 * ne02 - i02 * ne01 * ne00) / ne00;
+        const int64_t i00 = i - i03 * ne00 * ne01 * ne02 - i02 * ne01 * ne00 - i01 * ne00;
+        const size_t x_offset = i00 * nb00 + i01 * nb01 + i02 * nb02 + i03 * nb03;
+
+        const int64_t i13 = i / (ne10 * ne11 * ne12);
+        const int64_t i12 = (i - i13 * ne10 * ne11 * ne12) / (ne10 * ne11);
+        const int64_t i11 = (i - i13 * ne10 * ne11 * ne12 - i12 * ne10 * ne11) / ne10;
+        const int64_t i10 = i - i13 * ne10 * ne11 * ne12 - i12 * ne10 * ne11 - i11 * ne10;
+        const size_t dst_offset = (i10 / qk) * nb10 + i11 * nb11 + i12 * nb12 + i13 * nb13;
+
+        float xf[qk];
+        if (nb00 == sizeof(SrcScalar)) {
+            const SrcScalar * src_row = (const SrcScalar *) (cx + x_offset);
+            for (int j = 0; j < qk; ++j) {
+                xf[j] = ggml_sycl_src_to_f32(src_row[j]);
+            }
+        } else {
+            for (int j = 0; j < qk; ++j) {
+                const SrcScalar * src_val = (const SrcScalar *) (cx + x_offset + j * nb00);
+                xf[j] = ggml_sycl_src_to_f32(*src_val);
+            }
+        }
+
+        quantize_block((const char *) xf, cdst + dst_offset);
+    });
+}
+
+template <typename SrcScalar>
+static void ggml_sycl_quantize_rows_sycl(const char * cx, char * cdst, const ggml_tensor * src0, const ggml_tensor * src1,
+                                         const int64_t ne, const int64_t ne00, const int64_t ne01, const int64_t ne02,
+                                         const size_t nb00, const size_t nb01, const size_t nb02, const size_t nb03,
+                                         const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10,
+                                         const size_t nb11, const size_t nb12, const size_t nb13, queue_ptr stream) {
+    GGML_UNUSED(src0);
+    GGML_UNUSED(src1);
+
+    switch (src1->type) {
+        case GGML_TYPE_Q8_0:
+            ggml_sycl_quantize_rows_q<SrcScalar, cpy_blck_f32_q8_0, QK8_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01,
+                                                                            nb02, nb03, ne10, ne11, ne12, nb10, nb11,
+                                                                            nb12, nb13, stream);
+            break;
+        case GGML_TYPE_Q1_0:
+            ggml_sycl_quantize_rows_q<SrcScalar, cpy_blck_f32_q1_0, QK1_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01,
+                                                                            nb02, nb03, ne10, ne11, ne12, nb10, nb11,
+                                                                            nb12, nb13, stream);
+            break;
+        case GGML_TYPE_Q5_1:
+            ggml_sycl_quantize_rows_q<SrcScalar, cpy_blck_f32_q5_1, QK5_1>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01,
+                                                                            nb02, nb03, ne10, ne11, ne12, nb10, nb11,
+                                                                            nb12, nb13, stream);
+            break;
+        case GGML_TYPE_Q5_0:
+            ggml_sycl_quantize_rows_q<SrcScalar, cpy_blck_f32_q5_0, QK5_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01,
+                                                                            nb02, nb03, ne10, ne11, ne12, nb10, nb11,
+                                                                            nb12, nb13, stream);
+            break;
+        case GGML_TYPE_Q4_1:
+            ggml_sycl_quantize_rows_q<SrcScalar, cpy_blck_f32_q4_1, QK4_1>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01,
+                                                                            nb02, nb03, ne10, ne11, ne12, nb10, nb11,
+                                                                            nb12, nb13, stream);
+            break;
+        case GGML_TYPE_Q4_0:
+            ggml_sycl_quantize_rows_q<SrcScalar, cpy_blck_f32_q4_0, QK4_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01,
+                                                                            nb02, nb03, ne10, ne11, ne12, nb10, nb11,
+                                                                            nb12, nb13, stream);
+            break;
+        case GGML_TYPE_IQ4_NL:
+            ggml_sycl_quantize_rows_q<SrcScalar, cpy_blck_f32_iq4_nl, QK4_NL>(cx, cdst, ne, ne00, ne01, ne02, nb00,
+                                                                               nb01, nb02, nb03, ne10, ne11, ne12,
+                                                                               nb10, nb11, nb12, nb13, stream);
+            break;
+        case GGML_TYPE_MXFP4:
+            ggml_sycl_quantize_rows_q<SrcScalar, cpy_blck_f32_mxfp4, QK_MXFP4>(cx, cdst, ne, ne00, ne01, ne02, nb00,
+                                                                                nb01, nb02, nb03, ne10, ne11, ne12,
+                                                                                nb10, nb11, nb12, nb13, stream);
+            break;
+        case GGML_TYPE_NVFP4:
+            ggml_sycl_quantize_rows_q<SrcScalar, cpy_blck_f32_nvfp4, QK_NVFP4>(cx, cdst, ne, ne00, ne01, ne02, nb00,
+                                                                                nb01, nb02, nb03, ne10, ne11, ne12,
+                                                                                nb10, nb11, nb12, nb13, stream);
+            break;
+        default:
+            GGML_ABORT("unsupported quantized target type in sycl quantizer src1->type=%s\n",
+                       ggml_type_name(src1->type));
+    }
+}
+
 static void ggml_cpy_f16_f16_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
@@ -509,8 +869,269 @@ static void ggml_cpy_q4_1_q4_1(const char * cx, char * cdst, const int ne, const
         });
 }
 
+static void ggml_cpy_q1_0_q1_0(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_q1_0, QK1_0>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_mxfp4_mxfp4(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_mxfp4, QK_MXFP4>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_nvfp4_nvfp4(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_nvfp4, QK_NVFP4>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_q2_K_q2_K(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_q2_K, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_q3_K_q3_K(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_q3_K, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_q4_K_q4_K(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_q4_K, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_q5_K_q5_K(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_q5_K, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_q6_K_q6_K(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_q6_K, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_iq2_xxs_iq2_xxs(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_iq2_xxs, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_iq2_xs_iq2_xs(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_iq2_xs, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_iq2_s_iq2_s(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_iq2_s, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_iq3_xxs_iq3_xxs(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_iq3_xxs, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_iq1_s_iq1_s(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_iq1_s, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_iq1_m_iq1_m(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_iq1_m, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_iq4_nl_iq4_nl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_iq4_nl, QK4_NL>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_iq3_s_iq3_s(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_iq3_s, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_iq4_xs_iq4_xs(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                   const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                   const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                   const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = ceil_div(ne, SYCL_CPY_BLOCK_SIZE);
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE), sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)), [=](sycl::nd_item<3> item_ct1) {
+            cpy_q_q<block_iq4_xs, QK_K>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+#ifdef GGML_SYCL_HAS_BF16
+static void ggml_cpy_f32_bf16_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                  const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                  const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                  const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = (ne + SYCL_CPY_BLOCK_SIZE - 1) / SYCL_CPY_BLOCK_SIZE;
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
+                          sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) {
+            cpy_f32_f16<cpy_1_f32_bf16>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
+                                           nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_bf16_f32_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                  const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                  const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                  const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = (ne + SYCL_CPY_BLOCK_SIZE - 1) / SYCL_CPY_BLOCK_SIZE;
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
+                          sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) {
+            cpy_f32_f16<cpy_1_bf16_f32>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
+                                           nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_bf16_bf16_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                  const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                  const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                  const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = (ne + SYCL_CPY_BLOCK_SIZE - 1) / SYCL_CPY_BLOCK_SIZE;
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
+                          sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) {
+            cpy_f32_f16<cpy_1_bf16_bf16>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
+                                            nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_f16_bf16_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                  const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                  const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                  const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = (ne + SYCL_CPY_BLOCK_SIZE - 1) / SYCL_CPY_BLOCK_SIZE;
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
+                          sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) {
+            cpy_f32_f16<cpy_1_f16_bf16>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
+                                           nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+
+static void ggml_cpy_bf16_f16_sycl(const char * cx, char * cdst, const int ne, const int ne00, const int ne01,
+                                  const int ne02, const int nb00, const int nb01, const int nb02, const int nb03,
+                                  const int ne10, const int ne11, const int ne12, const int nb10, const int nb11,
+                                  const int nb12, const int nb13, queue_ptr stream) {
+    const int num_blocks = (ne + SYCL_CPY_BLOCK_SIZE - 1) / SYCL_CPY_BLOCK_SIZE;
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE),
+                          sycl::range<3>(1, 1, SYCL_CPY_BLOCK_SIZE)),
+        [=](sycl::nd_item<3> item_ct1) {
+            cpy_f32_f16<cpy_1_bf16_f16>(cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
+                                           nb10, nb11, nb12, nb13, item_ct1);
+        });
+}
+#endif
+
 void ggml_sycl_cpy(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1) try {
     // Unlike other operators ggml_sycl_cpy takes 2 distinct tensors instead of a dst ggml_tensor and rely on its src field
+    GGML_SYCL_DEBUG("ggml_sycl_cpy: src0->type=%s, src1->type=%s\n",
+        ggml_type_name(src0->type), ggml_type_name(src1->type));
     scope_op_debug_print scope_dbg_print(__func__, src1, /*num_src=*/0, debug_get_tensor_str("\tsrc0", src0));
     const int64_t ne = ggml_nelements(src0);
     GGML_ASSERT(ne == ggml_nelements(src1));
@@ -525,12 +1146,31 @@ void ggml_sycl_cpy(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, co
     if ((src0->type == src1->type) && (ggml_is_contiguous(src0) && ggml_is_contiguous(src1))) {
         GGML_SYCL_DEBUG("%s: memcpy path\n", __func__);
         main_stream->memcpy(src1_ddc, src0_ddc, ggml_nbytes(src0));
+    } else if (src0->type == GGML_TYPE_F32 && ggml_sycl_is_quantized_type(src1->type)) {
+        GGML_ASSERT(ggml_sycl_can_quantize_rows_sycl(src1->type));
+        ggml_sycl_quantize_rows_sycl<float>(src0_ddc, src1_ddc, src0, src1, ne, ne00, ne01, ne02, nb00, nb01,
+                                            nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_F16 && ggml_sycl_is_quantized_type(src1->type)) {
+        GGML_ASSERT(ggml_sycl_can_quantize_rows_sycl(src1->type));
+        ggml_sycl_quantize_rows_sycl<sycl::half>(src0_ddc, src1_ddc, src0, src1, ne, ne00, ne01, ne02, nb00,
+                                                 nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13,
+                                                 main_stream);
+#ifdef GGML_SYCL_HAS_BF16
+    } else if (src0->type == GGML_TYPE_BF16 && ggml_sycl_is_quantized_type(src1->type)) {
+        GGML_ASSERT(ggml_sycl_can_quantize_rows_sycl(src1->type));
+        ggml_sycl_quantize_rows_sycl<ggml_bf16_t>(src0_ddc, src1_ddc, src0, src1, ne, ne00, ne01, ne02,
+                                                  nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11,
+                                                  nb12, nb13, main_stream);
+#endif
     } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32) {
         ggml_cpy_f32_f32_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
                               nb11, nb12, nb13, main_stream);
     } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16) {
         ggml_cpy_f32_f16_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
                               nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_I32) {
+        ggml_cpy_f32_i32_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
+                              nb11, nb12, nb13, main_stream);
     } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_Q8_0) {
         ggml_cpy_f32_q8_0_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
                                nb11, nb12, nb13, main_stream);
@@ -546,12 +1186,24 @@ void ggml_sycl_cpy(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, co
     } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16) {
         ggml_cpy_f16_f16_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
                               nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_Q4_0) {
+        ggml_cpy_f16_q4_0_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02,
+                               nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_Q4_1) {
+        ggml_cpy_f16_q4_1_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02,
+                               nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_Q5_0) {
+        ggml_cpy_f16_q5_0_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02,
+                               nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
     } else if (src0->type == GGML_TYPE_I16 && src1->type == GGML_TYPE_I16) {
         ggml_cpy_i16_i16_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
                               nb11, nb12, nb13, main_stream);
     } else if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_I32) {
         ggml_cpy_i32_i32_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
                               nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_F32) {
+        ggml_cpy_i32_f32_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
+                              nb11, nb12, nb13, main_stream);
     } else if (src0->type == GGML_TYPE_Q4_0 && src1->type == GGML_TYPE_F32) {
         ggml_cpy_q4_0_f32_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
                                nb11, nb12, nb13, main_stream);
@@ -573,6 +1225,9 @@ void ggml_sycl_cpy(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, co
     } else if (src0->type == GGML_TYPE_Q5_1 && src1->type == GGML_TYPE_F32) {
         ggml_cpy_q5_1_f32_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
                                nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_MXFP4 && src1->type == GGML_TYPE_F32) {
+        ggml_cpy_mxfp4_f32_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
+                                nb10, nb11, nb12, nb13, main_stream);
     } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_IQ4_NL) {
         ggml_cpy_f32_iq4_nl_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12,
                                  nb10, nb11, nb12, nb13, main_stream);
@@ -586,6 +1241,57 @@ void ggml_sycl_cpy(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, co
         ggml_cpy_q4_0_q4_0(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
     } else if (src0->type == GGML_TYPE_Q4_1 && src1->type == GGML_TYPE_Q4_1) {
         ggml_cpy_q4_1_q4_1(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_Q1_0 && src1->type == GGML_TYPE_Q1_0) {
+        ggml_cpy_q1_0_q1_0(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_MXFP4 && src1->type == GGML_TYPE_MXFP4) {
+        ggml_cpy_mxfp4_mxfp4(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_NVFP4 && src1->type == GGML_TYPE_NVFP4) {
+        ggml_cpy_nvfp4_nvfp4(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_Q2_K && src1->type == GGML_TYPE_Q2_K) {
+        ggml_cpy_q2_K_q2_K(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_Q3_K && src1->type == GGML_TYPE_Q3_K) {
+        ggml_cpy_q3_K_q3_K(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_Q4_K && src1->type == GGML_TYPE_Q4_K) {
+        ggml_cpy_q4_K_q4_K(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_Q5_K && src1->type == GGML_TYPE_Q5_K) {
+        ggml_cpy_q5_K_q5_K(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_Q6_K && src1->type == GGML_TYPE_Q6_K) {
+        ggml_cpy_q6_K_q6_K(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_IQ2_XXS && src1->type == GGML_TYPE_IQ2_XXS) {
+        ggml_cpy_iq2_xxs_iq2_xxs(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_IQ2_XS && src1->type == GGML_TYPE_IQ2_XS) {
+        ggml_cpy_iq2_xs_iq2_xs(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_IQ2_S && src1->type == GGML_TYPE_IQ2_S) {
+        ggml_cpy_iq2_s_iq2_s(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_IQ3_XXS && src1->type == GGML_TYPE_IQ3_XXS) {
+        ggml_cpy_iq3_xxs_iq3_xxs(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_IQ1_S && src1->type == GGML_TYPE_IQ1_S) {
+        ggml_cpy_iq1_s_iq1_s(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_IQ1_M && src1->type == GGML_TYPE_IQ1_M) {
+        ggml_cpy_iq1_m_iq1_m(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_IQ4_NL && src1->type == GGML_TYPE_IQ4_NL) {
+        ggml_cpy_iq4_nl_iq4_nl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_IQ3_S && src1->type == GGML_TYPE_IQ3_S) {
+        ggml_cpy_iq3_s_iq3_s(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_IQ4_XS && src1->type == GGML_TYPE_IQ4_XS) {
+        ggml_cpy_iq4_xs_iq4_xs(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
+#ifdef GGML_SYCL_HAS_BF16
+    } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_BF16) {
+        ggml_cpy_f32_bf16_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
+                              nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32) {
+        ggml_cpy_bf16_f32_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
+                              nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_BF16) {
+        ggml_cpy_bf16_bf16_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
+                               nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_BF16) {
+        ggml_cpy_f16_bf16_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
+                              nb11, nb12, nb13, main_stream);
+    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F16) {
+        ggml_cpy_bf16_f16_sycl(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10,
+                              nb11, nb12, nb13, main_stream);
+#endif
     } else {
         GGML_LOG_ERROR("%s: unsupported type combination (%s to %s)\n", __func__, ggml_type_name(src0->type),
                        ggml_type_name(src1->type));

ggml/src/ggml-sycl/cpy.hpp

@@ -317,7 +317,7 @@ inline void cpy_blck_f32_nvfp4(const char * cxi, char * cdsti) {
 
         const uint8_t ue = ggml_fp32_to_ue4m3(amax / 6.0f);
         dsti->d[s] = ue;
-        const float d = ggml_ue4m3_to_fp32(ue);
+        const float d = ggml_sycl_ue4m3_to_fp32(ue);
 
         for (int j = 0; j < QK_NVFP4_SUB / 2; ++j) {
             const uint8_t x0 = best_index_mxfp4(xb[0 + j], d);

ggml/src/ggml-sycl/cross_entropy_loss.cpp

@@ -0,0 +1,255 @@
+#include "cross_entropy_loss.hpp"
+
+#include <cstdint>
+#include <cmath>
+
+template <bool has_shared>
+static __dpct_inline__ void cross_entropy_loss_f32_kernel(
+        const float * __restrict__ logits,
+        const float * __restrict__ labels,
+        float * __restrict__ row_loss,
+        const int nclasses,
+        const int nrows,
+        float * __restrict__ smem,
+        const sycl::nd_item<3> & item) {
+
+    const int row = item.get_group(2);
+    const int tid = item.get_local_id(2);
+
+    logits += (int64_t) row * nclasses;
+    labels += (int64_t) row * nclasses;
+
+    float max_logit = -INFINITY;
+    for (int i = tid; i < nclasses; i += WARP_SIZE) {
+        const float v = logits[i];
+        max_logit = sycl::fmax(max_logit, v);
+        if (has_shared) {
+            smem[i] = v;
+        }
+    }
+    max_logit = warp_reduce_max<WARP_SIZE>(max_logit);
+
+    float sum_exp = 0.0f;
+    for (int i = tid; i < nclasses; i += WARP_SIZE) {
+        const float v = has_shared ? smem[i] : logits[i];
+        sum_exp += sycl::exp(v - max_logit);
+    }
+    sum_exp = warp_reduce_sum<WARP_SIZE>(sum_exp);
+    const float log_sum = sycl::log(sum_exp);
+
+    float loss = 0.0f;
+    for (int i = tid; i < nclasses; i += WARP_SIZE) {
+        const float v = has_shared ? smem[i] : logits[i];
+        loss += (v - max_logit - log_sum) * labels[i];
+    }
+    loss = -warp_reduce_sum<WARP_SIZE>(loss) / (float) nrows;
+
+    if (tid == 0) {
+        row_loss[row] = loss;
+    }
+}
+
+template <bool has_shared>
+static __dpct_inline__ void cross_entropy_loss_back_f32_kernel(
+        const float * __restrict__ grad,
+        const float * __restrict__ logits,
+        const float * __restrict__ labels,
+        float * __restrict__ dst,
+        const int nclasses,
+        const int nrows,
+        float * __restrict__ smem,
+        const sycl::nd_item<3> & item) {
+
+    const int row = item.get_group(2);
+    const int tid = item.get_local_id(2);
+
+    logits += (int64_t) row * nclasses;
+    labels += (int64_t) row * nclasses;
+    dst    += (int64_t) row * nclasses;
+
+    float max_logit = -INFINITY;
+    for (int i = tid; i < nclasses; i += WARP_SIZE) {
+        const float v = logits[i];
+        max_logit = sycl::fmax(max_logit, v);
+        if (has_shared) {
+            smem[i] = v;
+        }
+    }
+    max_logit = warp_reduce_max<WARP_SIZE>(max_logit);
+
+    float sum_exp = 0.0f;
+    for (int i = tid; i < nclasses; i += WARP_SIZE) {
+        const float v = sycl::exp((has_shared ? smem[i] : logits[i]) - max_logit);
+        sum_exp += v;
+        if (has_shared) {
+            smem[i] = v;
+        } else {
+            dst[i] = v;
+        }
+    }
+    sum_exp = warp_reduce_sum<WARP_SIZE>(sum_exp);
+    const float inv_sum = 1.0f / sum_exp;
+
+    const float d_by_nrows = grad[0] / (float) nrows;
+    for (int i = tid; i < nclasses; i += WARP_SIZE) {
+        const float sm_num = has_shared ? smem[i] : dst[i];
+        dst[i] = (sm_num * inv_sum - labels[i]) * d_by_nrows;
+    }
+}
+
+static void cross_entropy_reduce_rows(
+        ggml_backend_sycl_context & ctx,
+        const float * row_loss,
+        float * dst,
+        const int64_t nrows) {
+    if (nrows == 1) {
+        SYCL_CHECK(CHECK_TRY_ERROR(
+            ctx.stream()->memcpy(dst, row_loss, sizeof(float))));
+        return;
+    }
+
+    ggml_sycl_pool_alloc<float> tmp_alloc(ctx.pool(), nrows);
+    float * tmp = tmp_alloc.get();
+    SYCL_CHECK(CHECK_TRY_ERROR(
+        ctx.stream()->memcpy(tmp, row_loss, nrows * sizeof(float))));
+
+    int64_t cur = nrows;
+    while (cur > 1) {
+        const int64_t out = (cur + WARP_SIZE - 1) / WARP_SIZE;
+        const sycl::range<3> block(1, 1, WARP_SIZE);
+        const sycl::range<3> grid(1, 1, out);
+        ctx.stream()->parallel_for(
+            sycl::nd_range<3>(grid * block, block),
+            [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                const int row = item.get_group(2);
+                const int tid = item.get_local_id(2);
+                const int64_t i = (int64_t) row * WARP_SIZE + tid;
+                float v = i < cur ? tmp[i] : 0.0f;
+                v = warp_reduce_sum<WARP_SIZE>(v);
+                if (tid == 0) {
+                    tmp[row] = v;
+                }
+            });
+        cur = out;
+    }
+
+    SYCL_CHECK(CHECK_TRY_ERROR(
+        ctx.stream()->memcpy(dst, tmp, sizeof(float))));
+}
+
+void ggml_sycl_cross_entropy_loss(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
+
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_contiguous(src0));
+    GGML_ASSERT(ggml_is_contiguous(src1));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+    GGML_ASSERT(ggml_are_same_shape(src0, src1));
+    GGML_ASSERT(ggml_is_scalar(dst));
+
+    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
+
+    const int64_t nclasses = src0->ne[0];
+    const int64_t nrows = ggml_nrows(src0);
+
+    const float * logits_d = (const float *) src0->data;
+    const float * labels_d = (const float *) src1->data;
+    float * dst_d = (float *) dst->data;
+
+    ggml_sycl_pool_alloc<float> row_loss_alloc(ctx.pool(), nrows);
+    float * row_loss = row_loss_alloc.get();
+
+    const sycl::range<3> block(1, 1, WARP_SIZE);
+    const sycl::range<3> grid(1, 1, nrows);
+    const size_t nbytes_shared = (size_t) nclasses * sizeof(float);
+    const size_t smpbo = ggml_sycl_info().devices[ctx.device].smpbo;
+
+    if (nbytes_shared <= smpbo) {
+        ctx.stream()->submit([&](sycl::handler & cgh) {
+            sycl::local_accessor<float, 1> smem(sycl::range<1>(nclasses), cgh);
+            cgh.parallel_for(
+                sycl::nd_range<3>(grid * block, block),
+                [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                    cross_entropy_loss_f32_kernel<true>(
+                        logits_d, labels_d, row_loss,
+                        (int) nclasses, (int) nrows,
+                        get_pointer(smem), item);
+                });
+        });
+    } else {
+        ctx.stream()->parallel_for(
+            sycl::nd_range<3>(grid * block, block),
+            [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                cross_entropy_loss_f32_kernel<false>(
+                    logits_d, labels_d, row_loss,
+                    (int) nclasses, (int) nrows,
+                    nullptr, item);
+            });
+    }
+
+    cross_entropy_reduce_rows(ctx, row_loss, dst_d, nrows);
+}
+
+void ggml_sycl_cross_entropy_loss_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/3);
+
+    const ggml_tensor * grad  = dst->src[0];
+    const ggml_tensor * src0f = dst->src[1];
+    const ggml_tensor * src1f = dst->src[2];
+
+    GGML_ASSERT(grad->type == GGML_TYPE_F32);
+    GGML_ASSERT(src0f->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1f->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+
+    GGML_ASSERT(ggml_is_scalar(grad));
+    GGML_ASSERT(ggml_is_contiguous(grad));
+    GGML_ASSERT(ggml_is_contiguous(src0f));
+    GGML_ASSERT(ggml_is_contiguous(src1f));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+    GGML_ASSERT(ggml_are_same_shape(src0f, src1f));
+    GGML_ASSERT(ggml_are_same_shape(src0f, dst));
+
+    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
+
+    const int64_t nclasses = src0f->ne[0];
+    const int64_t nrows = ggml_nrows(src0f);
+
+    const float * grad_d  = (const float *) grad->data;
+    const float * logits_d = (const float *) src0f->data;
+    const float * labels_d = (const float *) src1f->data;
+    float * dst_d = (float *) dst->data;
+
+    const sycl::range<3> block(1, 1, WARP_SIZE);
+    const sycl::range<3> grid(1, 1, nrows);
+    const size_t nbytes_shared = (size_t) nclasses * sizeof(float);
+    const size_t smpbo = ggml_sycl_info().devices[ctx.device].smpbo;
+
+    if (nbytes_shared <= smpbo) {
+        ctx.stream()->submit([&](sycl::handler & cgh) {
+            sycl::local_accessor<float, 1> smem(sycl::range<1>(nclasses), cgh);
+            cgh.parallel_for(
+                sycl::nd_range<3>(grid * block, block),
+                [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                    cross_entropy_loss_back_f32_kernel<true>(
+                        grad_d, logits_d, labels_d, dst_d,
+                        (int) nclasses, (int) nrows,
+                        get_pointer(smem), item);
+                });
+        });
+    } else {
+        ctx.stream()->parallel_for(
+            sycl::nd_range<3>(grid * block, block),
+            [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                cross_entropy_loss_back_f32_kernel<false>(
+                    grad_d, logits_d, labels_d, dst_d,
+                    (int) nclasses, (int) nrows,
+                    nullptr, item);
+            });
+    }
+}

ggml/src/ggml-sycl/cross_entropy_loss.hpp

@@ -0,0 +1,7 @@
+#pragma once
+
+#include "common.hpp"
+
+void ggml_sycl_cross_entropy_loss(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
+
+void ggml_sycl_cross_entropy_loss_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst);

ggml/src/ggml-sycl/dmmv.cpp

@@ -680,14 +680,14 @@ static void dequantize_mul_mat_vec_q4_k(const void *__restrict__ vx,
             q16[2] = q2[0] & 0x0f0f;
             q16[3] = q2[0] & 0xf0f0;
 
-            float4 s = {0.f, 0.f, 0.f, 0.f};
+            sycl::float4 s = {0.f, 0.f, 0.f, 0.f};
             float smin = 0;
             for (int l = 0; l < 2; ++l) {
-                s.x += y1[l] * q4[l+0]; s.y += y1[l+32] * q4[l+2];
-                s.z += y2[l] * q4[l+4]; s.w += y2[l+32] * q4[l+6];
+                s.x() += y1[l] * q4[l+0]; s.y() += y1[l+32] * q4[l+2];
+                s.z() += y2[l] * q4[l+4]; s.w() += y2[l+32] * q4[l+6];
                 smin += y1[l] * sc[2] + y1[l+32] * sc[3] + y2[l] * sc[6] + y2[l+32] * sc[7];
             }
-            tmp += dall * (s.x * sc[0] + s.y * sc[1] * 1.f/16.f + s.z * sc[4] + s.w * sc[5] * 1.f/16.f) - dmin * smin;
+            tmp += dall * (s.x() * sc[0] + s.y() * sc[1] * 1.f/16.f + s.z() * sc[4] + s.w() * sc[5] * 1.f/16.f) - dmin * smin;
 #endif
         }
 
@@ -835,14 +835,14 @@ static void dequantize_mul_mat_vec_q4_k_reorder(const void *__restrict__ vx,
             q16[2] = q2[0] & 0x0f0f;
             q16[3] = q2[0] & 0xf0f0;
 
-            float4 s = {0.f, 0.f, 0.f, 0.f};
+            sycl::float4 s = {0.f, 0.f, 0.f, 0.f};
             float smin = 0;
             for (int l = 0; l < 2; ++l) {
-                s.x += y1[l] * q4[l+0]; s.y += y1[l+32] * q4[l+2];
-                s.z += y2[l] * q4[l+4]; s.w += y2[l+32] * q4[l+6];
+                s.x() += y1[l] * q4[l+0]; s.y() += y1[l+32] * q4[l+2];
+                s.z() += y2[l] * q4[l+4]; s.w() += y2[l+32] * q4[l+6];
                 smin += y1[l] * sc[2] + y1[l+32] * sc[3] + y2[l] * sc[6] + y2[l+32] * sc[7];
             }
-            tmp += dall * (s.x * sc[0] + s.y * sc[1] * 1.f/16.f + s.z * sc[4] + s.w * sc[5] * 1.f/16.f) - dmin * smin;
+            tmp += dall * (s.x() * sc[0] + s.y() * sc[1] * 1.f/16.f + s.z() * sc[4] + s.w() * sc[5] * 1.f/16.f) - dmin * smin;
 #endif
         }
 
@@ -1126,7 +1126,7 @@ static void dequantize_mul_mat_vec_q5_k_reorder(const void *__restrict__ vx,
 
     // sum up partial sums and write back result
 #pragma unroll
-    for (int mask = QK_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+    for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) {
         tmp +=
             dpct::permute_sub_group_by_xor(item_ct1.get_sub_group(), tmp, mask);
     }
@@ -1762,10 +1762,13 @@ static void dequantize_mul_mat_vec_q5_K_sycl_reorder(const void *vx, const float
                                                      const int nrows,
                                                      dpct::queue_ptr stream) {
     GGML_ASSERT(ncols % QK_K == 0);
-    const sycl::range<3> block_dims(1, 1, QK_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_item<3> item_ct1) [[sycl::reqd_sub_group_size(QK_WARP_SIZE)]] {
+        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_reorder(vx, y, dst, ncols, nrows, item_ct1);
         });
 }

ggml/src/ggml-sycl/element_wise.cpp

@@ -9,9 +9,12 @@
 #define SYCL_LOCAL_ID_CALC(ITEM, IDX) \
     (ITEM.get_local_range(IDX) * ITEM.get_group(IDX) + ITEM.get_local_id(IDX))
 
-static void acc_f32(const float * x, const float * y, float * dst, const int64_t ne,
-        const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t ne13,
-        const int64_t s11, const int64_t s12, const int64_t s13, const int64_t offset) {
+static void acc_f32(const char * x, const char * y, float * dst, const int64_t ne,
+    const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3,
+    const int64_t nb00, const int64_t nb01, const int64_t nb02, const int64_t nb03,
+    const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t ne13,
+    const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
+    const int64_t s11, const int64_t s12, const int64_t s13, const int64_t offset) {
     auto item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
     const int64_t i = SYCL_LOCAL_ID_CALC(item_ct1, 2);
 
@@ -30,9 +33,18 @@ static void acc_f32(const float * x, const float * y, float * dst, const int64_t
     tmp -= i11 * s11;
     const int64_t i10 = tmp;
 
-    float val = x[i];
+    int64_t tmp_dst = i;
+    const int64_t i3 = tmp_dst / (ne2*ne1*ne0);
+    tmp_dst -= i3 * (ne2*ne1*ne0);
+    const int64_t i2 = tmp_dst / (ne1*ne0);
+    tmp_dst -= i2 * (ne1*ne0);
+    const int64_t i1 = tmp_dst / ne0;
+    tmp_dst -= i1 * ne0;
+    const int64_t i0 = tmp_dst;
+
+    float val = *(const float *) (x + i0*nb00 + i1*nb01 + i2*nb02 + i3*nb03);
     if (src1_idx >= 0 && i10 < ne10 && i11 < ne11 && i12 < ne12 && i13 < ne13) {
-        val += y[((i13*ne12 + i12) * ne11 + i11) * ne10 + i10];
+        val += *(const float *) (y + i10*nb10 + i11*nb11 + i12*nb12 + i13*nb13);
     }
     dst[i] = val;
 }
@@ -422,15 +434,24 @@ static void gated_op_fused_geglu_quick(const T * x, const T * g, T * dst, const
 }
 
 namespace ggml_sycl_detail {
-static void acc_f32_sycl(const float *x, const float *y, float *dst,
-                         const int64_t n_elements, const int64_t ne10, const int64_t ne11,
-                         const int64_t ne12, const int64_t ne13, const int64_t s1, const int64_t s2, const int64_t s3,
+static void acc_f32_sycl(const char *x, const char *y, float *dst,
+                         const int64_t n_elements,
+                         const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3,
+                         const int64_t nb00, const int64_t nb01, const int64_t nb02, const int64_t nb03,
+                         const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t ne13,
+                         const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
+                         const int64_t s1, const int64_t s2, const int64_t s3,
                          const int64_t offset, queue_ptr stream) {
     const int num_blocks = (n_elements + SYCL_ACC_BLOCK_SIZE - 1) / SYCL_ACC_BLOCK_SIZE;
     stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, num_blocks) * sycl::range<3>(1, 1, SYCL_ACC_BLOCK_SIZE),
                                            sycl::range<3>(1, 1, SYCL_ACC_BLOCK_SIZE)),
                          [=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                             acc_f32(x, y, dst, n_elements, ne10, ne11, ne12, ne13, s1, s2, s3, offset);
+                            acc_f32(x, y, dst, n_elements,
+                                ne0, ne1, ne2, ne3,
+                                nb00, nb01, nb02, nb03,
+                                ne10, ne11, ne12, ne13,
+                                nb10, nb11, nb12, nb13,
+                                s1, s2, s3, offset);
                          });
 }
 
@@ -843,8 +864,8 @@ static inline void ggml_sycl_op_acc(ggml_backend_sycl_context & ctx, ggml_tensor
     const ggml_tensor * src0 = dst->src[0];
     const ggml_tensor * src1 = dst->src[1];
 
-    const float * src0_d = (const float *) src0->data;
-    const float * src1_d = (const float *) src1->data;
+    const char  * src0_d = (const char  *) src0->data;
+    const char  * src1_d = (const char  *) src1->data;
     float       * dst_d  = (float       *)  dst->data;
 
     dpct::queue_ptr stream = ctx.stream();
@@ -853,17 +874,20 @@ static inline void ggml_sycl_op_acc(ggml_backend_sycl_context & ctx, ggml_tensor
     GGML_ASSERT(src1->type == GGML_TYPE_F32);
     GGML_ASSERT( dst->type == GGML_TYPE_F32);
 
-    GGML_ASSERT(ggml_is_contiguous(src1));
     GGML_ASSERT(dst->nb[0] == ggml_element_size(dst));
     GGML_ASSERT(ggml_is_contiguously_allocated(dst));
+    GGML_ASSERT(ggml_are_same_shape(src0, dst));
 
-    const int64_t s1     = dst->op_params[0] / sizeof(float);
-    const int64_t s2     = dst->op_params[1] / sizeof(float);
-    const int64_t s3     = dst->op_params[2] / sizeof(float);
-    const int64_t offset = dst->op_params[3] / sizeof(float);
+    const int64_t s1     = (int64_t) ((const int32_t *) dst->op_params)[0] / (int64_t) sizeof(float);
+    const int64_t s2     = (int64_t) ((const int32_t *) dst->op_params)[1] / (int64_t) sizeof(float);
+    const int64_t s3     = (int64_t) ((const int32_t *) dst->op_params)[2] / (int64_t) sizeof(float);
+    const int64_t offset = (int64_t) ((const int32_t *) dst->op_params)[3] / (int64_t) sizeof(float);
 
     ggml_sycl_detail::acc_f32_sycl(src0_d, src1_d, dst_d, ggml_nelements(dst),
+        dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
+        src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3],
         src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3],
+        src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3],
         s1, s2, s3, offset, stream);
 }
 

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

@@ -41,7 +41,7 @@
 #if SYCL_EXT_ONEAPI_VIRTUAL_MEM
 #    include <sycl/ext/oneapi/virtual_mem/physical_mem.hpp>
 #    include <sycl/ext/oneapi/virtual_mem/virtual_mem.hpp>
-#    define GGML_SYCL_USE_VMM
+#    define GGML_SYCL_SUPPORT_VMM
 #endif
 #include <sycl/half_type.hpp>
 
@@ -74,15 +74,16 @@
 #include "ggml-sycl/solve_tri.hpp"
 #include "ggml-sycl/gated_delta_net.hpp"
 #include "ggml-sycl/pool.hpp"
+#include "ggml-sycl/cross_entropy_loss.hpp"
 
 #define MEM_SIZE_2M	0x00200000
 #define MEM_SIZE_1G	0x40000000
 
 static bool g_sycl_loaded = false;
 int g_ggml_sycl_debug = 0;
-int g_ggml_sycl_disable_optimize = 0;
-int g_ggml_sycl_disable_graph = 0;
-int g_ggml_sycl_disable_dnn = 0;
+int g_ggml_sycl_enable_optimize = 1;
+int g_ggml_sycl_enable_graph = 0;
+int g_ggml_sycl_enable_dnn = 1;
 int g_ggml_sycl_enable_vmm = 1;
 int g_ggml_sycl_prioritize_dmmv = 0;
 int g_ggml_sycl_use_async_mem_op = 0;
@@ -117,7 +118,7 @@ static ggml_sycl_device_info ggml_sycl_init() {
         SYCL_CHECK(CHECK_TRY_ERROR(dpct::get_device_info(
             prop, device)));
 
-#if !defined(GGML_SYCL_USE_VMM)
+#if !defined(GGML_SYCL_SUPPORT_VMM)
         info.devices[i].vmm = 0;
 #else
         info.devices[i].vmm = device.has(sycl::aspect::ext_oneapi_virtual_mem);
@@ -265,14 +266,24 @@ void ggml_backend_sycl_print_sycl_devices() {
     print_device_opt_feature(device_count);
 }
 
+static const char* dev2dev_int2str(int dev2dev) {
+    if (dev2dev == DEV2DEV_MEMCPY_SYCL) {
+        return "SYCL API";
+    } else if (dev2dev == DEV2DEV_MEMCPY_L0) {
+        return "Level Zero API";
+    } else {
+        return "Unknown";
+    }
+}
+
 static void ggml_check_sycl() try {
     static bool initialized = false;
 
     if (!initialized) {
         g_ggml_sycl_debug = ggml_sycl_get_env("GGML_SYCL_DEBUG", 0);
-        g_ggml_sycl_disable_optimize = ggml_sycl_get_env("GGML_SYCL_DISABLE_OPT", 0);
-        g_ggml_sycl_disable_graph = ggml_sycl_get_env("GGML_SYCL_DISABLE_GRAPH", 1);
-        g_ggml_sycl_disable_dnn = ggml_sycl_get_env("GGML_SYCL_DISABLE_DNN", 0);
+        g_ggml_sycl_enable_optimize = ggml_sycl_get_env("GGML_SYCL_ENABLE_OPT", 1);
+        g_ggml_sycl_enable_graph = ggml_sycl_get_env("GGML_SYCL_ENABLE_GRAPH", 0);
+        g_ggml_sycl_enable_dnn = ggml_sycl_get_env("GGML_SYCL_ENABLE_DNN", 1);
         g_ggml_sycl_enable_vmm = ggml_sycl_get_env("GGML_SYCL_ENABLE_VMM", 1);
         g_ggml_sycl_prioritize_dmmv = ggml_sycl_get_env("GGML_SYCL_PRIORITIZE_DMMV", 0);
 
@@ -292,66 +303,56 @@ static void ggml_check_sycl() try {
         GGML_SYCL_DEBUG("[SYCL] call ggml_check_sycl\n");
 
         GGML_LOG_INFO("Build with Macros:\n");
-#if defined(GGML_SYCL_FORCE_MMQ)
-        GGML_LOG_INFO("  GGML_SYCL_FORCE_MMQ: yes\n");
-#else
-        GGML_LOG_INFO("  GGML_SYCL_FORCE_MMQ: no\n");
-#endif
-#if defined(GGML_SYCL_F16)
-        GGML_LOG_INFO("  GGML_SYCL_F16: yes\n");
-#else
-        GGML_LOG_INFO("  GGML_SYCL_F16: no\n");
-#endif
-#if defined(GGML_SYCL_GRAPH)
-        GGML_LOG_INFO("  GGML_SYCL_GRAPH: yes\n");
-#else
-        GGML_LOG_INFO("  GGML_SYCL_GRAPH: no\n");
-#endif
 #if defined(GGML_SYCL_DNNL)
         GGML_LOG_INFO("  GGML_SYCL_DNNL: yes\n");
 #else
         GGML_LOG_INFO("  GGML_SYCL_DNNL: no\n");
 #endif
+
+#if defined(GGML_SYCL_F16)
+        GGML_LOG_INFO("  GGML_SYCL_F16: yes\n");
+#else
+        GGML_LOG_INFO("  GGML_SYCL_F16: no\n");
+#endif
+
+#if defined(GGML_SYCL_FORCE_MMQ)
+        GGML_LOG_INFO("  GGML_SYCL_FORCE_MMQ: yes\n");
+#else
+        GGML_LOG_INFO("  GGML_SYCL_FORCE_MMQ: no\n");
+#endif
+
+#if defined(GGML_SYCL_GRAPH)
+        GGML_LOG_INFO("  GGML_SYCL_GRAPH: yes\n");
+#else
+        GGML_LOG_INFO("  GGML_SYCL_GRAPH: no\n");
+#endif
+
 #if defined(GGML_SYCL_SUPPORT_LEVEL_ZERO_API)
         GGML_LOG_INFO("  GGML_SYCL_SUPPORT_LEVEL_ZERO_API: yes\n");
 #else
         GGML_LOG_INFO("  GGML_SYCL_SUPPORT_LEVEL_ZERO_API: no\n");
 #endif
-#if defined(GGML_SYCL_USE_VMM)
-        GGML_LOG_INFO("  GGML_SYCL_USE_VMM: yes\n");
+#if defined(GGML_SYCL_SUPPORT_VMM)
+        GGML_LOG_INFO("  GGML_SYCL_SUPPORT_VMM: yes\n");
 #else
-        GGML_LOG_INFO("  GGML_SYCL_USE_VMM: no\n");
+        GGML_LOG_INFO("  GGML_SYCL_SUPPORT_VMM: no\n");
 #endif
 
         GGML_LOG_INFO("Running with Environment Variables:\n");
         GGML_LOG_INFO("  GGML_SYCL_DEBUG: %d\n", g_ggml_sycl_debug);
-        GGML_LOG_INFO("  GGML_SYCL_DISABLE_OPT: %d\n", g_ggml_sycl_disable_optimize);
-#ifdef GGML_SYCL_GRAPH
-        GGML_LOG_INFO("  GGML_SYCL_DISABLE_GRAPH: %d\n", g_ggml_sycl_disable_graph);
-#else
-        GGML_LOG_INFO("  GGML_SYCL_DISABLE_GRAPH: graph disabled by compile flag\n");
-#endif
+
 #ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API
-        GGML_LOG_INFO("  GGML_SYCL_USE_LEVEL_ZERO_API: %d\n", g_ggml_sycl_use_level_zero_api);
-        GGML_LOG_INFO("  GGML_SYCL_DEV2DEV_MEMCPY: %d\n", g_ggml_sycl_dev2dev_memcpy);
+        GGML_LOG_INFO("  GGML_SYCL_DEV2DEV_MEMCPY: %d (%s)\n", g_ggml_sycl_dev2dev_memcpy, dev2dev_int2str(g_ggml_sycl_dev2dev_memcpy));
 #else
-        GGML_LOG_INFO("  GGML_SYCL_USE_LEVEL_ZERO_API: Disable Level Zero API usage by compile flag\n");
-        GGML_LOG_INFO("  GGML_SYCL_DEV2DEV_MEMCPY: %d, enable to SYCL API since missing GGML_SYCL_SUPPORT_LEVEL_ZERO_API\n",
-                      g_ggml_sycl_dev2dev_memcpy);
+        GGML_LOG_INFO("  GGML_SYCL_DEV2DEV_MEMCPY: %d (%s), enable to SYCL API since missing GGML_SYCL_SUPPORT_LEVEL_ZERO_API\n",
+                      g_ggml_sycl_dev2dev_memcpy, dev2dev_int2str(g_ggml_sycl_dev2dev_memcpy));
 #endif
-#if GGML_SYCL_DNNL
-        GGML_LOG_INFO("  GGML_SYCL_DISABLE_DNN: %d\n", g_ggml_sycl_disable_dnn);
+
+#if defined(GGML_SYCL_DNNL)
+        GGML_LOG_INFO("  GGML_SYCL_ENABLE_DNN: %d\n", g_ggml_sycl_enable_dnn);
 #else
-        GGML_LOG_INFO("  GGML_SYCL_DISABLE_DNN: DNN disabled by compile flag\n");
+        GGML_LOG_INFO("  GGML_SYCL_ENABLE_DNN: DNN disabled by compile flag\n");
 #endif
-#if defined(GGML_SYCL_USE_VMM)
-        GGML_LOG_INFO("  GGML_SYCL_ENABLE_VMM: %d\n", g_ggml_sycl_enable_vmm);
-#else
-        GGML_LOG_INFO("  GGML_SYCL_ENABLE_VMM: virtual memory extension is not available\n");
-#endif
-        GGML_LOG_INFO("  GGML_SYCL_PRIORITIZE_DMMV: %d\n", g_ggml_sycl_prioritize_dmmv);
-        g_ggml_sycl_use_async_mem_op_requested = ggml_sycl_get_env("GGML_SYCL_USE_ASYNC_MEM_OP", 1);
-        GGML_LOG_INFO("  GGML_SYCL_USE_ASYNC_MEM_OP: %d\n", g_ggml_sycl_use_async_mem_op_requested);
 
 #ifdef SYCL_FLASH_ATTN
         GGML_LOG_INFO("  GGML_SYCL_ENABLE_FLASH_ATTN: %d\n", g_ggml_sycl_enable_flash_attention);
@@ -360,6 +361,31 @@ static void ggml_check_sycl() try {
             g_ggml_sycl_enable_flash_attention);
 #endif
 
+#ifdef GGML_SYCL_GRAPH
+        GGML_LOG_INFO("  GGML_SYCL_ENABLE_GRAPH: %d\n", g_ggml_sycl_enable_graph);
+#else
+        GGML_LOG_INFO("  GGML_SYCL_ENABLE_GRAPH: graph disabled by compile flag\n");
+#endif
+
+        GGML_LOG_INFO("  GGML_SYCL_ENABLE_OPT: %d\n", g_ggml_sycl_enable_optimize);
+
+#if defined(GGML_SYCL_SUPPORT_VMM)
+        GGML_LOG_INFO("  GGML_SYCL_ENABLE_VMM: %d\n", g_ggml_sycl_enable_vmm);
+#else
+        GGML_LOG_INFO("  GGML_SYCL_ENABLE_VMM: virtual memory extension is not available\n");
+#endif
+
+        GGML_LOG_INFO("  GGML_SYCL_PRIORITIZE_DMMV: %d\n", g_ggml_sycl_prioritize_dmmv);
+
+        g_ggml_sycl_use_async_mem_op_requested = ggml_sycl_get_env("GGML_SYCL_USE_ASYNC_MEM_OP", 1);
+        GGML_LOG_INFO("  GGML_SYCL_USE_ASYNC_MEM_OP: %d\n", g_ggml_sycl_use_async_mem_op_requested);
+
+#ifdef GGML_SYCL_SUPPORT_LEVEL_ZERO_API
+        GGML_LOG_INFO("  GGML_SYCL_USE_LEVEL_ZERO_API: %d\n", g_ggml_sycl_use_level_zero_api);
+#else
+        GGML_LOG_INFO("  GGML_SYCL_USE_LEVEL_ZERO_API: Disable Level Zero API usage by compile flag\n");
+#endif
+
         GGML_LOG_INFO("  GGML_SYCL_USM_SYSTEM: %d\n", g_ggml_sycl_usm_system);
 
 /* NOT REMOVE, keep it for next optimize for XMX.
@@ -373,7 +399,7 @@ static void ggml_check_sycl() try {
         // staging path while preserving queue ordering semantics. Graph support still depends on the extension being
         // available, but it no longer needs to control the non-graph fast path.
 #if defined(GGML_SYCL_GRAPH) && SYCL_EXT_ONEAPI_ASYNC_MEMORY_ALLOC
-        g_ggml_sycl_use_async_mem_op = g_ggml_sycl_use_async_mem_op_requested || !g_ggml_sycl_disable_graph;
+        g_ggml_sycl_use_async_mem_op = g_ggml_sycl_use_async_mem_op_requested || g_ggml_sycl_enable_graph;
         if (g_ggml_sycl_use_async_mem_op) {
             for (unsigned int i = 0; i < dpct::dev_mgr::instance().device_count(); ++i) {
                 if (!dpct::dev_mgr::instance().get_device(i).has(sycl::aspect::ext_oneapi_async_memory_alloc)) {
@@ -516,12 +542,14 @@ ggml_backend_sycl_buffer_init_tensor(ggml_backend_buffer_t buffer,
         return GGML_STATUS_SUCCESS;
     }
 
-    if (!g_ggml_sycl_disable_optimize) {
+    if (g_ggml_sycl_enable_optimize) {
         // set reorder extra buffer based on supported type
         switch (tensor->type) {
             case GGML_TYPE_Q4_0:
             case GGML_TYPE_Q8_0:
+            case GGML_TYPE_Q3_K:
             case GGML_TYPE_Q4_K:
+            case GGML_TYPE_Q5_K:
             case GGML_TYPE_Q6_K:{
                 ggml_tensor_extra_gpu * extra = new ggml_tensor_extra_gpu{};
                 tensor->extra                 = extra;
@@ -1562,7 +1590,7 @@ struct ggml_sycl_pool_leg : public ggml_sycl_pool {
 };
 
 // pool with virtual memory management
-#if defined(GGML_SYCL_USE_VMM)
+#if defined(GGML_SYCL_SUPPORT_VMM)
 struct ggml_sycl_pool_vmm : public ggml_sycl_pool {
     static const size_t SYCL_POOL_VMM_MAX_SIZE = 1ull << 35; // 32 GB
 
@@ -1674,7 +1702,7 @@ struct ggml_sycl_pool_vmm : public ggml_sycl_pool {
         GGML_ASSERT(ptr == reinterpret_cast<void *>(pool_addr + pool_used));
     }
 };
-#endif // defined(GGML_SYCL_USE_VMM)
+#endif // defined(GGML_SYCL_SUPPORT_VMM)
 
 struct ggml_sycl_pool_host : public ggml_sycl_pool {
     queue_ptr qptr;
@@ -1756,11 +1784,11 @@ std::unique_ptr<ggml_sycl_pool> ggml_backend_sycl_context::new_pool_for_host(que
 }
 
 std::unique_ptr<ggml_sycl_pool> ggml_backend_sycl_context::new_pool_for_device(queue_ptr qptr, int device) {
-#if defined(GGML_SYCL_USE_VMM)
+#if defined(GGML_SYCL_SUPPORT_VMM)
     if (g_ggml_sycl_enable_vmm && ggml_sycl_info().devices[device].vmm) {
         return std::unique_ptr<ggml_sycl_pool>(new ggml_sycl_pool_vmm(qptr, device));
     }
-#endif // defined(GGML_SYCL_USE_VMM)
+#endif // defined(GGML_SYCL_SUPPORT_VMM)
     return std::unique_ptr<ggml_sycl_pool>(new ggml_sycl_pool_leg(qptr, device));
 }
 
@@ -2088,11 +2116,148 @@ static int next_power_of_2(int x) {
     return n;
 }
 
+static void init_argsort_indices_padded(
+        int * idx,
+        const int nrows,
+        const int ncols_pad,
+        const sycl::nd_item<1> & item_ct1) {
+    const size_t gid = item_ct1.get_local_range(0) * item_ct1.get_group(0) + item_ct1.get_local_id(0);
+    const size_t total = (size_t) nrows * (size_t) ncols_pad;
+
+    if (gid >= total) {
+        return;
+    }
+
+    idx[gid] = (int) (gid % (size_t) ncols_pad);
+}
+
+template <ggml_sort_order order>
+static void argsort_f32_i32_global_pass(const float *            x,
+                                        int *                    idx,
+                                        const int                ncols,
+                                        const int                nrows,
+                                        const int                ncols_pad,
+                                        const int                j,
+                                        const int                k,
+                                        const sycl::nd_item<1> & item_ct1) {
+    const size_t gid   = item_ct1.get_local_range(0) * item_ct1.get_group(0) + item_ct1.get_local_id(0);
+    const size_t total = (size_t) nrows * (size_t) ncols_pad;
+
+    if (gid >= total) {
+        return;
+    }
+
+    const int row = (int) (gid / (size_t) ncols_pad);
+    const int col = (int) (gid % (size_t) ncols_pad);
+    const int ixj = col ^ j;
+
+    if (ixj <= col || ixj >= ncols_pad) {
+        return;
+    }
+
+    const size_t base  = (size_t) row * (size_t) ncols_pad;
+    const size_t pos_a = base + (size_t) col;
+    const size_t pos_b = base + (size_t) ixj;
+
+    const int a = idx[pos_a];
+    const int b = idx[pos_b];
+
+    bool do_swap = false;
+
+    if ((col & k) == 0) {
+        if (a >= ncols ||
+            (b < ncols &&
+             (order == GGML_SORT_ORDER_ASC ?
+                  x[(size_t) row * (size_t) ncols + (size_t) a] > x[(size_t) row * (size_t) ncols + (size_t) b] :
+                  x[(size_t) row * (size_t) ncols + (size_t) a] < x[(size_t) row * (size_t) ncols + (size_t) b]))) {
+            do_swap = true;
+        }
+    } else {
+        if (b >= ncols ||
+            (a < ncols &&
+             (order == GGML_SORT_ORDER_ASC ?
+                  x[(size_t) row * (size_t) ncols + (size_t) a] < x[(size_t) row * (size_t) ncols + (size_t) b] :
+                  x[(size_t) row * (size_t) ncols + (size_t) a] > x[(size_t) row * (size_t) ncols + (size_t) b]))) {
+            do_swap = true;
+        }
+    }
+
+    if (do_swap) {
+        idx[pos_a] = b;
+        idx[pos_b] = a;
+    }
+}
+
+static void copy_argsort_indices_unpadded(const int *              idx_padded,
+                                          int *                    dst,
+                                          const int                nrows,
+                                          const int                ncols,
+                                          const int                ncols_pad,
+                                          const sycl::nd_item<1> & item_ct1) {
+    const size_t gid   = item_ct1.get_local_range(0) * item_ct1.get_group(0) + item_ct1.get_local_id(0);
+    const size_t total = (size_t) nrows * (size_t) ncols;
+
+    if (gid >= total) {
+        return;
+    }
+
+    const int row = (int) (gid / (size_t) ncols);
+    const int col = (int) (gid % (size_t) ncols);
+
+    dst[(size_t) row * (size_t) ncols + (size_t) col] = idx_padded[(size_t) row * (size_t) ncols_pad + (size_t) col];
+}
+
 static void argsort_f32_i32_sycl(const float *x, int *dst, const int ncols,
                                  const int nrows, ggml_sort_order order,
-                                 queue_ptr stream, int device) {
+                                 queue_ptr stream, int device, ggml_sycl_pool & pool) {
     // bitonic sort requires ncols to be power of 2
     const int ncols_pad = next_power_of_2(ncols);
+    const size_t shared_mem = (size_t) ncols_pad * sizeof(int);
+    const size_t smpbo = ggml_sycl_info().devices[device].smpbo;
+
+    if (shared_mem > smpbo) {
+        ggml_sycl_pool_alloc<int> idx_padded_alloc(pool, (size_t) nrows * (size_t) ncols_pad);
+        int *                     idx_padded = idx_padded_alloc.get();
+
+        constexpr size_t block_size     = 256;
+        const size_t     total_padded   = (size_t) nrows * (size_t) ncols_pad;
+        const size_t     nblocks_padded = (total_padded + block_size - 1) / block_size;
+
+        stream->parallel_for(
+            sycl::nd_range<1>(sycl::range<1>(nblocks_padded * block_size), sycl::range<1>(block_size)),
+            [=](sycl::nd_item<1> item_ct1) { init_argsort_indices_padded(idx_padded, nrows, ncols_pad, item_ct1); });
+
+        for (int k = 2; k <= ncols_pad; k *= 2) {
+            for (int j = k / 2; j > 0; j /= 2) {
+                if (order == GGML_SORT_ORDER_ASC) {
+                    stream->parallel_for(
+                        sycl::nd_range<1>(sycl::range<1>(nblocks_padded * block_size), sycl::range<1>(block_size)),
+                        [=](sycl::nd_item<1> item_ct1) {
+                            argsort_f32_i32_global_pass<GGML_SORT_ORDER_ASC>(x, idx_padded, ncols, nrows, ncols_pad, j,
+                                                                             k, item_ct1);
+                        });
+                } else if (order == GGML_SORT_ORDER_DESC) {
+                    stream->parallel_for(
+                        sycl::nd_range<1>(sycl::range<1>(nblocks_padded * block_size), sycl::range<1>(block_size)),
+                        [=](sycl::nd_item<1> item_ct1) {
+                            argsort_f32_i32_global_pass<GGML_SORT_ORDER_DESC>(x, idx_padded, ncols, nrows, ncols_pad, j,
+                                                                              k, item_ct1);
+                        });
+                } else {
+                    GGML_ABORT("invalid sort order");
+                }
+            }
+        }
+
+        const size_t total   = (size_t) nrows * (size_t) ncols;
+        const size_t nblocks = (total + block_size - 1) / block_size;
+        stream->parallel_for(sycl::nd_range<1>(sycl::range<1>(nblocks * block_size), sycl::range<1>(block_size)),
+                             [=](sycl::nd_item<1> item_ct1) {
+                                 copy_argsort_indices_unpadded(idx_padded, dst, nrows, ncols, ncols_pad, item_ct1);
+                             });
+
+        return;
+    }
 
     int nth = 1;
     int max_block_size = ggml_sycl_info().max_work_group_sizes[device];
@@ -2105,8 +2270,6 @@ static void argsort_f32_i32_sycl(const float *x, int *dst, const int ncols,
 
     const sycl::range<3> block_dims(1, 1, nth);
     const sycl::range<3> block_nums(1, nrows, 1);
-    const size_t shared_mem = ncols_pad * sizeof(int);
-    GGML_ASSERT(shared_mem<=ggml_sycl_info().devices[device].smpbo);
 
     if (order == GGML_SORT_ORDER_ASC) {
         stream->submit([&](sycl::handler &cgh) {
@@ -2429,7 +2592,7 @@ inline void ggml_sycl_op_mul_mat_sycl(
 
 #if GGML_SYCL_DNNL && defined(GGML_SYCL_HAS_BF16)
     // Fast path for bf16 src0
-    if (src0->type == GGML_TYPE_BF16 && !g_ggml_sycl_disable_dnn && ggml_is_contiguous(src0) &&
+    if (src0->type == GGML_TYPE_BF16 && g_ggml_sycl_enable_dnn && ggml_is_contiguous(src0) &&
         row_diff == src0->ne[1]) {
         using bf16_t = sycl::ext::oneapi::bfloat16;
         ggml_sycl_pool_alloc<bf16_t> src1_as_bf16(ctx.pool(), src1_ncols*ne10);
@@ -2482,7 +2645,7 @@ inline void ggml_sycl_op_mul_mat_sycl(
                                          : src1_as_f16.get();
 
 #if GGML_SYCL_DNNL
-        if (!g_ggml_sycl_disable_dnn) {
+        if (g_ggml_sycl_enable_dnn) {
                 DnnlGemmWrapper::row_gemm(ctx,row_diff, src1_ncols , ne10, src0_ptr,
                                      DnnlGemmWrapper::to_dt<sycl::half>(), src1_ptr, DnnlGemmWrapper::to_dt<sycl::half>(),
                                       dst_dd_i, DnnlGemmWrapper::to_dt<float>(), stream);
@@ -2532,7 +2695,7 @@ inline void ggml_sycl_op_mul_mat_sycl(
             const int64_t gemm_flops = (int64_t)row_diff * src1_ncols * ne10;
             const bool use_mkl_direct = gemm_flops < 256 * 256 * 256;
 #if GGML_SYCL_DNNL
-            if (!g_ggml_sycl_disable_dnn && !use_mkl_direct) {
+            if (g_ggml_sycl_enable_dnn && !use_mkl_direct) {
                 DnnlGemmWrapper::row_gemm(ctx, row_diff, src1_ncols, ne10, src0_ddf_i,
                                           DnnlGemmWrapper::to_dt<float>(), src1_ddf1_i, DnnlGemmWrapper::to_dt<float>(),
                                           dst_dd_i, DnnlGemmWrapper::to_dt<float>(), stream);
@@ -2625,7 +2788,7 @@ inline void ggml_sycl_op_argsort(ggml_backend_sycl_context & ctx, ggml_tensor *
     enum ggml_sort_order order = (enum ggml_sort_order) dst->op_params[0];
 
     argsort_f32_i32_sycl(src0_dd, (int *)dst_dd, ncols, nrows, order,
-                         main_stream, ctx.device);
+                         main_stream, ctx.device, ctx.pool());
 }
 
 static void ggml_sycl_op_top_k(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
@@ -3352,7 +3515,7 @@ static void ggml_sycl_mul_mat_batched_sycl(ggml_backend_sycl_context & ctx, cons
     const int64_t r3 = ne13 / ne03;
 
 #if GGML_SYCL_DNNL
-    if (!g_ggml_sycl_disable_dnn) {
+    if (g_ggml_sycl_enable_dnn) {
             int64_t str_a0 = nb00 / type_size_src0;
             int64_t str_a1 = nb01 / type_size_src0;
             int64_t str_a2 = nb02 / type_size_src0;
@@ -3527,6 +3690,10 @@ inline bool ggml_sycl_supports_reorder_dmmv(enum ggml_type type) {
         case GGML_TYPE_Q1_0:
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q8_0:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
             return true;
         default:
             return false;
@@ -4092,12 +4259,12 @@ static bool reorder_qw(const ggml_tensor * src0, dpct::queue_ptr stream) {
 }
 
 static bool should_reorder_tensor(ggml_backend_sycl_context& ctx, const ggml_tensor * dst) {
-    return !g_ggml_sycl_disable_optimize && //allow optimize, controlled by $GGML_SYCL_DISABLE_OPT
-            ctx.opt_feature.reorder &&      //allow this device due to good perf, skip the devices with bad perf.
-            dst->op == GGML_OP_MUL_MAT &&   //limit to some supported cases of Q4_0, to do for more cases.
-            // ne[1] <= 8 so multi-column decode (spec / MTP verify) also bootstraps the reorder;
-            // all reorderable types have a _switch_ncols kernel.
-            dst->src[1]->ne[1] <= 8 && dst->src[1]->ne[2]==1 && dst->src[1]->ne[3]==1;
+    return g_ggml_sycl_enable_optimize && //allow optimize, controlled by $GGML_SYCL_ENABLE_OPT
+           ctx.opt_feature.reorder &&      //allow this device due to good perf, skip the devices with bad perf.
+           dst->op == GGML_OP_MUL_MAT &&   //limit to some supported cases of Q4_0, to do for more cases.
+           // ne[1] <= 8 so multi-column decode (spec / MTP verify) also bootstraps the reorder;
+           // all reorderable types have a _switch_ncols kernel.
+           dst->src[1]->ne[1] <= 8 && dst->src[1]->ne[2]==1 && dst->src[1]->ne[3]==1;
 }
 
 static void opt_for_reorder(ggml_backend_sycl_context * ctx, const ggml_tensor * src0, const ggml_tensor * /* src1 */,
@@ -4136,7 +4303,7 @@ static void opt_for_reorder(ggml_backend_sycl_context * ctx, const ggml_tensor *
 
 // Lazily reorder supported MoE expert weights once their fused path is used.
 static void opt_for_reorder_id(ggml_backend_sycl_context * ctx, const ggml_tensor * src0) {
-    if (g_ggml_sycl_disable_optimize || !ctx->opt_feature.reorder) {
+    if (!g_ggml_sycl_enable_optimize || !ctx->opt_feature.reorder) {
         return;
     }
     if (src0->type != GGML_TYPE_Q4_K && src0->type != GGML_TYPE_Q5_K && src0->type != GGML_TYPE_Q6_K) {
@@ -4604,6 +4771,11 @@ static void ggml_sycl_im2col_3d(ggml_backend_sycl_context & ctx, ggml_tensor * d
     ggml_sycl_op_im2col_3d(ctx, dst);
 }
 
+static void ggml_sycl_col2im_1d(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1);
+    ggml_sycl_op_col2im_1d(ctx, dst);
+}
+
 static void ggml_sycl_conv_3d(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
     ggml_sycl_op_conv_3d(ctx, dst);
@@ -4912,6 +5084,12 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg
         case GGML_OP_SOFT_MAX_BACK:
             ggml_sycl_op_soft_max_back(ctx, dst);
             break;
+        case GGML_OP_CROSS_ENTROPY_LOSS:
+            ggml_sycl_cross_entropy_loss(ctx, dst);
+            break;
+        case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
+            ggml_sycl_cross_entropy_loss_back(ctx, dst);
+            break;
         case GGML_OP_ROPE:
             ggml_sycl_rope(ctx, dst);
             break;
@@ -4924,6 +5102,9 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg
         case GGML_OP_IM2COL_3D:
             ggml_sycl_im2col_3d(ctx, dst);
             break;
+        case GGML_OP_COL2IM_1D:
+            ggml_sycl_col2im_1d(ctx, dst);
+            break;
         case GGML_OP_POOL_2D:
             ggml_sycl_pool2d(ctx, dst);
             break;
@@ -5204,7 +5385,10 @@ static ggml_status ggml_backend_sycl_graph_compute(ggml_backend_t backend, ggml_
     auto * sycl_ctx = static_cast<ggml_backend_sycl_context *>(backend->context);
 
 #ifdef GGML_SYCL_GRAPH
-    bool use_sycl_graph = !g_ggml_sycl_disable_graph && check_graph_compatibility(cgraph);
+    bool use_sycl_graph = false;
+    if (g_ggml_sycl_enable_graph) {
+        use_sycl_graph = check_graph_compatibility(cgraph);
+    }
     if (use_sycl_graph) {
         const bool graph_support = dpct::get_device(sycl_ctx->device).has(sycl::aspect::ext_oneapi_limited_graph);
         if (!graph_support) {
@@ -5470,7 +5654,6 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
                 // TODO: This specific configuration can fail with oneDNN and needs more debugging
                 if (!ggml_is_permuted(a) && ggml_is_permuted(b) && b->ne[2] > 1 && b->ne[3] > 1 &&
                     a->ne[0] > 128 && a->ne[2] == 1 && src0_type == GGML_TYPE_F16) {
-                        printf("zjy 2\n");
                     return false;
                 }
                 return true;
@@ -5538,70 +5721,99 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
             {
                 ggml_type src0_type = op->src[0]->type;
                 ggml_type src1_type = op->src[1]->type;
-                if (src0_type == src1_type && (ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1])) && src0_type != GGML_TYPE_BF16) {
-                    return true;
+
+                if (src0_type == GGML_TYPE_F16) {
+                    if (src1_type == GGML_TYPE_Q2_K ||
+                        src1_type == GGML_TYPE_Q3_K ||
+                        src1_type == GGML_TYPE_Q4_K ||
+                        src1_type == GGML_TYPE_Q5_K ||
+                        src1_type == GGML_TYPE_Q6_K ||
+                        src1_type == GGML_TYPE_IQ2_XXS ||
+                        src1_type == GGML_TYPE_IQ2_XS ||
+                        src1_type == GGML_TYPE_IQ2_S ||
+                        src1_type == GGML_TYPE_IQ3_XXS ||
+                        src1_type == GGML_TYPE_IQ1_S ||
+                        src1_type == GGML_TYPE_IQ1_M ||
+                        src1_type == GGML_TYPE_IQ3_S ||
+                        src1_type == GGML_TYPE_IQ4_XS) {
+                        return false;
+                    }
                 }
-                if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) {
-                    return true;
+
+                if (src0_type == GGML_TYPE_BF16) {
+                    if (src1_type == GGML_TYPE_Q4_0 || //big error in ut
+                        src1_type == GGML_TYPE_Q4_1 || //big error in ut
+                        src1_type == GGML_TYPE_Q8_0 || //big error in ut
+                        src1_type == GGML_TYPE_Q2_K ||
+                        src1_type == GGML_TYPE_Q3_K ||
+                        src1_type == GGML_TYPE_Q4_K ||
+                        src1_type == GGML_TYPE_Q5_K ||
+                        src1_type == GGML_TYPE_Q6_K ||
+                        src1_type == GGML_TYPE_IQ2_XXS ||
+                        src1_type == GGML_TYPE_IQ2_XS ||
+                        src1_type == GGML_TYPE_IQ2_S ||
+                        src1_type == GGML_TYPE_IQ3_XXS ||
+                        src1_type == GGML_TYPE_IQ1_S ||
+                        src1_type == GGML_TYPE_IQ1_M ||
+                        src1_type == GGML_TYPE_IQ3_S ||
+                        src1_type == GGML_TYPE_IQ4_XS) {
+                        return false;
+                    }
                 }
-                if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F16) {
-                    return true;
+
+                if (src0_type == GGML_TYPE_F32) {
+                    if (src1_type == GGML_TYPE_Q2_K ||
+                        src1_type == GGML_TYPE_Q3_K ||
+                        src1_type == GGML_TYPE_Q4_K ||
+                        src1_type == GGML_TYPE_Q5_K ||
+                        src1_type == GGML_TYPE_Q6_K ||
+                        src1_type == GGML_TYPE_IQ2_XXS ||
+                        src1_type == GGML_TYPE_IQ2_XS ||
+                        src1_type == GGML_TYPE_IQ2_S ||
+                        src1_type == GGML_TYPE_IQ3_XXS ||
+                        src1_type == GGML_TYPE_IQ1_S ||
+                        src1_type == GGML_TYPE_IQ1_M ||
+                        src1_type == GGML_TYPE_IQ3_S ||
+                        src1_type == GGML_TYPE_IQ4_XS) {
+                        return false;
+                    }
                 }
-                if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q8_0) {
-                    return true;
+
+                if (src1_type == GGML_TYPE_F32) {
+                    if (src0_type == GGML_TYPE_Q1_0 ||
+                        src0_type == GGML_TYPE_NVFP4 ||
+                        src0_type == GGML_TYPE_Q2_K ||
+                        src0_type == GGML_TYPE_Q3_K ||
+                        src0_type == GGML_TYPE_Q4_K ||
+                        src0_type == GGML_TYPE_Q5_K ||
+                        src0_type == GGML_TYPE_Q6_K ||
+                        src0_type == GGML_TYPE_IQ2_XXS ||
+                        src0_type == GGML_TYPE_IQ2_XS ||
+                        src0_type == GGML_TYPE_IQ2_S ||
+                        src0_type == GGML_TYPE_IQ3_XXS ||
+                        src0_type == GGML_TYPE_IQ1_S ||
+                        src0_type == GGML_TYPE_IQ1_M ||
+                        src0_type == GGML_TYPE_IQ3_S ||
+                        src0_type == GGML_TYPE_IQ4_NL ||
+                        src0_type == GGML_TYPE_IQ4_XS
+                    ) {
+                        return false;
+                    }
                 }
-                if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q4_0) {
-                    return true;
+
+                if (src0_type == src1_type) {
+                    if (src1_type == GGML_TYPE_IQ2_XXS ||
+                        src1_type == GGML_TYPE_IQ2_XS ||
+                        src1_type == GGML_TYPE_IQ2_S ||
+                        src1_type == GGML_TYPE_IQ3_XXS ||
+                        src1_type == GGML_TYPE_IQ3_S ||
+                        src1_type == GGML_TYPE_IQ1_S ||
+                        src1_type == GGML_TYPE_IQ1_M) {
+                        return false;
+                    }
                 }
-                if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q4_1) {
-                    return true;
-                }
-                if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) {
-                    return true;
-                }
-                if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) {
-                    return true;
-                }
-                if (src0_type == GGML_TYPE_Q8_0 && src1_type == GGML_TYPE_F32) {
-                    return true;
-                }
-                if (src0_type == GGML_TYPE_Q4_0 && src1_type == GGML_TYPE_F32) {
-                    return true;
-                }
-                if (src0_type == GGML_TYPE_Q4_1 && src1_type == GGML_TYPE_F32) {
-                    return true;
-                }
-                if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q5_0) {
-                    return true;
-                }
-                if (src0_type == GGML_TYPE_Q5_0 && src1_type == GGML_TYPE_F32) {
-                    return true;
-                }
-                if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q5_1) {
-                    return true;
-                }
-                if (src0_type == GGML_TYPE_Q5_1 && src1_type == GGML_TYPE_F32) {
-                    return true;
-                }
-                if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_IQ4_NL) {
-                    return true;
-                }
-                if(src0_type == GGML_TYPE_Q8_0 && src1_type == GGML_TYPE_Q8_0) {
-                    return true;
-                }
-                if(src0_type == GGML_TYPE_Q5_0 && src1_type == GGML_TYPE_Q5_0) {
-                    return true;
-                }
-                if(src0_type == GGML_TYPE_Q5_1 && src1_type == GGML_TYPE_Q5_1) {
-                    return true;
-                }
-                if(src0_type == GGML_TYPE_Q4_0 && src1_type == GGML_TYPE_Q4_0) {
-                    return true;
-                }
-                if(src0_type == GGML_TYPE_Q4_1 && src1_type == GGML_TYPE_Q4_1) {
-                    return true;
-                }
-                return false;
+
+                return true;
             }
         case GGML_OP_REPEAT_BACK:
             {
@@ -5643,7 +5855,7 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
         case GGML_OP_SCALE:
             return true;
         case GGML_OP_CONT:
-            return op->src[0]->type != GGML_TYPE_BF16;
+            return true;
         case GGML_OP_TRI:
             {
                 const ggml_tensor * src0 = op->src[0];
@@ -5666,6 +5878,14 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
         case GGML_OP_IM2COL_3D:
         case GGML_OP_UPSCALE:
             return true;
+        case GGML_OP_COL2IM_1D:
+            return ggml_is_contiguous(op->src[0]) &&
+                   (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16
+#ifdef GGML_SYCL_HAS_BF16
+                    || op->type == GGML_TYPE_BF16
+#endif
+                   ) &&
+                   op->src[0]->type == op->type;
         case GGML_OP_CONV_3D:
             return op->type == GGML_TYPE_F32 &&
                    (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) &&
@@ -5677,8 +5897,7 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
         case GGML_OP_MEAN:
             return ggml_is_contiguous(op->src[0]);
         case GGML_OP_ARGSORT:
-            return op->src[0]->ne[0] * sizeof(int) <=
-                   ggml_sycl_info().devices[device].smpbo;
+            return true;
         case GGML_OP_TOP_K: {
             const ggml_tensor * src0 = op->src[0];
             const int k = op->ne[0];
@@ -5690,9 +5909,8 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
         }
         case GGML_OP_POOL_2D:
         case GGML_OP_POOL_1D:
-            return true;
         case GGML_OP_ACC:
-            return ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]);
+            return true;
         case GGML_OP_PAD:
             if (ggml_get_op_params_i32(op, 8) != 0) {
                 return false;
@@ -5725,6 +5943,8 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
         case GGML_OP_FILL:
         case GGML_OP_CUMSUM:
         case GGML_OP_DIAG:
+        case GGML_OP_CROSS_ENTROPY_LOSS:
+        case GGML_OP_CROSS_ENTROPY_LOSS_BACK:
             return true;
         case GGML_OP_SOLVE_TRI:
             return op->src[0]->ne[0] <= SYCL_SOLVE_TRI_MAX_N && op->src[1]->ne[0] <= SYCL_SOLVE_TRI_MAX_K;

ggml/src/ggml-sycl/presets.hpp

@@ -19,6 +19,7 @@
 #define WARP_SIZE GGML_SYCL_WARP_SIZE
 #define MATRIX_ROW_PADDING 512 // last row of quant. matrices is a multiple of this to avoid out-of-bounds memory accesses
 
+#define SYCL_COL2IM_1D_BLOCK_SIZE 256
 #define SYCL_GELU_BLOCK_SIZE 256
 #define SYCL_SILU_BLOCK_SIZE 256
 #define SYCL_TANH_BLOCK_SIZE 256
@@ -62,7 +63,7 @@
 #endif
 
 #ifndef K_QUANTS_PER_ITERATION
-#define K_QUANTS_PER_ITERATION 2
+#define K_QUANTS_PER_ITERATION 1
 #else
 static_assert(K_QUANTS_PER_ITERATION == 1 || K_QUANTS_PER_ITERATION == 2, "K_QUANTS_PER_ITERATION must be 1 or 2");
 #endif

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

@@ -129,7 +129,7 @@ typedef struct VkPhysicalDeviceShaderMixedFloatDotProductFeaturesVALVE {
 #endif
 
 #define ROUNDUP_POW2(M, N) (((M) + (N) - 1) & ~((N) - 1))
-#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+#define CEIL_DIV(M, N) (((M) / (N)) + (((M) % (N)) != 0))
 static bool is_pow2(uint32_t x) { return x > 1 && (x & (x-1)) == 0; }
 
 #define VK_VENDOR_ID_AMD 0x1002
@@ -16308,7 +16308,18 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
     uint32_t submit_count = 0;
     uint64_t batch_flops = 0;
     uint64_t total_flops = 0;
-    uint64_t flops_per_submit = std::min(uint64_t(200'000'000'000), ctx->last_total_flops / 40u);
+    uint64_t flops_cap = 200'000'000'000ULL;
+
+    // On weaker AMD GPUs larger submissions can hit a driver timeout, submit more often to avoid this
+    if (ctx->device->vendor_id == VK_VENDOR_ID_AMD && ctx->device->shader_core_count > 0) {
+        if (ctx->device->architecture == AMD_GCN && ctx->device->shader_core_count < 32) {
+            flops_cap = 500'000'000ULL * ctx->device->shader_core_count;
+        } else if (ctx->device->architecture != AMD_GCN && ctx->device->shader_core_count < 24) {
+            flops_cap = 2'000'000'000ULL * ctx->device->shader_core_count;
+        }
+    }
+    uint64_t flops_per_submit = std::min(flops_cap, ctx->last_total_flops / 40u);
+
     for (int i = 0; i < cgraph->n_nodes; i++) {
         if (first_node_in_batch) {
             submit_node_idx = i;
@@ -17370,21 +17381,24 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
             return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_SET_ROWS:
             {
-                switch (op->type) {
-                    case GGML_TYPE_F32:
-                    case GGML_TYPE_F16:
-                    case GGML_TYPE_BF16:
-                    case GGML_TYPE_Q1_0:
-                    case GGML_TYPE_Q4_0:
-                    case GGML_TYPE_Q4_1:
-                    case GGML_TYPE_Q5_0:
-                    case GGML_TYPE_Q5_1:
-                    case GGML_TYPE_Q8_0:
-                    case GGML_TYPE_IQ4_NL:
-                        return true;
-                    default:
-                        return false;
+                if (op->src[0]->type == GGML_TYPE_F32) {
+                    switch (op->type) {
+                        case GGML_TYPE_F32:
+                        case GGML_TYPE_F16:
+                        case GGML_TYPE_BF16:
+                        case GGML_TYPE_Q1_0:
+                        case GGML_TYPE_Q4_0:
+                        case GGML_TYPE_Q4_1:
+                        case GGML_TYPE_Q5_0:
+                        case GGML_TYPE_Q5_1:
+                        case GGML_TYPE_Q8_0:
+                        case GGML_TYPE_IQ4_NL:
+                            return true;
+                        default:
+                            return false;
+                    }
                 }
+                return false;
             }
         case GGML_OP_CONT:
         case GGML_OP_CPY:

ggml/src/ggml.c

@@ -525,7 +525,11 @@ const char * ggml_commit(void) {
 
 #if defined(_MSC_VER) || defined(__MINGW32__)
 static int64_t timer_freq, timer_start;
-void ggml_time_init(void) {
+static BOOL CALLBACK ggml_time_init_once(PINIT_ONCE once, PVOID param, PVOID *ctx) {
+    UNUSED(once);
+    UNUSED(param);
+    UNUSED(ctx);
+
     LARGE_INTEGER t;
     QueryPerformanceFrequency(&t);
     timer_freq = t.QuadPart;
@@ -535,6 +539,12 @@ void ggml_time_init(void) {
     // We subtract the program start time to reduce the likelihood of that happening.
     QueryPerformanceCounter(&t);
     timer_start = t.QuadPart;
+
+    return TRUE;
+}
+void ggml_time_init(void) {
+    static INIT_ONCE once = INIT_ONCE_STATIC_INIT;
+    InitOnceExecuteOnce(&once, ggml_time_init_once, NULL, NULL);
 }
 int64_t ggml_time_ms(void) {
     LARGE_INTEGER t;
@@ -671,6 +681,14 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = {
         .to_float                 = (ggml_to_float_t) dequantize_row_q1_0,
         .from_float_ref           = (ggml_from_float_t) quantize_row_q1_0_ref,
     },
+    [GGML_TYPE_Q2_0] = {
+        .type_name                = "q2_0",
+        .blck_size                = QK2_0,
+        .type_size                = sizeof(block_q2_0),
+        .is_quantized             = true,
+        .to_float                 = (ggml_to_float_t) dequantize_row_q2_0,
+        .from_float_ref           = (ggml_from_float_t) quantize_row_q2_0_ref,
+    },
     [GGML_TYPE_Q4_0] = {
         .type_name                = "q4_0",
         .blck_size                = QK4_0,
@@ -1407,6 +1425,7 @@ enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype) {
         case GGML_FTYPE_MOSTLY_Q4_0:          wtype = GGML_TYPE_Q4_0;  break;
         case GGML_FTYPE_MOSTLY_Q4_1:          wtype = GGML_TYPE_Q4_1;  break;
         case GGML_FTYPE_MOSTLY_Q1_0:          wtype = GGML_TYPE_Q1_0;  break;
+        case GGML_FTYPE_MOSTLY_Q2_0:          wtype = GGML_TYPE_Q2_0;  break;
         case GGML_FTYPE_MOSTLY_Q5_0:          wtype = GGML_TYPE_Q5_0;  break;
         case GGML_FTYPE_MOSTLY_Q5_1:          wtype = GGML_TYPE_Q5_1;  break;
         case GGML_FTYPE_MOSTLY_Q8_0:          wtype = GGML_TYPE_Q8_0;  break;
@@ -3907,7 +3926,7 @@ struct ggml_tensor * ggml_set_rows(
     GGML_ASSERT(b->ne[2] % c->ne[1] == 0);
     GGML_ASSERT(b->ne[3] % c->ne[2] == 0);
     GGML_ASSERT(c->ne[3] == 1);
-    GGML_ASSERT(b->type == GGML_TYPE_F32);
+    GGML_ASSERT(b->type == GGML_TYPE_F32 || b->type == GGML_TYPE_F16);
     GGML_ASSERT(c->type == GGML_TYPE_I64 || c->type == GGML_TYPE_I32);
 
     GGML_ASSERT(ggml_is_contiguous_rows(a));
@@ -7409,6 +7428,10 @@ static int ggml_node_list_find_tensor(const struct ggml_cgraph * cgraph,
     return -1;
 }
 
+static bool ggml_is_constant(const struct ggml_tensor * tensor) {
+    return tensor->buffer != NULL && ggml_backend_buffer_get_usage(tensor->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && (tensor->flags & GGML_TENSOR_FLAG_PARAM) == 0;
+}
+
 bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph,
                                 const int *                node_idxs,
                                 int                        count,
@@ -7454,10 +7477,11 @@ bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph,
             return false;
         }
 
-        // if node is a view, check if the view_src and all it's parent view_srcs are within the subgraph
+        // if node is a view, check if the view_src and all its parent view_srcs are within the subgraph.
+        // external view sources are allowed only for weight tensors, which are constant for this graph execution.
         struct ggml_tensor * view_src = node->view_src;
         while (view_src) {
-            if (ggml_node_list_find_tensor(cgraph, node_idxs, count, view_src) == -1) {
+            if (ggml_node_list_find_tensor(cgraph, node_idxs, count, view_src) == -1 && !ggml_is_constant(view_src)) {
                 return false;
             }
             view_src = view_src->view_src;
@@ -7729,6 +7753,7 @@ size_t ggml_quantize_chunk(
 
     switch (type) {
         case GGML_TYPE_Q1_0:    result = quantize_q1_0   (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
+        case GGML_TYPE_Q2_0:    result = quantize_q2_0   (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
         case GGML_TYPE_Q4_0:    result = quantize_q4_0   (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
         case GGML_TYPE_Q4_1:    result = quantize_q4_1   (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;
         case GGML_TYPE_Q5_0:    result = quantize_q5_0   (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break;

gguf-py/gguf/constants.py

@@ -4533,6 +4533,7 @@ class GGMLQuantizationType(IntEnum):
     MXFP4   = 39
     NVFP4   = 40
     Q1_0    = 41
+    Q2_0    = 42
 
 
 class ExpertGatingFuncType(IntEnum):
@@ -4588,6 +4589,7 @@ class LlamaFileType(IntEnum):
     MOSTLY_MXFP4_MOE     = 38  # except 1d tensors
     MOSTLY_NVFP4         = 39  # except 1d tensors
     MOSTLY_Q1_0          = 40  # except 1d tensors
+    MOSTLY_Q2_0          = 41  # except 1d tensors
 
     GUESSED              = 1024  # not specified in the model file
 
@@ -4713,6 +4715,7 @@ GGML_QUANT_SIZES: dict[GGMLQuantizationType, tuple[int, int]] = {
     GGMLQuantizationType.MXFP4:   (32, 1 + 16),
     GGMLQuantizationType.NVFP4:   (64, 4 + 32),
     GGMLQuantizationType.Q1_0:    (128, 2 + 16),
+    GGMLQuantizationType.Q2_0:    (64, 2 + 16),
 }
 
 

include/llama.h

@@ -155,6 +155,7 @@ extern "C" {
         LLAMA_FTYPE_MOSTLY_MXFP4_MOE     = 38, // except 1d tensors
         LLAMA_FTYPE_MOSTLY_NVFP4         = 39, // except 1d tensors
         LLAMA_FTYPE_MOSTLY_Q1_0          = 40, // except 1d tensors
+        LLAMA_FTYPE_MOSTLY_Q2_0          = 41, // except 1d tensors
 
         LLAMA_FTYPE_GUESSED = 1024, // not specified in the model file
     };

src/llama-batch.cpp

@@ -379,6 +379,8 @@ bool llama_batch_allocr::init(
                     LLAMA_LOG_ERROR("%s: sequence %d positions are decreasing (not allowed)\n", __func__, seq_id);
                     return false;
                 }
+
+                cur_seq_pos[seq_id] = pos;
             }
         }
     }
@@ -505,7 +507,7 @@ llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) {
     return ubatch_add(idxs, idxs.size(), false);
 }
 
-llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential) {
+llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail) {
     if (sequential && has_cpl) {
         LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__);
 
@@ -548,7 +550,7 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential)
         }
     }
 
-    const uint32_t n_seqs = cur_seq_set.size();
+    uint32_t n_seqs = cur_seq_set.size();
 
     // we are done
     if (n_seqs == 0) {
@@ -569,7 +571,7 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential)
     std::vector<idx_vec_t> idxs_per_seq(n_seqs);
 
     while (true) {
-        // we can only add new n_seq_tokens tokens if all the sequence sets have at least one more unused token and
+        // we can only add new n_seq_tokens tokens if all the sequence sets have at least 1 more unused tokens and
         //   if we haven't reached n_ubatch
         bool can_expand = true;
 
@@ -600,6 +602,72 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential)
         }
     }
 
+    // if n_keep_tail > 0, keep only the seqs that either finish in this ubatch or have at least
+    //   n_keep_tail tokens remaining for a future ubatch, so that the trailing n_keep_tail tokens
+    //   of each seq are never split across ubatches
+    if (n_keep_tail > 0) {
+        GGML_ASSERT(n_ubatch > n_keep_tail);
+
+        auto n_remaining = [&](uint32_t s) {
+            return (uint32_t) (seq_set_map[cur_seq_set[s]].size() - cur_idx[s]);
+        };
+
+        // keep the longest prefix of seqs that satisfy the constraint, to preserve sequential seq ids
+        uint32_t n_keep = 0;
+        while (n_keep < n_seqs) {
+            const uint32_t remaining = n_remaining(n_keep);
+
+            if (remaining != 0 && remaining < n_keep_tail) {
+                break;
+            }
+
+            n_keep++;
+        }
+
+        // all seqs violate the constraint - resolve the first one directly and emit it alone
+        if (n_keep == 0) {
+            auto & idxs = idxs_per_seq[0];
+
+            const auto & seq_idxs = seq_set_map[cur_seq_set[0]];
+
+            if (idxs.size() + n_remaining(0) <= n_ubatch) {
+                // extend the seq to completion
+                while (n_remaining(0) > 0) {
+                    const int32_t idx = seq_idxs[cur_idx[0]];
+
+                    idxs.push_back(idx);
+
+                    used[idx] = true;
+                    ++n_used;
+
+                    ++cur_idx[0];
+                }
+            } else {
+                // truncate the seq so that at least n_keep_tail tokens remain
+                while (n_remaining(0) < n_keep_tail) {
+                    used[idxs.back()] = false;
+                    --n_used;
+
+                    idxs.pop_back();
+
+                    --cur_idx[0];
+                }
+            }
+
+            n_keep = 1;
+        }
+
+        // return the tokens of the deferred seqs back to the pool
+        for (uint32_t s = n_keep; s < n_seqs; ++s) {
+            for (const int32_t idx : idxs_per_seq[s]) {
+                used[idx] = false;
+                --n_used;
+            }
+        }
+
+        n_seqs = n_keep;
+    }
+
     // concat the per-sequence-set lists
     std::vector<int32_t> idxs;
 
@@ -814,7 +882,7 @@ void llama_batch_allocr::ubatch_print(const llama_ubatch & ubatch, int debug) {
         LLAMA_LOG_DEBUG("%s:   output     = %p\n", __func__, (void *) ubatch.output);
         LLAMA_LOG_DEBUG("%s:   n_outputs  = %d\n", __func__, n_outputs);
 
-        if (debug > 1) {
+        if (debug > 0) {
             int seq_id_max = 0;
             for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
                 for (int s = 0; s < ubatch.n_seq_id[i]; ++s) {

src/llama-batch.h

@@ -104,7 +104,8 @@ public:
 
     // make ubatches of equal-length sequences sets
     // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids
-    llama_ubatch split_equal(uint32_t n_ubatch, bool sequential);
+    // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch
+    llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail);
 
     // sequence-set-wise split - each ubatch contains a single sequence-set
     llama_ubatch split_seq(uint32_t n_ubatch);

src/llama-context.cpp

@@ -17,6 +17,7 @@
 #include <cstring>
 #include <limits>
 #include <stdexcept>
+#include <string>
 
 //
 // llama_context
@@ -30,6 +31,30 @@ static llm_graph_type ctx_type_to_graph_type(llama_context_type ctx_type) {
     throw std::runtime_error("Unsupported ctx type");
 }
 
+struct llm_fused_op_probe {
+    llm_fused_op op;
+    const char * name;
+    uint32_t n_tokens_per_seq;
+};
+
+static const llm_fused_op_probe llm_fused_op_flash_attn_probe = {
+    /*.op               =*/ LLM_FUSED_OP_FLASH_ATTN,
+    /*.name             =*/ "Flash Attention",
+    /*.n_tokens_per_seq =*/ 1,
+};
+
+static const llm_fused_op_probe llm_fused_op_gdn_ar_probe = {
+    /*.op               =*/ LLM_FUSED_OP_GDN_AR,
+    /*.name             =*/ "fused Gated Delta Net (autoregressive)",
+    /*.n_tokens_per_seq =*/ 1,
+};
+
+static const llm_fused_op_probe llm_fused_op_gdn_ch_probe = {
+    /*.op               =*/ LLM_FUSED_OP_GDN_CH,
+    /*.name             =*/ "fused Gated Delta Net (chunked)",
+    /*.n_tokens_per_seq =*/ 16,
+};
+
 llama_context::llama_context(
         const llama_model & model,
               llama_context_params params) :
@@ -436,6 +461,69 @@ llama_context::~llama_context() {
     ggml_opt_free(opt_ctx);
 }
 
+void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs) {
+    const char * func = __func__;
+    auto resolve = [&](const llm_fused_op_probe & probe, bool & enabled) {
+        if (!enabled) {
+            return;
+        }
+
+        const uint32_t n_tokens_probe = probe.n_tokens_per_seq*n_seqs;
+
+        auto * gf = graph_reserve(n_tokens_probe, n_seqs, n_tokens_probe, mctx, true);
+        if (!gf) {
+            throw std::runtime_error(std::string("failed to reserve graph for ") + probe.name + " check");
+        }
+
+        bool device_mismatch = false;
+        for (const auto & node : get_gf_res_reserve()->get_fused_nodes()) {
+            if (node.op != probe.op) {
+                continue;
+            }
+
+            GGML_ASSERT(node.il >= 0);
+
+            ggml_backend_t backend_fused = ggml_backend_sched_get_tensor_backend(sched.get(), node.tensor);
+            ggml_backend_dev_t device_fused = backend_fused ? ggml_backend_get_device(backend_fused) : nullptr;
+
+            // TODO: make this descriptor-specific; model.dev_layer() preserves the current behavior,
+            // but is still wrong for cases like --no-kv-offload.
+            ggml_backend_dev_t device_layer = model.dev_layer(node.il);
+
+            if (device_fused != device_layer) {
+                LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but %s "
+                        "is assigned to device %s (usually due to missing support)\n",
+                        func, node.il,
+                        device_layer ? ggml_backend_dev_name(device_layer) : "none",
+                        probe.name,
+                        device_fused ? ggml_backend_dev_name(device_fused) : "none");
+                device_mismatch = true;
+                break;
+            }
+        }
+
+        if (device_mismatch) {
+            enabled = false;
+            LLAMA_LOG_WARN("%s: %s not supported, set to disabled\n", func, probe.name);
+        } else {
+            enabled = true;
+            LLAMA_LOG_INFO("%s: %s enabled\n", func, probe.name);
+        }
+    };
+
+    if (cparams.auto_fa) {
+        resolve(llm_fused_op_flash_attn_probe, cparams.flash_attn);
+        cparams.auto_fa = false;
+    }
+
+    if (cparams.auto_fgdn) {
+        LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", func);
+        resolve(llm_fused_op_gdn_ar_probe, cparams.fused_gdn_ar);
+        resolve(llm_fused_op_gdn_ch_probe, cparams.fused_gdn_ch);
+        cparams.auto_fgdn = false;
+    }
+}
+
 void llama_context::sched_reserve() {
     if (!sched_need_reserve) {
         return;
@@ -475,128 +563,7 @@ void llama_context::sched_reserve() {
 
     LLAMA_LOG_DEBUG("%s: worst-case: n_tokens = %d, n_seqs = %d, n_outputs = %d\n", __func__, n_tokens, n_seqs, n_outputs);
 
-    // resolve automatic Flash Attention use
-    if (cparams.auto_fa) {
-        auto * gf = graph_reserve(1, n_seqs, n_outputs, mctx.get(), true);
-        if (!gf) {
-            throw std::runtime_error("failed to reserve graph for Flash Attention check");
-        }
-
-        const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FATTN) + 1;
-        bool fa_device_mismatch = false;
-        for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
-            ggml_tensor * n = ggml_graph_node(gf, i);
-            if (n->op != GGML_OP_FLASH_ATTN_EXT) {
-                continue;
-            }
-            ggml_backend_dev_t device_fa = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n));
-
-            // TODO: instead of the tensor names, use a map to keep track of which (FA) tensors belong to which layer
-            GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FATTN "-", prefix_len) == 0);
-            const int il = std::stoi(n->name + prefix_len);
-            ggml_backend_dev_t device_kv = model.dev_layer(il);
-            if (device_fa != device_kv) {
-                LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the Flash Attention tensor "
-                        "is assigned to device %s (usually due to missing support)\n",
-                        __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_fa));
-                // FIXME: fa_device_mismatch logic is wrong for --no-kv-offload, but this is broken anyways
-                fa_device_mismatch = true;
-                break;
-            }
-        }
-
-        if (fa_device_mismatch) {
-            cparams.flash_attn = false;
-            LLAMA_LOG_WARN("%s: Flash Attention was auto, set to disabled\n", __func__);
-        } else {
-            cparams.flash_attn = true;
-            LLAMA_LOG_INFO("%s: Flash Attention was auto, set to enabled\n", __func__);
-        }
-
-        cparams.auto_fa = false;
-    }
-
-    if (cparams.auto_fgdn) {
-        LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", __func__);
-
-        if (cparams.fused_gdn_ar) {
-            auto * gf = graph_reserve(1, n_seqs, n_outputs, mctx.get(), true);
-            if (!gf) {
-                throw std::runtime_error("failed to reserve graph for fused Gated Delta Net check (autoregressive)");
-            }
-
-            const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FGDN_AR) + 1;
-            bool gdn_device_mismatch = false;
-            for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
-                ggml_tensor * n = ggml_graph_node(gf, i);
-                if (n->op != GGML_OP_GATED_DELTA_NET) {
-                    continue;
-                }
-                ggml_backend_dev_t device_gdn = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n));
-
-                GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FGDN_AR "-", prefix_len) == 0);
-                const int il = std::stoi(n->name + prefix_len);
-                ggml_backend_dev_t device_kv = model.dev_layer(il);
-                if (device_gdn != device_kv) {
-                    LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the fused Gated Delta Net tensor "
-                            "is assigned to device %s (usually due to missing support)\n",
-                            __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_gdn));
-                    gdn_device_mismatch = true;
-                    break;
-                }
-            }
-
-            if (gdn_device_mismatch) {
-                cparams.fused_gdn_ar = false;
-                LLAMA_LOG_WARN("%s: fused Gated Delta Net (autoregressive) not supported, set to disabled\n", __func__);
-            } else {
-                LLAMA_LOG_INFO("%s: fused Gated Delta Net (autoregressive) enabled\n", __func__);
-            }
-        }
-
-        if (cparams.fused_gdn_ch) {
-            // more than one token in the batch per sequence in order to take the chunked path
-            // note: n_outputs must match n_tokens for embedding models with mean/rank pooling,
-            // because build_pooling creates inp_mean with shape [n_tokens, n_seqs] and multiplies
-            // it with t_embd which is reduced to [n_outputs, ...] via out_ids. if n_outputs != n_tokens,
-            // the ggml_mul_mat assertion fails.
-            const uint32_t n_tokens_ch = 16*n_seqs;
-            auto * gf = graph_reserve(n_tokens_ch, n_seqs, n_tokens_ch, mctx.get(), true);
-            if (!gf) {
-                throw std::runtime_error("failed to reserve graph for fused Gated Delta Net check (chunked)");
-            }
-
-            const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FGDN_CH) + 1;
-            bool gdn_device_mismatch = false;
-            for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
-                ggml_tensor * n = ggml_graph_node(gf, i);
-                if (n->op != GGML_OP_GATED_DELTA_NET) {
-                    continue;
-                }
-                ggml_backend_dev_t device_gdn = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n));
-
-                GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FGDN_CH "-", prefix_len) == 0);
-                const int il = std::stoi(n->name + prefix_len);
-                ggml_backend_dev_t device_kv = model.dev_layer(il);
-                if (device_gdn != device_kv) {
-                    LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the fused Gated Delta Net tensor "
-                            "is assigned to device %s (usually due to missing support)\n",
-                            __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_gdn));
-                    gdn_device_mismatch = true;
-                    break;
-                }
-            }
-
-            if (gdn_device_mismatch) {
-                cparams.fused_gdn_ch = false;
-                LLAMA_LOG_WARN("%s: fused Gated Delta Net (chunked) not supported, set to disabled\n", __func__);
-            } else {
-                LLAMA_LOG_INFO("%s: fused Gated Delta Net (chunked) enabled\n", __func__);
-            }
-        }
-
-        cparams.auto_fgdn = false;
-    }
+    resolve_fused_ops(mctx.get(), n_seqs);
 
     // reserve worst-case graph
     int n_splits_pp = -1;

src/llama-context.h

@@ -262,6 +262,10 @@ private:
 
     llm_graph_cb graph_get_cb() const;
 
+    // disable auto fused ops (Flash Attention, Gated Delta Net) whose op lands on a device
+    // that differs from the layer it belongs to (usually due to missing backend support)
+    void resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs);
+
     // TODO: read/write lora adapters and cvec
     size_t state_write_data(llama_io_write_i & io);
     size_t state_read_data (llama_io_read_i  & io);

src/llama-graph.cpp

@@ -63,26 +63,6 @@ static bool can_reuse_kq_mask(
 
 // impl
 
-static ggml_tensor * ggml_mul_mat_aux(
-        ggml_context * ctx,
-        ggml_tensor * cur,
-        ggml_tensor * rot) {
-    const auto n = rot->ne[0];
-
-    ggml_tensor * res;
-
-    if (!ggml_is_contiguous(cur)) {
-        res = ggml_cont_2d   (ctx, cur, n, ggml_nelements(cur)/n);
-    } else {
-        res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n);
-    }
-    res = ggml_mul_mat   (ctx, rot, res);
-    ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD);
-    res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]);
-
-    return res;
-}
-
 void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) {
     if (ubatch->token) {
         const int64_t n_tokens = ubatch->n_tokens;
@@ -881,6 +861,14 @@ void llm_graph_input_dsv4::set_input(const llama_ubatch * ubatch) {
     dsv4_set_comp_inputs(inp_hca, plan_hca, "hca", debug > 0, ubatch->n_tokens, n_stream);
     dsv4_set_comp_inputs(inp_lid, plan_lid, "lid", debug > 0, ubatch->n_tokens, n_stream);
 
+    if (inp_csa.k_rot && inp_csa.k_rot->buffer) {
+        mctx->get_csa()->set_input_k_rot(inp_csa.k_rot);
+    }
+
+    if (inp_hca.k_rot && inp_hca.k_rot->buffer) {
+        mctx->get_hca()->set_input_k_rot(inp_hca.k_rot);
+    }
+
     if (inp_lid.k_rot && inp_lid.k_rot->buffer) {
         mctx->get_lid()->set_input_k_rot(inp_lid.k_rot);
     }
@@ -1204,6 +1192,7 @@ void llm_graph_result::reset() {
     params = {};
 
     inputs.clear();
+    fused_nodes.clear();
 
     buf_compute_meta.resize(ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false));
 
@@ -1305,6 +1294,10 @@ llm_graph_input_i * llm_graph_result::add_input(llm_graph_input_ptr input) {
     return inputs.back().get();
 }
 
+void llm_graph_result::add_fused_node(llm_graph_fused_node result) {
+    fused_nodes.push_back(result);
+}
+
 void llm_graph_result::set_params(const llm_graph_params & params) {
     this->params = params;
 }
@@ -1364,6 +1357,8 @@ void llm_graph_context::cb(ggml_tensor * cur, const char * name, int il) const {
     }
 }
 
+
+
 ggml_tensor * llm_graph_context::build_cvec(
          ggml_tensor * cur,
                  int   il) const {
@@ -2414,7 +2409,7 @@ ggml_tensor * llm_graph_context::build_attn_mha(
 
         cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias,
                                   hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);
-        cb(cur, LLAMA_TENSOR_NAME_FATTN, il);
+        res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il});
 
         ggml_flash_attn_ext_add_sinks(cur, sinks);
         ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32);
@@ -2633,12 +2628,12 @@ ggml_tensor * llm_graph_context::build_attn(
     GGML_ASSERT(v_mla == nullptr);
 
     if (inp->self_k_rot) {
-        q_cur = ggml_mul_mat_aux(ctx0, q_cur, inp->self_k_rot);
-        k_cur = ggml_mul_mat_aux(ctx0, k_cur, inp->self_k_rot);
+        q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot);
+        k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot);
     }
 
     if (inp->self_v_rot) {
-        v_cur = ggml_mul_mat_aux(ctx0, v_cur, inp->self_v_rot);
+        v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot);
     }
 
     // these nodes are added to the graph together so that they are not reordered
@@ -2669,7 +2664,7 @@ ggml_tensor * llm_graph_context::build_attn(
     cb(cur, "kqv_out", il);
 
     if (inp->self_v_rot) {
-        cur = ggml_mul_mat_aux(ctx0, cur, inp->self_v_rot);
+        cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot);
     }
 
     if (wo) {
@@ -2874,14 +2869,14 @@ ggml_tensor * llm_graph_context::build_attn(
     auto * v_rot = is_swa ? inp->self_v_rot_swa : inp->self_v_rot;
 
     if (k_rot) {
-        q_cur = ggml_mul_mat_aux(ctx0, q_cur, k_rot);
+        q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot);
         if (k_cur) {
-            k_cur = ggml_mul_mat_aux(ctx0, k_cur, k_rot);
+            k_cur = llama_mul_mat_hadamard(ctx0, k_cur, k_rot);
         }
     }
     if (v_rot) {
         if (v_cur) {
-            v_cur = ggml_mul_mat_aux(ctx0, v_cur, v_rot);
+            v_cur = llama_mul_mat_hadamard(ctx0, v_cur, v_rot);
         }
     }
 
@@ -2924,7 +2919,7 @@ ggml_tensor * llm_graph_context::build_attn(
     cb(cur, "kqv_out", il);
 
     if (v_rot) {
-        cur = ggml_mul_mat_aux(ctx0, cur, v_rot);
+        cur = llama_mul_mat_hadamard(ctx0, cur, v_rot);
     }
 
     if (wo) {
@@ -3084,6 +3079,8 @@ llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const {
     dsv4_build_comp_inputs(ctx0, inp->inp_csa, mctx_cur->get_csa_plan(ubatch), "csa", n_stream);
     dsv4_build_comp_inputs(ctx0, inp->inp_hca, mctx_cur->get_hca_plan(ubatch), "hca", n_stream);
     dsv4_build_comp_inputs(ctx0, inp->inp_lid, mctx_cur->get_lid_plan(ubatch), "lid", n_stream);
+    inp->inp_csa.k_rot = mctx_cur->get_csa()->build_input_k_rot(ctx0);
+    inp->inp_hca.k_rot = mctx_cur->get_hca()->build_input_k_rot(ctx0);
     inp->inp_lid.k_rot = mctx_cur->get_lid()->build_input_k_rot(ctx0);
 
     return (llm_graph_input_dsv4 *) res->add_input(std::move(inp));

src/llama-graph.h

@@ -38,6 +38,12 @@ enum llm_graph_type {
     LLM_GRAPH_TYPE_DECODER_MTP,
 };
 
+enum llm_fused_op {
+    LLM_FUSED_OP_FLASH_ATTN,
+    LLM_FUSED_OP_GDN_AR,
+    LLM_FUSED_OP_GDN_CH,
+};
+
 enum llm_ffn_op_type : int {
     LLM_FFN_NONE = 0,           // sentinel: unset; archs must assign before use
     LLM_FFN_SILU,
@@ -775,6 +781,12 @@ struct llm_graph_params {
     }
 };
 
+struct llm_graph_fused_node {
+    llm_fused_op op;
+    ggml_tensor * tensor;
+    int il;
+};
+
 class llm_graph_result {
 public:
     llm_graph_result(int64_t max_nodes);
@@ -808,6 +820,10 @@ public:
 
     llm_graph_input_i * add_input(llm_graph_input_ptr input);
 
+    void add_fused_node(llm_graph_fused_node result);
+
+    const std::vector<llm_graph_fused_node> & get_fused_nodes() const { return fused_nodes; }
+
     void set_params(const llm_graph_params & params);
 
     // important graph nodes
@@ -826,6 +842,7 @@ public:
     std::map<llama_seq_id, ggml_tensor *> t_sampled_probs;
 
     std::vector<llm_graph_input_ptr> inputs;
+    std::vector<llm_graph_fused_node> fused_nodes;
 
     ggml_context_ptr ctx_compute;
 

src/llama-impl.h

@@ -54,6 +54,26 @@ static inline dst_t llama_cast(src_t v) {
     }
 }
 
+static inline ggml_tensor * llama_mul_mat_hadamard(
+        ggml_context * ctx,
+        ggml_tensor * cur,
+        ggml_tensor * rot) {
+    const auto n = rot->ne[0];
+
+    ggml_tensor * res;
+
+    if (!ggml_is_contiguous(cur)) {
+        res = ggml_cont_2d(ctx, cur, n, ggml_nelements(cur)/n);
+    } else {
+        res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n);
+    }
+    res = ggml_mul_mat(ctx, rot, res);
+    ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD);
+    res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]);
+
+    return res;
+}
+
 struct time_meas {
     time_meas(int64_t & t_acc, bool disable = false);
     ~time_meas();
@@ -83,7 +103,3 @@ std::string llama_format_tensor_shape(const std::vector<int64_t> & ne);
 std::string llama_format_tensor_shape(const struct ggml_tensor * t);
 
 std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i);
-
-#define LLAMA_TENSOR_NAME_FATTN   "__fattn__"
-#define LLAMA_TENSOR_NAME_FGDN_AR "__fgdn_ar__"
-#define LLAMA_TENSOR_NAME_FGDN_CH "__fgdn_ch__"

src/llama-kv-cache-dsa.cpp

@@ -113,7 +113,7 @@ llama_memory_context_ptr llama_kv_cache_dsa::init_batch(
 
         std::vector<llama_ubatch> ubatches;
         while (true) {
-            auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true);
+            auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0);
 
             if (ubatch.n_tokens == 0) {
                 break;

src/llama-kv-cache-dsv4.cpp

@@ -1110,7 +1110,7 @@ llama_memory_context_ptr llama_kv_cache_dsv4::init_batch(
             if (has_coupled) {
                 ubatch = balloc.split_seq(n_ubatch);
             } else {
-                ubatch = balloc.split_equal(n_ubatch, raw_per_seq || comp_per_seq);
+                ubatch = balloc.split_equal(n_ubatch, raw_per_seq || comp_per_seq, 0);
             }
 
             if (ubatch.n_tokens == 0) {

src/llama-kv-cache-iswa.cpp

@@ -206,7 +206,7 @@ llama_memory_context_ptr llama_kv_cache_iswa::init_batch(llama_batch_allocr & ba
 
         std::vector<llama_ubatch> ubatches;
         while (true) {
-            auto ubatch = balloc.split_equal(n_ubatch, !unified);
+            auto ubatch = balloc.split_equal(n_ubatch, !unified, 0);
 
             if (ubatch.n_tokens == 0) {
                 break;

src/llama-kv-cache.cpp

@@ -57,22 +57,6 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) {
     }
 }
 
-static ggml_tensor * ggml_mul_mat_aux(
-        ggml_context * ctx,
-        ggml_tensor * cur,
-        ggml_tensor * rot) {
-    const auto n = rot->ne[0];
-
-    ggml_tensor * res;
-
-    res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n);
-    res = ggml_mul_mat   (ctx, rot, res);
-    ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD);
-    res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]);
-
-    return res;
-}
-
 //
 // llama_kv_cache
 //
@@ -722,7 +706,7 @@ llama_memory_context_ptr llama_kv_cache::init_batch(
 
         std::vector<llama_ubatch> ubatches;
         while (true) {
-            auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true);
+            auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0);
 
             if (ubatch.n_tokens == 0) {
                 break;
@@ -1875,14 +1859,14 @@ ggml_tensor * llama_kv_cache::build_rope_shift(
         tmp = ggml_cast(ctx, cur, GGML_TYPE_F32);
 
         // rotate back
-        tmp = ggml_mul_mat_aux(ctx, tmp, rot);
+        tmp = llama_mul_mat_hadamard(ctx, tmp, rot);
 
         tmp = ggml_rope_ext(ctx, tmp,
                 shift, factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
                 yarn_ext_factor, yarn_attn_factor, yarn_beta_fast, yarn_beta_slow);
 
         // rotate fwd
-        tmp = ggml_mul_mat_aux(ctx, tmp, rot);
+        tmp = llama_mul_mat_hadamard(ctx, tmp, rot);
 
         tmp = ggml_cpy(ctx, tmp, cur);
     } else {

src/llama-memory-hybrid-iswa.cpp

@@ -77,15 +77,15 @@ llama_memory_context_ptr llama_memory_hybrid_iswa::init_batch(llama_batch_allocr
                 // if all tokens are output, split by sequence
                 ubatch = balloc.split_seq(n_ubatch);
             } else {
-                if (mem_recr->n_rs_seq > 0) {
-                    // [TAG_RECURRENT_ROLLBACK_SPLITS]
-                    // TODO: recurrent state rollback does not support equal splits
-                    ubatch = balloc.split_seq(n_ubatch);
-                } else {
-                    // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice)
-                    const bool unified = (mem_attn->get_base()->get_n_stream() == 1);
-                    ubatch = balloc.split_equal(n_ubatch, !unified);
-                }
+                // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice)
+                const bool unified = (mem_attn->get_base()->get_n_stream() == 1);
+
+                // [TAG_RECURRENT_ROLLBACK_SPLITS]
+                // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch
+                //   so that the rollback snapshots remain valid
+                const uint32_t n_rs_seq = mem_recr->n_rs_seq;
+
+                ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0);
             }
 
             if (ubatch.n_tokens == 0) {

src/llama-memory-hybrid.cpp

@@ -78,15 +78,15 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba
                 // if all tokens are output, split by sequence
                 ubatch = balloc.split_seq(n_ubatch);
             } else {
-                if (mem_recr->n_rs_seq > 0) {
-                    // [TAG_RECURRENT_ROLLBACK_SPLITS]
-                    // TODO: recurrent state rollback does not support equal splits
-                    ubatch = balloc.split_seq(n_ubatch);
-                } else {
-                    // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice)
-                    const bool unified = (mem_attn->get_n_stream() == 1);
-                    ubatch = balloc.split_equal(n_ubatch, !unified);
-                }
+                // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice)
+                const bool unified = (mem_attn->get_n_stream() == 1);
+
+                // [TAG_RECURRENT_ROLLBACK_SPLITS]
+                // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch
+                //   so that the rollback snapshots remain valid
+                const uint32_t n_rs_seq = mem_recr->n_rs_seq;
+
+                ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0);
             }
 
             if (ubatch.n_tokens == 0) {

src/llama-memory-recurrent.cpp

@@ -416,15 +416,12 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr &
                 // if all tokens are output, split by sequence
                 ubatch = balloc.split_seq(n_ubatch);
             } else {
-                if (n_rs_seq > 0) {
-                    // [TAG_RECURRENT_ROLLBACK_SPLITS]
-                    // TODO: recurrent state rollback does not support equal splits
-                    ubatch = balloc.split_seq(n_ubatch);
-                } else {
-                    // TODO: non-sequential equal split can be done if using unified KV cache
-                    //       for simplicity, we always use sequential equal split for now
-                    ubatch = balloc.split_equal(n_ubatch, true);
-                }
+                // TODO: non-sequential equal split can be done if using unified KV cache
+                //       for simplicity, we always use sequential equal split for now
+                // [TAG_RECURRENT_ROLLBACK_SPLITS]
+                // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch
+                //   so that the rollback snapshots remain valid
+                ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0);
             }
 
             if (ubatch.n_tokens == 0) {

src/llama-model-loader.cpp

@@ -37,6 +37,7 @@ const char * llama_ftype_name(llama_ftype ftype) {
         case LLAMA_FTYPE_MOSTLY_F16:       name = LLAMA_FTYPE_PREFIX "F16"; break;
         case LLAMA_FTYPE_MOSTLY_BF16:      name = LLAMA_FTYPE_PREFIX "BF16"; break;
         case LLAMA_FTYPE_MOSTLY_Q1_0:      name = LLAMA_FTYPE_PREFIX "Q1_0"; break;
+        case LLAMA_FTYPE_MOSTLY_Q2_0:      name = LLAMA_FTYPE_PREFIX "Q2_0"; break;
         case LLAMA_FTYPE_MOSTLY_Q4_0:      name = LLAMA_FTYPE_PREFIX "Q4_0"; break;
         case LLAMA_FTYPE_MOSTLY_Q4_1:      name = LLAMA_FTYPE_PREFIX "Q4_1"; break;
         case LLAMA_FTYPE_MOSTLY_Q5_0:      name = LLAMA_FTYPE_PREFIX "Q5_0"; break;
@@ -767,6 +768,7 @@ llama_model_loader::llama_model_loader(
             case GGML_TYPE_IQ3_S:   ftype = LLAMA_FTYPE_MOSTLY_IQ3_S;   break;
             case GGML_TYPE_NVFP4:   ftype = LLAMA_FTYPE_MOSTLY_NVFP4;   break;
             case GGML_TYPE_Q1_0:    ftype = LLAMA_FTYPE_MOSTLY_Q1_0;    break;
+            case GGML_TYPE_Q2_0:    ftype = LLAMA_FTYPE_MOSTLY_Q2_0;    break;
             default:
                 {
                     LLAMA_LOG_WARN("%s: unknown type %s\n", __func__, ggml_type_name(type_max));

src/llama-model.cpp

@@ -953,6 +953,8 @@ static buft_list_t make_gpu_buft_list(ggml_backend_dev_t dev, llama_split_mode s
             if (buft != nullptr) {
                 buft_list.emplace_back(dev, buft);
             }
+        } else {
+            throw std::runtime_error(format("device %s does not support split buffers", ggml_backend_dev_name(dev)));
         }
     }
 

src/llama-quant.cpp

@@ -380,6 +380,7 @@ static ggml_type tensor_type_fallback(quantize_state_impl & qs, const ggml_tenso
             case GGML_TYPE_IQ3_XXS:
             case GGML_TYPE_IQ3_S:   // types on the right: block size 32
             case GGML_TYPE_IQ4_XS:  return_type = GGML_TYPE_IQ4_NL; break;
+            case GGML_TYPE_Q2_0:
             case GGML_TYPE_Q2_K:
             case GGML_TYPE_Q3_K:
             case GGML_TYPE_TQ1_0:
@@ -480,7 +481,7 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type
             else if (ftype == LLAMA_FTYPE_MOSTLY_IQ3_XXS) {
                 new_type = GGML_TYPE_IQ3_S;
             }
-            else if (ftype == LLAMA_FTYPE_MOSTLY_TQ1_0 || ftype == LLAMA_FTYPE_MOSTLY_TQ2_0) {
+            else if (ftype == LLAMA_FTYPE_MOSTLY_TQ1_0 || ftype == LLAMA_FTYPE_MOSTLY_TQ2_0 || ftype == LLAMA_FTYPE_MOSTLY_Q2_0) {
                 new_type = GGML_TYPE_Q4_K;
             }
         }
@@ -800,6 +801,7 @@ ggml_type llama_ftype_get_default_type(llama_ftype ftype) {
         case LLAMA_FTYPE_MOSTLY_BF16: return GGML_TYPE_BF16;
         case LLAMA_FTYPE_ALL_F32:     return GGML_TYPE_F32;
         case LLAMA_FTYPE_MOSTLY_Q1_0: return GGML_TYPE_Q1_0;
+        case LLAMA_FTYPE_MOSTLY_Q2_0: return GGML_TYPE_Q2_0;
 
         case LLAMA_FTYPE_MOSTLY_MXFP4_MOE: return GGML_TYPE_MXFP4;
 

src/llama-vocab.cpp

@@ -887,9 +887,6 @@ struct llm_tokenizer_ugm : llm_tokenizer {
             // blob containing XOR-compressed compact double array (XCDA) entries
             uint32_t xcda_blob_size = *(const uint32_t *) &precompiled_charsmap[0];
             charsmap_offset += sizeof(xcda_blob_size);
-            if (xcda_blob_size + charsmap_offset >= precompiled_charsmap.size()) {
-                throw std::runtime_error("Index out of array bounds in precompiled charsmap!");
-            }
 
             // Next xcda_blob_size bytes contain entries of XOR-compressed compact
             // double array (XCDA). Each entry is bit-packed into a 32-bit integer.
@@ -1205,7 +1202,15 @@ private:
                 throw std::runtime_error("Index out of array bounds in precompiled charsmap!");
             }
             const char * prefix_replacement = &(tokenizer.prefix_replacements)[longest_prefix_offset];
-            return { prefix_replacement, strlen(prefix_replacement), longest_prefix_length };
+            size_t max_len = tokenizer.prefix_replacements_size - longest_prefix_offset;
+            size_t repl_len = 0;
+            while (repl_len < max_len && prefix_replacement[repl_len] != '\0') {
+                repl_len++;
+            }
+            if (repl_len == max_len) {
+                throw std::runtime_error("Unterminated string in precompiled charsmap!");
+            }
+            return { prefix_replacement, repl_len, longest_prefix_length };
         }
 
         // check if the input prefix contains a valid sequence of UTF-8 code units
@@ -2018,11 +2023,18 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                 const size_t n_precompiled_charsmap = gguf_get_arr_n(ctx, precompiled_charsmap_keyidx);
                 const char * pc = (const char *) gguf_get_arr_data(ctx, precompiled_charsmap_keyidx);
                 precompiled_charsmap.assign(pc, pc + n_precompiled_charsmap);
+                if (precompiled_charsmap.size() < sizeof(uint32_t)) {
+                    throw std::runtime_error("precompiled_charsmap too small for xcda_blob_size header!");
+                }
+                uint32_t * xcda_blob_size = (uint32_t *) &precompiled_charsmap[0];
+#if defined(__BYTE_ORDER__) && defined(__ORDER_BIG_ENDIAN__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+                *xcda_blob_size = __builtin_bswap32(*xcda_blob_size);
+#endif
+                if (*xcda_blob_size + sizeof(uint32_t) >= precompiled_charsmap.size()) {
+                    throw std::runtime_error("Index out of array bounds in precompiled charsmap!");
+                }
 #if defined(__BYTE_ORDER__) && defined(__ORDER_BIG_ENDIAN__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
                 // correct endianness of data in precompiled_charsmap binary blob
-                uint32_t * xcda_blob_size = (uint32_t *) &precompiled_charsmap[0];
-                *xcda_blob_size = __builtin_bswap32(*xcda_blob_size);
-                assert(*xcda_blob_size + sizeof(uint32_t) < n_precompiled_charsmap);
                 size_t xcda_array_size = *xcda_blob_size / sizeof(uint32_t);
                 uint32_t * xcda_array = (uint32_t *) &precompiled_charsmap[sizeof(uint32_t)];
                 for (size_t i = 0; i < xcda_array_size; ++i) {