nx_f32_qwen_layer.nx source

1// nx_f32_qwen_layer.nx -- a FULL sovereign f32 Qwen transformer layer (the Z-Image text-encoder + companion 2// chat model's repeating unit), assembled from the gated bricks. 3// 4// sd-server -> Nishi migration: the pre-norm GQA transformer layer (Llama/Qwen family). Per the real GGUF: 5// RMSNorm -> [Wq/Wk/Wv] -> GQA attention (32 q / 8 kv heads) -> Wo -> residual -> RMSNorm -> SwiGLU 6// (ffn_gate/ffn_up -> SiLU-gate -> ffn_down) -> residual. Composes nx_f32_rmsnorm + nx_f32_gqa_attention + 7// nx_f32_silu + inline linears. Stack this x n_layers (a loop) = the full encoder forward -> real prompt 8// embeddings. Same layer serves image-prompt-encoding AND companion chat (the shared-LLM reuse). 9// 10// x,out: flat *i64 f32 bits [n_tokens, hidden]. W{q}: [q_dim,hidden], W{k,v}: [kv_dim,hidden], Wo: [hidden,q_dim]; 11// ffn Wg/Wu: [ffn_dim,hidden], Wd: [hidden,ffn_dim]; an_g/fn_g: [hidden]. Scratch is internal (fork-per-use). 12// license_tier: ORIGINAL 13import "nx_syscalls.nx" 14import "nx_f32.nx" 15import "nx_f32_div.nx" 16import "nx_f32_cvt.nx" 17import "nx_f32_rmsnorm.nx" 18import "nx_f32_gqa_attention.nx" 19import "nx_f32_activations.nx" 20const K_MAGIC_1000000: i64 = 1000000 21 22func qw_linear(inp: *i64, n_tokens: i64, in_dim: i64, W: *i64, out_dim: i64, out: *i64) -> i64 { 23 var t: i64 = 0 24 while t < n_tokens { 25 var o: i64 = 0 26 while o < out_dim { 27 var acc: i64 = 0 28 var i: i64 = 0 29 while i < in_dim { acc = nx_f32_add(acc, nx_f32_mul(inp[t * in_dim + i], W[o * in_dim + i])); i = i + 1 } 30 out[t * out_dim + o] = acc 31 o = o + 1 32 } 33 t = t + 1 34 } 35 return 0 36} 37 38func qw_swiglu(inp: *i64, n_tokens: i64, hidden: i64, ffn_dim: i64, Wg: *i64, Wu: *i64, Wd: *i64, out: *i64, hb: *i64) -> i64 { 39 var t: i64 = 0 40 while t < n_tokens { 41 var f: i64 = 0 42 while f < ffn_dim { 43 var ag: i64 = 0 44 var au: i64 = 0 45 var i: i64 = 0 46 while i < hidden { ag = nx_f32_add(ag, nx_f32_mul(inp[t * hidden + i], Wg[f * hidden + i])); au = nx_f32_add(au, nx_f32_mul(inp[t * hidden + i], Wu[f * hidden + i])); i = i + 1 } 47 hb[f] = nx_f32_mul(nx_f32_silu(ag), au) 48 f = f + 1 49 } 50 var d: i64 = 0 51 while d < hidden { 52 var o: i64 = 0 53 var ff: i64 = 0 54 while ff < ffn_dim { o = nx_f32_add(o, nx_f32_mul(hb[ff], Wd[d * ffn_dim + ff])); ff = ff + 1 } 55 out[t * hidden + d] = o 56 d = d + 1 57 } 58 t = t + 1 59 } 60 return 0 61} 62 63// PER-HEAD RMSNorm over head_dim, applied in place to a [n_tokens, n_heads*head_dim] buffer. 64// This is Qwen3's QK-norm. Gain is an F32 vector of head_dim (MEASURED off the production GGUF: 65// blk.N.attn_q_norm.weight is 128x0 ggml_type=0 => 1-D, head_dim long, UNQUANTIZED). 66func qw_head_norm(buf: *i64, n_tokens: i64, n_heads: i64, head_dim: i64, gain: *i64, eps: i64) -> i64 { 67 if (gain as i64) == 0 { return 0 } // null gain => this norm is absent 68 let tmp: *i64 = sys_mmap(head_dim * 8) as *i64 69 var t: i64 = 0 70 while t < n_tokens { 71 var h: i64 = 0 72 while h < n_heads { 73 let bp: *i64 = ((buf as i64) + (t * n_heads * head_dim + h * head_dim) * 8) as *i64 74 nx_f32_rmsnorm(bp, gain, head_dim, eps, tmp) 75 var i2: i64 = 0 76 while i2 < head_dim { bp[i2] = tmp[i2]; i2 = i2 + 1 } 77 h = h + 1 78 } 79 t = t + 1 80 } 81 return 0 82} 83 84// nx_f32_qwen_layer_qk -- the layer WITH Qwen3 QK-norm. This is the real architecture. 85// 86// WHY (debt 1785449705): the 9-slot layer below omitted q_norm/k_norm entirely. Caught by ARITHMETIC, 87// not by a test: 36 layers x 9 slots = 324 but the GGUF's tensor_count is 398, and 36*11 + token_embd 88// + output_norm = 398 exactly. The two unaccounted tensors per block are attn_q_norm/attn_k_norm, and 89// they resolve in ALL 36 layers. Without them the forward pass yields finite, plausible, WRONG 90// embeddings -- nothing errors, so the damage surfaces as unexplained quality loss much later. 91// 92// Passing q_ng=0 and k_ng=0 reproduces the original pre-QK-norm behaviour EXACTLY, which is how 93// nx_f32_qwen_layer below stays contract-identical (rule 19) instead of being duplicated. 94func nx_f32_qwen_layer_qk(x: *i64, n_tokens: i64, hidden: i64, n_q_heads: i64, n_kv_heads: i64, head_dim: i64, ffn_dim: i64, 95 an_g: *i64, Wq: *i64, Wk: *i64, Wv: *i64, Wo: *i64, 96 fn_g: *i64, Wg: *i64, Wu: *i64, Wd: *i64, 97 q_ng: *i64, k_ng: *i64, eps: i64, out: *i64) -> i64 { 98 let q_dim: i64 = n_q_heads * head_dim 99 let kv_dim: i64 = n_kv_heads * head_dim 100 let hn: *i64 = sys_mmap(n_tokens * hidden * 8) as *i64 101 let Qb: *i64 = sys_mmap(n_tokens * q_dim * 8) as *i64 102 let Kb: *i64 = sys_mmap(n_tokens * kv_dim * 8) as *i64 103 let Vb: *i64 = sys_mmap(n_tokens * kv_dim * 8) as *i64 104 let attn: *i64 = sys_mmap(n_tokens * q_dim * 8) as *i64 105 let ao: *i64 = sys_mmap(n_tokens * hidden * 8) as *i64 106 let x1: *i64 = sys_mmap(n_tokens * hidden * 8) as *i64 107 let hn2: *i64 = sys_mmap(n_tokens * hidden * 8) as *i64 108 let ffn: *i64 = sys_mmap(n_tokens * hidden * 8) as *i64 109 let hb: *i64 = sys_mmap(ffn_dim * 8) as *i64 110 let scale: i64 = nx_i32_to_f32(1) 111 112 // ---- attention sublayer ---- 113 var t: i64 = 0 114 while t < n_tokens { 115 nx_f32_rmsnorm(((x as i64) + t * hidden * 8) as *i64, an_g, hidden, eps, ((hn as i64) + t * hidden * 8) as *i64) 116 t = t + 1 117 } 118 qw_linear(hn, n_tokens, hidden, Wq, q_dim, Qb) 119 qw_linear(hn, n_tokens, hidden, Wk, kv_dim, Kb) 120 qw_linear(hn, n_tokens, hidden, Wv, kv_dim, Vb) 121 // QK-NORM: per-head RMSNorm on Q and K AFTER projection, BEFORE attention. Q uses n_q_heads, 122 // K uses n_kv_heads (GQA: 32 vs 8) -- using the wrong head count here would silently mis-stride. 123 qw_head_norm(Qb, n_tokens, n_q_heads, head_dim, q_ng, eps) 124 qw_head_norm(Kb, n_tokens, n_kv_heads, head_dim, k_ng, eps) 125 nx_f32_gqa_attention(Qb, Kb, Vb, n_tokens, n_q_heads, n_kv_heads, head_dim, scale, attn) 126 qw_linear(attn, n_tokens, q_dim, Wo, hidden, ao) 127 var i: i64 = 0 128 while i < n_tokens * hidden { x1[i] = nx_f32_add(x[i], ao[i]); i = i + 1 } 129 130 // ---- FFN sublayer ---- 131 t = 0 132 while t < n_tokens { 133 nx_f32_rmsnorm(((x1 as i64) + t * hidden * 8) as *i64, fn_g, hidden, eps, ((hn2 as i64) + t * hidden * 8) as *i64) 134 t = t + 1 135 } 136 qw_swiglu(hn2, n_tokens, hidden, ffn_dim, Wg, Wu, Wd, ffn, hb) 137 i = 0 138 while i < n_tokens * hidden { out[i] = nx_f32_add(x1[i], ffn[i]); i = i + 1 } 139 return 0 140} 141 142// nx_f32_qwen_layer -- ORIGINAL 9-slot contract, preserved EXACTLY by delegating with null QK gains. 143// Kept so every existing caller and gate is untouched (rule 19). New work should call _qk and pass 144// the real attn_q_norm/attn_k_norm weights; this signature cannot express them. 145func nx_f32_qwen_layer(x: *i64, n_tokens: i64, hidden: i64, n_q_heads: i64, n_kv_heads: i64, head_dim: i64, ffn_dim: i64, 146 an_g: *i64, Wq: *i64, Wk: *i64, Wv: *i64, Wo: *i64, 147 fn_g: *i64, Wg: *i64, Wu: *i64, Wd: *i64, eps: i64, out: *i64) -> i64 { 148 return nx_f32_qwen_layer_qk(x, n_tokens, hidden, n_q_heads, n_kv_heads, head_dim, ffn_dim, 149 an_g, Wq, Wk, Wv, Wo, fn_g, Wg, Wu, Wd, 0 as *i64, 0 as *i64, eps, out) 150} 151 152// ===== Self-test (inline gate) ==================================== 153// Wo == 0 AND Wd == 0 -> both residual branch outputs are 0 -> out == x (bit-exact), any other weights. 154// nonzero Wo,Wd -> the layer contributes (out != x somewhere). 155func main() -> i64 { 156 let n_tokens: i64 = 2 157 let hidden: i64 = 4 158 let nqh: i64 = 2 159 let nkvh: i64 = 1 160 let hd: i64 = 2 161 let ffn_dim: i64 = 4 162 let q_dim: i64 = nqh * hd 163 let kv_dim: i64 = nkvh * hd 164 let x: *i64 = sys_mmap(n_tokens * hidden * 8) as *i64 165 let out: *i64 = sys_mmap(n_tokens * hidden * 8) as *i64 166 let ang: *i64 = sys_mmap(hidden * 8) as *i64 167 let fng: *i64 = sys_mmap(hidden * 8) as *i64 168 let Wq: *i64 = sys_mmap(q_dim * hidden * 8) as *i64 169 let Wk: *i64 = sys_mmap(kv_dim * hidden * 8) as *i64 170 let Wv: *i64 = sys_mmap(kv_dim * hidden * 8) as *i64 171 let Wo: *i64 = sys_mmap(hidden * q_dim * 8) as *i64 172 let Wg: *i64 = sys_mmap(ffn_dim * hidden * 8) as *i64 173 let Wu: *i64 = sys_mmap(ffn_dim * hidden * 8) as *i64 174 let Wd: *i64 = sys_mmap(hidden * ffn_dim * 8) as *i64 175 let one: i64 = nx_i32_to_f32(1) 176 let p1: i64 = nx_f32_div(one, nx_i32_to_f32(10)) 177 let eps: i64 = nx_f32_div(one, nx_i32_to_f32(K_MAGIC_1000000)) 178 179 var i: i64 = 0 180 while i < hidden { ang[i] = one; fng[i] = one; i = i + 1 } 181 while i < q_dim * hidden { Wq[i] = p1; i = i + 1 } 182 i = 0 183 while i < kv_dim * hidden { Wk[i] = p1; Wv[i] = p1; i = i + 1 } 184 i = 0 185 while i < ffn_dim * hidden { Wg[i] = p1; Wu[i] = p1; i = i + 1 } 186 i = 0 187 while i < n_tokens * hidden { x[i] = nx_i32_to_f32(i + 1); i = i + 1 } 188 189 // (a) Wo=0, Wd=0 -> identity 190 i = 0 191 while i < hidden * q_dim { Wo[i] = 0; i = i + 1 } 192 i = 0 193 while i < hidden * ffn_dim { Wd[i] = 0; i = i + 1 } 194 if nx_f32_qwen_layer(x, n_tokens, hidden, nqh, nkvh, hd, ffn_dim, ang, Wq, Wk, Wv, Wo, fng, Wg, Wu, Wd, eps, out) != 0 { return 10 } 195 i = 0 196 while i < n_tokens * hidden { if out[i] != x[i] { return 20 } i = i + 1 } 197 198 // (b) nonzero Wo,Wd -> contributes 199 i = 0 200 while i < hidden * q_dim { Wo[i] = p1; i = i + 1 } 201 i = 0 202 while i < hidden * ffn_dim { Wd[i] = p1; i = i + 1 } 203 if nx_f32_qwen_layer(x, n_tokens, hidden, nqh, nkvh, hd, ffn_dim, ang, Wq, Wk, Wv, Wo, fng, Wg, Wu, Wd, eps, out) != 0 { return 30 } 204 var diff: i64 = 0 205 i = 0 206 while i < n_tokens * hidden { if out[i] != x[i] { diff = 1 } i = i + 1 } 207 if diff == 0 { return 40 } 208 209 // ===== (c) QK-NORM (debt 1785449705) ========================================== 210 // Q and K are normalised INDEPENDENTLY, so they are tested INDEPENDENTLY. A combined tooth 211 // (both gains on vs both off) would pass even if one of the two were silently dropped -- the 212 // same weakest-slot trap that a mutation caught in nx_f32_qwen_encoder earlier today. 213 // c1 null gains reproduce the pre-QK-norm result BIT-EXACT -> else 50 214 // c2 q_ng alone moves the output -> else 51 (q-norm ignored) 215 // c3 k_ng alone moves the output -> else 52 (k-norm ignored) 216 // BITE-VERIFIED: dropping k_ng in the implementation yields exit 52. 217 let qng: *i64 = sys_mmap(hd * 8) as *i64 218 let kng: *i64 = sys_mmap(hd * 8) as *i64 219 let out2: *i64 = sys_mmap(n_tokens * hidden * 8) as *i64 220 let p2: i64 = nx_f32_div(one, nx_i32_to_f32(5)) // 0.2 -- a gain of 1.0 would be a no-op 221 i = 0 222 while i < hd { qng[i] = p2; kng[i] = p2; i = i + 1 } 223 224 if nx_f32_qwen_layer_qk(x, n_tokens, hidden, nqh, nkvh, hd, ffn_dim, ang, Wq, Wk, Wv, Wo, 225 fng, Wg, Wu, Wd, 0 as *i64, 0 as *i64, eps, out2) != 0 { return 49 } 226 i = 0 227 while i < n_tokens * hidden { if out[i] != out2[i] { return 50 } i = i + 1 } 228 229 if nx_f32_qwen_layer_qk(x, n_tokens, hidden, nqh, nkvh, hd, ffn_dim, ang, Wq, Wk, Wv, Wo, 230 fng, Wg, Wu, Wd, qng, 0 as *i64, eps, out2) != 0 { return 49 } 231 var movedq: i64 = 0 232 i = 0 233 while i < n_tokens * hidden { if out[i] != out2[i] { movedq = 1 } i = i + 1 } 234 if movedq == 0 { return 51 } 235 236 if nx_f32_qwen_layer_qk(x, n_tokens, hidden, nqh, nkvh, hd, ffn_dim, ang, Wq, Wk, Wv, Wo, 237 fng, Wg, Wu, Wd, 0 as *i64, kng, eps, out2) != 0 { return 49 } 238 var movedk: i64 = 0 239 i = 0 240 while i < n_tokens * hidden { if out[i] != out2[i] { movedk = 1 } i = i + 1 } 241 if movedk == 0 { return 52 } 242 243 return 0 244}