public_git/whisper.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/whisper.cpp.git synced 2026-10-20 17:23:00 +02:00
Code Issues Packages Projects Releases Wiki Activity

hexagon: matmul and flash-atten scalability updates (llama/29974)

Browse Source 
* hexagon: head-parallel flash_attn partitioning for row-split multicore

In row-split mode each core computes its output row shard of every
MUL_MAT, but flash_attn was previously partitioning by Q tokens
(flat qrow split) instead of by heads. This forced every core to
read the full KV cache (all n_kv_heads), negating the memory
bandwidth benefit of multicore on flash_attn.

Change both HMX and HVX flash_attn kernels to partition by KV heads
when n_kv_heads is divisible by n_cores: core i processes heads
[i*n_kv_heads/N, (i+1)*n_kv_heads/N) exclusively, reading only its
head shard of the KV cache. Falls back to the original token-block
split when n_kv_heads % n_cores != 0 (e.g. Gemma-4 with 2 KV heads
on 4 cores).

Controlled by GGML_HEXAGON_FA_HEAD_SPLIT (default 1 = on).
The flag is packed into bit 1 of the existing is_dst_fp32 kparams
byte to stay within the 128-byte kernel_params blob limit.

Measured gains at 4c row-split (PP t/s, ubatch=1024):
  Qwen3-0.6B:    6977 -> 11026  (+58%)
  llama-3.2-3B:  3717 ->  5522  (+49%)
  Qwen3.5-4B:    2739 ->  2855   (+4%)
  Gemma-4 MoE:   no change (MoE FFN dominates, fallback path)

TG is unchanged (flash_attn is a small fraction of decode time
relative to the matmul+barrier cost per layer).

* hex-fa: cleanup kern_params and head-split selection

* hex-fa: add -fa-head-split option to run.py

* hex-mdev: update matmul solver to account for reduced work in row-split scenarios

* hex-mmid: better work splitting by expers in multi-dev scenarios

* hex-fa: update HMX gating based on the model/n-hvx/ctx-len sweep

* hex-fa: precompute softcap/scale on the host

* hexagon: flatten matmul into 2d to use HMX in multi-sequence

* hex-mm: cleanup kparams and use collapse to 3/4D -> 2D mapping

* hex-mm: fix typo in collapse fallback

* hex-mm: another pass at consistent naming for act tensors

* hex-mm: add support for colapsing dims in fused matmuls

* hex-build: fix WoS build errors

* hex-mm: make sure to enforce dst stride in can_collapse

* hex-fa: add a onliner commit for head-split check

* hex-fa: remove unused local head_split var

* hex-fa: tighten up can_split checks

* hex-mm: update unfused paths to use act instead src1

* hex-mm: make sure to check all dsts for splitting

* hexagon: fix the second weight chunk address in the batched HMX matmul prologue

* hexagon: F16 activation and ragged N in the HMX matmul

* hex-mm: tighten the ragged/split checks in mdev cases

* hex-mm: enable MM fusion for F16 activations

* hex-mm: pass tiled sizes to the solver in fused paths

* hex-mmid: remove scalar divs from expert mapping loops

* hex-mmid: proper cacheline safety enforcement for mdev splits

* hex-mm: improve solver for mdev split scanarios and tail handling

* hex-mm: remove redundant checks

* hex-mm: fix fused HMX MUL_MAT_NX drops the final partial tile for quantized weights

* hex-mm: better handling of ragged shapes (removes scalar memset of vtcm)

---------

