Commit caa8191df for llama.cpp
commit caa8191df159f764fb0160e3d2fc6be008948aee
Author: Max Krasnyansky <maxk@qti.qualcomm.com>
Date: Mon Oct 26 08:21:14 2026 -0700
hexagon: MUL_MAT_ID row skipping optimizations and improvements (#31169)
* hex-mmid: vectorize skipped row zeroing and added zeroing to fused MM_ID_NX
* hex-mmid: remove scalar div
diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c
index 784bb37b9..b4b943058 100644
--- a/ggml/src/ggml-hexagon/htp/matmul-ops.c
+++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c
@@ -16,6 +16,7 @@
#include "hvx-dump.h"
#include "hvx-arith.h"
#include "hvx-reduce.h"
+#include "hvx-copy.h"
#define GGML_COMMON_DECL_C
#include "ggml-common.h"
@@ -1413,7 +1414,7 @@ static void hvx_mv_id_tiled(unsigned int nth, unsigned int ith, void * data) {
for (uint32_t ie1 = 0; ie1 < n_aids; ++ie1) { // for each expert
const int32_t eid = *(const int32_t *) ((const uint8_t *) src2->data + ie1 * src2->nb[0]);
- if (eid == -1) {
+ if (eid < 0) {
continue;
}
assert(eid >= 0 && eid < (int32_t) n_ids);
@@ -4111,7 +4112,7 @@ static int hvx_mm_matmul_id_nx(
return HTP_STATUS_OK;
}
-static inline void scan_expert_ids_n(
+static __attribute__((noinline)) uint32_t scan_expert_ids_n(
const struct htp_tensor * ids,
const uint32_t n_ids,
uint32_t n_as,
@@ -4121,15 +4122,17 @@ static inline void scan_expert_ids_n(
) {
const size_t ids_nb1 = ids->nb[1];
const uint8_t * ids_data = (const uint8_t *) ids->data;
+ uint32_t n_skipped = 0;
for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
for (uint32_t id = 0; id < n_ids; ++id) {
const int32_t i02 = row_ptr[id];
- if (i02 == -1) {
+ if (i02 < 0) {
+ n_skipped++;
continue;
}
- assert(i02 >= 0 && i02 < n_as);
+ assert(i02 >= 0 && i02 < (int32_t) n_as);
if (matrix_rows) {
matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
@@ -4137,31 +4140,10 @@ static inline void scan_expert_ids_n(
counts[i02] += 1;
}
}
+ return n_skipped;
}
-// the matmuls only write the rows that an expert owns, so the rows of the skipped slots are zeroed here
-static inline void zero_skipped_dst_rows(
- const struct htp_tensor * ids,
- uint32_t n_ids,
- const struct htp_tensor * dst,
- uint32_t ne0,
- size_t nb1,
- size_t nb2
-) {
- const uint8_t * ids_data = (const uint8_t *) ids->data;
-
- for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
- for (uint32_t id = 0; id < n_ids; ++id) {
- const int32_t i02 = *(const int32_t *) (ids_data + iid1 * ids->nb[1] + id * ids->nb[0]);
- if (i02 != -1) {
- continue;
- }
- memset((uint8_t *) (dst->data + id * nb1 + iid1 * nb2), 0, ne0 * sizeof(float));
- }
- }
-}
-
-static inline void scan_expert_ids(
+static __attribute__((noinline)) uint32_t scan_expert_ids(
const struct htp_tensor * ids,
uint32_t n_ids,
uint32_t n_as,
@@ -4173,23 +4155,25 @@ static inline void scan_expert_ids(
if (ids_nb0 == 4) {
switch (n_ids) {
- case 8: scan_expert_ids_n(ids, 8, n_as, counts, matrix_rows, mapping_stride); break;
- case 4: scan_expert_ids_n(ids, 4, n_as, counts, matrix_rows, mapping_stride); break;
- case 2: scan_expert_ids_n(ids, 2, n_as, counts, matrix_rows, mapping_stride); break;
- default: scan_expert_ids_n(ids, n_ids, n_as, counts, matrix_rows, mapping_stride); break;
+ case 8: return scan_expert_ids_n(ids, 8, n_as, counts, matrix_rows, mapping_stride);
+ case 4: return scan_expert_ids_n(ids, 4, n_as, counts, matrix_rows, mapping_stride);
+ case 2: return scan_expert_ids_n(ids, 2, n_as, counts, matrix_rows, mapping_stride);
+ default: return scan_expert_ids_n(ids, n_ids, n_as, counts, matrix_rows, mapping_stride);
}
} else {
// Strided fallback
const size_t ids_nb1 = ids->nb[1];
const uint8_t * ids_data = (const uint8_t *) ids->data;
+ uint32_t n_skipped = 0;
for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
const int32_t * row_ptr = (const int32_t *) (ids_data + iid1 * ids_nb1);
for (uint32_t id = 0; id < n_ids; ++id) {
const int32_t i02 = *(const int32_t *) ((const uint8_t *) row_ptr + id * ids_nb0);
- if (i02 == -1) {
+ if (i02 < 0) {
+ n_skipped++;
continue;
}
- assert(i02 >= 0 && i02 < n_as);
+ assert(i02 >= 0 && i02 < (int32_t) n_as);
if (matrix_rows) {
matrix_rows[i02 * mapping_stride + counts[i02]] = (struct mmid_row_mapping) { id, iid1 };
@@ -4197,48 +4181,83 @@ static inline void scan_expert_ids(
counts[i02] += 1;
}
}
+ return n_skipped;
}
}
-int op_matmul_id(struct htp_ops_context * octx) {
- htp_matmul_tensors_preamble;
-
- const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
- struct htp_mm_context mmctx_struct = {0};
- struct htp_mm_context * mmctx = &mmctx_struct;
-
- const int status = htp_mm_init_context(octx, kparams);
- if (status != HTP_STATUS_OK) {
- return status;
+static __attribute__((noinline)) uint32_t count_skipped_ids_single(const struct htp_tensor * ids, uint32_t n_ids) {
+ const uint8_t * ids_data = (const uint8_t *) ids->data;
+ const size_t ids_nb0 = ids->nb[0];
+ uint32_t n_skipped = 0;
+ for (uint32_t id = 0; id < n_ids; ++id) {
+ const int32_t i02 = *(const int32_t *) (ids_data + id * ids_nb0);
+ if (i02 < 0) {
+ n_skipped++;
+ }
}
+ return n_skipped;
+}
- struct htp_thread_trace * tr = &octx->ctx->trace[0];
- htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
-
- mmctx->octx = octx;
- mmctx->act = act;
+// Zero rows for skipped slots (ids < 0). Matmuls only write rows that an expert owns.
+static void zero_skipped_dst_rows(
+ struct htp_ops_context * octx,
+ const struct htp_tensor * ids,
+ uint32_t n_ids,
+ const struct htp_tensor * dst
+) {
+ const uint32_t elem_size = (dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float);
+ uint32_t col_start = 0;
+ uint32_t col_count = dst->ne[0];
- const struct htp_tensor * restrict ids = octx->src[2];
- if (htp_tensor_is_extended(ids) || htp_tensor_is_extended(act) || htp_tensor_is_extended(dst)) {
- return HTP_STATUS_NO_SUPPORT;
+ if (octx->ctx->mdev.count > 1) {
+ const bool can_split = htp_tensor_can_row_partition(dst, elem_size);
+ const uint32_t chunk_elems = (dst->type == HTP_TYPE_F16) ? 64 : 32;
+ const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+ dst->ne[0], can_split ? chunk_elems : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+ col_start = range.start;
+ col_count = range.count;
}
- const size_t src0_row_size = nb01;
- const size_t dst_row_size = nb1;
-
- const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);
+ if (col_count == 0) {
+ return;
+ }
- const uint32_t src0_nrows = ne01; // per expert
- const uint32_t act_nrows = ne11 * ne12 * ne13;
+ const uint8_t * ids_data = (const uint8_t *) ids->data;
+ const size_t ids_nb0 = ids->nb[0];
+ const size_t ids_nb1 = ids->nb[1];
+ const size_t dst_nb1 = dst->nb[1];
+ const size_t dst_nb2 = dst->nb[2];
+ const uint32_t col_offset_bytes = col_start * elem_size;
- // row groups
- const int n_ids = ids->ne[0]; // n_expert_used
- const int n_as = ne02; // n_expert
+ for (uint32_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) {
+ const uint8_t * ids_row = ids_data + iid1 * ids_nb1;
+ uint8_t * dst_row_group = (uint8_t *) dst->data + iid1 * dst_nb2 + col_offset_bytes;
+ for (uint32_t id = 0; id < n_ids; ++id) {
+ const int32_t i02 = *(const int32_t *) (ids_row + id * ids_nb0);
+ if (i02 < 0) {
+ hvx_splat_u(dst_row_group + id * dst_nb1, Q6_V_vzero(), col_count, elem_size);
+ }
+ }
+ }
+}
+static int htp_mmid_prepare_mapping(
+ struct htp_ops_context * octx,
+ const struct htp_tensor * ids,
+ uint32_t n_ids,
+ uint32_t n_as,
+ uint32_t act_nrows,
+ uint8_t ** out_mapping_buf,
+ uint32_t ** out_matrix_row_counts,
+ struct mmid_row_mapping ** out_matrix_rows,
+ uint32_t * out_mapping_stride,
+ uint32_t * out_n_skipped
+) {
uint8_t * mapping_buf = octx->ctx->ddr_spad_base;
uint32_t mapping_stride = 1;
uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
struct mmid_row_mapping * matrix_rows = NULL;
+ uint32_t n_skipped = 0;
if (act_nrows > 1) {
const size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
@@ -4246,8 +4265,8 @@ int op_matmul_id(struct htp_ops_context * octx) {
hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]);
- memset(matrix_row_counts, 0, matrix_row_counts_size);
- scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);
+ hvx_splat_u(matrix_row_counts, Q6_V_vzero(), n_as, sizeof(uint32_t));
+ n_skipped = scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);
uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as);
mapping_stride = max_count > 0 ? max_count : 1;
@@ -4265,11 +4284,72 @@ int op_matmul_id(struct htp_ops_context * octx) {
matrix_row_counts = (uint32_t *) mapping_buf;
matrix_rows = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size);
- memset(matrix_row_counts, 0, n_as * sizeof(uint32_t));
+ hvx_splat_u(matrix_row_counts, Q6_V_vzero(), n_as, sizeof(uint32_t));
scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
+ } else {
+ n_skipped = count_skipped_ids_single(ids, n_ids);
+ }
+
+ *out_mapping_buf = mapping_buf;
+ *out_matrix_row_counts = matrix_row_counts;
+ *out_matrix_rows = matrix_rows;
+ *out_mapping_stride = mapping_stride;
+ *out_n_skipped = n_skipped;
+
+ return HTP_STATUS_OK;
+}
+
+int op_matmul_id(struct htp_ops_context * octx) {
+ htp_matmul_tensors_preamble;
+
+ const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params;
+ struct htp_mm_context mmctx_struct = {0};
+ struct htp_mm_context * mmctx = &mmctx_struct;
+
+ const int status = htp_mm_init_context(octx, kparams);
+ if (status != HTP_STATUS_OK) {
+ return status;
}
- zero_skipped_dst_rows(ids, n_ids, dst, ne0, nb1, nb2);
+ struct htp_thread_trace * tr = &octx->ctx->trace[0];
+ htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0);
+
+ mmctx->octx = octx;
+ mmctx->act = act;
+
+ const struct htp_tensor * restrict ids = octx->src[2];
+ if (htp_tensor_is_extended(ids) || htp_tensor_is_extended(act) || htp_tensor_is_extended(dst)) {
+ return HTP_STATUS_NO_SUPPORT;
+ }
+
+ const size_t src0_row_size = nb01;
+ const size_t dst_row_size = nb1;
+
+ const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);
+
+ const uint32_t src0_nrows = ne01; // per expert
+ const uint32_t act_nrows = ne11 * ne12 * ne13;
+
+ // row groups
+ const int n_ids = ids->ne[0]; // n_expert_used
+ const int n_as = ne02; // n_expert
+
+ uint8_t * mapping_buf = NULL;
+ uint32_t * matrix_row_counts = NULL;
+ struct mmid_row_mapping * matrix_rows = NULL;
+ uint32_t mapping_stride = 1;
+ uint32_t n_skipped = 0;
+
+ int map_status = htp_mmid_prepare_mapping(octx, ids, n_ids, n_as, act_nrows,
+ &mapping_buf, &matrix_row_counts, &matrix_rows,
+ &mapping_stride, &n_skipped);
+ if (map_status != HTP_STATUS_OK) {
+ return map_status;
+ }
+
+ if (n_skipped > 0) {
+ zero_skipped_dst_rows(octx, ids, n_ids, dst);
+ }
mmctx->matrix_row_counts = matrix_row_counts;
mmctx->matrix_rows = matrix_rows;
@@ -4360,38 +4440,26 @@ int op_matmul_id_nx(struct htp_ops_context * octx) {
const int n_ids = ids->ne[0];
const int n_as = src0->ne[2];
- uint8_t * mapping_buf = octx->ctx->ddr_spad_base;
- uint32_t mapping_stride = 1;
- uint32_t * matrix_row_counts = (uint32_t *) mapping_buf;
+ uint8_t * mapping_buf = NULL;
+ uint32_t * matrix_row_counts = NULL;
struct mmid_row_mapping * matrix_rows = NULL;
+ uint32_t mapping_stride = 1;
+ uint32_t n_skipped = 0;
- if (act_nrows > 1) {
- const size_t matrix_row_counts_size = n_as * sizeof(uint32_t);
- assert(octx->ctx->ddr_spad_size >= matrix_row_counts_size);
-
- hex_l2fetch_block((const void *) ids->data, ids->ne[1] * ids->nb[1]);
-
- memset(matrix_row_counts, 0, matrix_row_counts_size);
- scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, NULL, 0);
-
- uint32_t max_count = hvx_reduce_max_i32((const uint8_t *) matrix_row_counts, n_as);
- mapping_stride = max_count > 0 ? max_count : 1;
-
- size_t matrix_row_map_size = n_as * mapping_stride * sizeof(struct mmid_row_mapping);
- const size_t total_map_size = matrix_row_counts_size + matrix_row_map_size;
+ int map_status = htp_mmid_prepare_mapping(octx, ids, n_ids, n_as, act_nrows,
+ &mapping_buf, &matrix_row_counts, &matrix_rows,
+ &mapping_stride, &n_skipped);
+ if (map_status != HTP_STATUS_OK) {
+ return map_status;
+ }
- if (total_map_size > octx->ctx->ddr_spad_size) {
- mapping_buf = memalign(128, total_map_size);
- if (!mapping_buf) {
- return HTP_STATUS_INTERNAL_ERR;
+ if (n_skipped > 0) {
+ for (uint32_t p = 0; p < n_weights; ++p) {
+ const struct htp_tensor * restrict d = octx->dsts[p];
+ if (d) {
+ zero_skipped_dst_rows(octx, ids, n_ids, d);
}
}
-
- matrix_row_counts = (uint32_t *) mapping_buf;
- matrix_rows = (struct mmid_row_mapping *) (mapping_buf + matrix_row_counts_size);
-
- memset(matrix_row_counts, 0, n_as * sizeof(uint32_t));
- scan_expert_ids(ids, n_ids, n_as, matrix_row_counts, matrix_rows, mapping_stride);
}
mmctx->matrix_row_counts = matrix_row_counts;