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).

Python function

split_batch_replicated

split_batch_replicated()

max.nn.split_batch_replicated(devices, input, input_row_offsets, input_row_offsets_int64, data_parallel_splits, prefix='')

source

Split a ragged token batch into data parallel batches.

This version takes a list of input and input_row_offsets replicated on each device. Also see split_input for a version of this method that takes a single ragged token batch.

The following example splits a ragged batch of 4 requests that is replicated across two device references. Each device holds a full copy of the input and its row offsets, and data_parallel_splits = [0, 2, 4] assigns the first two requests to device 0 and the last two to device 1:

from max.dtype import DType
from max.graph import DeviceRef, Graph, TensorType
from max.nn.data_parallelism import split_batch_replicated

cpu = DeviceRef.CPU()
devices = [DeviceRef.CPU(0), DeviceRef.CPU(1)]
with Graph(
    "split_batch_replicated_example",
    input_types=(
        TensorType(DType.int64, ["seq_len"], device=devices[0]),
        TensorType(DType.int64, ["seq_len"], device=devices[1]),
        TensorType(DType.uint32, ["offsets_len"], device=devices[0]),
        TensorType(DType.uint32, ["offsets_len"], device=devices[1]),
        TensorType(DType.uint32, ["offsets_len"], device=cpu),
        TensorType(DType.int64, [3], device=cpu),
    ),
) as graph:
    (
        input_0,
        input_1,
        offsets_0,
        offsets_1,
        input_row_offsets_int64,
        data_parallel_splits,
    ) = (v.tensor for v in graph.inputs)
    split_input, split_offsets = split_batch_replicated(
        devices,
        [input_0, input_1],
        [offsets_0, offsets_1],
        input_row_offsets_int64,
        data_parallel_splits,
    )
    graph.output(*split_input, *split_offsets)

This method places the outputs on the devices specified in devices.

Parameters:

• devices (list[DeviceRef]) – List of devices to split the input on.
• input (list[TensorValue]) – List of input token tensors of shape [total_seq_len]. The list must be the same length as the number of devices.
• input_row_offsets (list[TensorValue]) – Row offsets tensor indicating batch boundaries. The list must be the same length as the number of devices.
• input_row_offsets_int64 (TensorValue) – Row offsets tensor indicating batch boundaries. Must be located on CPU.
• data_parallel_splits (TensorValue) – Buffer containing batch splits for each device that must be located on CPU. The size of data_parallel_splits must be equal to the number of devices + 1.
• prefix (str)

Returns:

Tuple of (split_input, split_offsets) where split_input and split_offsets are lists of tensors, one per device.

Return type:

tuple[list[TensorValue], list[TensorValue]]