For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo module

gated_group_rmsnorm

Fused gated group-RMSNorm for the Mamba-2 mixer (norm_before_gate=False).

Portable GPU/CPU kernel; the immediate target is Apple M5 Metal (compute_capability == 5) batch-1 decode. Collapses the four-op decode-graph chain that HF Zamba2RMSNormGated lowers to -- cast(y -> f32), silu(gate) * y, group rms_norm (a reduction, a hard dispatch boundary), and * norm_weight + cast -- into ONE kernel dispatch per Mamba layer. On the Nemotron-H decode graph that reduction cannot fold into the surrounding elementwise ops, so the chain is ~3-4 serial dispatches x 21 layers; the win is removing those launches, not FLOPs (decode processes 1 x 7680 elements).

Math is the norm_before_gate=False branch of the dense combined scan (selective_scan.mojo:2643-2648), ported to a standalone varlen op. For a row n and group g over group_size-wide contiguous groups of the intermediate axis, with col = g * group_size + j:

gated = f32(y[n, col]) * silu(f32(gate[n, col]))          # f32
m2    = sum_j gated^2                                     # f32, over group
nf    = rsqrt(m2 / group_size + eps)                      # matches ops.rms_norm
out[n, col] = cast(f32(norm_weight[col]) * f32(cast(gated * nf, y.dtype)),
                   out.dtype)

The intermediate cast(gated * nf, y.dtype) before the fp32 norm_weight multiply reproduces the reference exactly: w * ops.cast(yf, input_dtype) then a final cast to the model dtype (nemotron_h.py::_gated_group_rmsnorm + the out_proj cast). norm_weight is fp32; the product is fp32; this op returns the model dtype so the downstream out_proj cast is a no-op.

Parallelization (Apple-M5 idiom -- mirrors apple/fa_prefill, apple/fp4_gemv): one warp (simdgroup) owns one (row, group) reduction; its lanes stride the group by WARP_SIZE, accumulate a partial sum-of-squares in fp32, and one warp.sum reduces to the group m2. No shared memory, no barrier() -- the two levers Apple silicon is most sensitive to (KB kernels/apple-m5-fa-prefill, patterns/apple-m5-gpu-performance-considerations). A group_size that is not a multiple of WARP_SIZE is handled by the strided while-loop (idle lanes add 0), so no divisibility constraint is needed. All memory access is through TileTensor indexing (.load/.store with Coord) -- no raw pointer arithmetic -- so a strided y/gate/output view (a split of the fused in-proj) is read correctly with no host-side stride plumbing.

On NVIDIA and AMD the launcher takes a vectorized single-pass path when every operand is 16-byte aligned and group_size is a multiple of the vector width: each lane loads y, gate and weight vectors once, keeps the gated values in registers, and reuses them after the group reduction. The vector width is chosen so the wider of y and gate is 16 bytes. The fp32 sum of squares is accumulated per lane in a different order than the scalar path, so nf can differ in the last ulp; the formulas are otherwise identical.

comptime values

VEC_BLK

comptime VEC_BLK = 128

Functions

• ​gated_group_rmsnorm_cpu:
• ​gated_group_rmsnorm_gpu: Enqueues the fused gated group-RMSNorm; one warp per (row, group).
• ​gated_group_rmsnorm_kernel:
• ​gated_group_rmsnorm_vec_kernel: