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

run_radix_sort_pairs_gpu

def run_radix_sort_pairs_gpu[dtype: DType, out_idx_type: DType, ascending: Bool = False, BLOCK_SIZE: Int = Int(256), NUM_BITS_PER_PASS: Int = Int(4)](ctx: DeviceContext, mut keys: DoubleBuffer[dtype], mut key_ids: DoubleBuffer[out_idx_type], skip_sort: Pointer[Scalar[DType.bool], _safe=False], in_shape: IndexList[element_type=in_shape.element_type])

Runs a multi-pass radix sort on key/index pairs across batches on the GPU using double buffering.

Parameters:

dtype (DType): DType - Data type of the keys to sort.
out_idx_type (DType): DType - Data type of the output indices.
ascending (Bool): Bool - Whether to sort in ascending order (default is descending).
BLOCK_SIZE (Int): Int - Number of threads per block (default 256, found empirically).
NUM_BITS_PER_PASS (Int): Int - Number of radix bits processed per pass (default 4).

Args:

ctx (DeviceContext): DeviceContext - The GPU device context for enqueuing kernels.
keys (DoubleBuffer[dtype]): DoubleBuffer[dtype] - Double buffer holding the keys to sort, swapped each pass.
key_ids (DoubleBuffer[out_idx_type]): DoubleBuffer[out_idx_type] - Double buffer holding the key indices, swapped each pass.
skip_sort (Pointer[Scalar[DType.bool], _safe=False]): UnsafePointer[Scalar[DType.bool]] - Per-batch flag indicating whether sorting is skipped.
in_shape (IndexList[element_type=in_shape.element_type]): IndexList - Shape of the input tensor as [batch_size, vocab_size].