nx_trig_kat.nx source
↩ module page · 55 lines · 2806 B
1// nx_trig_kat.nx -- KAT for nx_f32_cos/nx_f32_sin at rope-realistic arguments. RoPE at position p rotates the
2// i=0 (highest-freq) dim by angle p -> cos(p)/sin(p) for p up to seq_len. If the range reduction is only good
3// near 0, varied multi-token sequences get corrupted rope on the fast dims (masked by repetition = the residual
4// bug after the QKV-bias fix). Self-consistency check: cos^2+sin^2 == 1.0 (0x3F800000) for ALL args; a deviation
5// at larger args localizes a range-reduction bug. Also checks cos(0)=1, sin(0)=0. expect_exit: 0
6import "nx_syscalls.nx"
7import "nx_tier.nx"
8import "nx_f32.nx"
9import "nx_f32_sincos.nx"
10
11func pw(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
12func ph(v: i64) -> i64 { let b: *u8=sys_mmap(8); var i: i64=0; while i<8 { let nib: i64=(v>>((7-i)*4))&0xF; if nib<10 {b[i]=(48+nib) as u8} else {b[i]=(87+nib) as u8} i=i+1 } sys_write(1,b,8); return 0 }
13
14// |c^2+s^2 - 1.0| small? compare raw sum against 1.0 band [0.996, 1.004] ~ [0x3F7EF9DB, 0x3F8083127].
15func near_one(sum: i64) -> i64 {
16 let s: i64 = sum & 0xFFFFFFFF
17 if s < 0x3F7EF9DB { return 0 }
18 if s > 0x3F810625 { return 0 }
19 return 1
20}
21
22func one(label: *u8, x: i64) -> i64 {
23 let c: i64 = nx_f32_cos(x)
24 let s: i64 = nx_f32_sin(x)
25 let cc: i64 = nx_f32_mul(c, c)
26 let ss: i64 = nx_f32_mul(s, s)
27 let sum: i64 = nx_f32_add(cc, ss)
28 pw(" x="); pw(label)
29 pw(" cos="); ph(c); pw(" sin="); ph(s); pw(" cos2+sin2="); ph(sum)
30 if near_one(sum)==1 { pw(" OK\n"); return 1 }
31 pw(" <== BROKEN (should be ~0x3F800000)\n"); return 0
32}
33
34func main() -> i64 {
35 pw("=== nx_f32 trig KAT (rope needs cos/sin accurate to seq_len radians) ===\n" as *u8)
36 var ok: i64 = 0
37 var tot: i64 = 0
38 // cos(0)=1, sin(0)=0 sanity
39 let c0: i64 = nx_f32_cos(0x00000000)
40 let s0: i64 = nx_f32_sin(0x00000000)
41 pw(" cos(0)="); ph(c0); pw(" (want 3F800000) sin(0)="); ph(s0); pw(" (want 00000000)\n")
42 // identity across args: 0, 1, pi, 2pi, 7, 14, 30, 50
43 tot=tot+1; ok=ok+one("0.0 " as *u8, 0x00000000)
44 tot=tot+1; ok=ok+one("1.0 " as *u8, 0x3F800000)
45 tot=tot+1; ok=ok+one("pi " as *u8, 0x40490FDB)
46 tot=tot+1; ok=ok+one("2pi " as *u8, 0x40C90FDB)
47 tot=tot+1; ok=ok+one("7.0 " as *u8, 0x40E00000)
48 tot=tot+1; ok=ok+one("14.0" as *u8, 0x41600000)
49 tot=tot+1; ok=ok+one("30.0" as *u8, 0x41F00000)
50 tot=tot+1; ok=ok+one("50.0" as *u8, 0x42480000)
51 pw("trig KAT ok="); ph(ok); pw(" / "); ph(tot); pw("\n")
52 if ok==tot { pw("VERDICT: trig ACCURATE to 50 rad -> NOT the residual bug (look at attention/cache next)\n" as *u8); sys_exit(0); return 0 }
53 pw("VERDICT: trig BREAKS at large args -> range-reduction IS the residual rope bug. FIX nx_f32_sincos.\n" as *u8)
54 sys_exit(0); return 0
55}