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

shard_and_stack

def shard_and_stack[axis: Int](outputs: VariadicTensors[IOSpec[_, _].Output, static_specs=outputs.static_specs], inputs: VariadicTensors[IOSpec[_, _].Input, static_specs=inputs.static_specs], dev_ctxs_input: DeviceContextArray)

Shard weight tensors across multiple devices for tensor parallelism.

This operation takes multiple input tensors with identical shapes and shards them along a specified axis, distributing the shards to different devices (typically GPUs for tensor parallel inference).

Parameters:

axis (Int): The dimension along which to shard the weights.

Args:

outputs (VariadicTensors[IOSpec[_, _].Output, static_specs=outputs.static_specs]): Output tensors, one per device/shard.
inputs (VariadicTensors[IOSpec[_, _].Input, static_specs=inputs.static_specs]): Input tensors to be sharded, all with identical shapes.
dev_ctxs_input (DeviceContextArray): Device contexts for multi-device transfers.