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

DistributedAllGatherRMSNormQuantMXFP8

struct DistributedAllGatherRMSNormQuantMXFP8

Registers mo.composite.distributed.allgather_rms_norm_quant_mxfp8.

Implemented traits

AnyType, Deinitable, Movable

Methods

execute

static def execute[dtype: DType, quant_dtype: DType, scales_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_quant: VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_quant.static_specs], outputs_scale: VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_scale.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 that also emits an MXFP8 copy.

DistributedAllGatherRMSNorm plus outputs_quant/outputs_scale. The quantize rides the collective's epilogue on the same bf16 that lands in outputs_normed, so it is byte-identical to a standalone quantize.

Limitations: - Targets AMD (CDNA4): the fallback calls quantize_mx_amd and the scale layout is its rank-2 one -- what block_scaled_matmul_amd takes as a_scales, not the SM100 SF atom and not the preshuffled order the _preb kernel needs. - Maximum of 8 GPUs; requires P2P.

Parameters:

• ​dtype (DType): Element type of the bf16 input/normed/residual tensors.
• ​quant_dtype (DType): Quantized element type (float8_e4m3fn).
• ​scales_dtype (DType): Block-scale element type (float8_e8m0fnu).
• ​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): Devices per independent all-gather group; see DistributedAllGatherRMSNorm.

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_quant (VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_quant.static_specs]): Per-device MXFP8 copy of outputs_normed.
• ​outputs_scale (VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_scale.static_specs]): Per-device E8M0 block scales, [rows, cols / 32].
• ​outputs_residual (VariadicTensors[IOSpec[_, _].Output, static_specs=outputs_residual.static_specs]): Per-device gathered residual.
• ​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.
• ​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.
• ​dev_ctxs_input (DeviceContextArray): Device contexts for participating GPUs.