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).
StoredOutputs(flat_uuids, flat_rids, ...)	One batch's stored outputs, in the layout routing consumes.
TensorCommunicationManager(my_entity_id, ...)	Base class for inter-worker tensor transport.
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)

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()

Wait for pending writes and shut down the thread pool.

submit(read_info)

Non-blocking: enqueue a batch of READs.

On a device, 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()

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]])

Bases: object

Parameters:

• total_fanin (int)

• shard_dim (int)

• shards (dict[int, list[Tensor]])

consolidate()

Return type:

list[Tensor]

is_done()

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: int | str = '', rx_time: float | None = None)

Bases: object

Parameters:

• future (Future | None)

• graph_edges (list[GraphEdge])

• request_id (int | str)

• rx_time (float | None)

future: Future | None

graph_edges: list[GraphEdge]

request_id: int | str = ''

rx_time: float | None = None

class mstar.communication.tensors.LocalTransferEngine(hostname)

Bases: TensorTransferEngine

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

Parameters:

hostname (str)

get_async_reader(device)

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

Return type:

None

get_session_id()

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

Return type:

str

register_memory(ptr, nbytes)

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

Parameters:

• ptr (int)

• nbytes (int)

Return type:

int

unregister_memory(ptr)

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)

Bases: TensorCommunicationManager

Parameters:

• my_entity_id (str)

• hostname (str)

• device (str)

• communicator (BaseCommunicator)

• protocol (CommProtocol)

• metadata_server (str)

• tcp_transfer_device (str)

• enable_prof (bool)

register_for_send(request_id, tensor_infos, skip_cuda_sync=False)

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

register_for_send_uuids(per_request, skip_cuda_sync=False)

Uuid-driven form. Registration here reads nothing off a descriptor but the uuid – the tensor itself comes from the store – so taking uuids avoids rebuilding one per tensor.

Parameters:

• per_request (ParallelList[int | str, list[int]])

• skip_cuda_sync (bool)

start_read_tensors(request_id, graph_edges, graph_walk=None)

Parameters:

• request_id (int | str)

• graph_edges (list[GraphEdge])

• graph_walk (str | None)

Return type:

list[Future]

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

Bases: TensorTransferEngine

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

Parameters:

• hostname (str)

• protocol (CommProtocol)

• metadata_server (str)

• tcp_transfer_device (str)

get_async_reader(device)

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

Return type:

AsyncMooncakeReader

get_session_id()

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

Return type:

str

register_memory(ptr, nbytes)

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

Parameters:

• ptr (int)

• nbytes (int)

Return type:

int

unregister_memory(ptr)

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)

Bases: TensorCommunicationManager

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

Parameters:

• my_entity_id (str)

• hostname (str)

• device (str)

• communicator (BaseCommunicator)

• shm_dir (str | None)

• enable_prof (bool)

property needs_cpu_tensor: bool

Whether this transport sends from host memory, so the store keeps the producer’s host copies (cpu_tensors) for the send to reuse. False for transports that read device memory (TCP / RDMA).

register_for_send(request_id, tensor_infos, skip_cuda_sync=False)

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:

• request_id (int | str)

• tensor_infos (list[TensorPointerInfo])

• skip_cuda_sync (bool)

register_for_send_uuids(per_request, skip_cuda_sync=False)

Uuid-driven form. This path reads nothing off a descriptor but the uuid, so taking uuids avoids rebuilding one per tensor.

Parameters:

• per_request (ParallelList[int | str, list[int]])

• skip_cuda_sync (bool)

start_read_tensors(request_id, graph_edges, graph_walk=None)

Parameters:

• request_id (int | str)

• graph_edges (list[GraphEdge])

• graph_walk (str | None)

class mstar.communication.tensors.StoredOutputs(flat_uuids, flat_rids, signal_idxs, num_tensors)

Bases: NamedTuple

One batch’s stored outputs, in the layout routing consumes.

flat_uuids/flat_rids/signal_idxs are index-parallel, one entry per stored tensor. num_tensors has one entry per (request, signal) pair, request-major then signal order. Returned as columns rather than a dict keyed by request and signal because the only caller flattens it immediately – building the dict just to take it apart again cost more than the store itself.

Parameters:

• flat_uuids (list[int])

• flat_rids (list[int])

• signal_idxs (list[int])

• num_tensors (list[int])

flat_rids: list[int]

Alias for field number 1

flat_uuids: list[int]

Alias for field number 0

num_tensors: list[int]

Alias for field number 3

signal_idxs: list[int]

Alias for field number 2

class mstar.communication.tensors.TensorCommunicationManager(my_entity_id, my_session_id, device, communicator, transfer_engine, enable_prof=False)

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, register_for_send_uuids, start_read_tensors, and _cleanup_by_uuid.

