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
cluster_allreduce
def cluster_allreduce[dtype: DType, width: SIMDLength, //, combine_fn: def[dtype: DType, width: SIMDLength](SIMD[dtype, width], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width], cluster_size: Int, need_tail_sync: Bool = True](slot: Pointer[Scalar[dtype], address_space=AddressSpace.SHARED], vals: SIMD[dtype, width]) -> SIMD[dtype, width]
Combines one block-reduced vector across every CTA of a cluster.
The caller reduces within its own CTA first, leaving the CTA's values in thread 0; this function folds the CTAs together and returns the combined vector in every thread of the block. Every CTA gathers every peer rather than reducing to one and broadcasting back, and every CTA folds in rank order, so the result is bit-identical across the cluster -- callers may branch on it.
Give every CTA of the cluster the same slot allocation and pass the
allocation itself, never an offset into one: peer access maps this CTA's
address onto a peer's shared-memory window, and an offset breaks that
mapping even when it is a compile-time constant. The low width elements
are what the peers read; the elements above them carry the combined
result from thread 0 to the rest of the block.
With need_tail_sync (the default) a trailing cluster_sync retires the
slot, so the same slot may serve the next combine. A caller that combines
in a loop can drop the trailing sync and alternate between two slots
instead: a CTA that races ahead then writes the slot its peers are not
reading, and it cannot get further than one combine ahead of any peer.
Parameters:
- dtype (
DType): Element type of the combined vector; must be 32-bit. Inferred. - width (
SIMDLength): Number of elements combined. Inferred. - combine_fn (
def[dtype: DType, width: SIMDLength](SIMD[dtype, width], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]): Associative binary reduction (e.g. add, max, min). - cluster_size (
Int): Number of CTAs in the cluster. With 1 the cross-CTA traffic disappears and the call only publishes thread 0's values to the rest of the block. - need_tail_sync (
Bool): If True, retire the slot with a trailing sync.
Args:
- slot (
Pointer[Scalar[dtype], address_space=AddressSpace.SHARED]): Cluster-invariant shared-memory scratch of2 * widthelements. - vals (
SIMD[dtype, width]): This CTA's block-reduced values, valid in thread 0.
Returns:
SIMD[dtype, width]: The values combined across the cluster, in every thread.