ttnn.reduce_scatter

ttnn.reduce_scatter(input_tensor: ttnn.Tensor, dim: int, *, cluster_axis: int = None, subdevice_id: ttnn.SubDeviceId | None, memory_config: ttnn.MemoryConfig | None, intermediate_memory_config: ttnn.MemoryConfig | None, output_tensor: ttnn.Tensor | None, num_links: int = None, topology: ttnn.Topology = None, chunks_per_sync: int | None, num_workers_per_link: int | None, num_buffers_per_channel: int | None, compute_kernel_config: ttnn.DeviceComputeKernelConfig = None, use_l1_small_for_semaphores: bool = False) → ttnn.Tensor

Reduce-scatter operation across devices along a selected dimension and optional cluster axis. This operation reduces the mesh tensor across the devices in the mesh, along the specified dimension. It then scatters the reduced tensor back to the devices in the mesh, along the same dimension. When cluster axis is specified, we reduce and scatter along the cluster axis. When it is not specified, then we reduce and scatter across all devices in the mesh. When the layout is row-major or the scatter breaks apart tiles, we use the composite reduce-scatter implementation that falls back to all-broadcast.

Parameters:

• input_tensor (ttnn.Tensor)Input tensor to be reduced and scattered.

• dim (int)Dimension along which to reduce.

Keyword Arguments:

• cluster_axis (int, optional)The cluster axis to reduce across. Defaults to None.

• subdevice_id (ttnn.SubDeviceId, optional)Subdevice id for worker cores.

• memory_config (ttnn.MemoryConfig, optional)Output memory configuration.

• intermediate_memory_config (ttnn.MemoryConfig, optional)Memory configuration for intermediate buffer used within the operation.

• output_tensor (ttnn.Tensor, optional)Preallocated output tensor.

• num_links (int, optional)The number of links to use for the reduce-scatter operation. Defaults to None, for which the number of links is determined automatically.

• topology (ttnn.Topology, optional)Fabric topology. Defaults to None.

• chunks_per_sync (int, optional)Hyperparameter.

• num_workers_per_link (int, optional)Hyperparameter.

• num_buffers_per_channel (int, optional)Hyperparameter.

• compute_kernel_config (ttnn.DeviceComputeKernelConfig, optional)Compute kernel configuration for the reduction. Defaults to None.

• use_l1_small_for_semaphores (bool, optional)If True, allocate internal global semaphores in L1_SMALL instead of L1 to reduce L1 fragmentation. Defaults to False.

Returns:

ttnn.TensorThe reduced and scattered tensor, with output_shape = input_shape for all the unspecified dimensions, and output_shape[dim] = input_shape[dim] / num_devices, where num_devices is the number of devices along the cluster_axis if specified, else the total number of devices along the mesh.

Supported dtypes and layouts:

Dtypes	Layouts
BFLOAT16, BFLOAT8_B, FLOAT32	TILE, ROW_MAJOR

Input must be rank 2 or greater and the scatter dim must be divisible by the number of devices (for tiled inputs scattered on the last two dims, by num_devices * tile_size). The output preserves the input dtype. Row-major inputs, and tiled inputs whose scatter breaks a tile, are routed through the composite implementation.

Memory Support:

• Interleaved: DRAM and L1

• Sharded: WIDTH_SHARDED, HEIGHT_SHARDED and ND_SHARDED may reside in DRAM or L1; BLOCK_SHARDED must be in L1.

Example

# Each device holds a full-width slice; the scatter dim must be divisible by num_devices.
torch_input = torch.randn([1, 1, 32, 256], dtype=torch.bfloat16)
tt_input = ttnn.from_torch(
    torch_input,
    dtype=ttnn.bfloat16,
    layout=ttnn.TILE_LAYOUT,
    device=mesh_device,
    mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
)

# Sum-reduce across devices then scatter along dim 3: output[dim] = input[dim] / num_devices.
output = ttnn.reduce_scatter(tt_input, dim=3, cluster_axis=1)
logger.info(output.shape)  # [1, 1, 32, 128]