mstar.communication.tensors

Contents

mstar.communication.tensors#

Functions

create_tensor_communication_manager(...[, ...])

Select tensor transport backend based on protocol.

Classes

AsyncMooncakeReader(engine, device[, ...])

Background thread for non-blocking mooncake READ operations.

BufferedShards(total_fanin, shard_dim, shards)

FutureAndPointers(future, graph_edges[, ...])

LocalTransferEngine(hostname)

No-op engine for SHM / single-node — data is already in local GPU memory.

MooncakeCommunicationManager(my_entity_id, ...)

MooncakeTransferEngine(hostname, protocol[, ...])

Wraps mooncake.engine.TransferEngine for RDMA or TCP transport.

SharedMemoryCommunicationManager(...[, ...])

Tensor transport via file I/O to a tmpfs directory (/dev/shm).

TensorAndReferenceInfo(tensor[, ref_cnt, ...])

TensorCommunicationManager(my_entity_id, ...)

Base class for inter-worker tensor transport.

TensorStore()

TensorTransferEngine()

Abstract interface for low-level memory registration and async reads.

TransferReadInfo(source_session_id, ...[, fp])

class mstar.communication.tensors.AsyncMooncakeReader(engine, device, max_workers=3, max_batch_size=500, enable_prof=False)[source]#

Bases: object

Background thread for non-blocking mooncake READ operations.

Follows SGLang’s pattern: caller records CUDA event on default stream, submits write task to thread pool. Worker thread waits on event via dedicated CUDA stream, then does blocking mooncake PUTs. The default stream is never blocked by store writes.

Parameters:

  • max_workers (int)

  • enable_prof (bool)

shutdown()[source]#

Wait for pending writes and shut down the thread pool.

submit(read_info)[source]#

Non-blocking: enqueue a batch of READs.

Records a CUDA event on the current stream to ensure GPU data is ready before the background thread reads it.

Parameters:

read_info (list[TransferReadInfo])

Return type:

Future

wait_all()[source]#

Block until all pending writes complete. Re-raises exceptions.

class mstar.communication.tensors.BufferedShards(total_fanin: int, shard_dim: int, shards: dict[int, list[torch.Tensor]])[source]#

Bases: object

Parameters:
consolidate()[source]#
Return type:

list[Tensor]

is_done()[source]#
shard_dim: int#
shards: dict[int, list[Tensor]]#
total_fanin: int#
class mstar.communication.tensors.FutureAndPointers(future: concurrent.futures._base.Future | None, graph_edges: list[mstar.graph.base.GraphEdge], request_id: str = '', rx_time: float | None = None)[source]#

Bases: object

Parameters:
future: Future | None#
graph_edges: list[GraphEdge]#
request_id: str = ''#
rx_time: float | None = None#
class mstar.communication.tensors.LocalTransferEngine(hostname)[source]#

Bases: TensorTransferEngine

No-op engine for SHM / single-node — data is already in local GPU memory.

Parameters:

hostname (str)

get_async_reader(device)[source]#

Return an async reader for background transfers, or None if not needed.

Return type:

None

get_session_id()[source]#

Return the session ID for this engine (e.g., ‘hostname:port’).

Return type:

str

register_memory(ptr, nbytes)[source]#

Register a memory region for remote access. Returns 0 on success.

Parameters:
Return type:

int

unregister_memory(ptr)[source]#

Unregister a previously registered memory region. Returns 0 on success.

Parameters:

ptr (int)

Return type:

int

class mstar.communication.tensors.MooncakeCommunicationManager(my_entity_id, hostname, device, communicator, protocol=CommProtocol.RDMA, metadata_server='P2PHANDSHAKE', tcp_transfer_device='', enable_prof=False)[source]#

Bases: TensorCommunicationManager

Parameters:
register_for_send(request_id, tensor_infos, skip_cuda_sync=False)[source]#

Mark these tensors ready for remote consumers to RDMA-read.

skip_cuda_sync=True skips the default-stream sync this call normally issues to ensure the source tensors’ writes are visible before their addresses are shared with peers. Callers must have already synced on their own (e.g. before a batched loop) — meant to cut N serialized syncs to 1 when registering many tensors in a row.

If self.enable_prof is set, this should also update self.req_tx_info

