For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

pack_ptrs_array

def pack_ptrs_array[ptrs_layout: TensorLayout, //, ptr_type: DType, local_rank_only: Bool = False, n_gpus_per_node: Int = Int(1) if local_rank_only else ptrs_layout.static_shape[Int(0)]](_ptrs: TileTensor[DType.uint64, ptrs_layout, Storage=_ptrs.Storage, address_space=_ptrs.address_space, linear_idx_type=_ptrs.linear_idx_type], my_rank: Int32, out result: Array[Pointer[Scalar[ptr_type], MutUntrackedOrigin, _safe=False], n_gpus_per_node])

Pack the pointers into an inline array.

Reads device addresses from _ptrs and produces an Array of UnsafePointer[Scalar[ptr_type]] entries to pass to an EP kernel. When local_rank_only is set, every entry is filled with the address at my_rank instead of one address per rank.

Parameters:

ptrs_layout (TensorLayout): Layout of the _ptrs input tensor (inferred).
ptr_type (DType): Element DType the packed pointers point to.
local_rank_only (Bool): Whether to replicate only the local rank's pointer across all entries (defaults to False).
n_gpus_per_node (Int): Number of entries in the output array, one per local GPU (defaults to 1 when local_rank_only is set, otherwise ptrs_layout.static_shape[0]).

Args:

_ptrs (TileTensor[DType.uint64, ptrs_layout, Storage=_ptrs.Storage, address_space=_ptrs.address_space, linear_idx_type=_ptrs.linear_idx_type]): 1D uint64 tensor of device addresses, one per rank.
my_rank (Int32): Rank index of the calling device within the communicator.

Returns:

Array[Pointer[Scalar[ptr_type], MutUntrackedOrigin, _safe=False], n_gpus_per_node]