Skip to content

vllm_omni.distributed

Modules:

Name Description
omni_connectors
omni_coordinator
ray_utils

MooncakeConnector module-attribute 

MooncakeConnector = MooncakeStoreConnector

ConnectorSpec dataclass

Specification for a connector instance.

extra class-attribute instance-attribute 

extra: dict[str, Any] = field(default_factory=dict)

name instance-attribute 

name: str

MooncakeStoreConnector

Bases: OmniConnectorBase

Mooncake-based distributed connector for OmniConnector.

config instance-attribute 

config = config

host instance-attribute 

host = config.get('host', '127.0.0.1')

localbuf instance-attribute 

localbuf = config.get('localbuf', 64 * 1024 * 1024)

master instance-attribute 

master = config.get('master', '127.0.0.1:50051')

metadata instance-attribute 

metadata = config.get(
    "metadata_server", "http://127.0.0.1:8080/metadata"
)

pin instance-attribute 

pin: ReplicateConfig | None = None

proto instance-attribute 

proto = config.get('proto', 'tcp')

rdma instance-attribute 

rdma = config.get('rdma', '')

segment instance-attribute 

segment = config.get('segment', 512 * 1024 * 1024)

store instance-attribute 

store: MooncakeDistributedStore | None = None

cleanup 

cleanup(request_id: str) -> None

close 

close()

Clean shutdown.

get 

get(
    from_stage: str,
    to_stage: str,
    get_key: str,
    metadata: dict[str, Any] | None = None,
) -> tuple[Any, int] | None

health 

health() -> dict[str, Any]

put 

put(
    from_stage: str, to_stage: str, put_key: str, data: Any
) -> tuple[bool, int, dict[str, Any] | None]

MooncakeTransferEngineConnector

Bases: OmniConnectorBase

OmniConnector implementation using Mooncake Transfer Engine with a managed memory pool. Supports both CPU (Pinned) and GPU memory pools, and both RDMA and TCP protocols. Topology limitations (current implementation): Current design focuses on peer-to-peer communication between stages. - 1 sender → 1 receiver per key: After a successful RDMA write the sender immediately cleans up the buffer (cleanup()), so only the first receiver to pull a given key will succeed. Broadcast / multicast (1 sender → N receivers sharing the same data) is not yet supported. - 1 receiver → N senders: Supported via partial metadata. The manager constructs metadata with the target sender's source_host / source_port (computed from from_rank) and passes it to get(metadata=...). The connector detects that data_size is missing, queries the specified sender at the given address to fill it in, then performs the RDMA pull. This enables heterogeneous TP (sender TP > receiver TP) where a single receiver must pull KV shards from multiple sender ranks.

Future work
  • Support 1 sender → N receivers (e.g. reference-counted buffers, or explicit retain() / release() semantics so the buffer survives multiple pulls).

buffer_ttl_seconds instance-attribute 

buffer_ttl_seconds = float(
    config.get("buffer_ttl_seconds", _BUFFER_TTL_SECONDS)
)

can_put instance-attribute 

can_put = role == 'sender'

config instance-attribute 

config = config

device_name instance-attribute 

device_name = config.get('device_name', '')

engine_id instance-attribute 

engine_id = str(uuid.uuid4())

host instance-attribute 

host = self._get_local_ip()

pool_device instance-attribute 

pool_device = config.get('memory_pool_device', 'cpu')

pool_size instance-attribute 

pool_size = config.get('memory_pool_size', 1024 ** 3)

protocol instance-attribute 

protocol = config.get('protocol', 'rdma')

sender_host instance-attribute 

sender_host = config.get('sender_host', None)

sender_zmq_port instance-attribute 

sender_zmq_port = config.get('sender_zmq_port', None)

supports_raw_data class-attribute instance-attribute 

supports_raw_data: bool = True

zmq_port instance-attribute 

zmq_port = config.get('zmq_port', 50051)

cleanup 

cleanup(
    request_id: str,
    from_stage: str | None = None,
    to_stage: str | None = None,
) -> None

Release the producer-side buffer associated with the request.

Parameters:

Name Type Description Default
request_id str

The key used in put().

 required
