Repository: ggerganov/llama.cpp
File: ggml/src/ggml-hexagon/htp/flash-attn-ops.c (C)

Code before:
                HVX_Vector scores_shifted = Q6_Vqf32_vsub_VsfVsf(scores, M_new_vec);
                HVX_Vector P = hvx_vec_exp_f32(Q6_Vsf_equals_Vqf32(scores_shifted));

                HVX_Vector p_sum_vec = hvx_vec_reduce_sum_f32(P);
                float p_sum = hvx_vec_get_f32(p_sum_vec);
                S += p_sum;

                // 5. Accumulate V
                float __attribute__((aligned(VLEN))) p_arr[VLEN_FP32];
                *(HVX_Vector*)p_arr = P;

                for (int j = 0; j < VLEN_FP32; ++j) {
                    const uint32_t cur_ic = ic + j;
                    const uint8_t * v_ptr = v_base + cur_ic * size_v_row_padded;
                    hvx_mad_f32_f16_aa(VKQ32, v_ptr, DV, p_arr[j]);
                }
            }

            // Leftover
            for (; ic < current_block_size; ++ic) {
                float s_val;
                const uint8_t * k_ptr = k_base + ic * size_k_row_padded;

                if (q->type == HTP_TYPE_F32) {
                    hvx_dot_f32_f16_aa(&s_val, q_ptr_vtcm, k_ptr, DK, scale);
                } else {
                    hvx_dot_f16_f16_aa(&s_val, q_ptr_vtcm, k_ptr, DK, scale);
                }

                if (logit_softcap != 0.0f) {
                    s_val = logit_softcap * tanhf(s_val);
                }

                if (mask) {
                    const float m_val = m_base[ic];
                    s_val += slope * m_val;
                }

                const float Mold = M;
                float ms = 1.0f;
                float vs = 1.0f;

                if (s_val > M) {
                    M = s_val;
                    ms = expf(Mold - M);
                    hvx_scale_f32_aa((uint8_t *) VKQ32, (const uint8_t *) VKQ32, DV, ms);
                } else {
                    vs = expf(s_val - M);
                }

                const uint8_t * v_ptr = v_base + ic * size_v_row_padded;

Review comment (performance):
Here, we move the max reduction and the `e^(x-max)` update outside of the vector loop. This saves `FLASH_ATTN_BLOCK_SIZE / VLEN_FP32 - 1` calls to `hvx_vec_reduce_max_f32`, `hvx_scale_f32_aa`, and `expf` (assuming `expf` is computationally expensive).

Alternatively, we could optimize further by moving the softmax update after the leftover block, which would reduce expensive calculations even more and improve performance.

Code after:

                const float ms = expf(M_old - M_new);
                hvx_scale_f32_aa((uint8_t *) VKQ32, (const uint8_t *) VKQ32, DV, ms);

                HVX_Vector M_new_vec = hvx_vec_splat_f32(M_new);
                HVX_Vector p_sum_vec = hvx_vec_splat_f32(0.0f);
                for (uint32_t ic2 = 0, iv = 0; ic2 + VLEN_FP32 <= current_block_size; ic2 += VLEN_FP32, ++iv) {
                    HVX_Vector scores = scores_x4.v[iv];
                    HVX_Vector scores_shifted = Q6_Vqf32_vsub_VsfVsf(scores, M_new_vec);
                    HVX_Vector P = hvx_vec_exp_f32(Q6_Vsf_equals_Vqf32(scores_shifted));

                    p_sum_vec = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(p_sum_vec, P));

                    // 5. Accumulate V
                    float __attribute__((aligned(VLEN))) p_arr[VLEN_FP32];
                    *(HVX_Vector*)p_arr = P;

                    for (int j = 0; j < VLEN_FP32; ++j) {
                        const uint32_t cur_ic = ic2 + j;
                        const uint8_t * v_ptr = v_base + cur_ic * size_v_row_padded;
                        hvx_mad_f32_f16_aa(VKQ32, v_ptr, DV, p_arr[j]);
                    }
                }

                p_sum_vec = hvx_vec_reduce_sum_f32(p_sum_vec);
                S = S * ms + hvx_vec_get_f32(p_sum_vec);
            }

            // Leftover
            for (; ic < current_block_size; ++ic) {
                float s_val;
                const uint8_t * k_ptr = k_base + ic * size_k_row_padded;

                if (q->type == HTP_TYPE_F32) {
                    hvx_dot_f32_f16_aa(&s_val, q_ptr_vtcm, k_ptr, DK, scale);
                } else {
                    hvx_dot_f16_f16_aa(&s_val, q_ptr_vtcm, k_ptr, DK, scale);
                }

                if (logit_softcap != 0.0f) {
                    s_val = logit_softcap * tanhf(s_val);
                }

                if (mask) {
                    const float m_val = m_base[ic];
                    s_val += slope * m_val;
                }

                const float Mold = M;
                float ms = 1.0f;
                float vs = 1.0f;