For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

DistributedAllGatherRMSNorm

struct DistributedAllGatherRMSNorm

Registers the mo.composite.distributed.allgather_rms_norm graph op with the graph compiler.

Implemented traits

AnyType, Deinitable, Movable

Methods

execute

static def execute[dtype: DType, rank: Int, target: StringSpan[ImmStaticOrigin], _trace_name: StringSpan[ImmStaticOrigin], group_size: Int = Int(0)](outputs_normed: VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_normed.static_specs], outputs_residual: VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_residual.static_specs], inputs: VariadicTensors[IOSpec[_, _].Input, static_specs=inputs.static_specs], signal_buffers: VariadicTensors[IOSpec[_, _].MutableInput, static_specs=signal_buffers.static_specs], gammas: VariadicTensors[IOSpec[_, _].Input, static_specs=gammas.static_specs], epsilons: VariadicTensors[IOSpec[_, _].Input, static_specs=epsilons.static_specs], weight_offsets: VariadicTensors[IOSpec[_, _].Input, static_specs=weight_offsets.static_specs], dev_ctxs_input: DeviceContextArray)

Fused all-gather + RMSNorm (bf16 in/out, no quantization).

All-gathers inputs (per-device [shard_i, cols] row-shards) into the full replicated [rows, cols] stream and RMSNorms every gathered row in one launch: normed to outputs_normed, gathered residual to outputs_residual (both full and replicated).

Limitations: - Maximum of 8 GPUs supported (matches MAX_GPUS in comm/sync.mojo). - Requires P2P; within a group, shard shapes may differ only in the gathered axis.

Parameters:

• ​dtype (DType): Element type of the input/output tensors.
• ​rank (Int): Tensor rank of the inputs and outputs.
• ​target (StringSpan[ImmStaticOrigin]): Target device string for tracing.
• ​_trace_name (StringSpan[ImmStaticOrigin]): Trace name for profiling.
• ​group_size (Int): Number of contiguous devices per independent all-gather group; must be at least 2 and must evenly divide the total number of devices. Equal to num_devices for a full-world collective. The builder always sets it; the 0 attribute default is not a usable value.

Args:

• ​outputs_normed (VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_normed.static_specs]): Per-device normed output (the group's gathered [rows, cols], replicated within the group).
• ​outputs_residual (VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_residual.static_specs]): Per-device gathered residual (same shape as outputs_normed; the residual stream).
• ​inputs (VariadicTensors[IOSpec[_, _].Input, static_specs=inputs.static_specs]): Per-device input row-shards to gather.
• ​signal_buffers (VariadicTensors[IOSpec[_, _].MutableInput, static_specs=signal_buffers.static_specs]): Per-device synchronization buffers.
• ​gammas (VariadicTensors[IOSpec[_, _].Input, static_specs=gammas.static_specs]): Per-device RMSNorm gamma weights (in_dtype, length cols).
• ​epsilons (VariadicTensors[IOSpec[_, _].Input, static_specs=epsilons.static_specs]): Per-device RMSNorm epsilon scalars (float32).
• ​weight_offsets (VariadicTensors[IOSpec[_, _].Input, static_specs=weight_offsets.static_specs]): Per-device gamma offset scalars (in_dtype).
• ​dev_ctxs_input (DeviceContextArray): Device contexts for participating GPUs.