# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unified OmniConnector and KV cache transfer management.""" import json import struct import time from collections.abc import Callable from dataclasses import asdict, dataclass from typing import Any import torch from vllm.logger import init_logger from vllm_omni.diffusion.request import OmniDiffusionRequest from .factory import OmniConnectorFactory from .utils.config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec from .utils.env import expand_env_int from .utils.initialization import KV_RANK_PORT_STRIDE from .utils.kv_utils import ( KVTPTopology, build_rank_aware_recv_keys, build_rank_aware_send_keys, get_kv_target_ranks, get_local_tp_rank, get_omni_replica_id, get_tp_world_size, kv_zmq_port, merge_received_rank_shards, normalize_layer_kv, slice_layer_blocks, slice_received_rank_shard, ) logger = init_logger(__name__) LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor] _SAFE_TORCH_DTYPES = { name: dtype for name in ( "bool", "uint8", "int8", "int16", "int32", "int64", "float16", "float32", "float64", "bfloat16", "complex64", "complex128", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz", ) if isinstance((dtype := getattr(torch, name, None)), torch.dtype) } @dataclass class OmniKVCacheConfig: """Configuration for OmniKVTransferManager.""" connector_config: dict[str, Any] | None = None from_stage: str | None = None to_stage: str | None = None stage_id: str | int | None = None engine_input_source: list[str | int] | None = None need_recv_cache: bool = False need_send_cache: bool = False recv_timeout: float = 30.0 from_tp: int = 1 to_tp: int = 1 @dataclass class KVCacheTransferData: """Container for KV cache transfer data.""" request_id: str layer_blocks: dict[str, Any] block_ids: list[int] metadata: dict[str, Any] def to_dict(self) -> dict[str, Any]: """Convert to dictionary for serialization.""" return asdict(self) def _build_tensors_desc(self, *, cpu: bool) -> tuple[list[dict[str, Any]], list, int, torch.device | None]: """Iterate layer blocks and build tensor descriptors + data chunks. Returns ``(tensors_desc, chunks, total_bytes, device)``. *chunks* contains ``bytes`` when *cpu* is True, flat uint8 GPU tensors otherwise. """ tensors_desc: list[dict[str, Any]] = [] chunks: list = [] data_offset = 0 device = None for cache_name in ("key_cache", "value_cache"): for layer_idx, tensor in enumerate(self.layer_blocks.get(cache_name, [])): if tensor is None: tensors_desc.append({"n": f"{cache_name}_{layer_idx}", "x": True}) continue t = tensor.detach().contiguous() if cpu: t = t.cpu() elif device is None and getattr(t.device, "type", "cpu") != "cpu": device = t.device nbytes = t.numel() * t.element_size() tensors_desc.append( { "n": f"{cache_name}_{layer_idx}", "i": layer_idx, "d": str(t.dtype).removeprefix("torch."), "s": list(t.shape), "o": data_offset, "b": nbytes, } ) chunks.append(t.view(torch.uint8).numpy().tobytes() if cpu else t.view(torch.uint8).flatten()) data_offset += nbytes return tensors_desc, chunks, data_offset, device def _build_header_bytes(self, tensors_desc: list[dict[str, Any]]) -> bytes: header = json.dumps( { "rid": self.request_id, "bids": self.block_ids, "meta": self.metadata, "td": tensors_desc, "nl": len(self.layer_blocks.get("key_cache", [])), }, separators=(",", ":"), ).encode("utf-8") return struct.pack(">I", len(header)) + header def to_bytes(self) -> bytes: """Convert to compact binary format for fast transfer.""" tensors_desc, chunks, _, _ = self._build_tensors_desc(cpu=True) return b"".join([self._build_header_bytes(tensors_desc)] + chunks) def to_gpu_tensor(self) -> torch.Tensor: """Convert to a packed device tensor for raw-data connectors.""" tensors_desc, chunks, data_offset, device = self._build_tensors_desc(cpu=False) if device is None: raise RuntimeError("No device tensors found, use to_bytes() instead") header_prefix = self._build_header_bytes(tensors_desc) output = torch.empty(len(header_prefix) + data_offset, dtype=torch.uint8, device=device) output[: len(header_prefix)].copy_(torch.frombuffer(bytearray(header_prefix), dtype=torch.uint8)) pos = len(header_prefix) for t_flat in chunks: n = t_flat.numel() output[pos : pos + n].copy_(t_flat) pos += n return output @staticmethod def _load_header_from_memoryview(raw_mv: memoryview) -> tuple[dict[str, Any], memoryview]: if len(raw_mv) < 4: raise ValueError("Corrupted KV payload: missing 4-byte header length") header_len = struct.unpack(">I", raw_mv[:4])[0] if header_len > len(raw_mv) - 4: raise ValueError(f"Corrupted KV payload: header_len={header_len} exceeds buffer size={len(raw_mv)}") return json.loads(bytes(raw_mv[4 : 4 + header_len])), raw_mv[4 + header_len :] @staticmethod def _load_header_from_tensor(tensor: torch.Tensor) -> tuple[dict[str, Any], int]: if tensor.dtype != torch.uint8 or tensor.dim() != 1: raise ValueError("Packed device KV payload must be a 1-D uint8 tensor") total_bytes = int(tensor.numel()) if total_bytes < 4: raise ValueError("Corrupted KV payload: missing 4-byte header length") header_len = struct.unpack(">I", tensor[:4].cpu().numpy().tobytes())[0] if header_len > total_bytes - 4: raise ValueError(f"Corrupted KV payload: header_len={header_len} exceeds buffer size={total_bytes}") header_bytes = tensor[4 : 4 + header_len].cpu().numpy().tobytes() return json.loads(header_bytes), 4 + header_len @staticmethod def _validate_tensor_span(name: str, info: dict[str, Any], tensor_data_bytes: int) -> tuple[int, int]: offset = info["o"] nbytes = info["b"] if offset < 0 or nbytes < 0 or offset + nbytes > tensor_data_bytes: raise ValueError( f"Corrupted KV payload tensor span for {name}: " f"offset={offset}, bytes={nbytes}, tensor_data_bytes={tensor_data_bytes}" ) return offset, nbytes @staticmethod def _resolve_torch_dtype(dtype_name: Any) -> torch.dtype: torch_dtype = _SAFE_TORCH_DTYPES.get(str(dtype_name)) if torch_dtype is None: raise ValueError(f"Unsupported dtype in KV payload: {dtype_name}") return torch_dtype @staticmethod def _resolve_layer_idx(info: dict[str, Any], num_layers: int) -> int: layer_idx = info.get("i") if layer_idx is None: name = info.get("n") if isinstance(name, str) and name.startswith("key_cache_"): layer_idx = int(name.removeprefix("key_cache_")) elif isinstance(name, str) and name.startswith("value_cache_"): layer_idx = int(name.removeprefix("value_cache_")) else: raise ValueError(f"Invalid KV tensor name in payload: {name}") if not isinstance(layer_idx, int): raise ValueError(f"Invalid layer index in KV payload: {layer_idx}") if layer_idx < 0 or layer_idx >= num_layers: raise ValueError(f"Invalid layer index in KV payload: {layer_idx} (num_layers={num_layers})") return layer_idx @staticmethod def _populate_caches(header: dict[str, Any], get_tensor: callable) -> dict[str, Any]: """Shared deserialization loop for both CPU and GPU paths.""" num_layers = header["nl"] key_cache: list[torch.Tensor | None] = [None] * num_layers value_cache: list[torch.Tensor | None] = [None] * num_layers for info in header["td"]: if info.get("x"): continue name: str = info["n"] torch_dtype = KVCacheTransferData._resolve_torch_dtype(info["d"]) t = get_tensor(info).view(torch_dtype).reshape(info["s"]) layer_idx = KVCacheTransferData._resolve_layer_idx(info, num_layers) if name.startswith("key_cache_"): key_cache[layer_idx] = t elif name.startswith("value_cache_"): value_cache[layer_idx] = t return { "request_id": header["rid"], "layer_blocks": {"key_cache": key_cache, "value_cache": value_cache}, "block_ids": header["bids"], "metadata": header["meta"], } @staticmethod def from_bytes(raw: "bytes | bytearray | memoryview") -> dict[str, Any]: """Reconstruct KV cache data from the packed bytes format.""" raw_mv = memoryview(raw) if not isinstance(raw, memoryview) else raw header, tensor_data_mv = KVCacheTransferData._load_header_from_memoryview(raw_mv) data_len = len(tensor_data_mv) def _get(info: dict) -> torch.Tensor: offset, nbytes = KVCacheTransferData._validate_tensor_span(info["n"], info, data_len) return torch.frombuffer(tensor_data_mv, dtype=torch.uint8, offset=offset, count=nbytes) return KVCacheTransferData._populate_caches(header, _get) @staticmethod def from_bytes_device(tensor: torch.Tensor) -> dict[str, Any]: """Reconstruct KV cache data from a packed device tensor.""" header, data_start = KVCacheTransferData._load_header_from_tensor(tensor) data_len = int(tensor.numel()) - data_start def _get(info: dict) -> torch.Tensor: offset, nbytes = KVCacheTransferData._validate_tensor_span(info["n"], info, data_len) return tensor[data_start + offset : data_start + offset + nbytes].clone() return KVCacheTransferData._populate_caches(header, _get) @staticmethod def from_bytes_gpu(tensor: torch.Tensor) -> dict[str, Any]: """Compatibility alias for callers using the old GPU-specific name.""" return KVCacheTransferData.from_bytes_device(tensor) class OmniKVTransferManager: """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 """ def __init__(self, config: OmniKVCacheConfig): self.config = config self._connector = None # Pre-calculate send stages (from_stage, to_stage) self.send_stages = ( (str(config.from_stage), str(config.to_stage)) if config.from_stage and config.to_stage else (None, None) ) # Pre-calculate receive stages (from_stage, to_stage) recv_from = config.from_stage if config.engine_input_source: recv_from = config.engine_input_source[0] elif isinstance(config.stage_id, int): recv_from = config.stage_id - 1 self.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) ) local_rank = get_local_tp_rank() if config.from_tp <= 1 and config.to_tp <= 1: detected_tp = get_tp_world_size() from_tp = detected_tp to_tp = detected_tp else: from_tp = config.from_tp to_tp = config.to_tp self._tp_topo = KVTPTopology(source_tp_size=from_tp, target_tp_size=to_tp, local_rank=local_rank) # Injectable hooks (compatible with PR #2677 OmniConnectorModelRunnerMixin). self.kv_send_key_builder: Callable | None = None self.kv_recv_key_builder: Callable | None = None self.kv_payload_merger: Callable | None = None self.kv_payload_slicer: Callable | None = None # Base sender endpoint (rank-0 host/port) stored during # update_sender_info(). Used by the receive path to construct # per-rank metadata for heterogeneous TP without querying a registry. self._sender_base_host: str | None = None self._sender_base_zmq_port: int | None = None if config.need_send_cache and config.connector_config: try: _ = self.connector logger.info("Sender connector eagerly initialized") except Exception as e: logger.warning("Failed to eagerly initialize sender connector: %s", e) # ------------------------------------------------------------------ # # Factory helpers # ------------------------------------------------------------------ # @classmethod def _create(cls, cfg: dict | None) -> "OmniKVTransferManager": """Create manager from raw config dict.""" if not cfg or not isinstance(cfg, dict): return cls(OmniKVCacheConfig()) rank_mapping = cfg.get("rank_mapping", {}) if not isinstance(rank_mapping, dict): rank_mapping = {} return cls( OmniKVCacheConfig( connector_config=cfg.get("connector_config"), from_stage=cfg.get("omni_from_stage"), to_stage=cfg.get("omni_to_stage"), stage_id=cfg.get("stage_id"), engine_input_source=cfg.get("engine_input_source", []), need_recv_cache=cfg.get("need_recv_cache", False), need_send_cache=cfg.get("need_send_cache", False), recv_timeout=cfg.get("recv_timeout", 30.0), from_tp=int(rank_mapping.get("from_tp", 1)), to_tp=int(rank_mapping.get("to_tp", 1)), ) ) @classmethod def from_od_config(cls, config: Any) -> "OmniKVTransferManager": """Create from model or OmniDiffusion config.""" return cls._create(getattr(config, "omni_kv_config", None)) from_model_config = from_od_config @classmethod def from_vllm_config(cls, vllm_config: Any, model_config: Any) -> "OmniKVTransferManager": """Create from vllm config with fallback to kv_transfer_config.""" # Primary: omni_kv_config from model_config omni_kv = getattr(model_config, "omni_kv_config", None) if isinstance(omni_kv, dict): return cls._create(omni_kv) # Fallback: check kv_transfer_config kv_cfg = getattr(vllm_config, "kv_transfer_config", None) if kv_cfg: direct = getattr(kv_cfg, "omni_connector_config", None) if isinstance(direct, dict) and direct: return cls._create({"connector_config": direct}) extra = getattr(kv_cfg, "kv_connector_extra_config", None) if isinstance(extra, dict): omni = extra.get("omni_connector_config") if isinstance(omni, dict) and omni: return cls._create({"connector_config": omni}) return cls(OmniKVCacheConfig()) @property def connector(self): """Lazy initialization of connector.""" # If a previous initialization attempt failed, don't retry on every access. if self._connector is False: return None if self._connector is None: cfg = self.config.connector_config if cfg and (c_type := cfg.get("type")): try: c_extra = {k: v for k, v in cfg.items() if k != "type"} if c_type in TRANSFER_ENGINE_CONNECTOR_NAMES: base_port = expand_env_int(c_extra.get("zmq_port", 50051), "zmq_port") c_extra["from_stage"] = ( str(self.config.from_stage) if self.config.from_stage is not None else "0" ) c_extra["to_stage"] = str(self.config.to_stage) if self.config.to_stage is not None else "1" try: stage_int = int(self.config.from_stage) if self.config.from_stage is not None else 0 except (TypeError, ValueError): stage_int = 0 replica_id = get_omni_replica_id() zmq_port = kv_zmq_port( base_port, stage_int, self._tp_topo.local_rank, replica_id=replica_id, ) if self.config.need_send_cache: c_extra["role"] = "sender" c_extra["zmq_port"] = zmq_port elif self.config.need_recv_cache: c_extra["role"] = "receiver" # Receiver-side sender endpoints are request-scoped. # They are attached by the orchestrator as # kv_sender_info and applied in update_sender_info(). # Do not derive sender_host from this receiver's # local YAML host; on multi-node runs that points at # the wrong process. Explicit sender_* YAML values # are still preserved for standalone connector use. logger.info( "Initializing OmniConnector type=%s role=%s", c_type, c_extra.get("role", "N/A"), ) self._connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=c_type, extra=c_extra)) except Exception: logger.exception("Failed to initialize OmniConnector") self._connector = False return self._connector if self._connector else None get_connector = property(lambda self: self.connector) def _resolve_sender_info( self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None ) -> dict[str, Any] | None: if not sender_info: return None if "host" in sender_info: return sender_info if not isinstance(sender_info, dict): return None preferred_keys: list[str | int] = [] if sender_stage_id is None: recv_from, _ = self.recv_stages sender_stage_id = recv_from if sender_stage_id is not None: preferred_keys.append(sender_stage_id) preferred_keys.append(str(sender_stage_id)) try: preferred_keys.append(int(sender_stage_id)) except (TypeError, ValueError): pass for key in dict.fromkeys(preferred_keys): info = sender_info.get(key) if isinstance(info, dict) and "host" in info: return info candidates = [info for info in sender_info.values() if isinstance(info, dict) and "host" in info] if len(candidates) == 1: return candidates[0] if candidates: logger.warning( "Ambiguous sender_info for sender_stage_id=%s: " "expected caller to resolve a single sender entry, got %s", sender_stage_id, sender_info, ) return None @staticmethod def _clone_received_payload_tensors(data: dict[str, Any]) -> dict[str, Any]: if not isinstance(data, dict) or "layer_blocks" not in data: return data layer_blocks = data["layer_blocks"] for cache_name in ("key_cache", "value_cache"): cache_list = layer_blocks.get(cache_name, []) for idx, tensor in enumerate(cache_list): if isinstance(tensor, torch.Tensor): cache_list[idx] = tensor.clone() return data def _slice_transfer_data_for_target(self, kv_data: KVCacheTransferData, target_rank: int) -> KVCacheTransferData: """Pre-slice sender payload for one target rank when sender TP < receiver TP.""" topo = self._tp_topo ratio = topo.target_tp_size // topo.source_tp_size offset_in_sender = target_rank % ratio metadata = dict(kv_data.metadata) if isinstance(kv_data.metadata, dict) else {} metadata["tp_head_slice"] = { "applied": True, "side": "sender", "target_rank": target_rank, "source_rank": topo.local_rank, "from_tp": topo.source_tp_size, "to_tp": topo.target_tp_size, "offset_in_shard": offset_in_sender, "num_slices": ratio, } return KVCacheTransferData( request_id=kv_data.request_id, layer_blocks=slice_layer_blocks(kv_data.layer_blocks, offset_in_sender, ratio), block_ids=list(kv_data.block_ids), metadata=metadata, ) def _serialize_transfer_payload(self, kv_data: KVCacheTransferData) -> torch.Tensor | bytes | dict[str, Any]: """Serialize KV transfer data using the connector's fastest supported path.""" if getattr(self.connector, "supports_raw_data", False): try: return kv_data.to_gpu_tensor() except Exception: pass try: return kv_data.to_bytes() except Exception: return kv_data.to_dict() @staticmethod def _collect_request_kv_payload(req: Any) -> dict[str, object]: """Collect request-side KV objects for object broadcast.""" kv_payload: dict[str, object] = {} for attr in ("past_key_values", "kv_metadata"): val = getattr(req, attr, None) if val is not None: kv_payload[attr] = val if hasattr(req, "sampling_params") and req.sampling_params is not None: for key in list(vars(req.sampling_params).keys()): if key in ("past_key_values", "kv_metadata") or ( key.startswith("cfg_") and ( key.endswith("_past_key_values") or key.endswith("_kv_metadata") or key in ( "cfg_kv_request_ids", "cfg_active_branch", "cfg_branch_roles", "cfg_branch_past_key_values", "cfg_branch_kv_metadata", ) ) ): val = getattr(req.sampling_params, key, None) if val is not None: kv_payload[f"sp.{key}"] = val return kv_payload @staticmethod def _apply_request_kv_payload( req: Any, kv_payload: dict[str, object], target_device: torch.device | None = None, ) -> None: """Apply a broadcast KV payload back onto a request object.""" for attr in ("past_key_values", "kv_metadata"): val = kv_payload.get(attr) if val is not None: if target_device is not None: val = _move_to_device(val, target_device) setattr(req, attr, val) if hasattr(req, "sampling_params") and req.sampling_params is not None: for key, val in kv_payload.items(): if key.startswith("sp."): if target_device is not None: val = _move_to_device(val, target_device) setattr(req.sampling_params, key[3:], val) @staticmethod def _discover_cfg_branch_roles(req: Any) -> list[str]: """Discover CFG branch roles in a stable order.""" sampling_params = getattr(req, "sampling_params", None) if sampling_params is None: return [] roles: list[str] = [] branch_map = getattr(sampling_params, "cfg_branch_past_key_values", None) or {} for preferred_role in ("cfg_text", "cfg_img"): if ( preferred_role in branch_map or getattr(sampling_params, f"{preferred_role}_past_key_values", None) is not None ): roles.append(preferred_role) for role in branch_map.keys(): if role not in roles and branch_map.get(role) is not None: roles.append(role) for key in vars(sampling_params).keys(): if not (key.startswith("cfg_") and key.endswith("_past_key_values")): continue role = key.removesuffix("_past_key_values") if role in ("cfg_branch",) or role in roles: continue if getattr(sampling_params, key, None) is not None: roles.append(role) return roles @classmethod def _build_cfg_rank_local_payloads(cls, req: Any, cfg_size: int) -> list[dict[str, object] | None]: """Build per-cfg-rank payloads so each rank receives only its branch KV.""" full_payload = cls._collect_request_kv_payload(req) payloads: list[dict[str, object] | None] = [] main_payload = { key: value for key, value in full_payload.items() if key in ("past_key_values", "kv_metadata", "sp.past_key_values", "sp.kv_metadata") } branch_roles = cls._discover_cfg_branch_roles(req) if branch_roles: main_payload["sp.cfg_branch_roles"] = list(branch_roles) main_payload["sp.cfg_active_branch"] = None payloads.append(main_payload or None) sampling_params = getattr(req, "sampling_params", None) branch_map = getattr(sampling_params, "cfg_branch_past_key_values", None) or {} branch_metadata_map = getattr(sampling_params, "cfg_branch_kv_metadata", None) or {} for role in branch_roles: if sampling_params is None: payloads.append(None) continue branch_kv = branch_map.get(role) if branch_kv is None: branch_kv = getattr(sampling_params, f"{role}_past_key_values", None) branch_metadata = branch_metadata_map.get(role) if branch_metadata is None: branch_metadata = getattr(sampling_params, f"{role}_kv_metadata", None) if branch_kv is None: payloads.append(None) continue local_payload = dict(main_payload) local_payload["sp.cfg_active_branch"] = role local_payload["sp.cfg_branch_roles"] = list(branch_roles) local_payload["sp.cfg_branch_past_key_values"] = {role: branch_kv} local_payload[f"sp.{role}_past_key_values"] = branch_kv if branch_metadata is not None: local_payload["sp.cfg_branch_kv_metadata"] = {role: branch_metadata} local_payload[f"sp.{role}_kv_metadata"] = branch_metadata payloads.append(local_payload) while len(payloads) < cfg_size: payloads.append(None) return payloads[:cfg_size] def update_sender_info(self, 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. """ if not self.config.need_recv_cache: return actual_info = self._resolve_sender_info(sender_info, sender_stage_id=sender_stage_id) if not actual_info or "host" not in actual_info: logger.warning("Invalid sender_info format: %s", sender_info) return sender_host = actual_info.get("host") base_zmq_port = actual_info.get("zmq_port") # Store base sender info for per-rank metadata construction. self._sender_base_host = sender_host if base_zmq_port is not None: self._sender_base_zmq_port = int(base_zmq_port) # --- Default sender: offset to match this receiver's corresponding sender rank --- zmq_port = base_zmq_port if zmq_port is not None and self._tp_topo.local_rank > 0: zmq_port = int(zmq_port) + self._tp_topo.local_rank * KV_RANK_PORT_STRIDE if self.config.connector_config: self.config.connector_config["sender_host"] = sender_host self.config.connector_config["sender_zmq_port"] = zmq_port if self._connector and hasattr(self._connector, "update_sender_info"): try: self._connector.update_sender_info(sender_host, zmq_port) except Exception: if hasattr(self._connector, "sender_host"): self._connector.sender_host = sender_host if hasattr(self._connector, "sender_zmq_port"): self._connector.sender_zmq_port = zmq_port logger.info( "Sender info updated: host=%s, base_port=%s, adjusted_port=%s (local_rank=%s)", sender_host, base_zmq_port, zmq_port, self._tp_topo.local_rank, ) def handle_finished_requests_kv_transfer( self, 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. Args: finished_reqs: Dict mapping request_id to {block_ids, seq_len} kv_caches: List of KV cache (tensor or tuple) per layer block_size: Size of each cache block cache_dtype: Data type of the cache request_id_resolver: Optional function to resolve global request ID Returns: List of request IDs that were processed """ if not finished_reqs: return [] if not self.config.need_send_cache: return list(finished_reqs.keys()) if not self.connector: logger.warning("No connector available, skipping KV transfer but freeing resources") return list(finished_reqs.keys()) logger.debug(f"Processing KV transfer for {len(finished_reqs)} requests") extracted_ids = [] for req_id, data in finished_reqs.items(): try: seq_len = data.get("seq_len", 0) block_ids = data.get("block_ids", []) if not block_ids: logger.warning(f"Request {req_id} has no block IDs, skipping") continue custom_metadata = data.get("custom_metadata") # Extract KV cache from GPU blocks and keep it on-device when # possible so raw-data connectors can use the fast path. kv_data = self._extract_kv_cache( req_id, block_ids, seq_len, kv_caches, block_size, cache_dtype, custom_metadata ) if kv_data: # Resolve global request ID if available transfer_req_id = request_id_resolver(req_id) if request_id_resolver else req_id # Transfer to downstream stage via connector self._transfer_kv_cache(kv_data, transfer_req_id) except Exception as e: logger.error(f"Failed KV transfer for {req_id}: {e}") finally: extracted_ids.append(req_id) return extracted_ids def _extract_kv_cache( self, req_id: str, block_ids: list[int], seq_len: int, kv_caches: list[LayerKV], block_size: int, cache_dtype: str, custom_metadata: dict[str, Any] | None = None, ) -> KVCacheTransferData | None: """Extract KV cache from GPU blocks for a single request. Args: req_id: Request identifier block_ids: List of block IDs to extract seq_len: Sequence length kv_caches: List of KV cache (tensor or tuple) per layer block_size: Size of each cache block cache_dtype: Data type of the cache custom_metadata: Optional custom metadata to include Note: If key/value block counts differ, extraction uses only the overlapping block range. Extra key/value blocks are ignored, so returned KV may be partial. Returns: KVCacheTransferData if extraction successful, None otherwise """ num_layers = len(kv_caches) key_cache: list[torch.Tensor | None] = [None] * num_layers value_cache: list[torch.Tensor | None] = [None] * num_layers for layer_idx, layer_kv in enumerate(kv_caches): kv_pair = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=layer_idx) if kv_pair is None: continue key_blocks, value_blocks = kv_pair if key_blocks.shape[0] != value_blocks.shape[0]: logger.warning( f"Layer {layer_idx} for request {req_id} has mismatched KV block counts: " f"key={key_blocks.shape[0]}, value={value_blocks.shape[0]}; using shared range" ) # Validate block IDs - shape: [num_blocks, block_size, n_heads, head_dim] max_block = min(key_blocks.shape[0], value_blocks.shape[0]) - 1 valid_ids = [bid for bid in block_ids if 0 <= bid <= max_block] if not valid_ids: continue # Extract and reshape: [n_blocks, block_size, n_heads, head_dim] # -> [seq_len, n_heads, head_dim] selected_k = key_blocks[valid_ids] selected_v = value_blocks[valid_ids] flat_k = selected_k.flatten(0, 1) flat_v = selected_v.flatten(0, 1) if seq_len < flat_k.shape[0]: flat_k = flat_k[:seq_len] flat_v = flat_v[:seq_len] key_cache[layer_idx] = flat_k.detach().contiguous() value_cache[layer_idx] = flat_v.detach().contiguous() if not any(k is not None for k in key_cache): return None return KVCacheTransferData( request_id=req_id, layer_blocks={"key_cache": key_cache, "value_cache": value_cache}, block_ids=block_ids, metadata={ "block_size": block_size, "num_layers": num_layers, "dtype": str(cache_dtype), "seq_len": seq_len, **(custom_metadata or {}), }, ) def _transfer_kv_cache(self, kv_data: KVCacheTransferData, transfer_req_id: str) -> None: """Transfer KV cache data to downstream stage via OmniConnector. Args: kv_data: The extracted KV cache data transfer_req_id: The request ID to use for transfer """ from_stage, to_stage = self.send_stages if not from_stage or not to_stage: raise ValueError("Transfer stages (omni_from_stage, omni_to_stage) not configured") kv_data.request_id = transfer_req_id serialization_start = time.perf_counter() topo = self._tp_topo send_keys = build_rank_aware_send_keys( transfer_req_id, from_stage, to_stage, topo, hook=self.kv_send_key_builder ) sender_slice_active = ( topo.source_tp_size < topo.target_tp_size and len(send_keys) > 1 and not callable(self.kv_send_key_builder) ) per_key_payloads: list[tuple[str, torch.Tensor | bytes | dict[str, Any]]] = [] if sender_slice_active: target_ranks = get_kv_target_ranks(topo) if len(target_ranks) != len(send_keys): logger.warning( "Skip sender-side KV slicing because target rank count does not match send key count: " "target_ranks=%s send_keys=%s", len(target_ranks), len(send_keys), ) sender_slice_active = False else: for put_key, target_rank in zip(send_keys, target_ranks, strict=False): sliced_kv_data = self._slice_transfer_data_for_target(kv_data, target_rank) per_key_payloads.append((put_key, self._serialize_transfer_payload(sliced_kv_data))) if not per_key_payloads: transfer_data = self._serialize_transfer_payload(kv_data) per_key_payloads = [(put_key, transfer_data) for put_key in send_keys] serialization_ms = (time.perf_counter() - serialization_start) * 1000 logger.info("KV cache serialized for %s in %.1f ms", transfer_req_id, serialization_ms) transfer_start = time.perf_counter() total_size = 0 all_succeeded = True for put_key, transfer_data in per_key_payloads: success, size, _ = self._transfer_with_retry(from_stage, to_stage, put_key, transfer_data) total_size += size all_succeeded = all_succeeded and success elapsed = time.perf_counter() - transfer_start if all_succeeded: mbps = (total_size / 1024 / 1024) / elapsed if elapsed > 0 else 0 logger.info( "KV transfer OK: %s, %s bytes across %s key(s), %.3fs, %.1f MB/s", transfer_req_id, total_size, len(send_keys), elapsed, mbps, ) else: logger.error(f"KV transfer FAILED: {transfer_req_id}") def _transfer_with_retry( self, from_stage: str, to_stage: str, put_key: str, data: "dict[str, Any] | bytes | torch.Tensor", max_retries: int = 3, ) -> tuple[bool, int, dict[str, Any] | None]: """Transfer data with retry and exponential backoff. Args: from_stage: Source stage identifier to_stage: Target stage identifier put_key: Pre-built connector key (rank-aware when TP > 1) data: Data to transfer max_retries: Maximum number of retry attempts Returns: Tuple of (success, size, metadata) """ for attempt in range(max_retries): try: success, size, metadata = self.connector.put( from_stage=from_stage, to_stage=to_stage, put_key=put_key, data=data ) if success: return success, size, metadata logger.warning(f"Transfer attempt {attempt + 1} failed for {put_key}") except Exception as e: logger.warning(f"Transfer attempt {attempt + 1} exception: {e}") if attempt < max_retries - 1: time.sleep(0.1 * (2**attempt)) return False, 0, None @torch.inference_mode() def receive_kv_cache_for_request( self, request_id: str, target_device: torch.device | None = None, ) -> tuple[dict[str, Any] | None, int]: """Receive KV cache for a specific request. This implements the receiving logic from gpu_diffusion_model_runner.py. Args: request_id: The request ID to receive KV cache for target_device: Optional device to move tensors to Returns: Tuple of (data dict, size) if successful, (None, 0) otherwise """ # Check if we should receive KV cache based on config if not self.config.need_recv_cache: logger.debug("Skip receiving KV cache for %s (need_recv_cache=False)", request_id) return None, 0 if not self.connector: logger.warning("No connector available for receiving KV cache") return None, 0 from_stage, to_stage = self.recv_stages if not from_stage or not to_stage: logger.warning("Receive stages not configured") return None, 0 # Skip during warmup dummy run — no sender is available. if OmniDiffusionRequest.is_dummy_run_request_id(request_id): logger.info("Skip receiving KV cache for dummy warmup request") return None, 0 timeout = self.config.recv_timeout start_time = time.time() poll_interval = 0.01 max_poll_interval = 0.5 topo = self._tp_topo recv_key_pairs = build_rank_aware_recv_keys( request_id, from_stage, to_stage, topo, hook=self.kv_recv_key_builder ) pending_pairs = list(recv_key_pairs) received_payloads: dict[str, tuple[dict[str, Any], int]] = {} logger.info( "Wait for KV cache for request %s from stage %s to %s via %s key(s)...", request_id, from_stage, to_stage, len(recv_key_pairs), ) try: while True: link_start = time.perf_counter() for get_key, from_rank in list(pending_pairs): # Construct per-rank metadata so the connector queries # the correct sender endpoint (heterogeneous TP path). # When from_rank is None (TP<=1), metadata stays None # and the connector falls back to its default sender. rank_metadata: dict[str, Any] | None = None if from_rank is not None and self._sender_base_host and self._sender_base_zmq_port is not None: rank_metadata = { "source_host": self._sender_base_host, "source_port": self._sender_base_zmq_port + from_rank * KV_RANK_PORT_STRIDE, } result = self.connector.get( from_stage=from_stage, to_stage=to_stage, get_key=get_key, metadata=rank_metadata, ) if not result: continue raw_data, size = result managed_buffer = None if hasattr(raw_data, "tensor") and hasattr(raw_data, "release"): managed_buffer = raw_data try: buf_tensor = raw_data.tensor if getattr(buf_tensor.device, "type", "cpu") != "cpu": data = KVCacheTransferData.from_bytes_device(buf_tensor) raw_data.release() managed_buffer = None else: data = KVCacheTransferData.from_bytes(memoryview(buf_tensor.numpy())) data = self._clone_received_payload_tensors(data) raw_data.release() managed_buffer = None except Exception as e: logger.error("Failed to deserialize KV cache from ManagedBuffer: %s", e) if managed_buffer is not None: raw_data.release() return None, 0 elif isinstance(raw_data, (bytes, bytearray)): data = KVCacheTransferData.from_bytes(raw_data) elif isinstance(raw_data, torch.Tensor) and raw_data.dtype == torch.uint8 and raw_data.dim() == 1: if getattr(raw_data.device, "type", "cpu") != "cpu": data = KVCacheTransferData.from_bytes_device(raw_data) else: data = KVCacheTransferData.from_bytes(raw_data.numpy().tobytes()) else: data = raw_data received_payloads[get_key] = (data, size) pending_pairs.remove((get_key, from_rank)) if not pending_pairs and received_payloads: elapsed = time.time() - start_time link_ms = (time.perf_counter() - link_start) * 1000 ordered_payloads = [received_payloads[key][0] for key, _ in recv_key_pairs] total_size = sum(received_payloads[key][1] for key, _ in recv_key_pairs) if len(ordered_payloads) == 1: data = ordered_payloads[0] else: data = merge_received_rank_shards(ordered_payloads, merger=self.kv_payload_merger) data = slice_received_rank_shard(data, topo, slicer=self.kv_payload_slicer) try: if isinstance(data, dict) and "layer_blocks" in data: layer_blocks = data["layer_blocks"] for cache_list in [ layer_blocks.get("key_cache", []), layer_blocks.get("value_cache", []), ]: for i, tensor in enumerate(cache_list): if not isinstance(tensor, torch.Tensor): continue if target_device is not None and tensor.device != target_device: cache_list[i] = tensor.to(target_device).contiguous() except Exception: logger.exception("Failed to move KV cache tensors to target device") logger.info( "Successfully received KV cache for %s, %s bytes across %s key(s), wait=%.3fs, link=%.1fms", request_id, total_size, len(recv_key_pairs), elapsed, link_ms, ) return data, total_size if time.time() - start_time > timeout: logger.error(f"Timeout waiting for KV cache for request {request_id} after {timeout}s") return None, 0 time.sleep(poll_interval) poll_interval = min(poll_interval * 2, max_poll_interval) except Exception: logger.exception("Error receiving KV cache for %s", request_id) return None, 0 def apply_kv_cache_to_request(self, req: Any, data: dict[str, Any]) -> None: """Apply received KV cache data to a request object. Args: req: The request object to apply KV cache to data: The received KV cache data dictionary """ if isinstance(data, dict) and "layer_blocks" in data: layer_blocks = data["layer_blocks"] from types import SimpleNamespace kv_obj = SimpleNamespace(**layer_blocks) req.past_key_values = kv_obj # [Omni] Also attach to sampling_params for BagelPipeline compatibility # BagelPipeline checks req.sampling_params.past_key_values if hasattr(req, "sampling_params") and req.sampling_params is not None: req.sampling_params.past_key_values = kv_obj if "metadata" in data: req.kv_metadata = data["metadata"] if hasattr(req, "sampling_params") and req.sampling_params is not None: req.sampling_params.kv_metadata = data["metadata"] @staticmethod def _resolve_request_id(req: Any) -> str | None: """Resolve the logical request ID used for KV transfer lookups.""" return getattr(req, "request_id", None) # Legacy compatibility method def receive_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool: """Receive KV cache and populate request object (legacy interface). Args: req: Request object with request_id attribute target_device: Optional device to move tensors to Returns: True if successful, False otherwise """ kv_sender_info = getattr(req, "kv_sender_info", None) if kv_sender_info: self.update_sender_info(kv_sender_info, sender_stage_id=self.recv_stages[0]) request_id = self._resolve_request_id(req) if not request_id: logger.warning("Request has no ID, cannot receive KV cache") return False data, size = self.receive_kv_cache_for_request(request_id, target_device) if data: self.apply_kv_cache_to_request(req, data) return True return False def receive_multi_kv_cache( self, req: Any, cfg_kv_collect_func: Callable | None = None, target_device: torch.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. Args: req: Request object with request_id and sampling_params. cfg_kv_collect_func: Model-specific function for collecting CFG KV caches. Signature: (request_id, cfg_request_ids, kv_transfer_manager, target_device) -> dict[str, Any] target_device: Device to move tensors to. Returns: True if primary KV cache was received successfully. """ primary_ok = self.receive_kv_cache(req, target_device) cfg_ids = getattr(getattr(req, "sampling_params", None), "cfg_kv_request_ids", None) if cfg_ids and cfg_kv_collect_func: request_id = self._resolve_request_id(req) try: cfg_kvs = cfg_kv_collect_func( request_id, cfg_ids, self, target_device, ) if cfg_kvs and hasattr(req, "sampling_params") and req.sampling_params is not None: for key, value in cfg_kvs.items(): setattr(req.sampling_params, key, value) logger.info("Applied CFG KV caches: %s", list(cfg_kvs.keys())) except Exception: logger.exception("Failed to collect CFG KV caches for %s", request_id) return primary_ok def receive_multi_kv_cache_distributed( self, req: Any, cfg_kv_collect_func: Callable | None = None, target_device: torch.device | None = None, ) -> bool: """Distributed wrapper around :meth:`receive_multi_kv_cache`. TP-aware path selection: - world size 1: direct receive - TP active, cfg size 1: each rank independently receives - TP active, cfg size > 1: cfg-rank 0 receives, then broadcasts to peers that share the same TP rank - TP inactive: legacy rank-0 receive then world broadcast """ from vllm_omni.diffusion.distributed.parallel_state import ( get_cfg_group, get_classifier_free_guidance_rank, get_classifier_free_guidance_world_size, get_sequence_parallel_rank, get_sequence_parallel_world_size, get_sp_group, get_world_group, ) world = get_world_group() if world.world_size <= 1: return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device) topo = self._tp_topo tp_active = topo.source_tp_size > 1 or topo.target_tp_size > 1 cfg_size = 1 cfg_rank = 0 cfg_group = None sp_size = 1 sp_rank = 0 sp_group = None try: cfg_size = get_classifier_free_guidance_world_size() cfg_rank = get_classifier_free_guidance_rank() cfg_group = get_cfg_group() except Exception: cfg_size = 1 cfg_rank = 0 cfg_group = None try: sp_size = get_sequence_parallel_world_size() sp_rank = get_sequence_parallel_rank() sp_group = get_sp_group() except Exception: sp_size = 1 sp_rank = 0 sp_group = None if tp_active and (cfg_size <= 1 and sp_size <= 1): logger.info( "Rank-aware KV receive: rank %s independently receiving (from_tp=%s, to_tp=%s)", topo.local_rank, topo.source_tp_size, topo.target_tp_size, ) return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device) if tp_active and (cfg_size > 1 or sp_size > 1): kv_payload: dict[str, object] | None = None is_owner = cfg_rank == 0 and sp_rank == 0 # step1: only owner need to receive if is_owner: received = self.receive_multi_kv_cache( req, cfg_kv_collect_func, torch.device("cpu"), ) if received: if cfg_size > 1: # build cfg-local payloads cfg_rank_payloads = self._build_cfg_rank_local_payloads( req, cfg_size, ) kv_payload = cfg_rank_payloads[0] # send to other cfg for dst_cfg_rank in range(1, cfg_size): cfg_group.send_object( cfg_rank_payloads[dst_cfg_rank], dst_cfg_rank, ) elif sp_size > 1: kv_payload = self._collect_request_kv_payload(req) else: # Owner receive failed: send None sentinel to CFG followers to avoid deadlock if cfg_size > 1: for dst_cfg_rank in range(1, cfg_size): cfg_group.send_object(None, dst_cfg_rank) elif sp_rank == 0 and cfg_size > 1: kv_payload = cfg_group.recv_object(0) # sp broadcast if sp_size > 1 and sp_group is not None: kv_payload = sp_group.broadcast_object( kv_payload, src=0, ) if not kv_payload: return False self._apply_request_kv_payload(req, kv_payload, target_device) return True kv_payload: dict[str, object] | None = None if world.rank_in_group == 0: received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu")) if received: kv_payload = self._collect_request_kv_payload(req) kv_payload = world.broadcast_object(kv_payload, src=0) if not kv_payload: return False self._apply_request_kv_payload(req, kv_payload, target_device) return True def _move_to_device(obj: object, device: torch.device) -> object: """Recursively move tensors inside a KV cache object to *device*.""" if isinstance(obj, torch.Tensor): return obj.to(device).contiguous() if obj.device != device else obj if hasattr(obj, "__dict__"): for k, v in vars(obj).items(): setattr(obj, k, _move_to_device(v, device)) return obj if isinstance(obj, dict): return {k: _move_to_device(v, device) for k, v in obj.items()} if isinstance(obj, list): return [_move_to_device(v, device) for v in obj] return obj