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

DistributedReduceScatterRMSNorm

struct DistributedReduceScatterRMSNorm

Registers the mo.composite.distributed.reduce_scatter_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), has_residual: Bool = False](outputs_normed: VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_normed.static_specs], outputs_sum: VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_sum.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], residuals: VariadicTensors[IOSpec[_, _].Input, static_specs=residuals.static_specs], dev_ctxs_input: DeviceContextArray)

Fused reduce-scatter sum + RMSNorm + residual add (bf16, no quant).

Reduce-scatters inputs (one [rows, cols] tensor per device) along rows, adds residuals, and RMSNorm-normalizes each owned shard in the same launch, writing the normed shard to outputs_normed and the summed shard (the residual stream) to outputs_sum.

Under has_residual, residuals carries the TP-replicated residual stream. Each device adds only its own row shard of it, which is why such callers must NOT pre-add it on the group leader: the reduce-scatter sums across ranks, so a leader-side add lands once for the whole group and this per-rank add reproduces it exactly -- but only because the residual is bit-identical on every rank of the group. Folding it here deletes a full-width elementwise add that ran on the group leader alone.

Without has_residual the operands are still passed (the op's variadic groups must match in size) but never read, and both arms are exactly the reduce-scatter + norm this op was before the fold existed.

Limitations: - Maximum of 8 GPUs supported (matches MAX_GPUS in comm/sync.mojo). - Requires P2P, and identical tensor shapes within each group.

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 reduce-scatter 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.
• ​has_residual (Bool): Fold residuals into the reduce-scatter sum. Off leaves both arms byte-for-byte the plain reduce-scatter + norm.

Args:

• ​outputs_normed (VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_normed.static_specs]): Per-device normed output shards.
• ​outputs_sum (VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_sum.static_specs]): Per-device reduce-scatter sum shards (residual stream).
• ​inputs (VariadicTensors[IOSpec[_, _].Input, static_specs=inputs.static_specs]): Per-device input tensors to reduce and scatter.
• ​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).
• ​residuals (VariadicTensors[IOSpec[_, _].Input, static_specs=residuals.static_specs]): Per-device residual stream, same shape as inputs and bit-identical across each group. Read only under has_residual.
• ​dev_ctxs_input (DeviceContextArray): Device contexts for participating GPUs.