nx_f32_qwen_encoder.nx source

1// nx_f32_qwen_encoder.nx -- stack the sovereign Qwen layer into the FULL encoder forward. 2// 3// sd-server -> Nishi migration: the text encoder is n_layers of nx_f32_qwen_layer applied in sequence 4// (Qwen3-4B = 36 layers). This drives that stack (ping-pong buffers). Output = the final hidden states = 5// the real prompt embeddings that feed the DiT (cap_embedder) AND, exposed as chat, the companion -- the 6// shared-LLM reuse, sovereign. This is the encoder analogue of nx_f32_rectflow_denoise (which stacks the 7// DiT block). Real use passes per-layer weights; this gate loops one weight set to prove the stack mechanism. 8// license_tier: ORIGINAL 9import "nx_syscalls.nx" 10import "nx_f32.nx" 11import "nx_f32_div.nx" 12import "nx_f32_cvt.nx" 13import "nx_f32_qwen_layer.nx" 14const K_MAGIC_1000000: i64 = 1000000 15 16func nx_f32_qwen_encoder(x: *i64, n_tokens: i64, hidden: i64, n_q_heads: i64, n_kv_heads: i64, head_dim: i64, ffn_dim: i64, 17 an_g: *i64, Wq: *i64, Wk: *i64, Wv: *i64, Wo: *i64, 18 fn_g: *i64, Wg: *i64, Wu: *i64, Wd: *i64, eps: i64, n_layers: i64, out: *i64) -> i64 { 19 let nel: i64 = n_tokens * hidden 20 let a: *i64 = sys_mmap(nel * 8) as *i64 21 let b: *i64 = sys_mmap(nel * 8) as *i64 22 var i: i64 = 0 23 while i < nel { a[i] = x[i]; i = i + 1 } 24 var cur: *i64 = a 25 var nxt: *i64 = b 26 var L: i64 = 0 27 while L < n_layers { 28 nx_f32_qwen_layer(cur, n_tokens, hidden, n_q_heads, n_kv_heads, head_dim, ffn_dim, an_g, Wq, Wk, Wv, Wo, fn_g, Wg, Wu, Wd, eps, nxt) 29 let tmp: *i64 = cur 30 cur = nxt 31 nxt = tmp 32 L = L + 1 33 } 34 i = 0 35 while i < nel { out[i] = cur[i]; i = i + 1 } 36 return 0 37} 38 39// nx_f32_qwen_encoder_layers -- the FULL-SCALE ASSEMBLY entry point: PER-LAYER weights. 40// 41// WHY THIS EXISTS: nx_f32_qwen_encoder above takes ONE weight set and passes the SAME pointers on 42// every iteration (see the loop: an_g, Wq, ... are loop-invariant). That proves the stack MECHANISM, 43// and its own header says so -- "this gate loops one weight set to prove the stack mechanism" -- but 44// it structurally cannot run a real model, where all 36 blocks have DIFFERENT weights. A stack that 45// reuses layer 0's weights 36 times is not a 36-layer forward; it is layer 0 applied 36 times. 46// 47// Each *_l argument is an ARRAY OF n_layers POINTERS (stored as i64), so layer L uses its own tensors. 48// ADDITIVE by construction (rule 19): the original single-weight-set function is untouched, so its 49// inline gate and every existing caller keep working. Nothing is broken to add this. 50// 51// Real dims for Z-Image's Qwen3-4B text encoder, MEASURED off the production GGUF by 52// nx_zimage_gguf_arch (not assumed): n_layers=36, hidden=2560, n_q_heads=32, n_kv_heads=8, 53// head_dim=128, ffn_dim=9728. 54func nx_f32_qwen_encoder_layers(x: *i64, n_tokens: i64, hidden: i64, n_q_heads: i64, n_kv_heads: i64, head_dim: i64, ffn_dim: i64, 55 an_g_l: *i64, Wq_l: *i64, Wk_l: *i64, Wv_l: *i64, Wo_l: *i64, 56 fn_g_l: *i64, Wg_l: *i64, Wu_l: *i64, Wd_l: *i64, eps: i64, n_layers: i64, out: *i64) -> i64 { 57 let nel: i64 = n_tokens * hidden 58 let a: *i64 = sys_mmap(nel * 8) as *i64 59 let b: *i64 = sys_mmap(nel * 8) as *i64 60 var i: i64 = 0 61 while i < nel { a[i] = x[i]; i = i + 1 } 62 var cur: *i64 = a 63 var nxt: *i64 = b 64 var L: i64 = 0 65 while L < n_layers { 66 nx_f32_qwen_layer(cur, n_tokens, hidden, n_q_heads, n_kv_heads, head_dim, ffn_dim, 67 an_g_l[L] as *i64, Wq_l[L] as *i64, Wk_l[L] as *i64, Wv_l[L] as *i64, Wo_l[L] as *i64, 68 fn_g_l[L] as *i64, Wg_l[L] as *i64, Wu_l[L] as *i64, Wd_l[L] as *i64, eps, nxt) 69 let tmp: *i64 = cur 70 cur = nxt 71 nxt = tmp 72 L = L + 1 73 } 74 i = 0 75 while i < nel { out[i] = cur[i]; i = i + 1 } 76 return 0 77} 78 79// ===== Self-test (inline gate) ==================================== 80// Wo=0, Wd=0 -> each layer is identity -> out == x through ALL layers (bit-exact). 81// nonzero -> the stack transforms x AND is deterministic (two runs identical). 82func main() -> i64 { 83 let n_tokens: i64 = 2 84 let hidden: i64 = 4 85 let nqh: i64 = 2 86 let nkvh: i64 = 1 87 let hd: i64 = 2 88 let ffn_dim: i64 = 4 89 let q_dim: i64 = nqh * hd 90 let x: *i64 = sys_mmap(n_tokens * hidden * 8) as *i64 91 let o1: *i64 = sys_mmap(n_tokens * hidden * 8) as *i64 92 let o2: *i64 = sys_mmap(n_tokens * hidden * 8) as *i64 93 let ang: *i64 = sys_mmap(hidden * 8) as *i64 94 let fng: *i64 = sys_mmap(hidden * 8) as *i64 95 let Wq: *i64 = sys_mmap(q_dim * hidden * 8) as *i64 96 let Wk: *i64 = sys_mmap(hd * hidden * 8) as *i64 97 let Wv: *i64 = sys_mmap(hd * hidden * 8) as *i64 98 let Wo: *i64 = sys_mmap(hidden * q_dim * 8) as *i64 99 let Wg: *i64 = sys_mmap(ffn_dim * hidden * 8) as *i64 100 let Wu: *i64 = sys_mmap(ffn_dim * hidden * 8) as *i64 101 let Wd: *i64 = sys_mmap(hidden * ffn_dim * 8) as *i64 102 let one: i64 = nx_i32_to_f32(1) 103 let p1: i64 = nx_f32_div(one, nx_i32_to_f32(10)) 104 let eps: i64 = nx_f32_div(one, nx_i32_to_f32(K_MAGIC_1000000)) 105 106 var i: i64 = 0 107 while i < hidden { ang[i] = one; fng[i] = one; i = i + 1 } 108 i = 0 109 while i < q_dim * hidden { Wq[i] = p1; i = i + 1 } 110 i = 0 111 while i < hd * hidden { Wk[i] = p1; Wv[i] = p1; i = i + 1 } 112 i = 0 113 while i < ffn_dim * hidden { Wg[i] = p1; Wu[i] = p1; i = i + 1 } 114 i = 0 115 while i < n_tokens * hidden { x[i] = nx_i32_to_f32(i + 1); i = i + 1 } 116 117 // (a) identity stack (3 layers) 118 i = 0 119 while i < hidden * q_dim { Wo[i] = 0; i = i + 1 } 120 i = 0 121 while i < hidden * ffn_dim { Wd[i] = 0; i = i + 1 } 122 if nx_f32_qwen_encoder(x, n_tokens, hidden, nqh, nkvh, hd, ffn_dim, ang, Wq, Wk, Wv, Wo, fng, Wg, Wu, Wd, eps, 3, o1) != 0 { return 10 } 123 i = 0 124 while i < n_tokens * hidden { if o1[i] != x[i] { return 20 } i = i + 1 } 125 126 // (b) nonzero -> transforms + deterministic (3 layers, twice) 127 i = 0 128 while i < hidden * q_dim { Wo[i] = p1; i = i + 1 } 129 i = 0 130 while i < hidden * ffn_dim { Wd[i] = p1; i = i + 1 } 131 nx_f32_qwen_encoder(x, n_tokens, hidden, nqh, nkvh, hd, ffn_dim, ang, Wq, Wk, Wv, Wo, fng, Wg, Wu, Wd, eps, 3, o1) 132 nx_f32_qwen_encoder(x, n_tokens, hidden, nqh, nkvh, hd, ffn_dim, ang, Wq, Wk, Wv, Wo, fng, Wg, Wu, Wd, eps, 3, o2) 133 var diff: i64 = 0 134 i = 0 135 while i < n_tokens * hidden { 136 if o1[i] != o2[i] { return 30 } 137 if o1[i] != x[i] { diff = 1 } 138 i = i + 1 139 } 140 if diff == 0 { return 40 } 141 142 // ===== (c) PER-LAYER weights: the full-scale-assembly path ===================== 143 // c1 FAITHFULNESS -- same weight set bound to every layer must be BIT-EXACT equal to the original. 144 // c2 NON-VACUITY -- PER SLOT, because the realistic bug is ONE wrong index among nine. 145 // 146 // I got c2 wrong twice and the BITE TEST caught me both times; the reasoning is recorded so the 147 // next person does not repeat it: 148 // * naive tooth ([X,identity] vs all-X) does NOT kill `always read [0]` -- it still differs. 149 // * whole-function tooth ([X,identity] vs [X,X]) kills a TOTAL index failure, but MUTATION 150 // PROVED IT GREEN when only `Wd_l[L]` was broken to `Wd_l[0]`: the other slots still honoured 151 // L so the outputs differed anyway. A tooth is only as strong as its WEAKEST SLOT. 152 // Now each of the nine slots is probed independently; a wrong index in any one goes RED with a 153 // code that NAMES it (60+slot). Bite-verified: correct=0, Wd_l[0] mutation=68. 154 let LN: i64 = 2 155 let ang_l: *i64 = sys_mmap(LN * 8) as *i64 156 let fng_l: *i64 = sys_mmap(LN * 8) as *i64 157 let Wq_l: *i64 = sys_mmap(LN * 8) as *i64 158 let Wk_l: *i64 = sys_mmap(LN * 8) as *i64 159 let Wv_l: *i64 = sys_mmap(LN * 8) as *i64 160 let Wo_l: *i64 = sys_mmap(LN * 8) as *i64 161 let Wg_l: *i64 = sys_mmap(LN * 8) as *i64 162 let Wu_l: *i64 = sys_mmap(LN * 8) as *i64 163 let Wd_l: *i64 = sys_mmap(LN * 8) as *i64 164 165 var L2: i64 = 0 166 while L2 < LN { 167 ang_l[L2] = ang as i64; fng_l[L2] = fng as i64 168 Wq_l[L2] = Wq as i64; Wk_l[L2] = Wk as i64; Wv_l[L2] = Wv as i64; Wo_l[L2] = Wo as i64 169 Wg_l[L2] = Wg as i64; Wu_l[L2] = Wu as i64; Wd_l[L2] = Wd as i64 170 L2 = L2 + 1 171 } 172 nx_f32_qwen_encoder(x, n_tokens, hidden, nqh, nkvh, hd, ffn_dim, ang, Wq, Wk, Wv, Wo, fng, Wg, Wu, Wd, eps, LN, o1) 173 if nx_f32_qwen_encoder_layers(x, n_tokens, hidden, nqh, nkvh, hd, ffn_dim, 174 ang_l, Wq_l, Wk_l, Wv_l, Wo_l, fng_l, Wg_l, Wu_l, Wd_l, eps, LN, o2) != 0 { return 50 } 175 i = 0 176 while i < n_tokens * hidden { if o1[i] != o2[i] { return 51 } i = i + 1 } 177 178 let alt_g: *i64 = sys_mmap(hidden * 8) as *i64 179 let alt_Wq: *i64 = sys_mmap(q_dim * hidden * 8) as *i64 180 let alt_Wk: *i64 = sys_mmap(hd * hidden * 8) as *i64 181 let alt_Wo: *i64 = sys_mmap(hidden * q_dim * 8) as *i64 182 let alt_Wf: *i64 = sys_mmap(ffn_dim * hidden * 8) as *i64 183 let alt_Wd: *i64 = sys_mmap(hidden * ffn_dim * 8) as *i64 184 let p2: i64 = nx_f32_div(one, nx_i32_to_f32(5)) 185 i = 0 186 while i < hidden { alt_g[i] = p2; i = i + 1 } 187 i = 0 188 while i < q_dim * hidden { alt_Wq[i] = p2; i = i + 1 } 189 i = 0 190 while i < hd * hidden { alt_Wk[i] = p2; i = i + 1 } 191 i = 0 192 while i < hidden * q_dim { alt_Wo[i] = p2; i = i + 1 } 193 i = 0 194 while i < ffn_dim * hidden { alt_Wf[i] = p2; i = i + 1 } 195 i = 0 196 while i < hidden * ffn_dim { alt_Wd[i] = p2; i = i + 1 } 197 198 if nx_f32_qwen_encoder_layers(x, n_tokens, hidden, nqh, nkvh, hd, ffn_dim, 199 ang_l, Wq_l, Wk_l, Wv_l, Wo_l, fng_l, Wg_l, Wu_l, Wd_l, eps, LN, o1) != 0 { return 52 } 200 201 var slot: i64 = 0 202 while slot < 9 { 203 if slot == 0 { ang_l[1] = alt_g as i64 } 204 if slot == 1 { Wq_l[1] = alt_Wq as i64 } 205 if slot == 2 { Wk_l[1] = alt_Wk as i64 } 206 if slot == 3 { Wv_l[1] = alt_Wk as i64 } 207 if slot == 4 { Wo_l[1] = alt_Wo as i64 } 208 if slot == 5 { fng_l[1] = alt_g as i64 } 209 if slot == 6 { Wg_l[1] = alt_Wf as i64 } 210 if slot == 7 { Wu_l[1] = alt_Wf as i64 } 211 if slot == 8 { Wd_l[1] = alt_Wd as i64 } 212 213 if nx_f32_qwen_encoder_layers(x, n_tokens, hidden, nqh, nkvh, hd, ffn_dim, 214 ang_l, Wq_l, Wk_l, Wv_l, Wo_l, fng_l, Wg_l, Wu_l, Wd_l, eps, LN, o2) != 0 { return 54 } 215 var moved: i64 = 0 216 i = 0 217 while i < n_tokens * hidden { if o1[i] != o2[i] { moved = 1 } i = i + 1 } 218 if moved == 0 { return 60 + slot } 219 220 if slot == 0 { ang_l[1] = ang as i64 } 221 if slot == 1 { Wq_l[1] = Wq as i64 } 222 if slot == 2 { Wk_l[1] = Wk as i64 } 223 if slot == 3 { Wv_l[1] = Wv as i64 } 224 if slot == 4 { Wo_l[1] = Wo as i64 } 225 if slot == 5 { fng_l[1] = fng as i64 } 226 if slot == 6 { Wg_l[1] = Wg as i64 } 227 if slot == 7 { Wu_l[1] = Wu as i64 } 228 if slot == 8 { Wd_l[1] = Wd as i64 } 229 slot = slot + 1 230 } 231 232 return 0 233}