Commit 18a04f09c for llama.cpp
commit 18a04f09c24616898792bcfaa17f3550bdc78912
Author: Todor Boinovski <todorb@qti.qualcomm.com>
Date: Fri Sep 18 13:15:08 2026 -0700
hexagon: HMX flash-attention head_dim padding (support DK=DV=72) (#26539)
Allow HMX flash-attention to run with head_dim not a multiple of 64
(e.g. SigLIP head_dim=72), by operating on DK/DV rounded up to 64 with
zero-filled tail lanes.
diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
index 3f1495645..af8013b08 100644
--- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp
+++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
@@ -4000,7 +4000,9 @@ static bool ggml_hexagon_flash_attn_is_hmx_eligible(
const uint32_t DK = q->ne[0];
const uint32_t DV = v->ne[0];
- if (DK % 64 != 0 || DV % 64 != 0) {
+ // Head dims that are not multiples of 64 are handled by internally padding to
+ // DK_pad/DV_pad = round_up(.,64) and zero-filling the tail lanes.
+ if (DK % 8 != 0 || DV % 8 != 0) {
return false;
}
@@ -4073,8 +4075,13 @@ static bool ggml_hexagon_precompute_flash_attn_params(
// Check HMX eligibility
const struct ggml_tensor * sinks = op->src[4];
if (ggml_hexagon_flash_attn_is_hmx_eligible(sess, q, k, v, sinks)) {
+ // HMX tiles head_dim in units of 64; when DK/DV are not 64-aligned the kernel
+ // operates on padded dims with zero-filled tail lanes. VTCM budget and chunk-size
+ // are sized for the padded tiles.
+ 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, DV, neq1, nek1, sess->vtcm_size, sess->n_threads, kparams->is_q_fp32 != 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);
if (ret == 0) {
kparams->kernel_type = HTP_FA_KERNEL_HMX;
kparams->Br = Br;
@@ -4084,7 +4091,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, DV, Br, Bc, kparams->n_threads, kparams->u.hmx.pipeline != 0, kparams->is_q_fp32 != 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);
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
diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c
index 8a1caba22..75422f420 100644
--- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c
+++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c
@@ -108,6 +108,7 @@ struct hmx_fa_context {
// Dimensions
uint32_t DK, DV;
+ uint32_t DK_pad, DV_pad; // head_dim rounded up to 64 for HMX tiling
uint32_t n_kv; // kv_len
uint32_t n_kv_heads; // number of KV heads
uint32_t n_heads; // number of Q heads
@@ -652,7 +653,7 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
hvx_dequantize_row_q8_0_f16(row_k, row_k, factx->DK);
}
}
- hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK,
+ hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK_pad,
args->src_stride, start, end);
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
}
@@ -706,7 +707,7 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data)
hvx_dequantize_row_q8_0_f16(row_v, row_v, factx->DV);
}
}
- hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV,
+ hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV_pad,
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));
}
@@ -832,17 +833,22 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
const uint32_t kv_head = args->kv_head;
const uint32_t ib3 = args->ib3;
- assert(factx->DK == factx->DV);
-
const bool use_q_dma = (factx->vtcm_q_dma != NULL);
__fp16 * q_tiles = factx->vtcm_q_tiles;
+ const size_t DK_pad = factx->DK_pad;
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_q_dma;
- if (factx->is_q_fp32) {
+ if (DK_pad != DK) {
+ if (factx->is_q_fp32) {
+ hmx_fa_q_prep_fp32_pad(q_tiles, q_flat, start, end, g_rows_end, DK, DK_pad, G, args->n_rows_q, &factx->div_G, args->q_transposed);
+ } else {
+ hmx_fa_q_prep_fp16_pad(q_tiles, q_flat, start, end, g_rows_end, DK, DK_pad, G, args->n_rows_q, &factx->div_G, args->q_transposed);
+ }
+ } else 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;
@@ -858,7 +864,7 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
} 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);
+ q_start, kv_head, ib3, start, end, n_rows_g, G, DK, DK_pad, factx->is_q_fp32, &factx->div_G);
}
}
@@ -952,6 +958,8 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
const uint32_t kv_head = args->kv_head;
const uint32_t ib3 = args->ib3;
+ const size_t DV_pad = factx->DV_pad;
+
size_t q_idx = fastdiv(start, &factx->div_G);
size_t h_idx = fastmodulo(start, G, &factx->div_G);
@@ -961,7 +969,7 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
size_t r0 = r / HMX_FP16_TILE_N_ROWS;
size_t r1 = r % HMX_FP16_TILE_N_ROWS;
- const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV;
+ const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV_pad;
for (uint32_t d = 0; d < DV / 32; ++d) {
const HVX_Vector * in_tile = (const HVX_Vector *) (tile_row_base + d * HMX_FP16_TILE_N_ELMS);
@@ -972,6 +980,16 @@ static void fa_o_store_thread_f32(unsigned int n, unsigned int i, void * data) {
*(HVX_UVector *) (out + d * 32) = Q6_V_hi_W(vp);
}
}
+ // Ragged tail: DV not a multiple of 32 (e.g. 72 -> last 8 lanes). Partial vector-write
+ // for the remaining (DV % 32) floats.
+ const uint32_t d_tail = DV / 32;
+ const uint32_t rem = DV - d_tail * 32;
+ if (rem) {
+ const HVX_Vector * in_tile = (const HVX_Vector *) (tile_row_base + d_tail * HMX_FP16_TILE_N_ELMS);
+ HVX_VectorPair vp = hvx_vec_f16_to_f32_shuff(in_tile[r1 / 2]);
+ HVX_Vector vd = (r1 % 2 == 0) ? Q6_V_lo_W(vp) : Q6_V_hi_W(vp);
+ hvx_vec_store_u((void *) (out + d_tail * 32), rem * sizeof(float), vd);
+ }
h_idx++;
if (h_idx == G) {
@@ -1006,6 +1024,9 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
const uint32_t kv_head = args->kv_head;
const uint32_t ib3 = args->ib3;
+ // O-tiles use the padded head dim (DV_pad); dst holds the real DV lanes.
+ const size_t DV_pad = factx->DV_pad;
+
size_t q_idx = fastdiv(start, &factx->div_G);
size_t h_idx = fastmodulo(start, G, &factx->div_G);
@@ -1015,7 +1036,7 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
size_t r0 = r / HMX_FP16_TILE_N_ROWS;
size_t r1 = r % HMX_FP16_TILE_N_ROWS;
- const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV;
+ const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV_pad;
for (uint32_t d = 0; d < DV / 64; ++d) {
const __fp16 * in_dtile = tile_row_base + d * HMX_FP16_TILE_N_ELMS * 2;
@@ -1028,6 +1049,17 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
*(HVX_UVector *) (out + d * 64) = Q6_V_hi_W(vp);
}
}
+ // Ragged tail when DV is not a multiple of 64.
+ const uint32_t d_tail = DV / 64;
+ const uint32_t rem = DV - d_tail * 64;
+ if (rem) {
+ const __fp16 * in_dtile = tile_row_base + d_tail * 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);
+ HVX_Vector vd = (r1 % 2 == 0) ? Q6_V_lo_W(vp) : Q6_V_hi_W(vp);
+ hvx_vec_store_u((void *) (out + d_tail * 64), rem * sizeof(__fp16), vd);
+ }
h_idx++;
if (h_idx == G) {
@@ -1829,8 +1861,11 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
const uint32_t DK = neq0;
const uint32_t DV = nev0;
- // HMX requires head_dim to be multiple of 32
- if (DK % 32 != 0 || DV % 32 != 0) {
+ // HMX tiles head_dim in units of 64. head_dim need not be 64- (or 32-) aligned:
+ // we can operate on DK/DV rounded up to 64 with tail lanes [D, D_pad) zero-filled.
+ const uint32_t DK_pad = hex_round_up(DK, 64);
+ const uint32_t DV_pad = hex_round_up(DV, 64);
+ if (DK == 0 || DV == 0) {
return HTP_STATUS_NO_SUPPORT;
}
@@ -1847,6 +1882,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
factx.n_threads = kparams->n_threads;
factx.DK = DK;
factx.DV = DV;
+ factx.DK_pad = DK_pad;
+ factx.DV_pad = DV_pad;
factx.n_kv = nek1;
factx.n_kv_heads = n_kv_heads;
factx.n_heads = neq2;
@@ -1905,16 +1942,18 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
// ======== VTCM allocation (GQA-aware) ========
// K/V row sizes drive the DMA descriptors (not the VTCM layout) and are used
- // throughout the KV loop below.
+ // throughout the KV loop below. The DMA copies only the real DK/DV columns; the
+ // staging rows are padded to hold DK_pad/DV_pad columns (tail zero-filled below)
+ // so the HMX interleave/tile logic can operate on 64-aligned head dims.
const size_t size_k_row = htp_tensor_get_row_size(k->type, DK);
const size_t size_v_row = htp_tensor_get_row_size(v->type, DV);
- const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128);
- const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128);
+ const size_t size_k_row_padded = hex_round_up(DK_pad * sizeof(__fp16), 128);
+ const size_t size_v_row_padded = hex_round_up(DV_pad * sizeof(__fp16), 128);
// Build the VTCM layout once (shared with the host estimator) and place every
- // scratch buffer at its computed offset.
+ // scratch buffer at its computed offset. Padded head dims size the HMX tiles.
struct hmx_fa_vtcm_layout L;
- hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline, factx.is_q_fp32);
+ hmx_fa_vtcm_layout_build(&L, G, DK_pad, DV_pad, Br, Bc, n_threads, pipeline, factx.is_q_fp32);
if (L.total_bytes > ctx->vtcm_size) {
return HTP_STATUS_VTCM_TOO_SMALL;
@@ -1961,6 +2000,24 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
dma_cache_init(&factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, L.m_buf_slot_bytes, HMX_FA_DMA_CACHE_SIZE);
+ // Head-dim padding: the K/V DMA staging buffers and the flat-Q buffer are laid out
+ // with padded row strides (size_{k,v,q}_row_padded, covering D_pad columns) but the
+ // DMA only writes the real D columns per row. Zero the whole staging buffers once up
+ // front so tail lanes [D, D_pad) stay zero for all KV blocks. No-op when already aligned.
+ if (DK_pad != DK || DV_pad != DV) {
+ const size_t k_buf_bytes = (size_t) factx.Bc * size_k_row_padded;
+ const size_t v_buf_bytes = (size_t) factx.Bc * size_v_row_padded;
+ hvx_splat_u8_a((char *) factx.vtcm_k_fp16[0], 0, k_buf_bytes);
+ hvx_splat_u8_a((char *) factx.vtcm_k_fp16[1], 0, k_buf_bytes);
+ hvx_splat_u8_a((char *) factx.vtcm_v_fp16[0], 0, v_buf_bytes);
+ hvx_splat_u8_a((char *) factx.vtcm_v_fp16[1], 0, v_buf_bytes);
+ // Flat-Q DMA scratch
+ if (factx.vtcm_q_dma) {
+ const size_t q_dma_bytes = hex_align_up(factx.g_br * DK * (factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16)), 128);
+ hvx_splat_u8_a((char *) factx.vtcm_q_dma, 0, q_dma_bytes);
+ }
+ }
+
// ======== 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));
@@ -2072,7 +2129,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
qk_job[0].s_tiles = factx.vtcm_s_tiles[0];
qk_job[0].n_row_tiles = n_row_tiles;
qk_job[0].n_col_tiles = hmx_ceil_div(kv_rows0, HMX_FP16_TILE_N_COLS);
- qk_job[0].n_dot_tiles = DK / 32;
+ qk_job[0].n_dot_tiles = DK_pad / 32;
qk_job[0].n_tiles_per_bc = n_tiles_per_bc;
qk_job[0].hmx_scales = factx.vtcm_hmx_scales_qk;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[0]));
@@ -2116,7 +2173,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
ou_job[prev_buf].n_row_tiles_g_br = n_row_tiles_g_br;
ou_job[prev_buf].n_tiles_per_bc = n_tiles_per_bc;
- ou_job[prev_buf].DV = DV;
+ ou_job[prev_buf].DV = DV_pad;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[prev_buf]));
}
@@ -2134,7 +2191,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
qk_job[next_buf].s_tiles = factx.vtcm_s_tiles[next_buf];
qk_job[next_buf].n_row_tiles = n_row_tiles;
qk_job[next_buf].n_col_tiles = hmx_ceil_div(next_rows, HMX_FP16_TILE_N_COLS);
- qk_job[next_buf].n_dot_tiles = DK / 32;
+ qk_job[next_buf].n_dot_tiles = DK_pad / 32;
qk_job[next_buf].n_tiles_per_bc = n_tiles_per_bc;
qk_job[next_buf].hmx_scales = factx.vtcm_hmx_scales_qk;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[next_buf]));
@@ -2198,7 +2255,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
ou_job[0].n_col_tiles = last_cols;
ou_job[0].n_row_tiles_g_br = n_row_tiles_g_br;
ou_job[0].n_tiles_per_bc = n_tiles_per_bc;
- ou_job[0].DV = DV;
+ ou_job[0].DV = DV_pad;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[0]));
// Overlapped: run HVX build diag inv L while HMX is busy executing the update
@@ -2246,7 +2303,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
qk_job.s_tiles = factx.vtcm_s_tiles[0];
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_dot_tiles = (size_t) (DK_pad / 32);
qk_job.n_tiles_per_bc = n_tiles_per_bc;
qk_job.hmx_scales = factx.vtcm_hmx_scales_qk;
@@ -2302,7 +2359,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
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;
+ ou_job.DV = DV_pad;
hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
if (kv_blk + 1 == factx.n_kv_blocks) {
@@ -2380,7 +2437,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
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;
+ on_job.DV = DV_pad;
hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_norm_worker, &on_job));
hmx_queue_pop(ctx->hmx_queue);
}
diff --git a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h
index d6795bf0b..8fd299795 100644
--- a/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h
+++ b/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h
@@ -495,12 +495,140 @@ static inline void hmx_fa_q_prep_fp16(
}
+// Head-dim-padded Q-prep (f32). Used when DK is not a multiple of 64.
+static inline void hmx_fa_q_prep_fp32_pad(__fp16 * vtcm_q_tiles,
+ const uint8_t * temp_q_vtcm,
+ size_t start,
+ size_t end,
+ size_t g_rows_end,
+ size_t dk_in,
+ size_t dk_out,
+ size_t G,
+ size_t n_rows_q,
+ const struct fastdiv_values * div_G,
+ bool q_transposed) {
+ const uint32_t n_out_tiles = (uint32_t) (dk_out / 32);
+ 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_out;
+
+ if (r >= g_rows_end) {
+ for (uint32_t d = 0; d < n_out_tiles; ++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_UVector * pv_in0 = (const HVX_UVector *) (temp_q_vtcm + offset0 * dk_in * sizeof(float));
+ const HVX_UVector * pv_in1 = (r + 1 < g_rows_end) ? (const HVX_UVector *) (temp_q_vtcm + offset1 * dk_in * sizeof(float)) : NULL;
+
+ for (uint32_t d = 0; d < n_out_tiles; ++d) {
+ HVX_Vector * out_tile = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
+ const size_t base_lane = (size_t) d * 32;
+ const size_t real_lanes = (base_lane < dk_in) ? hex_smin(32, dk_in - base_lane) : 0;
+
+ if (real_lanes == 0) {
+ out_tile[r1 / 2] = Q6_V_vzero();
+ continue;
+ }
+
+ HVX_Vector v0 = pv_in0[d];
+ HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+ if (real_lanes < 32) {
+ // Straddle tile: keep the first real_lanes floats, zero the padded tail so
+ // the packed f16 lanes beyond DK are zero.
+ const HVX_VectorPred keep = Q6_Q_vsetq_R((uint32_t) (real_lanes * sizeof(float)));
+ v0 = Q6_V_vmux_QVV(keep, v0, Q6_V_vzero());
+ v1 = Q6_V_vmux_QVV(keep, v1, Q6_V_vzero());
+ }
+ out_tile[r1 / 2] = hvx_vec_f32_to_f16_shuff(v0, v1);
+ }
+ }
+}
+
+// Head-dim-padded Q-prep (f16). Used when DK is not a multiple of 64.
+static inline void hmx_fa_q_prep_fp16_pad(__fp16 * vtcm_q_tiles,
+ const uint8_t * temp_q_vtcm,
+ size_t start,
+ size_t end,
+ size_t g_rows_end,
+ size_t dk_in,
+ size_t dk_out,
+ size_t G,
+ size_t n_rows_q,
+ const struct fastdiv_values * div_G,
+ bool q_transposed) {
+ const uint32_t n_out_pairs = (uint32_t) (dk_out / 64);
+ 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_out;
+
+ if (r >= g_rows_end) {
+ for (uint32_t d = 0; d < n_out_pairs; ++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_UVector * pv_in0 = (const HVX_UVector *) (temp_q_vtcm + offset0 * dk_in * sizeof(__fp16));
+ const HVX_UVector * pv_in1 = (r + 1 < g_rows_end) ? (const HVX_UVector *) (temp_q_vtcm + offset1 * dk_in * sizeof(__fp16)) : NULL;
+
+ for (uint32_t d = 0; d < n_out_pairs; ++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;
+
+ const size_t base_lane = (size_t) d * 64;
+ const size_t real_lanes = (base_lane < dk_in) ? hex_smin(64, dk_in - base_lane) : 0;
+
+ if (real_lanes == 0) {
+ *pv_out0 = Q6_V_vzero();
+ *pv_out1 = Q6_V_vzero();
+ continue;
+ }
+
+ HVX_Vector v0 = pv_in0[d];
+ HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+ if (real_lanes < 64) {
+ const HVX_VectorPred keep = Q6_Q_vsetq_R((uint32_t) (real_lanes * sizeof(__fp16)));
+ v0 = Q6_V_vmux_QVV(keep, v0, Q6_V_vzero());
+ v1 = Q6_V_vmux_QVV(keep, v1, Q6_V_vzero());
+ }
+ HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+ *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,
+ size_t G, size_t dk_in, size_t dk_out, bool is_q_fp32,
const struct fastdiv_values * div_G
) {
for (size_t r = start; r < end; r += 2) {
@@ -518,33 +646,55 @@ static inline void hmx_fa_q_prep_fallback(
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;
+ __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * dk_out;
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;
+ for (uint32_t d = 0; d < dk_out / 32; ++d) {
+ HVX_Vector * out_tile = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
+ const size_t base_lane = (size_t) d * 32;
+ const size_t real_lanes = (base_lane < dk_in) ? hex_smin(32, dk_in - base_lane) : 0;
+
+ if (real_lanes == 0) {
+ out_tile[r1 / 2] = Q6_V_vzero();
+ continue;
+ }
+ HVX_Vector v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
+ HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+ if (real_lanes < 32) {
+ const HVX_VectorPred keep = Q6_Q_vsetq_R((uint32_t) (real_lanes * sizeof(float)));
+ v0 = Q6_V_vmux_QVV(keep, v0, Q6_V_vzero());
+ v1 = Q6_V_vmux_QVV(keep, v1, Q6_V_vzero());
+ }
+ out_tile[r1 / 2] = hvx_vec_f32_to_f16_shuff(v0, v1);
}
} 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);
-
+ for (uint32_t d = 0; d < dk_out / 64; ++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;
+ const size_t base_lane = (size_t) d * 64;
+ const size_t real_lanes = (base_lane < dk_in) ? hex_smin(64, dk_in - base_lane) : 0;
+
+ if (real_lanes == 0) {
+ *pv_out0 = Q6_V_vzero();
+ *pv_out1 = Q6_V_vzero();
+ continue;
+ }
+ HVX_Vector v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
+ HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+ if (real_lanes < 64) {
+ const HVX_VectorPred keep = Q6_Q_vsetq_R((uint32_t) (real_lanes * sizeof(__fp16)));
+ v0 = Q6_V_vmux_QVV(keep, v0, Q6_V_vzero());
+ v1 = Q6_V_vmux_QVV(keep, v1, Q6_V_vzero());
+ }
+ HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
*pv_out0 = Q6_V_lo_W(vp);
*pv_out1 = Q6_V_hi_W(vp);
}
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 30792e409..d8f4c3708 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -10713,6 +10713,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
}
+ // asymmetric head_dim (hsk != hsv) with one or both sides not 64-aligned
+ test_cases.emplace_back(new test_flash_attn_ext(72, 64, 4, {1, 1}, 256, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+ test_cases.emplace_back(new test_flash_attn_ext(64, 72, 4, {1, 1}, 256, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+
// mixed quant and Q1_0 test cases
test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0));
test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_F16));