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

split_batch()

max.nn.split_batch(devices, input, input_row_offsets, data_parallel_splits)

source

Split a ragged input batch into data parallel batches.

The following example splits a ragged batch of 4 requests across two device references, sending the first two requests to device 0 and the last two to device 1 via data_parallel_splits = [0, 2, 4]:

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

cpu = DeviceRef.CPU()
devices = [DeviceRef.CPU(0), DeviceRef.CPU(1)]
with Graph(
    "split_batch_example",
    input_types=(
        TensorType(DType.float32, ["total_seq_len", 8], device=cpu),
        TensorType(DType.uint32, ["offsets_len"], device=cpu),
        TensorType(DType.int64, [3], device=cpu),
    ),
) as graph:
    input, input_row_offsets, data_parallel_splits = (
        v.tensor for v in graph.inputs
    )
    split_input, split_offsets = split_batch(
        devices, input, input_row_offsets, data_parallel_splits
    )
    graph.output(*split_input, *split_offsets)

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

See split_batch_replicated() for a version of this method that takes replicated inputs and input_row_offsets for each device.

Parameters:

• input (TensorValue) – Input tensor of shape [total_seq_len, …].
• input_row_offsets (TensorValue) – Row offsets tensor indicating batch boundaries.
• 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.
• devices (list[DeviceRef])

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]]