char-fork/llama.cpp

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

Max Krasnyanskyhexagon: flash attention rework (optimizations, accuracy improvements, etc) (#25085)13e673863

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
9.8 KiB255 linesraw
1#ifndef HVX_EXP_H
2#define HVX_EXP_H
3
4#include <stdbool.h>
5#include <stdint.h>
6#include <math.h>
7
8#include "hvx-base.h"
9#include "hvx-floor.h"
10
11#define EXP_COEFF_5 (0x39506967)  // 0.000198757 = 1/(7!)
12#define EXP_COEFF_4 (0x3AB743CE)  // 0.0013982   = 1/(6!)
13#define EXP_COEFF_3 (0x3C088908)  // 0.00833345  = 1/(5!)
14#define EXP_COEFF_2 (0x3D2AA9C1)  // 0.416658    = 1/(4!)
15#define EXP_COEFF_1 (0x3E2AAAAA)  // 0.16666667  = 1/(3!)
16#define EXP_COEFF_0 (0x3F000000)  // 0.5         = 1/(2!)
17#define EXP_LOGN2   (0x3F317218)  // ln(2)   = 0.6931471805
18#define EXP_LOG2E   (0x3FB8AA3B)  // log2(e) = 1/ln(2) = 1.4426950408
19#define EXP_LOG2E_F 1.44269504f
20#define EXP_ONE     (0x3f800000)  // 1.0
21#define EXP_RANGE_R (0x42B17218)  // ln(FLT_MAX) approx = 88.7228
22#define EXP_RANGE_L (0xC2B00000)  // -88.0 (approx log(FLT_MIN))
23
24static inline HVX_Vector hvx_vec_exp_f32(HVX_Vector in_vec) {
25    HVX_Vector z_qf32_v;
26    HVX_Vector x_v;
27    HVX_Vector x_qf32_v;
28    HVX_Vector y_v;
29    HVX_Vector k_v;
30    HVX_Vector f_v;
31    HVX_Vector epsilon_v;
32    HVX_Vector log2e = Q6_V_vsplat_R(EXP_LOG2E);
33    HVX_Vector logn2 = Q6_V_vsplat_R(EXP_LOGN2);
34    HVX_Vector E_const;
35    HVX_Vector zero_v = Q6_V_vzero();
36
37    // exp(x) is approximated as follows:
38    //   f = floor(x/ln(2)) = floor(x*log2(e))
39    //   epsilon = x - f*ln(2)
40    //   exp(x) = exp(epsilon+f*ln(2))
41    //          = exp(epsilon)*exp(f*ln(2))
42    //          = exp(epsilon)*2^f
43    //
44    //   Since epsilon is close to zero, it can be approximated with its Taylor series:
45    //            exp(x) ~= 1+x+x^2/2!+x^3/3!+...+x^n/n!+...
46    //   Preserving the first eight elements, we get:
47    //            exp(x) ~= 1+x+e0*x^2+e1*x^3+e2*x^4+e3*x^5+e4*x^6+e5*x^7
48    //                   =  1+x+(E0+(E1+(E2+(E3+(E4+E5*x)*x)*x)*x)*x)*x^2
49
50    HVX_Vector temp_v = in_vec;
51
52    // Clamp inputs to (-88.0, 88.0) to avoid overflow/underflow
53    HVX_VectorPred pred_cap_right = Q6_Q_vcmp_gt_VsfVsf(in_vec, Q6_V_vsplat_R(EXP_RANGE_R));
54    HVX_VectorPred pred_cap_left  = Q6_Q_vcmp_gt_VsfVsf(Q6_V_vsplat_R(EXP_RANGE_L), in_vec);
55
56    in_vec = Q6_V_vmux_QVV(pred_cap_right, Q6_V_vsplat_R(EXP_RANGE_R), temp_v);
57    in_vec = Q6_V_vmux_QVV(pred_cap_left, Q6_V_vsplat_R(EXP_RANGE_L), in_vec);
58
59    epsilon_v = Q6_Vqf32_vmpy_VsfVsf(log2e, in_vec);
60    epsilon_v = Q6_Vsf_equals_Vqf32(epsilon_v);
61
62    //    f_v is the floating point result and k_v is the integer result
63    f_v = hvx_vec_floor_f32(epsilon_v);
64    k_v = hvx_vec_truncate_f32(f_v);
65
66    x_qf32_v = Q6_Vqf32_vadd_VsfVsf(in_vec, zero_v);
67
68    //  x = x - f_v * logn2;
69    epsilon_v = Q6_Vqf32_vmpy_VsfVsf(f_v, logn2);
70    x_qf32_v  = Q6_Vqf32_vsub_Vqf32Vqf32(x_qf32_v, epsilon_v);
71    // normalize before every QFloat's vmpy
72    x_qf32_v  = Q6_Vqf32_vadd_Vqf32Vsf(x_qf32_v, zero_v);
73
74    x_v = Q6_Vsf_equals_Vqf32(x_qf32_v);
75
76    // z = x * x;
77    z_qf32_v = Q6_Vqf32_vmpy_Vqf32Vqf32(x_qf32_v, x_qf32_v);
78    z_qf32_v = Q6_Vqf32_vadd_Vqf32Vsf(z_qf32_v, zero_v);
79
80    // y = E4 + E5 * x;
81    E_const = Q6_V_vsplat_R(EXP_COEFF_5);
82    y_v     = Q6_Vqf32_vmpy_VsfVsf(E_const, x_v);
83    E_const = Q6_V_vsplat_R(EXP_COEFF_4);
84    y_v     = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const);
85    y_v     = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v);
86
87    // y = E3 + y * x;
88    E_const = Q6_V_vsplat_R(EXP_COEFF_3);
89    y_v     = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v);
90    y_v     = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const);
91    y_v     = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v);
92
93    // y = E2 + y * x;
94    E_const = Q6_V_vsplat_R(EXP_COEFF_2);
95    y_v     = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v);
96    y_v     = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const);
97    y_v     = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v);
98
99    // y = E1 + y * x;
100    E_const = Q6_V_vsplat_R(EXP_COEFF_1);
101    y_v     = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v);
102    y_v     = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const);
103    y_v     = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v);
104
105    // y = E0 + y * x;
106    E_const = Q6_V_vsplat_R(EXP_COEFF_0);
107    y_v     = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, x_qf32_v);
108    y_v     = Q6_Vqf32_vadd_Vqf32Vsf(y_v, E_const);
109    y_v     = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v);
110
111    // y = x + y * z;
112    y_v = Q6_Vqf32_vmpy_Vqf32Vqf32(y_v, z_qf32_v);
113    y_v = Q6_Vqf32_vadd_Vqf32Vqf32(y_v, x_qf32_v);
114    y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, zero_v);
115
116    // y = y + 1.0;
117    y_v = Q6_Vqf32_vadd_Vqf32Vsf(y_v, Q6_V_vsplat_R(EXP_ONE));
118
119    // insert exponents
120    //        y = ldexpf(y, k);
121    //    y_v += k_v; // qf32
122    // modify exponent
123
124    y_v = Q6_Vsf_equals_Vqf32(y_v);
125
126    // add k_v to the exponent of y_v
127    HVX_Vector y_v_exponent = Q6_Vw_vasl_VwR(y_v, 1);
128
129    y_v_exponent = Q6_Vuw_vlsr_VuwR(y_v_exponent, IEEE_VSF_MANTLEN + 1);
130    y_v_exponent = Q6_Vw_vadd_VwVw(k_v, y_v_exponent);
131
132    // exponent cannot be negative; if overflow is detected, result is set to zero
133    HVX_VectorPred qy_v_negative_exponent = Q6_Q_vcmp_gt_VwVw(zero_v, y_v_exponent);
134
135    y_v = Q6_Vw_vaslacc_VwVwR(y_v, k_v, IEEE_VSF_MANTLEN);
136
137    y_v = Q6_V_vmux_QVV(qy_v_negative_exponent, zero_v, y_v);
138
139    return y_v;
140}
141
142static inline HVX_Vector hvx_vec_exp_f32_guard(HVX_Vector in_vec, HVX_Vector max_exp, HVX_Vector inf) {
143    const HVX_VectorPred pred0 = Q6_Q_vcmp_gt_VsfVsf(in_vec, max_exp);
144
145    HVX_Vector out = hvx_vec_exp_f32(in_vec);
146
147    return Q6_V_vmux_QVV(pred0, inf, out);
148}
149
150static inline void hvx_exp_f32(uint8_t * restrict dst, const uint8_t * restrict src, const int num_elems, bool negate) {
151    int left_over       = num_elems & (VLEN_FP32 - 1);
152    int num_elems_whole = num_elems - left_over;
153
154    int unaligned_addr = 0;
155    int unaligned_loop = 0;
156    if ((0 == hex_is_aligned((void *) src, VLEN)) || (0 == hex_is_aligned((void *) dst, VLEN))) {
157        unaligned_addr = 1;
158    }
159    // assert((0 == unaligned_addr) || (0 == num_elems_whole));
160    if ((1 == unaligned_addr) && (num_elems_whole != 0)) {
161        unaligned_loop = 1;
162    }
163
164    HVX_Vector vec_out = Q6_V_vzero();
165
166    static const float kInf    = INFINITY;
167    static const float kMaxExp = 88.7228f;
168
169    const HVX_Vector max_exp = hvx_vec_splat_f32(kMaxExp);
170    const HVX_Vector inf     = hvx_vec_splat_f32(kInf);
171
172    if (0 == unaligned_loop) {
173        HVX_Vector * p_vec_in1 = (HVX_Vector *) src;
174        HVX_Vector * p_vec_out = (HVX_Vector *) dst;
175
176        #pragma unroll(4)
177        for (int i = 0; i < num_elems_whole; i += VLEN_FP32) {
178            if (true == negate) {
179                HVX_Vector neg_vec_in = hvx_vec_neg_f32(*p_vec_in1++);
180                *p_vec_out++          = hvx_vec_exp_f32_guard(neg_vec_in, max_exp, inf);
181            } else {
182                *p_vec_out++ = hvx_vec_exp_f32_guard(*p_vec_in1++, max_exp, inf);
183            }
184        }
185    } else {
186        #pragma unroll(4)
187        for (int i = 0; i < num_elems_whole; i += VLEN_FP32) {
188            HVX_Vector in = *(HVX_UVector *) (src + i * SIZEOF_FP32);
189
190            if (true == negate) {
191                HVX_Vector neg_vec_in                    = hvx_vec_neg_f32(in);
192                *(HVX_UVector *) (dst + i * SIZEOF_FP32) = hvx_vec_exp_f32_guard(neg_vec_in, max_exp, inf);
193            } else {
194                *(HVX_UVector *) (dst + i * SIZEOF_FP32) = hvx_vec_exp_f32_guard(in, max_exp, inf);
195            }
196        }
197    }
198
199    if (left_over > 0) {
200        const float * srcf = (float *) src + num_elems_whole;
201        float *       dstf = (float *) dst + num_elems_whole;
202
203        HVX_Vector in = *(HVX_UVector *) srcf;
204
205        if (true == negate) {
206            HVX_Vector neg_vec_in = hvx_vec_neg_f32(in);
207
208            vec_out = hvx_vec_exp_f32_guard(neg_vec_in, max_exp, inf);
209        } else {
210            vec_out = hvx_vec_exp_f32_guard(in, max_exp, inf);
211        }
212
213        hvx_vec_store_u((void *) dstf, left_over * SIZEOF_FP32, vec_out);
214    }
215}
216
217static inline HVX_Vector hvx_vec_exp2_f16(HVX_Vector x_v) {
218    const HVX_Vector zero_v    = Q6_V_vzero();
219    const HVX_Vector half_hf_v = Q6_Vh_vsplat_R(0x3800);  // fp16 0.5
220
221    // Clamp input to prevent integer underflow in FP16-to-INT16 conversion
222    const HVX_Vector v_clamp_min = hvx_vec_splat_f16(-24.0f);
223    x_v = Q6_Vhf_vmax_VhfVhf(v_clamp_min, x_v);
224
225    // k = round_toward_neg_inf(x);  f = (float)k;  frac = x - f
226    HVX_Vector x_minus_half = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vsub_VhfVhf(x_v, half_hf_v));
227    HVX_Vector k_v          = Q6_Vh_equals_Vhf(x_minus_half);  // truncate to int16
228    HVX_Vector f_v          = Q6_Vhf_equals_Vh(k_v);           // back to fp16
229
230    HVX_Vector x_qf16 = Q6_Vqf16_vsub_VhfVhf(x_v, f_v);        // fractional part in qf16
231
232    // Horner: y = ((((E5*x + E4)*x + E3)*x + E2)*x + E1)*x + E0
233    HVX_Vector y = Q6_Vqf16_vmpy_Vqf16Vqf16(Q6_Vh_vsplat_R(0x5082), x_qf16); // E5*x
234    y            = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x157d));        // + E4
235    y            = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
236    y            = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x20ed));        // + E3
237    y            = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
238    y            = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x2b1b));        // + E2
239    y            = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
240    y            = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x33b0));        // + E1
241    y            = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
242    y            = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x398c));        // + E0
243    y            = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);                      // y = y * x
244    y            = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x3c00));        // + 1.0
245
246    // Combine polynomial (mantissa) with integer part (exponent): result = y * 2^k
247    y                          = Q6_Vhf_equals_Vqf16(y);
248    HVX_Vector y_exp           = Q6_Vuh_vlsr_VuhR(Q6_Vh_vasl_VhR(y, 1), 11);
249    y_exp                      = Q6_Vh_vadd_VhVh(k_v, y_exp);
250    HVX_VectorPred q_underflow = Q6_Q_vcmp_gt_VhVh(zero_v, y_exp);
251    y                          = Q6_Vh_vaslacc_VhVhR(y, k_v, 10);
252    return Q6_V_vmux_QVV(q_underflow, zero_v, y);
253}
254
255#endif /* HVX_EXP_H */