hey! this is a static git repo viewer. commit diffs, file history, older commits and other revisions won't load without javascript.

you are still able to view source files on branch/tag tips, but the experience is degraded.

char-fork/llama.cpp

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

Max Krasnyanskyhexagon: support for multi-device model split (aka row-split) (#28589)eafe15a5e2026-09-11

1#include <math.h>
2#include <stdint.h>
3#include <string.h>
4
5#include "hvx-utils.h"
6#include "hex-fastdiv.h"
7#include "hex-common.h"
8#include "hex-profile.h"
9
10#define GGML_COMMON_DECL_C
11#include "ggml-common.h"
12#include "htp-ctx.h"
13#include "htp-tensor.h"
14
15#ifndef MIN
16#define MIN(a, b) ((a) < (b) ? (a) : (b))
17#endif
18
19#define HTP_GDN_MAX_SV 128
20
21
22struct htp_gdn_context {
23    struct htp_ops_context * octx;
24    uint32_t rows_per_thread;
25    size_t   state_bytes;
26    uint8_t * vtcm_base;
27    size_t   vtcm_per_thread;
28    uint32_t row_start;
29    uint32_t nrows;
30};
31
32static inline HVX_Vector gdn_mul_dot_f32(float * restrict dst, const float * restrict mul, const float * restrict dot, uint32_t n) {
33    HVX_Vector acc = Q6_V_vzero();
34
35    const uint32_t epv  = 128 / sizeof(float);
36    const uint32_t nvec = n / epv;
37    const uint32_t nloe = n % epv;
38    for (uint32_t i = 0; i < nvec; ++i) {
39        HVX_Vector vd   = hvx_vmemu(dst + i * epv);
40        HVX_Vector vm   = hvx_vmem(mul + i * epv);
41        HVX_Vector vdot = hvx_vmem(dot + i * epv);
42        HVX_Vector out  = hvx_vec_mul_f32_f32(vd, vm);
43        hvx_vmemu(dst + i * epv) = out;
44        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(out, vdot));
45    }
46
47    if (nloe) {
48        const uint32_t off = nvec * epv;
49        HVX_Vector vd   = hvx_vmemu(dst + off);
50        HVX_Vector vm   = hvx_vmem(mul + off);
51        HVX_Vector vdot = hvx_vmem(dot + off);
52        HVX_Vector out  = hvx_vec_mul_f32_f32(vd, vm);
53        hvx_vec_store_u(dst + off, nloe * sizeof(float), out);
54        HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float));
55        HVX_Vector prod = hvx_vec_mul_f32_f32(out, vdot);
56        acc = hvx_vec_add_f32_f32(acc, Q6_V_vmux_QVV(mask, prod, Q6_V_vzero()));
57    }
58
59    return hvx_vec_reduce_sum_f32(acc);
60}
61
62static inline HVX_Vector gdn_mul_scalar_dot_f32(float * restrict dst, float mul, const float * restrict dot, uint32_t n) {
63    HVX_Vector acc = Q6_V_vzero();
64    const HVX_Vector vmul = hvx_vec_splat_f32(mul);
65
66    const uint32_t epv  = 128 / sizeof(float);
67    const uint32_t nvec = n / epv;
68    const uint32_t nloe = n % epv;
69    for (uint32_t i = 0; i < nvec; ++i) {
70        HVX_Vector vd   = hvx_vmemu(dst + i * epv);
71        HVX_Vector vdot = hvx_vmem(dot + i * epv);
72        HVX_Vector out  = hvx_vec_mul_f32_f32(vd, vmul);
73        hvx_vmemu(dst + i * epv) = out;
74        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(out, vdot));
75    }
76
77    if (nloe) {
78        const uint32_t off = nvec * epv;
79        HVX_Vector vd   = hvx_vmemu(dst + off);
80        HVX_Vector vdot = hvx_vmem(dot + off);
81        HVX_Vector out  = hvx_vec_mul_f32_f32(vd, vmul);
82        hvx_vec_store_u(dst + off, nloe * sizeof(float), out);
83        HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float));
84        HVX_Vector prod = hvx_vec_mul_f32_f32(out, vdot);
85        acc = hvx_vec_add_f32_f32(acc, Q6_V_vmux_QVV(mask, prod, Q6_V_vzero()));
86    }
87
88    return hvx_vec_reduce_sum_f32(acc);
89}
90
91static inline HVX_Vector gdn_add_scaled_dot_f32(float * restrict dst, const float * restrict src,
92        HVX_Vector vscale, const float * restrict dot, uint32_t n) {
93    HVX_Vector acc = Q6_V_vzero();
94
95    const uint32_t epv  = 128 / sizeof(float);
96    const uint32_t nvec = n / epv;
97    const uint32_t nloe = n % epv;
98    for (uint32_t i = 0; i < nvec; ++i) {
99        HVX_Vector vd   = hvx_vmemu(dst + i * epv);
100        HVX_Vector vs   = hvx_vmem(src + i * epv);
101        HVX_Vector vdot = hvx_vmem(dot + i * epv);
102        HVX_Vector out  = hvx_vec_add_f32_f32(vd, hvx_vec_mul_f32_f32(vs, vscale));
103        hvx_vmemu(dst + i * epv) = out;
104        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(out, vdot));
105    }
106
107    if (nloe) {
108        const uint32_t off = nvec * epv;
109        HVX_Vector vd   = hvx_vmemu(dst + off);
110        HVX_Vector vs   = hvx_vmem(src + off);
111        HVX_Vector vdot = hvx_vmem(dot + off);
112        HVX_Vector out  = hvx_vec_add_f32_f32(vd, hvx_vec_mul_f32_f32(vs, vscale));
113        hvx_vec_store_u(dst + off, nloe * sizeof(float), out);
114        HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float));
115        HVX_Vector prod = hvx_vec_mul_f32_f32(out, vdot);
116        acc = hvx_vec_add_f32_f32(acc, Q6_V_vmux_QVV(mask, prod, Q6_V_vzero()));
117    }
118
119    return hvx_vec_reduce_sum_f32(acc);
120}
121
122static inline void gdn_mul_dot4_f32(float * restrict dst0, float * restrict dst1,
123        float * restrict dst2, float * restrict dst3, const float * restrict mul,
124        const float * restrict dot, uint32_t n, float * restrict sums) {
125    HVX_Vector acc0 = Q6_V_vzero();
126    HVX_Vector acc1 = Q6_V_vzero();
127    HVX_Vector acc2 = Q6_V_vzero();
128    HVX_Vector acc3 = Q6_V_vzero();
129
130    const uint32_t epv = 128 / sizeof(float);
131    const uint32_t nvec = n / epv;
132    const uint32_t nloe = n % epv;
133    for (uint32_t i = 0; i < nvec; ++i) {
134        HVX_Vector vm = hvx_vmem(mul + i * epv);
135        HVX_Vector vdot = hvx_vmem(dot + i * epv);
136
137        HVX_Vector out0 = hvx_vec_mul_f32_f32(hvx_vmemu(dst0 + i * epv), vm);
138        HVX_Vector out1 = hvx_vec_mul_f32_f32(hvx_vmemu(dst1 + i * epv), vm);
139        HVX_Vector out2 = hvx_vec_mul_f32_f32(hvx_vmemu(dst2 + i * epv), vm);
140        HVX_Vector out3 = hvx_vec_mul_f32_f32(hvx_vmemu(dst3 + i * epv), vm);
141
142        hvx_vmemu(dst0 + i * epv) = out0;
143        hvx_vmemu(dst1 + i * epv) = out1;
144        hvx_vmemu(dst2 + i * epv) = out2;
145        hvx_vmemu(dst3 + i * epv) = out3;
146
147        acc0 = hvx_vec_add_f32_f32(acc0, hvx_vec_mul_f32_f32(out0, vdot));
148        acc1 = hvx_vec_add_f32_f32(acc1, hvx_vec_mul_f32_f32(out1, vdot));
149        acc2 = hvx_vec_add_f32_f32(acc2, hvx_vec_mul_f32_f32(out2, vdot));
150        acc3 = hvx_vec_add_f32_f32(acc3, hvx_vec_mul_f32_f32(out3, vdot));
151    }
152
153    if (nloe) {
154        const uint32_t off = nvec * epv;
155        HVX_Vector vm   = hvx_vmem(mul + off);
156        HVX_Vector vdot = hvx_vmem(dot + off);
157        HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float));
158        HVX_Vector zero = Q6_V_vzero();
159
160        HVX_Vector out0 = hvx_vec_mul_f32_f32(hvx_vmemu(dst0 + off), vm);
161        HVX_Vector out1 = hvx_vec_mul_f32_f32(hvx_vmemu(dst1 + off), vm);
162        HVX_Vector out2 = hvx_vec_mul_f32_f32(hvx_vmemu(dst2 + off), vm);
163        HVX_Vector out3 = hvx_vec_mul_f32_f32(hvx_vmemu(dst3 + off), vm);
164
165        hvx_vec_store_u(dst0 + off, nloe * sizeof(float), out0);
166        hvx_vec_store_u(dst1 + off, nloe * sizeof(float), out1);
167        hvx_vec_store_u(dst2 + off, nloe * sizeof(float), out2);
168        hvx_vec_store_u(dst3 + off, nloe * sizeof(float), out3);
169
170        acc0 = hvx_vec_add_f32_f32(acc0, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out0, vdot), zero));
171        acc1 = hvx_vec_add_f32_f32(acc1, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out1, vdot), zero));
172        acc2 = hvx_vec_add_f32_f32(acc2, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out2, vdot), zero));
173        acc3 = hvx_vec_add_f32_f32(acc3, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out3, vdot), zero));
174    }
175
176    HVX_Vector_x4 acc = { .v = { acc0, acc1, acc2, acc3 } };
177    hvx_vec_store_u(sums, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(acc));
178}
179
180static inline void gdn_mul_scalar_dot4_f32(float * restrict dst0, float * restrict dst1,
181        float * restrict dst2, float * restrict dst3, float mul,
182        const float * restrict dot, uint32_t n, float * restrict sums) {
183    HVX_Vector acc0 = Q6_V_vzero();
184    HVX_Vector acc1 = Q6_V_vzero();
185    HVX_Vector acc2 = Q6_V_vzero();
186    HVX_Vector acc3 = Q6_V_vzero();
187    const HVX_Vector vmul = hvx_vec_splat_f32(mul);
188
189    const uint32_t epv = 128 / sizeof(float);
190    const uint32_t nvec = n / epv;
191    const uint32_t nloe = n % epv;
192    for (uint32_t i = 0; i < nvec; ++i) {
193        HVX_Vector vdot = hvx_vmem(dot + i * epv);
194
195        HVX_Vector out0 = hvx_vec_mul_f32_f32(hvx_vmemu(dst0 + i * epv), vmul);
196        HVX_Vector out1 = hvx_vec_mul_f32_f32(hvx_vmemu(dst1 + i * epv), vmul);
197        HVX_Vector out2 = hvx_vec_mul_f32_f32(hvx_vmemu(dst2 + i * epv), vmul);
198        HVX_Vector out3 = hvx_vec_mul_f32_f32(hvx_vmemu(dst3 + i * epv), vmul);
199
200        hvx_vmemu(dst0 + i * epv) = out0;
201        hvx_vmemu(dst1 + i * epv) = out1;
202        hvx_vmemu(dst2 + i * epv) = out2;
203        hvx_vmemu(dst3 + i * epv) = out3;
204
205        acc0 = hvx_vec_add_f32_f32(acc0, hvx_vec_mul_f32_f32(out0, vdot));
206        acc1 = hvx_vec_add_f32_f32(acc1, hvx_vec_mul_f32_f32(out1, vdot));
207        acc2 = hvx_vec_add_f32_f32(acc2, hvx_vec_mul_f32_f32(out2, vdot));
208        acc3 = hvx_vec_add_f32_f32(acc3, hvx_vec_mul_f32_f32(out3, vdot));
209    }
210
211    if (nloe) {
212        const uint32_t off = nvec * epv;
213        HVX_Vector vdot = hvx_vmem(dot + off);
214        HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float));
215        HVX_Vector zero = Q6_V_vzero();
216
217        HVX_Vector out0 = hvx_vec_mul_f32_f32(hvx_vmemu(dst0 + off), vmul);
218        HVX_Vector out1 = hvx_vec_mul_f32_f32(hvx_vmemu(dst1 + off), vmul);
219        HVX_Vector out2 = hvx_vec_mul_f32_f32(hvx_vmemu(dst2 + off), vmul);
220        HVX_Vector out3 = hvx_vec_mul_f32_f32(hvx_vmemu(dst3 + off), vmul);
221
222        hvx_vec_store_u(dst0 + off, nloe * sizeof(float), out0);
223        hvx_vec_store_u(dst1 + off, nloe * sizeof(float), out1);
224        hvx_vec_store_u(dst2 + off, nloe * sizeof(float), out2);
225        hvx_vec_store_u(dst3 + off, nloe * sizeof(float), out3);
226
227        acc0 = hvx_vec_add_f32_f32(acc0, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out0, vdot), zero));
228        acc1 = hvx_vec_add_f32_f32(acc1, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out1, vdot), zero));
229        acc2 = hvx_vec_add_f32_f32(acc2, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out2, vdot), zero));
230        acc3 = hvx_vec_add_f32_f32(acc3, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out3, vdot), zero));
231    }
232
233    HVX_Vector_x4 acc = { .v = { acc0, acc1, acc2, acc3 } };
234    hvx_vec_store_u(sums, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(acc));
235}
236
237static inline void gdn_add_scaled_dot4_f32(float * restrict dst0, float * restrict dst1,
238        float * restrict dst2, float * restrict dst3, const float * restrict src,
239        const float * restrict scale, const float * restrict dot, uint32_t n,
240        float * restrict sums) {
241    HVX_Vector acc0 = Q6_V_vzero();
242    HVX_Vector acc1 = Q6_V_vzero();
243    HVX_Vector acc2 = Q6_V_vzero();
244    HVX_Vector acc3 = Q6_V_vzero();
245    const HVX_Vector scale0 = hvx_vec_splat_f32(scale[0]);
246    const HVX_Vector scale1 = hvx_vec_splat_f32(scale[1]);
247    const HVX_Vector scale2 = hvx_vec_splat_f32(scale[2]);
248    const HVX_Vector scale3 = hvx_vec_splat_f32(scale[3]);
249
250    const uint32_t epv = 128 / sizeof(float);
251    const uint32_t nvec = n / epv;
252    const uint32_t nloe = n % epv;
253    for (uint32_t i = 0; i < nvec; ++i) {
254        HVX_Vector vs = hvx_vmem(src + i * epv);
255        HVX_Vector vdot = hvx_vmem(dot + i * epv);
256
257        HVX_Vector out0 = hvx_vec_add_f32_f32(hvx_vmemu(dst0 + i * epv), hvx_vec_mul_f32_f32(vs, scale0));
258        HVX_Vector out1 = hvx_vec_add_f32_f32(hvx_vmemu(dst1 + i * epv), hvx_vec_mul_f32_f32(vs, scale1));
259        HVX_Vector out2 = hvx_vec_add_f32_f32(hvx_vmemu(dst2 + i * epv), hvx_vec_mul_f32_f32(vs, scale2));
260        HVX_Vector out3 = hvx_vec_add_f32_f32(hvx_vmemu(dst3 + i * epv), hvx_vec_mul_f32_f32(vs, scale3));
261
262        hvx_vmemu(dst0 + i * epv) = out0;
263        hvx_vmemu(dst1 + i * epv) = out1;
264        hvx_vmemu(dst2 + i * epv) = out2;
265        hvx_vmemu(dst3 + i * epv) = out3;
266
267        acc0 = hvx_vec_add_f32_f32(acc0, hvx_vec_mul_f32_f32(out0, vdot));
268        acc1 = hvx_vec_add_f32_f32(acc1, hvx_vec_mul_f32_f32(out1, vdot));
269        acc2 = hvx_vec_add_f32_f32(acc2, hvx_vec_mul_f32_f32(out2, vdot));
270        acc3 = hvx_vec_add_f32_f32(acc3, hvx_vec_mul_f32_f32(out3, vdot));
271    }
272
273    if (nloe) {
274        const uint32_t off = nvec * epv;
275        HVX_Vector vs = hvx_vmem(src + off);
276        HVX_Vector vdot = hvx_vmem(dot + off);
277        HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float));
278        HVX_Vector zero = Q6_V_vzero();
279
280        HVX_Vector out0 = hvx_vec_add_f32_f32(hvx_vmemu(dst0 + off), hvx_vec_mul_f32_f32(vs, scale0));
281        HVX_Vector out1 = hvx_vec_add_f32_f32(hvx_vmemu(dst1 + off), hvx_vec_mul_f32_f32(vs, scale1));
282        HVX_Vector out2 = hvx_vec_add_f32_f32(hvx_vmemu(dst2 + off), hvx_vec_mul_f32_f32(vs, scale2));
283        HVX_Vector out3 = hvx_vec_add_f32_f32(hvx_vmemu(dst3 + off), hvx_vec_mul_f32_f32(vs, scale3));
284
285        hvx_vec_store_u(dst0 + off, nloe * sizeof(float), out0);
286        hvx_vec_store_u(dst1 + off, nloe * sizeof(float), out1);
287        hvx_vec_store_u(dst2 + off, nloe * sizeof(float), out2);
288        hvx_vec_store_u(dst3 + off, nloe * sizeof(float), out3);
289
290        acc0 = hvx_vec_add_f32_f32(acc0, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out0, vdot), zero));
291        acc1 = hvx_vec_add_f32_f32(acc1, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out1, vdot), zero));
292        acc2 = hvx_vec_add_f32_f32(acc2, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out2, vdot), zero));
293        acc3 = hvx_vec_add_f32_f32(acc3, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out3, vdot), zero));
294    }
295
296    HVX_Vector_x4 acc = { .v = { acc0, acc1, acc2, acc3 } };
297    hvx_vec_store_u(sums, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(acc));
298}
299
300static inline void gdn_mul_dot8_f32(float * restrict dst0, float * restrict dst1,
301        float * restrict dst2, float * restrict dst3, float * restrict dst4,
302        float * restrict dst5, float * restrict dst6, float * restrict dst7,
303        const float * restrict mul, const float * restrict dot, uint32_t n,
304        float * restrict sums) {
305    HVX_Vector acc0 = Q6_V_vzero();
306    HVX_Vector acc1 = Q6_V_vzero();
307    HVX_Vector acc2 = Q6_V_vzero();
308    HVX_Vector acc3 = Q6_V_vzero();
309    HVX_Vector acc4 = Q6_V_vzero();
310    HVX_Vector acc5 = Q6_V_vzero();
311    HVX_Vector acc6 = Q6_V_vzero();
312    HVX_Vector acc7 = Q6_V_vzero();
313
314    const uint32_t epv = 128 / sizeof(float);
315    const uint32_t nvec = n / epv;
316    const uint32_t nloe = n % epv;
317    for (uint32_t i = 0; i < nvec; ++i) {
318        HVX_Vector vm = hvx_vmem(mul + i * epv);
319        HVX_Vector vdot = hvx_vmem(dot + i * epv);
320
321        HVX_Vector out0 = hvx_vec_mul_f32_f32(hvx_vmemu(dst0 + i * epv), vm);
322        HVX_Vector out1 = hvx_vec_mul_f32_f32(hvx_vmemu(dst1 + i * epv), vm);
323        HVX_Vector out2 = hvx_vec_mul_f32_f32(hvx_vmemu(dst2 + i * epv), vm);
324        HVX_Vector out3 = hvx_vec_mul_f32_f32(hvx_vmemu(dst3 + i * epv), vm);
325        HVX_Vector out4 = hvx_vec_mul_f32_f32(hvx_vmemu(dst4 + i * epv), vm);
326        HVX_Vector out5 = hvx_vec_mul_f32_f32(hvx_vmemu(dst5 + i * epv), vm);
327        HVX_Vector out6 = hvx_vec_mul_f32_f32(hvx_vmemu(dst6 + i * epv), vm);
328        HVX_Vector out7 = hvx_vec_mul_f32_f32(hvx_vmemu(dst7 + i * epv), vm);
329
330        hvx_vmemu(dst0 + i * epv) = out0;
331        hvx_vmemu(dst1 + i * epv) = out1;
332        hvx_vmemu(dst2 + i * epv) = out2;
333        hvx_vmemu(dst3 + i * epv) = out3;
334        hvx_vmemu(dst4 + i * epv) = out4;
335        hvx_vmemu(dst5 + i * epv) = out5;
336        hvx_vmemu(dst6 + i * epv) = out6;
337        hvx_vmemu(dst7 + i * epv) = out7;
338
339        acc0 = hvx_vec_add_f32_f32(acc0, hvx_vec_mul_f32_f32(out0, vdot));
340        acc1 = hvx_vec_add_f32_f32(acc1, hvx_vec_mul_f32_f32(out1, vdot));
341        acc2 = hvx_vec_add_f32_f32(acc2, hvx_vec_mul_f32_f32(out2, vdot));
342        acc3 = hvx_vec_add_f32_f32(acc3, hvx_vec_mul_f32_f32(out3, vdot));
343        acc4 = hvx_vec_add_f32_f32(acc4, hvx_vec_mul_f32_f32(out4, vdot));
344        acc5 = hvx_vec_add_f32_f32(acc5, hvx_vec_mul_f32_f32(out5, vdot));
345        acc6 = hvx_vec_add_f32_f32(acc6, hvx_vec_mul_f32_f32(out6, vdot));
346        acc7 = hvx_vec_add_f32_f32(acc7, hvx_vec_mul_f32_f32(out7, vdot));
347    }
348
349    if (nloe) {
350        const uint32_t off = nvec * epv;
351        HVX_Vector vm = hvx_vmem(mul + off);
352        HVX_Vector vdot = hvx_vmem(dot + off);
353        HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float));
354        HVX_Vector zero = Q6_V_vzero();
355
356        HVX_Vector out0 = hvx_vec_mul_f32_f32(hvx_vmemu(dst0 + off), vm);
357        HVX_Vector out1 = hvx_vec_mul_f32_f32(hvx_vmemu(dst1 + off), vm);
358        HVX_Vector out2 = hvx_vec_mul_f32_f32(hvx_vmemu(dst2 + off), vm);
359        HVX_Vector out3 = hvx_vec_mul_f32_f32(hvx_vmemu(dst3 + off), vm);
360        HVX_Vector out4 = hvx_vec_mul_f32_f32(hvx_vmemu(dst4 + off), vm);
361        HVX_Vector out5 = hvx_vec_mul_f32_f32(hvx_vmemu(dst5 + off), vm);
362        HVX_Vector out6 = hvx_vec_mul_f32_f32(hvx_vmemu(dst6 + off), vm);
363        HVX_Vector out7 = hvx_vec_mul_f32_f32(hvx_vmemu(dst7 + off), vm);
364
365        hvx_vec_store_u(dst0 + off, nloe * sizeof(float), out0);
366        hvx_vec_store_u(dst1 + off, nloe * sizeof(float), out1);
367        hvx_vec_store_u(dst2 + off, nloe * sizeof(float), out2);
368        hvx_vec_store_u(dst3 + off, nloe * sizeof(float), out3);
369        hvx_vec_store_u(dst4 + off, nloe * sizeof(float), out4);
370        hvx_vec_store_u(dst5 + off, nloe * sizeof(float), out5);
371        hvx_vec_store_u(dst6 + off, nloe * sizeof(float), out6);
372        hvx_vec_store_u(dst7 + off, nloe * sizeof(float), out7);
373
374        acc0 = hvx_vec_add_f32_f32(acc0, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out0, vdot), zero));
375        acc1 = hvx_vec_add_f32_f32(acc1, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out1, vdot), zero));
376        acc2 = hvx_vec_add_f32_f32(acc2, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out2, vdot), zero));
377        acc3 = hvx_vec_add_f32_f32(acc3, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out3, vdot), zero));
378        acc4 = hvx_vec_add_f32_f32(acc4, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out4, vdot), zero));
379        acc5 = hvx_vec_add_f32_f32(acc5, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out5, vdot), zero));
380        acc6 = hvx_vec_add_f32_f32(acc6, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out6, vdot), zero));
381        acc7 = hvx_vec_add_f32_f32(acc7, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out7, vdot), zero));
382    }
383
384    HVX_Vector_x4 accA = { .v = { acc0, acc1, acc2, acc3 } };
385    HVX_Vector_x4 accB = { .v = { acc4, acc5, acc6, acc7 } };
386    hvx_vec_store_u(sums + 0, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(accA));
387    hvx_vec_store_u(sums + 4, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(accB));
388}
389
390static inline void gdn_mul_scalar_dot8_f32(float * restrict dst0, float * restrict dst1,
391        float * restrict dst2, float * restrict dst3, float * restrict dst4,
392        float * restrict dst5, float * restrict dst6, float * restrict dst7,
393        float mul, const float * restrict dot, uint32_t n, float * restrict sums) {
394    HVX_Vector acc0 = Q6_V_vzero();
395    HVX_Vector acc1 = Q6_V_vzero();
396    HVX_Vector acc2 = Q6_V_vzero();
397    HVX_Vector acc3 = Q6_V_vzero();
398    HVX_Vector acc4 = Q6_V_vzero();
399    HVX_Vector acc5 = Q6_V_vzero();
400    HVX_Vector acc6 = Q6_V_vzero();
401    HVX_Vector acc7 = Q6_V_vzero();
402    const HVX_Vector vmul = hvx_vec_splat_f32(mul);
403
404    const uint32_t epv = 128 / sizeof(float);
405    const uint32_t nvec = n / epv;
406    const uint32_t nloe = n % epv;
407    for (uint32_t i = 0; i < nvec; ++i) {
408        HVX_Vector vdot = hvx_vmem(dot + i * epv);
409
410        HVX_Vector out0 = hvx_vec_mul_f32_f32(hvx_vmemu(dst0 + i * epv), vmul);
411        HVX_Vector out1 = hvx_vec_mul_f32_f32(hvx_vmemu(dst1 + i * epv), vmul);
412        HVX_Vector out2 = hvx_vec_mul_f32_f32(hvx_vmemu(dst2 + i * epv), vmul);
413        HVX_Vector out3 = hvx_vec_mul_f32_f32(hvx_vmemu(dst3 + i * epv), vmul);
414        HVX_Vector out4 = hvx_vec_mul_f32_f32(hvx_vmemu(dst4 + i * epv), vmul);
415        HVX_Vector out5 = hvx_vec_mul_f32_f32(hvx_vmemu(dst5 + i * epv), vmul);
416        HVX_Vector out6 = hvx_vec_mul_f32_f32(hvx_vmemu(dst6 + i * epv), vmul);
417        HVX_Vector out7 = hvx_vec_mul_f32_f32(hvx_vmemu(dst7 + i * epv), vmul);
418
419        hvx_vmemu(dst0 + i * epv) = out0;
420        hvx_vmemu(dst1 + i * epv) = out1;
421        hvx_vmemu(dst2 + i * epv) = out2;
422        hvx_vmemu(dst3 + i * epv) = out3;
423        hvx_vmemu(dst4 + i * epv) = out4;
424        hvx_vmemu(dst5 + i * epv) = out5;
425        hvx_vmemu(dst6 + i * epv) = out6;
426        hvx_vmemu(dst7 + i * epv) = out7;
427
428        acc0 = hvx_vec_add_f32_f32(acc0, hvx_vec_mul_f32_f32(out0, vdot));
429        acc1 = hvx_vec_add_f32_f32(acc1, hvx_vec_mul_f32_f32(out1, vdot));
430        acc2 = hvx_vec_add_f32_f32(acc2, hvx_vec_mul_f32_f32(out2, vdot));
431        acc3 = hvx_vec_add_f32_f32(acc3, hvx_vec_mul_f32_f32(out3, vdot));
432        acc4 = hvx_vec_add_f32_f32(acc4, hvx_vec_mul_f32_f32(out4, vdot));
433        acc5 = hvx_vec_add_f32_f32(acc5, hvx_vec_mul_f32_f32(out5, vdot));
434        acc6 = hvx_vec_add_f32_f32(acc6, hvx_vec_mul_f32_f32(out6, vdot));
435        acc7 = hvx_vec_add_f32_f32(acc7, hvx_vec_mul_f32_f32(out7, vdot));
436    }
437
438    if (nloe) {
439        const uint32_t off = nvec * epv;
440        HVX_Vector vdot = hvx_vmem(dot + off);
441        HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float));
442        HVX_Vector zero = Q6_V_vzero();
443
444        HVX_Vector out0 = hvx_vec_mul_f32_f32(hvx_vmemu(dst0 + off), vmul);
445        HVX_Vector out1 = hvx_vec_mul_f32_f32(hvx_vmemu(dst1 + off), vmul);
446        HVX_Vector out2 = hvx_vec_mul_f32_f32(hvx_vmemu(dst2 + off), vmul);
447        HVX_Vector out3 = hvx_vec_mul_f32_f32(hvx_vmemu(dst3 + off), vmul);
448        HVX_Vector out4 = hvx_vec_mul_f32_f32(hvx_vmemu(dst4 + off), vmul);
449        HVX_Vector out5 = hvx_vec_mul_f32_f32(hvx_vmemu(dst5 + off), vmul);
450        HVX_Vector out6 = hvx_vec_mul_f32_f32(hvx_vmemu(dst6 + off), vmul);
451        HVX_Vector out7 = hvx_vec_mul_f32_f32(hvx_vmemu(dst7 + off), vmul);
452
453        hvx_vec_store_u(dst0 + off, nloe * sizeof(float), out0);
454        hvx_vec_store_u(dst1 + off, nloe * sizeof(float), out1);
455        hvx_vec_store_u(dst2 + off, nloe * sizeof(float), out2);
456        hvx_vec_store_u(dst3 + off, nloe * sizeof(float), out3);
457        hvx_vec_store_u(dst4 + off, nloe * sizeof(float), out4);
458        hvx_vec_store_u(dst5 + off, nloe * sizeof(float), out5);
459        hvx_vec_store_u(dst6 + off, nloe * sizeof(float), out6);
460        hvx_vec_store_u(dst7 + off, nloe * sizeof(float), out7);
461
462        acc0 = hvx_vec_add_f32_f32(acc0, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out0, vdot), zero));
463        acc1 = hvx_vec_add_f32_f32(acc1, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out1, vdot), zero));
464        acc2 = hvx_vec_add_f32_f32(acc2, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out2, vdot), zero));
465        acc3 = hvx_vec_add_f32_f32(acc3, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out3, vdot), zero));
466        acc4 = hvx_vec_add_f32_f32(acc4, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out4, vdot), zero));
467        acc5 = hvx_vec_add_f32_f32(acc5, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out5, vdot), zero));
468        acc6 = hvx_vec_add_f32_f32(acc6, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out6, vdot), zero));
469        acc7 = hvx_vec_add_f32_f32(acc7, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out7, vdot), zero));
470    }
471
472    HVX_Vector_x4 accA = { .v = { acc0, acc1, acc2, acc3 } };
473    HVX_Vector_x4 accB = { .v = { acc4, acc5, acc6, acc7 } };
474    hvx_vec_store_u(sums + 0, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(accA));
475    hvx_vec_store_u(sums + 4, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(accB));
476}
477
478static inline void gdn_add_scaled_dot8_f32(float * restrict dst0, float * restrict dst1,
479        float * restrict dst2, float * restrict dst3, float * restrict dst4,
480        float * restrict dst5, float * restrict dst6, float * restrict dst7,
481        const float * restrict src, const float * restrict scale,
482        const float * restrict dot, uint32_t n, float * restrict sums) {
483    HVX_Vector acc0 = Q6_V_vzero();
484    HVX_Vector acc1 = Q6_V_vzero();
485    HVX_Vector acc2 = Q6_V_vzero();
486    HVX_Vector acc3 = Q6_V_vzero();
487    HVX_Vector acc4 = Q6_V_vzero();
488    HVX_Vector acc5 = Q6_V_vzero();
489    HVX_Vector acc6 = Q6_V_vzero();
490    HVX_Vector acc7 = Q6_V_vzero();
491    const HVX_Vector scale0 = hvx_vec_splat_f32(scale[0]);
492    const HVX_Vector scale1 = hvx_vec_splat_f32(scale[1]);
493    const HVX_Vector scale2 = hvx_vec_splat_f32(scale[2]);
494    const HVX_Vector scale3 = hvx_vec_splat_f32(scale[3]);
495    const HVX_Vector scale4 = hvx_vec_splat_f32(scale[4]);
496    const HVX_Vector scale5 = hvx_vec_splat_f32(scale[5]);
497    const HVX_Vector scale6 = hvx_vec_splat_f32(scale[6]);
498    const HVX_Vector scale7 = hvx_vec_splat_f32(scale[7]);
499
500    const uint32_t epv = 128 / sizeof(float);
501    const uint32_t nvec = n / epv;
502    const uint32_t nloe = n % epv;
503    for (uint32_t i = 0; i < nvec; ++i) {
504        HVX_Vector vs = hvx_vmem(src + i * epv);
505        HVX_Vector vdot = hvx_vmem(dot + i * epv);
506
507        HVX_Vector out0 = hvx_vec_add_f32_f32(hvx_vmemu(dst0 + i * epv), hvx_vec_mul_f32_f32(vs, scale0));
508        HVX_Vector out1 = hvx_vec_add_f32_f32(hvx_vmemu(dst1 + i * epv), hvx_vec_mul_f32_f32(vs, scale1));
509        HVX_Vector out2 = hvx_vec_add_f32_f32(hvx_vmemu(dst2 + i * epv), hvx_vec_mul_f32_f32(vs, scale2));
510        HVX_Vector out3 = hvx_vec_add_f32_f32(hvx_vmemu(dst3 + i * epv), hvx_vec_mul_f32_f32(vs, scale3));
511        HVX_Vector out4 = hvx_vec_add_f32_f32(hvx_vmemu(dst4 + i * epv), hvx_vec_mul_f32_f32(vs, scale4));
512        HVX_Vector out5 = hvx_vec_add_f32_f32(hvx_vmemu(dst5 + i * epv), hvx_vec_mul_f32_f32(vs, scale5));
513        HVX_Vector out6 = hvx_vec_add_f32_f32(hvx_vmemu(dst6 + i * epv), hvx_vec_mul_f32_f32(vs, scale6));
514        HVX_Vector out7 = hvx_vec_add_f32_f32(hvx_vmemu(dst7 + i * epv), hvx_vec_mul_f32_f32(vs, scale7));
515
516        hvx_vmemu(dst0 + i * epv) = out0;
517        hvx_vmemu(dst1 + i * epv) = out1;
518        hvx_vmemu(dst2 + i * epv) = out2;
519        hvx_vmemu(dst3 + i * epv) = out3;
520        hvx_vmemu(dst4 + i * epv) = out4;
521        hvx_vmemu(dst5 + i * epv) = out5;
522        hvx_vmemu(dst6 + i * epv) = out6;
523        hvx_vmemu(dst7 + i * epv) = out7;
524
525        acc0 = hvx_vec_add_f32_f32(acc0, hvx_vec_mul_f32_f32(out0, vdot));
526        acc1 = hvx_vec_add_f32_f32(acc1, hvx_vec_mul_f32_f32(out1, vdot));
527        acc2 = hvx_vec_add_f32_f32(acc2, hvx_vec_mul_f32_f32(out2, vdot));
528        acc3 = hvx_vec_add_f32_f32(acc3, hvx_vec_mul_f32_f32(out3, vdot));
529        acc4 = hvx_vec_add_f32_f32(acc4, hvx_vec_mul_f32_f32(out4, vdot));
530        acc5 = hvx_vec_add_f32_f32(acc5, hvx_vec_mul_f32_f32(out5, vdot));
531        acc6 = hvx_vec_add_f32_f32(acc6, hvx_vec_mul_f32_f32(out6, vdot));
532        acc7 = hvx_vec_add_f32_f32(acc7, hvx_vec_mul_f32_f32(out7, vdot));
533    }
534
535    if (nloe) {
536        const uint32_t off = nvec * epv;
537        HVX_Vector vs = hvx_vmem(src + off);
538        HVX_Vector vdot = hvx_vmem(dot + off);
539        HVX_VectorPred mask = Q6_Q_vsetq2_R(nloe * sizeof(float));
540        HVX_Vector zero = Q6_V_vzero();
541
542        HVX_Vector out0 = hvx_vec_add_f32_f32(hvx_vmemu(dst0 + off), hvx_vec_mul_f32_f32(vs, scale0));
543        HVX_Vector out1 = hvx_vec_add_f32_f32(hvx_vmemu(dst1 + off), hvx_vec_mul_f32_f32(vs, scale1));
544        HVX_Vector out2 = hvx_vec_add_f32_f32(hvx_vmemu(dst2 + off), hvx_vec_mul_f32_f32(vs, scale2));
545        HVX_Vector out3 = hvx_vec_add_f32_f32(hvx_vmemu(dst3 + off), hvx_vec_mul_f32_f32(vs, scale3));
546        HVX_Vector out4 = hvx_vec_add_f32_f32(hvx_vmemu(dst4 + off), hvx_vec_mul_f32_f32(vs, scale4));
547        HVX_Vector out5 = hvx_vec_add_f32_f32(hvx_vmemu(dst5 + off), hvx_vec_mul_f32_f32(vs, scale5));
548        HVX_Vector out6 = hvx_vec_add_f32_f32(hvx_vmemu(dst6 + off), hvx_vec_mul_f32_f32(vs, scale6));
549        HVX_Vector out7 = hvx_vec_add_f32_f32(hvx_vmemu(dst7 + off), hvx_vec_mul_f32_f32(vs, scale7));
550
551        hvx_vec_store_u(dst0 + off, nloe * sizeof(float), out0);
552        hvx_vec_store_u(dst1 + off, nloe * sizeof(float), out1);
553        hvx_vec_store_u(dst2 + off, nloe * sizeof(float), out2);
554        hvx_vec_store_u(dst3 + off, nloe * sizeof(float), out3);
555        hvx_vec_store_u(dst4 + off, nloe * sizeof(float), out4);
556        hvx_vec_store_u(dst5 + off, nloe * sizeof(float), out5);
557        hvx_vec_store_u(dst6 + off, nloe * sizeof(float), out6);
558        hvx_vec_store_u(dst7 + off, nloe * sizeof(float), out7);
559
560        acc0 = hvx_vec_add_f32_f32(acc0, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out0, vdot), zero));
561        acc1 = hvx_vec_add_f32_f32(acc1, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out1, vdot), zero));
562        acc2 = hvx_vec_add_f32_f32(acc2, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out2, vdot), zero));
563        acc3 = hvx_vec_add_f32_f32(acc3, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out3, vdot), zero));
564        acc4 = hvx_vec_add_f32_f32(acc4, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out4, vdot), zero));
565        acc5 = hvx_vec_add_f32_f32(acc5, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out5, vdot), zero));
566        acc6 = hvx_vec_add_f32_f32(acc6, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out6, vdot), zero));
567        acc7 = hvx_vec_add_f32_f32(acc7, Q6_V_vmux_QVV(mask, hvx_vec_mul_f32_f32(out7, vdot), zero));
568    }
569
570    HVX_Vector_x4 accA = { .v = { acc0, acc1, acc2, acc3 } };
571    HVX_Vector_x4 accB = { .v = { acc4, acc5, acc6, acc7 } };
572    hvx_vec_store_u(sums + 0, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(accA));
573    hvx_vec_store_u(sums + 4, 4 * sizeof(float), hvx_vec_reduce_sum_f32x4(accB));
574}
575
576static void gated_delta_net_f32_pp_thread(unsigned int nth, unsigned int ith, void * data) {
577    struct htp_gdn_context * gctx = (struct htp_gdn_context *) data;
578    struct htp_ops_context * octx = gctx->octx;
579
580    const struct htp_tensor * q     = octx->src[0];
581    const struct htp_tensor * k     = octx->src[1];
582    const struct htp_tensor * v     = octx->src[2];
583    const struct htp_tensor * g     = octx->src[3];
584    const struct htp_tensor * beta  = octx->src[4];
585    const struct htp_tensor * state = octx->src[5];
586    const struct htp_tensor * dst   = octx->dst;
587
588    const uint32_t S_v      = v->ne[0];
589    const uint32_t H        = v->ne[1];
590    const uint32_t n_tokens = v->ne[2];
591    const uint32_t n_seqs   = v->ne[3];
592    const uint32_t K        = octx->op_params[0];
593
594    const uint32_t row_end = gctx->row_start + gctx->nrows;
595
596    if (ith >= gctx->nrows) {
597        return;
598    }
599
600    const uint32_t rq3 = n_seqs / q->ne[3];
601    const uint32_t rk3 = n_seqs / k->ne[3];
602    const float scale = 1.0f / sqrtf((float) S_v);
603
604    float * dst_base       = (float *) (uintptr_t) dst->data;
605    float * state_out_base = dst_base + (uint64_t) S_v * H * n_tokens * n_seqs;
606    const float * state_in_base = (const float *) (uintptr_t) state->data;
607
608    const bool kda = (g->ne[0] == S_v);
609    float local_gate[HTP_GDN_MAX_SV] __attribute__((aligned(128)));
610    float local_q[HTP_GDN_MAX_SV] __attribute__((aligned(128)));
611    float local_k[HTP_GDN_MAX_SV] __attribute__((aligned(128)));
612    float local_sums[32] __attribute__((aligned(128)));
613
614    dma_queue * dma = octx->ctx->dma[ith];
615    size_t state_aligned = (size_t) S_v * S_v * sizeof(float);
616    state_aligned = (state_aligned + 127) & ~(size_t)127;
617    float * s_work[2];
618    s_work[0] = (float *) (gctx->vtcm_base + gctx->vtcm_per_thread * ith);
619    s_work[1] = s_work[0] + state_aligned / sizeof(float);
620
621    struct fastdiv_values fd_H = init_fastdiv_values(H);
622    struct fastdiv_values fd_q1 = init_fastdiv_values(q->ne[1]);
623    struct fastdiv_values fd_k1 = init_fastdiv_values(k->ne[1]);
624    struct fastdiv_values fd_rq3 = init_fastdiv_values(rq3);
625    struct fastdiv_values fd_rk3 = init_fastdiv_values(rk3);
626
627    const uint64_t state_seq_stride = state->nb[3] / sizeof(float);
628    const uint64_t state_size_per_snap = (uint64_t) S_v * S_v * H * n_seqs;
629
630    uint32_t ir_prefetch = gctx->row_start + ith;
631    int spad_idx = 0;
632
633    // Prefetch preamble (up to 2 steps)
634    for (int k = 0; k < 2 && ir_prefetch < row_end; k++) {
635        const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H);
636        const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H);
637        const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v;
638        // final state lands in snapshot slot 0 (most-recent-first ordering)
639        float * ps_out = state_out_base + ((uint64_t) piv3 * H + piv1) * S_v * S_v;
640
641        // Push dummy write-back
642        dma_queue_push(dma, dma_make_ptr(ps_out, s_work[spad_idx]),
643                       S_v * sizeof(float), S_v * sizeof(float),
644                       S_v * sizeof(float), 0);
645
646        // Push fetch
647        dma_queue_push(dma, dma_make_ptr(s_work[spad_idx], ps_in),
648                       S_v * sizeof(float), S_v * sizeof(float),
649                       S_v * sizeof(float), S_v);
650
651        ir_prefetch += nth;
652        spad_idx ^= 1;
653    }
654
655    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
656    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (gctx->row_start + ith));
657
658    int curr_spad_idx = 0;
659    for (uint32_t ir = gctx->row_start + ith; ir < row_end; ir += nth) {
660        dma_queue_pop(dma);
661        dma_queue_pop(dma);
662
663        float * s_work_curr = s_work[curr_spad_idx];
664
665        const uint32_t iv1 = fastmodulo(ir, H, &fd_H);
666        const uint32_t iv3 = fastdiv(ir, &fd_H);
667
668        const uint32_t iq1 = fastmodulo(iv1, q->ne[1], &fd_q1);
669        const uint32_t ik1 = fastmodulo(iv1, k->ne[1], &fd_k1);
670        const uint32_t iq3 = fastdiv(iv3, &fd_rq3);
671        const uint32_t ik3 = fastdiv(iv3, &fd_rk3);
672
673        // final state lands in snapshot slot 0 (most-recent-first ordering)
674        float * s_out = state_out_base + ((uint64_t) iv3 * H + iv1) * S_v * S_v;
675
676        float * attn_data = dst_base + ((uint64_t) iv3 * n_tokens * H + iv1) * S_v;
677
678        for (uint32_t t = 0; t < n_tokens; ++t) {
679            const float * q_t = (const float *) ((const uint8_t *) (uintptr_t) q->data +
680                    (uint64_t) iq3 * q->nb[3] + (uint64_t) t * q->nb[2] + (uint64_t) iq1 * q->nb[1]);
681            const float * k_t = (const float *) ((const uint8_t *) (uintptr_t) k->data +
682                    (uint64_t) ik3 * k->nb[3] + (uint64_t) t * k->nb[2] + (uint64_t) ik1 * k->nb[1]);
683            const float * v_t = (const float *) ((const uint8_t *) (uintptr_t) v->data +
684                    (uint64_t) iv3 * v->nb[3] + (uint64_t) t * v->nb[2] + (uint64_t) iv1 * v->nb[1]);
685            const float * g_t = (const float *) ((const uint8_t *) (uintptr_t) g->data +
686                    (uint64_t) iv3 * g->nb[3] + (uint64_t) t * g->nb[2] + (uint64_t) iv1 * g->nb[1]);
687            const float beta_val = *(const float *) ((const uint8_t *) (uintptr_t) beta->data +
688                    (uint64_t) iv3 * beta->nb[3] + (uint64_t) t * beta->nb[2] + (uint64_t) iv1 * beta->nb[1]);
689
690            hvx_copy_f32_au((uint8_t *) local_q, (const uint8_t *) q_t, S_v);
691            hvx_copy_f32_au((uint8_t *) local_k, (const uint8_t *) k_t, S_v);
692
693            if (kda) {
694                hvx_exp_f32((uint8_t *) local_gate, (const uint8_t *) g_t, S_v, false);
695
696                uint32_t j = 0;
697                for (; j + 8 <= S_v; j += 8) {
698                    float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v;
699                    float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v;
700                    float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v;
701                    float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v;
702                    float * row4 = s_work_curr + (uint64_t) (j + 4) * S_v;
703                    float * row5 = s_work_curr + (uint64_t) (j + 5) * S_v;
704                    float * row6 = s_work_curr + (uint64_t) (j + 6) * S_v;
705                    float * row7 = s_work_curr + (uint64_t) (j + 7) * S_v;
706                    gdn_mul_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7,
707                                     local_gate, local_k, S_v, local_sums);
708
709                    float local_delta_b[32] __attribute__((aligned(128)));
710                    HVX_Vector vv_t = hvx_vmemu(v_t + j);
711                    HVX_Vector v_local_sums = hvx_vmem(local_sums);
712                    HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums);
713                    hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val));
714
715                    gdn_add_scaled_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7,
716                                            local_k, local_delta_b, local_q, S_v, local_sums);
717
718                    HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale));
719                    hvx_vec_store_u(attn_data + j, 8 * sizeof(float), res_attn);
720                }
721                for (; j + 4 <= S_v; j += 4) {
722                    float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v;
723                    float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v;
724                    float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v;
725                    float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v;
726                    gdn_mul_dot4_f32(row0, row1, row2, row3, local_gate, local_k, S_v, local_sums);
727
728                    float local_delta_b[32] __attribute__((aligned(128)));
729                    HVX_Vector vv_t = hvx_vmemu(v_t + j);
730                    HVX_Vector v_local_sums = hvx_vmem(local_sums);
731                    HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums);
732                    hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val));
733
734                    gdn_add_scaled_dot4_f32(row0, row1, row2, row3, local_k, local_delta_b, local_q, S_v, local_sums);
735
736                    HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale));
737                    hvx_vec_store_u(attn_data + j, 4 * sizeof(float), res_attn);
738                }
739                HVX_Vector vscale_splat = hvx_vec_splat_f32(scale);
740                for (; j < S_v; ++j) {
741                    float * row = s_work_curr + (uint64_t) j * S_v;
742                    HVX_Vector vsum = gdn_mul_dot_f32(row, local_gate, local_k, S_v);
743                    HVX_Vector vv_t = hvx_vec_splat_f32(v_t[j]);
744                    HVX_Vector vdj = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv_t, vsum), hvx_vec_splat_f32(beta_val));
745                    HVX_Vector vres = gdn_add_scaled_dot_f32(row, local_k, vdj, local_q, S_v);
746                    attn_data[j] = hvx_vec_get_f32(hvx_vec_mul_f32_f32(vres, vscale_splat));
747                }
748            } else {
749                const float gate = expf(g_t[0]);
750                uint32_t j = 0;
751                for (; j + 8 <= S_v; j += 8) {
752                    float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v;
753                    float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v;
754                    float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v;
755                    float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v;
756                    float * row4 = s_work_curr + (uint64_t) (j + 4) * S_v;
757                    float * row5 = s_work_curr + (uint64_t) (j + 5) * S_v;
758                    float * row6 = s_work_curr + (uint64_t) (j + 6) * S_v;
759                    float * row7 = s_work_curr + (uint64_t) (j + 7) * S_v;
760                    gdn_mul_scalar_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7,
761                                            gate, local_k, S_v, local_sums);
762
763                    float local_delta_b[32] __attribute__((aligned(128)));
764                    HVX_Vector vv_t = hvx_vmemu(v_t + j);
765                    HVX_Vector v_local_sums = hvx_vmem(local_sums);
766                    HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums);
767                    hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val));
768
769                    gdn_add_scaled_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7,
770                                            local_k, local_delta_b, local_q, S_v, local_sums);
771
772                    HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale));
773                    hvx_vec_store_u(attn_data + j, 8 * sizeof(float), res_attn);
774                }
775                for (; j + 4 <= S_v; j += 4) {
776                    float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v;
777                    float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v;
778                    float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v;
779                    float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v;
780                    gdn_mul_scalar_dot4_f32(row0, row1, row2, row3, gate, local_k, S_v, local_sums);
781
782                    float local_delta_b[32] __attribute__((aligned(128)));
783                    HVX_Vector vv_t = hvx_vmemu(v_t + j);
784                    HVX_Vector v_local_sums = hvx_vmem(local_sums);
785                    HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums);
786                    hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val));
787
788                    gdn_add_scaled_dot4_f32(row0, row1, row2, row3, local_k, local_delta_b, local_q, S_v, local_sums);
789
790                    HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale));
791                    hvx_vec_store_u(attn_data + j, 4 * sizeof(float), res_attn);
792                }
793                HVX_Vector vscale_splat = hvx_vec_splat_f32(scale);
794                for (; j < S_v; ++j) {
795                    float * row = s_work_curr + (uint64_t) j * S_v;
796                    HVX_Vector vsum = gdn_mul_scalar_dot_f32(row, gate, local_k, S_v);
797                    HVX_Vector vv_t = hvx_vec_splat_f32(v_t[j]);
798                    HVX_Vector vdj = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv_t, vsum), hvx_vec_splat_f32(beta_val));
799                    HVX_Vector vres = gdn_add_scaled_dot_f32(row, local_k, vdj, local_q, S_v);
800                    attn_data[j] = hvx_vec_get_f32(hvx_vec_mul_f32_f32(vres, vscale_splat));
801                }
802            }
803
804            if (K > 1) {
805                // snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back.
806                const int64_t target_slot = (int64_t) n_tokens - 1 - (int64_t) t;
807                if (target_slot >= 0 && target_slot < (int64_t) K) {
808                    float * curr_state_o = state_out_base + (uint64_t) target_slot * state_size_per_snap + ((uint64_t) iv3 * H + iv1) * S_v * S_v;
809                    if (curr_state_o != s_out) {
810                        hvx_copy_f32_uu((uint8_t *) curr_state_o, (const uint8_t *) s_work_curr, S_v * S_v);
811                    }
812                }
813            }
814
815            attn_data += (uint64_t) S_v * H;
816        }
817
818        // Push real write-back
819        dma_queue_push(dma, dma_make_ptr(s_out, s_work_curr),
820                       S_v * sizeof(float), S_v * sizeof(float),
821                       S_v * sizeof(float), S_v);
822
823        // Prefetch next block (if any)
824        if (ir_prefetch < row_end) {
825            const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H);
826            const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H);
827            const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v;
828
829            dma_queue_push(dma, dma_make_ptr(s_work[spad_idx], ps_in),
830                           S_v * sizeof(float), S_v * sizeof(float),
831                           S_v * sizeof(float), S_v);
832
833            ir_prefetch += nth;
834            spad_idx ^= 1;
835        }
836
837        curr_spad_idx ^= 1;
838    }
839    dma_queue_flush(dma);
840    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row_end);
841}
842
843
844static void gated_delta_net_f32_tg_thread(unsigned int nth, unsigned int ith, void * data) {
845    struct htp_gdn_context * gctx = (struct htp_gdn_context *) data;
846    struct htp_ops_context * octx = gctx->octx;
847
848    const struct htp_tensor * q     = octx->src[0];
849    const struct htp_tensor * k     = octx->src[1];
850    const struct htp_tensor * v     = octx->src[2];
851    const struct htp_tensor * g     = octx->src[3];
852    const struct htp_tensor * beta  = octx->src[4];
853    const struct htp_tensor * state = octx->src[5];
854    const struct htp_tensor * dst   = octx->dst;
855
856    const uint32_t S_v      = v->ne[0];
857    const uint32_t H        = v->ne[1];
858    const uint32_t n_seqs   = v->ne[3];
859
860    const uint32_t row_end = gctx->row_start + gctx->nrows;
861
862    if (ith >= gctx->nrows) {
863        return;
864    }
865
866    const uint32_t rq3 = n_seqs / q->ne[3];
867    const uint32_t rk3 = n_seqs / k->ne[3];
868    const float scale = 1.0f / sqrtf((float) S_v);
869
870    float * dst_base       = (float *) (uintptr_t) dst->data;
871    float * state_out_base = dst_base + (uint64_t) S_v * H * n_seqs;
872    const float * state_in_base = (const float *) (uintptr_t) state->data;
873
874    const bool kda = (g->ne[0] == S_v);
875    float local_gate[HTP_GDN_MAX_SV] __attribute__((aligned(128)));
876    float local_q[HTP_GDN_MAX_SV] __attribute__((aligned(128)));
877    float local_k[HTP_GDN_MAX_SV] __attribute__((aligned(128)));
878    float local_sums[32] __attribute__((aligned(128)));
879
880    dma_queue * dma = octx->ctx->dma[ith];
881    size_t state_aligned = (size_t) S_v * S_v * sizeof(float);
882    state_aligned = (state_aligned + 127) & ~(size_t)127;
883    float * s_work[2];
884    s_work[0] = (float *) (gctx->vtcm_base + gctx->vtcm_per_thread * ith);
885    s_work[1] = s_work[0] + state_aligned / sizeof(float);
886
887    struct fastdiv_values fd_H = init_fastdiv_values(H);
888    struct fastdiv_values fd_q1 = init_fastdiv_values(q->ne[1]);
889    struct fastdiv_values fd_k1 = init_fastdiv_values(k->ne[1]);
890    struct fastdiv_values fd_rq3 = init_fastdiv_values(rq3);
891    struct fastdiv_values fd_rk3 = init_fastdiv_values(rk3);
892
893    const uint64_t state_seq_stride = state->nb[3] / sizeof(float);
894
895    uint32_t ir_prefetch = gctx->row_start + ith;
896    int spad_idx = 0;
897
898    // Prefetch preamble (up to 2 steps)
899    for (int k = 0; k < 2 && ir_prefetch < row_end; k++) {
900        const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H);
901        const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H);
902        const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v;
903        // final state lands in snapshot slot 0 (most-recent-first ordering)
904        float * ps_out = state_out_base + ((uint64_t) piv3 * H + piv1) * S_v * S_v;
905
906        // Push dummy write-back
907        dma_queue_push(dma, dma_make_ptr(ps_out, s_work[spad_idx]),
908                       S_v * sizeof(float), S_v * sizeof(float),
909                       S_v * sizeof(float), 0);
910
911        // Push fetch
912        dma_queue_push(dma, dma_make_ptr(s_work[spad_idx], ps_in),
913                       S_v * sizeof(float), S_v * sizeof(float),
914                       S_v * sizeof(float), S_v);
915
916        ir_prefetch += nth;
917        spad_idx ^= 1;
918    }
919
920    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
921    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (gctx->row_start + ith));
922
923    int curr_spad_idx = 0;
924    for (uint32_t ir = gctx->row_start + ith; ir < row_end; ir += nth) {
925        dma_queue_pop(dma);
926        dma_queue_pop(dma);
927
928        float * s_work_curr = s_work[curr_spad_idx];
929
930        const uint32_t iv1 = fastmodulo(ir, H, &fd_H);
931        const uint32_t iv3 = fastdiv(ir, &fd_H);
932
933        const uint32_t iq1 = fastmodulo(iv1, q->ne[1], &fd_q1);
934        const uint32_t ik1 = fastmodulo(iv1, k->ne[1], &fd_k1);
935        const uint32_t iq3 = fastdiv(iv3, &fd_rq3);
936        const uint32_t ik3 = fastdiv(iv3, &fd_rk3);
937
938        // final state lands in snapshot slot 0 (most-recent-first ordering)
939        float * s_out = state_out_base + ((uint64_t) iv3 * H + iv1) * S_v * S_v;
940
941        float * attn_data = dst_base + ((uint64_t) iv3 * H + iv1) * S_v;
942
943        const float * q_t = (const float *) ((const uint8_t *) (uintptr_t) q->data +
944                (uint64_t) iq3 * q->nb[3] + (uint64_t) iq1 * q->nb[1]);
945        const float * k_t = (const float *) ((const uint8_t *) (uintptr_t) k->data +
946                (uint64_t) ik3 * k->nb[3] + (uint64_t) ik1 * k->nb[1]);
947        const float * v_t = (const float *) ((const uint8_t *) (uintptr_t) v->data +
948                (uint64_t) iv3 * v->nb[3] + (uint64_t) iv1 * v->nb[1]);
949        const float * g_t = (const float *) ((const uint8_t *) (uintptr_t) g->data +
950                (uint64_t) iv3 * g->nb[3] + (uint64_t) iv1 * g->nb[1]);
951        const float beta_val = *(const float *) ((const uint8_t *) (uintptr_t) beta->data +
952                (uint64_t) iv3 * beta->nb[3] + (uint64_t) iv1 * beta->nb[1]);
953
954        hvx_copy_f32_au((uint8_t *) local_q, (const uint8_t *) q_t, S_v);
955        hvx_copy_f32_au((uint8_t *) local_k, (const uint8_t *) k_t, S_v);
956
957        if (kda) {
958            hvx_exp_f32((uint8_t *) local_gate, (const uint8_t *) g_t, S_v, false);
959
960            uint32_t j = 0;
961            for (; j + 8 <= S_v; j += 8) {
962                float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v;
963                float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v;
964                float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v;
965                float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v;
966                float * row4 = s_work_curr + (uint64_t) (j + 4) * S_v;
967                float * row5 = s_work_curr + (uint64_t) (j + 5) * S_v;
968                float * row6 = s_work_curr + (uint64_t) (j + 6) * S_v;
969                float * row7 = s_work_curr + (uint64_t) (j + 7) * S_v;
970                gdn_mul_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7,
971                                 local_gate, local_k, S_v, local_sums);
972
973                float local_delta_b[32] __attribute__((aligned(128)));
974                HVX_Vector vv_t = hvx_vmemu(v_t + j);
975                HVX_Vector v_local_sums = hvx_vmem(local_sums);
976                HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums);
977                hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val));
978
979                gdn_add_scaled_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7,
980                                        local_k, local_delta_b, local_q, S_v, local_sums);
981
982                HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale));
983                hvx_vec_store_u(attn_data + j, 8 * sizeof(float), res_attn);
984            }
985            for (; j + 4 <= S_v; j += 4) {
986                float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v;
987                float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v;
988                float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v;
989                float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v;
990                gdn_mul_dot4_f32(row0, row1, row2, row3, local_gate, local_k, S_v, local_sums);
991
992                float local_delta_b[32] __attribute__((aligned(128)));
993                HVX_Vector vv_t = hvx_vmemu(v_t + j);
994                HVX_Vector v_local_sums = hvx_vmem(local_sums);
995                HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums);
996                hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val));
997
998                gdn_add_scaled_dot4_f32(row0, row1, row2, row3, local_k, local_delta_b, local_q, S_v, local_sums);
999
1000                HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale));
1001                hvx_vec_store_u(attn_data + j, 4 * sizeof(float), res_attn);
1002            }
1003            HVX_Vector vscale_splat = hvx_vec_splat_f32(scale);
1004            for (; j < S_v; ++j) {
1005                float * row = s_work_curr + (uint64_t) j * S_v;
1006                HVX_Vector vsum = gdn_mul_dot_f32(row, local_gate, local_k, S_v);
1007                HVX_Vector vv_t = hvx_vec_splat_f32(v_t[j]);
1008                HVX_Vector vdj = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv_t, vsum), hvx_vec_splat_f32(beta_val));
1009                HVX_Vector vres = gdn_add_scaled_dot_f32(row, local_k, vdj, local_q, S_v);
1010                attn_data[j] = hvx_vec_get_f32(hvx_vec_mul_f32_f32(vres, vscale_splat));
1011            }
1012        } else {
1013            const float gate = expf(g_t[0]);
1014            uint32_t j = 0;
1015            for (; j + 8 <= S_v; j += 8) {
1016                float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v;
1017                float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v;
1018                float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v;
1019                float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v;
1020                float * row4 = s_work_curr + (uint64_t) (j + 4) * S_v;
1021                float * row5 = s_work_curr + (uint64_t) (j + 5) * S_v;
1022                float * row6 = s_work_curr + (uint64_t) (j + 6) * S_v;
1023                float * row7 = s_work_curr + (uint64_t) (j + 7) * S_v;
1024                gdn_mul_scalar_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7,
1025                                        gate, local_k, S_v, local_sums);
1026
1027                float local_delta_b[32] __attribute__((aligned(128)));
1028                HVX_Vector vv_t = hvx_vmemu(v_t + j);
1029                HVX_Vector v_local_sums = hvx_vmem(local_sums);
1030                HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums);
1031                hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val));
1032
1033                gdn_add_scaled_dot8_f32(row0, row1, row2, row3, row4, row5, row6, row7,
1034                                        local_k, local_delta_b, local_q, S_v, local_sums);
1035
1036                HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale));
1037                hvx_vec_store_u(attn_data + j, 8 * sizeof(float), res_attn);
1038            }
1039            for (; j + 4 <= S_v; j += 4) {
1040                float * row0 = s_work_curr + (uint64_t) (j + 0) * S_v;
1041                float * row1 = s_work_curr + (uint64_t) (j + 1) * S_v;
1042                float * row2 = s_work_curr + (uint64_t) (j + 2) * S_v;
1043                float * row3 = s_work_curr + (uint64_t) (j + 3) * S_v;
1044                gdn_mul_scalar_dot4_f32(row0, row1, row2, row3, gate, local_k, S_v, local_sums);
1045
1046                float local_delta_b[32] __attribute__((aligned(128)));
1047                HVX_Vector vv_t = hvx_vmemu(v_t + j);
1048                HVX_Vector v_local_sums = hvx_vmem(local_sums);
1049                HVX_Vector diff = hvx_vec_sub_f32_f32(vv_t, v_local_sums);
1050                hvx_vmem(local_delta_b) = hvx_vec_mul_f32_f32(diff, hvx_vec_splat_f32(beta_val));
1051
1052                gdn_add_scaled_dot4_f32(row0, row1, row2, row3, local_k, local_delta_b, local_q, S_v, local_sums);
1053
1054                HVX_Vector res_attn = hvx_vec_mul_f32_f32(hvx_vmem(local_sums), hvx_vec_splat_f32(scale));
1055                hvx_vec_store_u(attn_data + j, 4 * sizeof(float), res_attn);
1056            }
1057            HVX_Vector vscale_splat = hvx_vec_splat_f32(scale);
1058            for (; j < S_v; ++j) {
1059                float * row = s_work_curr + (uint64_t) j * S_v;
1060                HVX_Vector vsum = gdn_mul_scalar_dot_f32(row, gate, local_k, S_v);
1061                HVX_Vector vv_t = hvx_vec_splat_f32(v_t[j]);
1062                HVX_Vector vdj = hvx_vec_mul_f32_f32(hvx_vec_sub_f32_f32(vv_t, vsum), hvx_vec_splat_f32(beta_val));
1063                HVX_Vector vres = gdn_add_scaled_dot_f32(row, local_k, vdj, local_q, S_v);
1064                attn_data[j] = hvx_vec_get_f32(hvx_vec_mul_f32_f32(vres, vscale_splat));
1065            }
1066        }
1067
1068        // Push real write-back
1069        dma_queue_push(dma, dma_make_ptr(s_out, s_work_curr),
1070                       S_v * sizeof(float), S_v * sizeof(float),
1071                       S_v * sizeof(float), S_v);
1072
1073        // Prefetch next block (if any)
1074        if (ir_prefetch < row_end) {
1075            const uint32_t piv1 = fastmodulo(ir_prefetch, H, &fd_H);
1076            const uint32_t piv3 = fastdiv(ir_prefetch, &fd_H);
1077            const float * ps_in = state_in_base + (uint64_t) piv3 * state_seq_stride + (uint64_t) piv1 * S_v * S_v;
1078
1079            dma_queue_push(dma, dma_make_ptr(s_work[spad_idx], ps_in),
1080                           S_v * sizeof(float), S_v * sizeof(float),
1081                           S_v * sizeof(float), S_v);
1082
1083            ir_prefetch += nth;
1084            spad_idx ^= 1;
1085        }
1086
1087        curr_spad_idx ^= 1;
1088    }
1089    dma_queue_flush(dma);
1090    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row_end);
1091}
1092
1093
1094int op_gated_delta_net(struct htp_ops_context * octx) {
1095    const struct htp_tensor * q     = octx->src[0];
1096    const struct htp_tensor * k     = octx->src[1];
1097    const struct htp_tensor * v     = octx->src[2];
1098    const struct htp_tensor * g     = octx->src[3];
1099    const struct htp_tensor * beta  = octx->src[4];
1100    const struct htp_tensor * state = octx->src[5];
1101    const struct htp_tensor * dst   = octx->dst;
1102
1103    if (q->type != HTP_TYPE_F32 || k->type != HTP_TYPE_F32 || v->type != HTP_TYPE_F32 ||
1104        g->type != HTP_TYPE_F32 || beta->type != HTP_TYPE_F32 || state->type != HTP_TYPE_F32 ||
1105        dst->type != HTP_TYPE_F32) {
1106        return HTP_STATUS_NO_SUPPORT;
1107    }
1108
1109    const uint32_t S_v      = v->ne[0];
1110    const uint32_t H        = v->ne[1];
1111    const uint32_t n_tokens = v->ne[2];
1112    const uint32_t n_seqs   = v->ne[3];
1113    const uint32_t K        = octx->op_params[0];
1114
1115    if (S_v == 0 || S_v > HTP_GDN_MAX_SV || H == 0 || n_tokens == 0 || n_seqs == 0) {
1116        return HTP_STATUS_NO_SUPPORT;
1117    }
1118    if ((g->ne[0] != 1 && g->ne[0] != S_v) || beta->ne[0] != 1) {
1119        return HTP_STATUS_NO_SUPPORT;
1120    }
1121    if (q->ne[0] != S_v || k->ne[0] != S_v || q->ne[1] == 0 || k->ne[1] == 0 ||
1122        q->ne[2] != n_tokens || k->ne[2] != n_tokens || q->ne[3] == 0 || k->ne[3] == 0 ||
1123        (n_seqs % q->ne[3]) != 0 || (n_seqs % k->ne[3]) != 0) {
1124        return HTP_STATUS_NO_SUPPORT;
1125    }
1126    // state holds s0 only: [S_v, S_v, H, n_seqs]
1127    if (state->ne[0] != S_v || state->ne[1] != S_v || state->ne[2] != H || state->ne[3] != n_seqs) {
1128        return HTP_STATUS_NO_SUPPORT;
1129    }
1130    if (dst->ne[0] != S_v * H || dst->ne[1] != n_tokens * n_seqs + S_v * n_seqs * K) {
1131        return HTP_STATUS_NO_SUPPORT;
1132    }
1133
1134    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
1135        return HTP_STATUS_OK;
1136    }
1137
1138    const uint32_t total_rows = H * n_seqs;
1139
1140    uint32_t row_start = 0;
1141    uint32_t nrows     = total_rows;
1142
1143    if (octx->ctx->mdev.count > 1) {
1144        const uint32_t head_bytes = S_v * sizeof(float);
1145        const uint32_t rows_per_chunk = (head_bytes > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(head_bytes, HEX_L2_LINE_SIZE)) : 1;
1146        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0,
1147                                                                             octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
1148        row_start = range.start;
1149        nrows     = range.count;
1150    }
1151
1152    if (nrows == 0) {
1153        return HTP_STATUS_OK;
1154    }
1155
1156    const uint32_t n_threads = octx->n_threads;
1157
1158    struct htp_gdn_context gctx;
1159    gctx.octx = octx;
1160    gctx.row_start       = row_start;
1161    gctx.nrows           = nrows;
1162    gctx.rows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
1163    gctx.state_bytes = (size_t) S_v * S_v * sizeof(float);
1164
1165    size_t state_aligned = (size_t) S_v * S_v * sizeof(float);
1166    state_aligned = (state_aligned + 127) & ~(size_t)127;
1167
1168    assert(octx->ctx->vtcm_size >= 2 * state_aligned * n_threads);
1169
1170    gctx.vtcm_base = octx->ctx->vtcm_base;
1171    gctx.vtcm_per_thread = 2 * state_aligned;
1172
1173    FARF(HIGH, "gated-delta-net-f32: q(%ux%ux%ux%u) k(%ux%ux%ux%u) v(%ux%ux%ux%u) state(%ux%ux%ux%u) -> (%ux%ux%ux%u) : "
1174         "vtcm-size %zu n_threads %u\n",
1175         q->ne[0], q->ne[1], q->ne[2], q->ne[3],
1176         k->ne[0], k->ne[1], k->ne[2], k->ne[3],
1177         v->ne[0], v->ne[1], v->ne[2], v->ne[3],
1178         state->ne[0], state->ne[1], state->ne[2], state->ne[3],
1179         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
1180         gctx.vtcm_per_thread * octx->n_threads, octx->n_threads);
1181
1182    if (n_tokens == 1) {
1183        work_queue_run(octx->ctx->work_queue, gated_delta_net_f32_tg_thread, &gctx, n_threads);
1184    } else {
1185        work_queue_run(octx->ctx->work_queue, gated_delta_net_f32_pp_thread, &gctx, n_threads);
1186    }
1187
1188    return HTP_STATUS_OK;
1189}