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

reducescatter_rmsnorm

Fused reduce-scatter + RMSNorm (bf16 in/out, no FP8).

Reduce-scatters a [rows, cols] stream across ngpus GPUs and RMSNorms each owned row, keeping the sum in registers so there is no HBM round-trip between the two. The sum is rounded to in_dtype before the norm so the result is bit-for-bit the standalone reduce-scatter -> bf16 shard -> rms_norm path (norming the wider f32 sum silently shifts model behavior). Emits two [rank_units, cols] shards: sum_out (the reduce-scatter sum / residual stream) and normed_out (its RMSNorm). Inherently multiply_before_cast=True.

Factored from the 2-stage AR kernel (allreduce_residual_rmsnorm.mojo); keep the reduce/norm math in sync with it.

comptime values

RS_NORM_FUSE_THRESHOLD

comptime RS_NORM_FUSE_THRESHOLD = (Int(786432) * size_of[DType.bfloat16]())

Functions

• ​reducescatter_rmsnorm: Fused reduce-scatter + RMSNorm across ngpus GPUs (bf16 in/out).

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit