For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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.

Functions

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

Functions​

Functions