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
allgather_rmsnorm
Fused all-gather + RMSNorm (bf16 in/out, no FP8).
All-gathers ngpus row-shards into the full replicated [rows, cols] stream and
RMSNorms every gathered row in one launch (no separate-norm HBM round-trip).
Emits sum_out (gathered residual) and normed_out (its RMSNorm), both full and
replicated; mbc=True. A gathered row is a verbatim copy, so sum_out is
bit-for-bit a standalone allgather (no f32 peer-sum, unlike reduce-scatter).
Factored from allreduce_residual_rmsnorm.mojo's Stage-2 gather + the norm math
of reducescatter_rmsnorm.mojo -- keep in sync. One block per GLOBAL row so each
row's block.sum runs parallel across blocks, not ngpus serial within one.
comptime values
AG_NORM_CALIBRATED_COLS
comptime AG_NORM_CALIBRATED_COLS = 6144
AG_NORM_FUSE_THRESHOLD
comptime AG_NORM_FUSE_THRESHOLD = (Int(786432) * size_of[DType.bfloat16]())
Functions
-
allgather_rmsnorm: Fused all-gather + RMSNorm acrossngpusGPUs (bf16 in/out).
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!