Unified OmniConnector and KV cache transfer management.
dataclass ¶
Container for KV cache transfer data.
staticmethod ¶
Reconstruct KV cache data from the packed bytes format.
staticmethod ¶
Reconstruct KV cache data from a packed device tensor.
staticmethod ¶
Compatibility alias for callers using the old GPU-specific name.
Convert to a packed device tensor for raw-data connectors.
Bases: RuntimeError
Payload consumed from connector but post-get processing failed — sync retry impossible.
dataclass ¶
Configuration for OmniKVTransferManager.
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
enable_kv_async_prefetch: bool = False
class-attribute instance-attribute ¶
Unified management for OmniConnector and KV cache transfer.
This class encapsulates all KV cache related operations: - Connector initialization and lazy creation - KV cache extraction from GPU blocks - KV cache transfer with retry logic - KV cache receiving with timeout
class-attribute instance-attribute ¶
get_connector = property(lambda self: self.connector)
instance-attribute ¶
recv_stages = (
(str(recv_from), str(config.stage_id))
if recv_from is not None and config.stage_id is not None
else (None, None)
)
instance-attribute ¶
send_stages = (
(str(config.from_stage), str(config.to_stage))
if config.from_stage and config.to_stage
else (None, None)
)
property ¶
Lazily build the parallel topology snapshot on first access.
Consume prefetched KV → apply → distribute; sync-receive on miss (no retry on consume-error).
Consume a prefetched KV payload; (None, 0) on miss/recoverable failure.
Raises KVPrefetchConsumeError when consumed but post-get failed (no fallback).
distribute_kv_cache(
req: Any,
target_device: device | None = None,
*,
received: bool = False,
) -> dict[str, Any] | None
Distribute KV LEADER→FOLLOWER. No-op for LOCAL/pure TP.
classmethod ¶
from_od_config(config: Any) -> OmniKVTransferManager
Create from model or OmniDiffusion config.
classmethod ¶
from_vllm_config(
vllm_config: Any, model_config: Any
) -> OmniKVTransferManager
Create from vllm config with fallback to kv_transfer_config.
handle_finished_requests_kv_transfer(
finished_reqs: dict[str, dict[str, Any]],
kv_caches: list[LayerKV],
block_size: int,
cache_dtype: str,
request_id_resolver: Callable[[str], str] | None = None,
) -> list[str]
Handle KV cache transfer for finished requests.
This method extracts KV cache from GPU blocks and transfers them to the downstream stage via the connector.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
finished_reqs | dict[str, dict[str, Any]] | Dict mapping request_id to {block_ids, seq_len} | required |
kv_caches | list[LayerKV] | List of KV cache (tensor or tuple) per layer | required |
block_size | int | Size of each cache block | required |
cache_dtype | str | Data type of the cache | required |
request_id_resolver | Callable[[str], str] | None | Optional function to resolve global request ID | None |
Returns:
| Type | Description |
|---|---|
list[str] | List of request IDs that were processed |
receive_kv_cache_for_request(
request_id: str,
target_device: device | None = None,
*,
sender_info: dict[str, Any] | None = None,
) -> tuple[dict[str, Any] | None, int]
Receive KV cache for request_id; returns (data dict, size) or (None, 0).
sender_info overrides instance _sender_base_* (per-call endpoint).
receive_multi_kv_cache(
req: Any,
cfg_kv_collect_func: Callable | None = None,
target_device: device | None = None,
) -> bool
Receive primary KV cache and optional CFG companion KV caches.
First receives the primary KV cache (existing logic). Then, if the request carries cfg_kv_request_ids and a model-specific cfg_kv_collect_func is provided, calls it to fetch and attach the companion KV caches to sampling_params.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
req | Any | Request object with request_id and sampling_params. | required |
cfg_kv_collect_func | Callable | None | Model-specific function for collecting CFG KV caches. Signature: (request_id, cfg_request_ids, kv_transfer_manager, target_device) -> dict[str, Any] | None |
target_device | device | None | Device to move tensors to. | None |
Returns:
| Type | Description |
|---|---|
bool | True if primary KV cache was received successfully. |
receive_multi_kv_cache_distributed(
req: Any,
cfg_kv_collect_func: Callable | None = None,
target_device: device | None = None,
) -> bool
Synchronous receive: fetch → distribute → apply (FOLLOWER skips fetch).
Cancel pending prefetches and stop the executor (call on teardown).
start_prefetch(
kv_prefetch_jobs: dict[str, Any] | None,
target_device: device | None = None,
) -> None
Kick off a background KV load (non-blocking). No-op unless prefetch enabled.
update_sender_info(
sender_info: dict[str, Any],
sender_stage_id: str | int | None = None,
) -> None
Update receiver-side sender info before loading remote KV cache.
The orchestrator always reports rank-0's ZMQ port. When TP > 1 the receiver must offset the port so that each TP rank connects to the corresponding sender rank's port.
The base host/port are also stored so that the receive path can construct per-rank metadata for heterogeneous TP scenarios.