request_id is the worker’s integer rid handle. The api-server data worker drives this same class with string request ids, which behaves identically because nothing here does more than use the id as an opaque key – only rid_to_str cares, and its default (str) is the identity on strings.

Parameters:

• my_entity_id (str)

• my_session_id (str)

• device (str)

• communicator (BaseCommunicator)

• transfer_engine (TensorTransferEngine)

• enable_prof (bool)

ack_unread_tensors(request_id, graph_edges)

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:

• request_id (int | str)

• graph_edges (list[GraphEdge])

buffered_shards: dict[int | str, dict[str, BufferedShards]]

cleanup_collectable(collectable, registered)

Transport-side teardown for tensors the bookkeeper has already dereferenced to zero and forgotten – what a batched dereference, or the graph runtime releasing a node’s consumed inputs, hands back.

Parameters:

• collectable (list[int])

• registered (list[bool])

cleanup_request(request_id, force=False)

Teardown for a request’s tensor state.

Default (soft): drop tensors that are safe to GC, defer any still referenced or persisted. Does NOT force-drop persisted signals — that is the conductor-coordinated hard cleanup, see force_cleanup_request().

Parameters:

• request_id (int | str)

• force (bool)

dereference(uuid, n=1)

Parameters:

• uuid (int)

• n (int)

dereference_batch(uuids, counts)

dereference_batch_uniform where the count differs per uuid: a TENSOR_RECEIVED ack covers as many reads as the edge fanned out to.

Parameters:

• uuids (list[int])

• counts (list[int])

dereference_batch_uniform(uuids, n=1)

dereference for a batch, in one crossing into the bookkeeper.

Per uuid the loop above pays one crossing for the decrement and one for can_gc, plus is_registered and forget_tensor on anything it frees. Here the bookkeeper does all four and hands back only what became collectable, leaving this side the part it cannot do: the shm files, the arena slots and the memory unregistration.

Parameters:

• uuids (list[int])

• n (int)

force_cleanup_request(request_id)

Unconditional teardown: drop every tensor for the request and unlink its SHM, ignoring ref counts and persist markers. Safe only after every reader has confirmed (READS_DONE) it has no in-flight reads for the request. Also reclaims non-persisted buffers whose readers drained before ACKing (which the soft path would otherwise defer forever).

Parameters:

request_id (int | str)

get_ready_tensors(graph_walk=None)

Parameters:

graph_walk (str | None)

Return type:

dict[int | str, list[GraphEdge]]

get_rx_info(request_id)

Parameters:

request_id (int | str)

Return type:

list[RxInfo]

get_tensor(uuid)

Parameters:

uuid (int)

Return type:

Tensor

get_tx_info(request_id)

Parameters:

request_id (int | str)

Return type:

list[TxInfo]

has_inflight_reads(request_id)

Whether any async read for this request is still touching a remote segment. Synchronous (SHM) reads complete inside start_read_tensors with future=None, so they never count here; only outstanding async (RDMA) futures do. Gates the READS_DONE ACK during a drain.

Parameters:

request_id (int | str)

Return type:

bool

increment_ref(uuid, n=1)

Parameters:

• uuid (int)

• n (int)

increment_ref_batch_uniform(uuids, n=1)

increment_ref_batch for a uniform count, which is what the safety hold on a stored output batch always is.

Parameters:

• uuids (list[int])

• n (int)

property needs_cpu_tensor: bool

Whether this transport sends from host memory, so the store keeps the producer’s host copies (cpu_tensors) for the send to reuse. False for transports that read device memory (TCP / RDMA).

pending: list[FutureAndPointers]

read_finished: dict[int | str, set[int]]

refresh_shm_placement(edges)

Re-read arena placement from the store onto descriptors about to be sent.

Registration stamps shm_segment / shm_offset onto the descriptor the STORE holds, and that is the only copy it can reach – register_for_send_uuids is handed uuids, not objects. A SHARDED edge does not carry the store’s copy: fanout_graph_edges clones it per destination to rewrite dims/nbytes/offset, and it clones during routing, which runs BEFORE staging. So the clone never sees the stamp and ships shm_segment=None, which a consumer reads as “the producer spilled this one to a per-uuid file” – and then opens a file the arena producer never wrote.

Called at each point an edge leaves this process, after register_for_send* has run. Only the two placement fields are copied: everything else on the clone is this destination’s slice and must not be overwritten.

Parameters:

edges (list[GraphEdge])

Return type:

None

abstractmethod register_for_send(request_id, tensor_infos, skip_cuda_sync=False)

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:

• request_id (int | str)

• tensor_infos (list[TensorPointerInfo])

• skip_cuda_sync (bool)

