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!