etha.tensor_bus.pair_state

State.

Classes

M2MMap	Mesh to mesh topology using M2MMap (shape-independent).
PairState	State of a registered Pair.

Module Contents

class etha.tensor_bus.pair_state.M2MMap

Bases: msgspec.Struct

Mesh to mesh topology using M2MMap (shape-independent).

M2MMap structure: dict[src_rank, dict[src_idx, list[tuple[dst_rank, dst_idx]]]]

m2m_map: dict[int, dict[tuple, list[tuple[int, tuple]]]] | None = None

source_num_slicers: list[int] | None = None

source_partial_reductions: list[tuple[int, str]] = []

target_num_slicers: list[int] | None = None

class etha.tensor_bus.pair_state.PairState

Bases: msgspec.Struct

State of a registered Pair.

PairState is created once per pair via register_pair() and represents the communication topology. It contains NO tensor data or execution state. All tensor data and execution plans are now stored in BatchState.

local_group: torch.distributed.ProcessGroup

local_is_first: bool

local_name: str

local_ranks: list[int]

m2m_recv: M2MMap | None = None

m2m_send: M2MMap | None = None

pair_group: torch.distributed.ProcessGroup

pair_name: str

pair_size: int

remote_name: str

remote_ranks: list[int]

source_partial_groups: list[tuple[torch.distributed.ProcessGroup, str]] | None = None