abstractmethod register_for_send_uuids(per_request, skip_cuda_sync=False)

Mark every request’s tensors ready for remote consumers, given uuids rather than descriptors.

Registration reads exactly one field off a descriptor – uuid – and takes the tensor itself from the store. Under the Rust bookkeeper, rebuilding a descriptor so the callee can read its uuid costs ~2.2us each (15 attribute crossings plus the dataclass), and the caller already has the uuids. skip_cuda_sync as for register_for_send.

Parameters:

• per_request (ParallelList[int | str, list[int]])

• skip_cuda_sync (bool)

register_request(request_id, sharding_config)

Parameters:

• request_id (int | str)

• sharding_config (ShardingConfig)

req_rx_info: dict[int | str, dict[tuple[str, str], RxInfo]]

req_tx_info: dict[int | str, dict[str, TxInfo]]

rid_to_str: Callable[[int | str], str]

set_output_ref_counts(safety_hold_uuids, routed_edges)

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 increment_ref_batch_uniform after store_and_return_tensor_info_batch. routed_edges is the flat list of all edges that will actually be consumed (local ingestion, remote send, persist, emit, streaming).

Parameters:

• safety_hold_uuids (set[int])

• routed_edges (list[GraphEdge])

set_persist(uuid, persist)

Parameters:

• uuid (int)

• persist (bool)

sharding_configs: dict[int | str, ShardingConfig]

stamps_shm_placement: bool = False

Whether this transport writes placement onto a descriptor AFTER routing has already copied it. Only the SHM arena does (shm_segment / shm_offset); the per-uuid file and Mooncake paths derive everything a consumer needs from uuid and address, so they leave this off and refresh_shm_placement stays free for them.

abstractmethod start_read_tensors(request_id, graph_edges, graph_walk=None)

Parameters:

• request_id (int | str)

• graph_edges (list[GraphEdge])

• graph_walk (str | None)

Return type:

list[Future]

store_and_return_tensor_info(request_id, tensors, node_name=None, graph_walk=None, skip_cuda_sync=False, cpu_tensors=None)

Store one request’s outputs and return their descriptors.

cpu_tensors mirrors tensors with host copies the caller already made; they are kept only if needs_cpu_tensor.

Parameters:

• request_id (int | str)

• tensors (dict[str, list[Tensor]])

• node_name (str | None)

• graph_walk (str | None)

• skip_cuda_sync (bool)

• cpu_tensors (dict[str, list[Tensor]] | None)

Return type:

dict[str, list[TensorPointerInfo]]

store_and_return_tensor_info_batch(request_ids, outputs, signals, node_name=None, graph_walk=None, skip_cuda_sync=False, cpu_tensors=None)

Store a whole output batch and return the columns routing needs.

One pass: the flat columns are filled as the uuids are minted, so no per-request or per-signal dict is built to be taken apart again. The mint order is therefore request-major then signal order, which is the order the columns are in – that is what keeps them index-parallel.

Which form is marshalled depends on the backend. A bookkeeper that pays per tensor to marshal (the Rust one) gets columns of primitives; one that stores the descriptor by reference (the Python one) gets descriptors, since building columns for it would only add work.

Outputs under a name no signal carries are not stored at all: no edge will ever read them, and a stored tensor nothing references is only freed when its request is torn down.

cpu_tensors mirrors outputs with host copies the caller already made (e.g. the stop check’s), matched as in _host_copies_by_tensor; kept only if needs_cpu_tensor.

Parameters:

• request_ids (list[int | str])

• outputs (dict[int | str, dict[str, list[Tensor]]])

• signals (list[str])

• node_name (str | None)

• graph_walk (str | None)

• skip_cuda_sync (bool)

• cpu_tensors (dict[int | str, dict[str, list[Tensor]]] | None)

Return type:

StoredOutputs

uuid_to_edge_name: dict[int, str]

uuid_to_shard_dim: dict[int, int | None]

class mstar.communication.tensors.TensorTransferEngine

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)

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

Return type:

AsyncMooncakeReader | None

abstractmethod get_session_id()

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

Return type:

str

abstractmethod register_memory(ptr, nbytes)

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

Parameters:

• ptr (int)

• nbytes (int)

Return type:

int

abstractmethod unregister_memory(ptr)

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)

Bases: object

Parameters:

• source_session_id (str)

• local_ptr (int)

• remote_ptr (int)

• nbytes (int)

• fp (FutureAndPointers | None)

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)

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:

• protocol (CommProtocol)

• my_entity_id (str)

• hostname (str)

• device (str)

• communicator (BaseCommunicator)

• metadata_server (str)

• tcp_transfer_device (str)

• shm_dir (str | None)

• enable_prof (bool)

Return type:

TensorCommunicationManager