nx_gen_swiglu_verify.nx source
↩ module page · 106 lines · 4482 B
1// nx_gen_swiglu_verify.nx -- SOVEREIGN SwiGLU feed-forward activation, verified vs the oracle.
2//
3// ffn_act == silu(ffn_w1) * ffn_w3 elementwise
4//
5// SwiGLU is the FFN gate in Z-Image, Flux, SD3/MMDiT and every Llama-class text encoder, so this
6// organ is named for the OP and takes the model as an argument. Nothing about it is Z-Image.
7//
8// Usage: nx_gen_swiglu_verify <model_id> <tap_scope>
9// model_id fixture directory under the gen fixture root, e.g. zimage_turbo_2602_q8
10// tap_scope block scope the oracle tap emitted, e.g. layers.0
11//
12// Every SHAPE is read from the oracle's manifest.tsv. Pointing this at a different model or
13// block needs no rebuild, and an organ that cannot silently keep a stale dimension cannot
14// silently verify the wrong model.
15//
16// This is the cheapest real differential test in the lane -- pure elementwise, no weight
17// matrices, no GGUF. Its job is to prove the harness end to end so the expensive organs above
18// it fail for their own reasons rather than the harness's.
19// license_tier: ORIGINAL
20
21import "nx_syscalls.nx"
22import "nx_le.nx"
23import "nx_f32.nx"
24import "nx_f32_div.nx"
25import "nx_f32_cvt.nx"
26import "nx_f32_exp.nx"
27import "nx_f32_activations.nx"
28import "nx_strconv.nx"
29import "nx_genfix.nx"
30import "nx_genver.nx"
31
32// Rows to verify. The op is elementwise, so any subset of tokens is a sound check and the
33// full 768 would only cost time. Raised here, never silently: the count is reported.
34const K_ROWS_CHECK: i64 = 4
35
36func main(argc: i64, argv: *i64) -> i64 {
37 if argc < 3 {
38 nx_genver_emit("usage_model_and_scope_required" as *u8, argc)
39 return 2
40 }
41 let model: *u8 = argv[1] as *u8
42 let scope: *u8 = argv[2] as *u8
43
44 let n_w1: *u8 = sys_mmap(256)
45 let n_w3: *u8 = sys_mmap(256)
46 let n_ac: *u8 = sys_mmap(256)
47 let l_w1: i64 = nx_genfix_name(n_w1, scope, "ffn_w1" as *u8, 6)
48 let l_w3: i64 = nx_genfix_name(n_w3, scope, "ffn_w3" as *u8, 6)
49 let l_ac: i64 = nx_genfix_name(n_ac, scope, "ffn_act" as *u8, 7)
50
51 let ne1: *i64 = sys_mmap(64) as *i64
52 let ne3: *i64 = sys_mmap(64) as *i64
53 let nea: *i64 = sys_mmap(64) as *i64
54 let c_w1: i64 = nx_genfix_dims(model, n_w1, l_w1, ne1)
55 let c_w3: i64 = nx_genfix_dims(model, n_w3, l_w3, ne3)
56 let c_ac: i64 = nx_genfix_dims(model, n_ac, l_ac, nea)
57
58 // A missing tap is a different fault from a numeric mismatch; refuse rather than compare.
59 if c_w1 < 0 { nx_genver_emit("missing_ffn_w1" as *u8, 1); return 30 }
60 if c_w3 < 0 { nx_genver_emit("missing_ffn_w3" as *u8, 1); return 31 }
61 if c_ac < 0 { nx_genver_emit("missing_ffn_act" as *u8, 1); return 32 }
62
63 let fd_dim: i64 = ne1[0] // contiguous dim: ffn hidden width
64 let n_tok: i64 = ne1[1] // token count
65 if ne3[0] != fd_dim { nx_genver_emit("shape_mismatch_w3_d0" as *u8, ne3[0]); return 33 }
66 if nea[0] != fd_dim { nx_genver_emit("shape_mismatch_act_d0" as *u8, nea[0]); return 34 }
67 if ne3[1] != n_tok { nx_genver_emit("shape_mismatch_w3_d1" as *u8, ne3[1]); return 35 }
68 if nea[1] != n_tok { nx_genver_emit("shape_mismatch_act_d1" as *u8, nea[1]); return 36 }
69
70 let w1: *u8 = nx_genfix_load(model, n_w1, l_w1, c_w1)
71 if (w1 as i64) == 0 { nx_genver_emit("load_failed_ffn_w1" as *u8, 1); return 40 }
72 let w3: *u8 = nx_genfix_load(model, n_w3, l_w3, c_w3)
73 if (w3 as i64) == 0 { nx_genver_emit("load_failed_ffn_w3" as *u8, 1); return 41 }
74 let ag: *u8 = nx_genfix_load(model, n_ac, l_ac, c_ac)
75 if (ag as i64) == 0 { nx_genver_emit("load_failed_ffn_act" as *u8, 1); return 42 }
76
77 var rows: i64 = K_ROWS_CHECK
78 if n_tok < rows { rows = n_tok }
79
80 nx_genver_emit("dim" as *u8, fd_dim)
81 nx_genver_emit("tokens_total" as *u8, n_tok)
82 nx_genver_emit("tokens_checked" as *u8, rows)
83
84 let tol: *i64 = sys_mmap(64) as *i64
85 nx_genver_tols(tol)
86 let c: *i64 = sys_mmap(128) as *i64
87 nx_genver_init(c, 4)
88
89 var t: i64 = 0
90 while t < rows {
91 var i: i64 = 0
92 while i < fd_dim {
93 let flat: i64 = t * fd_dim + i
94 let off: i64 = flat * 4
95 let a: i64 = nx_le_read_u32(w1, off)
96 let b: i64 = nx_le_read_u32(w3, off)
97 let want: i64 = nx_le_read_u32(ag, off)
98 let got: i64 = nx_f32_mul(nx_f32_silu(a), b)
99 nx_genver_tally(c, tol, got, want, flat)
100 i = i + 1
101 }
102 t = t + 1
103 }
104
105 return nx_genver_report(c)
106}