core.resharding.nvshmem_copy_service.nvshmem_types#

Module Contents#

Classes#

SendRequest

Container for a send operation request.

ReceiveRequest

Container for a receive operation request.

WorkloadGroup

Container for a group of send requests to a specific destination PE.

ScheduledBatch

Metadata for a scheduled communication batch.

WorkloadSummary

Summary of a workload group for communication with other PEs.

TransferMetadata

GPU-resident metadata for communication tasks.

Data#

API#

core.resharding.nvshmem_copy_service.nvshmem_types.MAX_SEGMENT_SIZE#

None

core.resharding.nvshmem_copy_service.nvshmem_types.MAX_TASKS_PER_BATCH#

10000

class core.resharding.nvshmem_copy_service.nvshmem_types.SendRequest#

Container for a send operation request.

task_id: int#

None

src_tensor: Any#

None

src_pos: int#

None

size: int#

None

dest_pe: int#

None

class core.resharding.nvshmem_copy_service.nvshmem_types.ReceiveRequest#

Container for a receive operation request.

task_id: int#

None

dest_tensor: Any#

None

dest_pos: int#

None

size: int#

None

src_pe: int#

None

class core.resharding.nvshmem_copy_service.nvshmem_types.WorkloadGroup#

Container for a group of send requests to a specific destination PE.

dest_pe: int#

None

tasks: List[core.resharding.nvshmem_copy_service.nvshmem_types.SendRequest]#

‘field(…)’

total_size: int#

0

class core.resharding.nvshmem_copy_service.nvshmem_types.ScheduledBatch#

Metadata for a scheduled communication batch.

src_pe: int#

None

dest_pe: int#

None

batch_index: int#

None

iteration: int#

None

gpu_plan: Any#

None

tasks: List[core.resharding.nvshmem_copy_service.nvshmem_types.SendRequest]#

‘field(…)’

total_size: int#

0

tasks_summary: Any#

None

class core.resharding.nvshmem_copy_service.nvshmem_types.WorkloadSummary#

Summary of a workload group for communication with other PEs.

total_size: int#

None

task_ids: List[int]#

None

task_sizes: List[int]#

None

class core.resharding.nvshmem_copy_service.nvshmem_types.TransferMetadata#

GPU-resident metadata for communication tasks.

ptrs: Any#

None

sizes: Any#

None

num_tasks: int#

None

total_size: int#

None