RMSNorm Layer - Root Mean Square Normalization

Implements RMSNorm as used in BitNet and Llama architectures.

Mathematical Definition

RMSNorm normalizes a vector by its root mean square:

RMS(x) = sqrt(1/n * Σᵢ xᵢ²)
y = (x / RMS(x)) * γ

Where:

• x: input vector (n elements)
• γ: learned scale parameter (n elements)
• ε: small constant for numerical stability (typically 1e-5 or 1e-6)

Comparison with LayerNorm

LayerNorm: y = γ * (x - μ) / σ + β

• Requires mean (μ) and variance (σ) calculation
• Has bias term (β)
• More expensive (two passes + bias)

RMSNorm: y = γ * x / RMS(x)

• No mean subtraction (assumes mean ≈ 0 from prior layers)
• No bias term
• Simpler and faster (single pass)

Performance Advantages

1. Simpler computation: No mean calculation
2. Better gradients: Avoids mean-centering instability
3. Fused implementation: RMS + scale in single kernel

References

• Llama: https://github.com/facebookresearch/llama (RMSNorm implementation)
• Paper: "Root Mean Square Layer Normalization" (Zhang & Sennrich, 2019)
• llama.cpp: ggml/src/ggml.c (ggml_rms_norm implementation)

opaque Hesper.Layers.RMSNorm.preparedHitsRef : IO.Ref Nat

Counters for PreparedDispatch fast-path vs slow-path

opaque Hesper.Layers.RMSNorm.preparedMissesRef : IO.Ref Nat

Layer Configuration

structure Hesper.Layers.RMSNorm.Config : Type

RMSNorm layer configuration

• dim : Nat
• eps : Float

instance Hesper.Layers.RMSNorm.instReprConfig : Repr Config

Equations

• Hesper.Layers.RMSNorm.instReprConfig = { reprPrec := Hesper.Layers.RMSNorm.instReprConfig.repr }

def Hesper.Layers.RMSNorm.instReprConfig.repr : Config → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

GPU Kernel Implementation

def Hesper.Layers.RMSNorm.rmsNormKernel (config : Config) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

RMSNorm kernel using workgroup reduction for RMS calculation

Algorithm:

1. Each thread loads one element and computes x²
2. Parallel reduction to compute sum(x²) in workgroup shared memory
3. Compute RMS = sqrt(sum / n + ε)
4. Each thread normalizes: y = (x / RMS) * scale

Workgroup strategy:

• Workgroup size: 256 threads (typical for hidden dims like 768, 1024, 2048)
• For larger dims: Multiple workgroups, each handling 256 elements
• Reduction uses shared memory for efficiency

@param config RMSNorm configuration

Equations

• One or more equations did not get rendered due to their size.

Fused Single-Pass Kernel (multi-row)

def Hesper.Layers.RMSNorm.rmsNormFusedKernel (config : Config) (numRows : Nat) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

Fused RMSNorm kernel: compute RMS + apply normalization in one dispatch. Each workgroup handles one row. Threads use strided loops for both RMS accumulation and normalization, supporting dim >> workgroupSize.

Reduces dispatch count from 2 to 1 per RMSNorm call.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Layers.RMSNorm.rmsNormAddBatchRowsKernel (config : Config) (numRows : Nat) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

Row-major batched fused post-norm+residual: output[r,i] = RMSNorm(layer_out[r])[i]*scale[i] + residual[r,i]. Same reduction structure as rmsNormFusedKernel (1 WG/row, warp-shuffle block sum); fuses the separate residual-add dispatch and skips the intermediate normed-row round trip. Fusion pattern after ggml-metal's kernel_rms_norm_fuse_impl (F=3; llama.cpp, MIT).

Equations

• One or more equations did not get rendered due to their size.

Fused post-norm (Gemma 4 style): RMSNorm(layer_out) + residual

def Hesper.Layers.RMSNorm.rmsNormThenAddKernel (config : Config) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

Computes output = RMSNorm(layer_out) * scale + residual in one kernel.

This is Gemma 4's post-norm shape: the normalisation is applied to the attention / FFN output FIRST, and only then the pre-block residual is added back in. Replaces the two-dispatch pattern RMSNorm.forward → residualAddKernel used after attention and FFN.

Dispatch: 1 workgroup × workgroupSize threads.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Layers.RMSNorm.rmsNormThenAddThenBroadcastScaleKernel (config : Config) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

rmsNormThenAddKernel with an extra broadcast scalar multiply at the tail: output[i] = (RMSNorm(layer_out)[i] * scale[i] + residual[i]) * out_scale[0].

Saves a separate layerOutScale dispatch when the block has a per-layer output scale (Gemma 4 dense FFN blocks). Structurally identical to rmsNormThenAddKernel except for the tail multiply.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Layers.RMSNorm.rmsNormThenAddBatchKernel (config : Config) (seqLen : Nat) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

Batched version of rmsNormThenAddKernel: processes seqLen rows in a single dispatch. Each workgroup handles one row i of the [dim, seqLen] column-major batch buffer (element (d, i) lives at buf[i*dim + d], matching columnExtractKernel / columnInsertKernel).