from_stage str | None

Optional source stage. When both from_stage and to_stage are provided the method applies _make_key() internally, so callers can pass the same raw key they used in put() / get() without knowing the internal format. When omitted the request_id is used as-is (suitable for internal callers that already hold the transformed key).

 None
to_stage str | None

Optional destination stage (see from_stage).

 None

close 

close() -> None

Gracefully shutdown the connector and release all resources. This method should be called when the connector is no longer needed. Idempotent: safe to call multiple times.

get 

get(
    from_stage: str,
    to_stage: str,
    get_key: str,
    metadata: dict[str, Any] | None = None,
) -> tuple[Any, int] | None

Consumer Side. Allocates from local pool and pulls data via RDMA.

Metadata resolution:

  1. metadata provided with data_size → use directly (RDMA pull).
  2. metadata provided with source_host/source_port but without data_size → query that specific sender for data_size / is_fast_path, then RDMA pull. This is the heterogeneous-TP path where the manager knows the target sender endpoint but not the payload size.
  3. metadata=None → query the default sender (set via update_sender_info()) for the full metadata.

Returns:

Type Description
tuple[Any, int] | None

(data, size) on success, None on failure.
tuple[Any, int] | None
  • is_fast_path=True (tensor or bytes payload): Returns (ManagedBuffer, size). CALLER MUST call ManagedBuffer.release() after consuming.
tuple[Any, int] | None
  • is_fast_path=False (serialized Python object): Returns (DeserializedObject, size). Buffer is auto-released internally after deserialization.

get_connection_info 

get_connection_info() -> dict[str, Any]

Get connection info for this connector (useful for sender to share with receivers).

health 

health() -> dict[str, Any]

put 

put(
    from_stage: str, to_stage: str, put_key: str, data: Any
) -> tuple[bool, int, dict[str, Any] | None]

Producer Side. Exposes data for RDMA transfer.

Behavior by data type: - ManagedBuffer (from this pool): Zero-copy, is_fast_path=True - ManagedBuffer (different pool): Fallback to tensor copy path - torch.Tensor / bytes: Copy to pool, is_fast_path=True - Other (dict, etc.): Serialize to bytes, copy to pool, is_fast_path=False (receiver will deserialize automatically)

update_sender_info 

update_sender_info(
    sender_host: str,
    sender_zmq_port: int,
    sender_rank: int | None = None,
) -> None

Inject a sender's ZMQ endpoint into the receiver connector.

When sender_rank is None (default), sets the single default sender used by get() when no rank is specified — this preserves backward-compatible 1:1 semantics.

When sender_rank is an integer, the endpoint is stored in a per-rank registry for internal use (e.g. by _query_metadata_from_sender(sender_rank=R)).

MoriTransferEngineConnector

Bases: OmniConnectorBase

OmniConnector backed by the Mori RDMA IOEngine.

Topology: 1 sender ↔ 1 receiver per key (same as MooncakeTransferEngineConnector).

allocator instance-attribute 

allocator = BufferAllocator(self.pool_size, alignment=4096)

backend_type instance-attribute 

backend_type = backend_type_cfg

base_ptr instance-attribute 

base_ptr = self.pool.data_ptr()

can_put instance-attribute 

can_put = role == 'sender'

config instance-attribute 

config = config

device_name instance-attribute 

device_name = ''

engine instance-attribute 

engine = IOEngine(self.engine_key, engine_config)

engine_desc instance-attribute 

engine_desc: EngineDesc = self.engine.get_engine_desc()

engine_desc_packed instance-attribute 

engine_desc_packed: bytes = self.engine_desc.pack()

engine_key instance-attribute 

engine_key = f"omni-{role}-{uuid.uuid4().hex[:8]}-pid{os.getpid()}-{self.host}"

host instance-attribute 

host = self._get_local_ip()

pool instance-attribute 

pool = torch.empty(
    self.pool_size, dtype=torch.uint8
).pin_memory()

pool_device instance-attribute 

pool_device = config.get('memory_pool_device', 'cpu')

pool_mem_desc instance-attribute 

pool_mem_desc: MemoryDesc = (
    self.engine.register_torch_tensor(self.pool)
)

