char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Max Krasnyanskyhexagon: MUL_MAT, MUL_MAT_ID, FLASH_ATTN and GDN cleanup and optimizations for latest models (#23989)8f7f3bf14

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
1.8 KiB47 linesraw
1#ifndef HVX_FLASH_ATTN_H
2#define HVX_FLASH_ATTN_H
3
4#include <math.h>
5#include "hvx-utils.h"
6
7// Scalar helper to compute a single ALiBi slope.
8static inline float alibi_slope(uint32_t h, uint32_t n_head_log2, float m0, float m1) {
9    return (h < n_head_log2) ? powf(m0, h + 1) : powf(m1, 2 * (h - n_head_log2) + 1);
10}
11
12// Vectorized helper to compute 32 ALiBi slopes starting from (kv_head * G).
13static inline HVX_Vector hvx_alibi_slopes(
14    uint32_t kv_head,
15    uint32_t G,
16    uint32_t n_head_log2,
17    float m0,
18    float m1
19) {
20    static const float ramp_32[32] __attribute__((aligned(128))) = {
21        0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f,
22        8.0f, 9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f,
23        16.0f, 17.0f, 18.0f, 19.0f, 20.0f, 21.0f, 22.0f, 23.0f,
24        24.0f, 25.0f, 26.0f, 27.0f, 28.0f, 29.0f, 30.0f, 31.0f
25    };
26    HVX_Vector v_ramp = hvx_vmem(ramp_32);
27    HVX_Vector v_h_base = hvx_vec_splat_f32((float)(kv_head * G));
28    HVX_Vector v_h = hvx_vec_add_f32_f32(v_h_base, v_ramp);
29
30    // Compute exponent_m0: h + 1
31    HVX_Vector v_exp_m0 = hvx_vec_add_f32_f32(v_h, hvx_vec_splat_f32(1.0f));
32
33    // Compute exponent_m1: 2 * (h - n_head_log2) + 1
34    HVX_Vector v_n_head_log2 = hvx_vec_splat_f32((float)n_head_log2);
35    HVX_Vector v_h_minus = hvx_vec_sub_f32_f32(v_h, v_n_head_log2);
36    HVX_Vector v_exp_m1 = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(hvx_vec_splat_f32(2.0f), v_h_minus), hvx_vec_splat_f32(1.0f));
37
38    // Compute powers
39    HVX_Vector v_pow_m0 = hvx_vec_pow_const_base_f32(m0, v_exp_m0);
40    HVX_Vector v_pow_m1 = hvx_vec_pow_const_base_f32(m1, v_exp_m1);
41
42    // Select based on h < n_head_log2
43    HVX_VectorPred p_cond = Q6_Q_vcmp_gt_VsfVsf(v_n_head_log2, v_h); // v_n_head_log2 > v_h <=> h < n_head_log2
44    return Q6_V_vmux_QVV(p_cond, v_pow_m0, v_pow_m1);
45}
46
47#endif /* HVX_FLASH_ATTN_H */