Co-authored-by: ebateni <ebateni@qti.qualcomm.com>
Co-authored-by: Jhen-Jie Hong <iainst0409@gmail.com>
Co-authored-by: Yiwei Shao <yiwei@aizip.ai>
This commit is contained in:
Max Krasnyansky
2026-10-19 09:47:09 +02:00
committed by Daniel Bevenius
co-authored by ebateni Jhen-Jie Hong Yiwei Shao
parent 0c7be69637
commit 8ab093a587
7 changed files with 990 additions and 531 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-hexagon/ggml-hexagon.cpp
+273 -65
View File
@@ -1,3 +1,5 @@
#define _USE_MATH_DEFINES
#include <assert.h>
#include <inttypes.h>
#include <stdio.h>
@@ -24,6 +26,10 @@
#include <cmath>
#include <initializer_list>
#ifndef M_LOG2E
# define M_LOG2E 1.44269504088896340736
#endif
#ifdef _WIN32
# define WIN32_LEAN_AND_MEAN
# ifndef NOMINMAX
@@ -101,6 +107,7 @@ static bool opt_dma64 = false;
static int opt_mm_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
static int opt_fa_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
static int opt_fa_head_split = 1; // 1 = partition flash_attn by KV heads in multicore (default on), 0 = token-based (original)
static int opt_gdn_select = 2; // 2 = HMX -> HVX, 1 = HVX, 0 = CPU (unsupported)
static int opt_ar_select = 2; // 2 = fused ALLREDUCE+ADD (default), 1 = unfused ALLREDUCE, 0 = fallback to CPY+FENCE
static int opt_ar_scatter = 1; // 1 = reduce-scatter the fused ALLREDUCE+ADD (default), 0 = full reduction
@@ -387,7 +394,7 @@ static void ggml_hexagon_precompute_sort_params(
static void ggml_hexagon_precompute_fused_mmnx_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
const struct ggml_tensor * src1,
const struct ggml_tensor * act,
int32_t n_weights,
struct htp_mm_kernel_params * kparams
);
@@ -395,7 +402,7 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
static void ggml_hexagon_precompute_fused_mmidnx_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
const struct ggml_tensor * src1,
const struct ggml_tensor * act,
const struct ggml_tensor * dst,
int32_t n_weights,
struct htp_mm_kernel_params * kparams
@@ -412,6 +419,10 @@ static bool ggml_hexagon_precompute_allreduce_params(
struct htp_allreduce_kernel_params * kparams
);
static bool ggml_hexagon_rows_stride(const int64_t * ne, const size_t * nb, size_t * stride);
static bool ggml_hexagon_matmul_can_collapse(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst);
static ggml_tensor ggml_hexagon_tensor_collapse_rows(const struct ggml_tensor * t);
static bool mm_is_hmx_eligible(const ggml_tensor * t);
static htp_op_code op_remap_to_htp(const ggml_tensor * t);
static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams);
@@ -444,6 +455,13 @@ static inline bool ggml_hexagon_tensors_overlap(const struct ggml_tensor * a, co
return a0 < b1 && b0 < a1;
}
static inline bool ggml_hexagon_can_row_partition(const struct ggml_tensor * t) {
if (t->ne[1] > 1 && (t->nb[1] & 127) != 0) return false;
if (t->ne[2] > 1 && (t->nb[2] & 127) != 0) return false;
if (t->ne[3] > 1 && (t->nb[3] & 127) != 0) return false;
return true;
}
struct htp_opnode;
struct ggml_hexagon_opbatch;
@@ -3410,6 +3428,10 @@ struct ggml_hexagon_opbatch {
return false;
}
if (orig_kparams->collapse != kparams.collapse) {
return false;
}
const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3];
const bool can_fuse = (kparams.n_hmx > 0) || (src1_nrows == 1);
if (!can_fuse) return false;
@@ -3493,8 +3515,20 @@ struct ggml_hexagon_opbatch {
return false;
}
const struct htp_mm_kernel_params * orig_kparams = (const struct htp_mm_kernel_params *) last_node.kernel_params;
const bool collapse = orig_kparams->collapse && ggml_hexagon_matmul_can_collapse(w_in, x, d_in);
if (orig_kparams->collapse && !collapse) {
return false;
}
struct htp_mm_kernel_params kparams;
ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, curr_n + 1, &kparams);
if (collapse) {
const ggml_tensor x_collapsed = ggml_hexagon_tensor_collapse_rows(x);
ggml_hexagon_precompute_fused_mmnx_params(sess, w0, &x_collapsed, curr_n + 1, &kparams);
kparams.collapse = 1;
} else {
ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, curr_n + 1, &kparams);
}
if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) {
return false;
}
@@ -3543,9 +3577,23 @@ struct ggml_hexagon_opbatch {
const ggml_tensor * w0 = last_node.src0();
const ggml_tensor * x = last_node.src1();
const ggml_tensor * w1 = node.src0();
const ggml_tensor * dst_0 = last_node.dst();
const ggml_tensor * dst_1 = node.dst();
const struct htp_mm_kernel_params * orig_kparams = (const struct htp_mm_kernel_params *) last_node.kernel_params;
const bool collapse = orig_kparams->collapse && ggml_hexagon_matmul_can_collapse(w1, x, dst_1);
if (orig_kparams->collapse && !collapse) {
return false;
}
struct htp_mm_kernel_params kparams;
ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, 2, &kparams);
if (collapse) {
const ggml_tensor x_collapsed = ggml_hexagon_tensor_collapse_rows(x);
ggml_hexagon_precompute_fused_mmnx_params(sess, w0, &x_collapsed, 2, &kparams);
kparams.collapse = 1;
} else {
ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, 2, &kparams);
}
if (!is_supported_mul_mat_nx_kernel(w0, &kparams)) {
return false;
}
@@ -3559,9 +3607,6 @@ struct ggml_hexagon_opbatch {
return false;
}
const ggml_tensor * dst_0 = last_node.dst();
const ggml_tensor * dst_1 = node.dst();
last_node.opcode = HTP_OP_MUL_MAT_NX;
last_node.name = "MUL_MAT_NX";
last_node.inputs.clear();
@@ -4831,15 +4876,48 @@ static bool ggml_hexagon_flash_attn_is_hmx_eligible(
return false;
}
// Fall back to HVX for small token counts if head dimension is small (DK <= 128)
const uint32_t neq1 = q->ne[1];
if (DK <= 128 && neq1 < 5) {
return false;
GGML_UNUSED(sinks);
// Explicit force mode
if (opt_fa_select > 2) {
return true;
}
return true;
const uint32_t M = q->ne[1];
GGML_UNUSED(sinks);
// Prefill or batched decode
if (M > 1) {
return true;
}
// Compute-bound head dim
if (DK >= 256) {
return true;
}
const uint32_t n_head = q->ne[2];
const uint32_t n_kv_heads = k->ne[2];
const uint32_t G = n_kv_heads > 0 ? n_head / n_kv_heads : 1;
const uint32_t S = k->ne[1];
// Tile alignment for 32-row HMX tiles
const bool is_tile_aligned = (G > 0 && (32 % G == 0));
if (!is_tile_aligned) {
if (sess->n_threads >= 6) {
return false;
}
return S >= 1024;
}
// Context depth crossover
uint32_t s_cross = 512;
if (DK <= 64) {
s_cross = (sess->n_threads >= 8) ? 2048 : ((sess->n_threads >= 6) ? 768 : 512);
} else {
s_cross = (sess->n_threads >= 8) ? 1024 : 512;
}
return S >= s_cross;
}
static bool ggml_hexagon_precompute_flash_attn_params(
@@ -4857,7 +4935,6 @@ static bool ggml_hexagon_precompute_flash_attn_params(
const struct ggml_tensor * k = op->src[1];
const struct ggml_tensor * v = op->src[2];
const struct ggml_tensor * mask = op->src[3];
const struct ggml_tensor * dst = op;
const uint32_t neq0 = q->ne[0]; // head_dim (DK)
const uint32_t neq1 = q->ne[1]; // n_tokens
@@ -4888,8 +4965,8 @@ static bool ggml_hexagon_precompute_flash_attn_params(
kparams->max_bias = max_bias;
kparams->logit_softcap = logit_softcap;
kparams->is_q_fp32 = (q->type == GGML_TYPE_F32) ? 1 : 0;
kparams->is_dst_fp32 = (dst->type == GGML_TYPE_F32) ? 1 : 0;
kparams->head_split = (opt_fa_head_split != 0) ? 1 : 0;
kparams->flags = 0;
kparams->G = G;
const uint32_t n_head = q->ne[2];
@@ -4906,9 +4983,16 @@ static bool ggml_hexagon_precompute_flash_attn_params(
const uint32_t DK_pad = hex_round_up(DK, 64);
const uint32_t DV_pad = hex_round_up(DV, 64);
size_t Br = 0, Bc = 0;
int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, sess->vtcm_size, sess->n_threads, kparams->is_q_fp32 != 0, sinks != nullptr, n_head);
int ret = hmx_fa_find_chunk_size(&Br, &Bc, G, DK_pad, DV_pad, neq1, nek1, sess->vtcm_size, sess->n_threads, (q->type == GGML_TYPE_F32), sinks != nullptr, n_head);
if (ret == 0) {
kparams->kernel_type = HTP_FA_KERNEL_HMX;
if (logit_softcap == 0.0f) {
kparams->scale = scale * (float) M_LOG2E;
kparams->logit_softcap = 0.0f;
} else {
kparams->scale = scale;
kparams->logit_softcap = logit_softcap * (float) M_LOG2E;
}
kparams->Br = Br;
kparams->Bc = Bc;
kparams->n_kv_blocks = (nek1 + Bc - 1) / Bc;
@@ -4916,7 +5000,7 @@ static bool ggml_hexagon_precompute_flash_attn_params(
kparams->u.hmx.g_br = hex_align_up(G * Br, 32);
kparams->u.hmx.pipeline = (kparams->n_kv_blocks >= 3 && sess->n_threads >= 2) ? 1 : 0;
kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK_pad, DV_pad, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0, kparams->is_q_fp32 != 0, sinks != nullptr, n_head);
kparams->vtcm_size = hmx_fa_compute_vtcm_usage(G, DK_pad, DV_pad, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0, (q->type == GGML_TYPE_F32), sinks != nullptr, n_head);
const size_t row_vec_bytes = hex_align_up(Bc * sizeof(uint16_t), 256);
kparams->u.hmx.row_buf_stride = row_vec_bytes / 128; // HVX vector is 128 bytes
@@ -4943,11 +5027,11 @@ static bool ggml_hexagon_precompute_flash_attn_params(
kparams->n_kv_blocks = (k->ne[1] + 64 - 1) / 64;
kparams->n_threads = sess->n_threads;
const size_t size_q_row_padded = hex_round_up(q->ne[0] * (kparams->is_q_fp32 ? 4 : 2), 128);
const size_t size_q_row_padded = hex_round_up(q->ne[0] * ((q->type == GGML_TYPE_F32) ? 4 : 2), 128);
const size_t size_k_row_padded = hex_round_up(k->ne[0] * 2, 128);
const size_t size_v_row_padded = hex_round_up(v->ne[0] * 2, 128);
kparams->vtcm_size = hvx_fa_compute_vtcm_usage(DK, DV, kparams->is_q_fp32 != 0, mask != nullptr, sinks != nullptr, n_head, sess->n_threads);
kparams->vtcm_size = hvx_fa_compute_vtcm_usage(DK, DV, (q->type == GGML_TYPE_F32), mask != nullptr, sinks != nullptr, n_head, sess->n_threads);
kparams->u.hvx.size_q_row_padded = size_q_row_padded;
kparams->u.hvx.size_k_row_padded = size_k_row_padded;
@@ -5102,7 +5186,7 @@ static bool ggml_hexagon_matmul_is_hmx_eligible(
bool is_matmul_id,
bool is_batched
) {
if (src1->type != GGML_TYPE_F32) {
if (src1->type != GGML_TYPE_F32 && (src1->type != GGML_TYPE_F16 || is_matmul_id)) {
return false;
}
@@ -5111,8 +5195,8 @@ static bool ggml_hexagon_matmul_is_hmx_eligible(
const int ne12 = src1->ne[2];
const int wtype = src0->type;
// HMX weight tile requires N to be 32-aligned.
if (ne01_padded % 32 != 0) {
// HMX weight tiles accept non-32-aligned N for non-matmul_id.
if (ne01_padded % 32 != 0 && is_matmul_id) {
return false;
}
@@ -5151,7 +5235,7 @@ static bool ggml_hexagon_matmul_is_hmx_eligible(
static bool ggml_hexagon_precompute_hmx_mm_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
const struct ggml_tensor * src1,
const struct ggml_tensor * act,
const struct ggml_tensor * dst,
int wtype,
int ne00_padded,
@@ -5167,9 +5251,22 @@ static bool ggml_hexagon_precompute_hmx_mm_params(
struct htp_mm_kernel_params * kparams
) {
const int aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
const bool pipeline = is_matmul_id ? false : htp_mm_hmx_pipeline(ne11);
const int n_threads = (int)sess->n_threads;
const int ne10 = src1->ne[0];
const int ne10 = act->ne[0];
int m_for_solver = ne11;
int m_for_solver_padded = ne11_padded;
// matmul_id partitions by expert; regular matmul partitions M rows (ne11) across devices
if (!is_matmul_id && sess->mdev.count > 1 && ((uint32_t) ne11 >= sess->mdev.count)) {
// when dst is null, padded dims are used for estimate which are 128-byte aligned
const bool dst_row_split = dst ? ggml_hexagon_can_row_partition(dst) : true;
const bool act_row_split = ggml_hexagon_can_row_partition(act);
if (dst_row_split && act_row_split) {
m_for_solver = (ne11 + (int) sess->mdev.count - 1) / (int) sess->mdev.count;
m_for_solver_padded = hex_round_up(std::max(m_for_solver, 32), 32);
}
}
const bool pipeline = is_matmul_id ? false : htp_mm_hmx_pipeline(m_for_solver);
const bool is_batched_val = is_matmul_id ? false : is_batched;
const int group_size = (ne02 > 0 ? ne12 / ne02 : 1);
@@ -5182,15 +5279,25 @@ static bool ggml_hexagon_precompute_hmx_mm_params(
if (is_batched_val && wtype == GGML_TYPE_F16 && group_size > 1) {
// Try grouped path first
if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, ne11, group_size, n_threads, pipeline, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, m_for_solver, group_size, n_threads, pipeline, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
use_grouped = true;
}
}
if (!use_grouped) {
// Fallback to simple 2D path (group_size = 1)
const int m_id_rows = (dst && is_matmul_id) ? (int) ((size_t) dst->ne[1] * dst->ne[2]) : 0;
if (!htp_mm_hmx_solve_2d_params(wtype, ne00_padded, m_id_rows, ne01_padded, ne11_padded, ne11, n_threads, pipeline, is_matmul_id, aligned_tile_size, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
int m_id_rows = 0;
if (dst && is_matmul_id) {
const int n_experts = ne02 > 0 ? ne02 : 1;
const size_t total_expert_rows = (size_t) dst->ne[1] * dst->ne[2];
int m_per_expert = (int) ((total_expert_rows + n_experts - 1) / n_experts);
if (sess->mdev.count > 1 && ggml_hexagon_can_row_partition(dst)) {
m_per_expert = (m_per_expert + (int) sess->mdev.count - 1) / (int) sess->mdev.count;
}
m_id_rows = hex_round_up(std::max(m_per_expert, 32), 32);
}
const uint32_t cost_m = is_matmul_id ? (uint32_t) m_id_rows : (uint32_t) m_for_solver;
if (!htp_mm_hmx_solve_2d_params(wtype, ne00_padded, m_id_rows, ne01_padded, m_for_solver_padded, cost_m, n_threads, pipeline, is_matmul_id, aligned_tile_size, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
return false;
}
}
@@ -5203,14 +5310,14 @@ static bool ggml_hexagon_precompute_hmx_mm_params(
kparams->n_act_threads = act_threads_selected;
kparams->tile_size = htp_mm_get_weight_tile_size(wtype);
kparams->aligned_tile_size = aligned_tile_size;
kparams->src1_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
kparams->act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
kparams->vtcm_size = vtcm_size;
kparams->vtcm_src0_size = 0;
kparams->div_n_act_threads = init_fastdiv_values(act_threads_selected);
kparams->div_ne00_padded = init_fastdiv_values(ne00_padded);
kparams->vtcm_src1_size = 0;
kparams->vtcm_src2_size = (int32_t) src2_size;
kparams->vtcm_dst_size = 0;
kparams->vtcm_act_size = 0;
kparams->vtcm_bias_size = (int32_t) src2_size;
kparams->vtcm_dst_size = 0;
if (is_batched && !is_matmul_id) {
kparams->kernel_type = HTP_MM_KERNEL_HMX_F16_BATCHED;
@@ -5247,6 +5354,9 @@ static void ggml_hexagon_precompute_hvx_mm_params(
kparams->n_hmx = 0;
kparams->n_threads = sess->n_threads;
GGML_UNUSED(ne02);
GGML_UNUSED(ne03);
const bool is_quant = (wtype != GGML_TYPE_F16 && wtype != GGML_TYPE_F32);
const int src1_nrows = ne11 * ne12 * ne13;
@@ -5259,7 +5369,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
if (is_matmul_id) {
kparams->kernel_type = (src1_nrows < (int) sess->n_threads) ? HTP_MM_KERNEL_HVX_QUANT_BLOCK : HTP_MM_KERNEL_HVX_QUANT_ROW;
kparams->src1_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
kparams->act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
struct htp_mm_hvx_vtcm_layout L;
uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
@@ -5267,7 +5377,7 @@ static void ggml_hexagon_precompute_hvx_mm_params(
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
htp_mm_hvx_vtcm_layout_build(
&L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
0, src0->nb[1], kparams->src1_row_size, 0, d, true, false
0, src0->nb[1], kparams->act_row_size, 0, d, true, false
);
if (L.total_bytes <= vtcm_budget) {
best_n_prefetch = d;
@@ -5281,13 +5391,14 @@ static void ggml_hexagon_precompute_hvx_mm_params(
kparams->n_prefetch = best_n_prefetch;
kparams->vtcm_size = L.total_bytes;
kparams->vtcm_src0_size = L.src0_bytes;
kparams->vtcm_src1_size = L.src1_bytes;
kparams->vtcm_act_size = L.act_bytes;
kparams->vtcm_bias_size = 0;
kparams->vtcm_dst_size = L.dst_bytes;
goto done_quant;
} else {
bool try_tiled = (k_align && opt_mm_select >= 1);
if (try_tiled) {
kparams->src1_row_size = htp_mm_weight_has_offset(wtype)
kparams->act_row_size = htp_mm_weight_has_offset(wtype)
? htp_mm_q8_1_tiled_row_size(ne10)
: htp_mm_q8_0_tiled_row_size(ne10);
if (src1_nrows < (int) sess->n_threads) {
@@ -5319,8 +5430,8 @@ static void ggml_hexagon_precompute_hvx_mm_params(
kparams->m_chunk = (m_chunk < (uint32_t) src1_nrows) ? m_chunk : 0;
kparams->vtcm_size = L.total_bytes;
kparams->vtcm_src0_size = L.src0_bytes;
kparams->vtcm_src1_size = L.src1_bytes;
kparams->vtcm_src2_size = L.src2_bytes;
kparams->vtcm_act_size = L.act_bytes;
kparams->vtcm_bias_size = L.bias_bytes;
kparams->vtcm_dst_size = L.dst_bytes;
goto done_quant;
}
@@ -5341,11 +5452,11 @@ static void ggml_hexagon_precompute_hvx_mm_params(
&L, &m_chunk)) {
kparams->kernel_type = HTP_MM_KERNEL_HVX_F16_F16_VTCM;
kparams->m_chunk = (m_chunk < (uint32_t) src1_nrows) ? m_chunk : 0;
kparams->src1_row_size = hex_round_up(ne10 * 2, 128);
kparams->act_row_size = hex_round_up(ne10 * 2, 128);
kparams->vtcm_size = L.total_bytes;
kparams->vtcm_src0_size = L.src0_bytes;
kparams->vtcm_src1_size = L.src1_bytes;
kparams->vtcm_src2_size = L.src2_bytes;
kparams->vtcm_act_size = L.act_bytes;
kparams->vtcm_bias_size = L.bias_bytes;
kparams->vtcm_dst_size = L.dst_bytes;
kparams->n_prefetch = 16;
return;
@@ -5363,11 +5474,11 @@ static void ggml_hexagon_precompute_hvx_mm_params(
&L, &m_chunk)) {
kparams->kernel_type = HTP_MM_KERNEL_HVX_F32_F32_VTCM;
kparams->m_chunk = (m_chunk < (uint32_t) src1_nrows) ? m_chunk : 0;
kparams->src1_row_size = hex_round_up(ne10 * 4, 128);
kparams->act_row_size = hex_round_up(ne10 * 4, 128);
kparams->vtcm_size = L.total_bytes;
kparams->vtcm_src0_size = L.src0_bytes;
kparams->vtcm_src1_size = L.src1_bytes;
kparams->vtcm_src2_size = L.src2_bytes;
kparams->vtcm_act_size = L.act_bytes;
kparams->vtcm_bias_size = L.bias_bytes;
kparams->vtcm_dst_size = L.dst_bytes;
kparams->n_prefetch = 16;
return;
@@ -5404,6 +5515,9 @@ static void ggml_hexagon_precompute_matmul_params_impl(
const int ne00_padded = is_repack ? hex_round_up(ne00, 32) : ne00;
const int ne01_padded = is_repack ? hex_round_up(ne01, 32) : ne01;
const int ne11_padded = hex_round_up(ne11, 32);
// VTCM has to hold whole 32-row weight tiles, so size for the rounded-up N
// even when the tensor itself is ragged.
const int ne01_tiled = hex_round_up(ne01_padded, 32);
const bool is_matmul_id = (dst->op == GGML_OP_MUL_MAT_ID);
const bool is_batched = (ne02 * ne03 > 1 || ne12 * ne13 > 1);
@@ -5413,7 +5527,7 @@ static void ggml_hexagon_precompute_matmul_params_impl(
// Check HMX eligibility and try precomputing HMX parameters
bool hmx_enabled = (sess->n_hmx > 0) && (opt_mm_select >= 2);
if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, src1, dst, ne01_padded, is_matmul_id, is_batched)) {
if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, dst, wtype, ne00_padded, ne01_padded, ne02, ne11, ne12, ne11_padded, is_matmul_id, is_batched, src2_size, vtcm_budget, kparams)) {
if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, dst, wtype, ne00_padded, ne01_tiled, ne02, ne11, ne12, ne11_padded, is_matmul_id, is_batched, src2_size, vtcm_budget, kparams)) {
goto finalize;
}
}
@@ -5429,6 +5543,66 @@ finalize:
kparams->div_ne12 = init_fastdiv_values(ne12);
}
// The rows of dims 1..3 can be walked with one stride (size-1 dims skipped); returns that stride
static bool ggml_hexagon_rows_stride(const int64_t * ne, const size_t * nb, size_t * stride) {
size_t s = 0, next = 0;
for (int i = 1; i < GGML_MAX_DIMS; i++) {
if (ne[i] == 1) continue;
if (s == 0) { s = nb[i]; next = s * ne[i]; continue; }
if (nb[i] != next) return false;
next *= ne[i];
}
*stride = s ? s : nb[1];
return true;
}
// A 2D weight applied to a batched activation whose rows are evenly strided is the same matmul over ne11 * ne12 * ne13 rows
static bool ggml_hexagon_matmul_can_collapse(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
size_t s1, sd;
return (dst->op == GGML_OP_MUL_MAT || dst->op == GGML_OP_ADD) &&
src0->ne[2] == 1 && src0->ne[3] == 1 && src1->ne[2] * src1->ne[3] > 1 &&
src1->nb[0] == ggml_type_size(src1->type) && ggml_hexagon_rows_stride(src1->ne, src1->nb, &s1) &&
dst->nb[0] == ggml_type_size(dst->type) && ggml_hexagon_rows_stride(dst->ne, dst->nb, &sd);
}
static bool ggml_hexagon_matmul_add_can_collapse(
const struct ggml_tensor * src0,
const struct ggml_tensor * src1,
const struct ggml_tensor * src2,
const struct ggml_tensor * dst
) {
if (!ggml_hexagon_matmul_can_collapse(src0, src1, dst)) {
return false;
}
if (!src2) {
return true;
}
const int64_t src2_nrows = src2->ne[1] * src2->ne[2] * src2->ne[3];
if (src2_nrows == 1) {
return true;
}
size_t s2;
return src2->nb[0] == ggml_type_size(src2->type) &&
src2->ne[0] == dst->ne[0] &&
src2->ne[1] == dst->ne[1] &&
src2->ne[2] == dst->ne[2] &&
src2->ne[3] == dst->ne[3] &&
ggml_hexagon_rows_stride(src2->ne, src2->nb, &s2);
}
static ggml_tensor ggml_hexagon_tensor_collapse_rows(const struct ggml_tensor * t) {
size_t stride = 0;
ggml_hexagon_rows_stride(t->ne, t->nb, &stride);
ggml_tensor c = *t;
c.ne[1] = t->ne[1] * t->ne[2] * t->ne[3];
c.ne[2] = 1;
c.ne[3] = 1;
c.nb[1] = stride;
c.nb[2] = c.nb[1] * c.ne[1];
c.nb[3] = c.nb[2];
return c;
}
static void ggml_hexagon_precompute_matmul_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
@@ -5436,6 +5610,13 @@ static void ggml_hexagon_precompute_matmul_params(
const struct ggml_tensor * dst,
struct htp_mm_kernel_params * kparams
) {
if (ggml_hexagon_matmul_can_collapse(src0, src1, dst)) {
const ggml_tensor src1_collapsed = ggml_hexagon_tensor_collapse_rows(src1);
const ggml_tensor dst_collapsed = ggml_hexagon_tensor_collapse_rows(dst);
ggml_hexagon_precompute_matmul_params_impl(sess, src0, &src1_collapsed, &dst_collapsed, 0, 0, kparams);
kparams->collapse = 1;
return;
}
ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, 0, kparams);
}
@@ -5447,6 +5628,18 @@ static void ggml_hexagon_precompute_fused_matmul_add_params(
const struct ggml_tensor * dst,
struct htp_mm_kernel_params * kparams
) {
if (ggml_hexagon_matmul_add_can_collapse(src0, src1, src2, dst)) {
const ggml_tensor src1_collapsed = ggml_hexagon_tensor_collapse_rows(src1);
const ggml_tensor dst_collapsed = ggml_hexagon_tensor_collapse_rows(dst);
const ggml_tensor src2_collapsed = (src2 && (src2->ne[1] * src2->ne[2] * src2->ne[3] > 1))
? ggml_hexagon_tensor_collapse_rows(src2)
: (src2 ? *src2 : ggml_tensor{});
const struct ggml_tensor * p_src2 = src2 ? &src2_collapsed : nullptr;
const size_t src2_size = p_src2 ? hex_round_up(ggml_nbytes(p_src2), 128) : 0;
ggml_hexagon_precompute_matmul_params_impl(sess, src0, &src1_collapsed, &dst_collapsed, p_src2 ? p_src2->nb[1] : 0, src2_size, kparams);
kparams->collapse = 1;
return;
}
const size_t src2_size = src2 ? hex_round_up(ggml_nbytes(src2), 128) : 0;
ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, src2 ? src2->nb[1] : 0, src2_size, kparams);
}
@@ -6066,7 +6259,7 @@ static void ggml_hexagon_precompute_sort_params(
static void ggml_hexagon_precompute_fused_mmnx_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0, // W0
const struct ggml_tensor * src1, // x
const struct ggml_tensor * act, // x
int32_t n_weights,
struct htp_mm_kernel_params * kparams
) {
@@ -6078,23 +6271,24 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
const int ne02 = src0->ne[2];
const int ne03 = src0->ne[3];
const int ne10 = src1->ne[0];
const int ne11 = src1->ne[1];
const int ne12 = src1->ne[2];
const int ne13 = src1->ne[3];
const int ne10 = act->ne[0];
const int ne11 = act->ne[1];
const int ne12 = act->ne[2];
const int ne13 = act->ne[3];
const int wtype = src0->type;
const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype);
const int ne00_padded = is_repack ? hex_round_up(ne00, 32) : ne00;
const int ne01_padded = is_repack ? hex_round_up(ne01, 32) : ne01;
const int ne11_padded = hex_round_up(ne11, 32);
const int ne01_tiled = hex_round_up(ne01_padded, 32);
const size_t vtcm_budget = sess->vtcm_size;
const bool is_batched = (ne02 * ne03 > 1 || ne12 * ne13 > 1);
bool hmx_enabled = (sess->n_hmx > 0) && (opt_mm_select >= 2);
if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, src1, nullptr, ne01_padded, false, is_batched)) {
if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, nullptr, wtype, ne00_padded, ne01_padded, ne02, ne11, ne12, ne11_padded, false, is_batched, 0, vtcm_budget, kparams)) {
if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, act, nullptr, ne01_padded, false, is_batched)) {
if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, act, nullptr, wtype, ne00_padded, ne01_tiled, ne02, ne11, ne12, ne11_padded, false, is_batched, 0, vtcm_budget, kparams)) {
kparams->n_weights = n_weights;
goto finalize;
}
@@ -6106,20 +6300,20 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
}
{
const int src1_nrows = ne11 * ne12 * ne13;
const size_t src1_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
const int act_nrows = ne11 * ne12 * ne13;
const size_t act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
const size_t src0_row_size = src0->nb[1];
uint32_t best_n_prefetch = 16;
if (is_repack) {
const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
const uint32_t max_prefetch = (act_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
best_n_prefetch = 2;
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, src1_row_size, 0, d, false, true
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, act_nrows, sess->n_threads,
0, src0_row_size, act_row_size, 0, d, false, true
);
if (L.total_bytes <= sess->vtcm_size) {
best_n_prefetch = d;
@@ -6133,14 +6327,16 @@ static void ggml_hexagon_precompute_fused_mmnx_params(
// Test tiled first
htp_mm_hvx_vtcm_layout_build(
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true
&L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, act_nrows, sess->n_threads,
0, src0_row_size, act_row_size, 0, best_n_prefetch, false, true
);
if (try_tiled && L.total_bytes <= sess->vtcm_size) {
kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
kparams->act_row_size = act_row_size;
kparams->vtcm_src0_size = L.src0_bytes;
kparams->vtcm_src1_size = L.src1_bytes;
kparams->vtcm_act_size = L.act_bytes;
kparams->vtcm_bias_size = 0;
kparams->vtcm_dst_size = L.dst_bytes;
kparams->vtcm_size = L.total_bytes;
kparams->n_prefetch = best_n_prefetch;
@@ -6162,12 +6358,12 @@ finalize:
static void ggml_hexagon_precompute_fused_mmidnx_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0, // W0
const struct ggml_tensor * src1, // x
const struct ggml_tensor * act, // x
const struct ggml_tensor * dst, // dst0
int32_t n_weights,
struct htp_mm_kernel_params * kparams
) {
ggml_hexagon_precompute_matmul_params_impl(sess, src0, src1, dst, 0, 0, kparams);
ggml_hexagon_precompute_matmul_params_impl(sess, src0, act, dst, 0, 0, kparams);
kparams->n_weights = n_weights;
}
@@ -7137,6 +7333,14 @@ static bool mm_is_hmx_eligible(const ggml_tensor * t) {
const ggml_tensor * src0 = t->src[0];
const ggml_tensor * src1 = t->src[1];
if (ggml_hexagon_matmul_can_collapse(src0, src1, t)) {
const ggml_tensor src1_c = ggml_hexagon_tensor_collapse_rows(src1);
const int wtype = src0->type;
const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype);
const int ne01_padded = is_repack ? hex_round_up(src0->ne[1], 32) : src0->ne[1];
return ggml_hexagon_matmul_is_hmx_eligible(src0, &src1_c, t, ne01_padded, false, false);
}
const int wtype = src0->type;
const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype);
const bool is_matmul_id = (t->op == GGML_OP_MUL_MAT_ID);
@@ -7176,13 +7380,15 @@ static bool is_mergeable_mul_mat(const ggml_tensor * t) {
const ggml_tensor * src0 = t->src[0];
const ggml_tensor * src1 = t->src[1];
if (src1->type != GGML_TYPE_F32) return false;
if (src0->ne[2] != 1 || src0->ne[3] != 1) return false;
if (mm_is_hmx_eligible(t)) {
return ggml_hexagon_is_hmx_weight_type(src0->type);
}
// HVX path requires F32 activations and repacked weights (except Q6_K)
if (src1->type != GGML_TYPE_F32) return false;
return ggml_hexagon_is_repack_type(src0->type) && src0->type != GGML_TYPE_Q6_K;
}
@@ -8683,6 +8889,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
const char * str_nhmx = getenv("GGML_HEXAGON_NHMX");
const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT");
const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT");
const char * str_fa_head_split = getenv("GGML_HEXAGON_FA_HEAD_SPLIT");
const char * str_gdn_select = getenv("GGML_HEXAGON_GDN_SELECT");
const char * str_ar_select = getenv("GGML_HEXAGON_AR_SELECT");
const char * str_ar_scatter = getenv("GGML_HEXAGON_AR_SCATTER");
@@ -8737,6 +8944,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
opt_nhmx = str_nhmx ? atoi(str_nhmx) : opt_nhmx;
opt_mm_select = str_mm_select ? atoi(str_mm_select) : opt_mm_select;
opt_fa_select = str_fa_select ? atoi(str_fa_select) : opt_fa_select;
opt_fa_head_split = str_fa_head_split ? atoi(str_fa_head_split) : opt_fa_head_split;
opt_gdn_select = str_gdn_select ? atoi(str_gdn_select) : opt_gdn_select;
opt_ar_select = str_ar_select ? atoi(str_ar_select) : opt_ar_select;
opt_ar_scatter = str_ar_scatter ? atoi(str_ar_scatter) : opt_ar_scatter;
ggml/src/ggml-hexagon/htp/flash-attn-ops.c
+60 -29
View File
@@ -1869,8 +1869,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
factx.Bc = kparams->Bc;
factx.g_br = kparams->u.hmx.g_br;
factx.n_kv_blocks = kparams->n_kv_blocks;
factx.is_q_fp32 = (kparams->is_q_fp32 != 0);
factx.is_dst_fp32 = (kparams->is_dst_fp32 != 0);
factx.is_q_fp32 = (q->type == HTP_TYPE_F32);
factx.is_dst_fp32 = (dst->type == HTP_TYPE_F32);
factx.pipeline = (kparams->u.hmx.pipeline != 0);
factx.mask_broadcast = (kparams->u.hmx.mask_broadcast != 0);
if (mask) {
@@ -1879,13 +1879,12 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
}
factx.has_softcap = (kparams->logit_softcap != 0.0f);
if (!factx.has_softcap) {
factx.scale = (__fp16) (kparams->scale * EXP_LOG2E_F); // log2(e)
} else {
factx.scale = (__fp16) kparams->scale;
}
factx.scale = (__fp16) kparams->scale;
factx.max_bias = kparams->max_bias;
factx.logit_softcap = factx.has_softcap ? (__fp16) (kparams->logit_softcap * EXP_LOG2E_F) : 0;
factx.logit_softcap = 0;
if (factx.has_softcap) {
factx.logit_softcap = (__fp16) kparams->logit_softcap;
}
factx.n_head_log2 = kparams->n_head_log2;
factx.m0 = kparams->m0;
@@ -1898,22 +1897,36 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
const uint32_t n_threads = factx.n_threads;
const uint32_t G = factx.G;
// Multi-device: split Q blocks across devices
// Multi-device: prefer head-parallel partitioning (each core owns a disjoint head
// shard), falling back to Q-block (token) split when heads don't divide evenly.
const uint32_t n_q_blocks = (neq1 + Br - 1) / Br;
uint32_t q_start_min = 0;
uint32_t q_start_max = neq1;
uint32_t q_start_min = 0;
uint32_t q_start_max = neq1;
uint32_t kv_head_min = 0;
uint32_t kv_head_max = n_kv_heads;
if (octx->ctx->mdev.count > 1) {
const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(n_q_blocks, htp_tensor_mdev_data_aligned(dst) ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
const uint32_t block_start = range.start;
const uint32_t block_end = range.start + range.count;
const uint32_t mdev_count = octx->ctx->mdev.count;
const uint32_t mdev_idx = octx->ctx->mdev.idx;
const uint32_t dst_e_size = (dst->type == HTP_TYPE_F32) ? sizeof(float) : sizeof(__fp16);
const bool can_split = htp_tensor_can_row_partition(dst, dst_e_size);
if (block_start >= block_end) {
return HTP_STATUS_OK;
if (kparams->head_split && can_split && n_kv_heads >= mdev_count && n_kv_heads % mdev_count == 0) {
const uint32_t kv_per_core = n_kv_heads / mdev_count;
kv_head_min = mdev_idx * kv_per_core;
kv_head_max = kv_head_min + kv_per_core;
} else {
const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(n_q_blocks, can_split ? 1 : 0, mdev_idx, mdev_count, &octx->ctx->mdev.count_div);
const uint32_t block_start = range.start;
const uint32_t block_end = range.start + range.count;
if (block_start >= block_end) {
return HTP_STATUS_OK;
}
q_start_min = block_start * Br;
q_start_max = MIN(block_end * Br, neq1);
}
q_start_min = block_start * Br;
q_start_max = MIN(block_end * Br, neq1);
}
// ======== VTCM allocation (GQA-aware) ========
@@ -2032,7 +2045,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
const size_t g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
const size_t n_row_tiles = g_br_actual / HMX_FP16_TILE_N_ROWS;
for (uint32_t kv_head = 0; kv_head < n_kv_heads; ++kv_head) {
for (uint32_t kv_head = kv_head_min; kv_head < kv_head_max; ++kv_head) {
const uint32_t ik2 = kv_head;
const uint32_t ik3 = fastdiv(ib3, &kparams->broadcast_rk3);
const uint32_t iv2 = kv_head;
@@ -2040,7 +2053,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
// 1. Push Q and KV DMAs for the very first iteration.
// Subsequent iterations are enqueued early at the end of the previous iteration.
if (ib3 == 0 && q_start == q_start_min && kv_head == 0) {
if (ib3 == 0 && q_start == q_start_min && kv_head == kv_head_min) {
const dma_addr_t q_ptr = q->data + q_start * q->nb[1] +
(kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3];
const size_t q_row_bytes = q_transposed ? n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
@@ -2358,8 +2371,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
uint32_t next_kv_head = kv_head + 1;
uint32_t next_q_start = q_start;
uint32_t next_ib3 = ib3;
if (next_kv_head >= n_kv_heads) {
next_kv_head = 0;
if (next_kv_head >= kv_head_max) {
next_kv_head = kv_head_min;
next_q_start = q_start + Br;
if (next_q_start >= q_start_max) {
next_q_start = q_start_min;
@@ -2478,7 +2491,7 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
factx.src3_div3 = kparams->src3_div3;
}
factx.is_q_fp32 = (kparams->is_q_fp32 != 0);
factx.is_q_fp32 = (q->type == HTP_TYPE_F32);
factx.size_q_row_padded = kparams->u.hvx.size_q_row_padded;
factx.size_k_row_padded = kparams->u.hvx.size_k_row_padded;
factx.size_v_row_padded = kparams->u.hvx.size_v_row_padded;
@@ -2488,7 +2501,10 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
factx.scale = kparams->scale;
factx.max_bias = kparams->max_bias;
factx.has_softcap = (kparams->logit_softcap != 0.0f);
factx.logit_softcap = factx.has_softcap ? (__fp16) kparams->logit_softcap : 0;
factx.logit_softcap = 0;
if (factx.has_softcap) {
factx.logit_softcap = (__fp16) kparams->logit_softcap;
}
factx.n_head_log2 = kparams->n_head_log2;
factx.m0 = kparams->m0;
@@ -2512,10 +2528,25 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
uint32_t qrows = total_qrows;
if (octx->ctx->mdev.count > 1) {
const bool can_split = htp_tensor_mdev_data_aligned(dst) && ((dst->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0);
const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_qrows, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
qrow_start = range.start;
qrows = range.count;
const uint32_t mdev_count = octx->ctx->mdev.count;
const uint32_t mdev_idx = octx->ctx->mdev.idx;
const uint32_t n_kv_heads = k->ne[2];
const uint32_t dst_e_size = (dst->type == HTP_TYPE_F32) ? sizeof(float) : sizeof(__fp16);
const bool can_split = htp_tensor_can_row_partition(dst, dst_e_size);
// head range is contiguous in flat row space only when neq3 == 1
if (kparams->head_split && can_split && neq3 == 1 && n_kv_heads >= mdev_count && n_kv_heads % mdev_count == 0) {
const uint32_t G = kparams->G;
const uint32_t kv_per_core = n_kv_heads / mdev_count;
const uint32_t heads_per_core = kv_per_core * G;
const uint32_t head_start = mdev_idx * heads_per_core;
qrow_start = head_start * neq1;
qrows = heads_per_core * neq1;
} else {
const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_qrows, can_split ? 1 : 0, mdev_idx, mdev_count, &octx->ctx->mdev.count_div);
qrow_start = range.start;
qrows = range.count;
}
}
if (qrows == 0) {
ggml/src/ggml-hexagon/htp/flash-attn-ops.h
+2 -2
View File
@@ -34,8 +34,8 @@ enum htp_fa_kernel_type {
struct htp_fa_kernel_params {
uint8_t kernel_type; // enum htp_fa_kernel_type
uint8_t is_q_fp32; // 1 = Q type is F32, 0 = F16
uint8_t is_dst_fp32; // 1 = dst type is F32, 0 = F16
uint8_t head_split; // 1 = partition by KV heads in multicore, 0 = token partition
uint8_t flags; // reserved
uint8_t n_threads; // Number of threads to run
// Common parameters
ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h
+70 -5
View File
@@ -745,7 +745,7 @@ void convert_f16_weight_to_fp16_tiles_task(
const uint8_t *r0 = state->src + row0 * state->row_stride;
const uint8_t *r1 = state->src + row1 * state->row_stride;
HVX_Vector v0 = hvx_vmemu((const __fp16 *)(r0 + byte_off));
HVX_Vector v0 = (row0 < state->n_cols) ? hvx_vmemu((const __fp16 *)(r0 + byte_off)) : Q6_V_vzero();
HVX_Vector v1 = (row1 < state->n_cols) ? hvx_vmemu((const __fp16 *)(r1 + byte_off)) : Q6_V_vzero();
Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v0);
@@ -788,7 +788,7 @@ void quantize_f32_weight_to_fp16_tiles_task(
const uint8_t *r0 = state->src + row0 * state->row_stride;
const uint8_t *r1 = state->src + row1 * state->row_stride;
HVX_Vector v0_f32 = hvx_vmem((const float *)(r0 + byte_off));
HVX_Vector v0_f32 = (row0 < state->n_cols) ? hvx_vmem((const float *)(r0 + byte_off)) : Q6_V_vzero();
HVX_Vector v1_f32 = (row1 < state->n_cols) ? hvx_vmem((const float *)(r1 + byte_off)) : Q6_V_vzero();
HVX_Vector v_out = hvx_vec_f32_to_f16(v0_f32, v1_f32);
@@ -988,9 +988,7 @@ static void transfer_output_chunk_fp16_to_fp32_col_chunk(
uint32_t src2_stride,
uint32_t dst_cols
) {
assert(c_len % HTP_MM_HMX_TILE_N_COLS == 0);
assert(total_n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
const size_t tile_row_stride = (total_n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;
const size_t tile_row_stride = hmx_ceil_div(total_n_cols, HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;
const HVX_Vector one = hvx_vec_splat_f16(1.0);
@@ -1137,6 +1135,73 @@ static void transfer_activation_row_pair_fp32_to_fp16(
}
}
// F16-input variant of transfer_activation_row_pair_fp32_to_fp16, for F16 activation
// (src1). Same shape as the F16 Q-prep in hmx-fa-kernels.h: one 128-byte load carries 64 f16
// = two tile columns, and Q6_W_vshuff_VVR interleaves the two rows straight into the HMX tile
// layout, so no F32 round-trip is needed. Rows are only 64-byte aligned when k_block is an odd
// multiple of the tile width, hence the unaligned load type.
static void transfer_activation_row_pair_f16_to_f16(__fp16 * restrict vtcm_dst,
const __fp16 * restrict row0,
const __fp16 * restrict row1,
uint32_t r,
uint32_t k_block,
uint32_t k_valid,
bool row0_valid,
bool row1_valid) {
uint32_t r0 = r / HTP_MM_HMX_TILE_N_ROWS; // tile row index
uint32_t r1 = r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
const uint32_t n_tile_cols = k_block / HTP_MM_HMX_TILE_N_COLS;
__fp16 * restrict tile_row = vtcm_dst + (size_t) r0 * n_tile_cols * HTP_MM_HMX_TILE_N_ELMS;
const HVX_UVector * pv0 = row0_valid ? (const HVX_UVector *) row0 : NULL;
const HVX_UVector * pv1 = row1_valid ? (const HVX_UVector *) row1 : NULL;
uint32_t c = 0;
for (; c + 64 <= k_valid; c += 64) {
HVX_Vector v0 = pv0 ? pv0[c / 64] : Q6_V_vzero();
HVX_Vector v1 = pv1 ? pv1[c / 64] : Q6_V_vzero();
HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS;
HVX_Vector * tile0 = (HVX_Vector *) (tile_row + (size_t) c0 * HTP_MM_HMX_TILE_N_ELMS);
HVX_Vector * tile1 = (HVX_Vector *) (tile_row + (size_t) (c0 + 1) * HTP_MM_HMX_TILE_N_ELMS);
tile0[r1 / 2] = Q6_V_lo_W(vp);
tile1[r1 / 2] = Q6_V_hi_W(vp);
}
// Tail: fewer than 64 valid columns left, plus the k_valid..k_block padding that HMX will
// still multiply, so it has to be written as zeros.
for (; c < k_block; c += 64) {
HVX_Vector v0 = Q6_V_vzero();
HVX_Vector v1 = Q6_V_vzero();
if (c < k_valid) {
uint32_t rem = k_valid - c; // 1..63 valid f16 lanes
HVX_VectorPred mask = Q6_Q_vsetq2_R(rem * sizeof(__fp16));
if (pv0) {
v0 = Q6_V_vmux_QVV(mask, pv0[c / 64], Q6_V_vzero());
}
if (pv1) {
v1 = Q6_V_vmux_QVV(mask, pv1[c / 64], Q6_V_vzero());
}
}
HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS;
HVX_Vector * tile0 = (HVX_Vector *) (tile_row + (size_t) c0 * HTP_MM_HMX_TILE_N_ELMS);
tile0[r1 / 2] = Q6_V_lo_W(vp);
if (c0 + 1 < n_tile_cols) {
HVX_Vector * tile1 = (HVX_Vector *) (tile_row + (size_t) (c0 + 1) * HTP_MM_HMX_TILE_N_ELMS);
tile1[r1 / 2] = Q6_V_hi_W(vp);
}
}
}
static void transfer_activation_row_pair_fp32_to_fp16_col_chunk(
__fp16 *restrict vtcm_dst,
const float *restrict row0, // offset by c_first
ggml/src/ggml-hexagon/htp/hmx-utils.h
+18 -10
View File
@@ -73,19 +73,20 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
for (uint32_t r = start_row; r < end_row; r += 2) {
const uint32_t ct = r / HMX_FP16_TILE_N_ROWS;
const uint32_t local_r = r % HMX_FP16_TILE_N_ROWS;
const bool row0_valid = r < n_cols;
const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_cols;
const HVX_Vector v_off0 = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4));
const HVX_Vector v_off1 = Q6_Vw_vadd_VwVw(v_off0, v_scat_step);
__fp16 * tile_base = vtcm_dst + (size_t) ct * n_k_tiles * HMX_FP16_TILE_N_ELMS;
const uint8_t * p0 = (const uint8_t *) (vtcm_src + r * src_stride);
const uint8_t * p0 = row0_valid ? (const uint8_t *) (vtcm_src + r * src_stride) : NULL;
const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL;
assert(hex_is_aligned(p0, 128));
assert(hex_is_aligned(p1, 128));
assert(!p0 || hex_is_aligned(p0, 128));
assert(!p1 || hex_is_aligned(p1, 128));
assert(c_byte_step % 128 == 0);
if (p1) {
if (p0 && p1) {
for (uint32_t i = 0; i < n_c_iters; ++i) {
HVX_Vector v0 = hvx_vmem(p0); p0 += c_byte_step;
HVX_Vector v1 = hvx_vmem(p1); p1 += c_byte_step;
@@ -96,9 +97,12 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
} else {
const HVX_Vector vzero = Q6_V_vzero();
for (uint32_t i = 0; i < n_c_iters; ++i) {
HVX_Vector v0 = hvx_vmem(p0); p0 += c_byte_step;
HVX_Vector v0 = p0 ? hvx_vmem(p0) : vzero;
if (p0) p0 += c_byte_step;
HVX_Vector v1 = p1 ? hvx_vmem(p1) : vzero;
if (p1) p1 += c_byte_step;
Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0);
Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, vzero);
Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, v1);
tile_base += dst_step;
}
}
@@ -113,15 +117,16 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
for (uint32_t r = start_row; r < end_row; r += 2) {
const uint32_t ct = r / HMX_FP16_TILE_N_ROWS;
const uint32_t local_r = r % HMX_FP16_TILE_N_ROWS;
const bool row0_valid = r < n_cols;
const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_cols;
const HVX_Vector v_off0 = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4));
const HVX_Vector v_off1 = Q6_Vw_vadd_VwVw(v_off0, v_scat_step);
__fp16 * tile_base = vtcm_dst + (size_t) ct * n_k_tiles * HMX_FP16_TILE_N_ELMS;
const uint8_t * p0 = (const uint8_t *) (vtcm_src + r * src_stride);
const uint8_t * p0 = row0_valid ? (const uint8_t *) (vtcm_src + r * src_stride) : NULL;
const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL;
if (p1) {
if (p0 && p1) {
for (uint32_t i = 0; i < n_c_iters; ++i) {
HVX_Vector v0 = hvx_vmemu(p0); p0 += c_byte_step;
HVX_Vector v1 = hvx_vmemu(p1); p1 += c_byte_step;
@@ -132,9 +137,12 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
} else {
const HVX_Vector vzero = Q6_V_vzero();
for (uint32_t i = 0; i < n_c_iters; ++i) {
HVX_Vector v0 = hvx_vmemu(p0); p0 += c_byte_step;
HVX_Vector v0 = p0 ? hvx_vmemu(p0) : vzero;
if (p0) p0 += c_byte_step;
HVX_Vector v1 = p1 ? hvx_vmemu(p1) : vzero;
if (p1) p1 += c_byte_step;
Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0);
Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, vzero);
Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, v1);
tile_base += dst_step;
}
}
ggml/src/ggml-hexagon/htp/matmul-ops.c
+481 -335
View File
File diff suppressed because it is too large Load Diff
ggml/src/ggml-hexagon/htp/matmul-ops.h
+86 -85
View File
@@ -87,24 +87,26 @@ enum htp_mm_kernel_type {
// Op-specific struct for precomputed matmul params
struct htp_mm_kernel_params {
int32_t kernel_type; // enum htp_mm_kernel_type
int32_t pipeline; // 1 = pipelined execution, 0 = standard
uint8_t kernel_type; // enum htp_mm_kernel_type
uint8_t pipeline; // 1 = pipelined execution, 0 = standard
uint8_t collapse; // 1 = collapse outer dims into 2D, 0 = standard
uint8_t n_hmx; // 1 = use HMX, 0 = use HVX
uint8_t n_threads; // Number of threads to spawn
uint8_t n_act_threads; // Number of threads for activation preparation
uint8_t n_prefetch; // Prefetch lookahead buffers/rows in VTCM
uint8_t n_weights; // Number of weights for fused NX
int32_t m_chunk; // Row chunk size (M chunk)
int32_t n_chunk; // Col chunk size (N chunk)
int32_t n_threads; // Number of threads to spawn
int32_t n_act_threads; // Number of threads for activation preparation
int32_t n_hmx; // 1 = use HMX, 0 = use HVX
int32_t n_prefetch; // Prefetch lookahead buffers/rows in VTCM
int32_t tile_size; // Weight tile size
int32_t aligned_tile_size; // Aligned weight tile size (padded to 128)
int32_t src1_row_size; // Row size for quantized activation
int32_t act_row_size; // Row size for activation scratchpad
int32_t vtcm_size; // Total required scratchpad size in VTCM
int32_t vtcm_src0_size; // src0 scratchpad size in VTCM
int32_t vtcm_src1_size; // src1 scratchpad size in VTCM
int32_t vtcm_src2_size; // src2 scratchpad size in VTCM (fused only)
int32_t vtcm_src3_size; // src3 scratchpad size in VTCM (fused only)
int32_t vtcm_act_size; // activation scratchpad size in VTCM
int32_t vtcm_bias_size; // bias scratchpad size in VTCM (fused only)
int32_t vtcm_dst_size; // dst scratchpad size in VTCM
int32_t n_weights; // Number of weights for fused NX
// Precomputed division values
struct fastdiv_values div_ne12_ne1;
@@ -147,6 +149,7 @@ static inline int htp_mm_hmx_compute_chunks(size_t vtcm_total,
const size_t usable = vtcm_total - overhead;
size_t best_cost = SIZE_MAX;
size_t best_tail_waste = SIZE_MAX;
size_t best_mn = 0;
size_t best_m = 0, best_n = 0;
@@ -173,12 +176,17 @@ static inline int htp_mm_hmx_compute_chunks(size_t vtcm_total,
size_t mblocks = ((size_t) m + mc - 1) / mc;
size_t nblocks = ((size_t) n + nc - 1) / nc;
size_t cost = mblocks * m_block_cost + nblocks * n_block_cost;
size_t rem = n % nc;
size_t tail_waste = (rem == 0) ? 0 : (nc - rem);
size_t mn = mc * nc;
if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
best_cost = cost;
best_mn = mn;
best_m = mc;
best_n = nc;
if (cost < best_cost ||
(cost == best_cost && tail_waste < best_tail_waste) ||
(cost == best_cost && tail_waste == best_tail_waste && mn > best_mn)) {
best_cost = cost;
best_tail_waste = tail_waste;
best_mn = mn;
best_m = mc;
best_n = nc;
}
}
@@ -349,7 +357,7 @@ struct htp_mm_hmx_vtcm_layout {
size_t off_dst[2]; // [1] is only used when pipelined
size_t off_scratch[2]; // dequantization scratch pads
size_t off_scales; // HMX scales (256 bytes)
size_t off_src2; // src2 bias in VTCM
size_t off_bias; // bias in VTCM
// Cached sizes of regions for HMX kernel use
size_t weight_area_bytes;
@@ -358,25 +366,23 @@ struct htp_mm_hmx_vtcm_layout {
size_t output_area_bytes;
size_t scratch_bytes[2];
size_t act_head_stride;
size_t src2_bytes;
size_t bias_bytes;
size_t total_bytes;
};
struct htp_mm_hvx_vtcm_layout {
// Byte offsets from vtcm_base for each region
size_t off_src1; // vtcm_src1 (activation)
size_t off_act; // vtcm_act (activation)
size_t off_src0; // vtcm_src0 (weight/Wk)
size_t off_src2; // vtcm_src2 (Wq / fused only)
size_t off_src3; // vtcm_src3 (Wv / fused only)
size_t off_bias; // vtcm_bias (bias / fused add only)
size_t off_dst; // vtcm_dst (output scratch)
size_t off_act_raw; // vtcm_act_raw (raw activation DMA staging)
// Cached sizes
size_t src0_bytes;
size_t src1_bytes;
size_t src2_bytes;
size_t src3_bytes;
size_t act_bytes;
size_t bias_bytes;
size_t dst_bytes;
size_t act_raw_bytes;
@@ -394,7 +400,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
bool pipeline,
uint32_t act_threads,
uint32_t aligned_tile_size,
size_t src2_size
size_t bias_size
) {
size_t off = 0;
@@ -411,7 +417,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
size_t off_group_a = 0;
VTCM_LAYOUT_ALLOC(off_group_a, off_act, activation_area_size);
VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src2, hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE), src2_size > 0);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_bias, hex_align_up(bias_size, HTP_MM_HMX_TILE_SIZE), bias_size > 0);
// Group B: Compute-only buffers (starts at off_group_a)
size_t off_group_b = off_group_a;
@@ -439,7 +445,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
L->scratch_bytes[0] = scratch_area_size;
L->scratch_bytes[1] = scratch_area_size;
L->act_head_stride = act_head_stride;
L->src2_bytes = src2_size;
L->bias_bytes = bias_size;
off = off_group_a + hex_smax(group_b_size, group_c_size);
} else {
@@ -463,7 +469,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
size_t off_group_a = 0;
VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data
VTCM_LAYOUT_ALLOC(off_group_a, off_act, act_area_size);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src2, hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE), src2_size > 0);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_bias, hex_align_up(bias_size, HTP_MM_HMX_TILE_SIZE), bias_size > 0);
// Group B: Compute-only buffers (starts at off_group_a)
size_t off_group_b = off_group_a;
@@ -491,7 +497,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
L->scratch_bytes[0] = scratch0_size;
L->scratch_bytes[1] = scratch1_size;
L->act_head_stride = 0;
L->src2_bytes = src2_size;
L->bias_bytes = bias_size;
off = off_group_a + hex_smax(group_b_size, group_c_size);
}
@@ -504,21 +510,20 @@ static inline void htp_mm_hvx_vtcm_layout_build(
int kernel_type,
int wtype,
uint32_t ne10, // k
uint32_t src1_nrows, // m_total
uint32_t act_nrows, // m_total
uint32_t n_threads,
size_t dst_row_size,
size_t src0_row_size,
size_t src1_row_size,
size_t src2_row_size,
size_t act_row_size,
size_t bias_row_size,
uint32_t n_prefetch,
bool is_matmul_id,
bool is_fused_nx
) {
(void)src1_row_size;
(void)act_row_size;
size_t src0_sz = 0;
size_t src1_sz = 0;
size_t src2_sz = src2_row_size > 0 ? htp_mm_round_up(src2_row_size, 128) : 0;
size_t src3_sz = 0;
size_t act_sz = 0;
size_t bias_sz = bias_row_size > 0 ? htp_mm_round_up(bias_row_size, 128) : 0;
size_t dst_sz = 0;
size_t act_raw_sz = 0;
@@ -544,22 +549,21 @@ static inline void htp_mm_hvx_vtcm_layout_build(
}
size_t tiled_act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
size_t act_sz = hex_round_up(tiled_act_row_size * src1_nrows, 128);
size_t q_act_sz = hex_round_up(tiled_act_row_size * act_nrows, 128);
size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
src0_sz = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer
src1_sz = act_sz; // quantized activation buffer
src2_sz = 0;
src3_sz = 0;
act_sz = q_act_sz; // quantized activation buffer
bias_sz = 0;
dst_sz = 0;
act_raw_sz = hex_round_up(raw_row_size * src1_nrows, 128);
act_raw_sz = hex_round_up(raw_row_size * act_nrows, 128);
} else if (is_matmul_id) {
const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
const size_t src1_row_size_tiled = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10)
: htp_mm_q8_0_tiled_row_size(ne10);
const size_t act_row_size_tiled = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10)
: htp_mm_q8_0_tiled_row_size(ne10);
size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
src1_sz = htp_mm_round_up(src1_row_size_tiled * src1_nrows, 256);
act_sz = htp_mm_round_up(act_row_size_tiled * act_nrows, 256);
if (is_repack) {
const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
@@ -573,25 +577,24 @@ static inline void htp_mm_hvx_vtcm_layout_build(
src0_sz = src0_sz_per_thread * n_threads;
dst_sz = 0;
src2_sz = 0;
src3_sz = 0;
act_raw_sz = hex_round_up(raw_row_size * src1_nrows, 128);
bias_sz = 0;
act_raw_sz = hex_round_up(raw_row_size * act_nrows, 128);
} else {
const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
const size_t dst_nrows = (src1_nrows > 1) ? 0 : 1;
const size_t dst_nrows = (act_nrows > 1) ? 0 : 1;
switch (kernel_type) {
case HTP_MM_KERNEL_HVX_F16_F16_VTCM: {
size_t f16_src1_row_size = htp_mm_round_up(ne10 * 2, 128);
src1_sz = htp_mm_round_up(f16_src1_row_size * src1_nrows, 256);
size_t f16_act_row_size = htp_mm_round_up(ne10 * 2, 128);
act_sz = htp_mm_round_up(f16_act_row_size * act_nrows, 256);
src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
act_raw_sz = hex_round_up(hex_round_up(ne10 * sizeof(float), 128) * src1_nrows, 128);
act_raw_sz = hex_round_up(hex_round_up(ne10 * sizeof(float), 128) * act_nrows, 128);
break;
}
case HTP_MM_KERNEL_HVX_F32_F32_VTCM: {
size_t f32_src1_row_size = htp_mm_round_up(ne10 * 4, 128);
src1_sz = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256);
size_t f32_act_row_size = htp_mm_round_up(ne10 * 4, 128);
act_sz = htp_mm_round_up(f32_act_row_size * act_nrows, 256);
src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
act_raw_sz = 0;
@@ -599,10 +602,10 @@ static inline void htp_mm_hvx_vtcm_layout_build(
}
case HTP_MM_KERNEL_HVX_QUANT_BLOCK:
case HTP_MM_KERNEL_HVX_QUANT_ROW: {
size_t q_src1_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
size_t q_act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
src1_sz = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
act_sz = htp_mm_round_up(q_act_row_size * act_nrows, 256);
src0_sz = src0_sz * n_threads;
@@ -614,10 +617,10 @@ static inline void htp_mm_hvx_vtcm_layout_build(
src0_sz = repacked_vtcm_size * n_threads;
}
size_t dst_slice_per_thread = (dst_nrows > 0 && src1_nrows == 1) ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0;
size_t dst_slice_per_thread = (dst_nrows > 0 && act_nrows == 1) ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0;
dst_sz = dst_slice_per_thread * n_threads;
size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
act_raw_sz = hex_round_up(raw_row_size * src1_nrows, 128);
act_raw_sz = hex_round_up(raw_row_size * act_nrows, 128);
break;
}
default:
@@ -627,9 +630,8 @@ static inline void htp_mm_hvx_vtcm_layout_build(
// Group A: Persistent buffers across chunk compute
size_t off_group_a = 0;
VTCM_LAYOUT_ALLOC(off_group_a, off_src1, src1_sz);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src2, src2_sz, src2_sz > 0);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src3, src3_sz, src3_sz > 0);
VTCM_LAYOUT_ALLOC(off_group_a, off_act, act_sz);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_bias, bias_sz, bias_sz > 0);
// Group B: Compute-only buffers (starts at off_group_a)
size_t off_group_b = off_group_a;
@@ -643,9 +645,8 @@ static inline void htp_mm_hvx_vtcm_layout_build(
const size_t group_c_size = off_group_c - off_group_a;
L->src0_bytes = src0_sz;
L->src1_bytes = src1_sz;
L->src2_bytes = src2_sz;
L->src3_bytes = src3_sz;
L->act_bytes = act_sz;
L->bias_bytes = bias_sz;
L->dst_bytes = dst_sz;
L->act_raw_bytes = act_raw_sz;
L->total_bytes = off_group_a + hex_smax(group_b_size, group_c_size);
@@ -655,12 +656,12 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
int kernel_type,
int wtype,
uint32_t ne10,
uint32_t src1_nrows,
uint32_t act_nrows,
uint32_t n_threads,
size_t dst_row_size,
size_t src0_row_size,
size_t src1_row_size,
size_t src2_row_size,
size_t act_row_size,
size_t bias_row_size,
uint32_t n_prefetch,
size_t vtcm_budget,
struct htp_mm_hvx_vtcm_layout * L_out,
@@ -668,17 +669,17 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
) {
struct htp_mm_hvx_vtcm_layout L;
htp_mm_hvx_vtcm_layout_build(
&L, kernel_type, wtype, ne10, src1_nrows, n_threads,
dst_row_size, src0_row_size, src1_row_size, src2_row_size, n_prefetch, false, false
&L, kernel_type, wtype, ne10, act_nrows, n_threads,
dst_row_size, src0_row_size, act_row_size, bias_row_size, n_prefetch, false, false
);
if (L.total_bytes <= vtcm_budget) {
*L_out = L;
*m_chunk_out = src1_nrows;
*m_chunk_out = act_nrows;
return true;
}
const size_t fixed_bytes = L.src0_bytes + L.src2_bytes + L.dst_bytes;
const size_t fixed_bytes = L.src0_bytes + L.bias_bytes + L.dst_bytes;
if (vtcm_budget <= fixed_bytes) {
return false;
}
@@ -707,8 +708,8 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
if (m_chunk > 1) {
m_chunk &= ~1U;
}
if (m_chunk > src1_nrows) {
m_chunk = src1_nrows;
if (m_chunk > act_nrows) {
m_chunk = act_nrows;
}
if (m_chunk < 1) {
return false;
@@ -716,14 +717,14 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
htp_mm_hvx_vtcm_layout_build(
&L, kernel_type, wtype, ne10, m_chunk, n_threads,
dst_row_size, src0_row_size, src1_row_size, src2_row_size, n_prefetch, false, false
dst_row_size, src0_row_size, act_row_size, bias_row_size, n_prefetch, false, false
);
while (m_chunk > 2 && L.total_bytes > vtcm_budget) {
m_chunk -= 2;
htp_mm_hvx_vtcm_layout_build(
&L, kernel_type, wtype, ne10, m_chunk, n_threads,
dst_row_size, src0_row_size, src1_row_size, src2_row_size, n_prefetch, false, false
dst_row_size, src0_row_size, act_row_size, bias_row_size, n_prefetch, false, false
);
}
@@ -737,18 +738,18 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
}
static inline size_t htp_mm_hmx_get_2d_vtcm_size(
int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size, size_t src2_size
int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size, size_t bias_size
) {
struct htp_mm_hmx_vtcm_layout L;
htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, pipeline, act_threads, aligned_tile_size, src2_size);
htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, pipeline, act_threads, aligned_tile_size, bias_size);
return L.total_bytes;
}
static inline size_t htp_mm_hmx_get_batched_vtcm_size(
int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool pipeline, uint32_t act_threads, size_t src2_size) {
int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool pipeline, uint32_t act_threads, size_t bias_size) {
(void)pipeline;
struct htp_mm_hmx_vtcm_layout L;
htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, false, act_threads, 0, src2_size);
htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, false, act_threads, 0, bias_size);
return L.total_bytes;
}
@@ -760,7 +761,7 @@ static inline bool htp_mm_hmx_solve_batched_params(
uint32_t group_size,
int n_threads,
bool pipeline,
size_t src2_size,
size_t bias_size,
size_t vtcm_budget,
size_t * m_chunk_out,
size_t * n_chunk_out,
@@ -775,7 +776,7 @@ static inline bool htp_mm_hmx_solve_batched_params(
int act_threads = n_threads;
while (act_threads >= 1) {
size_t group_overhead = htp_mm_hmx_get_batched_overhead() + (src2_size > 0 ? hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE) : 0);
size_t group_overhead = htp_mm_hmx_get_batched_overhead() + (bias_size > 0 ? hex_align_up(bias_size, HTP_MM_HMX_TILE_SIZE) : 0);
size_t group_size_per_n, group_size_per_m, group_size_per_mn;
htp_mm_hmx_get_batched_chunk_costs(k, group_size, &group_size_per_n, &group_size_per_m, &group_size_per_mn);
@@ -785,8 +786,8 @@ static inline bool htp_mm_hmx_solve_batched_params(
if (htp_mm_hmx_compute_chunks(vtcm_budget, group_overhead, group_size_per_n, group_size_per_m, group_size_per_mn, hex_align_up(ne11, 32), ne01_padded,
(size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) ne11 * HTP_MM_HMX_COST_A_CONVERT,
&m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, pipeline, act_threads, src2_size);
&m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, pipeline, act_threads, bias_size);
if (exact_size <= vtcm_budget) {
size_t mblocks = ((size_t) ne11 + m_chunk_candidate - 1) / m_chunk_candidate;
if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) {
@@ -826,7 +827,7 @@ static inline bool htp_mm_hmx_solve_2d_params(
bool pipeline,
bool is_matmul_id,
uint32_t aligned_tile_size,
size_t src2_size,
size_t bias_size,
size_t vtcm_budget,
size_t * m_chunk_out,
size_t * n_chunk_out,
@@ -843,7 +844,7 @@ static inline bool htp_mm_hmx_solve_2d_params(
int act_threads = n_threads;
while (act_threads >= 1) {
size_t simple_2d_overhead = htp_mm_hmx_get_2d_overhead(pipeline, is_matmul_id) + (src2_size > 0 ? hex_align_up(src2_size, HTP_MM_HMX_TILE_SIZE) : 0);
size_t simple_2d_overhead = htp_mm_hmx_get_2d_overhead(pipeline, is_matmul_id) + (bias_size > 0 ? hex_align_up(bias_size, HTP_MM_HMX_TILE_SIZE) : 0);
size_t simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn;
htp_mm_hmx_get_2d_chunk_costs(wtype, k, pipeline, aligned_tile_size, &simple_2d_size_per_n, &simple_2d_size_per_m, &simple_2d_size_per_mn);
@@ -854,7 +855,7 @@ static inline bool htp_mm_hmx_solve_2d_params(
if (htp_mm_hmx_compute_chunks(vtcm_budget, simple_2d_overhead, simple_2d_size_per_n, simple_2d_size_per_m, simple_2d_size_per_mn, m_for_chunks, ne01_padded,
(size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) m_for_cost * HTP_MM_HMX_COST_A_CONVERT,
&m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
size_t exact_size = htp_mm_hmx_get_2d_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, pipeline, is_matmul_id ? 0 : act_threads, aligned_tile_size, src2_size);
size_t exact_size = htp_mm_hmx_get_2d_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, pipeline, is_matmul_id ? 0 : act_threads, aligned_tile_size, bias_size);
if (exact_size <= vtcm_budget) {
size_t mblocks = ((size_t) m_for_cost + m_chunk_candidate - 1) / m_chunk_candidate;
if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) {