pool_mem_desc_packed instance-attribute 

pool_mem_desc_packed: bytes = self.pool_mem_desc.pack()

pool_size instance-attribute 

pool_size = config.get('memory_pool_size', 1024 ** 3)

role instance-attribute 

role = role

sender_host instance-attribute 

sender_host = config.get('sender_host', None)

sender_zmq_port instance-attribute 

sender_zmq_port = config.get('sender_zmq_port', None)

stage_id instance-attribute 

stage_id = config.get('stage_id', -1)

supports_raw_data class-attribute instance-attribute 

supports_raw_data: bool = True

zmq_ctx instance-attribute 

zmq_ctx = zmq.Context()

zmq_port instance-attribute 

zmq_port = config.get('zmq_port', 50051)

cleanup 

cleanup(
    request_id: str,
    from_stage: str | None = None,
    to_stage: str | None = None,
) -> None

close 

close() -> None

get 

get(
    from_stage: str,
    to_stage: str,
    get_key: str,
    metadata: dict[str, Any] | None = None,
) -> tuple[Any, int] | None

get_connection_info 

get_connection_info() -> dict[str, Any]

health 

health() -> dict[str, Any]

put 

put(
    from_stage: str, to_stage: str, put_key: str, data: Any
) -> tuple[bool, int, dict[str, Any] | None]

update_sender_info 

update_sender_info(
    sender_host: str, sender_zmq_port: int
) -> None

Inject the sender's ZMQ endpoint into the receiver connector.

NixlConnector

Bases: OmniConnectorBase

OmniConnector backed by vLLM's native NIXL wrapper.

This connector intentionally depends on vLLM's optional NIXL integration (vllm.distributed.nixl_utils). It transfers raw tensor payloads directly through NIXL READ operations. Non-tensor Python payloads are serialized with OmniSerializer, packed into a uint8 CPU tensor, and moved through the same NIXL path.

A NIXL READ needs the producer's agent metadata and memory descriptors, which put returns to its caller. Consumers use the ZMQ control plane to claim connector-managed payloads before READ and acknowledge completion, including when metadata is forwarded directly. Producers without a fixed zmq_port use an ephemeral listener advertised in that metadata. Callers passing metadata=None need a configured or request-specific producer endpoint to fetch and claim metadata by key.

config instance-attribute 

config = dict(config or {})

stage_id instance-attribute 

stage_id = int(self.config.get('stage_id', 0))

supports_raw_data class-attribute instance-attribute 

supports_raw_data: bool = True

abandon_get 

abandon_get(get_key: str) -> None

cleanup 

cleanup(request_id: str) -> None

close 

close() -> None

get 

get(
    from_stage: str,
    to_stage: str,
    get_key: str,
    metadata: dict[str, Any] | None = None,
) -> tuple[Any, int] | None

get_connection_info 

get_connection_info() -> dict[str, Any]

Endpoint a consumer needs to reach this producer's handshake socket.

get_with_deadline 

get_with_deadline(
    from_stage,
    to_stage,
    get_key,
    metadata=None,
    *,
    deadline,
)

Bound control-plane and DMA waits without freeing active READ buffers.

health 

health() -> dict[str, Any]

put 

put(
    from_stage: str, to_stage: str, put_key: str, data: Any
) -> tuple[bool, int, dict[str, Any] | None]

update_sender_info 

update_sender_info(
    sender_host: str, sender_zmq_port: int
) -> None

Register the producer handshake endpoint on the consumer side.

Used when the endpoint is only known after both stages have started and therefore cannot be baked into the connector config.

OmniConnectorBase

Bases: ABC

Base class for all OmniConnectors.

supports_raw_data class-attribute instance-attribute 

supports_raw_data: bool = False

abandon_get 

abandon_get(get_key: str) -> None

Retire unresolved discovery attempts; never cancel an active DMA READ.

cleanup abstractmethod 

cleanup(request_id: str) -> None

Clean up resources for a request.

close abstractmethod 

close() -> None

Release resources held by this connector.

