etha.comm.ir

Intermediate Representation for tensor transfer operations.

Attributes

logger

Classes

Bucket	Bucket for transfer operations (byte-based buffer).
BucketEntry	Bucket offset entry (byte-based).
Chunk	Unified chunk for all transfer operations.

Module Contents

class etha.comm.ir.Bucket

Bases: etha.comm.transfer.Transferable

Bucket for transfer operations (byte-based buffer).

finalize() → None

Finalize communication and cleanup.

is_complete() → bool

Check if communication is complete.

Returns:

True if complete, False otherwise.

launch() → bool

Launch communication operation.

Returns:

True if launched, False if still waiting for buffer to be ready.

prepare() → None

Prepare buffer for communication.

Source-side Partial reduce and dtype cast both live inside Chunk.prepare; this method only assembles entries into the bucket buffer.

buffer_ready_event: torch.cuda.Event | None = None

device: torch.device | None = None

entries: list[BucketEntry]

key: tuple

total_bytes: int

class etha.comm.ir.BucketEntry

Bucket offset entry (byte-based).

chunk: Chunk

nbytes: int

offset: int

class etha.comm.ir.Chunk

Bases: etha.comm.transfer.Transferable

Unified chunk for all transfer operations.

finalize() → None

Finalize communication and cleanup.

prepare(contiguous: bool = True) → None

Prepare source/target buffer.

Source side performs (in order): slice → in-place all-reduce on Partial sub-groups (in source dtype) → cast to transfer_dtype. Reducing before the cast matches DTensor Partial → Replicate semantics; running the all-reduce in the (possibly lower-precision) wire dtype would change numerical results.

property bucket_key: tuple

Return bucket grouping key.

cell_key is added for SHADOW Partial chunks only — they all share dst_ranks=() and would otherwise bundle across cells, making the bucket’s all_reduce sequence per-rank-specific and out of sync with peer ranks whose matching cells live in separate buckets. PRIMARY chunks already differ by dst_ranks across cells.

chunk_shape: tuple[int, Ellipsis]

dst_idx: tuple | None = None

slice_tuples: tuple[slice, Ellipsis] = ()

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

src_idx: tuple

src_slice_tuples: tuple[slice, Ellipsis] | None = None

tensor: torch.Tensor | None = None

transfer_dtype: torch.dtype | None = None

etha.comm.ir.logger