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 function

lamport_allreduce_rmsnorm

def lamport_allreduce_rmsnorm[dtype: DType, ngpus: Int, *, pdl: Bool = True, early_launch: Bool = True](my_rank: Int, src: Pointer[Scalar[dtype], ImmutAnyOrigin, _safe=False], dst: Pointer[Scalar[dtype], MutAnyOrigin, _safe=False], gamma: Pointer[Scalar[dtype], ImmutAnyOrigin, _safe=False], rank_sigs: Array[Pointer[Signal, MutAnyOrigin, _safe=False], Int(8)], rows: Int, cols: Int, epsilon: Scalar[dtype], ctx: DeviceContext)

Per-device fused Lamport allreduce (sum) + RMSNorm (high-perf protocol).

src/dst are [rows, cols] (tokens x hidden); RMSNorm runs over cols. Uses the device-resident Signal.lamport_region / lamport_state (so the signal buffers must be lamport_init-ed once before the first call, then one host sync). The generation flag advances in-kernel, so repeated calls -- including CUDA-graph replay -- rotate the buffers correctly with no per-call fill or phase argument.

Parameters:

dtype (DType): Element type (bf16/fp16/fp32 transport).
ngpus (Int): Number of participating GPUs.
pdl (Bool): If True (default) launch with PDL and early-trigger the consumer.
early_launch (Bool): If True (default) and pdl, call launch_dependent_grids() at kernel start so the consumer overlaps the whole kernel. If False, the trigger moves to the END (after the output is written), so the consumer overlaps only the tail bookkeeping + grid teardown.