← Back to C-Kernel-Engine Docs Doxygen Source Documentation
 
Loading...
Searching...
No Matches
rwkv7_decode.c
Go to the documentation of this file.
1/* RWKV-7 T=1 FP32 decode reference. See include/ckernel_rwkv7_decode.h. */
2#include "ckernel_rwkv7_decode.h"
3
4#include <math.h>
5#include <stddef.h>
6#include <stdint.h>
7
8/* Compare byte ranges without undefined relational comparisons of pointers. */
9static int overlaps(const void *a, size_t an, const void *b, size_t bn) {
10 uintptr_t x = (uintptr_t)a, y = (uintptr_t)b;
11 if (an > UINTPTR_MAX - x || bn > UINTPTR_MAX - y) return 1;
12 return x < y + bn && y < x + an;
13}
14
15int ck_rwkv7_token_shift_lerp(const float *x,
16 const float *shift_in,
17 const float *mu,
18 float *xm,
19 float *shift_out,
20 int dim) {
21 if (!x || !shift_in || !mu || !xm || !shift_out || dim <= 0) {
22 return -1;
23 }
24 size_t bytes = (size_t)dim * sizeof(float);
25 if (overlaps(xm, bytes, shift_out, bytes)) return -1;
26 const float *inputs[] = {x, shift_in, mu};
27 for (int j = 0; j < 3; ++j) {
28 if ((xm != inputs[j] && overlaps(xm, bytes, inputs[j], bytes)) ||
29 (shift_out != inputs[j] && overlaps(shift_out, bytes, inputs[j], bytes))) return -1;
30 }
31 for (int i = 0; i < dim; ++i) {
32 const float previous = shift_in[i];
33 const float current = x[i];
34 const float m = mu[i];
35 xm[i] = current + m * (previous - current);
36 shift_out[i] = current;
37 }
38 return 0;
39}
40
41int ck_rwkv7_time_mix_shifts(const float *x,
42 const float *shift_in,
43 const float *mu_r,
44 const float *mu_w,
45 const float *mu_k,
46 const float *mu_v,
47 const float *mu_a,
48 const float *mu_g,
49 float *xr,
50 float *xw,
51 float *xk,
52 float *xv,
53 float *xa,
54 float *xg,
55 float *shift_out,
56 int dim) {
57 if (!x || !shift_in || !mu_r || !mu_w || !mu_k || !mu_v || !mu_a ||
58 !mu_g || !xr || !xw || !xk || !xv || !xa || !xg || !shift_out ||
59 dim <= 0) {
60 return -1;
61 }
62 const size_t bytes = (size_t)dim * sizeof(float);
63 float *outs[] = {xr, xw, xk, xv, xa, xg, shift_out};
64 for (int a = 0; a < 7; ++a) {
65 for (int b = a + 1; b < 7; ++b) {
66 if (overlaps(outs[a], bytes, outs[b], bytes)) return -1;
67 }
68 }
69 const float *inputs[] = {x, shift_in, mu_r, mu_w, mu_k, mu_v, mu_a, mu_g};
70 for (int j = 0; j < 8; ++j) {
71 for (int o = 0; o < 7; ++o) {
72 if (outs[o] == inputs[j] &&
73 !(outs[o] == shift_out && inputs[j] == shift_in)) {
74 return -1;
75 }
76 if (outs[o] != inputs[j] &&
77 overlaps(outs[o], bytes, inputs[j], bytes)) {
78 return -1;
79 }
80 }
81 }
82 /* Fused per-index pass: each lane reads x[i]/shift_in[i] once, so the
83 * shared shift is written exactly once instead of six times. */
84 for (int i = 0; i < dim; ++i) {
85 const float current = x[i];
86 const float previous = shift_in[i];
87 const float d = previous - current;
88 xr[i] = current + mu_r[i] * d;
89 xw[i] = current + mu_w[i] * d;
90 xk[i] = current + mu_k[i] * d;
91 xv[i] = current + mu_v[i] * d;
92 xa[i] = current + mu_a[i] * d;
93 xg[i] = current + mu_g[i] * d;
94 shift_out[i] = current;
95 }
96 return 0;
97}
98
99int ck_rwkv7_norm_kk(const float *kk_in,
100 float *kk_out,
101 int num_heads,
102 int head_dim) {
103 if (!kk_in || !kk_out || num_heads <= 0 || head_dim <= 0 ||
104 num_heads > 256 || head_dim > 512) {
105 return -1;
106 }
107 for (int h = 0; h < num_heads; ++h) {
108 const float *src = kk_in + (size_t)h * (size_t)head_dim;
109 float *dst = kk_out + (size_t)h * (size_t)head_dim;
110 float norm = 0.0f;
111 for (int i = 0; i < head_dim; ++i) {
112 norm += src[i] * src[i];
113 }
114 norm = sqrtf(norm);
115 /* Divisor max(norm,1e-12); an all-zero row stays zero. */
116 if (norm < 1e-12f) {
117 norm = 1e-12f;
118 }
119 const float inv = 1.0f / norm;
120 for (int i = 0; i < head_dim; ++i) {
121 dst[i] = src[i] * inv;
122 }
123 }
124 return 0;
125}
126
127int ck_rwkv7_scale_k(const float *k_in,
128 const float *a,
129 const float *k_a,
130 float *k_out,
131 int dim) {
132 if (!k_in || !a || !k_a || !k_out || dim <= 0) {
133 return -1;
134 }
135 for (int i = 0; i < dim; ++i) {
136 k_out[i] = k_in[i] * (1.0f + (k_a[i] - 1.0f) * a[i]);
137 }
138 return 0;
139}
140
141int ck_rwkv7_apply_kk_factor(const float *k,
142 const float *k_k,
143 float *kk_raw,
144 int dim) {
145 if (!k || !k_k || !kk_raw || dim <= 0) {
146 return -1;
147 }
148 for (int i = 0; i < dim; ++i) {
149 kk_raw[i] = k[i] * k_k[i];
150 }
151 return 0;
152}
153
154static inline float ck_rwkv7_sigmoid(float x) {
155 return 1.0f / (1.0f + expf(-x));
156}
157
158int ck_rwkv7_lora_gates(const float *xw,
159 const float *xa,
160 const float *xg,
161 const float *w1,
162 const float *w2,
163 const float *w0,
164 const float *a1,
165 const float *a2,
166 const float *a0,
167 const float *g1,
168 const float *g2,
169 int dim,
170 int rank_w,
171 int rank_a,
172 int rank_g,
173 float *w_log,
174 float *a_out,
175 float *g_out) {
176 if (!xw || !xa || !xg || !w1 || !w2 || !w0 || !a1 || !a2 || !a0 || !g1 ||
177 !g2 || !w_log || !a_out || !g_out || dim <= 0 || rank_w <= 0 ||
178 rank_a <= 0 || rank_g <= 0 || rank_w > 512 || rank_a > 512 ||
179 rank_g > 512) {
180 return -1;
181 }
182 /* w: tanh between factors, sigmoid after (+w0), scaled by -e^-0.5. */
183 {
184 float tmp[512];
185 for (int r = 0; r < rank_w; ++r) {
186 float acc = 0.0f;
187 for (int i = 0; i < dim; ++i) {
188 acc += xw[i] * w1[(size_t)i * (size_t)rank_w + (size_t)r];
189 }
190 tmp[r] = tanhf(acc);
191 }
192 for (int i = 0; i < dim; ++i) {
193 float acc = w0[i];
194 for (int r = 0; r < rank_w; ++r) {
195 acc += tmp[r] * w2[(size_t)r * (size_t)dim + (size_t)i];
196 }
197 w_log[i] = -0.6065306597126334f * ck_rwkv7_sigmoid(acc);
198 }
199 }
200 /* a: no activation between, sigmoid after (+a0). */
201 {
202 float tmp[512];
203 for (int r = 0; r < rank_a; ++r) {
204 float acc = 0.0f;
205 for (int i = 0; i < dim; ++i) {
206 acc += xa[i] * a1[(size_t)i * (size_t)rank_a + (size_t)r];
207 }
208 tmp[r] = acc;
209 }
210 for (int i = 0; i < dim; ++i) {
211 float acc = a0[i];
212 for (int r = 0; r < rank_a; ++r) {
213 acc += tmp[r] * a2[(size_t)r * (size_t)dim + (size_t)i];
214 }
215 a_out[i] = ck_rwkv7_sigmoid(acc);
216 }
217 }
218 /* g: sigmoid between factors, nothing after, no bias. */
219 {
220 float tmp[512];
221 for (int r = 0; r < rank_g; ++r) {
222 float acc = 0.0f;
223 for (int i = 0; i < dim; ++i) {
224 acc += xg[i] * g1[(size_t)i * (size_t)rank_g + (size_t)r];
225 }
226 tmp[r] = ck_rwkv7_sigmoid(acc);
227 }
228 for (int i = 0; i < dim; ++i) {
229 float acc = 0.0f;
230 for (int r = 0; r < rank_g; ++r) {
231 acc += tmp[r] * g2[(size_t)r * (size_t)dim + (size_t)i];
232 }
233 g_out[i] = acc;
234 }
235 }
236 return 0;
237}
238
239int ck_rwkv7_vgate_vmix(const float *v_in,
240 const float *v_first,
241 const float *xv,
242 const float *v1,
243 const float *v2,
244 const float *v0,
245 int dim,
246 int rank_v,
247 float *v_out,
248 float *v_first_out,
249 float *v_gate) {
250 if (!v_in || !v_out || !v_first_out || dim <= 0) {
251 return -1;
252 }
253 if (xv == NULL) {
254 /* Layer 0: produce v_first, no gate read. */
255 for (int i = 0; i < dim; ++i) {
256 v_out[i] = v_in[i];
257 v_first_out[i] = v_in[i];
258 }
259 return 0;
260 }
261 if (!v_first || !v1 || !v2 || !v0 || rank_v <= 0 || rank_v > 512) {
262 return -1;
263 }
264 float tmp[512];
265 for (int r = 0; r < rank_v; ++r) {
266 float acc = 0.0f;
267 for (int i = 0; i < dim; ++i) {
268 acc += xv[i] * v1[(size_t)i * (size_t)rank_v + (size_t)r];
269 }
270 tmp[r] = acc;
271 }
272 for (int i = 0; i < dim; ++i) {
273 float acc = v0[i];
274 for (int r = 0; r < rank_v; ++r) {
275 acc += tmp[r] * v2[(size_t)r * (size_t)dim + (size_t)i];
276 }
277 const float gate = ck_rwkv7_sigmoid(acc);
278 if (v_gate) {
279 v_gate[i] = gate;
280 }
281 v_out[i] = v_in[i] + (v_first[i] - v_in[i]) * gate;
282 v_first_out[i] = v_first[i];
283 }
284 return 0;
285}
286
287int ck_rwkv7_state_step(const float *w_log,
288 const float *k,
289 const float *v,
290 const float *kk,
291 const float *a,
292 const float *state_in,
293 float *state_out,
294 int head_dim) {
295 if (!w_log || !k || !v || !kk || !a || !state_in || !state_out ||
296 head_dim <= 0 || head_dim > 512) {
297 return -1;
298 }
299 float decay[512];
300 float bvec[512];
301 for (int j = 0; j < head_dim; ++j) {
302 decay[j] = expf(w_log[j]);
303 bvec[j] = kk[j] * a[j];
304 }
305 for (int i = 0; i < head_dim; ++i) {
306 float t = 0.0f;
307 for (int l = 0; l < head_dim; ++l) {
308 t += state_in[(size_t)i * (size_t)head_dim + (size_t)l] * kk[l];
309 }
310 for (int j = 0; j < head_dim; ++j) {
311 state_out[(size_t)i * (size_t)head_dim + (size_t)j] =
312 state_in[(size_t)i * (size_t)head_dim + (size_t)j] * decay[j] -
313 t * bvec[j] + v[i] * k[j];
314 }
315 }
316 return 0;
317}
318
319int ck_rwkv7_state_readout(const float *state,
320 const float *r,
321 float *y,
322 int head_dim) {
323 if (!state || !r || !y || head_dim <= 0 || head_dim > 512) {
324 return -1;
325 }
326 for (int i = 0; i < head_dim; ++i) {
327 float acc = 0.0f;
328 for (int l = 0; l < head_dim; ++l) {
329 acc += state[(size_t)i * (size_t)head_dim + (size_t)l] * r[l];
330 }
331 y[i] = acc;
332 }
333 return 0;
334}
335
336int ck_rwkv7_wkv_decode(const float *r,
337 const float *w_log,
338 const float *k,
339 const float *v,
340 const float *kk,
341 const float *a,
342 const float *g,
343 const float *rk,
344 const float *ln_w,
345 const float *ln_b,
346 float norm_eps,
347 const float *state_in,
348 float *state_out,
349 float *y_out,
350 int num_heads,
351 int head_dim) {
352 if (!r || !w_log || !k || !v || !kk || !a || !g || !rk || !ln_w || !ln_b ||
353 !state_in || !state_out || !y_out || num_heads <= 0 || head_dim <= 0 ||
354 !(norm_eps > 0.0f) || !isfinite(norm_eps) ||
355 !isfinite(norm_eps * (float)head_dim)) {
356 return -1;
357 }
358 if (num_heads > 256 || head_dim > 512) {
359 return -1;
360 }
361
362 /* Per-head generalised delta-rule step + raw readout (pre-norm). */
363 for (int h = 0; h < num_heads; ++h) {
364 const float *rh = r + (size_t)h * (size_t)head_dim;
365 const float *kh = k + (size_t)h * (size_t)head_dim;
366 const float *vh = v + (size_t)h * (size_t)head_dim;
367 const float *wh = w_log + (size_t)h * (size_t)head_dim;
368 const float *khh = kk + (size_t)h * (size_t)head_dim;
369 const float *ah = a + (size_t)h * (size_t)head_dim;
370 const float *sin = state_in + (size_t)h * (size_t)head_dim * (size_t)head_dim;
371 float *sout = state_out + (size_t)h * (size_t)head_dim * (size_t)head_dim;
372 float *yh = y_out + (size_t)h * (size_t)head_dim;
373
374 if (ck_rwkv7_state_step(wh, kh, vh, khh, ah, sin, sout, head_dim) != 0 ||
375 ck_rwkv7_state_readout(sout, rh, yh, head_dim) != 0) {
376 return -1;
377 }
378 }
379
380 /* GroupNorm over C with H groups, eps = norm_eps*head_dim. */
381 const float eps = norm_eps * (float)head_dim;
382 for (int h = 0; h < num_heads; ++h) {
383 float *yh = y_out + (size_t)h * (size_t)head_dim;
384 const float *lwh = ln_w + (size_t)h * (size_t)head_dim;
385 const float *lbh = ln_b + (size_t)h * (size_t)head_dim;
386 float mean = 0.0f;
387 for (int i = 0; i < head_dim; ++i) {
388 mean += yh[i];
389 }
390 mean /= (float)head_dim;
391 float var = 0.0f;
392 for (int i = 0; i < head_dim; ++i) {
393 const float d = yh[i] - mean;
394 var += d * d;
395 }
396 var /= (float)head_dim;
397 const float inv = 1.0f / sqrtf(var + eps);
398 for (int i = 0; i < head_dim; ++i) {
399 yh[i] = (yh[i] - mean) * inv * lwh[i] + lbh[i];
400 }
401 }
402
403 /* Per-head bonus s_h = sum(r_h*k_h*rk_h); y += s_h*v; then y *= g. */
404 const int dim = num_heads * head_dim;
405 for (int h = 0; h < num_heads; ++h) {
406 const float *rh = r + (size_t)h * (size_t)head_dim;
407 const float *kh = k + (size_t)h * (size_t)head_dim;
408 const float *vh = v + (size_t)h * (size_t)head_dim;
409 const float *rkh = rk + (size_t)h * (size_t)head_dim;
410 float *yh = y_out + (size_t)h * (size_t)head_dim;
411 float s = 0.0f;
412 for (int i = 0; i < head_dim; ++i) {
413 s += rh[i] * kh[i] * rkh[i];
414 }
415 for (int i = 0; i < head_dim; ++i) {
416 yh[i] += s * vh[i];
417 }
418 }
419 for (int i = 0; i < dim; ++i) {
420 y_out[i] *= g[i];
421 }
422 return 0;
423}
424
425size_t ck_rwkv7_channelmix_scratch_elems(int dim, int hidden_dim) {
426 if (dim <= 0 || hidden_dim <= 0 || hidden_dim > 8192) {
427 return 0;
428 }
429 return (size_t)dim + (size_t)hidden_dim;
430}
431
432int ck_rwkv7_channelmix_decode(const float *x,
433 const float *shift_in,
434 const float *mu,
435 const float *Wk,
436 const float *Wv,
437 float *shift_out,
438 float *out,
439 int dim,
440 int hidden_dim,
441 float *scratch,
442 size_t scratch_elems) {
443 if (!x || !shift_in || !mu || !Wk || !Wv || !shift_out || !out || dim <= 0 ||
444 hidden_dim <= 0 || hidden_dim > 8192) {
445 return -1;
446 }
447 if (shift_out == shift_in) {
448 return -1;
449 }
450 if (!scratch ||
451 scratch_elems < (size_t)dim + (size_t)hidden_dim ||
452 scratch == (const float *)x || scratch == (const float *)shift_in ||
453 scratch == (const float *)mu || scratch == (const float *)Wk ||
454 scratch == (const float *)Wv || scratch == (const float *)shift_out ||
455 scratch == (const float *)out) {
456 return -1;
457 }
458 const size_t bytes = (size_t)dim * sizeof(float);
459 if ((size_t)hidden_dim > SIZE_MAX / bytes) return -1;
460 const size_t weight_bytes = bytes * (size_t)hidden_dim;
461 const size_t scratch_bytes = ((size_t)dim + (size_t)hidden_dim) * sizeof(float);
462 if (overlaps(shift_out, bytes, out, bytes)) return -1;
463 const float *inputs[] = {x, shift_in, mu, Wk, Wv};
464 const size_t sizes[] = {bytes, bytes, bytes, weight_bytes, weight_bytes};
465 for (int j = 0; j < 5; ++j) {
466 if (overlaps(scratch, scratch_bytes, inputs[j], sizes[j])) return -1;
467 if (!(j == 0 && shift_out == x) &&
468 overlaps(shift_out, bytes, inputs[j], sizes[j])) return -1;
469 if (!(j == 0 && out == x) &&
470 overlaps(out, bytes, inputs[j], sizes[j])) return -1;
471 }
472 if (overlaps(scratch, scratch_bytes, out, bytes) ||
473 overlaps(scratch, scratch_bytes, shift_out, bytes)) return -1;
474 /* Hoisted token-shift mix; per-index read-before-write keeps
475 * shift_out == x aliasing safe. */
476 float *xk = scratch;
477 float *hidden = scratch + dim;
478 for (int i = 0; i < dim; ++i) {
479 xk[i] = x[i] + mu[i] * (shift_in[i] - x[i]);
480 shift_out[i] = x[i];
481 }
482 /* hidden[j] = relu(Wk[j,:] @ xk)^2, Wk [K,C] row-major. */
483 for (int j = 0; j < hidden_dim; ++j) {
484 const float *row = Wk + (size_t)j * (size_t)dim;
485 float acc = 0.0f;
486 for (int i = 0; i < dim; ++i) {
487 acc += row[i] * xk[i];
488 }
489 const float rel = acc > 0.0f ? acc : 0.0f;
490 hidden[j] = rel * rel;
491 }
492 /* out = Wv @ hidden, Wv [C,K] row-major; runs after the last xk read,
493 * so out == x aliasing is safe. */
494 for (int i = 0; i < dim; ++i) {
495 const float *row = Wv + (size_t)i * (size_t)hidden_dim;
496 float acc = 0.0f;
497 for (int j = 0; j < hidden_dim; ++j) {
498 acc += row[j] * hidden[j];
499 }
500 out[i] = acc;
501 }
502 return 0;
503}
504
505size_t ck_rwkv7_time_mix_scratch_elems(int dim) {
506 if (dim <= 0 || (size_t)dim > SIZE_MAX / (7u * sizeof(float))) {
507 return 0;
508 }
509 return (size_t)dim * 7u;
510}int ck_rwkv7_time_mix_tail(const float *r,
511 const float *k_raw,
512 const float *v_in,
513 const float *v_first_in,
514 const float *xw,
515 const float *xa,
516 const float *xg,
517 const float *xv,
518 const float *w1,
519 const float *w2,
520 const float *w0,
521 const float *a1,
522 const float *a2,
523 const float *a0,
524 const float *g1,
525 const float *g2,
526 const float *v1,
527 const float *v2,
528 const float *v0,
529 const float *k_k,
530 const float *k_a,
531 const float *rk,
532 const float *ln_w,
533 const float *ln_b,
534 float norm_eps,
535 const float *state_in,
536 float *state_out,
537 float *y_preproj,
538 float *v_first_out,
539 int dim,
540 int num_heads,
541 int head_dim,
542 int rank_w,
543 int rank_a,
544 int rank_g,
545 int rank_v,
546 float *scratch,
547 size_t scratch_elems) {
548 const int layer0 = (xv == NULL);
549 if (!r || !k_raw || !v_in || !xw || !xa || !xg || !w1 || !w2 || !w0 ||
550 !a1 || !a2 || !a0 || !g1 || !g2 || !k_k || !k_a || !rk || !ln_w ||
551 !ln_b || !state_in || !state_out || !y_preproj || !v_first_out ||
552 !scratch || dim <= 0 || num_heads <= 0 || head_dim <= 0 ||
553 num_heads > 256 || head_dim > 512 || rank_w <= 0 || rank_w > 512 ||
554 rank_a <= 0 || rank_a > 512 || rank_g <= 0 || rank_g > 512 ||
555 !(norm_eps > 0.0f) || !isfinite(norm_eps) ||
556 !isfinite(norm_eps * (float)head_dim)) {
557 return -1;
558 }
559 if ((size_t)dim != (size_t)num_heads * (size_t)head_dim) {
560 return -1;
561 }
562 if (layer0) {
563 /* Layer 0 produces v_first: no residual v state exists yet, so the
564 * residual inputs must be NULL (produce-only path never reads them). */
565 if (v_first_in || v1 || v2 || v0) {
566 return -1;
567 }
568 } else {
569 if (!v_first_in || !v1 || !v2 || !v0 || rank_v <= 0 || rank_v > 512) {
570 return -1;
571 }
572 }
573 if (scratch_elems < (size_t)dim * 7u) {
574 return -1;
575 }
576 const size_t vec_bytes = (size_t)dim * sizeof(float);
577 if (vec_bytes / sizeof(float) != (size_t)dim) return -1;
578 if ((size_t)head_dim > SIZE_MAX / vec_bytes) return -1;
579 const size_t state_bytes = vec_bytes * (size_t)head_dim;
580 const size_t scratch_bytes = vec_bytes * 7u;
581 /* Scratch must be disjoint from every other argument. */
582 const void *ptrs[] = {r, k_raw, v_in, v_first_in, xw, xa, xg, xv,
583 w1, w2, w0, a1, a2, a0, g1, g2, v1, v2, v0,
584 k_k, k_a, rk, ln_w, ln_b, state_in, state_out,
585 y_preproj, v_first_out};
586 /* LoRA factor sizes: down [C,R], up [R,C], bias [C]. State is [H,N,N]. */
587 size_t lens[] = {(size_t)dim, (size_t)dim, (size_t)dim, (size_t)dim,
588 (size_t)dim, (size_t)dim, (size_t)dim, (size_t)dim,
589 (size_t)dim * (size_t)rank_w, (size_t)rank_w * (size_t)dim, (size_t)dim,
590 (size_t)dim * (size_t)rank_a, (size_t)rank_a * (size_t)dim, (size_t)dim,
591 (size_t)dim * (size_t)rank_g, (size_t)rank_g * (size_t)dim,
592 0, 0, 0,
593 (size_t)dim, (size_t)dim,
594 (size_t)num_heads * (size_t)head_dim, (size_t)dim, (size_t)dim,
595 state_bytes / sizeof(float), state_bytes / sizeof(float),
596 (size_t)dim, (size_t)dim};
597 if (!layer0) {
598 lens[16] = (size_t)dim * (size_t)rank_v;
599 lens[17] = (size_t)rank_v * (size_t)dim;
600 lens[18] = (size_t)dim;
601 }
602 for (int j = 0; j < 28; ++j) {
603 if (!ptrs[j] || lens[j] == 0) continue;
604 if (lens[j] > SIZE_MAX / sizeof(float)) return -1;
605 if (overlaps(scratch, scratch_bytes, ptrs[j], lens[j] * sizeof(float))) return -1;
606 }
607 /* Outputs disjoint from each other; state_out may == state_in,
608 * v_first_out may == v_first_in. All other output-input overlaps fail. */
609 if (overlaps(state_out, state_bytes, y_preproj, vec_bytes) ||
610 overlaps(state_out, state_bytes, v_first_out, vec_bytes) ||
611 overlaps(y_preproj, vec_bytes, v_first_out, vec_bytes)) {
612 return -1;
613 }
614 const void *ro_inputs[] = {r, k_raw, v_in, xw, xa, xg, xv, w1, w2, w0,
615 a1, a2, a0, g1, g2, v1, v2, v0, k_k, k_a,
616 rk, ln_w, ln_b};
617 size_t ro_lens[] = {(size_t)dim, (size_t)dim, (size_t)dim, (size_t)dim,
618 (size_t)dim, (size_t)dim, (size_t)dim,
619 (size_t)dim * (size_t)rank_w, (size_t)rank_w * (size_t)dim, (size_t)dim,
620 (size_t)dim * (size_t)rank_a, (size_t)rank_a * (size_t)dim, (size_t)dim,
621 (size_t)dim * (size_t)rank_g, (size_t)rank_g * (size_t)dim,
622 0, 0, 0,
623 (size_t)dim, (size_t)dim, (size_t)dim, (size_t)dim, (size_t)dim};
624 if (!layer0) {
625 ro_lens[15] = (size_t)dim * (size_t)rank_v;
626 ro_lens[16] = (size_t)rank_v * (size_t)dim;
627 ro_lens[17] = (size_t)dim;
628 }
629 for (int j = 0; j < 23; ++j) {
630 if (!ro_inputs[j] || ro_lens[j] == 0) continue;
631 size_t n = ro_lens[j];
632 if (overlaps(y_preproj, vec_bytes, ro_inputs[j], n * sizeof(float))) return -1;
633 if (overlaps(v_first_out, vec_bytes, ro_inputs[j], n * sizeof(float))) return -1;
634 /* state_out has no self-match in ro_inputs (state_in is not listed),
635 * so the check always applies; exact state_out==state_in aliasing is
636 * handled by the dedicated state_in comparison below. */
637 if (overlaps(state_out, state_bytes, ro_inputs[j], n * sizeof(float))) return -1;
638 }
639 /* y/v_first outputs must also be disjoint from the state and v_first
640 * input streams (ro_inputs above covers r..ln_b only). y has no allowed
641 * alias; v_first_out may exactly alias v_first_in (checked below). */
642 if (overlaps(y_preproj, vec_bytes, state_in, state_bytes)) return -1;
643 if (v_first_in && overlaps(y_preproj, vec_bytes, v_first_in, vec_bytes)) {
644 return -1;
645 }
646 if (overlaps(v_first_out, vec_bytes, state_in, state_bytes)) return -1;
647 if (v_first_in && overlaps(state_out, state_bytes, v_first_in, vec_bytes)) {
648 return -1;
649 }
650 if (v_first_in && v_first_out != v_first_in &&
651 overlaps(v_first_out, vec_bytes, v_first_in, vec_bytes)) {
652 return -1;
653 }
654 if (state_out != state_in &&
655 overlaps(state_out, state_bytes, state_in, state_bytes)) {
656 return -1;
657 }
658
659 float *w_log = scratch;
660 float *a_buf = scratch + (size_t)dim;
661 float *g_buf = scratch + (size_t)dim * 2u;
662 float *v_buf = scratch + (size_t)dim * 3u;
663 float *kk_raw = scratch + (size_t)dim * 4u;
664 float *kk = scratch + (size_t)dim * 5u;
665 float *k_sc = scratch + (size_t)dim * 6u;
666
667 if (ck_rwkv7_lora_gates(xw, xa, xg, w1, w2, w0, a1, a2, a0, g1, g2,
668 dim, rank_w, rank_a, rank_g,
669 w_log, a_buf, g_buf) != 0) {
670 return -1;
671 }
672 if (layer0) {
673 if (ck_rwkv7_vgate_vmix(v_in, NULL, NULL, NULL, NULL, NULL,
674 dim, 1, v_buf, v_first_out, NULL) != 0) {
675 return -1;
676 }
677 } else {
678 if (ck_rwkv7_vgate_vmix(v_in, v_first_in, xv, v1, v2, v0,
679 dim, rank_v, v_buf, v_first_out, NULL) != 0) {
680 return -1;
681 }
682 }
683 if (ck_rwkv7_apply_kk_factor(k_raw, k_k, kk_raw, dim) != 0) {
684 return -1;
685 }
686 if (ck_rwkv7_norm_kk(kk_raw, kk, num_heads, head_dim) != 0) {
687 return -1;
688 }
689 if (ck_rwkv7_scale_k(k_raw, a_buf, k_a, k_sc, dim) != 0) {
690 return -1;
691 }
692 if (ck_rwkv7_wkv_decode(r, w_log, k_sc, v_buf, kk, a_buf, g_buf, rk,
693 ln_w, ln_b, norm_eps, state_in, state_out,
694 y_preproj, num_heads, head_dim) != 0) {
695 return -1;
696 }
697 return 0;
698}
699
700size_t ck_rwkv7_time_mix_block_scratch_elems(int dim) {
701 if (dim <= 0 || (size_t)dim > SIZE_MAX / (19u * sizeof(float))) {
702 return 0;
703 }
704 return (size_t)dim * 19u;
705}
706
707int ck_rwkv7_time_mix_block(const float *h,
708 const float *att_shift_in,
709 const float *v_first_in,
710 const float *state_in,
711 const float *ln_g,
712 const float *ln_b,
713 const float *mu_r,
714 const float *mu_w,
715 const float *mu_k,
716 const float *mu_v,
717 const float *mu_a,
718 const float *mu_g,
719 const float *Wr,
720 const float *Wk,
721 const float *Wv,
722 const float *Wo,
723 const float *w1,
724 const float *w2,
725 const float *w0,
726 const float *a1,
727 const float *a2,
728 const float *a0,
729 const float *g1,
730 const float *g2,
731 const float *v1,
732 const float *v2,
733 const float *v0,
734 const float *k_k,
735 const float *k_a,
736 const float *rk,
737 const float *gn_w,
738 const float *gn_b,
739 float norm_eps,
740 int dim,
741 int num_heads,
742 int head_dim,
743 int rank_w,
744 int rank_a,
745 int rank_g,
746 int rank_v,
747 int is_layer_zero,
748 ck_rwkv7_gemv_fn gemv,
749 ck_rwkv7_layer_norm_fn layer_norm,
750 void *ctx,
751 float *h_out,
752 float *att_shift_out,
753 float *state_out,
754 float *v_first_out,
755 float *scratch,
756 size_t scratch_elems) {
757 const int layer0 = is_layer_zero ? 1 : 0;
758 if (!h || !att_shift_in || !ln_g || !ln_b || !mu_r || !mu_w || !mu_k ||
759 !mu_v || !mu_a || !mu_g || !Wr || !Wk || !Wv || !Wo || !w1 || !w2 ||
760 !w0 || !a1 || !a2 || !a0 || !g1 || !g2 || !k_k || !k_a || !rk ||
761 !gn_w || !gn_b || !state_in || !state_out || !h_out ||
762 !att_shift_out || !v_first_out || !gemv || !layer_norm || !scratch ||
763 dim <= 0 || num_heads <= 0 || head_dim <= 0 || num_heads > 256 ||
764 head_dim > 512 || rank_w <= 0 || rank_w > 512 || rank_a <= 0 ||
765 rank_a > 512 || rank_g <= 0 || rank_g > 512 ||
766 !(norm_eps > 0.0f) || !isfinite(norm_eps) ||
767 !isfinite(norm_eps * (float)head_dim) ||
768 (is_layer_zero != 0 && is_layer_zero != 1)) {
769 return -1;
770 }
771 if ((size_t)dim != (size_t)num_heads * (size_t)head_dim) {
772 return -1;
773 }
774 if (layer0) {
775 if (v_first_in || v1 || v2 || v0) return -1;
776 } else if (!v_first_in || !v1 || !v2 || !v0 || rank_v <= 0 ||
777 rank_v > 512) {
778 return -1;
779 }
780 if (scratch_elems < (size_t)dim * 19u) {
781 return -1;
782 }
783 const size_t vec_bytes = (size_t)dim * sizeof(float);
784 if ((size_t)head_dim > SIZE_MAX / vec_bytes) return -1;
785 const size_t state_bytes = vec_bytes * (size_t)head_dim;
786 const size_t scratch_bytes = vec_bytes * 19u;
787
788 /* Output aliasing: each output may exactly alias its own input stream;
789 * the four outputs must otherwise be pairwise disjoint and disjoint
790 * from every input. */
791 float *outs[] = {h_out, att_shift_out, state_out, v_first_out};
792 const void *allow[] = {h, att_shift_in, state_in, v_first_in};
793 const size_t out_bytes[] = {vec_bytes, vec_bytes, state_bytes, vec_bytes};
794 for (int a = 0; a < 4; ++a) {
795 for (int b = a + 1; b < 4; ++b) {
796 if (overlaps(outs[a], out_bytes[a], outs[b], out_bytes[b])) return -1;
797 }
798 }
799 const void *ro[] = {h, att_shift_in, v_first_in, state_in, ln_g, ln_b,
800 mu_r, mu_w, mu_k, mu_v, mu_a, mu_g,
801 Wr, Wk, Wv, Wo,
802 w1, w2, w0, a1, a2, a0, g1, g2, v1, v2, v0,
803 k_k, k_a, rk, gn_w, gn_b};
804 size_t ro_e[] = {(size_t)dim, (size_t)dim, (size_t)dim, 0,
805 (size_t)dim, (size_t)dim,
806 (size_t)dim, (size_t)dim, (size_t)dim, (size_t)dim,
807 (size_t)dim, (size_t)dim,
808 (size_t)dim * (size_t)dim, (size_t)dim * (size_t)dim,
809 (size_t)dim * (size_t)dim, (size_t)dim * (size_t)dim,
810 (size_t)dim * (size_t)rank_w, (size_t)rank_w * (size_t)dim, (size_t)dim,
811 (size_t)dim * (size_t)rank_a, (size_t)rank_a * (size_t)dim, (size_t)dim,
812 (size_t)dim * (size_t)rank_g, (size_t)rank_g * (size_t)dim,
813 0, 0, 0,
814 (size_t)dim, (size_t)dim, (size_t)dim, (size_t)dim, (size_t)dim};
815 ro_e[3] = state_bytes / sizeof(float);
816 if (!layer0) {
817 ro_e[24] = (size_t)dim * (size_t)rank_v;
818 ro_e[25] = (size_t)rank_v * (size_t)dim;
819 ro_e[26] = (size_t)dim;
820 }
821 for (int j = 0; j < 32; ++j) {
822 if (!ro[j] || ro_e[j] == 0) continue;
823 if (ro_e[j] > SIZE_MAX / sizeof(float)) return -1;
824 const size_t nbytes = ro_e[j] * sizeof(float);
825 if (overlaps(scratch, scratch_bytes, ro[j], nbytes)) return -1;
826 for (int o = 0; o < 4; ++o) {
827 if (outs[o] == ro[j] && ro[j] == allow[o]) continue;
828 if (overlaps(outs[o], out_bytes[o], ro[j], nbytes)) return -1;
829 }
830 }
831 for (int o = 0; o < 4; ++o) {
832 if (overlaps(scratch, scratch_bytes, outs[o], out_bytes[o])) return -1;
833 }
834
835 float *attn_in = scratch;
836 float *xr = scratch + (size_t)dim;
837 float *xw = scratch + (size_t)dim * 2u;
838 float *xk = scratch + (size_t)dim * 3u;
839 float *xvm = scratch + (size_t)dim * 4u;
840 float *xa = scratch + (size_t)dim * 5u;
841 float *xg = scratch + (size_t)dim * 6u;
842 float *r = scratch + (size_t)dim * 7u;
843 float *k_raw = scratch + (size_t)dim * 8u;
844 float *v_in = scratch + (size_t)dim * 9u;
845 float *y_pre = scratch + (size_t)dim * 10u;
846 float *attn_out = scratch + (size_t)dim * 11u;
847 float *tail_scratch = scratch + (size_t)dim * 12u;
848
849 if (layer_norm(h, ln_g, ln_b, attn_in, dim, ctx) != 0) {
850 return -1;
851 }
852 if (ck_rwkv7_time_mix_shifts(attn_in, att_shift_in, mu_r, mu_w, mu_k,
853 mu_v, mu_a, mu_g, xr, xw, xk, xvm, xa,
854 xg, att_shift_out, dim) != 0) {
855 return -1;
856 }
857 if (gemv(Wr, xr, r, ctx) != 0 || gemv(Wk, xk, k_raw, ctx) != 0 ||
858 gemv(Wv, xvm, v_in, ctx) != 0) {
859 return -1;
860 }
861 if (ck_rwkv7_time_mix_tail(r, k_raw, v_in, v_first_in, xw, xa, xg,
862 layer0 ? NULL : xvm, w1, w2, w0, a1, a2, a0,
863 g1, g2, layer0 ? NULL : v1,
864 layer0 ? NULL : v2, layer0 ? NULL : v0, k_k,
865 k_a, rk, gn_w, gn_b, norm_eps, state_in,
866 state_out, y_pre, v_first_out, dim,
867 num_heads, head_dim, rank_w, rank_a, rank_g,
868 layer0 ? 1 : rank_v, tail_scratch,
869 (size_t)dim * 7u) != 0) {
870 return -1;
871 }
872 if (gemv(Wo, y_pre, attn_out, ctx) != 0) {
873 return -1;
874 }
875 for (int i = 0; i < dim; ++i) {
876 h_out[i] = h[i] + attn_out[i];
877 }
878 return 0;
879}
ckernel_rwkv7_decode.h
ck_rwkv7_gemv_fn
int(* ck_rwkv7_gemv_fn)(const float *W, const float *x, float *y, void *ctx)
Definition ckernel_rwkv7_decode.h:270
ck_rwkv7_layer_norm_fn
int(* ck_rwkv7_layer_norm_fn)(const float *x, const float *gamma, const float *beta, float *y, int dim, void *ctx)
Definition ckernel_rwkv7_decode.h:274
ck_rwkv7_lora_gates
int ck_rwkv7_lora_gates(const float *xw, const float *xa, const float *xg, const float *w1, const float *w2, const float *w0, const float *a1, const float *a2, const float *a0, const float *g1, const float *g2, int dim, int rank_w, int rank_a, int rank_g, float *w_log, float *a_out, float *g_out)
Definition rwkv7_decode.c:158
ck_rwkv7_apply_kk_factor
int ck_rwkv7_apply_kk_factor(const float *k, const float *k_k, float *kk_raw, int dim)
Definition rwkv7_decode.c:141
overlaps
static int overlaps(const void *a, size_t an, const void *b, size_t bn)
Definition rwkv7_decode.c:9
ck_rwkv7_vgate_vmix
int ck_rwkv7_vgate_vmix(const float *v_in, const float *v_first, const float *xv, const float *v1, const float *v2, const float *v0, int dim, int rank_v, float *v_out, float *v_first_out, float *v_gate)
Definition rwkv7_decode.c:239
ck_rwkv7_channelmix_decode
int ck_rwkv7_channelmix_decode(const float *x, const float *shift_in, const float *mu, const float *Wk, const float *Wv, float *shift_out, float *out, int dim, int hidden_dim, float *scratch, size_t scratch_elems)
Definition rwkv7_decode.c:432
ck_rwkv7_time_mix_scratch_elems
size_t ck_rwkv7_time_mix_scratch_elems(int dim)
Definition rwkv7_decode.c:505
ck_rwkv7_time_mix_block_scratch_elems
size_t ck_rwkv7_time_mix_block_scratch_elems(int dim)
Definition rwkv7_decode.c:700
ck_rwkv7_scale_k
int ck_rwkv7_scale_k(const float *k_in, const float *a, const float *k_a, float *k_out, int dim)
Definition rwkv7_decode.c:127
ck_rwkv7_state_readout
int ck_rwkv7_state_readout(const float *state, const float *r, float *y, int head_dim)
Definition rwkv7_decode.c:319
ck_rwkv7_token_shift_lerp
int ck_rwkv7_token_shift_lerp(const float *x, const float *shift_in, const float *mu, float *xm, float *shift_out, int dim)
Definition rwkv7_decode.c:15
ck_rwkv7_time_mix_shifts
int ck_rwkv7_time_mix_shifts(const float *x, const float *shift_in, const float *mu_r, const float *mu_w, const float *mu_k, const float *mu_v, const float *mu_a, const float *mu_g, float *xr, float *xw, float *xk, float *xv, float *xa, float *xg, float *shift_out, int dim)
Definition rwkv7_decode.c:41
ck_rwkv7_channelmix_scratch_elems
size_t ck_rwkv7_channelmix_scratch_elems(int dim, int hidden_dim)
Definition rwkv7_decode.c:425
ck_rwkv7_state_step
int ck_rwkv7_state_step(const float *w_log, const float *k, const float *v, const float *kk, const float *a, const float *state_in, float *state_out, int head_dim)
Definition rwkv7_decode.c:287
ck_rwkv7_time_mix_block
int ck_rwkv7_time_mix_block(const float *h, const float *att_shift_in, const float *v_first_in, const float *state_in, const float *ln_g, const float *ln_b, const float *mu_r, const float *mu_w, const float *mu_k, const float *mu_v, const float *mu_a, const float *mu_g, const float *Wr, const float *Wk, const float *Wv, const float *Wo, const float *w1, const float *w2, const float *w0, const float *a1, const float *a2, const float *a0, const float *g1, const float *g2, const float *v1, const float *v2, const float *v0, const float *k_k, const float *k_a, const float *rk, const float *gn_w, const float *gn_b, float norm_eps, int dim, int num_heads, int head_dim, int rank_w, int rank_a, int rank_g, int rank_v, int is_layer_zero, ck_rwkv7_gemv_fn gemv, ck_rwkv7_layer_norm_fn layer_norm, void *ctx, float *h_out, float *att_shift_out, float *state_out, float *v_first_out, float *scratch, size_t scratch_elems)
Definition rwkv7_decode.c:707
ck_rwkv7_sigmoid
static float ck_rwkv7_sigmoid(float x)
Definition rwkv7_decode.c:154
ck_rwkv7_wkv_decode
int ck_rwkv7_wkv_decode(const float *r, const float *w_log, const float *k, const float *v, const float *kk, const float *a, const float *g, const float *rk, const float *ln_w, const float *ln_b, float norm_eps, const float *state_in, float *state_out, float *y_out, int num_heads, int head_dim)
Definition rwkv7_decode.c:336
ck_rwkv7_norm_kk
int ck_rwkv7_norm_kk(const float *kk_in, float *kk_out, int num_heads, int head_dim)
Definition rwkv7_decode.c:99
ck_rwkv7_time_mix_tail
int ck_rwkv7_time_mix_tail(const float *r, const float *k_raw, const float *v_in, const float *v_first_in, const float *xw, const float *xa, const float *xg, const float *xv, const float *w1, const float *w2, const float *w0, const float *a1, const float *a2, const float *a0, const float *g1, const float *g2, const float *v1, const float *v2, const float *v0, const float *k_k, const float *k_a, const float *rk, const float *ln_w, const float *ln_b, float norm_eps, const float *state_in, float *state_out, float *y_preproj, float *v_first_out, int dim, int num_heads, int head_dim, int rank_w, int rank_a, int rank_g, int rank_v, float *scratch, size_t scratch_elems)
Definition rwkv7_decode.c:510