nx_f32_rmsnorm.nx source
↩ module page · 86 lines · 2948 B
1// nx_f32_rmsnorm.nx -- bits-up f32 RMSNorm.
2//
3// L7 / L8 composition brick. The first ML normalization primitive
4// fully written in our bits-up IEEE 754 binary32 substrate:
5//
6// ss = sum_i (x[i] * x[i])
7// mean_ss = ss / n
8// denom = sqrt(mean_ss + eps)
9// inv_denom = 1 / denom
10// out[i] = x[i] * gamma[i] * inv_denom
11//
12// Composes:
13// L4 f32: mul, add, sqrt, div, cvt (i32 -> f32 for the n constant)
14//
15// No libm. No compiler-builtin float. Every op routes through the
16// bits-up substrate from this session's arc.
17//
18// Reference (the standard RMSNorm definition):
19// Zhang + Sennrich 2019, "Root Mean Square Layer Normalization"
20// (paper PDF; the math is a published recipe, no code borrowed)
21//
22// genealogy_id: zhang_sennrich_2019_rms_norm + ieee754_f32_compose
23// lineage_id: substrate_f32_rmsnorm_v1
24
25import "nx_syscalls.nx"
26import "nx_tier.nx"
27import "nx_f32.nx"
28import "nx_f32_div.nx"
29import "nx_f32_cvt.nx"
30
31const NX_F32_RMSN_ONE: i64 = 0x3F800000 // 1.0
32const NX_F32_RMSN_OK: nx_int = 0
33const NX_F32_RMSN_ERR_BAD_DIM: nx_int = 1
34const NX_F32_RMSN_N_VERDICTS: nx_int = 2
35
36func nx_f32_rmsnorm_verdict_is_valid(v: nx_int) -> nx_int {
37 if v < 0 { return 0 }
38 if v >= NX_F32_RMSN_N_VERDICTS { return 0 }
39 return 1
40}
41
42// Args:
43// x pointer to n f32 input values (raw bit patterns)
44// gamma pointer to n f32 learned scale values (raw)
45// n hidden dimension (positive)
46// eps_f32 small additive constant (raw f32 bits); pass 0 to disable
47// out pointer to n f32 output slots (raw)
48//
49// Pass-1 reduction + pass-2 broadcast. In-place permitted (out == x).
50
51func nx_f32_rmsnorm(x: *i64, gamma: *i64, n: nx_int,
52 eps_f32: i64, out: *i64) -> nx_int {
53 if n <= 0 { return NX_F32_RMSN_ERR_BAD_DIM }
54
55 // Pass 1: accumulate sum-of-squares in f32. HARDWARE __f32_mul/add
56 // (2026-07-10): IEEE-754, and the SAME serial accumulation order as the
57 // software nx_f32_* -> BIT-IDENTICAL, ~1 instr vs ~30. Called ~49x/token
58 // (2 norms/layer + final) x n=896 -> a real decode-hot loop. sqrt stays
59 // software (no __f32_sqrt intrinsic; once/call = negligible).
60 var ss: i64 = 0
61 var i: nx_int = 0
62 while i < n {
63 let xi: i64 = x[i]
64 ss = __f32_add(ss, __f32_mul(xi, xi))
65 i = i + 1
66 }
67
68 // mean = ss / n
69 let n_f32: i64 = nx_i32_to_f32(n)
70 let mean_ss: i64 = __f32_div(ss, n_f32)
71
72 // denom = sqrt(mean + eps)
73 let mean_plus_eps: i64 = __f32_add(mean_ss, eps_f32)
74 let denom: i64 = nx_f32_sqrt(mean_plus_eps)
75 let inv_denom: i64 = __f32_div(NX_F32_RMSN_ONE, denom)
76
77 // Pass 2: scale each x[i] by gamma[i] * inv_denom.
78 var j: nx_int = 0
79 while j < n {
80 let xj: i64 = x[j]
81 let gj: i64 = gamma[j]
82 out[j] = __f32_mul(__f32_mul(xj, gj), inv_denom)
83 j = j + 1
84 }
85 return NX_F32_RMSN_OK
86}