nx_rmsnorm.nx source

1// nx_rmsnorm.nx -- Root Mean Square Layer Normalisation. 2// 3// Closes a substrate gap: prior to this commit, NishiLang had matmul 4// + attention + softmax + quantization shipped, but NO normalisation 5// primitive. That makes the substrate unable to run a transformer 6// block end-to-end (every modern transformer normalises between 7// attention and FFN sublayers). 8// 9// RMSNorm vs LayerNorm choice (Zhang & Sennrich 2019): 10// 11// LayerNorm (Ba/Kiros/Hinton 2016): 12// y = (x - mean(x)) / sqrt(var(x) + eps) * gamma + beta 13// 14// RMSNorm (Zhang & Sennrich 2019, simpler): 15// y = x / sqrt(mean(x^2) + eps) * gamma 16// 17// RMSNorm drops the centering term (mean subtraction). Faster 18// (~50% speedup per Zhang 2019 Table 2) at same downstream quality 19// on translation + speech recognition + summarisation benchmarks. 20// Modern transformer architectures use RMSNorm exclusively: 21// 22// Llama 1/2/3, Mistral, Mixtral, Qwen, Phi-3, Gemma 23// Z-Image, Stable Diffusion 3, Flux 24// Mamba state-space models 25// 26// LayerNorm support remains queued -- the math composes against 27// the same nx_isqrt_q10 with one extra centering pass. 28// 29// ===== Math ======================================================= 30// 31// For each token i, hidden dim D, input x[i, :] of length D: 32// 33// sum_sq = sum_{d=0}^{D-1} x[i, d]^2 34// mean_sq = sum_sq / D 35// rms = sqrt(mean_sq + eps) 36// y[i, d] = x[i, d] / rms * gamma[d] 37// 38// gamma is the per-channel learned scale vector (length D), 39// initialised to 1.0 (Q10=1024) at network construction time. 40// 41// ===== Q-format =================================================== 42// 43// All tensors are i64 in Q10 by substrate convention: 44// * x[i, d] in Q10 (range typically [-30 Q10, +30 Q10]) 45// * gamma[d] in Q10 (typically near Q10=1024 at init) 46// * y[i, d] in Q10 47// 48// sum_sq compute: x^2 of two Q10 values has Q20 scale. We track 49// in raw integer (sum is summed integer products) and divide back 50// after sqrt: rms in Q10 = nx_isqrt_q10(mean_sq + eps). 51// 52// eps_q10 = 1 (smallest non-zero Q10 stabiliser). In float 53// references eps = 1e-6; in Q10 the smallest representable non- 54// zero is ~1e-3, so eps_q10=1 is the tightest stabiliser we can 55// give. Matches GGML's eps handling for q-format quantized 56// inference. 57// 58// Per the bits-up cardinal: composes against 59// * nx_tensor.NxTensor (L1 canonical container) 60// * nx_isqrt.nx_isqrt_q10 (L2 canonical sqrt primitive) 61// * nx_loop.LoopVerdict (bounded-loop discipline) 62// 63// genealogy_id: zhang_sennrich_2019_rms_norm + ba_kiros_hinton_2016_layer_norm + 64// touvron_2023_llama_rms_norm_adoption 65// lineage_id: substrate_rmsnorm_v1 66 67// nx_safety_envelope: 68// intended_use: AUTO_APPLIED -- primitive-specific tuning queued 69// sil_target: SIL1 70// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail] 71// verdict: NOT_YET_EVALUATED 72 73import "nx_syscalls.nx" 74import "nx_tier.nx" 75import "nx_loop.nx" 76import "nx_tensor.nx" 77import "nx_isqrt.nx" 78 79// ===== Constants ================================================== 80 81const NX_RMSN_Q10: nx_int = 1024 82const NX_RMSN_EPS_Q10: nx_int = 1 83 84// ===== Sealed-enum: RmsNormVerdict ================================ 85 86const NX_RMSN_OK: nx_int = 0 87const NX_RMSN_ERR_BAD_DTYPE: nx_int = 1 88const NX_RMSN_ERR_BAD_NDIM: nx_int = 2 89const NX_RMSN_ERR_SHAPE_MISMATCH: nx_int = 3 90const NX_RMSN_ERR_NOT_CONTIGUOUS: nx_int = 4 91const NX_RMSN_ERR_BAD_GAMMA_LEN: nx_int = 5 92const NX_RMSN_N_VERDICTS: nx_int = 6 93 94func nx_rmsn_verdict_is_valid(v: nx_int) -> nx_int { 95 if v < 0 { return 0 } 96 if v >= NX_RMSN_N_VERDICTS { return 0 } 97 return 1 98} 99 100// ===== Forward pass ============================================== 101// 102// x: [n_tokens, hidden_dim] Q10 input 103// gamma: [hidden_dim] Q10 scale (learned) 104// out: [n_tokens, hidden_dim] Q10 output 105// 106// In-place is supported: caller can pass out == x. 107 108func nx_rmsnorm_forward(x: *NxTensor, gamma: *i64, out: *NxTensor) -> nx_int { 109 if x.dtype != NX_DT_I64 { return NX_RMSN_ERR_BAD_DTYPE } 110 if out.dtype != NX_DT_I64 { return NX_RMSN_ERR_BAD_DTYPE } 111 if x.ndim != 2 { return NX_RMSN_ERR_BAD_NDIM } 112 if out.ndim != 2 { return NX_RMSN_ERR_BAD_NDIM } 113 if x.shape[0] != out.shape[0] { return NX_RMSN_ERR_SHAPE_MISMATCH } 114 if x.shape[1] != out.shape[1] { return NX_RMSN_ERR_SHAPE_MISMATCH } 115 if nx_t_is_contiguous(x) == 0 { return NX_RMSN_ERR_NOT_CONTIGUOUS } 116 if nx_t_is_contiguous(out) == 0 { return NX_RMSN_ERR_NOT_CONTIGUOUS } 117 118 let n_tok: nx_int = x.shape[0] 119 let d: nx_int = x.shape[1] 120 let px: *i64 = x.storage as *i64 121 let po: *i64 = out.storage as *i64 122 123 var i: nx_int = 0 124 var iter: nx_int = 0 125 var verdict: nx_int = NX_LOOP_RUNNING 126 let BUDGET: nx_int = n_tok 127 while verdict == NX_LOOP_RUNNING && iter < BUDGET { 128 // Pass 1: sum of squares (raw Q20 accumulation). 129 var sum_sq: i64 = 0 130 var c: nx_int = 0 131 var iter_c: nx_int = 0 132 var verdict_c: nx_int = NX_LOOP_RUNNING 133 let BUDGET_C: nx_int = d 134 while verdict_c == NX_LOOP_RUNNING && iter_c < BUDGET_C { 135 let v: i64 = px[i * d + c] 136 sum_sq = sum_sq + v * v 137 c = c + 1 138 iter_c = iter_c + 1 139 } 140 // mean(x^2) in Q20 / D = Q20. We want Q10 mean for sqrt: 141 // mean_sq_q20 = sum_sq / d 142 // mean_sq_q10 = mean_sq_q20 / Q10 = sum_sq / (d * Q10) 143 let mean_sq_q10: i64 = sum_sq / (d * NX_RMSN_Q10) 144 let rms_q10: i64 = nx_isqrt_q10(mean_sq_q10 + NX_RMSN_EPS_Q10) 145 if rms_q10 <= 0 { verdict = NX_LOOP_DONE_EXIT } 146 147 // Pass 2: normalise + scale by gamma. 148 if verdict == NX_LOOP_RUNNING { 149 var c2: nx_int = 0 150 var iter_c2: nx_int = 0 151 var verdict_c2: nx_int = NX_LOOP_RUNNING 152 while verdict_c2 == NX_LOOP_RUNNING && iter_c2 < BUDGET_C { 153 let xv: i64 = px[i * d + c2] 154 // norm_q10 = x_q10 * Q10 / rms_q10 155 let norm: i64 = (xv * NX_RMSN_Q10) / rms_q10 156 // scaled = norm * gamma_q10 / Q10 157 po[i * d + c2] = (norm * gamma[c2]) / NX_RMSN_Q10 158 c2 = c2 + 1 159 iter_c2 = iter_c2 + 1 160 } 161 } 162 i = i + 1 163 iter = iter + 1 164 } 165 return NX_RMSN_OK 166} 167 168// ===== Convenience: build a unit-gamma vector ==================== 169// 170// Caller-friendly factory. Allocates a length-D array all set to 171// Q10=1024 (the standard RMSNorm initialisation). 172 173func nx_rmsnorm_gamma_unit(d: nx_int) -> *i64 { 174 let g: *i64 = sys_mmap(d * 8) as *i64 175 var i: nx_int = 0 176 var iter: nx_int = 0 177 var verdict: nx_int = NX_LOOP_RUNNING 178 let BUDGET: nx_int = d 179 while verdict == NX_LOOP_RUNNING && iter < BUDGET { 180 g[i] = NX_RMSN_Q10 181 i = i + 1 182 iter = iter + 1 183 } 184 return g 185} 186 187// ===== Self-test ================================================== 188// 189// Closed-form invariants: 190// 191// (a) Constant input vector: RMS(c, c, ..., c) = c. After RMSNorm 192// with gamma=1, y[d] = c / c = 1 (Q10=1024) for every d. 193// 194// (b) Scale invariance: RMSNorm(k*x) = RMSNorm(x). Verify on 195// two parallel tokens scaled 2x apart. 196// 197// (c) Gamma scaling: with gamma[d] = 2 (in Q10=2048), output 198// should be 2x what gamma=1 produces. 199 200func main() -> i64 { 201 let n_tok: nx_int = 2 202 let d: nx_int = 8 203 204 let sh: *nx_int = sys_mmap(2 * 8) as *nx_int 205 sh[0] = n_tok; sh[1] = d 206 207 let err: *nx_int = sys_mmap(8) as *nx_int 208 err[0] = 0 209 let xt: *NxTensor = nx_t_alloc(NX_DT_I64, sh, 2, err) 210 let yt: *NxTensor = nx_t_alloc(NX_DT_I64, sh, 2, err) 211 if err[0] != 0 { return 5 } 212 213 // --- (a) Constant token --- 214 // Fill token 0 with all 5*Q10 = 5120. After RMSNorm with gamma=1, 215 // every output should be Q10 = 1024 (since x/RMS(x) = 1 for 216 // constant input). 217 let px: *i64 = xt.storage as *i64 218 let py: *i64 = yt.storage as *i64 219 var c: nx_int = 0 220 while c < d { px[c] = 5 * NX_RMSN_Q10; c = c + 1 } 221 222 // Token 1: scaled 4x version of token 0 -- tests scale invariance. 223 var c2: nx_int = 0 224 while c2 < d { px[d + c2] = 20 * NX_RMSN_Q10; c2 = c2 + 1 } 225 226 let gamma_unit: *i64 = nx_rmsnorm_gamma_unit(d) 227 let v: nx_int = nx_rmsnorm_forward(xt, gamma_unit, yt) 228 if v != NX_RMSN_OK { return 10 + v } 229 230 // Verify token 0 output = Q10=1024 (allow +/- 2 for rounding). 231 var ci: nx_int = 0 232 while ci < d { 233 let drift: i64 = py[ci] - NX_RMSN_Q10 234 if drift > 2 { return 20 } 235 if drift < -2 { return 21 } 236 ci = ci + 1 237 } 238 239 // --- (b) Scale invariance: token 1 should produce same output. --- 240 var ck: nx_int = 0 241 while ck < d { 242 let drift: i64 = py[d + ck] - NX_RMSN_Q10 243 if drift > 2 { return 30 } 244 if drift < -2 { return 31 } 245 ck = ck + 1 246 } 247 248 // --- (c) Gamma scaling: 2x gamma -> 2x output --- 249 var cj: nx_int = 0 250 while cj < d { gamma_unit[cj] = 2 * NX_RMSN_Q10; cj = cj + 1 } 251 nx_rmsnorm_forward(xt, gamma_unit, yt) 252 var cm: nx_int = 0 253 while cm < d { 254 let drift: i64 = py[cm] - 2 * NX_RMSN_Q10 255 if drift > 4 { return 40 } 256 if drift < -4 { return 41 } 257 cm = cm + 1 258 } 259 260 // --- (d) Verdict gate --- 261 var k: nx_int = 0 262 while k < NX_RMSN_N_VERDICTS { 263 if nx_rmsn_verdict_is_valid(k) != 1 { return 50 + k } 264 k = k + 1 265 } 266 267 return 0 268}