nx_nofloat_linattn_gate.nx source

1// nx_nofloat_linattn_gate.nx -- gate for the no-float linear-attention recurrent kernel (KDA foundation). 2// T1 recurrent output MATCHES the direct linear-attention formula o_t = sum_{s<=t}(q_t.k_s) v_s 3// T2 CAUSAL: o_t depends only on s<=t (truncating the sequence at t leaves o_t unchanged) 4// T3 O(1)-memory state accumulates correctly (S after T = sum k_s (x) v_s) 5// T4 DETERMINISTIC bit-identical (the exceed: float linear-attn drifts by accumulation order) 6// license_tier: ORIGINAL No hw writes (Rule 26). expect_exit: 0 7import "nx_nofloat_linattn.nx" 8import "nx_gate_verdict.nx" 9import "nx_syscalls.nx" 10 11const LG_T: i64 = 4 12const LG_D: i64 = 4 13const LG_Q: i64 = 65536 14const LG_QBITS: i64 = 16 15 16// direct linear attention: o_t[j] = sum_{s<=t} (sum_i q_t[i] k_s[i]) v_s[j] (Q16) 17func lg_direct(qm: *i64, km: *i64, vm: *i64, om: *i64, t: i64, d: i64) -> i64 { 18 var ti: i64 = 0 19 while ti < t { 20 var j: i64 = 0 21 while j < d { 22 var acc: i64 = 0 23 var s: i64 = 0 24 while s <= ti { 25 // qk = q_t . k_s 26 var qk: i64 = 0 27 var i: i64 = 0 28 while i < d { qk = qk + qm[ti*d+i] * km[s*d+i]; i = i + 1 } 29 qk = qk >> LG_QBITS 30 acc = acc + ((qk * vm[s*d+j]) >> LG_QBITS) 31 s = s + 1 32 } 33 om[ti*d+j] = acc 34 j = j + 1 35 } 36 ti = ti + 1 37 } 38 return 0 39} 40 41func main() -> i64 { 42 let ctr: *i64 = gv_ctr() 43 gv_head("nx_nofloat_linattn gate -- linear-attention recurrent kernel (KDA foundation, no-float)" as *u8) 44 let t: i64 = LG_T 45 let d: i64 = LG_D 46 let qm: *i64 = sys_mmap(64*8) as *i64 47 let km: *i64 = sys_mmap(64*8) as *i64 48 let vm: *i64 = sys_mmap(64*8) as *i64 49 // deterministic pseudo-values in Q16 (small, distinct): fill Q,K,V 50 var i: i64 = 0 51 while i < t*d { 52 qm[i] = ((i * 7) % 5 + 1) * (LG_Q / 8) // 0.125..0.625 53 km[i] = ((i * 3) % 4 + 1) * (LG_Q / 8) 54 vm[i] = ((i * 5) % 6 + 1) * (LG_Q / 8) 55 i = i + 1 56 } 57 let orec: *i64 = sys_mmap(64*8) as *i64 58 let odir: *i64 = sys_mmap(64*8) as *i64 59 la_forward(qm, km, vm, orec, t, d) 60 lg_direct(qm, km, vm, odir, t, d) 61 62 // T1: recurrent == direct (within tiny Q16 rounding tolerance) 63 var t1: i64 = 1 64 i = 0 65 while i < t*d { var e: i64 = orec[i] - odir[i]; if e < 0 { e = 0 - e } if e > 4 { t1 = 0 } i = i + 1 } 66 gv_check("T1 recurrent output MATCHES direct linear-attention formula (<=4 ulp)" as *u8, t1, ctr) 67 68 // T2: causal -- o_1 (t=2 seq) unchanged when sequence truncated to length 2 69 let orec2: *i64 = sys_mmap(64*8) as *i64 70 la_forward(qm, km, vm, orec2, 2, d) 71 var t2: i64 = 1 72 i = 0 73 while i < 2*d { if orec2[i] != orec[i] { t2 = 0 } i = i + 1 } 74 gv_check("T2 CAUSAL (truncating seq to length 2 leaves o_0,o_1 identical)" as *u8, t2, ctr) 75 76 // T3: state after full pass = sum k_s (x) v_s 77 let s: *i64 = sys_mmap(LA_DMAX*LA_DMAX*8) as *i64 78 la_zero(s, d) 79 var ts: i64 = 0 80 while ts < t { la_accum(s, ((km as i64)+ts*d*8) as *i64, ((vm as i64)+ts*d*8) as *i64, d); ts = ts + 1 } 81 // reference S[0][0] = sum_s k_s[0]*v_s[0] >>16 82 var ref00: i64 = 0 83 ts = 0 84 while ts < t { ref00 = ref00 + ((km[ts*d+0] * vm[ts*d+0]) >> LG_QBITS); ts = ts + 1 } 85 var t3: i64 = 0 86 if s[0] == ref00 { t3 = 1 } 87 gv_check("T3 O(1)-memory state accumulates (S[0][0] = sum k_s[0] v_s[0])" as *u8, t3, ctr) 88 89 // T4: deterministic 90 let ob: *i64 = sys_mmap(64*8) as *i64 91 la_forward(qm, km, vm, ob, t, d) 92 var t4: i64 = 1 93 i = 0 94 while i < t*d { if ob[i] != orec[i] { t4 = 0 } i = i + 1 } 95 gv_check("T4 DETERMINISTIC bit-identical (the exceed vs order-dependent float linear-attn)" as *u8, t4, ctr) 96 97 let rc: i64 = gv_verdict("NOFLOAT-LINATTN-GATE" as *u8, ctr, "linear-attention recurrent kernel (KDA foundation): matches direct formula, causal, O(1)-state, bit-exact deterministic" as *u8) 98 sys_exit(rc) 99 return rc 100}