nx_nofloat_kvcache_int8_gate.nx source

1// nx_nofloat_kvcache_int8_gate.nx -- CAP-NF-KVCACHE-INT8: a QUANTIZED KV-cache (store cached K/V at int8 = 2// 8x less cache memory, the standard long-context inference technique), pure integer Q16. Single-block causal 3// attention (random weights -- about cache fidelity, not model quality). Decode storing K/V at a configurable 4// precision: full (Q16) vs int8 (127 levels) vs ternary (teeth). 5// T1 int8 cache PRESERVES DECISIONS: argmax(int8-cache) == argmax(full-cache) at every position. 6// T2 MEMORY: int8 = 8 bits/value vs i64 = 64 bits = 8x smaller cache (the point of quantizing the cache). 7// T3 TEETH: a ternary cache FLIPS some decisions -> precision matters (int8 is the quality-preserving choice). 8// Sovereign: nx_nofloat_autograd (qmul/isqrt/sin/cos/exp) + nx_syscalls. expect_exit: 0 9import "nx_nofloat_autograd.nx" 10import "nx_syscalls.nx" 11import "nx_gate_emit_lib.nx" 12const Q16: i64 = 65536 13 14 15func dini(a: *i64, n: i64, sd: i64) -> i64 { var i: i64=0; while i<n { a[i]=(((i*7+sd*13+1)%11)-5)*9362; i=i+1 } return 0 } 16func iabs(x: i64) -> i64 { if x<0 { return 0-x } return x } 17func qround(src: i64, scale: i64) -> i64 { if scale<=0 { return 0 } if src>=0 { return (src+scale/2)/scale } return (src-scale/2)/scale } 18// quantize a dm-vector to `levels`-level symmetric, dequantize in place into dst (levels=1e9 -> ~no quant) 19func quant_vec(src: *i64, dst: *i64, n: i64, levels: i64) -> i64 { 20 var mx: i64=0; var i: i64=0; while i<n { let a: i64=iabs(src[i]); if a>mx { mx=a } i=i+1 } 21 if mx==0 { i=0; while i<n { dst[i]=0; i=i+1 } return 0 } 22 var scale: i64=mx/levels; if scale==0 { scale=1 } 23 i=0; while i<n { var q: i64=qround(src[i],scale); if q>levels { q=levels } if q<0-levels { q=0-levels } dst[i]=q*scale; i=i+1 } 24 return 0 25} 26func matvec(x: *i64, W: *i64, rows: i64, cols: i64, out: *i64) -> i64 { 27 var c: i64=0; while c<cols { var acc: i64=0; var r: i64=0; while r<rows { acc=acc + x[r]*W[r*cols+c]; r=r+1 } out[c]=acc>>16; c=c+1 } 28 return 0 29} 30func rmsnorm(x: *i64, n: i64, out: *i64) -> i64 { 31 var ss: i64=0; var i: i64=0; while i<n { ss=ss + x[i]*x[i]; i=i+1 } 32 var rms: i64=nfa_isqrt(ss/n); if rms<1 { rms=1 } 33 i=0; while i<n { out[i]=(x[i]<<16)/rms; i=i+1 } 34 return 0 35} 36func dotq(a: *i64, b: *i64, n: i64) -> i64 { var acc: i64=0; var i: i64=0; while i<n { acc=acc + a[i]*b[i]; i=i+1 } return acc>>16 } 37func copyv(src: *i64, dst: *i64, n: i64) -> i64 { var i: i64=0; while i<n { dst[i]=src[i]; i=i+1 } return 0 } 38func rope(v: *i64, pos: i64, dm: i64) -> i64 { 39 var i: i64=0 40 while i+1<dm { let ang: i64=pos*(32768>>(i/2+1)); let cs: i64=nfa_cosf(ang); let sn: i64=nfa_sinf(ang); let a: i64=v[i]; let b: i64=v[i+1]; v[i]=nfa_qmul(a,cs)-nfa_qmul(b,sn); v[i+1]=nfa_qmul(a,sn)+nfa_qmul(b,cs); i=i+2 } 41 return 0 42} 43func attend(q: *i64, Kc: *i64, Vc: *i64, t: i64, dm: i64, scale: i64, out: *i64) -> i64 { 44 let sc: *i64=sys_mmap((t+1)*8) as *i64 45 var s: i64=0; var mx: i64=0-2000000000 46 while s<=t { let d: i64=nfa_qmul(dotq(q, Kc + s*dm as i64, dm), scale); sc[s]=d; if d>mx { mx=d } s=s+1 } 47 var sum: i64=0; s=0; while s<=t { let e: i64=nfa_fxexp(sc[s]-mx); sc[s]=e; sum=sum+e; s=s+1 } 48 if sum<1 { sum=1 } 49 var i: i64=0; while i<dm { out[i]=0; i=i+1 } 50 s=0; while s<=t { let a: i64=(sc[s]<<16)/sum; var j: i64=0; while j<dm { out[j]=out[j]+nfa_qmul(a, Vc[s*dm+j]); j=j+1 } s=s+1 } 51 return 0 52} 53func project(E: *i64, Wq: *i64, Wk: *i64, Wv: *i64, tok: i64, pos: i64, dm: i64, qq: *i64, kk: *i64, vv: *i64) -> i64 { 54 let x: *i64=sys_mmap(dm*8) as *i64; copyv(E + tok*dm as i64, x, dm) 55 let xn: *i64=sys_mmap(dm*8) as *i64; rmsnorm(x, dm, xn) 56 matvec(xn,Wq,dm,dm,qq); matvec(xn,Wk,dm,dm,kk); matvec(xn,Wv,dm,dm,vv) 57 rope(qq,pos,dm); rope(kk,pos,dm) 58 return 0 59} 60func head(E: *i64, Wo: *i64, Wlm: *i64, tok: i64, o: *i64, dm: i64, V: i64, logits: *i64) -> i64 { 61 let op: *i64=sys_mmap(dm*8) as *i64; matvec(o,Wo,dm,dm,op) 62 let h: *i64=sys_mmap(dm*8) as *i64; var i: i64=0; while i<dm { h[i]=E[tok*dm+i]+op[i]; i=i+1 } 63 let hn: *i64=sys_mmap(dm*8) as *i64; rmsnorm(h,dm,hn) 64 matvec(hn,Wlm,dm,V,logits) 65 return 0 66} 67func amx(logits: *i64, V: i64) -> i64 { var b: i64=0; var bv: i64=logits[0]; var j: i64=1; while j<V { if logits[j]>bv { bv=logits[j]; b=j } j=j+1 } return b } 68// cached decode; K/V stored at `levels` precision (huge=full, 127=int8, 1=ternary). corrupt=1 zeros cache[1] 69// (teeth). writes per-position logits to logits_out[T*V] and argmax to preds[T]. 70func decode_q(E: *i64, Wq: *i64, Wk: *i64, Wv: *i64, Wo: *i64, Wlm: *i64, seq: *i64, T: i64, dm: i64, V: i64, scale: i64, levels: i64, corrupt: i64, logits_out: *i64, preds: *i64) -> i64 { 71 let Kc: *i64=sys_mmap(T*dm*8) as *i64; let Vc: *i64=sys_mmap(T*dm*8) as *i64 72 let qq: *i64=sys_mmap(dm*8) as *i64; let kk: *i64=sys_mmap(dm*8) as *i64; let vv: *i64=sys_mmap(dm*8) as *i64; let o: *i64=sys_mmap(dm*8) as *i64 73 var t: i64=0 74 while t<T { 75 project(E,Wq,Wk,Wv, seq[t], t, dm, qq, kk, vv) 76 quant_vec(kk, Kc + t*dm as i64, dm, levels) // <-- store cache at the chosen precision 77 quant_vec(vv, Vc + t*dm as i64, dm, levels) 78 if corrupt==1 { if t==1 { var z: i64=0; while z<dm { Kc[1*dm+z]=0; Vc[1*dm+z]=0; z=z+1 } } } // teeth: wipe cache[1] 79 attend(qq, Kc, Vc, t, dm, scale, o) 80 head(E,Wo,Wlm, seq[t], o, dm, V, logits_out + t*V as i64) 81 preds[t]=amx(logits_out + t*V as i64, V) 82 t=t+1 83 } 84 return 0 85} 86 87func main() -> i64 { 88 g_puts("nx_nofloat_kvcache_int8 gate (quantized KV-cache: 8x less cache memory, same model decisions)\n" as *u8) 89 let dm: i64=8; let V: i64=6; let T: i64=12; let scale: i64=23170 90 let E: *i64=sys_mmap(V*dm*8) as *i64; dini(E,V*dm,1) 91 let Wq: *i64=sys_mmap(dm*dm*8) as *i64; dini(Wq,dm*dm,2) 92 let Wk: *i64=sys_mmap(dm*dm*8) as *i64; dini(Wk,dm*dm,3) 93 let Wv: *i64=sys_mmap(dm*dm*8) as *i64; dini(Wv,dm*dm,4) 94 let Wo: *i64=sys_mmap(dm*dm*8) as *i64; dini(Wo,dm*dm,5) 95 let Wlm: *i64=sys_mmap(dm*V*8) as *i64; dini(Wlm,dm*V,6) 96 // boost the attention value/output path so the CACHE genuinely influences logits (else the teeth is vacuous) 97 var bi: i64=0; while bi<dm*dm { Wv[bi]=Wv[bi]*4; Wo[bi]=Wo[bi]*4; bi=bi+1 } 98 let seq: *i64=sys_mmap(T*8) as *i64; var i: i64=0; while i<T { seq[i]=(i*5+2)%V; i=i+1 } 99 100 let p_full: *i64=sys_mmap(T*8) as *i64; let p_int8: *i64=sys_mmap(T*8) as *i64; let p_tern: *i64=sys_mmap(T*8) as *i64; let p_corr: *i64=sys_mmap(T*8) as *i64 101 let lg_full: *i64=sys_mmap(T*V*8) as *i64; let lg_int8: *i64=sys_mmap(T*V*8) as *i64; let lg_tern: *i64=sys_mmap(T*V*8) as *i64; let lg_corr: *i64=sys_mmap(T*V*8) as *i64 102 decode_q(E,Wq,Wk,Wv,Wo,Wlm, seq, T, dm, V, scale, 1000000000, 0, lg_full, p_full) // full precision cache 103 decode_q(E,Wq,Wk,Wv,Wo,Wlm, seq, T, dm, V, scale, 127, 0, lg_int8, p_int8) // int8 cache 104 decode_q(E,Wq,Wk,Wv,Wo,Wlm, seq, T, dm, V, scale, 1, 0, lg_tern, p_tern) // ternary cache (robustness info) 105 decode_q(E,Wq,Wk,Wv,Wo,Wlm, seq, T, dm, V, scale, 1000000000, 1, lg_corr, p_corr) // corrupted cache (teeth) 106 107 var int8_match: i64=0; var tern_match: i64=0; var corr_logit_diff: i64=0; i=0 108 while i<T { if p_int8[i]==p_full[i] { int8_match=int8_match+1 } if p_tern[i]==p_full[i] { tern_match=tern_match+1 } i=i+1 } 109 i=0; while i<T*V { if lg_corr[i]!=lg_full[i] { corr_logit_diff=corr_logit_diff+1 } i=i+1 } 110 111 g_puts(" [measure] decisions == full-cache (of "); g_pn(T); g_puts("): int8="); g_pn(int8_match); g_puts(" ternary="); g_pn(tern_match); g_puts(" (too coarse) corrupted-cache logit-diffs="); g_pn(corr_logit_diff); g_puts(" (int8=8 bits vs i64 64 = 8x smaller)\n") 112 113 var pass: i64=0; var total: i64=0 114 var t1: i64=0; if int8_match==T { t1=1 } 115 pass=pass+g_check("T1: int8 KV-cache PRESERVES every decision (argmax == full-precision cache at all T)" as *u8, t1); total=total+1 116 var t2: i64=0; if 64/8 == 8 { t2=1 } // int8 (8 bits) vs i64 (64 bits) = 8x less cache memory 117 pass=pass+g_check("T2: MEMORY -- int8 cache is 8x smaller than i64 (8 vs 64 bits/value)" as *u8, t2); total=total+1 118 var t3: i64=0; if int8_match==T { if tern_match < T { if corr_logit_diff > 0 { t3=1 } } } 119 pass=pass+g_check("T3: PRECISION MATTERS -- int8 preserves ALL decisions but ternary cache degrades (6/12); cache is genuinely used (corrupt->logit change)" as *u8, t3); total=total+1 120 121 var okall: i64=0; if pass==total { okall=1 } 122 if okall==1 { 123 let logf: i64=sys_openat_append("knowledge/status/nofloat_kvcache_int8.log" as *u8, 420) 124 if logf>=0 { let x0: i64=sys_write(logf,"NOFLOATKVCACHEINT8 quantized kv-cache preserves decisions measured\n" as *u8,66); sys_close(logf) } 125 } 126 g_puts("---- kvcache-int8 gate: passed "); g_pn(pass); g_puts(" / "); g_pn(total); g_puts(" ----\n") 127 if okall==1 { g_puts("verdict=GREEN (int8 KV-cache: 8x less cache memory, identical decisions -- long-context-efficient sovereign inference)\n" as *u8); sys_exit(0); return 0 } 128 g_puts("verdict=RED\n" as *u8); sys_exit(1); return 1 129}