mstar.communication.tensors

Contents

mstar.communication.tensors#

Module Attributes

MOONCAKE_PATCHED_VERSION

Earlier Mooncake releases let any host that reaches the TCP data port read or write arbitrary process memory (CVE-2026-103764).

Functions

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

Select tensor transport backend based on protocol.

ensure_deployment_token()

Mint this launch's token unless a launcher set one.

Classes

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

Background thread for non-blocking mooncake READ operations.

BufferedShards(total_fanin, shard_dim, shards)

FailedRead(request_id, graph_edges, error)

Edges of a request that will not be delivered because one of its reads failed.

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.

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

Tensor reads that say, when they fail, which transfers may still be writing into their buffers.

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

Exceptions

AsyncReadError(message, running, unknown)

A Mooncake read that failed.

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

Bases: object

Background thread for non-blocking mooncake READ operations.

The caller records a CUDA event on the default stream and submits the write task to a 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)

shutdown()[source]#

Wait for pending writes and shut down the thread pool.

submit(read_info)[source]#

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()[source]#

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

exception mstar.communication.tensors.AsyncReadError(message, running, unknown)[source]#

Bases: RuntimeError

A Mooncake read that failed. running maps the local pointer of each buffer whose transfer had not reported COMPLETED or FAILED to its batch id; the reader stops polling them when it raises, and the manager that keeps the buffers takes them over. unknown holds pointers whose transfer state cannot be learned: a failed submit (which may have posted part of the transfer) or a status this code does not know.

Parameters:
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.FailedRead(request_id, graph_edges, error)[source]#

Bases: object

Edges of a request that will not be delivered because one of its reads failed. error is that read’s exception.

Parameters:
error: BaseException#
graph_edges: list[GraphEdge]#
request_id: int | str#
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, transfer_uuids: list[int] = <factory>)[source]#

Bases: object

Parameters:
count_read(info)[source]#

Count one tensor this read moved over the transport.

Parameters:

info (TensorPointerInfo)

Return type:

None

future: Future | None#
graph_edges: list[GraphEdge]#
read_bytes = 0#
read_more = None#
read_src = None#
read_tensors = 0#
request_id: int | str = ''#
rx_time: float | None = None#
transfer_uuids: list[int]#
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

mstar.communication.tensors.MOONCAKE_PATCHED_VERSION = '0.3.12.post1'#

Earlier Mooncake releases let any host that reaches the TCP data port read or write arbitrary process memory (CVE-2026-103764). NVD lists “before 0.3.13” because it points at a later refactor; the fix (Mooncake #2314) shipped in 0.3.12.post1, see kvcache-ai/Mooncake#4626.

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

register_for_send_uuids(per_request, skip_cuda_sync=False)[source]#

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:
start_read_tensors(request_id, graph_edges, graph_walk=None)[source]#
Parameters:
Return type:

list[Future]

property transport_label: str#

str(object=’’) -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.__str__() (if defined) or repr(object). encoding defaults to sys.getdefaultencoding(). errors defaults to ‘strict’.

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:
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)[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_for_send_uuids(per_request, skip_cuda_sync=False)[source]#

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

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

Unlink every tensor file this manager wrote that is still waiting for its TENSOR_RECEIVED or REMOVE_REQUEST, for a process that is stopping. A second call finds nothing left.

Return type:

None

class mstar.communication.tensors.StoredOutputs(flat_uuids, flat_rids, signal_idxs, num_tensors)[source]#

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_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)[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, 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:
ack_unread_tensors(request_id, graph_edges)[source]#

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

The producer holds these buffers until it gets the TENSOR_RECEIVED ack; emit it here so it can reclaim them. request_id is the wire string: a removed request has no handle on a worker.

Parameters:
buffered_shards: dict[int | str, dict[tuple[str, str], BufferedShards]]#
cleanup_collectable(collectable, registered)[source]#

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:
cleanup_request(request_id, force=False)[source]#

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

Reads of the request still in pending that failed, or have not finished, keep their buffers and are not acked: their transfers may still write into the buffers freed below.

Parameters:
dereference(uuid, n=1)[source]#
Parameters:
dereference_batch(uuids, counts)[source]#

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

Parameters:
dereference_batch_uniform(uuids, n=1)[source]#

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:
drop_host_copies()[source]#

Forget the host copies stored with a step’s outputs once that step’s sends are done. They are views of buffers the producer reuses on its next step, and a tensor can be sent in a later step than the one that stored it (a loop’s accumulated output is sent when the loop ends).

Return type:

None

enable_transfer_metrics()[source]#

Count what this manager reads from now on (see EdgeTransfers). Called once, before or while reads run on the reading thread.

Return type:

EdgeTransfers

force_cleanup_request(request_id)[source]#

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)[source]#
Parameters:

graph_walk (str | None)

Return type:

dict[int | str, list[GraphEdge]]

get_rx_info(request_id)[source]#
Parameters:

request_id (int | str)

Return type:

list[RxInfo]

get_tensor(uuid)[source]#
Parameters:

uuid (int)

Return type:

Tensor

get_tx_info(request_id)[source]#
Parameters:

request_id (int | str)

Return type:

list[TxInfo]

has_inflight_reads(request_id)[source]#

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. So does a failed read whose transfers have not all stopped: they may still be reading the producer’s memory.

Parameters:

request_id (int | str)

Return type:

bool

increment_ref(uuid, n=1)[source]#
Parameters:
increment_ref_batch_uniform(uuids, n=1)[source]#

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

Parameters:
names_unique_per_launch = False#

True when no other launch can write a file under this manager’s names (see SharedMemoryCommunicationManager); irrelevant to transports that write no files.

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

needs_flat_layout: bool = False#

Whether a consumer reads a stored tensor as nbytes contiguous bytes from its address (Mooncake RDMA / TCP), so the store must hold a layout that one flat read reproduces. False for transports that serialize a contiguous copy at send time (SHM). An attribute, not a property: every store call reads it.

pending: list[FutureAndPointers]#
read_finished: dict[int | str, set[int]]#
refresh_shm_placement(edges)[source]#

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)[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:
abstractmethod register_for_send_uuids(per_request, skip_cuda_sync=False)[source]#

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:
register_request(request_id, sharding_config)[source]#
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)[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 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:
set_persist(uuid, persist)[source]#
Parameters:
sharding_configs: dict[int | str, ShardingConfig]#
slice_origins: dict[int | str, dict[int, int]]#
slice_uuids: dict[int | str, dict[tuple[int, int, int, tuple[int, ...]], int]]#
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)[source]#
Parameters:
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)[source]#

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:
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, ready_event=None)[source]#

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.

skip_cuda_sync=True as in store_and_return_tensor_info: the producer’s writes must be complete, unless ready_event (the producing step’s completion event) is given, which a flat-layout copy waits on before it reads them.

Parameters:
Return type:

StoredOutputs

store_row_outputs_batch(request_ids, outputs, signals, row_views, row_request_ids, node_name=None, graph_walk=None, host_rows=None)[source]#

store_and_return_tensor_info_batch for a step whose per-request outputs are views of row-addressed batch tensors (a decode step: row i of new_token is request i).

The rows of one batch tensor differ only in address, so their descriptors are built from one row’s shape, stride and dtype instead of being read off every view, and the host copies come from the stop check’s row buffers (host_rows) by index rather than being matched tensor by tensor. Every (request, signal) entry must be exactly one tensor that is a row view, or the result is None and the caller takes the general path; nothing is recorded before that is settled. What is stored is what the general path would store: one uuid per (request, signal), request-major, the same row views.

Parameters:
Return type:

StoredOutputs | None

take_failed_reads()[source]#

Reads dropped by get_ready_tensors() since the last call. The caller fails their requests; their other reads are dropped as they finish, until cleanup_request().

Return type:

list[FailedRead]

transfer_stats: EdgeTransfers | None = None#
transport_label: str = 'SHM'#

the protocol the bytes move over.

Type:

The transport label

Unlink the shared-memory files this process wrote for its peers and still owns. A no-op here: only the SHM file transport (and the arena’s spills) writes any.

Return type:

None

uuid_to_edge_name: dict[int, str]#
uuid_to_shard_dim: dict[int, int | 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.TrackedMooncakeReader(engine, device, max_workers=3, device_filter='')[source]#

Bases: AsyncMooncakeReader

Tensor reads that say, when they fail, which transfers may still be writing into their buffers.

Each tensor is its own batch, read with batch_transfer_async_read and polled with transfer_check_status, which reports on a batch’s first task only. The sync read gives up on a timeout or on one failed tensor without waiting for the rest of its batch, and returns no handle to them. Retries and the deadline are the sync read’s.

Parameters:
  • max_workers (int)

  • device_filter (str)

transfer_status(batch_id)[source]#

transfer_check_status of one batch this reader handed over.

Parameters:

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

mstar.communication.tensors.ensure_deployment_token()[source]#

Mint this launch’s token unless a launcher set one. Called by the API server, which every other process of the deployment descends from (spawn children inherit the environment).

A token an API server minted is never kept by the next launch that finds it, a server built again in this process or started by one that holds it: on the same endpoint that launch would write the names a stopping worker of the old one may still unlink.

Return type:

str