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 struct
LamportAllreduceRMSNorm
struct LamportAllreduceRMSNorm
Per-rank fused Lamport allreduce + RMSNorm (high-perf protocol).
Built on comm.allreduce_lamport_rmsnorm; the Lamport comm region is
embedded in Signal (Signal.lamport_region), so the caller sizes/
initializes the signal buffers as sizeof(Signal) bytes. ngpus is
inferred from the number of signal buffers passed in (must be in [2, 8]).
Shape constraints (also checked inside lamport_allreduce_rmsnorm, but
surfaced here for op-level diagnostics):
cols % atomic_width == 0(whole 128-bit Lamport packs only;atomic_width = 16 / size_of[dtype], so bf16 needscols % 8 == 0).cols / atomic_width <= BLOCK_SIZE(one pack per thread / row;BLOCK_SIZE = floor(max_tpb / WARP_SIZE) * WARP_SIZE, e.g. 1024 on Hopper/Blackwell, capping bf16 hidden at 8192).
Implemented traits
Methods
execute
static def execute[target: StringSlice[ImmStaticOrigin], my_rank: Int, pdl: Bool = True, early_launch: Bool = True](output: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec], act: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=act.static_spec], gamma: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gamma.static_spec], signal_buffers: VariadicTensors[IOSpec[_, _].MutableInput, static_specs=signal_buffers.static_specs], ctx: DeviceContext)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!