nx_native_rmsnorm_bench.nx source

1// nx_native_rmsnorm_bench.nx -- LLM-relevant RMSNorm kernel. 2// 3// Third dominant kernel in transformer-forward (matmul + softmax + 4// RMSNorm = the trio). RMSNorm = rsqrt(mean(x^2) + eps) * x. 5// Per-token, per-residual-stream-position, per-layer (so for a 6// 22-layer Llama with 2048 tokens, 22 * 2048 = ~45k RMSNorm calls). 7// 8// Workload: RMSNorm over a 2048-element residual stream, 5000 9// iterations. Per iteration: dot product of x with itself, divide 10// by N, integer sqrt approximation, then multiply each element. 11// 12// The substrate uses Q10 fixed-point. Integer sqrt is via Newton 13// iteration; we measure the dot product + divide + multiply-back 14// loop, skipping the sqrt step. 15 16import "nx_syscalls.nx" 17const K_MAGIC_2048: i64 = 2048 18const K_MAGIC_5000: i64 = 5000 19 20func main() -> i64 { 21 let N: i64 = K_MAGIC_2048 22 let ITERS: i64 = K_MAGIC_5000 23 24 let buf: *u8 = sys_mmap(N * 8) 25 let x: *i64 = buf as *i64 26 27 // Init residual stream with a pattern 28 var k: i64 = 0 29 while k < N { 30 x[k] = (k * 31 + 5) & 1023 31 k = k + 1 32 } 33 34 var checksum: i64 = 0 35 var iter: i64 = 0 36 while iter < ITERS { 37 // Step 1: dot(x, x) 38 var ss: i64 = 0 39 var i: i64 = 0 40 while i < N { 41 ss = ss + x[i] * x[i] 42 i = i + 1 43 } 44 45 // Step 2: mean -> "scale" (substrate skips real sqrt; uses 46 // approximation table or Newton. Here: shift to approximate) 47 let scale: i64 = (ss / N) >> 4 48 49 // Step 3: multiply each element by approximate inv-rms. 50 // For checksum we accumulate (x[i] * scale) >> 10 51 var s: i64 = 0 52 var j: i64 = 0 53 while j < N { 54 s = s + ((x[j] * scale) >> 10) 55 j = j + 1 56 } 57 58 checksum = checksum + (s & 255) 59 iter = iter + 1 60 } 61 62 return checksum 63}