outputBatch[i, d] = RMSNorm(layerOutBatch[i, :])[d] * scale[d] + residualBatch[i, d]

Replaces the per-token (extract, forwardNormThenAdd, insert) chain in forwardPrefillBatch for post-attention and post-FFN residual+norm. Dispatch: (seqLen, 1, 1) workgroups × workgroupSize threads.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Layers.RMSNorm.rmsNormThenAddBatchInPlaceKernel (config : Config) (seqLen : Nat) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

IN-PLACE variant of rmsNormThenAddBatchKernel: residual_io += RMSNorm(layer_out) * scale. A single read_write binding — the two-binding form with residual == output is writable-storage aliasing (legal on CUDA, VALIDATION ERROR on WebGPU). Per-element read→add→write by the same thread at the same index (the RMS reduction only reads layer_out), so in-place is race-free.

Equations

• One or more equations did not get rendered due to their size.

Fused residual-add + RMSNorm

def Hesper.Layers.RMSNorm.residualAddRmsNormKernel (config : Config) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

Compute residualOut = a + b and output = RMSNorm(residualOut) * scale in one dispatch. Replaces the two-kernel pattern `residualAddKernel

• RMSNorm.forward` used after attention/FFN.

Dispatch: 1 workgroup × workgroupSize threads per row. Each row independently loads a[row]+b[row], reduces to RMS, and writes both residualOut[row] (for the next residual chain) and output[row] (the normalised result).

Equations

• One or more equations did not get rendered due to their size.

Optimized Two-Pass Kernel

def Hesper.Layers.RMSNorm.rmsComputeKernel (config : Config) (numRows : Nat) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

First pass: Compute RMS value

This kernel computes the RMS value using parallel reduction. Result is a single scalar stored in a 1-element buffer.

@param config RMSNorm configuration

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Layers.RMSNorm.rmsApplyKernel (config : Config) (numRows : Nat) : WGSL.Monad.ShaderM Unit

Second pass: Apply normalization using precomputed RMS

This kernel normalizes the input using the RMS computed in the first pass.

@param config RMSNorm configuration

Equations

• One or more equations did not get rendered due to their size.

Fused RMSNorm + Q8_1 quantize

Eliminates the VRAM round-trip between RMSNorm output and the matmul's Q8_1 quantize phase by combining both into a single dispatch.

Algorithm: Phase 1 — Single-WG cooperative RMSNorm reduction - 256 threads stride over D=config.dim input elements (each lane processes D / 256 elements via the strided loop). - Tree reduction in shared memory yields the sum-of-squares. - All threads compute invRms = rsqrt(sumSq/D + eps) from scratch[0]. Phase 2 — Per-block Q8_1 quantize (10 strided passes for D=2560) Per pass p, lane tid handles input element tid + p*256. The 256 lanes split into 8 warps of 32 lanes each. Warp w owns the 32-element block at 8*p + w of the output.

Inside each warp (lanes l=0..31, owning one Q8_1 block):
  - x_normed[l] = inputBuf[elemIdx] * scale[elemIdx] * invRms
  - amax = subgroupMax(|x_normed[l]|)
  - d = amax / 127
  - q[l] = round(x_normed[l] / d)  (clamped to int8)
  - shared_q[warpId * 32 + l] = q[l]; barrier
  - lane l divisible by 4 packs 4 quants into one u32 → output

