From 18a04f09c24616898792bcfaa17f3550bdc78912 Mon Sep 17 00:00:00 2001 From: Todor Boinovski Date: Fri, 18 Sep 2026 13:15:08 -0700 Subject: [PATCH] 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. --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 13 +- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 101 +++++++++--- ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h | 176 +++++++++++++++++++-- tests/test-backend-ops.cpp | 4 + 4 files changed, 256 insertions(+), 38 deletions(-) diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 3f14956453..af8013b08a 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 8a1caba22b..75422f4200 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 d6795bf0b5..8fd299795c 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); + 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; - HVX_Vector * out_tile = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS); - out_tile[r1 / 2] = v_hf; + 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 30792e4096..d8f4c3708b 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10713,6 +10713,10 @@ static std::vector> 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));