nx_blockfloat_mha_gate.nx source
↩ module page · 174 lines · 9792 B
1// nx_blockfloat_mha_gate.nx -- BLOCK-FLOAT weights through REAL MULTI-HEAD attention. Composes the existing
2// gradcheck-verified multi-head forward (mha_fwd/mha_out from nx_nofloat_multihead_gate: RMSNorm -> Q/K/V -> split
3// H heads -> per-head RoPE + scaled causal-softmax + A.V -> concat -> out-proj -> residual) UNCHANGED, run with the
4// weight matrices (Wq,Wk,Wv,Wo) BLOCK-FLOAT quantized. Extends the single-head block result to multi-head -- the
5// defining transformer feature -- proving block-float runs real multi-head attention better than per-tensor INT8.
6// 1 multi-head attention RUNS with block-float weights (non-trivial output)
7// 2 EXCEED: block-float MHA error vs full-precision < per-tensor INT8
8// 3 DETERMINISTIC: block-float MHA run twice == bit-identical
9// 4 per-tensor is LOSSY (output != full-precision) -- block-float's win is over real info loss
10// expect_exit: 0 license_tier: ORIGINAL
11import "nx_nofloat_autograd.nx"
12import "nx_syscalls.nx"
13import "nx_gate_verdict.nx"
14
15func bb_puts(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
16func bb_pn(v: i64) -> i64 {
17 let b: *u8 = sys_mmap(28); var m: i64 = v
18 if m < 0 { m = 0 - m; sys_write(1, "-" as *u8, 1) }
19 let t: *u8 = sys_mmap(28); var k: i64 = 0
20 if m == 0 { t[0] = 48 as u8; k = 1 }
21 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
22 var i: i64 = 0
23 while i < k { b[i] = t[k - 1 - i]; i = i + 1 }
24 sys_write(1, b, k); return 0
25}
26func bb_chk(name: *u8, ok: i64) -> i64 {
27 if ok == 1 { bb_puts(" PASS " as *u8); bb_puts(name); bb_puts("\n" as *u8); return 1 }
28 bb_puts(" FAIL " as *u8); bb_puts(name); bb_puts("\n" as *u8); return 0
29}
30func bb_abs(x: i64) -> i64 { if x < 0 { return 0 - x } return x }
31func bf_bitlen(x: i64) -> i64 { var b: i64 = 0; var m: i64 = x; while m > 0 { m = m >> 1; b = b + 1 } return b }
32func bf_scale_abs(W: *i64, off: i64, B: i64, MB: i64) -> i64 {
33 var amax: i64 = 0; var i: i64 = 0
34 while i < B { let a: i64 = bb_abs(W[off + i]); if a > amax { amax = a } i = i + 1 }
35 var e: i64 = bf_bitlen(amax) - MB
36 if e < 0 { e = 0 }
37 return e
38}
39func bf_quant_mat(W: *i64, rows: i64, cols: i64, B: i64, MB: i64, out: *i64) -> i64 {
40 var r: i64 = 0
41 while r < rows {
42 var blk: i64 = 0
43 while blk < cols / B {
44 let off: i64 = r * cols + blk * B
45 let e: i64 = bf_scale_abs(W, off, B, MB); let sc: i64 = 1 << e
46 var i: i64 = 0
47 while i < B { out[off + i] = (W[off + i] / sc) * sc; i = i + 1 }
48 blk = blk + 1
49 }
50 r = r + 1
51 }
52 return 0
53}
54func bf_quant_pt(W: *i64, n: i64, MB: i64, out: *i64) -> i64 {
55 let eg: i64 = bf_scale_abs(W, 0, n, MB); let sc: i64 = 1 << eg
56 var i: i64 = 0
57 while i < n { out[i] = (W[i] / sc) * sc; i = i + 1 }
58 return 0
59}
60// mixed-magnitude fill (large block next to small block) so per-block scaling matters; sgn varies the matrices.
61func fill_mixed(W: *i64, n: i64, B: i64, sgn: i64) -> i64 {
62 var i: i64 = 0
63 while i < n {
64 let blk: i64 = i / B; let pos: i64 = i % B
65 if (blk % 2) == 0 { W[i] = sgn * (65536 - pos * 16384) }
66 else { W[i] = sgn * (1280 + pos * 256) }
67 i = i + 1
68 }
69 return 0
70}
71
72// ---- the existing multi-head attention forward (copied verbatim from nx_nofloat_multihead_gate) ----
73func mha_fwd(tape: *i64, vals: *i64, st: *i64, W: *i64, X: *i64, T: i64, dm: i64, hd: i64, H: i64, scale: i64, leaves: *i64) -> i64 {
74 let Wq: *i64 = W[0] as *i64; let Wk: *i64 = W[1] as *i64; let Wv: *i64 = W[2] as *i64; let Wo: *i64 = W[3] as *i64
75 st[0] = 0; st[1] = 0
76 let nX: i64 = nfa_leaf(tape, vals, st, T, dm, X, 0)
77 let nWq: i64 = nfa_leaf(tape, vals, st, dm, dm, Wq, 0)
78 let nWk: i64 = nfa_leaf(tape, vals, st, dm, dm, Wk, 0)
79 let nWv: i64 = nfa_leaf(tape, vals, st, dm, dm, Wv, 0)
80 let nWo: i64 = nfa_leaf(tape, vals, st, dm, dm, Wo, 0)
81 let nXn: i64 = nfa_rmsnorm_rows(tape, vals, st, nX)
82 let nQ: i64 = nfa_matmul(tape, vals, st, nXn, nWq)
83 let nK: i64 = nfa_matmul(tape, vals, st, nXn, nWk)
84 let nV: i64 = nfa_matmul(tape, vals, st, nXn, nWv)
85 var Oacc: i64 = 0 - 1
86 var hh: i64 = 0
87 while hh < H {
88 let nQh: i64 = nfa_slice_cols(tape, vals, st, nQ, hh * hd, hd)
89 let nKh: i64 = nfa_slice_cols(tape, vals, st, nK, hh * hd, hd)
90 let nVh: i64 = nfa_slice_cols(tape, vals, st, nV, hh * hd, hd)
91 let nQr: i64 = nfa_rope(tape, vals, st, nQh)
92 let nKr: i64 = nfa_rope(tape, vals, st, nKh)
93 let nS: i64 = nfa_matmul_nt(tape, vals, st, nQr, nKr)
94 let nSs: i64 = nfa_cmul(tape, vals, st, nS, scale)
95 let nA: i64 = nfa_softmax_rows(tape, vals, st, nSs, 1)
96 let nOh: i64 = nfa_matmul(tape, vals, st, nA, nVh)
97 if Oacc < 0 { Oacc = nOh } else { Oacc = nfa_concat_cols(tape, vals, st, Oacc, nOh) }
98 hh = hh + 1
99 }
100 let nOp: i64 = nfa_matmul(tape, vals, st, Oacc, nWo)
101 let nOut: i64 = nfa_vadd(tape, vals, st, nX, nOp)
102 leaves[0] = nWq; leaves[1] = nWk; leaves[2] = nWv; leaves[3] = nWo
103 return nOut
104}
105func mha_out(tape: *i64, vals: *i64, st: *i64, W: *i64, X: *i64, T: i64, dm: i64, hd: i64, H: i64, scale: i64, outv: *i64) -> i64 {
106 let lv: *i64 = sys_mmap(4 * 8) as *i64
107 let nOut: i64 = mha_fwd(tape, vals, st, W, X, T, dm, hd, H, scale, lv)
108 var i: i64 = 0
109 while i < T * dm { outv[i] = nfa_val(tape, vals, nOut, i); i = i + 1 }
110 return 0
111}
112func l1err(a: *i64, b: *i64, n: i64) -> i64 { var e: i64 = 0; var i: i64 = 0; while i < n { e = e + bb_abs(a[i] - b[i]); i = i + 1 } return e }
113
114func main() -> i64 {
115 bb_puts("=== BLOCK-FLOAT weights through REAL MULTI-HEAD attention (existing mha_fwd, unchanged) ===\n" as *u8)
116 let tape: *i64 = sys_mmap(1024 * 7 * 8) as *i64
117 let vals: *i64 = sys_mmap(16384 * 8) as *i64
118 let st: *i64 = sys_mmap(2 * 8) as *i64
119 let T: i64 = 2; let dm: i64 = 4; let H: i64 = 2; let hd: i64 = 2; let scale: i64 = 46341; let MB: i64 = 4; let B: i64 = 2
120
121 let X: *i64 = sys_mmap(T * dm * 8) as *i64
122 X[0] = 32768; X[1] = 0 - 16384; X[2] = 49152; X[3] = 24576; X[4] = 0 - 32768; X[5] = 65536; X[6] = 16384; X[7] = 0 - 8192
123 let Wq: *i64 = sys_mmap(dm * dm * 8) as *i64; fill_mixed(Wq, dm * dm, B, 1)
124 let Wk: *i64 = sys_mmap(dm * dm * 8) as *i64; fill_mixed(Wk, dm * dm, B, 0 - 1)
125 let Wv: *i64 = sys_mmap(dm * dm * 8) as *i64; fill_mixed(Wv, dm * dm, B, 1)
126 let Wo: *i64 = sys_mmap(dm * dm * 8) as *i64; fill_mixed(Wo, dm * dm, B, 0 - 1)
127
128 let qq: *i64 = sys_mmap(dm*dm*8) as *i64; bf_quant_mat(Wq, dm, dm, B, MB, qq)
129 let qk: *i64 = sys_mmap(dm*dm*8) as *i64; bf_quant_mat(Wk, dm, dm, B, MB, qk)
130 let qv: *i64 = sys_mmap(dm*dm*8) as *i64; bf_quant_mat(Wv, dm, dm, B, MB, qv)
131 let qo: *i64 = sys_mmap(dm*dm*8) as *i64; bf_quant_mat(Wo, dm, dm, B, MB, qo)
132 let pq: *i64 = sys_mmap(dm*dm*8) as *i64; bf_quant_pt(Wq, dm*dm, MB, pq)
133 let pk: *i64 = sys_mmap(dm*dm*8) as *i64; bf_quant_pt(Wk, dm*dm, MB, pk)
134 let pv: *i64 = sys_mmap(dm*dm*8) as *i64; bf_quant_pt(Wv, dm*dm, MB, pv)
135 let po: *i64 = sys_mmap(dm*dm*8) as *i64; bf_quant_pt(Wo, dm*dm, MB, po)
136
137 let Wfull: *i64 = sys_mmap(4*8) as *i64; Wfull[0] = Wq as i64; Wfull[1] = Wk as i64; Wfull[2] = Wv as i64; Wfull[3] = Wo as i64
138 let Wbf: *i64 = sys_mmap(4*8) as *i64; Wbf[0] = qq as i64; Wbf[1] = qk as i64; Wbf[2] = qv as i64; Wbf[3] = qo as i64
139 let Wpt: *i64 = sys_mmap(4*8) as *i64; Wpt[0] = pq as i64; Wpt[1] = pk as i64; Wpt[2] = pv as i64; Wpt[3] = po as i64
140
141 let Yr: *i64 = sys_mmap(T*dm*8) as *i64; mha_out(tape, vals, st, Wfull, X, T, dm, hd, H, scale, Yr)
142 let Yb: *i64 = sys_mmap(T*dm*8) as *i64; mha_out(tape, vals, st, Wbf, X, T, dm, hd, H, scale, Yb)
143 let Yp: *i64 = sys_mmap(T*dm*8) as *i64; mha_out(tape, vals, st, Wpt, X, T, dm, hd, H, scale, Yp)
144 let err_bf: i64 = l1err(Yb, Yr, T*dm)
145 let err_pt: i64 = l1err(Yp, Yr, T*dm)
146
147 bb_puts(" "); bb_pn(H); bb_puts("-head MHA out error (Q16 L1): BLOCK-FLOAT="); bb_pn(err_bf); bb_puts(" PER-TENSOR="); bb_pn(err_pt); bb_puts("\n" as *u8)
148
149 let Yb2: *i64 = sys_mmap(T*dm*8) as *i64; mha_out(tape, vals, st, Wbf, X, T, dm, hd, H, scale, Yb2)
150 var det: i64 = 1; var d: i64 = 0
151 while d < T*dm { if Yb2[d] != Yb[d] { det = 0 } d = d + 1 }
152 var nontriv: i64 = 0; var z: i64 = 0
153 while z < T*dm { if Yr[z] != 0 { nontriv = 1 } z = z + 1 }
154 var real_pt: i64 = 0; var p: i64 = 0
155 while p < T*dm { if Yp[p] != Yr[p] { real_pt = 1 } p = p + 1 }
156
157 var pass: i64 = 0; var total: i64 = 0
158 total = total + 1; pass = pass + bb_chk("T1 real MULTI-HEAD attention RUNS with block-float weights" as *u8, nontriv)
159 var t2: i64 = 0; if err_bf < err_pt { t2 = 1 }
160 total = total + 1; pass = pass + bb_chk("T2 EXCEED: block-float MHA error < per-tensor (dynamic-range win)" as *u8, t2)
161 total = total + 1; pass = pass + bb_chk("T3 DETERMINISTIC: block-float MHA twice == bit-identical" as *u8, det)
162 total = total + 1; pass = pass + bb_chk("T4 per-tensor is LOSSY (output != full-precision) -- win over real info loss" as *u8, real_pt)
163
164 bb_puts("NX-BLOCKFLOAT-MHA-GATE "); bb_pn(pass); bb_puts(" / "); bb_pn(total)
165 // MIGRATED onto nx_gate_verdict by nx_gate_dry_apply (D001, minimal form): every check
166 // row above is untouched, so the PASS/FAIL vector cannot change; only the hand-rolled
167 // verdict emission is replaced by the ONE shared base class. Proven by nx_gate_migrate verify.
168 let ctr__dry: *i64 = gv_ctr()
169 ctr__dry[0] = pass
170 ctr__dry[1] = total
171 let rc__dry: i64 = gv_verdict("BLOCKFLOAT-MHA-GATE" as *u8, ctr__dry, "block-float runs real multi-head attention: FP8 range + bit-exact determinism)" as *u8)
172 sys_exit(rc__dry)
173 return rc__dry
174}