start_read_tensors(request_id, graph_edges, graph_walk=None)[source]#
Parameters:
Return type:

list[Future]

class mstar.communication.tensors.MooncakeTransferEngine(hostname, protocol, metadata_server='P2PHANDSHAKE', tcp_transfer_device='')[source]#

Bases: TensorTransferEngine

Wraps mooncake.engine.TransferEngine for RDMA or TCP transport.

Parameters:
get_async_reader(device)[source]#

Return an async reader for background transfers, or None if not needed.

Return type:

AsyncMooncakeReader

get_session_id()[source]#

Return the session ID for this engine (e.g., ‘hostname:port’).

Return type:

str

register_memory(ptr, nbytes)[source]#

Register a memory region for remote access. Returns 0 on success.

Parameters:
Return type:

int

unregister_memory(ptr)[source]#

Unregister a previously registered memory region. Returns 0 on success.

Parameters:

ptr (int)

Return type:

int

class mstar.communication.tensors.SharedMemoryCommunicationManager(my_entity_id, hostname, device, communicator, shm_dir=None, enable_prof=False)[source]#

Bases: TensorCommunicationManager

Tensor transport via file I/O to a tmpfs directory (/dev/shm).

Parameters:
register_for_send(request_id, tensor_infos, skip_cuda_sync=False)[source]#

Mark these tensors ready for remote consumers to RDMA-read.

skip_cuda_sync=True skips the default-stream sync this call normally issues to ensure the source tensors’ writes are visible before their addresses are shared with peers. Callers must have already synced on their own (e.g. before a batched loop) — meant to cut N serialized syncs to 1 when registering many tensors in a row.

If self.enable_prof is set, this should also update self.req_tx_info

Parameters:
start_read_tensors(request_id, graph_edges, graph_walk=None)[source]#
Parameters:
class mstar.communication.tensors.TensorAndReferenceInfo(tensor: torch.Tensor, ref_cnt: int = 0, persist: bool = False, mem_registered: bool = False)[source]#

Bases: object

Parameters:
mem_registered: bool = False#
persist: bool = False#
ref_cnt: int = 0#
tensor: Tensor#
class mstar.communication.tensors.TensorCommunicationManager(my_entity_id, my_session_id, device, communicator, transfer_engine, enable_prof=False)[source]#

Bases: ABC

Base class for inter-worker tensor transport.

Holds common attributes and shared method implementations. Subclasses only need to override __init__, register_for_send, start_read_tensors, and _cleanup_by_uuid.

Parameters:
ack_unread_tensors(request_id, graph_edges)[source]#

Ack result tensors for an already-removed request without reading them.

The producer holds these output buffers until it gets the TENSOR_RECEIVED ack; emit it here so it can reclaim them.

Parameters:
buffered_shards: dict[str, dict[str, BufferedShards]]#
cleanup_request(request_id)[source]#
Parameters:

request_id (str)

dereference(request_id, uuid, n=1)[source]#
Parameters:
get_ready_tensors(graph_walk=None)[source]#
Parameters:

graph_walk (str | None)

Return type:

dict[str, list[GraphEdge]]

get_rx_info(request_id)[source]#
Parameters:

request_id (str)

Return type:

list[RxInfo]

get_tensor(request_id, uuid)[source]#
Parameters:
Return type:

Tensor

get_tx_info(request_id)[source]#
Parameters:

request_id (str)

Return type:

list[TxInfo]

increment_ref(request_id, uuid, n=1)[source]#
Parameters:
pending: list[FutureAndPointers]#
read_finished: dict[str, set[str]]#
abstractmethod register_for_send(request_id, tensor_infos, skip_cuda_sync=False)[source]#

Mark these tensors ready for remote consumers to RDMA-read.

skip_cuda_sync=True skips the default-stream sync this call normally issues to ensure the source tensors’ writes are visible before their addresses are shared with peers. Callers must have already synced on their own (e.g. before a batched loop) — meant to cut N serialized syncs to 1 when registering many tensors in a row.

If self.enable_prof is set, this should also update self.req_tx_info

Parameters:
register_request(request_id, sharding_config)[source]#
Parameters:

  • request_id (str)

  • sharding_config (ShardingConfig)

req_rx_info: dict[str, dict[tuple[str, str], RxInfo]]#
req_tx_info: dict[str, dict[str, TxInfo]]#
set_output_ref_counts(request_id, safety_hold_uuids, routed_edges)[source]#

