hexagon: new vtcm layouts and improved pipelines for MUL_MAT, MUL_MAT_ID and FLASH_ATTN_EXT (#25425)

* 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

@@ -719,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");
         }
     }
@@ -1241,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"),

common/common.h

@@ -644,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;
+}

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

@@ -4709,10 +4709,16 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
             } break;
         case GGML_OP_SET_ROWS:
             {
-                return (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 ||
-                       op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 ||
-                       op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL) &&
-                       op->src[0]->type == GGML_TYPE_F32 &&
+                return (
+                           (
+                               (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 ||
+                               op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 ||
+                               op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL) &&
+                               op->src[0]->type == GGML_TYPE_F32
+                           ) || (
+                               op->type == GGML_TYPE_F16 && op->src[0]->type == GGML_TYPE_F16
+                           )
+                       ) &&
                        (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32);
             } break;
         case GGML_OP_SET:

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);
 

tools/cli/CMakeLists.txt

@@ -2,11 +2,13 @@
 
 set(TARGET llama-cli-impl)
 
-add_library(${TARGET} cli.cpp)
+add_library(${TARGET} cli.cpp
+                      cli-client.cpp
+                      cli-context.cpp)
 set_target_properties(${TARGET} PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS ON)
 
 target_include_directories(${TARGET} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} ../server)
-target_link_libraries(${TARGET} PUBLIC server-context llama-common ${CMAKE_THREAD_LIBS_INIT})
+target_link_libraries(${TARGET} PUBLIC llama-server-impl llama-common ${CMAKE_THREAD_LIBS_INIT})
 
 if(LLAMA_TOOLS_INSTALL)
     install(TARGETS ${TARGET} LIBRARY)

tools/cli/cli-client.cpp

@@ -0,0 +1,130 @@
+#include "cli-client.h"
+
+#include "http.h"
+
+#include <algorithm>
+#include <chrono>
+#include <thread>
+
+// generation can stall for a long time during prompt processing, so the
+// read timeout must be generous
+static constexpr time_t CLI_HTTP_READ_TIMEOUT_SEC = 3600;
+
+// upper bound for the accumulated response body kept for error reporting
+static constexpr size_t CLI_HTTP_MAX_ERROR_BODY = 1024 * 1024;
+
+// returns the path with the base url's path prefix prepended (if any)
+static std::string join_path(const common_http_url & parts, const std::string & path) {
+    if (parts.path.empty() || parts.path == "/") {
+        return path;
+    }
+    std::string prefix = parts.path;
+    if (prefix.back() == '/') {
+        prefix.pop_back();
+    }
+    return prefix + path;
+}
+
+std::string cli_client::get(const std::string & path) {
+    auto [cli, parts] = common_http_client(server_base);
+    cli.set_read_timeout(CLI_HTTP_READ_TIMEOUT_SEC, 0);
+    auto path_with_model = path + (model.empty() ? "" : ("?model=" + model));
+    auto res = cli.Get(join_path(parts, path_with_model));
+    if (!res) {
+        throw std::runtime_error("failed to connect to " + server_base + ": " + httplib::to_string(res.error()));
+    }
+    if (res->status < 200 || res->status >= 300) {
+        throw std::runtime_error("GET " + path + " failed with status " + std::to_string(res->status) + ": " + res->body);
+    }
+    return res->body;
+}
+
+std::string cli_client::post(const std::string & path, const std::string & body) {
+    auto [cli, parts] = common_http_client(server_base);
+    cli.set_read_timeout(CLI_HTTP_READ_TIMEOUT_SEC, 0);
+    auto res = cli.Post(join_path(parts, path), body, "application/json");
+    if (!res) {
+        throw std::runtime_error("failed to connect to " + server_base + ": " + httplib::to_string(res.error()));
+    }
+    if (res->status < 200 || res->status >= 300) {
+        throw std::runtime_error("POST " + path + " failed with status " + std::to_string(res->status) + ": " + res->body);
+    }
+    return res->body;
+}
+
+std::string cli_client::post_sse(const std::string & path,
+                                  const std::string & body,
+                                  const std::function<bool()> & should_stop,
+                                  const std::function<void(const std::string &)> & on_data) {
+    auto [cli, parts] = common_http_client(server_base);
+    cli.set_read_timeout(CLI_HTTP_READ_TIMEOUT_SEC, 0);
+
+    std::string pending;  // buffer for incomplete SSE lines
+    std::string raw_body; // accumulated body, used only for error reporting
+
+    auto receiver = [&](const char * data, size_t len) -> bool {
+        if (should_stop()) {
+            return false; // aborts the request
+        }
+        if (raw_body.size() < CLI_HTTP_MAX_ERROR_BODY) {
+            raw_body.append(data, std::min(len, CLI_HTTP_MAX_ERROR_BODY - raw_body.size()));
+        }
+        pending.append(data, len);
+        size_t pos;
+        while ((pos = pending.find('\n')) != std::string::npos) {
+            std::string line = pending.substr(0, pos);
+            pending.erase(0, pos + 1);
+            if (!line.empty() && line.back() == '\r') {
+                line.pop_back();
+            }
+            if (line.rfind("data: ", 0) != 0) {
+                continue;
+            }
+            std::string payload = line.substr(6);
+            if (payload == "[DONE]") {
+                continue;
+            }
+            on_data(payload);
+        }
+        return true;
+    };
+
+    httplib::Headers headers = {{"Accept", "text/event-stream"}};
+    auto res = cli.Post(join_path(parts, path), headers, body, "application/json", receiver);
+
+    if (!res) {
+        if (res.error() == httplib::Error::Canceled && should_stop()) {
+            return ""; // cancelled by the user
+        }
+        return "failed to connect to " + server_base + ": " + httplib::to_string(res.error());
+    }
+    if (res->status < 200 || res->status >= 300) {
+        if (!raw_body.empty()) {
+            return raw_body;
+        }
+        return "request failed with status " + std::to_string(res->status);
+    }
+    return "";
+}
+
+bool cli_client::wait_health(const std::function<bool()> & is_aborted) {
+    int connect_attempts = 0;
+    while (!is_aborted()) {
+        auto [cli, parts] = common_http_client(server_base);
+        cli.set_connection_timeout(1, 0);
+        auto res = cli.Get(join_path(parts, "/health"));
+        if (res) {
+            if (res->status == 200) {
+                return true;
+            }
+            // any other status means the server is up but not ready yet
+            // (e.g. 503 while the model is still loading)
+        } else if (++connect_attempts >= 10) {
+            last_error = "failed to connect to " + server_base + ": " + httplib::to_string(res.error());
+            return false;
+        }
+        std::this_thread::sleep_for(std::chrono::milliseconds(300));
+    }
+    last_error = "aborted while waiting for the server to become ready";
+    return false;
+}

tools/cli/cli-client.h

@@ -0,0 +1,33 @@
+#pragma once
+
+#include <functional>
+#include <string>
+
+// openai-like client for CLI
+struct cli_client {
+    std::string server_base; // base url, for example "http://127.0.0.1:8080"
+    std::string last_error;  // set when wait_health() fails
+
+    std::string model; // optional, set when the server has multiple models (router mode)
+
+    // simple GET request, returns the raw response body
+    // throws std::runtime_error on transport error or non-2xx status
+    std::string get(const std::string & path);
+
+    // simple POST request, returns the raw response body
+    // throws std::runtime_error on transport error or non-2xx status
+    std::string post(const std::string & path, const std::string & body);
+
+    // POST request with an SSE streaming response
+    // on_data is invoked per "data:" event with the raw event payload
+    // returns after the stream is finished (empty string on graceful exit)
+    // otherwise, the raw error response body
+    std::string post_sse(const std::string & path,
+                          const std::string & body,
+                          const std::function<bool()> & should_stop,
+                          const std::function<void(const std::string &)> & on_data);
+
+    // poll /health until the server is ready to accept requests
+    // returns false if is_aborted returned true or the server is unreachable
+    bool wait_health(const std::function<bool()> & is_aborted);
+};

tools/cli/cli-context.cpp

