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

allreduce_lamport_rmsnorm

Fused one-shot Lamport allreduce + RMSNorm (high-perf protocol).

Grafts the barrier-free Lamport comm protocol from comm/lamport.mojo + comm/allreduce.mojo (in-register local seed, poll-all-peers arrival, 3-way generation rotation, fused clear, the device-resident Signal.lamport_region / lamport_state) onto a row-blocked RMSNorm epilogue so the whole TP allreduce+RMSNorm producer is a single kernel.

Layout: each block owns whole rows (tokens). With the 128-bit Lamport pack (atomic_width = 16B / sizeof(dtype) = 8 for bf16) one thread owns exactly one pack = one column group, so a row's cols/atomic_width packs map to that many threads in one block. After the allreduce reduces each pack in-register, the block reduces the row's sum-of-squares (block.sum) and normalizes.

PDL is on by default: wait for the producer, then trigger the consumer immediately (early trigger) so a following GEMM overlaps this kernel; consumer correctness is its own wait_on_dependent_grids. The signal buffers must be sentinel-initialized once (lamport_init) before the first call.

Functions

• ​lamport_allreduce_rmsnorm: Per-device fused Lamport allreduce (sum) + RMSNorm (high-perf protocol).

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