Lane 0 of each warp writes the d|s header (s=0; subsequent
Q4_K dp4a path doesn't use s).

This stays within Hesper's hand-written-kernel inventory but is wired through the IR as Prim.rmsNormQ8_1Quantize, so callers see it as a proper compiler op. The Stage 3 follow-up will replace the body with auto-generated reduce + double-reduce-epilogue lowering.

def Hesper.Layers.RMSNorm.fusedRMSNormQ8_1Kernel (config : Config) (numRows : Nat := 1) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

Equations

• One or more equations did not get rendered due to their size.

High-Level API

structure Hesper.Layers.RMSNorm.RMSNorm (BufT : Type) (CacheT : Type := Unit) : Type

RMSNorm layer structure

• config : Config
• scale : BufT
• prepared : IO.Ref (Option CacheT)

def Hesper.Layers.RMSNorm.create {β : Type} [GPUBackend β] (ctx : β) (config : Config) (scaleData : ByteArray) : IO (RMSNorm (GPUBackend.Buf β) (GPUBackend.CachedDispatch β))

Create RMSNorm layer from GGUF tensors

@param device WebGPU device @param config Layer configuration @param scaleData Raw scale data from GGUF (Float32 or FP16)

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Hesper.Layers.RMSNorm.forward {β : Type} [GPUBackend β] (ctx : β) (layer : RMSNorm (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) (inputBuf outputBuf : GPUBackend.Buf β) (numRows : Nat := 1) (workgroupSize : Nat := 256) (preAllocRmsBuf : Option (GPUBackend.Buf β) := none) (refOverride : Option (IO.Ref (Option (GPUBackend.CachedDispatch β))) := none) : IO Unit

Execute forward pass (single-kernel version)

@param device WebGPU device @param layer RMSNorm layer @param inputBuf GPU buffer containing input (Float32) @param outputBuf GPU buffer for output (Float32)

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Layers.RMSNorm.forwardNormThenAddBatchRows {β : Type} [GPUBackend β] (ctx : β) (layer : RMSNorm (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) (layerOutBuf residualBuf outputBuf : GPUBackend.Buf β) (numRows : Nat) (workgroupSize : Nat := 256) : IO Unit

Dispatch rmsNormAddBatchRowsKernel (row-major [numRows, dim] buffers).

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Hesper.Layers.RMSNorm.forwardNormThenAdd {β : Type} [GPUBackend β] (ctx : β) (layer : RMSNorm (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) (layerOutBuf residualBuf outputBuf : GPUBackend.Buf β) (preparedRef : IO.Ref (Option (GPUBackend.CachedDispatch β))) (workgroupSize : Nat := 256) : IO Unit

Fused post-norm: output = RMSNorm(layer_out) * scale + residual. Gemma 4's post-attention / post-FFN pattern in a single dispatch.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Layers.RMSNorm.forwardNormThenAddThenBroadcastScale {β : Type} [GPUBackend β] (ctx : β) (layer : RMSNorm (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) (layerOutBuf residualBuf outScaleBuf outputBuf : GPUBackend.Buf β) (preparedRef : IO.Ref (Option (GPUBackend.CachedDispatch β))) (workgroupSize : Nat := 256) : IO Unit

Variant of forwardNormThenAdd that also applies an extra broadcast scalar outScaleBuf[0] at the tail. Saves a separate layerOutScale dispatch.

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Hesper.Layers.RMSNorm.forwardNormThenAddBatch {β : Type} [GPUBackend β] (ctx : β) (layer : RMSNorm (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) (layerOutBuf residualBuf outputBuf : GPUBackend.Buf β) (seqLen : Nat) (preparedRef : IO.Ref (Option (GPUBackend.CachedDispatch β))) (workgroupSize : Nat := 256) : IO Unit

Batched post-norm+residual: processes all seqLen rows in one dispatch. The input/output buffers are [dim, seqLen] column-major (same stride as columnExtractKernel), so this eliminates the per-token extract → forwardNormThenAdd → insert chain.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Layers.RMSNorm.forwardNormThenAddBatchInPlace {β : Type} [GPUBackend β] (ctx : β) (layer : RMSNorm (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) (layerOutBuf residualIoBuf : GPUBackend.Buf β) (seqLen : Nat) (preparedRef : IO.Ref (Option (GPUBackend.CachedDispatch β))) (workgroupSize : Nat := 256) : IO Unit

In-place driver for rmsNormThenAddBatchInPlaceKernel (residual buffer updated in place).

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Hesper.Layers.RMSNorm.forwardResidualAdd {β : Type} [GPUBackend β] (ctx : β) (layer : RMSNorm (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) (aBuf bBuf residualOutBuf outputBuf : GPUBackend.Buf β) (preparedRef : IO.Ref (Option (GPUBackend.CachedDispatch β))) (workgroupSize : Nat := 256) : IO Unit

Fused residual-add + RMSNorm forward. Computes residualOut = a + b and output = RMSNorm(residualOut) * scale in a single kernel — replaces the two-dispatch pattern residualAddKernel + RMSNorm.forward.

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Hesper.Layers.RMSNorm.forwardTwoPass {β : Type} [GPUBackend β] (ctx : β) (layer : RMSNorm (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) (inputBuf outputBuf rmsTempBuf : GPUBackend.Buf β) (numRows : Nat := 1) (workgroupSize : Nat := 256) : IO Unit

Execute forward pass (two-pass optimized version)

This version uses separate kernels for RMS computation and normalization, enabling better workgroup reduction optimization.

@param device WebGPU device @param layer RMSNorm layer @param inputBuf GPU buffer containing input (Float32) @param outputBuf GPU buffer for output (Float32) @param rmsTempBuf Temporary 1-element buffer for RMS value

Equations

• One or more equations did not get rendered due to their size.

Integration with GGUF Reader

def Hesper.Layers.RMSNorm.fromGGUF {β : Type} {α : Sort u_1} [GPUBackend β] (ctx : β) (gguf : α) (tensorName : String) (config : Config) : IO (RMSNorm (GPUBackend.Buf β) (GPUBackend.CachedDispatch β))

Create RMSNorm layer directly from GGUF file

This extracts the scale (weight) tensor for a specific layer.

Example tensor names in GGUF:

• blk.0.attn_norm.weight - Pre-attention RMSNorm
• blk.0.ffn_norm.weight - Pre-FFN RMSNorm
• output_norm.weight - Final output RMSNorm

@param device WebGPU device @param gguf Loaded GGUF file @param tensorName Name of the scale tensor in GGUF @param config Layer configuration

Equations

• One or more equations did not get rendered due to their size.