code wiki / _hdl_build / nx_nofloat_linattn.nx
nx_nofloat_linattn.nx source
↩ module page · 52 lines · 2387 B
1// nx_nofloat_linattn.nx -- SOVEREIGN no-float LINEAR ATTENTION recurrent kernel = the foundation of Kimi
2// Delta Attention (KDA, the K3 namesake; operator 2026-07-19 "state of the art"). Linear attention replaces
3// softmax(QK^T)V with a RECURRENT matrix state: S_t = S_{t-1} + k_t (x) v_t (outer product, d x d);
4// o_t = q_t S_t. This is O(1) memory per step (no growing KV cache) -> the 75% KV cut + 6x decode @1M that
5// KDA delivers. Done in PURE INTEGER Q16 (outer product + matvec, accumulate-then-shift) = BIT-EXACT
6// deterministic linear attention -- a float linear-attn stack drifts by accumulation order; ours does not.
7// KDA extends THIS with (a) the delta rule S_t = S_{t-1}(I - b k k^T) + b k v^T and (b) channel-wise gating
8// alpha_t (diagonal forget) -- both follow-ons on this proven kernel. Grounded arXiv 2510.26692 (Kimi Linear).
9// license_tier: ORIGINAL No hw writes (Rule 26).
10import "nx_syscalls.nx"
11
12const LA_QBITS: i64 = 16
13const LA_DMAX: i64 = 16 // max head dim
14
15// zero a d x d state
16func la_zero(s: *i64, d: i64) -> i64 { var i: i64 = 0; while i < d*d { s[i] = 0; i = i + 1 } return 0 }
17// S += k (x) v : S[i][j] += k[i]*v[j] (Q16 * Q16 >> 16)
18func la_accum(s: *i64, k: *i64, v: *i64, d: i64) -> i64 {
19 var i: i64 = 0
20 while i < d {
21 var j: i64 = 0
22 let ki: i64 = k[i]
23 while j < d { s[i*d+j] = s[i*d+j] + ((ki * v[j]) >> LA_QBITS); j = j + 1 }
24 i = i + 1
25 }
26 return 0
27}
28// o = q S : o[j] = sum_i q[i] S[i][j] (accumulate i64, >>16)
29func la_readout(s: *i64, q: *i64, o: *i64, d: i64) -> i64 {
30 var j: i64 = 0
31 while j < d {
32 var acc: i64 = 0
33 var i: i64 = 0
34 while i < d { acc = acc + q[i] * s[i*d+j]; i = i + 1 }
35 o[j] = acc >> LA_QBITS
36 j = j + 1
37 }
38 return 0
39}
40// RECURRENT linear attention over a length-T sequence: Q,K,V are T x d (row-major, Q16).
41// writes O (T x d): o_t = q_t * (sum_{s<=t} k_s (x) v_s). O(d^2) state, O(1) per step.
42func la_forward(qm: *i64, km: *i64, vm: *i64, om: *i64, t: i64, d: i64) -> i64 {
43 let s: *i64 = sys_mmap(LA_DMAX*LA_DMAX*8) as *i64
44 la_zero(s, d)
45 var ti: i64 = 0
46 while ti < t {
47 la_accum(s, ((km as i64) + ti*d*8) as *i64, ((vm as i64) + ti*d*8) as *i64, d)
48 la_readout(s, ((qm as i64) + ti*d*8) as *i64, ((om as i64) + ti*d*8) as *i64, d)
49 ti = ti + 1
50 }
51 return 0
52}