Adjust ref counts from the safety hold (1) to the actual fanout.

Called after process_node_outputs determines the real routing. safety_hold_uuids is the set of UUIDs that were given ref=1 by store_and_populate_graph_edges(skip_ref_count=True). routed_edges is the flat list of all edges that will actually be consumed (local ingestion, remote send, persist, emit, streaming).

Parameters:
set_persist(request_id, uuid, persist)[source]#
Parameters:
sharding_configs: dict[str, ShardingConfig]#
abstractmethod start_read_tensors(request_id, graph_edges, graph_walk=None)[source]#
Parameters:
Return type:

list[Future]

store_and_populate_graph_edges(request_id, tensors, graph_edges, node_name=None, graph_walk=None, skip_cuda_sync=False, skip_ref_count=False)[source]#
Parameters:
store_and_return_tensor_info(request_id, tensors, node_name=None, graph_walk=None, skip_cuda_sync=False)[source]#
Parameters:
Return type:

dict[str, list[TensorPointerInfo]]

uuid_to_edge_name: dict[str, str]#
uuid_to_shard_dim: dict[str, int | None]#
class mstar.communication.tensors.TensorStore[source]#

Bases: object

can_gc(request_id, uuid)[source]#
Parameters:
Return type:

bool

check_uuid_presence(request_id, uuid)[source]#
Parameters:
dereference(request_id, uuid, n=1)[source]#
Parameters:
get_all_uuids(request_id)[source]#
Parameters:

request_id (str)

Return type:

list[str]

get_tensor(request_id, uuid)[source]#
Parameters:
Return type:

Tensor

increment_ref(request_id, uuid, n=1)[source]#
Parameters:
is_registered(request_id, uuid)[source]#
Parameters:
per_req_tensors: dict[str, dict[str, TensorAndReferenceInfo]]#
put_tensor(request_id, uuid, tensor)[source]#
Parameters:
remove_tensor(request_id, uuid)[source]#
Parameters:
set_metadata(request_id, uuid, persist=None, mem_registered=None)[source]#
Parameters:

  • request_id (str)

  • uuid (str)

  • persist (bool | None)

  • mem_registered (bool | None)

class mstar.communication.tensors.TensorTransferEngine[source]#

Bases: ABC

Abstract interface for low-level memory registration and async reads.

Wraps the transport-specific engine (Mooncake RDMA, local no-op, etc.) so that higher-level code (PagedAllocationManager, TensorCommunicationManager) never imports or depends on a specific transport library.

abstractmethod get_async_reader(device)[source]#

Return an async reader for background transfers, or None if not needed.

Return type:

AsyncMooncakeReader | None

abstractmethod get_session_id()[source]#

Return the session ID for this engine (e.g., ‘hostname:port’).

Return type:

str

abstractmethod register_memory(ptr, nbytes)[source]#

Register a memory region for remote access. Returns 0 on success.

Parameters:
Return type:

int

abstractmethod unregister_memory(ptr)[source]#

Unregister a previously registered memory region. Returns 0 on success.

Parameters:

ptr (int)

Return type:

int

class mstar.communication.tensors.TransferReadInfo(source_session_id: str, local_ptr: int, remote_ptr: int, nbytes: int, fp: mstar.communication.tensors.FutureAndPointers | None = None)[source]#

Bases: object

Parameters:
fp: FutureAndPointers | None = None#
local_ptr: int#
nbytes: int#
remote_ptr: int#
source_session_id: str#
mstar.communication.tensors.create_tensor_communication_manager(protocol, my_entity_id, hostname, device, communicator, metadata_server='P2PHANDSHAKE', tcp_transfer_device='', shm_dir=None, enable_prof=False)[source]#

Select tensor transport backend based on protocol.

For the SHM protocol, MSTAR_SHM_ARENA selects the implementation (see docs/environment_variables.rst): 0 (default) — per-uuid files; 1 — the Rust shared-memory arena (raises if the mstar_rust extension is missing); AUTO — the arena when the extension imports, files otherwise. The flag must match across the deployment: the arena location rides in the tensor descriptors, so a file-transport consumer cannot read an arena producer.

Parameters:
Return type:

TensorCommunicationManager

Contents