code wiki / (root) / nx_rmsnorm.nx

nx_rmsnorm.nx

buildroot/runtime/nx_rmsnorm.nx

9571 B268 linesdepth 4pulls 6 transitivereach 25 importersview sourcekind tool
docsdependenciesstructsconstsfunctions

about

nx_rmsnorm.nx -- Root Mean Square Layer Normalisation. Closes a substrate gap: prior to this commit, NishiLang had matmul + attention + softmax + quantization shipped, but NO normalisation primitive. That makes the substrate unable to run a transformer block end-to-end (every modern transformer normalises between attention and FFN sublayers). RMSNorm vs LayerNorm choice (Zhang & Sennrich 2019): LayerNorm (Ba/Kiros/Hinton 2016): y = (x - mean(x)) / sqrt(var(x) + eps) * gamma + beta RMSNorm (Zhang & Sennrich 2019, simpler): y = x / sqrt(mean(x^2) + eps) * gamma RMSNorm drops the centering term (mean subtraction). Faster (~50% speedup per Zhang 2019 Table 2) at same downstream quality on translation + speech recognition + summarisation benchmarks. Modern transformer architectures use RMSNorm exclusively: Llama 1/2/3, Mistral, Mixtral, Qwen, Phi-3, Gemma Z-Image, Stable Diffusion 3, Flux Mamba state-space models LayerNorm support remains queued -- the math composes against the same nx_isqrt_q10 with one extra centering pass. ===== Math ======================================================= For each token i, hidden dim D, input x[i, :] of length D: sum_sq = sum_{d=0}^{D-1} x[i, d]^2 mean_sq = sum_sq / D rms = sqrt(mean_sq + eps) y[i, d] = x[i, d] / rms * gamma[d] gamma is the per-channel learned scale vector (length D), initialised to 1.0 (Q10=1024) at network construction time.

dependencies 5 imports · 3 importers

nx_syscalls.nx nx_tier.nx nx_loop.nx nx_tensor.nx nx_isqrt.nx nx_rmsnorm.nx nx_llm_run.nx nx_transformer_block.nx nx_transformer_stack.nx

imports: nx_syscalls.nxnx_tier.nxnx_loop.nxnx_tensor.nxnx_isqrt.nx

imported by: nx_llm_run.nxnx_transformer_block.nxnx_transformer_stack.nx

call flow from main pre-order; caps 40 nodes / depth 6 declared; ↻ = already shown

main sys_mmap nx_t_alloc nx_dt_is_valid nx_dt_element_bytes sys_mmap ↻ nx_t_compute_strides_rowma nx_rmsnorm_gamma_unit sys_mmap ↻ nx_rmsnorm_forward nx_t_is_contiguous nx_isqrt_q10 nx_isqrt nx_rmsn_verdict_is_valid

structs

none

consts

81const NX_RMSN_Q10: nx_int = 1024
82const NX_RMSN_EPS_Q10: nx_int = 1
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

functions

94func nx_rmsn_verdict_is_valid(v: nx_int) -> nx_int
called by 1: main
108func nx_rmsnorm_forward(x: *NxTensor, gamma: *i64, out: *NxTensor) -> nx_int
173func nx_rmsnorm_gamma_unit(d: nx_int) -> *i64
called by 1: main calls 1: sys_mmap
200func main() -> i64