nx_rmsnorm.nx
buildroot/runtime/nx_rmsnorm.nx
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
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
structs
| none |
consts
| 81 | const NX_RMSN_Q10: nx_int = 1024 |
| 82 | const NX_RMSN_EPS_Q10: nx_int = 1 |
| 86 | const NX_RMSN_OK: nx_int = 0 |
| 87 | const NX_RMSN_ERR_BAD_DTYPE: nx_int = 1 |
| 88 | const NX_RMSN_ERR_BAD_NDIM: nx_int = 2 |
| 89 | const NX_RMSN_ERR_SHAPE_MISMATCH: nx_int = 3 |
| 90 | const NX_RMSN_ERR_NOT_CONTIGUOUS: nx_int = 4 |
| 91 | const NX_RMSN_ERR_BAD_GAMMA_LEN: nx_int = 5 |
| 92 | const NX_RMSN_N_VERDICTS: nx_int = 6 |
functions
| 94 | func nx_rmsn_verdict_is_valid(v: nx_int) -> nx_int called by 1: main |
| 108 | func nx_rmsnorm_forward(x: *NxTensor, gamma: *i64, out: *NxTensor) -> nx_int called by 3: mainnx_transformer_block_forwardnx_transformer_stack_forward calls 2: nx_t_is_contiguousnx_isqrt_q10 |
| 173 | func nx_rmsnorm_gamma_unit(d: nx_int) -> *i64 |
| 200 | func main() -> i64 |