Subclasses must implement this to clean up transport-specific resources (connections, memory pools, threads, etc.). Implementations should be idempotent (safe to call multiple times).

deserialize_obj staticmethod 

deserialize_obj(data: bytes) -> Any

Deserialize bytes to Python object using centralized serializer.

get abstractmethod 

get(
    from_stage: str,
    to_stage: str,
    get_key: str,
    metadata: dict[str, Any] | None = None,
) -> tuple[Any, int] | None

Retrieve Python object and payload size (bytes).

Parameters:

Name Type Description Default
from_stage str

Source stage identifier

 required
to_stage str

Destination stage identifier

 required
get_key str

Unique request identifier

 required
metadata dict[str, Any] | None

Optional transport-specific metadata. When provided, the connector uses it directly (e.g. source_host, source_port, data_size) instead of querying the sender. For heterogeneous TP the manager may supply partial metadata (host/port only); the connector will query the sender at that address to fill in data_size.

 None

Returns:

Type Description
tuple[Any, int] | None

Tuple of (Python object, serialized byte size) if found, None otherwise

get_with_deadline 

get_with_deadline(
    from_stage: str,
    to_stage: str,
    get_key: str,
    metadata: dict[str, Any] | None = None,
    *,
    deadline: float,
) -> tuple[Any, int] | None

Receive with a monotonic deadline; blocking backends must override this.

health abstractmethod 

health() -> dict[str, Any]

Return health status and metrics.

put abstractmethod 

put(
    from_stage: str, to_stage: str, put_key: str, data: Any
) -> tuple[bool, int, dict[str, Any] | None]

Store Python object, internal serialization handled by connector.

Parameters:

Name Type Description Default
from_stage str

Source stage identifier

 required
to_stage str

Destination stage identifier

 required
put_key str

Unique request identifier

 required
data Any

Python object to store

 required

Returns:

Name Type Description
tuple tuple[bool, int, dict[str, Any] | None]

(success: bool, serialized_size: int, metadata: Optional[dict]) Metadata may contain transport-specific handles or inline data.

reap_consumed 

reap_consumed() -> None

Release producer bookkeeping for payloads consumed by another process.

serialize_obj staticmethod 

serialize_obj(obj: Any) -> bytes

Serialize a Python object to bytes using centralized serializer.

OmniConnectorFactory

Factory for creating OmniConnectors.

create_connector classmethod 

create_connector(spec: ConnectorSpec) -> OmniConnectorBase

Create a connector from specification.

list_registered_connectors classmethod 

list_registered_connectors() -> list[str]

List all registered connector names.

register_connector classmethod 

register_connector(
    name: str,
    constructor: Callable[
        [dict[str, Any]], OmniConnectorBase
    ],
) -> None

Register a connector constructor.

OmniTransferConfig dataclass

Top-level configuration for OmniConnector system. Members: connectors: A dictionary of connectors, keyed by (from_stage, to_stage). default_connector: The default connector to use if no connector is specified for an edge.

connectors class-attribute instance-attribute 

connectors: dict[tuple[str, str], ConnectorSpec] = field(
    default_factory=dict
)

default_connector class-attribute instance-attribute 

default_connector: ConnectorSpec | None = None

get_connector_for_edge 

get_connector_for_edge(
    from_stage: str, to_stage: str
) -> ConnectorSpec | None

Get connector spec for a specific edge.

has_connector_for_edge 

has_connector_for_edge(
    from_stage: str, to_stage: str
) -> bool

Check if there's a connector configured for the edge.

SharedMemoryConnector

Bases: OmniConnectorBase

Key-addressed local shared-memory connector.

