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
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]], 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]]): Pointer[Scalar[.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].