@@ -0,0 +1,622 @@
+#include "cli-context.h"
+#include "cli-ui.h"
+
+#include "arg.h"
+#include "base64.hpp"
+#include "log.h"
+#include "console.h"
+
+#define JSON_ASSERT GGML_ASSERT
+#include <nlohmann/json.hpp>
+
+#include <algorithm>
+#include <cctype>
+#include <filesystem>
+#include <fstream>
+#include <map>
+#include <set>
+
+using json = nlohmann::ordered_json;
+
+struct cli_context_impl {
+    json messages      = json::array();
+    json pending_media = json::array(); // staged multimodal content parts
+};
+
+cli_context::cli_context(const common_params & params) : params(params), impl(new cli_context_impl()) {}
+
+cli_context::~cli_context() {
+    shutdown();
+}
+
+std::atomic<bool> & cli_context::interrupted() {
+    static std::atomic<bool> flag = false;
+    return flag;
+}
+
+static bool should_stop() {
+    return cli_context::interrupted().load();
+}
+
+static constexpr size_t FILE_GLOB_MAX_RESULTS = 100;
+
+const char * LLAMA_ASCII_LOGO = R"(
+▄▄ ▄▄
+██ ██
+██ ██  ▀▀█▄ ███▄███▄  ▀▀█▄    ▄████ ████▄ ████▄
+██ ██ ▄█▀██ ██ ██ ██ ▄█▀██    ██    ██ ██ ██ ██
+██ ██ ▀█▄██ ██ ██ ██ ▀█▄██ ██ ▀████ ████▀ ████▀
+                                    ██    ██
+                                    ▀▀    ▀▀
+)";
+
+// number of values an arg consumes on the command line
+static int arg_num_values(const common_arg & opt) {
+    if (opt.value_hint_2 != nullptr) {
+        return 2;
+    }
+    if (opt.value_hint != nullptr) {
+        return 1;
+    }
+    return 0;
+}
+
+static std::string format_error_message(const json & err) {
+    if (err.contains("error") && err.at("error").is_object()) {
+        const auto & e = err.at("error");
+        if (e.contains("message") && e.at("message").is_string()) {
+            return e.at("message").get<std::string>();
+        }
+    }
+    return err.dump();
+}
+
+// err is the raw response body of a failed request; it may or may not be JSON
+static std::string format_error_message(const std::string & err) {
+    json parsed = json::parse(err, nullptr, false);
+    if (!parsed.is_discarded()) {
+        return format_error_message(parsed);
+    }
+    return err;
+}
+
+static std::string media_type_from_ext(const std::string & fname) {
+    std::string ext = std::filesystem::path(fname).extension().string();
+    std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { return std::tolower(c); });
+    if (ext == ".wav" || ext == ".mp3") {
+        return "audio";
+    }
+    if (ext == ".mp4" || ext == ".avi" || ext == ".mkv" || ext == ".mov" || ext == ".webm") {
+        return "video";
+    }
+    return "image";
+}
+
+bool cli_context::init() {
+    ui::init(params);
+
+    std::optional<ui::spinner> spinner;
+
+    bool use_external_server = !params.server_base.empty();
+    if (use_external_server) {
+        std::string base = params.server_base;
+        while (!base.empty() && base.back() == '/') {
+            base.pop_back();
+        }
+        client.server_base = base;
+
+        spinner.emplace("Connecting to server at " + base);
+    } else {
+        if (params.model.path.empty() && params.model.url.empty() &&
+                params.model.hf_repo.empty() && params.model.docker_repo.empty()) {
+            ui::show_error(
+                "no model specified",
+                "use -m <file.gguf> or -hf <user/repo> to run a local model,\n"
+                "or --server-base <url> to connect to a running llama-server"
+            );
+            return false;
+        }
+
+        spinner.emplace("\n\nLoading model...");
+
+        server.emplace();
+        if (!server->start(params)) {
+            ui::show_error("server start failed");
+            return false;
+        }
+        if (!server->wait_ready(should_stop)) {
+            if (!should_stop()) {
+                ui::show_error("the server exited before becoming ready");
+            }
+            return false;
+        }
+        client.server_base = server->address();
+    }
+
+    // for --server-base this is the main availability check; for a spawned
+    // server it is a cheap sanity check on top of the ready signal
+    auto is_aborted = [this]() {
+        return should_stop() || (server && !server->alive());
+    };
+    bool healthy = false;
+    try {
+        healthy = client.wait_health(is_aborted);
+    } catch (const std::exception & e) {
+        client.last_error = e.what();
+    }
+    if (!healthy) {
+        if (!should_stop()) {
+            ui::show_error(client.last_error);
+        }
+        return false;
+    }
+
+    if (use_external_server) {
+        spinner.reset();
+        if (!list_and_ask_models()) {
+            return false;
+        }
+        // restore the spinner for the next step
+        spinner.emplace("Waiting for server...");
+    }
+
+    fetch_server_props();
+
+    return true;
+}
+
+void cli_context::fetch_server_props() {
+    try {
+        json props = json::parse(client.get("/props"));
+        model_name = props.value("model_alias", "");
+        if (model_name.empty()) {
+            const std::string path = props.value("model_path", "");
+            if (!path.empty()) {
+                model_name = std::filesystem::path(path).filename().string();
+            }
+        }
+        model_ftype = props.value("model_ftype", "");
+        build_info = props.value("build_info", "");
+        if (props.contains("modalities") && props.at("modalities").is_object()) {
+            const auto & modalities = props.at("modalities");
+            has_vision = modalities.value("vision", false);
+            has_audio  = modalities.value("audio", false);
+            has_video  = modalities.value("video", false);
+        }
+    } catch (const std::exception & e) {
+        // /props can be disabled on remote servers; not fatal
+        LOG_DBG("failed to fetch /props: %s\n", e.what());
+    }
+}
+
+bool cli_context::list_and_ask_models() {
+    json resp = json::parse(client.get("/v1/models"));
+    if (!resp.contains("data") || !resp.at("data").is_array()) {
+        throw std::runtime_error("invalid response from /v1/models");
+    }
+    std::vector<std::string> models;
+    std::vector<std::string> models_display;
+    for (const auto & m : resp.at("data")) {
+        if (!m.contains("id") || !m.at("id").is_string()) {
+            continue;
+        }
+        std::string name = m.at("id").get<std::string>();
+        std::string display = name;
+        if (m.contains("aliases") && m.at("aliases").is_array()) {
+            std::vector<std::string> aliases;
+            for (const auto & a : m.at("aliases")) {
+                if (a.is_string()) {
+                    aliases.push_back(a.get<std::string>());
+                }
+            }
+            if (!aliases.empty()) {
+                display += " (" + string_join(aliases, ", ") + ")";
+            }
+        }
+        models.push_back(name);
+        models_display.push_back(display);
+    }
+
+    // only one model: use it without asking
+    if (models.size() == 1) {
+        model_name = models[0];
+        client.model = model_name;
+        return true;
+    }
+
+    std::string message = "\nAvailable models:";
+    for (size_t i = 0; i < models_display.size(); ++i) {
+        message += "\n  " + std::to_string(i + 1) + ". " + models_display[i];
+    }
+    message += "\n";
+    ui::show_message(message);
+    std::string selection;
+    while (selection.empty()) {
+        if (should_stop()) {
+            return false;
+        }
+        ui::user_turn user_turn;
+        selection = user_turn.read_input(false, "Select model by number: ");
+        if (selection.empty()) {
+            continue;
+        }
+        try {
+            size_t idx = std::stoul(selection);
+            if (idx > 0 && idx <= models.size()) {
+                model_name = models[idx - 1];
+                client.model = model_name;
+                ui::show_message("Selected model: " + model_name);
+                break;
+            }
+        } catch (...) {
+            // ignore
+        }
+        ui::show_error("Invalid selection. Please enter a valid number.");
+        selection.clear();
+        continue;
+    }
+    return true;
+}
+
+void cli_context::add_system_prompt() {
+    if (!params.system_prompt.empty()) {
+        impl->messages.push_back({
+            {"role",    "system"},
+            {"content", params.system_prompt}
+        });
+    }
+}
+
+void cli_context::push_user_message(const std::string & text) {
+    json content;
+    if (impl->pending_media.empty()) {
+        content = text;
+    } else {
+        // multimodal message: media parts first, then the text
+        content = impl->pending_media;
+        content.push_back({
+            {"type", "text"},
+            {"text", text}
+        });
+        impl->pending_media = json::array();
+    }
+    impl->messages.push_back({
+        {"role",    "user"},
+        {"content", content}
+    });
+}
+
+bool cli_context::stage_media_file(const std::string & fname, const std::string & type) {
+    std::ifstream file(fname, std::ios::binary);
+    if (!file) {
+        return false;
+    }
+    std::string data((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
+    std::string encoded = base64::encode(data);
+
+    if (type == "audio") {
+        std::string ext = std::filesystem::path(fname).extension().string();
+        std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { return std::tolower(c); });
+        impl->pending_media.push_back({
+            {"type", "input_audio"},
+            {"input_audio", {
+                {"data",   encoded},
+                {"format", ext == ".mp3" ? "mp3" : "wav"}
+            }}
+        });
+    } else if (type == "video") {
+        impl->pending_media.push_back({
+            {"type", "input_video"},
+            {"input_video", {
+                {"data", encoded}
+            }}
+        });
+    } else {
+        // the server detects the actual image type from the data
+        impl->pending_media.push_back({
+            {"type", "image_url"},
+            {"image_url", {
+                {"url", "data:image/unknown;base64," + encoded}
+            }}
+        });
+    }
+    return true;
+}
+
+bool cli_context::generate_completion(std::string & assistant_content, cli_timings & timings) {
+    json body = {
+        {"messages",          impl->messages},
+        {"stream",            true},
+        // in order to get timings even when we cancel mid-way
+        {"timings_per_token", true},
+    };
+    if (!client.model.empty()) {
+        body["model"] = client.model;
+    }
+
+    bool stream_error = false;
+
+    ui::assistant_turn a;
+
+    std::string err = client.post_sse("/v1/chat/completions", body.dump(), should_stop, [&](const std::string & payload) {
+        json chunk = json::parse(payload, nullptr, false);
+        if (chunk.is_discarded()) {
+            return;
+        }
+        if (chunk.contains("error")) {
+            stream_error = true;
+            ui::show_error(format_error_message(chunk));
+            return;
+        }
+        if (chunk.contains("timings")) {
+            const auto & t = chunk.at("timings");
+            timings.prompt_per_second    = t.value("prompt_per_second",    0.0);
+            timings.predicted_per_second = t.value("predicted_per_second", 0.0);
+        }
+        if (!chunk.contains("choices") || !chunk.at("choices").is_array() || chunk.at("choices").empty()) {
+            return;
+        }
+        const auto & choice = chunk.at("choices").at(0);
+        if (!choice.contains("delta")) {
+            return;
+        }
+        const auto & delta = choice.at("delta");
+        if (delta.contains("reasoning_content") && delta.at("reasoning_content").is_string()) {
+            const std::string text = delta.at("reasoning_content").get<std::string>();
+            if (!text.empty()) {
+                a.push(ui::ASSISTANT_DISPLAY_MODE_REASONING, text);
+            }
+        }
+        if (delta.contains("content") && delta.at("content").is_string()) {
+            const std::string text = delta.at("content").get<std::string>();
+            if (!text.empty()) {
+                assistant_content += text;
+                a.push(ui::ASSISTANT_DISPLAY_MODE_CONTENT, text);
+            }
+        }
+    });
+
+    cli_context::interrupted().store(false);
+
+    if (!err.empty()) {
+        ui::show_error(format_error_message(err));
+        return false;
+    }
+    return !stream_error;
+}
+
+int cli_context::run() {
+    add_system_prompt();
+
+    std::string modalities = "text";
+    if (has_vision) {
+        modalities += ", vision";
+    }
+    if (has_audio) {
+        modalities += ", audio";
+    }
+    if (has_video) {
+        modalities += ", video";
+    }
+
+    std::string banner;
+    banner += "\n";
+    banner += LLAMA_ASCII_LOGO;
+    banner += "\n";
+    banner += "build      : " + build_info + "\n";
+    banner += "model      : " + model_name + "\n";
+    if (!model_ftype.empty()) {
+        banner += "ftype      : " + model_ftype + "\n";
+    }
+    banner += "modalities : " + modalities + "\n";
+    if (!params.system_prompt.empty()) {
+        banner += "using custom system prompt\n";
+    }
+    banner += "\n";
+    banner += "available commands:\n";
+    banner += "  /exit or Ctrl+C     stop or exit\n";
+    banner += "  /regen              regenerate the last response\n";
+    banner += "  /clear              clear the chat history\n";
+    banner += "  /read <file>        add a text file\n";
+    banner += "  /glob <pattern>     add text files using globbing pattern\n";
+    if (has_vision) {
+        banner += "  /image <file>       add an image file\n";
+    }
+    if (has_audio) {
+        banner += "  /audio <file>       add an audio file\n";
+    }
+    if (has_video) {
+        banner += "  /video <file>       add a video file\n";
+    }
+    banner += "\n";
+
+    ui::show_message(banner);
+
+    // interactive loop
+    std::string cur_msg;
+
+    auto add_text_file = [&](const std::string & fname) -> bool {
+        std::ifstream file(fname, std::ios::binary);
+        if (!file) {
+            ui::show_error(string_format("file does not exist or cannot be opened: '%s'", fname.c_str()));
+            return false;
+        }
+        std::string content((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
+        cur_msg += "--- File: ";
+        cur_msg += fname;
+        cur_msg += " ---\n";
+        cur_msg += content;
+        ui::show_message(string_format("Loaded text from '%s'", fname.c_str()));
+        return true;
+    };
+
+    while (true) {
+        std::string buffer;
+        {
+            ui::user_turn user_turn;
+
+            if (params.prompt.empty()) {
+                buffer = user_turn.read_input(params.multiline_input);
+            } else {
+                // process input prompt from args
+                for (auto & fname : params.image) {
+                    if (!stage_media_file(fname, media_type_from_ext(fname))) {
+                        ui::show_error(string_format("file does not exist or cannot be opened: '%s'", fname.c_str()));
+                        break;
+                    }
+                    ui::show_message(string_format("Loaded media from '%s'", fname.c_str()));
+                }
+                buffer = params.prompt;
+                user_turn.echo(buffer);
+                params.prompt.clear(); // only use it once
+            }
+        }
+
+        if (should_stop()) {
+            cli_context::interrupted().store(false);
+            break;
+        }
+
+        // remove trailing newline
+        if (!buffer.empty() && buffer.back() == '\n') {
+            buffer.pop_back();
+        }
+
+        // skip empty messages
+        if (buffer.empty()) {
+            continue;
+        }
+
+        bool add_user_msg = true;
+
+        // process commands
+        if (string_starts_with(buffer, "/exit")) {
+            break;
+        } else if (string_starts_with(buffer, "/regen")) {
+            if (impl->messages.size() >= 2) {
+                size_t last_idx = impl->messages.size() - 1;
+                impl->messages.erase(last_idx);
+                add_user_msg = false;
+            } else {
+                ui::show_error("No message to regenerate.");
+                continue;
+            }
+        } else if (string_starts_with(buffer, "/clear")) {
+            impl->messages.clear();
+            add_system_prompt();
+
+            impl->pending_media = json::array();
+            ui::show_message("Chat history cleared.");
+            continue;
+        } else if (
+                (string_starts_with(buffer, "/image ") && has_vision) ||
+                (string_starts_with(buffer, "/audio ") && has_audio) ||
+                (string_starts_with(buffer, "/video ") && has_video)) {
+            std::string type = buffer.substr(1, 5);
+            // just in case (bad copy-paste for example), we strip all trailing/leading spaces
+            std::string fname = string_strip(buffer.substr(7));
+            if (!stage_media_file(fname, type)) {
+                ui::show_error(string_format("file does not exist or cannot be opened: '%s'", fname.c_str()));
+                continue;
+            }
+            ui::show_message(string_format("Loaded media from '%s'", fname.c_str()));
+            continue;
+        } else if (string_starts_with(buffer, "/read ")) {
+            std::string fname = string_strip(buffer.substr(6));
+            add_text_file(fname);
+            continue;
+        } else if (string_starts_with(buffer, "/glob ")) {
+            std::error_code ec;
+            size_t count = 0;
+            auto curdir = std::filesystem::current_path();
+            std::string pattern = string_strip(buffer.substr(6));
+            std::filesystem::path rel_path;
+
+            auto startglob = pattern.find_first_of("![*?");
+            if (startglob != std::string::npos && startglob != 0) {
+                auto endpath = pattern.substr(0, startglob).find_last_of('/');
+                if (endpath != std::string::npos) {
+                    std::string rel_pattern = pattern.substr(0, endpath);
+#if !defined(_WIN32)
+                    if (string_starts_with(rel_pattern, '~')) {
+                        const char * home = std::getenv("HOME");
+                        if (home && home[0]) {
+                            rel_pattern = home + rel_pattern.substr(1);
+                        }
+                    }
+#endif
+                    rel_path = rel_pattern;
+                    pattern.erase(0, endpath + 1);
+                    curdir /= rel_path;
+                }
+            }
+
+            for (const auto & entry : std::filesystem::recursive_directory_iterator(curdir,
+                    std::filesystem::directory_options::skip_permission_denied, ec)) {
+                if (!entry.is_regular_file()) {
+                    continue;
+                }
+
+                std::string rel = std::filesystem::relative(entry.path(), curdir, ec).string();
+                if (ec) {
+                    ec.clear();
+                    continue;
+                }
+                std::replace(rel.begin(), rel.end(), '\\', '/');
+
+                if (!glob_match(pattern, rel)) {
+                    continue;
+                }
+
+                if (!add_text_file((rel_path / rel).string())) {
+                    continue;
+                }
+
+                if (++count >= FILE_GLOB_MAX_RESULTS) {
+                    ui::show_error(string_format("Maximum number of globbed files allowed (%zu) reached.", FILE_GLOB_MAX_RESULTS));
+                    break;
+                }
+            }
+            continue;
+        } else {
+            // not a command
+            cur_msg += buffer;
+        }
+
+        // generate response
+        if (add_user_msg) {
+            push_user_message(cur_msg);
+            cur_msg.clear();
+        }
+        cli_timings timings;
+        std::string assistant_content;
+        generate_completion(assistant_content, timings);
+        impl->messages.push_back({
+            {"role",    "assistant"},
+            {"content", assistant_content}
+        });
+
+        if (params.show_timings) {
+            ui::show_info(string_format(
+                "\n[ Prompt: %.1f t/s | Generation: %.1f t/s ]",
+                timings.prompt_per_second,
+                timings.predicted_per_second
+            ));
+        }
+
+        if (params.single_turn) {
+            break;
+        }
+    }
+
+    ui::show_message("\n\nExiting...");
+
+    return 0;
+}
+
+void cli_context::shutdown() {
+    if (server) {
+        server->stop();
+        server.reset();
+    }
+}

tools/cli/cli-context.h

@@ -0,0 +1,66 @@
+#pragma once
+
+#include "common.h"
+
+#include "cli-client.h"
+#include "cli-server.h"
+
+#include <atomic>
+#include <memory>
+#include <optional>
+#include <string>
+
+struct cli_timings {
+    double prompt_per_second    = 0.0;
+    double predicted_per_second = 0.0;
+};
+
+struct cli_context_impl;
+
+struct cli_context {
+    common_params params;
+
+    cli_client client;                // always initialized
+    std::optional<cli_server> server; // only set when no --server-base is given
+
+    // properties of the connected server
+    // will be populated by fetch_server_props()
+    std::string model_name;
+    std::string model_ftype;
+    std::string build_info;
+    bool has_vision = false;
+    bool has_audio  = false;
+    bool has_video  = false;
+
+    cli_context(const common_params & params);
+    ~cli_context();
+
+    // connect to --server-base or spawn a local llama-server child;
+    // argc/argv are needed to forward the server-relevant args to the child
+    bool init();
+
+    // run the interactive chat loop, returns the process exit code
+    int run();
+
+    // stop the local server child (if any)
+    void shutdown();
+
+    // set by the SIGINT handler; cleared once the interrupt has been handled
+    static std::atomic<bool> & interrupted();
+
+private:
+    bool generate_completion(std::string & assistant_content, cli_timings & timings);
+    void fetch_server_props();
+    void add_system_prompt();
+    void push_user_message(const std::string & text);
+
+    // check if server have multiple models (router mode)
+    // if yes, list them then ask; do nothing otherwise
+    bool list_and_ask_models();
+
+    // read a file and stage it as a multimodal content part; type is one of
+    // "image", "audio", "video"; returns false if the file cannot be read
+    bool stage_media_file(const std::string & fname, const std::string & type);
+
+    std::unique_ptr<cli_context_impl> impl;
+};

tools/cli/cli-server.h

@@ -0,0 +1,89 @@
+#pragma once
+
+#include <thread>
+
+#include "http.h"
+
+// llama_server will be available as a dynamic library symbol
+int llama_server(common_params & params, int argc, char ** argv);
+void llama_server_terminate();
+
+struct cli_server {
+    std::thread th;
+    int port = -1;
+    std::atomic<bool> is_alive = false;
+    std::atomic<bool> is_stopping = false;
+
+    ~cli_server() {
+        stop();
+    }
+
+    void stop() {
+        if (is_stopping.exchange(true)) {
+            return;
+        }
+        if (alive()) {
+            llama_server_terminate();
+        }
+        if (th.joinable()) {
+            th.join();
+        }
+    }
+
+    // spawn llama-server in a thread and interact with it via a random port
+    bool start(common_params & params) {
+        port = common_http_get_free_port();
+        if (port <= 0) {
+            fprintf(stderr, "failed to get a free port\n");
+            exit(1);
+        }
+
+        is_alive.store(true, std::memory_order_release);
+
+        common_params server_params = params; // copy
+        server_params.port = port;
+
+        th = std::thread([this, server_params]() mutable {
+            // argc / argv are only used in router mode, we can skip them for now
+            int res = llama_server(server_params, 0, nullptr);
+            if (res != 0) {
+                fprintf(stderr, "llama_server exited with code %d\n", res);
+            }
+            is_alive.store(false, std::memory_order_release);
+        });
+
+        return true;
+    }
+
+    std::string address() const {
+        return "http://127.0.0.1:" + std::to_string(port);
+    }
+
+    bool wait_ready(std::function<bool()> should_stop) {
+        if (!alive()) {
+            return false;
+        }
+        while (!should_stop()) {
+            auto [cli, parts] = common_http_client(address());
+            cli.set_connection_timeout(1, 0);
+            auto res = cli.Get("/health");
+            if (res) {
+                if (res->status == 200) {
+                    return true;
+                }
+                // any other status means the server is up but not ready yet
+                // (e.g. 503 while the model is still loading)
+            }
+            if (!alive()) {
+                // in case server die permanently
+                return false;
+            }
+            std::this_thread::sleep_for(std::chrono::milliseconds(200));
+        }
+        return true;
+    }
+
+    bool alive() const {
+        return is_alive.load(std::memory_order_acquire);
+    }
+};

tools/cli/cli-ui.h

@@ -0,0 +1,251 @@
+#pragma once
+
+#include "common.h"
+#include "console.h"
+
+#include <array>
+#include <algorithm>
+#include <cctype>
+#include <filesystem>
+#include <string_view>
+
+// TODO?: Make this reusable, enums, docs
+static const std::array<std::string_view, 8> cmds = {
+    "/audio ",
+    "/clear",
+    "/exit",
+    "/glob ",
+    "/image ",
+    "/read ",
+    "/regen",
+    "/video ",
+};
+
+static std::vector<std::pair<std::string, size_t>> auto_completion_callback(std::string_view line, size_t cursor_byte_pos) {
+    std::vector<std::pair<std::string, size_t>> matches;
+    std::string cmd;
+
+    if (line.length() > 1 && line.front() == '/' && !std::any_of(cmds.begin(), cmds.end(), [line](std::string_view prefix) {
+        return string_starts_with(line, prefix);
+    })) {
+        auto it = cmds.begin();
+
+        while ((it = std::find_if(it, cmds.end(), [line](std::string_view cmd_line) {
+            return string_starts_with(cmd_line, line);
+        })) != cmds.end()) {
+            matches.emplace_back(*it, it->length());
+            ++it;
+        }
+    } else {
+        auto it = std::find_if(cmds.begin(), cmds.end(), [line](std::string_view prefix) {
+            return prefix.back() == ' ' && string_starts_with(line, prefix);
+        });
+
+        if (it != cmds.end()) {
+            cmd = *it;
+        }
+    }
+
+    if (!cmd.empty() && cmd != "/glob " && line.length() >= cmd.length() && cursor_byte_pos >= cmd.length()) {
+        const std::string path_prefix  = std::string(line.substr(cmd.length(), cursor_byte_pos - cmd.length()));
+        const std::string path_postfix = std::string(line.substr(cursor_byte_pos));
+        auto cur_dir = std::filesystem::current_path();
+        std::string cur_dir_str = cur_dir.string();
+        std::string expanded_prefix = path_prefix;
+
+#if !defined(_WIN32)
+        if (string_starts_with(path_prefix, '~')) {
+            const char * home = std::getenv("HOME");
+            if (home && home[0]) {
+                expanded_prefix = home + path_prefix.substr(1);
+            }
+        }
+        if (string_starts_with(expanded_prefix, '/')) {
+#else
+        if (std::isalpha(static_cast<unsigned char>(expanded_prefix[0])) && expanded_prefix.find(':') == 1) {
+#endif
+            cur_dir = std::filesystem::path(expanded_prefix).parent_path();
+            cur_dir_str.clear();
+        } else if (!path_prefix.empty()) {
+            cur_dir /= std::filesystem::path(path_prefix).parent_path();
+        }
+
+        std::error_code ec;
+        for (const auto & entry : std::filesystem::directory_iterator(cur_dir, ec)) {
+            if (ec) {
+                break;
+            }
+            if (!entry.exists(ec)) {
+                ec.clear();
+                continue;
+            }
+
+            const std::string path_full = entry.path().string();
+            std::string path_entry = !cur_dir_str.empty() && string_starts_with(path_full, cur_dir_str) ? path_full.substr(cur_dir_str.length() + 1) : path_full;
+
+            if (entry.is_directory(ec)) {
+                path_entry.push_back(std::filesystem::path::preferred_separator);
+            }
+
+            if (expanded_prefix.empty() || string_starts_with(path_entry, expanded_prefix)) {
+                const std::string updated_line = cmd + path_entry;
+                matches.emplace_back(updated_line + path_postfix, updated_line.length());
+            }
+
+            if (ec) {
+                ec.clear();
+            }
+        }
+
+        if (matches.empty()) {
+            const std::string updated_line = cmd + path_prefix;
+            matches.emplace_back(updated_line + path_postfix, updated_line.length());
+        }
+
+        // Add the longest common prefix
+        if (!expanded_prefix.empty() && matches.size() > 1) {
+            const std::string_view match0(matches[0].first);
+            const std::string_view match1(matches[1].first);
+            auto it = std::mismatch(match0.begin(), match0.end(), match1.begin(), match1.end());
+            size_t len = it.first - match0.begin();
+
+            for (size_t i = 2; i < matches.size(); ++i) {
+                const std::string_view matchi(matches[i].first);
+                auto cmp = std::mismatch(match0.begin(), match0.end(), matchi.begin(), matchi.end());
+                len = std::min(len, static_cast<size_t>(cmp.first - match0.begin()));
+            }
+
+            const std::string updated_line = std::string(match0.substr(0, len));
+            matches.emplace_back(updated_line + path_postfix, updated_line.length());
+        }
+
+        std::sort(matches.begin(), matches.end(), [](const auto & a, const auto & b) {
+            return a.first.compare(0, a.second, b.first, 0, b.second) < 0;
+        });
+    }
+
+    return matches;
+}
+
+// note: make this view implementation generic, so that we can move to TUI in the future if we want to
+namespace ui {
+    static void init(const common_params & params) {
+        // TODO: avoid using atexit() here by making `console` a singleton
+        console::init(params.simple_io, params.use_color);
+        atexit([]() { console::cleanup(); });
+
+        console::set_completion_callback(auto_completion_callback);
+    }
+
+    struct spinner {
+        spinner(const std::string & message) {
+            if (!message.empty()) {
+                console::log("%s ", message.c_str());
+            }
+            console::spinner::start();
+        }
+        ~spinner() {
+            console::spinner::stop();
+        }
+    };
+
+    struct user_turn {
+        user_turn() {
+            console::set_display(DISPLAY_TYPE_USER_INPUT);
+        }
+        ~user_turn() {
+            console::set_display(DISPLAY_TYPE_RESET);
+        }
+        void echo(const std::string & buffer) {
+            if (buffer.size() > 500) {
+                console::log("\n> %s ... (truncated)\n", buffer.substr(0, 500).c_str());
+            } else {
+                console::log("\n> %s\n", buffer.c_str());
+            }
+        }
+        std::string read_input(bool multiline_input, const char * prompt = nullptr) {
+            if (prompt) {
+                console::log("%s", prompt);
+            } else {
+                console::log("\n> ");
+            }
+            std::string buffer;
+            std::string line;
+            bool another_line = true;
+            do {
+                another_line = console::readline(line, multiline_input);
+                buffer += line;
+            } while (another_line);
+            return buffer;
+        }
+    };
+
+    enum assistant_display_mode {
+        ASSISTANT_DISPLAY_MODE_REASONING,
+        ASSISTANT_DISPLAY_MODE_CONTENT,
+    };
+    struct assistant_turn {
+        assistant_display_mode mode = ASSISTANT_DISPLAY_MODE_CONTENT;
+        bool trailing_newline = true;
+        bool is_inside_reasoning = false;
+        assistant_turn() {
+            console::set_display(DISPLAY_TYPE_RESET);
+        }
+        ~assistant_turn() {
+            console::set_display(DISPLAY_TYPE_RESET);
+            add_newline_if_needed();
+        }
+        void push(assistant_display_mode m, const std::string & buffer) {
+            if (m != mode) {
+                add_newline_if_needed();
+                switch (m) {
+                    case ASSISTANT_DISPLAY_MODE_CONTENT:
+                        {
+                            if (is_inside_reasoning) {
+                                console::log("[End thinking]\n\n");
+                                is_inside_reasoning = false;
+                            }
+                            console::set_display(DISPLAY_TYPE_RESET);
+                        } break;
+                    case ASSISTANT_DISPLAY_MODE_REASONING:
+                        {
+                            console::set_display(DISPLAY_TYPE_REASONING);
+                            is_inside_reasoning = true;
+                            console::log("\n[Start thinking]\n\n");
+                        } break;
+                }
+            }
+            mode = m;
+            if (buffer.empty()) {
+                return;
+            }
+            trailing_newline = buffer.back() == '\n';
+            console::log("%s", buffer.c_str());
+            console::flush();
+        }
+        void add_newline_if_needed() {
+            if (!trailing_newline) {
+                console::log("\n");
+                console::flush();
+            }
+        }
+    };
+
+    static void show_error(const std::string & title, const std::string & message = "") {
+        console::spinner::stop();
+        console::error("Error: %s\n", title.c_str());
+        if (!message.empty()) {
+            console::log("%s\n", message.c_str());
+        }
+    }
+
+    static void show_message(const std::string & message) {
+        console::log("%s\n", message.c_str());
+    }
+
+    static void show_info(const std::string & message) {
+        console::set_display(DISPLAY_TYPE_INFO);
+        console::log("%s\n", message.c_str());
+        console::set_display(DISPLAY_TYPE_RESET);
+    }
+}

tools/cli/cli.cpp

@@ -1,20 +1,9 @@
-#include "chat.h"
-#include "common.h"
 #include "arg.h"
-#include "console.h"
-#include "fit.h"
-// #include "log.h"
+#include "common.h"
+#include "log.h"
 
-#include "server-common.h"
-#include "server-context.h"
-#include "server-task.h"
+#include "cli-context.h"
 
-#include <array>
-#include <atomic>
-#include <algorithm>
-#include <filesystem>
-#include <fstream>
-#include <thread>
 #include <signal.h>
 
 #if defined(_WIN32)
@@ -25,342 +14,19 @@
 #include <windows.h>
 #endif
 
-const char * LLAMA_ASCII_LOGO = R"(
-▄▄ ▄▄
-██ ██
-██ ██  ▀▀█▄ ███▄███▄  ▀▀█▄    ▄████ ████▄ ████▄
-██ ██ ▄█▀██ ██ ██ ██ ▄█▀██    ██    ██ ██ ██ ██
-██ ██ ▀█▄██ ██ ██ ██ ▀█▄██ ██ ▀████ ████▀ ████▀
-                                    ██    ██
-                                    ▀▀    ▀▀
-)";
-
-static std::atomic<bool> g_is_interrupted = false;
-static bool should_stop() {
-    return g_is_interrupted.load();
-}
-
 #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) || defined (_WIN32)
 static void signal_handler(int) {
-    if (g_is_interrupted.load()) {
+    if (cli_context::interrupted().load()) {
         // second Ctrl+C - exit immediately
         // make sure to clear colors before exiting (not using LOG or console.cpp here to avoid deadlock)
         fprintf(stdout, "\033[0m\n");
         fflush(stdout);
         std::exit(130);
     }
-    g_is_interrupted.store(true);
+    cli_context::interrupted().store(true);
 }
 #endif
 
-struct cli_context {
-    server_context ctx_server;
-    json messages = json::array();
-    std::vector<raw_buffer> input_files;
-    task_params defaults;
-    bool verbose_prompt;
-
-    // thread for showing "loading" animation
-    std::atomic<bool> loading_show;
-
-    cli_context(const common_params & params) {
-        defaults.sampling    = params.sampling;
-        defaults.speculative = params.speculative;
-        defaults.n_keep      = params.n_keep;
-        defaults.n_predict   = params.n_predict;
-        defaults.antiprompt  = params.antiprompt;
-
-        defaults.stream = true; // make sure we always use streaming mode
-        defaults.timings_per_token = true; // in order to get timings even when we cancel mid-way
-        // defaults.return_progress = true; // TODO: show progress
-
-        verbose_prompt = params.verbose_prompt;
-    }
-
-    std::string generate_completion(result_timings & out_timings) {
-        server_response_reader rd = ctx_server.get_response_reader();
-        auto chat_params = format_chat();
-        {
-            // TODO: reduce some copies here in the future
-            server_task task = server_task(SERVER_TASK_TYPE_COMPLETION);
-            task.id         = rd.get_new_id();
-            task.index      = 0;
-            task.params     = defaults;           // copy
-            task.cli_prompt = chat_params.prompt; // copy
-            task.cli_files  = input_files;        // copy
-            task.cli        = true;
-
-            // chat template settings
-            task.params.chat_parser_params = common_chat_parser_params(chat_params);
-            task.params.chat_parser_params.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK;
-            if (!chat_params.parser.empty()) {
-                task.params.chat_parser_params.parser.load(chat_params.parser);
-            }
-
-            // Copy the preserved tokens into the sampling params
-            const llama_vocab * vocab = llama_model_get_vocab(
-                llama_get_model(ctx_server.get_llama_context()));
-            for (const auto & token : chat_params.preserved_tokens) {
-                auto ids = common_tokenize(vocab, token, false, true);
-                if (ids.size() == 1) {
-                    task.params.sampling.preserved_tokens.insert(ids[0]);
-                }
-            }
-
-            // reasoning budget sampler
-            if (!chat_params.thinking_end_tag.empty()) {
-                task.params.sampling.reasoning_budget_tokens = defaults.sampling.reasoning_budget_tokens;
-                task.params.sampling.generation_prompt = chat_params.generation_prompt;
-
-                if (!chat_params.thinking_start_tag.empty()) {
-                    task.params.sampling.reasoning_budget_start =
-                        common_tokenize(vocab, chat_params.thinking_start_tag, false, true);
-                }
-                task.params.sampling.reasoning_budget_end =
-                    common_tokenize(vocab, chat_params.thinking_end_tag, false, true);
-                task.params.sampling.reasoning_budget_forced =
-                    common_tokenize(vocab, defaults.sampling.reasoning_budget_message + chat_params.thinking_end_tag, false, true);
-            }
-
-            rd.post_task({std::move(task)});
-        }
-
-        if (verbose_prompt) {
-            console::set_display(DISPLAY_TYPE_PROMPT);
-            console::log("%s\n\n", chat_params.prompt.c_str());
-            console::set_display(DISPLAY_TYPE_RESET);
-        }
-
-        // wait for first result
-        console::spinner::start();
-        server_task_result_ptr result = rd.next(should_stop);
-
-        while (true) {
-            auto res_partial = dynamic_cast<server_task_result_cmpl_partial *>(result.get());
-            if (res_partial && res_partial->is_begin) {
-                // this is the "send 200 status to client" signal in streaming mode
-                // skip, do not stop the spinner
-                result = rd.next(should_stop);
-            } else {
-                console::spinner::stop();
-                break;
-            }
-        }
-
-        std::string curr_content;
-        bool is_thinking = false;
-
-        while (result) {
-            if (should_stop()) {
-                break;
-            }
-            if (result->is_error()) {
-                json err_data = result->to_json();
-                if (err_data.contains("message")) {
-                    console::error("Error: %s\n", err_data["message"].get<std::string>().c_str());
-                } else {
-                    console::error("Error: %s\n", err_data.dump().c_str());
-                }
-                return curr_content;
-            }
-            auto res_partial = dynamic_cast<server_task_result_cmpl_partial *>(result.get());
-            if (res_partial) {
-                out_timings = std::move(res_partial->timings);
-                for (const auto & diff : res_partial->oaicompat_msg_diffs) {
-                    if (!diff.content_delta.empty()) {
-                        if (is_thinking) {
-                            console::log("\n[End thinking]\n\n");
-                            console::set_display(DISPLAY_TYPE_RESET);
-                            is_thinking = false;
-                        }
-                        curr_content += diff.content_delta;
-                        console::log("%s", diff.content_delta.c_str());
-                        console::flush();
-                    }
-                    if (!diff.reasoning_content_delta.empty()) {
-                        console::set_display(DISPLAY_TYPE_REASONING);
-                        if (!is_thinking) {
-                            console::log("[Start thinking]\n");
-                        }
-                        is_thinking = true;
-                        console::log("%s", diff.reasoning_content_delta.c_str());
-                        console::flush();
-                    }
-                }
-            }
-            auto res_final = dynamic_cast<server_task_result_cmpl_final *>(result.get());
-            if (res_final) {
-                out_timings = std::move(res_final->timings);
-                break;
-            }
-            result = rd.next(should_stop);
-        }
-        g_is_interrupted.store(false);
-        // server_response_reader automatically cancels pending tasks upon destruction
-        return curr_content;
-    }
-
-    // TODO: support remote files in the future (http, https, etc)
-    std::string load_input_file(const std::string & fname, bool is_media) {
-        std::ifstream file = fs_open_ifstream(fname, std::ios::binary);
-        if (!file) {
-            return "";
-        }
-        if (is_media) {
-            raw_buffer buf;
-            buf.assign((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
-            input_files.push_back(std::move(buf));
-            return get_media_marker();
-        } else {
-            std::string content((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
-            return content;
-        }
-    }
-
-    common_chat_params format_chat() {
-        auto meta = ctx_server.get_meta();
-        auto & chat_params = meta.chat_params;
-
-        auto caps = common_chat_templates_get_caps(chat_params.tmpls.get());
-
-        common_chat_templates_inputs inputs;
-        inputs.messages              = common_chat_msgs_parse_oaicompat(messages);
-        inputs.tools                 = {}; // TODO
-        inputs.tool_choice           = COMMON_CHAT_TOOL_CHOICE_NONE;
-        inputs.json_schema           = ""; // TODO
-        inputs.grammar               = ""; // TODO
-        inputs.use_jinja             = chat_params.use_jinja;
-        inputs.parallel_tool_calls   = caps["supports_parallel_tool_calls"];
-        inputs.add_generation_prompt = true;
-        inputs.reasoning_format      = COMMON_REASONING_FORMAT_DEEPSEEK;
-        inputs.force_pure_content    = chat_params.force_pure_content;
-        inputs.enable_thinking       = chat_params.enable_thinking ? common_chat_templates_support_enable_thinking(chat_params.tmpls.get()) : false;
-
-        // Apply chat template to the list of messages
-        return common_chat_templates_apply(chat_params.tmpls.get(), inputs);
-    }
-};
-
-// TODO?: Make this reusable, enums, docs
-static const std::array<std::string_view, 8> cmds = {
-    "/audio ",
-    "/clear",
-    "/exit",
-    "/glob ",
-    "/image ",
-    "/read ",
-    "/regen",
-    "/video ",
-};
-
-static std::vector<std::pair<std::string, size_t>> auto_completion_callback(std::string_view line, size_t cursor_byte_pos) {
-    std::vector<std::pair<std::string, size_t>> matches;
-    std::string cmd;
-
-    if (line.length() > 1 && line.front() == '/' && !std::any_of(cmds.begin(), cmds.end(), [line](std::string_view prefix) {
-        return string_starts_with(line, prefix);
-    })) {
-        auto it = cmds.begin();
-
-        while ((it = std::find_if(it, cmds.end(), [line](std::string_view cmd_line) {
-            return string_starts_with(cmd_line, line);
-        })) != cmds.end()) {
-            matches.emplace_back(*it, it->length());
-            ++it;
-        }
-    } else {
-        auto it = std::find_if(cmds.begin(), cmds.end(), [line](std::string_view prefix) {
-            return prefix.back() == ' ' && string_starts_with(line, prefix);
-        });
-
-        if (it != cmds.end()) {
-            cmd = *it;
-        }
-    }
-
-    if (!cmd.empty() && cmd != "/glob " && line.length() >= cmd.length() && cursor_byte_pos >= cmd.length()) {
-        const std::string path_prefix  = std::string(line.substr(cmd.length(), cursor_byte_pos - cmd.length()));
-        const std::string path_postfix = std::string(line.substr(cursor_byte_pos));
-        auto cur_dir = std::filesystem::current_path();
-        std::string cur_dir_str = cur_dir.string();
-        std::string expanded_prefix = path_prefix;
-
-#if !defined(_WIN32)
-        if (string_starts_with(path_prefix, '~')) {
-            const char * home = std::getenv("HOME");
-            if (home && home[0]) {
-                expanded_prefix = home + path_prefix.substr(1);
-            }
-        }
-        if (string_starts_with(expanded_prefix, '/')) {
-#else
-        if (std::isalpha(expanded_prefix[0]) && expanded_prefix.find(':') == 1) {
-#endif
-            cur_dir = std::filesystem::path(expanded_prefix).parent_path();
-            cur_dir_str.clear();
-        } else if (!path_prefix.empty()) {
-            cur_dir /= std::filesystem::path(path_prefix).parent_path();
-        }
-
-        std::error_code ec;
-        for (const auto & entry : std::filesystem::directory_iterator(cur_dir, ec)) {
-            if (ec) {
-                break;
-            }
-            if (!entry.exists(ec)) {
-                ec.clear();
-                continue;
-            }
-
-            const std::string path_full = entry.path().string();
-            std::string path_entry = !cur_dir_str.empty() && string_starts_with(path_full, cur_dir_str) ? path_full.substr(cur_dir_str.length() + 1) : path_full;
-
-            if (entry.is_directory(ec)) {
-                path_entry.push_back(std::filesystem::path::preferred_separator);
-            }
-
-            if (expanded_prefix.empty() || string_starts_with(path_entry, expanded_prefix)) {
-                const std::string updated_line = cmd + path_entry;
-                matches.emplace_back(updated_line + path_postfix, updated_line.length());
-            }
-
-            if (ec) {
-                ec.clear();
-            }
-        }
-
-        if (matches.empty()) {
-            const std::string updated_line = cmd + path_prefix;
-            matches.emplace_back(updated_line + path_postfix, updated_line.length());
-        }
-
-        // Add the longest common prefix
-        if (!expanded_prefix.empty() && matches.size() > 1) {
-            const std::string_view match0(matches[0].first);
-            const std::string_view match1(matches[1].first);
-            auto it = std::mismatch(match0.begin(), match0.end(), match1.begin(), match1.end());
-            size_t len = it.first - match0.begin();
-
-            for (size_t i = 2; i < matches.size(); ++i) {
-                const std::string_view matchi(matches[i].first);
-                auto cmp = std::mismatch(match0.begin(), match0.end(), matchi.begin(), matchi.end());
-                len = std::min(len, static_cast<size_t>(cmp.first - match0.begin()));
-            }
-
-            const std::string updated_line = std::string(match0.substr(0, len));
-            matches.emplace_back(updated_line + path_postfix, updated_line.length());
-        }
-
-        std::sort(matches.begin(), matches.end(), [](const auto & a, const auto & b) {
-            return a.first.compare(0, a.second, b.first, 0, b.second) < 0;
-        });
-    }
-
-    return matches;
-}
-
-static constexpr size_t FILE_GLOB_MAX_RESULTS = 100;
-
 // satisfies -Wmissing-declarations
 int llama_cli(int argc, char ** argv);
 
@@ -375,25 +41,6 @@ int llama_cli(int argc, char ** argv) {
         return 1;
     }
 
-    // TODO: maybe support it later?
-    if (params.conversation_mode == COMMON_CONVERSATION_MODE_DISABLED) {
-        console::error("--no-conversation is not supported by llama-cli\n");
-        console::error("please use llama-completion instead\n");
-    }
-
-    // struct that contains llama context and inference
-    cli_context ctx_cli(params);
-
-    llama_backend_init();
-    llama_numa_init(params.numa);
-
-    // TODO: avoid using atexit() here by making `console` a singleton
-    console::init(params.simple_io, params.use_color);
-    atexit([]() { console::cleanup(); });
-
-    console::set_display(DISPLAY_TYPE_RESET);
-    console::set_completion_callback(auto_completion_callback);
-
 #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
     struct sigaction sigint_action;
     sigint_action.sa_handler = signal_handler;
@@ -408,276 +55,11 @@ int llama_cli(int argc, char ** argv) {
     SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
 #endif
 
-    console::log("\nLoading model... "); // followed by loading animation
-    console::spinner::start();
-    if (!ctx_cli.ctx_server.load_model(params)) {
-        console::spinner::stop();
-        console::error("\nFailed to load the model\n");
+    cli_context ctx_cli(params);
+
+    if (!ctx_cli.init()) {
         return 1;
     }
 
-    ctx_cli.defaults.sampling = params.sampling;
-
-    console::spinner::stop();
-    console::log("\n");
-
-    std::thread inference_thread([&ctx_cli]() {
-        ctx_cli.ctx_server.start_loop();
-    });
-
-    auto inf = ctx_cli.ctx_server.get_meta();
-    std::string modalities = "text";
-    if (inf.has_inp_image) {
-        modalities += ", vision";
-    }
-    if (inf.has_inp_audio) {
-        modalities += ", audio";
-    }
-
-    auto add_system_prompt = [&]() {
-        if (!params.system_prompt.empty()) {
-            ctx_cli.messages.push_back({
-                {"role",    "system"},
-                {"content", params.system_prompt}
-            });
-        }
-    };
-    add_system_prompt();
-
-    console::log("\n");
-    console::log("%s\n", LLAMA_ASCII_LOGO);
-    console::log("build      : %s\n", inf.build_info.c_str());
-    console::log("model      : %s\n", inf.model_name.c_str());
-    if (!inf.model_ftype.empty()) {
-        console::log("ftype      : %s\n", inf.model_ftype.c_str());
-    }
-    console::log("modalities : %s\n", modalities.c_str());
-    if (!params.system_prompt.empty()) {
-        console::log("using custom system prompt\n");
-    }
-    console::log("\n");
-    console::log("available commands:\n");
-    console::log("  /exit or Ctrl+C     stop or exit\n");
-    console::log("  /regen              regenerate the last response\n");
-    console::log("  /clear              clear the chat history\n");
-    console::log("  /read <file>        add a text file\n");
-    console::log("  /glob <pattern>     add text files using globbing pattern\n");
-    if (inf.has_inp_image) {
-        console::log("  /image <file>       add an image file\n");
-    }
-    if (inf.has_inp_audio) {
-        console::log("  /audio <file>       add an audio file\n");
-    }
-    if (inf.has_inp_video) {
-        console::log("  /video <file>       add a video file\n");
-    }
-    console::log("\n");
-
-    // interactive loop
-    std::string cur_msg;
-
-    auto add_text_file = [&](const std::string & fname) -> bool {
-        std::string marker = ctx_cli.load_input_file(fname, false);
-        if (marker.empty()) {
-            console::error("file does not exist or cannot be opened: '%s'\n", fname.c_str());
-            return false;
-        }
-        if (inf.fim_sep_token != LLAMA_TOKEN_NULL) {
-            cur_msg += common_token_to_piece(ctx_cli.ctx_server.get_llama_context(), inf.fim_sep_token, true);
-            cur_msg += fname;
-            cur_msg.push_back('\n');
-        } else {
-            cur_msg += "--- File: ";
-            cur_msg += fname;
-            cur_msg += " ---\n";
-        }
-        cur_msg += marker;
-        console::log("Loaded text from '%s'\n", fname.c_str());
-        return true;
-    };
-
-    while (true) {
-        std::string buffer;
-        console::set_display(DISPLAY_TYPE_USER_INPUT);
-        if (params.prompt.empty()) {
-            console::log("\n> ");
-            std::string line;
-            bool another_line = true;
-            do {
-                another_line = console::readline(line, params.multiline_input);
-                buffer += line;
-            } while (another_line);
-        } else {
-            // process input prompt from args
-            for (auto & fname : params.image) {
-                std::string marker = ctx_cli.load_input_file(fname, true);
-                if (marker.empty()) {
-                    console::error("file does not exist or cannot be opened: '%s'\n", fname.c_str());
-                    break;
-                }
-                console::log("Loaded media from '%s'\n", fname.c_str());
-                cur_msg += marker;
-            }
-            buffer = params.prompt;
-            if (buffer.size() > 500) {
-                console::log("\n> %s ... (truncated)\n", buffer.substr(0, 500).c_str());
-            } else {
-                console::log("\n> %s\n", buffer.c_str());
-            }
-            params.prompt.clear(); // only use it once
-        }
-        console::set_display(DISPLAY_TYPE_RESET);
-        console::log("\n");
-
-        if (should_stop()) {
-            g_is_interrupted.store(false);
-            break;
-        }
-
-        // remove trailing newline
-        if (!buffer.empty() &&buffer.back() == '\n') {
-            buffer.pop_back();
-        }
-
-        // skip empty messages
-        if (buffer.empty()) {
-            continue;
-        }
-
-        bool add_user_msg = true;
-
-        // process commands
-        if (string_starts_with(buffer, "/exit")) {
-            break;
-        } else if (string_starts_with(buffer, "/regen")) {
-            if (ctx_cli.messages.size() >= 2) {
-                size_t last_idx = ctx_cli.messages.size() - 1;
-                ctx_cli.messages.erase(last_idx);
-                add_user_msg = false;
-            } else {
-                console::error("No message to regenerate.\n");
-                continue;
-            }
-        } else if (string_starts_with(buffer, "/clear")) {
-            ctx_cli.messages.clear();
-            add_system_prompt();
-
-            ctx_cli.input_files.clear();
-            console::log("Chat history cleared.\n");
-            continue;
-        } else if (
-                (string_starts_with(buffer, "/image ") && inf.has_inp_image) ||
-                (string_starts_with(buffer, "/audio ") && inf.has_inp_audio) ||
-                (string_starts_with(buffer, "/video ") && inf.has_inp_video)) {
-            // just in case (bad copy-paste for example), we strip all trailing/leading spaces
-            std::string fname = string_strip(buffer.substr(7));
-            std::string marker = ctx_cli.load_input_file(fname, true);
-            if (marker.empty()) {
-                console::error("file does not exist or cannot be opened: '%s'\n", fname.c_str());
-                continue;
-            }
-            cur_msg += marker;
-            console::log("Loaded media from '%s'\n", fname.c_str());
-            continue;
-        } else if (string_starts_with(buffer, "/read ")) {
-            std::string fname = string_strip(buffer.substr(6));
-            add_text_file(fname);
-            continue;
-        } else if (string_starts_with(buffer, "/glob ")) {
-            std::error_code ec;
-            size_t count = 0;
-            auto curdir = std::filesystem::current_path();
-            std::string pattern = string_strip(buffer.substr(6));
-            std::filesystem::path rel_path;
-
-            auto startglob = pattern.find_first_of("![*?");
-            if (startglob != std::string::npos && startglob != 0) {
-                auto endpath = pattern.substr(0, startglob).find_last_of('/');
-                if (endpath != std::string::npos) {
-                    std::string rel_pattern = pattern.substr(0, endpath);
-#if !defined(_WIN32)
-                    if (string_starts_with(rel_pattern, '~')) {
-                        const char * home = std::getenv("HOME");
-                        if (home && home[0]) {
-                            rel_pattern = home + rel_pattern.substr(1);
-                        }
-                    }
-#endif
-                    rel_path = rel_pattern;
-                    pattern.erase(0, endpath + 1);
-                    curdir /= rel_path;
-                }
-            }
-
-            for (const auto & entry : std::filesystem::recursive_directory_iterator(curdir,
-                    std::filesystem::directory_options::skip_permission_denied, ec)) {
-                if (!entry.is_regular_file()) {
-                    continue;
-                }
-
-                std::string rel = std::filesystem::relative(entry.path(), curdir, ec).string();
-                if (ec) {
-                    ec.clear();
-                    continue;
-                }
-                std::replace(rel.begin(), rel.end(), '\\', '/');
-
-                if (!glob_match(pattern, rel)) {
-                    continue;
-                }
-
-                if (!add_text_file((rel_path / rel).string())) {
-                    continue;
-                }
-
-                if (++count >= FILE_GLOB_MAX_RESULTS) {
-                    console::error("Maximum number of globbed files allowed (%zu) reached.\n", FILE_GLOB_MAX_RESULTS);
-                    break;
-                }
-            }
-            continue;
-        } else {
-            // not a command
-            cur_msg += buffer;
-        }
-
-        // generate response
-        if (add_user_msg) {
-            ctx_cli.messages.push_back({
-                {"role",    "user"},
-                {"content", cur_msg}
-            });
-            cur_msg.clear();
-        }
-        result_timings timings;
-        std::string assistant_content = ctx_cli.generate_completion(timings);
-        ctx_cli.messages.push_back({
-            {"role",    "assistant"},
-            {"content", assistant_content}
-        });
-        console::log("\n");
-
-        if (params.show_timings) {
-            console::set_display(DISPLAY_TYPE_INFO);
-            console::log("\n");
-            console::log("[ Prompt: %.1f t/s | Generation: %.1f t/s ]\n", timings.prompt_per_second, timings.predicted_per_second);
-            console::set_display(DISPLAY_TYPE_RESET);
-        }
-
-        if (params.single_turn) {
-            break;
-        }
-    }
-
-    console::set_display(DISPLAY_TYPE_RESET);
-
-    console::log("\nExiting...\n");
-    ctx_cli.ctx_server.terminate();
-    inference_thread.join();
-
-    // bump the log level to display timings
-    common_log_set_verbosity_thold(LOG_LEVEL_INFO);
-    common_memory_breakdown_print(ctx_cli.ctx_server.get_llama_context());
-
-    return 0;
+    return ctx_cli.run();
 }

tools/server/server-context.cpp

@@ -4521,6 +4521,7 @@ void server_routes::init_routes() {
             { "default_generation_settings", default_generation_settings_for_props },
             { "total_slots",                 params.n_parallel },
             { "model_alias",                 meta->model_name },
+            { "model_ftype",                 meta->model_ftype },
             { "model_path",                  meta->model_path },
             { "modalities",                  json {
                 {"vision", meta->has_inp_image},

tools/server/server-models.cpp

@@ -7,6 +7,7 @@
 #include "build-info.h"
 #include "preset.h"
 #include "download.h"
+#include "http.h"
 
 #include <cpp-httplib/httplib.h> // TODO: remove this once we use HTTP client from download.h
 #include <optional>
@@ -28,14 +29,7 @@
 #include <sstream>
 #include <cstring>
 
-#ifdef _WIN32
-#include <winsock2.h>
-#include <windows.h>
-#else
-#include <sys/socket.h>
-#include <netinet/in.h>
-#include <arpa/inet.h>
-#include <unistd.h>
+#ifndef _WIN32
 extern char **environ;
 #endif
 
@@ -716,66 +710,6 @@ std::optional<server_model_meta> server_models::get_meta(const std::string & nam
     return std::nullopt;
 }
 
-static int 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;
-}
-
 // helper to convert vector<string> to char **
 // pointers are only valid as long as the original vector is valid
 static std::vector<char *> to_char_ptr_array(const std::vector<std::string> & vec) {
@@ -879,7 +813,7 @@ void server_models::load(const std::string & name, const load_options & opts) {
     // prepare new instance info
     instance_t inst;
     inst.meta             = meta;
-    inst.meta.port        = get_free_port();
+    inst.meta.port        = common_http_get_free_port();
     inst.meta.status      = SERVER_MODEL_STATUS_LOADING;
     inst.meta.loaded_info = json{};
     inst.meta.last_used   = ggml_time_ms();

tools/server/server.cpp

@@ -36,6 +36,19 @@ static inline void signal_handler(int signal) {
     shutdown_handler(signal);
 }
 
+// satisfies -Wmissing-declarations (used by llama command)
+int llama_server(int argc, char ** argv);
+
+// to be used via CLI (argc / argv are used by router mode only)
+int llama_server(common_params & params, int argc, char ** argv);
+void llama_server_terminate();
+void llama_server_terminate() {
+    if (shutdown_handler) {
+        shutdown_handler(0);
+    }
+}
+
+
 // wrapper function that handles exceptions and logs errors
 // this is to make sure handler_t never throws exceptions; instead, it returns an error response
 static server_http_context::handler_t ex_wrapper(server_http_context::handler_t func) {
@@ -72,9 +85,6 @@ static server_http_context::handler_t ex_wrapper(server_http_context::handler_t
     };
 }
 
-// satisfies -Wmissing-declarations
-int llama_server(int argc, char ** argv);
-
 int llama_server(int argc, char ** argv) {
     std::setlocale(LC_NUMERIC, "C");
 
@@ -94,16 +104,26 @@ int llama_server(int argc, char ** argv) {
     llama_backend_init();
     llama_numa_init(params.numa);
 
+    return llama_server(params, argc, argv);
+}
+
+int llama_server(common_params & params, int argc, char ** argv) {
+    bool is_run_by_cli = (argv == nullptr);
+
     common_models_handler models_handler;
-    try {
-        models_handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER);
-        if (common_models_handler_is_preset_repo(models_handler)) {
-            // apply the preset and start the server in router mode
-            common_models_handler_apply(models_handler, params);
+
+    // note: router mode also accepts -hf remote-preset, so we need to check that first
+    if (!is_run_by_cli && !params.model.hf_repo.empty()) {
+        try {
+            models_handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER);
+            if (common_models_handler_is_preset_repo(models_handler)) {
+                // apply the preset and start the server in router mode
+                common_models_handler_apply(models_handler, params);
+            }
+        } catch (const std::exception & e) {
+            SRV_ERR("failed to fetch model metadata: %s\n", e.what());
+            return 1;
         }
-    } catch (const std::exception & e) {
-        SRV_ERR("failed to fetch model metadata: %s\n", e.what());
-        return 1;
     }
 
     // router server never loads a model and must not touch the GPU
@@ -321,8 +341,9 @@ int llama_server(int argc, char ** argv) {
 
     if (child.is_child() && child.get_mode() == SERVER_CHILD_MODE_DOWNLOAD) {
         return child.run_download(params);
-    } else if (!is_router_server) {
+    } else if (!is_router_server && !is_run_by_cli) {
         // single-model mode (NOT spawned by router)
+        // if this is invoked by CLI, model downloading should be already handled
         try {
             common_models_handler_apply(models_handler, params);
         } catch (const std::exception & e) {
@@ -411,20 +432,22 @@ int llama_server(int argc, char ** argv) {
         };
     }
 
-    // TODO: refactor in common/console
+    // register signal handler if not running by CLI
+    if (!is_run_by_cli) {
 #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
-    struct sigaction sigint_action;
-    sigint_action.sa_handler = signal_handler;
-    sigemptyset (&sigint_action.sa_mask);
-    sigint_action.sa_flags = 0;
-    sigaction(SIGINT, &sigint_action, NULL);
-    sigaction(SIGTERM, &sigint_action, NULL);
+        struct sigaction sigint_action;
+        sigint_action.sa_handler = signal_handler;
+        sigemptyset (&sigint_action.sa_mask);
+        sigint_action.sa_flags = 0;
+        sigaction(SIGINT, &sigint_action, NULL);
+        sigaction(SIGTERM, &sigint_action, NULL);
 #elif defined (_WIN32)
-    auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {
-        return (ctrl_type == CTRL_C_EVENT) ? (signal_handler(SIGINT), true) : false;
-    };
-    SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
+        auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {
+            return (ctrl_type == CTRL_C_EVENT) ? (signal_handler(SIGINT), true) : false;
+        };
+        SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
 #endif
+    }
 
     SRV_INF("listening on %s\n", ctx_http.listening_address.c_str());