nx_nofloat_attention_gate.nx source
↩ module page · 164 lines · 7690 B
1// nx_nofloat_attention_gate.nx -- ASSEMBLE the trio into one DETERMINISTIC transformer attention layer.
2// Composes CAP-NOFLOAT-LAYERNORM + CAP-NOFLOAT-GEMM (matmul) + CAP-NOFLOAT-SOFTMAX, all integer Q16:
3// LayerNorm(X) -> Q,K,V = Xn*Wq/Wk/Wv -> scores = QK^T * 1/sqrt(D) -> softmax per row -> out = attn*V + X.
4// Every op integer -> the WHOLE LAYER is bit-exact DETERMINISTIC. Proven by running the full forward with
5// the score-dot and output-weighted-sum in FORWARD then REVERSE order: the output must be bit-identical
6// (float attention CANNOT promise this). This is the unit a coding-LLM is stacked from -- now sovereign +
7// reproducible. No hw writes (Rule 26). expect_exit: 0 license_tier: ORIGINAL
8import "nx_syscalls.nx"
9import "nx_gate_verdict.nx"
10
11const Q16: i64 = 65536
12const LOG2E: i64 = 94548
13const PC0: i64 = 65536
14const PC1: i64 = 45426
15const PC2: i64 = 15743
16const PC3: i64 = 4367
17
18func at_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
19func at_num(v: i64) -> i64 { let b: *u8=sys_mmap(28); var m: i64=v; if m<0{m=0-m;sys_write(1,"-" as *u8,1)} let t: *u8=sys_mmap(28); var k: i64=0; if m==0{t[0]=48 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} var i: i64=0; while i<k{b[i]=t[k-1-i];i=i+1} sys_write(1,b,k); return 0 }
20func qmul(a: i64, b: i64) -> i64 { return (a*b) >> 16 }
21func absd(a: i64, b: i64) -> i64 { if a>b { return a-b } return b-a }
22
23func isqrt(v: i64) -> i64 {
24 if v <= 0 { return 0 }
25 if v < 4 { return 1 }
26 var x: i64 = v
27 var y: i64 = (x + 1) >> 1
28 var go: i64 = 1
29 while go==1 { if y < x { x = y; y = (x + v/x) >> 1 } else { go = 0 } }
30 return x
31}
32func fx_exp(x: i64) -> i64 {
33 var xm: i64 = 0 - x
34 if x > 0 { xm = 0 }
35 let ym: i64 = (xm * LOG2E) >> 16
36 let yi: i64 = ym >> 16
37 let yf: i64 = ym - (yi << 16)
38 let g: i64 = Q16 - yf
39 var t: i64 = PC3
40 t = PC2 + ((g * t) >> 16)
41 t = PC1 + ((g * t) >> 16)
42 t = PC0 + ((g * t) >> 16)
43 t = t >> 1
44 if yi >= 31 { return 0 }
45 return t >> yi
46}
47func ln_row(x: *i64, off: i64, D: i64) -> i64 {
48 var s: i64=0; var i: i64=0
49 while i<D { s=s+x[off+i]; i=i+1 }
50 let mean: i64 = s/D
51 var vs: i64=0; i=0
52 while i<D { let d: i64=x[off+i]-mean; vs=vs+((d*d)>>16); i=i+1 }
53 let varq: i64=vs/D
54 let sd: i64=isqrt((varq+7)<<16)
55 if sd<=0 { return 0 }
56 i=0; while i<D { x[off+i]=((x[off+i]-mean)<<16)/sd; i=i+1 }
57 return 0
58}
59func proj(src: *i64, W: *i64, dst: *i64, T: i64, D: i64) -> i64 {
60 var i: i64=0
61 while i<T { var d: i64=0
62 while d<D { var s: i64=0; var k: i64=0
63 while k<D { s=s+qmul(src[i*D+k], W[k*D+d]); k=k+1 }
64 dst[i*D+d]=s; d=d+1 }
65 i=i+1 }
66 return 0
67}
68func softmax_row(sc: *i64, at: *i64, row: i64, T: i64) -> i64 {
69 var m: i64=sc[row*T]; var j: i64=1
70 while j<T { if sc[row*T+j]>m { m=sc[row*T+j] } j=j+1 }
71 var s: i64=0; j=0
72 while j<T { let e: i64=fx_exp(sc[row*T+j]-m); at[row*T+j]=e; s=s+e; j=j+1 }
73 if s<=0 { s=1 }
74 j=0; while j<T { at[row*T+j]=(at[row*T+j]<<16)/s; j=j+1 }
75 return s
76}
77func attn_fwd(X: *i64, Wq: *i64, Wk: *i64, Wv: *i64, Xn: *i64, Q: *i64, K: *i64, V: *i64, sc: *i64, at: *i64, out: *i64, T: i64, D: i64, scale: i64, rev: i64) -> i64 {
78 var i: i64=0
79 while i<T*D { Xn[i]=X[i]; i=i+1 }
80 i=0; while i<T { ln_row(Xn, i*D, D); i=i+1 }
81 proj(Xn, Wq, Q, T, D)
82 proj(Xn, Wk, K, T, D)
83 proj(Xn, Wv, V, T, D)
84 i=0
85 while i<T { var j: i64=0
86 while j<T {
87 var s: i64=0
88 if rev==0 { var d: i64=0; while d<D { s=s+qmul(Q[i*D+d], K[j*D+d]); d=d+1 } }
89 else { var d: i64=D-1; while d>=0 { s=s+qmul(Q[i*D+d], K[j*D+d]); d=d-1 } }
90 sc[i*T+j]=qmul(s, scale)
91 j=j+1 }
92 i=i+1 }
93 i=0; while i<T { softmax_row(sc, at, i, T); i=i+1 }
94 i=0
95 while i<T { var d: i64=0
96 while d<D {
97 var s: i64=0
98 if rev==0 { var j: i64=0; while j<T { s=s+qmul(at[i*T+j], V[j*D+d]); j=j+1 } }
99 else { var j: i64=T-1; while j>=0 { s=s+qmul(at[i*T+j], V[j*D+d]); j=j-1 } }
100 out[i*D+d]=s + X[i*D+d]
101 d=d+1 }
102 i=i+1 }
103 return 0
104}
105
106func main() -> i64 {
107 at_puts("ASSEMBLE the trio: one DETERMINISTIC no-float attention layer (LayerNorm->QKV->softmax->*V+res)\n\n" as *u8)
108 let T: i64=4
109 let D: i64=4
110 let scale: i64=32768 // 1/sqrt(4) = 0.5
111 let X: *i64 = sys_mmap(T*D*8) as *i64
112 let Wq: *i64 = sys_mmap(D*D*8) as *i64
113 let Wk: *i64 = sys_mmap(D*D*8) as *i64
114 let Wv: *i64 = sys_mmap(D*D*8) as *i64
115 let Xn: *i64 = sys_mmap(T*D*8) as *i64
116 let Q: *i64 = sys_mmap(T*D*8) as *i64
117 let K: *i64 = sys_mmap(T*D*8) as *i64
118 let V: *i64 = sys_mmap(T*D*8) as *i64
119 let sc: *i64 = sys_mmap(T*T*8) as *i64
120 let at0: *i64 = sys_mmap(T*T*8) as *i64
121 let at1: *i64 = sys_mmap(T*T*8) as *i64
122 let out0: *i64 = sys_mmap(T*D*8) as *i64
123 let out1: *i64 = sys_mmap(T*D*8) as *i64
124
125 var i: i64=0
126 while i<T*D { X[i]=(((i*2+ (i/D)*3)%7)-3)*16384; i=i+1 }
127 i=0
128 while i<D*D { Wq[i]=(((i+1)%5)-2)*8192; Wk[i]=(((i*2+1)%5)-2)*8192; Wv[i]=(((i+3)%5)-2)*8192; i=i+1 }
129
130 attn_fwd(X, Wq, Wk, Wv, Xn, Q, K, V, sc, at0, out0, T, D, scale, 0)
131 attn_fwd(X, Wq, Wk, Wv, Xn, Q, K, V, sc, at1, out1, T, D, scale, 1)
132
133 var detmis: i64=0
134 i=0; while i<T*D { if out0[i]!=out1[i] { detmis=detmis+1 } i=i+1 }
135 // attention row sums (softmax validity)
136 var rowsum_ok: i64=1
137 i=0
138 while i<T { var s: i64=0; var j: i64=0; while j<T { s=s+at0[i*T+j]; j=j+1 } if absd(s, Q16)>2000 { rowsum_ok=0 } i=i+1 }
139 // non-trivial: output differs from input somewhere
140 var changed: i64=0
141 i=0; while i<T*D { if out0[i]!=X[i] { changed=changed+1 } i=i+1 }
142
143 at_puts(" attn row0 = ["); var j: i64=0; while j<T { at_num(at0[j]); if j<T-1 { at_puts(", ") } j=j+1 } at_puts("] (Q16, sums to ~65536)\n");
144 at_puts(" out token0 = ["); i=0; while i<D { at_num(out0[i]); if i<D-1 { at_puts(", ") } i=i+1 } at_puts("]\n");
145 at_puts(" determinism mism (fwd-order vs rev-order full forward) = "); at_num(detmis); at_puts(" / "); at_num(T*D); at_puts(" output cells changed by attention = "); at_num(changed); at_puts("\n\n");
146
147 var pass: i64=0
148 var ttl: i64=0
149 ttl=ttl+1; at_puts(" T1 attention layer ran end-to-end (LayerNorm+QKV+softmax+*V+residual composed): "); if changed>0 { pass=pass+1; at_puts("PASS\n") } else { at_puts("FAIL\n") }
150 ttl=ttl+1; at_puts(" T2 softmax valid: every attention row sums to ~1.0 (65536): "); if rowsum_ok==1 { pass=pass+1; at_puts("PASS\n") } else { at_puts("FAIL\n") }
151 ttl=ttl+1; at_puts(" T3 DETERMINISTIC: full forward bit-identical under fwd vs rev sum order (0 mism): "); if detmis==0 { pass=pass+1; at_puts("PASS\n") } else { at_puts("FAIL\n") }
152 ttl=ttl+1; at_puts(" T4 the block TRANSFORMED the input (attention did real work, not identity): "); if changed>=T { pass=pass+1; at_puts("PASS\n") } else { at_puts("FAIL\n") }
153
154 at_puts("NX-NOFLOAT-ATTENTION-GATE passed "); at_num(pass); at_puts("/"); at_num(ttl)
155 // MIGRATED onto nx_gate_verdict by nx_gate_dry_apply (D001, minimal form): every check
156 // row above is untouched, so the PASS/FAIL vector cannot change; only the hand-rolled
157 // verdict emission is replaced by the ONE shared base class. Proven by nx_gate_migrate verify.
158 let ctr__dry: *i64 = gv_ctr()
159 ctr__dry[0] = pass
160 ctr__dry[1] = ttl
161 let rc__dry: i64 = gv_verdict("NOFLOAT-ATTENTION-GATE" as *u8, ctr__dry, "a full no-float attention LAYER, deterministic -- the transformer unit a coding-LLM is stacked from)" as *u8)
162 sys_exit(rc__dry)
163 return rc__dry
164}