vllm-omni/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

# 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