Struct DeviceMesh 

pub struct DeviceMesh { /* private fields */ }

A 2-D logical device mesh: tensor-parallel dimension × pipeline-parallel dimension.

Devices are stored in row-major order: device at (tp_rank, pp_rank) has linear index tp_rank + pp_rank * tp_size.

Implementations

impl DeviceMesh

pub fn tensor_parallel(n: usize) -> Self

Create a 1-D tensor-parallel mesh of n simulated devices.

pub fn new(tp_size: usize, pp_size: usize) -> Self

Create a 2-D (tp_size × pp_size) mesh.

Total device count is tp_size * pp_size.

pub fn size(&self) -> usize

Total number of devices in the mesh.

pub fn get(&self, tp_rank: usize, pp_rank: usize) -> Option<&DeviceInfo>

Get the device at tensor-parallel rank tp_rank and pipeline-parallel rank pp_rank.

Returns None if either rank is out of bounds.

pub fn tp_group(&self, pp_rank: usize) -> Vec<&DeviceInfo>

All devices in the tensor-parallel group for a given pp_rank.

Returns an empty vec if pp_rank is out of range.

pub fn pp_group(&self, tp_rank: usize) -> Vec<&DeviceInfo>

All devices in the pipeline-parallel group for a given tp_rank.

Returns an empty vec if tp_rank is out of range.