SHM is a local-only transport: it reads/writes POSIX shared memory segments identified purely by key. It does not understand remote-transport metadata such as source_host / source_port (that is the RDMA connector's job). When such metadata is passed in, the connector silently falls back to key-based lookup.

config instance-attribute 

config = config

stage_id instance-attribute 

stage_id = config.get('stage_id', -1)

cleanup 

cleanup(request_id: str) -> None

Unlink the exact key passed to put(), never a request-id prefix.

close 

close() -> None

Unlink all remaining tracked SHM segments.

get 

get(
    from_stage: str,
    to_stage: str,
    get_key: str,
    metadata=None,
) -> tuple[Any, int] | None

get_with_deadline 

get_with_deadline(
    from_stage,
    to_stage,
    get_key,
    metadata=None,
    *,
    deadline,
)

health 

health() -> dict[str, Any]

put 

put(
    from_stage: str, to_stage: str, put_key: str, data: Any
) -> tuple[bool, int, dict[str, Any] | None]

reap_consumed 

reap_consumed() -> None

Bounded round-robin sweep; receivers unlink SHM in another process.

YuanrongConnector

Bases: OmniConnectorBase

Datasystem-based distributed connector for OmniConnector.

client instance-attribute 

client = None

config instance-attribute 

config = config

get_sub_timeout_ms instance-attribute 

get_sub_timeout_ms = max(
    0, int(self.config.get("get_sub_timeout_ms", 1000))
)

set_param instance-attribute 

set_param = SetParam()

cleanup 

cleanup(request_id: str) -> None

close 

close() -> None

get 

get(
    from_stage: str,
    to_stage: str,
    get_key: str,
    metadata: dict[str, Any] | None = None,
) -> tuple[Any, int] | None

health 

health() -> dict[str, Any]

put 

put(
    from_stage: str, to_stage: str, put_key: str, data: Any
) -> tuple[bool, int, dict[str, Any] | None]

YuanrongTransferEngineConnector

Bases: OmniConnectorBase

Pull-based connector backed by Yuanrong transfer_engine.

allocator instance-attribute 

allocator = BufferAllocator(self.pool_size, alignment=4096)

base_ptr instance-attribute 

base_ptr = int(self.pool.data_ptr())

can_put instance-attribute 

can_put = role == 'sender'

config instance-attribute 

config = {
    key: expand_env_value(value)
    for key, value in config.items()
}

device_name instance-attribute 

device_name = _resolve_device_name(
    config.get("device_name", "auto"), self.protocol
)

engine instance-attribute 

engine = TransferEngine()

host instance-attribute 

host = (
    self._get_local_ip()
    if host_config.lower() in AUTO_HOST_VALUES
    else host_config
)

pool instance-attribute 

pool = torch.empty(
    self.pool_size,
    dtype=torch.uint8,
    device=self.pool_device,
)

pool_device instance-attribute 

pool_device = _resolve_pool_device(
    config.get("memory_pool_device", "npu")
)

pool_size instance-attribute 

pool_size = int(config.get('memory_pool_size', 1024 ** 3))

protocol instance-attribute 

protocol = str(config.get('protocol', 'ascend'))

rpc_port instance-attribute 

rpc_port = self._resolve_port(
    config.get("rpc_port"), self.host, "rpc_port"
)

sender_host instance-attribute 

sender_host = config.get('sender_host')

sender_zmq_port instance-attribute 

sender_zmq_port = self._resolve_optional_port(
    sender_zmq_port, "sender_zmq_port"
)

supports_raw_data class-attribute instance-attribute 

supports_raw_data: bool = True

zmq_ctx instance-attribute 

zmq_ctx = zmq.Context()

zmq_port instance-attribute 

zmq_port = self._resolve_port(
    config.get("zmq_port"), self.host, "zmq_port"
)

cleanup 

cleanup(
    request_id: str,
    from_stage: str | None = None,
    to_stage: str | None = None,
) -> None

close 

close() -> None

get 

get(
    from_stage: str,
    to_stage: str,
    get_key: str,
    metadata: dict[str, Any] | None = None,
) -> tuple[Any, int] | None

get_connection_info 

get_connection_info() -> dict[str, Any]

health 

health() -> dict[str, Any]

put 

put(
    from_stage: str, to_stage: str, put_key: str, data: Any
) -> tuple[bool, int, dict[str, Any] | None]

update_sender_info 

update_sender_info(
    sender_host: str, sender_zmq_port: int
) -> None

load_omni_transfer_config 

load_omni_transfer_config(
    config_path: str | Path | None = None,
    config_dict: dict[str, Any] | None = None,
) -> OmniTransferConfig | None

Load OmniTransferConfig from file or dict.