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 function

allgather_rmsnorm_quant

def allgather_rmsnorm_quant[in_dtype: DType, ngpus: Int, in_layout: TensorLayout, in_origin: Origin[mut=in_origin.mut], //, quant_epilogue: def[width: Int](row: Int, col: Int, val: SIMD[in_dtype, width]) capturing thin -> None, domain_id: Int = Int(0)](input_buffers: Array[TileTensor[in_dtype, in_layout, in_origin], ngpus], normed_out: TileTensor[in_dtype, Storage=normed_out.Storage, address_space=normed_out.address_space, linear_idx_type=normed_out.linear_idx_type], sum_out: TileTensor[in_dtype, Storage=sum_out.Storage, address_space=sum_out.address_space, linear_idx_type=sum_out.linear_idx_type], gamma: TileTensor[in_dtype, Storage=gamma.Storage, address_space=gamma.address_space, linear_idx_type=gamma.linear_idx_type], epsilon: Float32, weight_offset: Scalar[in_dtype], rank_sigs: Array[Pointer[Signal, MutAnyOrigin], Int(8)], ctx: DeviceContext, local_rank: Optional[Int] = None)

allgather_rmsnorm that also hands each normed value to an epilogue.

The epilogue sees the bf16 that lands in normed_out, so a folded-in quantizer emits the same bytes one launch fewer. Caller-supplied because comm cannot depend on linalg (which already depends on comm).

Parameters:

• ​in_dtype (DType): Input/output data type (bf16).
• ​ngpus (Int): Number of GPUs participating.
• ​in_layout (TensorLayout): Layout of the input shard TileTensors.
• ​in_origin (Origin[mut=in_origin.mut]): Origin of the input shard TileTensors.
• ​quant_epilogue (def[width: Int](row: Int, col: Int, val: SIMD[in_dtype, width]) capturing thin -> None): Normed-value epilogue; see _allgather_rmsnorm_impl.
• ​domain_id (Int): Barrier counter bank; see _allgather_rmsnorm_impl.

Args:

• ​input_buffers (Array[TileTensor[in_dtype, in_layout, in_origin], ngpus]): Per-GPU input row-shards as TileTensors.
• ​normed_out (TileTensor[in_dtype, Storage=normed_out.Storage, address_space=normed_out.address_space, linear_idx_type=normed_out.linear_idx_type]): This GPU's full normed output [rows, cols].
• ​sum_out (TileTensor[in_dtype, Storage=sum_out.Storage, address_space=sum_out.address_space, linear_idx_type=sum_out.linear_idx_type]): This GPU's full gathered residual [rows, cols].
• ​gamma (TileTensor[in_dtype, Storage=gamma.Storage, address_space=gamma.address_space, linear_idx_type=gamma.linear_idx_type]): RMSNorm gamma weights (1D TileTensor of length cols).
• ​epsilon (Float32): RMSNorm epsilon for numerical stability.
• ​weight_offset (Scalar[in_dtype]): Additive offset for gamma weights.
• ​rank_sigs (Array[Pointer[Signal, MutAnyOrigin], Int(8)]): Per-GPU signal pointers for synchronization.
• ​ctx (DeviceContext): Device context for this GPU.
• ​local_rank (Optional[Int]): Optional group-local rank of THIS GPU.