From ceefd16daa2093730bf3340315fa2b84b286ddbc Mon Sep 17 00:00:00 2001 From: Jhen-Jie Hong Date: Fri, 16 Oct 2026 16:51:14 +0800 Subject: [PATCH] hexagon: correct HMX FA sinks (#30581) * hexagon: correct HMX FA sinks * hex-fa: vectorize sinks in hmx kernels --------- Co-authored-by: Max Krasnyansky --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 126 ++++++++++++--------- 1 file changed, 73 insertions(+), 53 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 4b2734695..a0473c13d 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -755,40 +755,11 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) { const size_t m_start = i * m_bytes_per_t; const size_t m_end = hex_smin(m_start + m_bytes_per_t, col_vec_bytes); - if (factx->sinks) { - const float * sinks_data = factx->vtcm_sinks; - float * m_vec = (float *) factx->vtcm_m_vec; - const size_t r_start = l_start / sizeof(float); - const size_t r_end = l_end / sizeof(float); - const float scale_factor = EXP_LOG2E_F; - - const HVX_Vector v_scale = hvx_vec_splat_f32(scale_factor); - - for (size_t r = r_start; r < r_end; r += 32) { - HVX_VectorAlias local_m; - for (size_t j = 0; j < 32; ++j) { - size_t curr_r = r + j; - if (curr_r < n_rows_g) { - const size_t h_idx = fastmodulo(curr_r, G, &factx->div_G); - const size_t head = args->kv_head * G + h_idx; - local_m.fp32[j] = sinks_data[head]; - } else { - local_m.fp32[j] = HTP_FA_M_INITIAL_VAL; - } - } - HVX_Vector v_scaled = HVX_OP_MUL_F32(local_m.v, v_scale); - *(HVX_Vector *) (m_vec + r) = v_scaled; - } - if (l_start < col_vec_bytes) { - hvx_splat_u8_a((char *) factx->vtcm_l_vec + l_start, 0, l_end - l_start); - } - } else { - if (l_start < col_vec_bytes) { - hvx_splat_u8_a((char *) factx->vtcm_l_vec + l_start, 0, l_end - l_start); - } - if (m_start < col_vec_bytes) { - hvx_splat_f32_a((char *) factx->vtcm_m_vec + m_start, HTP_FA_M_INITIAL_VAL, (m_end - m_start) / sizeof(float)); - } + if (l_start < col_vec_bytes) { + hvx_splat_u8_a((char *) factx->vtcm_l_vec + l_start, 0, l_end - l_start); + } + if (m_start < col_vec_bytes) { + hvx_splat_f32_a((char *) factx->vtcm_m_vec + m_start, HTP_FA_M_INITIAL_VAL, (m_end - m_start) / sizeof(float)); } // Zero the whole rescale region: vtcm_d_tiles[0], the optional vtcm_d_tiles[1] @@ -1447,21 +1418,10 @@ static inline void fa_softmax_impl( HVX_Vector v_rowsum_acc_f32_0 = Q6_V_lo_W(rowsum_acc_pair); HVX_Vector v_rowsum_acc_f32_1 = Q6_V_hi_W(rowsum_acc_pair); - HVX_Vector v_l_curr0; - HVX_Vector v_l_curr1; - if (args->kv_start == 0 && factx->sinks != NULL) { - // First KV block with sinks: m_prev holds the seeded sink value (not -inf), - // so exp_m_diff = exp2(sink - m_curr) is the sink's contribution to the - // denominator. l_prev is 0 here, so add exp_m_diff directly instead of - // multiplying the (uninitialized) l_prev term. - v_l_curr0 = HVX_OP_ADD_F32(exp_m_diff0, v_rowsum_acc_f32_0); - v_l_curr1 = HVX_OP_ADD_F32(exp_m_diff1, v_rowsum_acc_f32_1); - } else { - HVX_Vector l_prev_v0 = factx->vtcm_l_vec[r_vec_idx * 2 + 0]; - HVX_Vector l_prev_v1 = factx->vtcm_l_vec[r_vec_idx * 2 + 1]; - v_l_curr0 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v0, exp_m_diff0), v_rowsum_acc_f32_0); - v_l_curr1 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v1, exp_m_diff1), v_rowsum_acc_f32_1); - } + HVX_Vector l_prev_v0 = factx->vtcm_l_vec[r_vec_idx * 2 + 0]; + HVX_Vector l_prev_v1 = factx->vtcm_l_vec[r_vec_idx * 2 + 1]; + HVX_Vector v_l_curr0 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v0, exp_m_diff0), v_rowsum_acc_f32_0); + HVX_Vector v_l_curr1 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v1, exp_m_diff1), v_rowsum_acc_f32_1); factx->vtcm_m_vec[r_vec_idx * 2 + 0] = v_m_curr0; factx->vtcm_m_vec[r_vec_idx * 2 + 1] = v_m_curr1; @@ -1532,9 +1492,66 @@ static void fa_softmax_thread(unsigned int n, unsigned int i, void * data) { fa_softmax_impl(n, i, data, has_mask, mask_broadcast, is_g1, has_alibi, has_softcap); } +static inline HVX_Vector hvx_fa_load_sinks(const float * sinks_data, size_t idx, size_t count) { + const HVX_Vector * v_src = (const HVX_Vector *) sinks_data + (idx >> 5); + const size_t lane_off = idx & 31; + const size_t byte_off = lane_off * sizeof(float); + return (lane_off + count > 32) + ? Q6_V_valign_VVR(v_src[1], v_src[0], byte_off) + : Q6_V_vror_VR(v_src[0], byte_off); +} + static __attribute__((noinline)) void fa_build_d_diag_inv_l(struct hmx_fa_context * factx, size_t n_row_tiles, - size_t n_row_tiles_g_br) { + size_t n_row_tiles_g_br, + uint32_t kv_head, + size_t n_rows_g) { + (void) n_row_tiles_g_br; + + // sinks: l += exp(sink - m * ln2) in f32, m and l are in the log2 domain + if (factx->sinks) { + const float * sinks_data = factx->vtcm_sinks; + const size_t G = factx->G; + const size_t head_base = (size_t) kv_head * G; + const HVX_Vector v_ln2 = Q6_V_vsplat_R(EXP_LOGN2); + const HVX_Vector v_head_sinks = (G <= 32) ? hvx_fa_load_sinks(sinks_data, head_base, G) : Q6_V_vzero(); + HVX_Vector v_sink = Q6_V_vzero(); + size_t g_prev = SIZE_MAX; + + for (size_t i = 0; i < n_row_tiles; ++i) { + // the sink pattern of a tile only depends on (i * 32) % G, rebuild it when that changes + const size_t r0 = i * 32; + const size_t g0 = fastmodulo(r0, G, &factx->div_G); + if (g0 != g_prev) { + const size_t rem_g = G - g0; + if (G <= 32) { + const HVX_Vector v_lo = Q6_V_vror_VR(v_head_sinks, g0 * sizeof(float)); + const HVX_Vector v_hi = Q6_V_vror_VR(v_head_sinks, 128 - rem_g * sizeof(float)); + v_sink = Q6_V_vmux_QVV(Q6_Q_vsetq2_R(rem_g * sizeof(float)), v_lo, v_hi); + for (size_t len = G; len < 32; len <<= 1) { + const HVX_VectorPred q_keep = Q6_Q_vsetq_R(len * sizeof(float)); + const HVX_Vector v_shifted = Q6_V_vror_VR(v_sink, 128 - len * sizeof(float)); + v_sink = Q6_V_vmux_QVV(q_keep, v_sink, v_shifted); + } + } else if (rem_g >= 32) { + v_sink = hvx_fa_load_sinks(sinks_data, head_base + g0, 32); + } else { + const HVX_Vector v_lo = hvx_fa_load_sinks(sinks_data, head_base + g0, rem_g); + const HVX_Vector v_hi = hvx_fa_load_sinks(sinks_data, head_base, 32 - rem_g); + v_sink = Q6_V_vmux_QVV(Q6_Q_vsetq_R(rem_g * sizeof(float)), v_lo, Q6_V_vror_VR(v_hi, 128 - rem_g * sizeof(float))); + } + g_prev = g0; + } + + HVX_Vector v_d = HVX_OP_SUB_F32(v_sink, HVX_OP_MUL_F32(factx->vtcm_m_vec[i], v_ln2)); + if (r0 + 32 > n_rows_g) { + const HVX_VectorPred q_valid = Q6_Q_vsetq_R((n_rows_g - r0) * sizeof(float)); + v_d = Q6_V_vmux_QVV(q_valid, v_d, Q6_V_vzero()); + } + factx->vtcm_l_vec[i] = HVX_OP_ADD_F32(factx->vtcm_l_vec[i], hvx_vec_exp_f32(v_d)); + } + } + const HVX_Vector v_offsets = *(const HVX_Vector *) d_tile_scatter_offsets; const HVX_VectorPred q_32_mask = Q6_Q_vsetq_R(32 * sizeof(__fp16)); const HVX_Vector one = hvx_vec_splat_f32(1.0f); @@ -2010,7 +2027,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // ======== 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)); + { + union { __fp16 f; uint16_t u; } sc = { .f = factx.scale }; + hmx_init_column_scales(factx.vtcm_hmx_scales_qk, Q6_V_vsplat_R(sc.u)); + } // ======== DMA setup ======== dma_queue * const dma_q = ctx->dma[0]; @@ -2250,7 +2270,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // Overlapped: run HVX build diag inv L while HMX is busy executing the update htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start); - fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br); + fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br, kv_head, n_rows_g); htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start); hmx_queue_pop(hmx_q); @@ -2355,7 +2375,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { if (kv_blk + 1 == factx.n_kv_blocks) { // Overlapped: run HVX build diag inv L while HMX is busy executing the update htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start); - fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br); + fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br, kv_head, n_rows_g); htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start); } hmx_queue_pop(ctx->hmx_queue);