Zhou TaichangandClaude Fable 5
2c4e9a945b
Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1971 lines
82 KiB
Python
1971 lines
82 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Unified OmniConnector and KV cache transfer management."""
|
|
|
|
import enum
|
|
import json
|
|
import struct
|
|
import time
|
|
from collections.abc import Callable
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
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 vllm_omni.diffusion.sched.interface import KVPrefetchJob
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
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,
|
|
)
|
|
from .utils.serialization import OmniSerializer
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
|
|
|
|
|
|
class KVPrefetchConsumeError(RuntimeError):
|
|
"""Payload consumed from connector but post-get processing failed — sync retry impossible."""
|
|
|
|
|
|
class ReceiveRole(enum.Enum):
|
|
"""How a rank obtains KV: LOCAL pulls its own shard, LEADER pulls then
|
|
distributes, FOLLOWER receives via collective."""
|
|
|
|
LOCAL = "local"
|
|
LEADER = "leader"
|
|
FOLLOWER = "follower"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _TransferTopoConfig:
|
|
"""Snapshot of the current rank's parallel topology for KV receive routing."""
|
|
|
|
role: ReceiveRole
|
|
tp_active: bool
|
|
cfg_size: int
|
|
cfg_rank: int
|
|
cfg_group: Any
|
|
sp_size: int
|
|
sp_rank: int
|
|
sp_group: Any
|
|
world: Any
|
|
# The stage's tensor-model-parallel group (vLLM parallel state). Only the
|
|
# ranks in this group execute the same request in lockstep — DP replicas
|
|
# and PP stages in ``world`` do not — so per-request consensus exchanges
|
|
# must use this group, never ``world``.
|
|
tp_group: Any = None
|
|
tp_size: int = 1
|
|
|
|
@property
|
|
def is_leader(self) -> bool:
|
|
return self.role == ReceiveRole.LEADER
|
|
|
|
@property
|
|
def is_follower(self) -> bool:
|
|
return self.role == ReceiveRole.FOLLOWER
|
|
|
|
@property
|
|
def is_local(self) -> bool:
|
|
return self.role == ReceiveRole.LOCAL
|
|
|
|
@property
|
|
def cfg_active(self) -> bool:
|
|
return self.cfg_size > 1
|
|
|
|
@property
|
|
def sp_active(self) -> bool:
|
|
return self.sp_size > 1
|
|
|
|
|
|
# Placeholder for the heavy primary KV in the side-payload dict; receiver swaps in the rebuilt object from the blob.
|
|
_KV_PLACEHOLDER = "__kv_placeholder__"
|
|
|
|
_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
|
|
enable_kv_async_prefetch: bool = False
|
|
kv_prefetch_min_free_mem_ratio: float = 0.0
|
|
|
|
|
|
@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, *, async_prefetch: bool = False):
|
|
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
|
|
|
|
# Prefetch
|
|
self._async_prefetch = async_prefetch
|
|
self._prefetch_min_free_mem_ratio: float = max(0.0, config.kv_prefetch_min_free_mem_ratio)
|
|
# Single-worker: serial mode ensures at most one outstanding prefetch
|
|
# and avoids stream-creation races on _bg_copy_stream.
|
|
self._prefetch_executor: ThreadPoolExecutor | None = (
|
|
ThreadPoolExecutor(max_workers=1, thread_name_prefix="kv-prefetch") if async_prefetch else None
|
|
)
|
|
|
|
self._prefetch_futures: dict[str, Any] = {}
|
|
self._bg_copy_stream: current_omni_platform.Stream | None = None
|
|
|
|
self._topo_config: _TransferTopoConfig | 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)
|
|
|
|
@property
|
|
def tp_topology(self) -> KVTPTopology:
|
|
"""Return the effective topology after config and runtime detection."""
|
|
return self._tp_topo
|
|
|
|
@property
|
|
def topo_config(self) -> _TransferTopoConfig:
|
|
"""Lazily build the parallel topology snapshot on first access."""
|
|
if self._topo_config is None:
|
|
self._topo_config = self._build_topo_config()
|
|
return self._topo_config
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Factory helpers
|
|
# ------------------------------------------------------------------ #
|
|
|
|
@classmethod
|
|
def _create(cls, cfg: dict | None, *, async_prefetch: bool = False) -> "OmniKVTransferManager":
|
|
"""Create manager from raw config dict."""
|
|
if not cfg or not isinstance(cfg, dict):
|
|
return cls(OmniKVCacheConfig(), async_prefetch=async_prefetch)
|
|
|
|
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)),
|
|
enable_kv_async_prefetch=async_prefetch,
|
|
kv_prefetch_min_free_mem_ratio=cfg.get("kv_prefetch_min_free_mem_ratio", 0.0),
|
|
),
|
|
async_prefetch=async_prefetch,
|
|
)
|
|
|
|
@staticmethod
|
|
def _receiver_pool_on_device(omni_kv: Any) -> bool:
|
|
"""True when the receiver connector pool is on a device (non-CPU)."""
|
|
if not isinstance(omni_kv, dict):
|
|
return False
|
|
cc = omni_kv.get("connector_config")
|
|
if not isinstance(cc, dict):
|
|
return False
|
|
return str(cc.get("memory_pool_device", "cpu")).lower() != "cpu"
|
|
|
|
@classmethod
|
|
def _resolve_async_prefetch(cls, omni_kv: Any, *, has_companion: bool) -> bool:
|
|
"""Effective prefetch flag: requested by config and not force-disabled."""
|
|
if not isinstance(omni_kv, dict) or not bool(omni_kv.get("enable_kv_async_prefetch", False)):
|
|
return False
|
|
return not (has_companion or cls._receiver_pool_on_device(omni_kv))
|
|
|
|
@classmethod
|
|
def from_od_config(cls, config: Any) -> "OmniKVTransferManager":
|
|
"""Create from model or OmniDiffusion config."""
|
|
omni_kv = getattr(config, "omni_kv_config", None)
|
|
has_companion = getattr(config, "cfg_kv_collect_func", None) is not None
|
|
async_prefetch = cls._resolve_async_prefetch(omni_kv, has_companion=has_companion)
|
|
return cls._create(omni_kv, async_prefetch=async_prefetch)
|
|
|
|
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 on model_config — same shape as from_od_config.
|
|
if isinstance(getattr(model_config, "omni_kv_config", None), dict):
|
|
return cls.from_od_config(model_config)
|
|
|
|
connector_cfg = cls._connector_cfg_from_kv_transfer(getattr(vllm_config, "kv_transfer_config", None))
|
|
if connector_cfg is not None:
|
|
return cls._create({"connector_config": connector_cfg})
|
|
return cls(OmniKVCacheConfig())
|
|
|
|
@staticmethod
|
|
def _connector_cfg_from_kv_transfer(kv_cfg: Any) -> dict | None:
|
|
"""Extract the omni connector config dict nested in a kv_transfer_config."""
|
|
if not kv_cfg:
|
|
return None
|
|
direct = getattr(kv_cfg, "omni_connector_config", None)
|
|
if isinstance(direct, dict) and direct:
|
|
return 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 omni
|
|
return None
|
|
|
|
@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 _build_topo_config(self) -> _TransferTopoConfig:
|
|
"""Build the parallel topology snapshot (called once at init)."""
|
|
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,
|
|
)
|
|
|
|
try:
|
|
world = get_world_group()
|
|
world_size = world.world_size
|
|
world_rank = world.rank_in_group
|
|
except Exception:
|
|
logger.exception("World group unavailable; defaulting to LOCAL")
|
|
return _TransferTopoConfig(
|
|
role=ReceiveRole.LOCAL,
|
|
tp_active=False,
|
|
cfg_size=1,
|
|
cfg_rank=0,
|
|
cfg_group=None,
|
|
sp_size=1,
|
|
sp_rank=0,
|
|
sp_group=None,
|
|
world=None,
|
|
)
|
|
|
|
if world_size <= 1:
|
|
return _TransferTopoConfig(
|
|
role=ReceiveRole.LOCAL,
|
|
tp_active=False,
|
|
cfg_size=1,
|
|
cfg_rank=0,
|
|
cfg_group=None,
|
|
sp_size=1,
|
|
sp_rank=0,
|
|
sp_group=None,
|
|
world=world,
|
|
)
|
|
|
|
tp_active = self._tp_topo.source_tp_size > 1 or self._tp_topo.target_tp_size > 1
|
|
|
|
# CFG/SP groups may legitimately be uninitialized on pure-TP or
|
|
# world-only distributed paths; treat that as "not active" (size 1)
|
|
# rather than raising, so those paths can still take the local/world
|
|
# receive route.
|
|
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, cfg_rank, cfg_group = 1, 0, None # CFG-parallel not enabled
|
|
|
|
try:
|
|
sp_size = get_sequence_parallel_world_size()
|
|
sp_rank = get_sequence_parallel_rank()
|
|
sp_group = get_sp_group()
|
|
except Exception:
|
|
sp_size, sp_rank, sp_group = 1, 0, None # SP not enabled
|
|
|
|
if tp_active and cfg_size <= 1 and sp_size <= 1:
|
|
role = ReceiveRole.LOCAL
|
|
elif tp_active and (cfg_size > 1 or sp_size > 1):
|
|
role = ReceiveRole.LEADER if (cfg_rank == 0 and sp_rank == 0) else ReceiveRole.FOLLOWER
|
|
else:
|
|
role = ReceiveRole.LEADER if world_rank == 0 else ReceiveRole.FOLLOWER
|
|
|
|
# The stage's TP group comes from vLLM parallel state (the TP-sharded
|
|
# LM path); it may legitimately be uninitialized on non-TP stages.
|
|
try:
|
|
from vllm.distributed.parallel_state import get_tp_group
|
|
|
|
tp_group = get_tp_group()
|
|
tp_size = tp_group.world_size
|
|
except Exception:
|
|
tp_group, tp_size = None, 1 # vLLM TP not initialized
|
|
|
|
return _TransferTopoConfig(
|
|
role=role,
|
|
tp_active=tp_active,
|
|
cfg_size=cfg_size,
|
|
cfg_rank=cfg_rank,
|
|
cfg_group=cfg_group,
|
|
sp_size=sp_size,
|
|
sp_rank=sp_rank,
|
|
sp_group=sp_group,
|
|
world=world,
|
|
tp_group=tp_group,
|
|
tp_size=tp_size,
|
|
)
|
|
|
|
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
|
|
|
|
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 _clear_request_kv_payload(req: Any) -> None:
|
|
"""Drop every request-side KV field a receive may have applied.
|
|
|
|
Enumerates fields via ``_collect_request_kv_payload`` so clearing stays
|
|
symmetric with what apply/broadcast set — the pipeline then takes the
|
|
same no-KV path as a rank that never received.
|
|
"""
|
|
for key in OmniKVTransferManager._collect_request_kv_payload(req):
|
|
if key.startswith("sp."):
|
|
setattr(req.sampling_params, key[3:], None)
|
|
else:
|
|
setattr(req, key, None)
|
|
|
|
@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)
|
|
|
|
padding = None
|
|
if not branch_roles and cfg_size == 2 and main_payload:
|
|
# AR didn't split KV by branch; cfg follower needs the same
|
|
# positive KV for shared-prefix reuse. Only safe for cfg_size == 2:
|
|
# for cfg_size > 2 every rank must get its own KV, which we can't
|
|
# assign here, so padding stays None and the runner raises.
|
|
# Shallow dict copy: shares PKV tensors with main_payload. Safe
|
|
# only because the send path (cfg_scatter -> pack -> to_gpu_tensor)
|
|
# reads tensors to serialize, never mutates in place. Any async/
|
|
# KV-pool path that mutates req.past_key_values mid-scatter must
|
|
# clone tensors here instead.
|
|
padding = dict(main_payload)
|
|
while len(payloads) < cfg_size:
|
|
payloads.append(padding)
|
|
|
|
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, block_size=block_size)
|
|
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
|
|
|
|
def _resolve_sender_base(self, sender_info: dict[str, Any] | None) -> tuple[str | None, int | None]:
|
|
"""Resolve (host, base_zmq_port) from sender_info."""
|
|
actual = self._resolve_sender_info(sender_info, sender_stage_id=self.recv_stages[0])
|
|
if not actual or "host" not in actual:
|
|
return None, None
|
|
port = actual.get("zmq_port")
|
|
return actual.get("host"), (int(port) if port is not None else None)
|
|
|
|
def _has_enough_prefetch_device_mem(self, device: torch.device | None) -> bool:
|
|
"""False when free-device-memory fraction is below the prefetch threshold.
|
|
|
|
CUDA/NPU only; other platforms report unconstrained (throttle skipped).
|
|
"""
|
|
ratio = self._prefetch_min_free_mem_ratio
|
|
if ratio <= 0.0 or device is None:
|
|
return True
|
|
try:
|
|
free, total = current_omni_platform.get_device_memory(device)
|
|
except Exception:
|
|
return True
|
|
return total <= 0 or (free / total) >= ratio
|
|
|
|
@staticmethod
|
|
def _release_pool_buffers(buffers: list[Any]) -> None:
|
|
"""Release all pool buffers in *buffers* and clear the list."""
|
|
if not buffers:
|
|
return
|
|
for buf in buffers:
|
|
try:
|
|
buf.release()
|
|
except Exception:
|
|
logger.exception("Failed to release KV pool buffer")
|
|
buffers.clear()
|
|
|
|
def start_prefetch(self, kv_prefetch_job: KVPrefetchJob | None, target_device: torch.device | None = None) -> None:
|
|
"""Kick off a background KV load (non-blocking). No-op unless prefetch enabled."""
|
|
if not (self._async_prefetch and self.config.need_recv_cache) or not kv_prefetch_job:
|
|
return
|
|
# Followers receive via collective distribute; bg pull would consume the owner's payload.
|
|
if self.topo_config.is_follower:
|
|
return
|
|
rid = kv_prefetch_job.get("request_id")
|
|
if not rid:
|
|
return
|
|
# Memory pressure → skip prefetch; sync receive handles it.
|
|
if not self._has_enough_prefetch_device_mem(target_device):
|
|
logger.warning(
|
|
"Skip KV prefetch for %s: device free mem below %.2f", rid, self._prefetch_min_free_mem_ratio
|
|
)
|
|
return
|
|
# Serial mode: at most one outstanding prefetch; drop any leftover request's future.
|
|
for stale_rid in [k for k in self._prefetch_futures if k != rid]:
|
|
self._discard_future(stale_rid)
|
|
if rid in self._prefetch_futures:
|
|
return
|
|
sender_info = kv_prefetch_job.get("kv_sender_info")
|
|
if not sender_info:
|
|
# No explicit endpoint → bg receive would target wrong sender under multi-replica.
|
|
logger.debug("Skip KV prefetch for %s: stub has no kv_sender_info", rid)
|
|
return
|
|
try:
|
|
self._prefetch_futures[rid] = self._prefetch_executor.submit(
|
|
self._prefetch_payload, rid, sender_info, target_device
|
|
)
|
|
except Exception:
|
|
logger.exception("Failed to submit KV prefetch for %s", rid)
|
|
|
|
def _prefetch_payload(
|
|
self,
|
|
request_id: str,
|
|
sender_info: dict[str, Any] | None,
|
|
target_device: torch.device | None,
|
|
) -> tuple[dict[str, Any] | None, int]:
|
|
"""Bg-thread body: get + deserialize + H2D on the dedicated ``_bg_copy_stream``.
|
|
|
|
Raises on failure (payload may be consumed → no sync retry).
|
|
"""
|
|
try:
|
|
on_device = target_device is not None and target_device.type != "cpu"
|
|
if on_device:
|
|
# bg thread doesn't inherit the main thread's current device.
|
|
torch.accelerator.set_device_index(target_device.index)
|
|
if self._bg_copy_stream is None:
|
|
assert self._prefetch_executor._max_workers == 1
|
|
self._bg_copy_stream = current_omni_platform.Stream()
|
|
with current_omni_platform.stream(self._bg_copy_stream):
|
|
data, size = self.receive_kv_cache_for_request(
|
|
request_id, target_device=target_device, sender_info=sender_info
|
|
)
|
|
else:
|
|
data, size = self.receive_kv_cache_for_request(
|
|
request_id, target_device=target_device, sender_info=sender_info
|
|
)
|
|
except Exception:
|
|
logger.exception("KV prefetch payload failed for %s (payload may be lost)", request_id)
|
|
raise
|
|
|
|
if data is None:
|
|
return None, 0
|
|
return data, size
|
|
|
|
def consume_prefetched_kv(self, req: Any) -> tuple[dict[str, Any] | None, int]:
|
|
"""Consume a prefetched KV payload; (None, 0) on miss/recoverable failure.
|
|
|
|
Raises ``KVPrefetchConsumeError`` when consumed but post-get failed (no fallback).
|
|
"""
|
|
request_id = self._resolve_request_id(req)
|
|
if not request_id:
|
|
return None, 0
|
|
|
|
fut = self._prefetch_futures.pop(request_id, None)
|
|
# Serial mode: any other request still in the table is an orphan — drop it.
|
|
for stale_rid in list(self._prefetch_futures):
|
|
self._discard_future(stale_rid)
|
|
|
|
if fut is None:
|
|
return None, 0
|
|
|
|
try:
|
|
return fut.result()
|
|
except KVPrefetchConsumeError:
|
|
logger.exception("KV load failed for %s (payload consumed, cannot retry)", request_id)
|
|
raise
|
|
except Exception:
|
|
logger.exception("KV load failed for %s; falling back to sync receive", request_id)
|
|
return None, 0
|
|
|
|
def _discard_future(self, request_id: str) -> None:
|
|
"""Cancel an unstarted prefetch or attach a callback to drop a running one."""
|
|
fut = self._prefetch_futures.pop(request_id, None)
|
|
if fut is None:
|
|
return
|
|
if not fut.cancel():
|
|
fut.add_done_callback(_drop_prefetch_result)
|
|
|
|
def shutdown_prefetch(self) -> None:
|
|
"""Cancel pending prefetches and stop the executor (call on teardown)."""
|
|
for rid in list(self._prefetch_futures):
|
|
self._discard_future(rid)
|
|
if self._prefetch_executor is not None:
|
|
try:
|
|
self._prefetch_executor.shutdown(wait=True, cancel_futures=True)
|
|
except Exception:
|
|
logger.exception("Failed to shut down KV prefetch executor")
|
|
self._prefetch_executor = None
|
|
|
|
@staticmethod
|
|
def _record_stream_for_prefetched(data: dict[str, Any]) -> None:
|
|
"""``record_stream(current_stream)`` on GPU tensors in *data*.
|
|
|
|
Protects prefetch-thread tensors from allocator reuse on LEADER collective paths.
|
|
"""
|
|
if not isinstance(data, dict) or "layer_blocks" not in data:
|
|
return
|
|
if current_omni_platform.get_device_count() < 1:
|
|
return
|
|
current_stream = current_omni_platform.current_stream()
|
|
layer_blocks = data["layer_blocks"]
|
|
for cache_list in (layer_blocks.get("key_cache", []), layer_blocks.get("value_cache", [])):
|
|
for tensor in cache_list:
|
|
if isinstance(tensor, torch.Tensor) and tensor.device.type != "cpu":
|
|
tensor.record_stream(current_stream)
|
|
|
|
@torch.inference_mode()
|
|
def receive_kv_cache_for_request(
|
|
self,
|
|
request_id: str,
|
|
target_device: torch.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).
|
|
"""
|
|
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]] = {}
|
|
deferred_memory: list[Any] = []
|
|
|
|
# Per-call sender_info takes precedence over instance fields.
|
|
if sender_info is not None:
|
|
base_host, base_port = self._resolve_sender_base(sender_info)
|
|
else:
|
|
base_host, base_port = self._sender_base_host, self._sender_base_zmq_port
|
|
|
|
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):
|
|
rank_metadata: dict[str, Any] | None = None
|
|
build_meta = from_rank is not None or sender_info is not None
|
|
if build_meta and base_host and base_port is not None:
|
|
rank_metadata = {
|
|
"source_host": base_host,
|
|
"source_port": base_port + (from_rank or 0) * KV_RANK_PORT_STRIDE,
|
|
}
|
|
elif sender_info is not None:
|
|
logger.error(
|
|
"KV receive for %s: unresolved per-call sender_info (host=%s, port=%s); aborting",
|
|
request_id,
|
|
base_host,
|
|
base_port,
|
|
)
|
|
return None, 0
|
|
|
|
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
|
|
|
|
if hasattr(raw_data, "tensor") and hasattr(raw_data, "release"):
|
|
try:
|
|
buf_tensor = raw_data.tensor
|
|
if getattr(buf_tensor.device, "type", "cpu") != "cpu":
|
|
data = KVCacheTransferData.from_bytes_device(buf_tensor)
|
|
raw_data.release()
|
|
else:
|
|
data = KVCacheTransferData.from_bytes(memoryview(buf_tensor.numpy()))
|
|
deferred_memory.append(raw_data)
|
|
except Exception as e:
|
|
logger.error("Failed to deserialize KV cache from ManagedBuffer: %s", e)
|
|
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)
|
|
|
|
needs_clone = bool(deferred_memory)
|
|
try:
|
|
if isinstance(data, dict) and "layer_blocks" in data:
|
|
layer_blocks = data["layer_blocks"]
|
|
cache_lists = [
|
|
layer_blocks.get("key_cache", []),
|
|
layer_blocks.get("value_cache", []),
|
|
]
|
|
for cache_list in cache_lists:
|
|
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()
|
|
elif needs_clone:
|
|
cache_list[i] = tensor.clone()
|
|
except Exception as exc:
|
|
logger.exception("Failed to detach/move KV cache tensors for %s", request_id)
|
|
raise KVPrefetchConsumeError(
|
|
f"Post-get processing failed for {request_id} (payload already consumed)"
|
|
) from exc
|
|
|
|
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 KVPrefetchConsumeError:
|
|
raise
|
|
except Exception:
|
|
logger.exception("Error receiving KV cache for %s", request_id)
|
|
return None, 0
|
|
finally:
|
|
self._release_pool_buffers(deferred_memory)
|
|
|
|
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 _extract_primary_kv_obj(self, kv_payload: dict[str, Any]) -> Any | None:
|
|
"""Return the primary KV object (``SimpleNamespace(key_cache, value_cache)``)."""
|
|
for key in ("past_key_values", "sp.past_key_values"):
|
|
obj = kv_payload.get(key)
|
|
if obj is not None and hasattr(obj, "key_cache"):
|
|
return obj
|
|
return None
|
|
|
|
def _pack_kv_payload(self, kv_payload: dict[str, Any] | None, device: torch.device) -> torch.Tensor | None:
|
|
"""Pack a KV payload into one uint8 tensor on device
|
|
Layout: ``[4B side_len][msgpack side_bytes][gpu blob]``.
|
|
Returns None on failure."""
|
|
if not kv_payload:
|
|
return None
|
|
primary = self._extract_primary_kv_obj(kv_payload)
|
|
if primary is None:
|
|
return None
|
|
|
|
key_cache = _move_to_device(getattr(primary, "key_cache", []) or [], device)
|
|
value_cache = _move_to_device(getattr(primary, "value_cache", []) or [], device)
|
|
if not any(isinstance(t, torch.Tensor) for t in key_cache):
|
|
return None
|
|
td = KVCacheTransferData(
|
|
request_id="",
|
|
layer_blocks={"key_cache": key_cache, "value_cache": value_cache},
|
|
block_ids=[],
|
|
metadata={},
|
|
)
|
|
try:
|
|
blob = td.to_gpu_tensor()
|
|
except Exception:
|
|
logger.exception("D2D: failed to pack primary KV; falling back to object transport")
|
|
return None
|
|
|
|
# Side-payload: replace the heavy primary with a sentinel, msgpack the rest.
|
|
side = dict(kv_payload)
|
|
for key in ("past_key_values", "sp.past_key_values"):
|
|
if side.get(key) is primary:
|
|
side[key] = _KV_PLACEHOLDER
|
|
|
|
try:
|
|
side_bytes = OmniSerializer.serialize(side)
|
|
except Exception:
|
|
logger.exception("D2D: failed to serialize side-payload; falling back to object transport")
|
|
return None
|
|
side_len = len(side_bytes)
|
|
side_tensor = torch.frombuffer(bytearray(side_bytes), dtype=torch.uint8).to(device)
|
|
|
|
# [4B side_len][side_bytes][blob]
|
|
header = torch.tensor([side_len], dtype=torch.int32, device=device).view(torch.uint8)
|
|
combined = torch.cat([header, side_tensor, blob])
|
|
return combined
|
|
|
|
def _unpack_kv_payload(self, combined: torch.Tensor) -> dict[str, Any] | None:
|
|
"""Inverse of :meth:`_pack_kv_payload`."""
|
|
if combined is None or combined.numel() == 0:
|
|
return None
|
|
|
|
offset = 4
|
|
side_len = int(combined[:offset].cpu().view(torch.int32).item())
|
|
if side_len <= 0 and combined.numel() > offset:
|
|
# Corrupt/zeroed blob (sender-side failure) → treat as miss; raising
|
|
# would crash only the receiver and desync the group.
|
|
logger.warning(
|
|
"D2D: corrupt KV payload (side_len=%d, total_bytes=%d, likely sender-side failure); treating as miss",
|
|
side_len,
|
|
combined.numel(),
|
|
)
|
|
return None
|
|
side_bytes = combined[offset : offset + side_len].cpu().numpy().tobytes()
|
|
offset += side_len
|
|
side: dict[str, Any] = OmniSerializer.deserialize(side_bytes)
|
|
|
|
# Rebuild primary KV from the trailing blob and restore it in place.
|
|
primary: Any = None
|
|
blob = combined[offset:]
|
|
if blob.numel() > 0:
|
|
from types import SimpleNamespace
|
|
|
|
data = KVCacheTransferData.from_bytes_device(blob)
|
|
lb = data["layer_blocks"]
|
|
primary = SimpleNamespace(key_cache=lb.get("key_cache"), value_cache=lb.get("value_cache"))
|
|
|
|
for key in ("past_key_values", "sp.past_key_values"):
|
|
if side.get(key) == _KV_PLACEHOLDER:
|
|
side[key] = primary
|
|
return side
|
|
|
|
def _broadcast_kv_payload(
|
|
self, group: Any, kv_payload: dict[str, Any] | None, device: torch.device, src: int = 0
|
|
) -> dict[str, Any] | None:
|
|
"""Broadcast KV payload via packed tensor; falls back to ``broadcast_object`` on packing failure.
|
|
|
|
The packed-tensor path uses a two-phase collective (size then data).
|
|
Both phases must always be called together — a rank entering only
|
|
one phase will deadlock the group.
|
|
"""
|
|
is_src = group.rank_in_group == src
|
|
if is_src:
|
|
combined = self._pack_kv_payload(kv_payload, device)
|
|
if combined is None:
|
|
group.broadcast_object(kv_payload, src=src) # packing failed -> object
|
|
return kv_payload
|
|
group.broadcast_object(int(combined.numel()), src=src)
|
|
group.broadcast(combined, src=src)
|
|
return kv_payload
|
|
|
|
size_obj = group.broadcast_object(None, src=src)
|
|
if size_obj is None or not isinstance(size_obj, int):
|
|
return size_obj # src fell back to object
|
|
combined = torch.empty(size_obj, dtype=torch.uint8, device=device)
|
|
group.broadcast(combined, src=src)
|
|
return self._unpack_kv_payload(combined)
|
|
|
|
def _send_kv_payload(self, group: Any, kv_payload: dict[str, Any] | None, dst: int, device: torch.device) -> None:
|
|
"""Point-to-point send; falls back to ``send_object`` on packing failure."""
|
|
combined = self._pack_kv_payload(kv_payload, device)
|
|
if combined is None:
|
|
group.send_object(kv_payload, dst)
|
|
return
|
|
size = int(combined.numel())
|
|
group.send_object(size, dst)
|
|
try:
|
|
group.send(combined, dst)
|
|
except Exception:
|
|
logger.exception("D2D: send(combined) to rank %d failed; sending zero sentinel", dst)
|
|
combined.zero_()
|
|
group.send(combined, dst)
|
|
raise
|
|
|
|
def _recv_kv_payload(self, group: Any, src: int, device: torch.device) -> dict[str, Any] | None:
|
|
"""Inverse of :meth:`_send_kv_payload`."""
|
|
size_or_obj = group.recv_object(src)
|
|
if not isinstance(size_or_obj, int):
|
|
return size_or_obj
|
|
combined = group.recv(torch.Size([size_or_obj]), torch.uint8, src)
|
|
combined = combined.to(device)
|
|
return self._unpack_kv_payload(combined)
|
|
|
|
def _tp_local_receive_consensus(self, req: Any, received: bool) -> bool:
|
|
"""Identical-KV-state verdict across pure-TP ranks (LOCAL role).
|
|
|
|
In pure TP every rank fetches its own KV shard with an independent
|
|
timeout and ``distribute_kv_cache`` is a no-op, so a single rank's
|
|
silent miss sends only that rank into the pipeline's no-KV prefill —
|
|
it then enters TP collectives its peers never join, and the NCCL
|
|
watchdog kills the worker (#5627). Each rank reports a signature of
|
|
its applied KV state — the primary-receive flag plus the set of
|
|
applied payload roles, so partially-received CFG companion KV
|
|
(``cfg_text``/``cfg_img``) diverging between ranks is caught too. If
|
|
the signatures differ, every rank drops its KV so all take the same
|
|
fallback path together; identical-but-partial state is kept, since
|
|
symmetric absence is consistent (e.g. Bagel's empty unconditional
|
|
branch).
|
|
|
|
The exchange runs over the stage's tensor-parallel group only: those
|
|
are the ranks executing this request in lockstep. ``world`` would
|
|
also contain DP replicas / PP stages working on different requests —
|
|
mixing their flags corrupts the verdict, and a replica skipping the
|
|
exchange (dummy request) while another enters it would deadlock.
|
|
Within the TP group the skip guards below depend only on topology,
|
|
config, and the request id — identical on every TP rank — so ranks
|
|
always agree on whether the exchange happens.
|
|
"""
|
|
pt = self.topo_config
|
|
group = pt.tp_group
|
|
if not (pt.is_local and pt.tp_active) or group is None or pt.tp_size <= 1:
|
|
return received
|
|
if not self.config.need_recv_cache:
|
|
return received
|
|
request_id = self._resolve_request_id(req)
|
|
if request_id is None or OmniDiffusionRequest.is_dummy_run_request_id(request_id):
|
|
return received
|
|
|
|
signature = (bool(received), tuple(sorted(self._collect_request_kv_payload(req))))
|
|
if group.rank_in_group == 0:
|
|
signatures = [signature]
|
|
signatures.extend(group.recv_object(src=src) for src in range(1, pt.tp_size))
|
|
consistent = all(sig == signature for sig in signatures)
|
|
group.broadcast_object(consistent, src=0)
|
|
else:
|
|
group.send_object(signature, dst=0)
|
|
consistent = bool(group.broadcast_object(None, src=0))
|
|
|
|
if not consistent:
|
|
logger.warning(
|
|
"KV state for %s diverged across TP ranks (this rank: received=%s, roles=%s); "
|
|
"dropping KV on every rank so all take the no-KV fallback together "
|
|
"instead of desyncing TP collectives",
|
|
request_id,
|
|
signature[0],
|
|
list(signature[1]),
|
|
)
|
|
self._clear_request_kv_payload(req)
|
|
return False
|
|
return received
|
|
|
|
def receive_multi_kv_cache_distributed(
|
|
self,
|
|
req: Any,
|
|
cfg_kv_collect_func: Callable | None = None,
|
|
target_device: torch.device | None = None,
|
|
) -> bool:
|
|
"""Synchronous receive: fetch → distribute → apply (FOLLOWER skips fetch)."""
|
|
received = False
|
|
if not self.topo_config.is_follower:
|
|
received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device)
|
|
received = self._tp_local_receive_consensus(req, received)
|
|
kv_payload = self.distribute_kv_cache(req, target_device, received=received)
|
|
if kv_payload is not None:
|
|
self._apply_request_kv_payload(req, kv_payload, target_device)
|
|
return received or kv_payload is not None
|
|
|
|
def consume_and_distribute_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool:
|
|
"""Consume prefetched KV → apply → distribute; sync-receive on miss (no retry on consume-error)."""
|
|
received = False
|
|
payload_consumed = False
|
|
if self._async_prefetch and not self.topo_config.is_follower:
|
|
try:
|
|
data, _ = self.consume_prefetched_kv(req)
|
|
if data is not None:
|
|
self._record_stream_for_prefetched(data)
|
|
self.apply_kv_cache_to_request(req, data)
|
|
received = True
|
|
except KVPrefetchConsumeError:
|
|
logger.error(
|
|
"KV prefetch consumed payload for %s but post-get failed; request cannot recover",
|
|
self._resolve_request_id(req),
|
|
)
|
|
payload_consumed = True
|
|
# Only fall back to sync receive on miss, not when payload was
|
|
# already consumed from the connector (sync receive would block
|
|
# until timeout as the data is gone).
|
|
if not received and not self.topo_config.is_follower and not payload_consumed:
|
|
logger.debug("KV prefetch miss for %s; falling back to sync receive", self._resolve_request_id(req))
|
|
received = self.receive_multi_kv_cache(req, None, target_device)
|
|
received = self._tp_local_receive_consensus(req, received)
|
|
kv_payload = self.distribute_kv_cache(req, target_device, received=received)
|
|
if kv_payload is not None:
|
|
self._apply_request_kv_payload(req, kv_payload, target_device)
|
|
return received or kv_payload is not None
|
|
|
|
def _cfg_scatter(
|
|
self,
|
|
group: Any,
|
|
payloads: list[Any],
|
|
cfg_size: int,
|
|
device: torch.device | None,
|
|
request_id: str | None = None,
|
|
) -> None:
|
|
"""Scatter one payload per CFG follower."""
|
|
for dst in range(1, cfg_size):
|
|
try:
|
|
self._send_kv_payload(group, payloads[dst], dst, device)
|
|
except Exception:
|
|
logger.warning("CFG scatter to rank %d failed for %s; that rank runs without KV", dst, request_id)
|
|
continue
|
|
|
|
def distribute_kv_cache(
|
|
self, req: Any, target_device: torch.device | None = None, *, received: bool = False
|
|
) -> dict[str, Any] | None:
|
|
"""Distribute KV LEADER→FOLLOWER. No-op for LOCAL/pure TP."""
|
|
pt = self.topo_config
|
|
|
|
# LOCAL: no distribution.
|
|
if pt.is_local:
|
|
return None
|
|
|
|
device = target_device if target_device is not None else torch.device("cpu")
|
|
# TP + CFG/SP
|
|
if pt.tp_active and (pt.cfg_active or pt.sp_active):
|
|
kv_payload: dict[str, object] | None = None
|
|
if pt.is_leader:
|
|
if pt.cfg_active:
|
|
payloads = (
|
|
self._build_cfg_rank_local_payloads(req, pt.cfg_size) if received else [None] * pt.cfg_size
|
|
)
|
|
kv_payload = payloads[0]
|
|
self._cfg_scatter(pt.cfg_group, payloads, pt.cfg_size, device, self._resolve_request_id(req))
|
|
elif pt.sp_active:
|
|
kv_payload = self._collect_request_kv_payload(req) if received else None
|
|
if pt.sp_active and pt.sp_group is not None:
|
|
kv_payload = self._broadcast_kv_payload(pt.sp_group, kv_payload, device, src=0)
|
|
elif pt.is_follower:
|
|
if pt.cfg_active and pt.sp_rank == 0:
|
|
kv_payload = self._recv_kv_payload(pt.cfg_group, 0, device)
|
|
if pt.sp_active and pt.sp_group is not None:
|
|
kv_payload = self._broadcast_kv_payload(pt.sp_group, kv_payload, device, src=0)
|
|
else:
|
|
logger.error("distribute_kv_cache: unexpected role %s in TP+CFG/SP path", pt.role)
|
|
return None
|
|
return kv_payload or None
|
|
|
|
# TP inactive, world broadcast
|
|
kv_payload = None
|
|
if pt.is_leader and received:
|
|
kv_payload = self._collect_request_kv_payload(req)
|
|
kv_payload = self._broadcast_kv_payload(pt.world, kv_payload, device, src=0)
|
|
return kv_payload or None
|
|
|
|
|
|
def _drop_prefetch_result(fut: Any) -> None:
|
|
try:
|
|
if not fut.cancelled():
|
|
fut.result()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
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
|