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_allgather
def cluster_allgather[dtype: DType, width: SIMDLength, //, cluster_size: Int, need_tail_sync: Bool = True](slot: Pointer[Scalar[dtype], address_space=AddressSpace.SHARED], vals: SIMD[dtype, width]) -> SIMD[dtype, (width * SIMDLength(next_power_of_two(cluster_size)))]
Gathers one block-reduced vector from every CTA of a cluster.
cluster_allreduce folds the per-CTA vectors into one; this keeps them
apart, for callers that need each rank's contribution -- a prefix over
the ranks of a cluster, for example. The slot rules are the same: give
every CTA the same allocation and pass the allocation itself, never an
offset into one. The low width elements are what the peers read; the
elements above them carry the gathered table from thread 0 to the rest
of the block, which peers never touch.
need_tail_sync works as on cluster_allreduce: keep the default to
reuse one slot across consecutive calls, or drop it and alternate
between two slots.
Parameters:
- dtype (
DType): Element type of the gathered vector; must be 32-bit. Inferred. - width (
SIMDLength): Number of elements each CTA contributes. Inferred. - 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 ofwidth * (1 + next_power_of_two(cluster_size))elements. - vals (
SIMD[dtype, width]): This CTA's block-reduced values, valid in thread 0.
Returns:
SIMD[dtype, (width * SIMDLength(next_power_of_two(cluster_size)))]: The per-rank values in every thread, rank-major: element
r * width + i holds rank r's vals[i]. Elements for ranks at or
beyond cluster_size are zero.