ttnn.reduce_to_root

ttnn.reduce_to_root(input_tensor_l: ttnn.Tensor, input_tensor_s: ttnn.Tensor, input_tensor_m: ttnn.Tensor, root_coord: ttnn.MeshCoordinate, *, scale_fp32: float = 1.0, topology: ttnn.Topology = ttnn.Topology.Linear, output_tensor_l: ttnn.Tensor = None, output_tensor_s: ttnn.Tensor = None, output_tensor_m: ttnn.Tensor = None, intermediate_tensor: ttnn.Tensor = None, input_mux_cores: List[ttnn.CoreCoord] = None) → List[ttnn.Tensor]

Reduce-to-root operation. Performs sdpa tree reduction across 4 devices and stores the output on the root device only.

Parameters:

• input_tensor_l (ttnn.Tensor) – the SDPA values (l) state tensor, sharded across the 4 devices.

• input_tensor_s (ttnn.Tensor) – the SDPA running-sum (s) state tensor, sharded across the 4 devices.

• input_tensor_m (ttnn.Tensor) – the SDPA running-max (m) state tensor, sharded across the 4 devices.

• root_coord (ttnn.MeshCoordinate) – Coordinate of the root device. Should be (1, 0) for the 4-device setup.

Keyword Arguments:

• scale_fp32 (float, optional) – scale applied during the reduction. Defaults to 1.0.

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

• output_tensor_l (ttnn.Tensor, optional) – Preallocated output tensor for values. Defaults to None.

• output_tensor_s (ttnn.Tensor, optional) – Preallocated output tensor for sum. Defaults to None.

• output_tensor_m (ttnn.Tensor, optional) – Preallocated output tensor for max. Defaults to None.

• intermediate_tensor (ttnn.Tensor, optional) – Preallocated intermediate tensor. Defaults to None.

• input_mux_cores (List[ttnn.CoreCoord], optional) – the 4 mux core coordinates used for the reduction. Defaults to None.

Returns:

List[ttnn.Tensor] – the reduced (l, s, m) tensors, each with the same spec as the corresponding input. The results are valid only on the root device.

Supported dtypes and layouts:

Tensor	Dtypes	Layouts
input_tensor_l / _s / _m	BFLOAT16	TILE

reduce_to_root operates on a fixed 4-device line topology with the root at (1, 0). All three input tensors must be sharded; each output preserves the spec of its corresponding input.

Memory Support:

• Sharded: required (L1)

Example

# reduce_to_root runs an SDPA tree-reduction across a fixed 4-device line and stores the
# result on the root device only. The three state tensors (values l, running-sum s,
# running-max m) must be TILE-laid-out, WIDTH_SHARDED, and resident in L1.
num_devices = 4
num_cores = 8
tile = ttnn.Tile((8, 32))

# 8 shard cores laid out as two rows of four.
shard_grid = ttnn.CoreRangeSet(
    {
        ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(0, 3)),
        ttnn.CoreRange(ttnn.CoreCoord(1, 0), ttnn.CoreCoord(1, 3)),
    }
)

def make_state(width_per_core):
    # Per-device tensor is [8, width_per_core * num_cores]; stack one per device and shard
    # dim 0 across the 4-device line so each device holds its own [8, width] slice.
    shard_spec = ttnn.ShardSpec(shard_grid, [8, width_per_core], ttnn.ShardOrientation.ROW_MAJOR)
    mem_config = ttnn.MemoryConfig(
        ttnn.types.TensorMemoryLayout.WIDTH_SHARDED, ttnn.types.BufferType.L1, shard_spec
    )
    per_device = torch.stack(
        [torch.randn([8, width_per_core * num_cores], dtype=torch.bfloat16) for _ in range(num_devices)], dim=0
    )
    return ttnn.from_torch(
        per_device,
        device=mesh_device,
        layout=ttnn.TILE_LAYOUT,
        tile=tile,
        dtype=ttnn.bfloat16,
        memory_config=mem_config,
        mesh_mapper=ttnn.ShardTensorToMesh(mesh_device, dim=0),
    )

l = make_state(128)
s = make_state(32)
m = make_state(32)

# Reduce the three states along the line to the root coordinate; outputs match the input specs.
root_coord = ttnn.MeshCoordinate(1, 0)
out_l, out_s, out_m = ttnn.reduce_to_root(l, s, m, root_coord, scale_fp32=1.0, topology=ttnn.Topology.Linear)
logger.info(f"{out_l.shape} {out_s.shape} {out_m.shape}")