mstar.communication.tensors#
Module Attributes
Earlier Mooncake releases let any host that reaches the TCP data port read or write arbitrary process memory (CVE-2026-103764). |
Functions
|
Select tensor transport backend based on protocol. |
Mint this launch's token unless a launcher set one. |
Classes
|
Background thread for non-blocking mooncake READ operations. |
|
|
|
Edges of a request that will not be delivered because one of its reads failed. |
|
|
|
No-op engine for SHM / single-node — data is already in local GPU memory. |
|
|
|
Wraps mooncake.engine.TransferEngine for RDMA or TCP transport. |
|
Tensor transport via file I/O to a tmpfs directory ( |
|
One batch's stored outputs, in the layout routing consumes. |
|
Base class for inter-worker tensor transport. |
Abstract interface for low-level memory registration and async reads. |
|
|
Tensor reads that say, when they fail, which transfers may still be writing into their buffers. |
|
Exceptions
|
A Mooncake read that failed. |
- class mstar.communication.tensors.AsyncMooncakeReader(engine, device, max_workers=3, max_batch_size=500)[source]#
Bases:
objectBackground 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)
- 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
- exception mstar.communication.tensors.AsyncReadError(message, running, unknown)[source]#
Bases:
RuntimeErrorA Mooncake read that failed.
runningmaps 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.unknownholds 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.
- class mstar.communication.tensors.BufferedShards(total_fanin: int, shard_dim: int, shards: dict[int, list[torch.Tensor]])[source]#
Bases:
object
- class mstar.communication.tensors.FailedRead(request_id, graph_edges, error)[source]#
Bases:
objectEdges of a request that will not be delivered because one of its reads failed.
erroris that read’s exception.- Parameters:
error (BaseException)
- error: BaseException#
- 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
- read_bytes = 0#
- read_more = None#
- read_src = None#
- read_tensors = 0#
- class mstar.communication.tensors.LocalTransferEngine(hostname)[source]#
Bases:
TensorTransferEngineNo-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:
- 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:
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)[source]#
Mark these tensors ready for remote consumers to RDMA-read.
skip_cuda_sync=Trueskips 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.
- 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:
TensorTransferEngineWraps mooncake.engine.TransferEngine for RDMA or TCP transport.
- Parameters:
hostname (str)
protocol (CommProtocol)
metadata_server (str)
tcp_transfer_device (str)
- get_async_reader(device)[source]#
Return an async reader for background transfers, or None if not needed.
- Return type:
- get_session_id()[source]#
Return the session ID for this engine (e.g., ‘hostname:port’).
- Return type:
Bases:
TensorCommunicationManagerTensor transport via file I/O to a tmpfs directory (
/dev/shm).- Parameters:
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).
Mark these tensors ready for remote consumers to RDMA-read.
skip_cuda_sync=Trueskips 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:
tensor_infos (list[TensorPointerInfo])
skip_cuda_sync (bool)
Uuid-driven form. This path reads nothing off a descriptor but the uuid, so taking uuids avoids rebuilding one per tensor.
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:
NamedTupleOne batch’s stored outputs, in the layout routing consumes.
flat_uuids/flat_rids/signal_idxsare index-parallel, one entry per stored tensor.num_tensorshas 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:
- class mstar.communication.tensors.TensorCommunicationManager(my_entity_id, my_session_id, device, communicator, transfer_engine, enable_prof=False)[source]#
Bases:
ABCBase 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_idis 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 – onlyrid_to_strcares, 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)[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_idis the wire string: a removed request has no handle on a worker.
- 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.
- 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
pendingthat failed, or have not finished, keep their buffers and are not acked: their transfers may still write into the buffers freed below.
- dereference_batch(uuids, counts)[source]#
dereference_batch_uniformwhere the count differs per uuid: a TENSOR_RECEIVED ack covers as many reads as the edge fanned out to.
- dereference_batch_uniform(uuids, n=1)[source]#
dereferencefor a batch, in one crossing into the bookkeeper.Per uuid the loop above pays one crossing for the decrement and one for
can_gc, plusis_registeredandforget_tensoron 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.
- 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:
- 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).
- 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_tensorswithfuture=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.
- increment_ref_batch_uniform(uuids, n=1)[source]#
increment_ref_batchfor a uniform count, which is what the safety hold on a stored output batch always is.
- 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
nbytescontiguous bytes from itsaddress(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]#
- refresh_shm_placement(edges)[source]#
Re-read arena placement from the store onto descriptors about to be sent.
Registration stamps
shm_segment/shm_offsetonto the descriptor the STORE holds, and that is the only copy it can reach –register_for_send_uuidsis handed uuids, not objects. A SHARDED edge does not carry the store’s copy:fanout_graph_edgesclones 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 shipsshm_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.
- 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=Trueskips 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:
tensor_infos (list[TensorPointerInfo])
skip_cuda_sync (bool)
- 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_syncas forregister_for_send.
- 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_outputsdetermines the real routing.safety_hold_uuidsis the set of UUIDs that were given ref=1 byincrement_ref_batch_uniformafterstore_and_return_tensor_info_batch.routed_edgesis the flat list of all edges that will actually be consumed (local ingestion, remote send, persist, emit, streaming).
- 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 fromuuidandaddress, so they leave this off andrefresh_shm_placementstays free for them.
- 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_tensorsmirrorstensorswith host copies the caller already made; they are kept only ifneeds_cpu_tensor.
- 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_tensorsmirrorsoutputswith host copies the caller already made (e.g. the stop check’s), matched as in_host_copies_by_tensor; kept only ifneeds_cpu_tensor.skip_cuda_sync=Trueas instore_and_return_tensor_info: the producer’s writes must be complete, unlessready_event(the producing step’s completion event) is given, which a flat-layout copy waits on before it reads them.
- 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_batchfor a step whose per-request outputs are views of row-addressed batch tensors (a decode step: row i ofnew_tokenis 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.
- 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, untilcleanup_request().- Return type:
- transfer_stats: EdgeTransfers | None = None#
- class mstar.communication.tensors.TensorTransferEngine[source]#
Bases:
ABCAbstract 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:
- class mstar.communication.tensors.TrackedMooncakeReader(engine, device, max_workers=3, device_filter='')[source]#
Bases:
AsyncMooncakeReaderTensor 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_readand polled withtransfer_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.
- 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:
source_session_id (str)
local_ptr (int)
remote_ptr (int)
nbytes (int)
fp (FutureAndPointers | None)
- fp: FutureAndPointers | None = None#
- 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_ARENAselects the implementation (seedocs/environment_variables.rst):0(default) — per-uuid files;1— the Rust shared-memory arena (raises if themstar_rustextension 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:
- 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 (
spawnchildren 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: