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

scatter_pull_kernel

def scatter_pull_kernel[dtype: DType, BLOCK_SIZE: Int, ngpus: Int, tp_size: Int, dp_size: Int, simd_width: Int = simd_width_of[dtype, get_gpu_target()]()](output_ptr: Pointer[Scalar[dtype], MutAnyOrigin], input_ptrs: Array[Pointer[Scalar[dtype], ImmutAnyOrigin], dp_size], chunk_num_elems: Array[Int32, dp_size], rank_sigs: Array[Pointer[Signal, MutAnyOrigin], Int(8)], my_rank: Int32)

Pull-based scatter+broadcast: each GPU reads its chunk from root.

Each GPU determines its replica index (my_rank // tp_size), then copies from input_ptrs[replica] on the root GPU to its own output buffer.