nx_gen_adaln_gate_verify.nx source
↩ module page · 144 lines · 6078 B
1// nx_gen_adaln_gate_verify.nx -- SOVEREIGN adaLN gated residual join, verified vs the oracle.
2//
3// out == inner * tanh(gate) + residual
4//
5// where `gate` is one chunk of a block's adaLN modulation vector, broadcast across all tokens.
6// This is the join at the end of each half of a modulated DiT block.
7//
8// Usage: nx_gen_adaln_gate_verify <model> <inner> <adaln> <n_chunks> <gate_idx> <residual> <out> [rows]
9//
10// The three tensors are NAMED, not derived from a scope with fixed op names. The first version
11// hardcoded attn_norm2/blk_in/mid, which meant it could only ever check the ATTENTION half of a
12// block -- the identical join at the end of the FFN half needed a second organ. Naming them makes
13// one organ cover both halves, and any architecture that gates a residual this way.
14//
15// The CHUNK COUNT AND INDEX ARE ARGUMENTS because they are architecture facts, not universal
16// ones: Z-Image packs 4 chunks [scale_msa|gate_msa|scale_mlp|gate_mlp], Flux-style blocks pack
17// 6 [shift|scale|gate] x2. Baking either in would make this organ silently verify the wrong
18// slice the first time it met another model -- and a wrong slice still produces plausible
19// numbers, so it would not announce itself.
20//
21// This organ needs NO model weights: every input is an oracle fixture. That makes it the
22// cheapest possible check of the chunk ORDER, an assumption every downstream organ inherits.
23// license_tier: ORIGINAL
24
25import "nx_syscalls.nx"
26import "nx_le.nx"
27import "nx_f32.nx"
28import "nx_f32_div.nx"
29import "nx_f32_cvt.nx"
30import "nx_f32_exp.nx"
31import "nx_f32_activations.nx"
32import "nx_strconv.nx"
33import "nx_genfix.nx"
34import "nx_genver.nx"
35
36const K_ROWS_CHECK: i64 = 4
37
38func zg_strlen(s: *u8) -> i64 {
39 var n: i64 = 0
40 while s[n] != (0 as u8) { n = n + 1 }
41 return n
42}
43
44func main(argc: i64, argv: *i64) -> i64 {
45 if argc < 8 {
46 nx_genver_emit("usage_model_inner_adaln_nchunks_gateidx_residual_out" as *u8, argc)
47 return 2
48 }
49 let model: *u8 = argv[1] as *u8
50 let errp: *i64 = sys_mmap(32) as *i64
51 errp[0] = 0
52 let n_chunks: i64 = nx_strconv_parse_i64(argv[4] as *u8, errp)
53 if errp[0] != 0 { nx_genver_emit("bad_n_chunks" as *u8, 1); return 3 }
54 let gate_idx: i64 = nx_strconv_parse_i64(argv[5] as *u8, errp)
55 if errp[0] != 0 { nx_genver_emit("bad_gate_idx" as *u8, 1); return 4 }
56 if n_chunks <= 0 { nx_genver_emit("bad_n_chunks" as *u8, n_chunks); return 5 }
57 if gate_idx < 0 { nx_genver_emit("bad_gate_idx" as *u8, gate_idx); return 6 }
58 if gate_idx >= n_chunks { nx_genver_emit("gate_idx_out_of_range" as *u8, gate_idx); return 7 }
59
60 let n_in: *u8 = argv[2] as *u8
61 let n_am: *u8 = argv[3] as *u8
62 let n_rs: *u8 = argv[6] as *u8
63 let n_ot: *u8 = argv[7] as *u8
64 let l_in: i64 = zg_strlen(n_in)
65 let l_am: i64 = zg_strlen(n_am)
66 let l_rs: i64 = zg_strlen(n_rs)
67 let l_ot: i64 = zg_strlen(n_ot)
68
69 let ne_am: *i64 = sys_mmap(64) as *i64
70 let ne_in: *i64 = sys_mmap(64) as *i64
71 let ne_rs: *i64 = sys_mmap(64) as *i64
72 let ne_ot: *i64 = sys_mmap(64) as *i64
73 let c_am: i64 = nx_genfix_dims(model, n_am, l_am, ne_am)
74 let c_in: i64 = nx_genfix_dims(model, n_in, l_in, ne_in)
75 let c_rs: i64 = nx_genfix_dims(model, n_rs, l_rs, ne_rs)
76 let c_ot: i64 = nx_genfix_dims(model, n_ot, l_ot, ne_ot)
77 if c_am < 0 { nx_genver_emit("missing_adaln_m" as *u8, 1); return 30 }
78 if c_in < 0 { nx_genver_emit("missing_attn_norm2" as *u8, 1); return 31 }
79 if c_rs < 0 { nx_genver_emit("missing_blk_in" as *u8, 1); return 32 }
80 if c_ot < 0 { nx_genver_emit("missing_mid" as *u8, 1); return 33 }
81
82 let d: i64 = ne_in[0]
83 let n_tok: i64 = ne_in[1]
84 if ne_rs[0] != d { nx_genver_emit("shape_mismatch_residual" as *u8, ne_rs[0]); return 34 }
85 if ne_ot[0] != d { nx_genver_emit("shape_mismatch_out" as *u8, ne_ot[0]); return 35 }
86
87 // The modulation vector must be exactly n_chunks wide in d. If it is not, the caller has
88 // named the wrong chunk count for this architecture -- catch it here rather than let a
89 // misaligned slice produce plausible-but-wrong numbers.
90 if c_am != n_chunks * d {
91 nx_genver_emit("adaln_width_not_nchunks_times_d" as *u8, c_am)
92 nx_genver_emit("expected" as *u8, n_chunks * d)
93 return 36
94 }
95
96 let am: *u8 = nx_genfix_load(model, n_am, l_am, c_am)
97 if (am as i64) == 0 { nx_genver_emit("load_failed_adaln_m" as *u8, 1); return 40 }
98 let inr: *u8 = nx_genfix_load(model, n_in, l_in, c_in)
99 if (inr as i64) == 0 { nx_genver_emit("load_failed_attn_norm2" as *u8, 1); return 41 }
100 let rs: *u8 = nx_genfix_load(model, n_rs, l_rs, c_rs)
101 if (rs as i64) == 0 { nx_genver_emit("load_failed_blk_in" as *u8, 1); return 42 }
102 let ot: *u8 = nx_genfix_load(model, n_ot, l_ot, c_ot)
103 if (ot as i64) == 0 { nx_genver_emit("load_failed_mid" as *u8, 1); return 43 }
104
105 var rows: i64 = K_ROWS_CHECK
106 if n_tok < rows { rows = n_tok }
107
108 nx_genver_emit("dim" as *u8, d)
109 nx_genver_emit("tokens_total" as *u8, n_tok)
110 nx_genver_emit("tokens_checked" as *u8, rows)
111 nx_genver_emit("gate_chunk" as *u8, gate_idx)
112
113 // tanh(gate) depends only on the channel -- hoist it out of the token loop.
114 let gate_base: i64 = gate_idx * d
115 let tg: *i64 = sys_mmap(d * 8) as *i64
116 var g: i64 = 0
117 while g < d {
118 tg[g] = nx_f32_tanh(nx_le_read_u32(am, (gate_base + g) * 4))
119 g = g + 1
120 }
121
122 let tol: *i64 = sys_mmap(64) as *i64
123 nx_genver_tols(tol)
124 let c: *i64 = sys_mmap(128) as *i64
125 nx_genver_init(c, 4)
126
127 var t: i64 = 0
128 while t < rows {
129 var i: i64 = 0
130 while i < d {
131 let flat: i64 = t * d + i
132 let off: i64 = flat * 4
133 let a: i64 = nx_le_read_u32(inr, off)
134 let r: i64 = nx_le_read_u32(rs, off)
135 let want: i64 = nx_le_read_u32(ot, off)
136 let got: i64 = __f32_add(__f32_mul(a, tg[i]), r)
137 nx_genver_tally(c, tol, got, want, flat)
138 i = i + 1
139 }
140 t = t + 1
141 }
142
143 return nx_genver_report(c)
144}