Commit ceefd16da for llama.cpp
commit ceefd16daa2093730bf3340315fa2b84b286ddbc
Author: Jhen-Jie Hong <iainst0409@gmail.com>
Date: Fri Oct 16 16:51:14 2026 +0800
hexagon: correct HMX FA sinks (#30581)
* hexagon: correct HMX FA sinks
* hex-fa: vectorize sinks in hmx kernels
---------
Co-authored-by: Max Krasnyansky <maxk@qti.qualcomm.com>
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);